dean-research-tools 0.1.0__tar.gz

dean_research_tools-0.1.0/LICENSE

@@ -0,0 +1,7 @@
+Copyright 2026 Dean MacGregor
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

dean_research_tools-0.1.0/PKG-INFO

@@ -0,0 +1,33 @@
+Metadata-Version: 2.4
+Name: dean_research_tools
+Version: 0.1.0
+Summary: Retriever tools package for research task and content lookup
+Author-email: Dean MacGregor <powertrading121@gmail.com>
+License: MIT
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Software Development :: Libraries
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: cohere>=5.0.0
+Requires-Dist: langchain-core>=0.3.0
+Requires-Dist: orjson>=3.9.0
+Requires-Dist: pgvector>=0.3.0
+Requires-Dist: psycopg[binary]>=3.2.0
+Requires-Dist: pydantic>=2.7.0
+Provides-Extra: dev
+Requires-Dist: build>=1.2.2; extra == "dev"
+Requires-Dist: pytest>=8.2.0; extra == "dev"
+Requires-Dist: twine>=5.1.1; extra == "dev"
+Dynamic: license-file
+
+# dean_research_tools
+
+I made this library so I could use it in different environments without copying and pasting. I'm not intending to have it be used by a lot of people.
+
+If a lot of people want to use it then that's fine, there's no secret sauce, but you're on your own.

dean_research_tools-0.1.0/README.md

@@ -0,0 +1,5 @@
+# dean_research_tools
+
+I made this library so I could use it in different environments without copying and pasting. I'm not intending to have it be used by a lot of people.
+
+If a lot of people want to use it then that's fine, there's no secret sauce, but you're on your own.

dean_research_tools-0.1.0/pyproject.toml

@@ -0,0 +1,48 @@
+[build-system]
+requires = ["setuptools>=69", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "dean_research_tools"
+version = "0.1.0"
+description = "Retriever tools package for research task and content lookup"
+readme = "README.md"
+requires-python = ">=3.11"
+license = { text = "MIT" }
+authors = [
+  { name = "Dean MacGregor", email="powertrading121@gmail.com"}
+]
+
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Developers",
+  "License :: OSI Approved :: MIT License",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Topic :: Software Development :: Libraries",
+]
+dependencies = [
+  "cohere>=5.0.0",
+  "langchain-core>=0.3.0",
+  "orjson>=3.9.0",
+  "pgvector>=0.3.0",
+  "psycopg[binary]>=3.2.0",
+  "pydantic>=2.7.0",
+]
+
+[project.optional-dependencies]
+dev = [
+  "build>=1.2.2",
+  "pytest>=8.2.0",
+  "twine>=5.1.1",
+]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]

dean_research_tools-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

dean_research_tools-0.1.0/src/dean_research_tools/__init__.py

@@ -0,0 +1,6 @@
+"""Retriever tools package."""
+
+from dean_research_tools.config import load_settings
+from dean_research_tools.retriever import PGTools
+
+__all__ = ["load_settings", "PGTools"]

dean_research_tools-0.1.0/src/dean_research_tools/config.py

@@ -0,0 +1,94 @@
+"""Configuration for the researcher tools package."""
+
+from __future__ import annotations
+
+import os
+
+from pydantic import BaseModel, ConfigDict, SecretStr
+
+
+class Settings(BaseModel):
+    """Environment-backed settings for DB and Azure OpenAI."""
+
+    model_config = ConfigDict(frozen=True)
+
+    db_host: str
+    db_port: int
+    db_name: str
+    db_user: str
+    db_password: str
+    azure_openai_api_key: SecretStr
+    azure_embedding_endpoint: str
+    azure_openai_embedding_deployment: str
+
+    @property
+    def db_dsn(self) -> str:
+        return (
+            f"host={self.db_host} "
+            f"port={self.db_port} "
+            f"dbname={self.db_name} "
+            f"user={self.db_user} "
+            f"password={self.db_password} "
+            f"connect_timeout=300 "
+            f"keepalives=1 "
+            f"keepalives_idle=30 "
+            f"keepalives_interval=10 "
+            f"keepalives_count=5"
+        )
+
+
+def _require_env(name: str) -> str:
+    value = os.getenv(name, "").strip()
+    if not value:
+        raise ValueError(f"Missing required environment variable: {name}")
+    return value
+
+
+def _research_env(name: str) -> str:
+    return f"RESEARCH_{name.upper()}"
+
+
+def load_settings(
+    *,
+    db_host: str | None = None,
+    db_port: int | None = None,
+    db_name: str | None = None,
+    db_user: str | None = None,
+    db_password: str | None = None,
+    azure_openai_api_key: str | SecretStr | None = None,
+    azure_embedding_endpoint: str | None = None,
+    azure_openai_embedding_deployment: str | None = None,
+) -> Settings:
+    """Load settings from keyword overrides or RESEARCH_* environment variables."""
+
+    resolved_db_host = db_host or _require_env(_research_env("db_host"))
+    resolved_db_port = db_port or int(os.getenv(_research_env("db_port"), "5432"))
+    resolved_db_name = db_name or _require_env(_research_env("db_name"))
+    resolved_db_user = db_user or _require_env(_research_env("db_user"))
+    resolved_db_password = db_password or _require_env(_research_env("db_password"))
+    resolved_embedding_endpoint = azure_embedding_endpoint or _require_env(
+        _research_env("azure_embedding_endpoint")
+    )
+    resolved_embedding_deployment = azure_openai_embedding_deployment or _require_env(
+        _research_env("azure_openai_embedding_deployment")
+    )
+
+    if azure_openai_api_key is None:
+        resolved_api_key = SecretStr(
+            _require_env(_research_env("azure_openai_api_key"))
+        )
+    elif isinstance(azure_openai_api_key, SecretStr):
+        resolved_api_key = azure_openai_api_key
+    else:
+        resolved_api_key = SecretStr(azure_openai_api_key)
+
+    return Settings(
+        db_host=resolved_db_host,
+        db_port=resolved_db_port,
+        db_name=resolved_db_name,
+        db_user=resolved_db_user,
+        db_password=resolved_db_password,
+        azure_openai_api_key=resolved_api_key,
+        azure_embedding_endpoint=resolved_embedding_endpoint,
+        azure_openai_embedding_deployment=resolved_embedding_deployment,
+    )

dean_research_tools-0.1.0/src/dean_research_tools/embeddings.py

@@ -0,0 +1,42 @@
+"""Embedding model construction for retrieval queries."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from cohere import AsyncClientV2
+
+from dean_research_tools.config import Settings
+
+
+class EmbeddingsModel:
+    def __init__(self, settings: Settings):
+        self.settings = settings
+        self.co_client = AsyncClientV2(
+            api_key=self.settings.azure_openai_api_key.get_secret_value(),
+            base_url=self.settings.azure_embedding_endpoint,
+        )
+
+    async def embed_texts(
+        self,
+        texts: list[str],
+        input_type: Literal["search_document", "search_query"] = "search_query",
+    ) -> list[list[float]] | None:
+        results = []
+        chunked = self.chunk_list(texts)
+        for chunk in chunked:
+            result = await self.co_client.embed(
+                model=self.settings.azure_openai_embedding_deployment,
+                texts=chunk,
+                input_type=input_type,
+                embedding_types=["float"],
+            )
+            flt = result.embeddings.float_
+            if flt is None:
+                return None
+            results.extend(flt)
+        return results
+
+    @staticmethod
+    def chunk_list(lst, size=96):
+        return [lst[i : i + size] for i in range(0, len(lst), size)]

dean_research_tools-0.1.0/src/dean_research_tools/models.py

@@ -0,0 +1,90 @@
+from pydantic import BaseModel, Field
+
+
+class SemanticContentSearchInput(BaseModel):
+    query: str = Field(
+        description="The search query string to be embedded with cohere and used for vector similarity search on text content."
+    )
+    task_id: int | None = Field(
+        default=None,
+        description=(
+            "Optional task scope. Omit this for corpus-wide searches. "
+            "Use it when the question is restricted to a browser task."
+        ),
+    )
+    url_part: str | None = Field(
+        default=None,
+        description=(
+            "Optional URL scope. Omit this for corpus-wide searches. "
+            "It will filter for content whose URL contains this substring."
+        ),
+    )
+    doc_id: int | None = Field(
+        default=None,
+        description="Optional doc_id to filter the search results. If provided, only content from this document will be considered.",
+    )
+    top_k: int = Field(
+        default=5,
+        description="The number of top results to return (default is 5).",
+    )
+    min_score: float = Field(
+        default=0,
+        description="Minimum score threshold for filtering results (between 0.0 and 1.0)",
+    )
+
+
+class SemanticTaskSearchInput(BaseModel):
+    query: str = Field(
+        description="The search query string to be embedded with cohere and used for vector similarity search on previous browser tasks."
+    )
+    top_k: int = Field(
+        default=5,
+        description="The number of top results to return (default is 5).",
+    )
+    min_score: float = Field(
+        default=0,
+        description="Minimum score threshold for filtering results (between 0.0 and 1.0)",
+    )
+
+
+class KeywordContentSearchInput(BaseModel):
+    keywords: list[str] = Field(
+        description="The keywords to search for in the text_chunk of the browser_content table. The search will be case-insensitive and will match any text_chunk that contains any of the keywords."
+    )
+    task_id: int | None = Field(
+        default=None,
+        description="The task_id by which to filter the search results. If this isn't provided then url_part must be provided.",
+    )
+    url_part: str | None = Field(
+        default=None,
+        description=(
+            "Optional URL scope. Omit this for corpus-wide searches. "
+            "It will filter for content whose URL contains this substring."
+        ),
+    )
+    doc_id: int | None = Field(
+        default=None,
+        description="Optional doc_id to further filter the search results.",
+    )
+    pages: list[int] | None = Field(
+        default=None,
+        description="Optional list of page numbers to filter the search results. Can only be provided if doc_id is also provided.",
+    )
+    top_k: int = Field(
+        default=5,
+        description="The number of top results to return (default is 5).",
+    )
+
+
+class GetTaskInput(BaseModel):
+    task_id: int = Field(
+        description="The task_id from which to retrieve the associated task details."
+    )
+
+
+class BrowserUseInput(BaseModel):
+    objective: str = Field(
+        description="""Invoke an agent using the browser_use library to use a real browser to collect content. The agent needs to be given a detailed objective of what content
+        it needs to ingest. The agent will then use the browser to collect content and ingest it into the database. The agent will return a task_id which can be used
+        to retrieve the content with your existing tools."""
+    )

dean_research_tools-0.1.0/src/dean_research_tools/retriever.py

@@ -0,0 +1,346 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from inspect import ismethod
+from typing import TYPE_CHECKING, Literal, Sequence, get_args
+from weakref import WeakSet
+
+import orjson
+from langchain_core.tools import StructuredTool
+from pgvector import Vector
+from pgvector.psycopg import register_vector_async
+from psycopg import AsyncConnection
+from psycopg.rows import dict_row
+from psycopg.sql import SQL
+
+import dean_research_tools.models as models
+from dean_research_tools.config import Settings, load_settings
+from dean_research_tools.embeddings import EmbeddingsModel
+
+_registered_pgvector = WeakSet()
+
+
+def _to_pascal_inputs(s: str) -> str:
+    return "".join(word.capitalize() for word in s.split("_")) + "Input"
+
+
+async def ensure_pgvector_registered(conn: AsyncConnection) -> None:
+    if conn in _registered_pgvector:
+        return
+    await register_vector_async(conn)
+    _registered_pgvector.add(conn)
+
+
+if TYPE_CHECKING:
+    from psycopg_pool import AsyncConnectionPool
+
+
+AVAILABLE_TOOLS = Literal[
+    "semantic_content_search",
+    "keyword_content_search",
+    "get_task",
+    "semantic_task_search",
+    "browser_use",
+]
+
+
+class PGTools:
+    def __init__(
+        self,
+        *,
+        include: Sequence[AVAILABLE_TOOLS] | None = None,
+        exclude: Sequence[AVAILABLE_TOOLS] | None = None,
+        settings: Settings | None = None,
+        conn: AsyncConnection | None = None,
+        pool: AsyncConnectionPool | None = None,
+    ):
+        self.settings = settings or load_settings()
+        self.include = include
+        self.exclude = exclude
+        if conn is not None and pool is not None:
+            raise ValueError("Cannot specify both conn and pool")
+        self.conn = conn
+        self.pool = pool
+        self._close_conn = conn is None
+        self.embeddings = EmbeddingsModel(self.settings)
+
+    @asynccontextmanager
+    async def _get_cur(self):
+        if self.conn is not None:
+            await ensure_pgvector_registered(self.conn)
+            async with self.conn.cursor(row_factory=dict_row) as cur:
+                yield cur
+        elif self.pool is not None:
+            async with self.pool.connection() as conn:
+                await ensure_pgvector_registered(conn)
+                async with conn.cursor(row_factory=dict_row) as cur:
+                    yield cur
+        else:
+            raise ValueError("No connection or pool provided")
+
+    def _get_all_tools(self) -> list[StructuredTool]:
+        tools = []
+        for tool_name in dir(self):
+            if tool_name[0] == "_" or not ismethod(getattr(self, tool_name)):
+                continue
+            kwargs = {
+                "coroutine": getattr(self, tool_name),
+                "args_schema": getattr(models, _to_pascal_inputs(tool_name)),
+                "description": getattr(self, tool_name).__doc__.replace(
+                    "RETURN_DIRECT", ""
+                ),
+            }
+            if "RETURN_DIRECT" in getattr(self, tool_name).__doc__:
+                kwargs["return_direct"] = True
+            tools.append(StructuredTool.from_function(**kwargs))
+        return tools
+
+    def _get_tools(
+        self,
+    ) -> list[StructuredTool]:
+        tools_to_get: tuple[AVAILABLE_TOOLS] = tuple([])
+        if self.include is not None and self.exclude is not None:
+            raise ValueError("Cannot specify both include and exclude")
+        elif self.include is None and self.exclude is None:
+            tools_to_get: tuple[AVAILABLE_TOOLS] = get_args(AVAILABLE_TOOLS)
+        elif self.include is not None:
+            tools_to_get = tuple(
+                x for x in get_args(AVAILABLE_TOOLS) if x in self.include
+            )
+        elif self.exclude is not None:
+            tools_to_get = tuple(
+                t for t in get_args(AVAILABLE_TOOLS) if t not in self.exclude
+            )
+
+        return [t for t in self._get_all_tools() if t.name in tools_to_get]
+
+    async def __aenter__(self) -> list[StructuredTool]:
+        if self.conn is None and self.pool is None:
+            self.conn = await AsyncConnection.connect(self.settings.db_dsn)
+            await ensure_pgvector_registered(self.conn)
+            await self.conn.set_autocommit(True)
+
+        return self._get_tools()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self.conn is not None and self._close_conn:
+            await self.conn.close()
+            self.conn = None
+
+    async def semantic_content_search(
+        self,
+        query: str,
+        task_id: int | None = None,
+        url_part: str | None = None,
+        doc_id: int | None = None,
+        top_k: int = 5,
+        min_score: float = 0,
+    ) -> str:
+        """Search for relevant content in the browser_content table using vector similarity."""
+        if self.conn is None:
+            raise RuntimeError("PGTools must be used as an async context manager")
+        if not query.strip():
+            raise ValueError("query must not be empty")
+        if top_k < 1:
+            raise ValueError("top_k must be >= 1")
+        if not (0.0 <= min_score <= 1.0):
+            raise ValueError("min_score must be between 0.0 and 1.0")
+
+        query_embedding = await self.embeddings.embed_texts(
+            [query], input_type="search_query"
+        )
+        if query_embedding is None:
+            raise ValueError("Failed to generate query embedding")
+        query_embedding = query_embedding[0]
+
+        select = SQL("""
+            WITH emb AS (
+                SELECT %s::vector AS embedding
+                )
+            SELECT
+                bt.id as task_id,
+                bc.chunk_id,
+                bc.doc_id,
+                (bc.chunk_meta->>'page')::int AS page,
+                (bc.chunk_meta->>'chunk_index')::int AS chunk_index,
+                bc.text_chunk,
+                COALESCE(bd.url, '') AS url,
+                COALESCE(bd.title, '') AS title,
+                (1 - (bc.embedding <=> emb.embedding))::double precision AS score
+            FROM ai_proj.browser_content bc
+            INNER JOIN ai_proj.browser_docs bd using (doc_id)
+            INNER JOIN ai_proj.browser_tasks bt using (id)
+            CROSS JOIN emb
+            """)
+        values: list[Vector | int | str] = [Vector(query_embedding)]
+        wheres = [
+            SQL("""
+            WHERE bc.embedding IS NOT NULL
+              AND bc.text_chunk IS NOT NULL
+            """)
+        ]
+
+        if task_id is not None:
+            wheres.append(SQL("AND bt.id = %s "))
+            values.append(task_id)
+
+        if doc_id is not None:
+            wheres.append(SQL("AND bc.doc_id = %s "))
+            values.append(doc_id)
+
+        if url_part is not None:
+            wheres.append(SQL("AND bd.url ILIKE '%%' || %s || '%%'"))
+            values.append(url_part)
+
+        limit = SQL("""
+            ORDER BY bc.embedding <=> emb.embedding
+            LIMIT %s
+            """)
+        values.append(top_k)
+        sql = select + SQL(" ").join(wheres) + limit
+        async with self.conn.cursor(row_factory=dict_row) as cur:
+            await cur.execute(sql, values)
+            res = await cur.fetchall()
+            return orjson.dumps(res).decode("utf-8")
+
+    async def keyword_content_search(
+        self,
+        keywords: list[str],
+        task_id: int | None = None,
+        url_part: str | None = None,
+        doc_id: int | None = None,
+        pages: list[int] | None = None,
+        top_k: int = 5,
+    ) -> str:
+        """Search for relevant content in the browser_content table using keywords, it is case-insensitive.
+        Must provide either a task_id or url_part to filter results."""
+        if not keywords:
+            raise ValueError("keywords must not be empty")
+        if not task_id and not url_part:
+            raise ValueError("Either task_id or url_part must be provided")
+        if top_k < 1:
+            raise ValueError("top_k must be >= 1")
+        if pages and not doc_id:
+            raise ValueError("doc_id must be provided if pages are specified")
+
+        select = SQL("""
+            SELECT
+                bt.id as task_id,
+                bc.chunk_id,
+                bc.doc_id,
+                (bc.chunk_meta->>'page')::int AS page,
+                (bc.chunk_meta->>'chunk_index')::int AS chunk_index,
+                bc.text_chunk,
+                COALESCE(bd.url, '') AS url,
+                COALESCE(bd.title, '') AS title
+            FROM ai_proj.browser_content bc
+            INNER JOIN ai_proj.browser_docs bd using (doc_id)
+            INNER JOIN ai_proj.browser_tasks bt using (id)
+            """)
+
+        wheres = [
+            SQL("""
+            WHERE bc.text_chunk IS NOT NULL
+              AND EXISTS (
+                  SELECT 1
+                  FROM unnest(%s::text[]) AS kw
+                  WHERE bc.text_chunk ILIKE '%%' || kw || '%%'
+              )
+        """)
+        ]
+        values: list[str | int | list[str] | list[int]] = [keywords]
+
+        if task_id is not None:
+            wheres.append(SQL("AND bt.id = %s"))
+            values.append(task_id)
+        if doc_id is not None:
+            wheres.append(SQL("AND bc.doc_id = %s"))
+            values.append(doc_id)
+            if pages:
+                wheres.append(SQL("AND (bc.chunk_meta->>'page')::int = ANY(%s)"))
+                values.append(pages)
+        if url_part is not None:
+            wheres.append(SQL("AND bd.url ILIKE '%%' || %s || '%%'"))
+            values.append(url_part)
+
+        limit = SQL("""
+                    ORDER BY doc_id, (bc.chunk_meta->>'page')::int, (bc.chunk_meta->>'chunk_index')::int
+                    LIMIT %s
+                    """)
+        values.append(top_k)
+        sql = select + SQL(" ").join(wheres) + limit
+
+        async with self._get_cur() as cur:
+            await cur.execute(sql, values)
+            res = await cur.fetchall()
+            return orjson.dumps(res).decode("utf-8")
+
+    async def get_task(self, task_id: int) -> str:
+        """Retrieve the starting_task for a given task_id from the browser_tasks table.
+        This is the same info from semantic_task_search."""
+
+        sql = """
+            SELECT starting_task
+            FROM ai_proj.browser_tasks
+            WHERE id = %s
+        """
+        async with self._get_cur() as cur:
+            await cur.execute(sql, [task_id])
+            res = await cur.fetchone()
+            return orjson.dumps(res).decode("utf-8")
+
+    async def semantic_task_search(
+        self,
+        query: str,
+        top_k: int = 5,
+        min_score: float = 0,
+    ) -> str:
+        """Search for relevant tasks in the browser_tasks table using vector similarity."""
+        if not query.strip():
+            raise ValueError("query must not be empty")
+        if top_k < 1:
+            raise ValueError("top_k must be >= 1")
+        if not (0.0 <= min_score <= 1.0):
+            raise ValueError("min_score must be between 0.0 and 1.0")
+
+        query_embedding = await self.embeddings.embed_texts(
+            [query], input_type="search_query"
+        )
+        if query_embedding is None:
+            raise ValueError("Failed to generate query embedding")
+        query_embedding = query_embedding[0]
+
+        sql = SQL("""
+            WITH emb AS (
+                SELECT %s::vector AS embedding
+                )
+            SELECT
+                bt.id as task_id,
+                bt.starting_task,
+                (1 - (bt.embedding <=> emb.embedding))::double precision AS score
+            FROM ai_proj.browser_tasks bt
+            CROSS JOIN emb
+            ORDER BY bt.embedding <=> emb.embedding
+            LIMIT %s
+            """)
+        values = [Vector(query_embedding), top_k]
+        async with self._get_cur() as cur:
+            await cur.execute(sql, values)
+            res = await cur.fetchall()
+            return orjson.dumps(res).decode("utf-8")
+
+    async def browser_use(self, objective: str) -> str:
+        """Use a real browser to collect content and ingest it into the database.RETURN_DIRECT"""
+
+        sql = SQL("""
+            INSERT INTO ai_proj.browser_tasks (starting_task)
+            VALUES (%s)
+            RETURNING id as task_id
+        """)
+        values = [objective]
+        # TODO invoking this needs to pause the research agent until the task is completed and then
+        # return the results of the task.
+        async with self._get_cur() as cur:
+            await cur.execute(sql, values)
+            res = await cur.fetchall()
+            return orjson.dumps(res).decode("utf-8")

dean_research_tools-0.1.0/src/dean_research_tools.egg-info/PKG-INFO

@@ -0,0 +1,33 @@
+Metadata-Version: 2.4
+Name: dean_research_tools
+Version: 0.1.0
+Summary: Retriever tools package for research task and content lookup
+Author-email: Dean MacGregor <powertrading121@gmail.com>
+License: MIT
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Software Development :: Libraries
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: cohere>=5.0.0
+Requires-Dist: langchain-core>=0.3.0
+Requires-Dist: orjson>=3.9.0
+Requires-Dist: pgvector>=0.3.0
+Requires-Dist: psycopg[binary]>=3.2.0
+Requires-Dist: pydantic>=2.7.0
+Provides-Extra: dev
+Requires-Dist: build>=1.2.2; extra == "dev"
+Requires-Dist: pytest>=8.2.0; extra == "dev"
+Requires-Dist: twine>=5.1.1; extra == "dev"
+Dynamic: license-file
+
+# dean_research_tools
+
+I made this library so I could use it in different environments without copying and pasting. I'm not intending to have it be used by a lot of people.
+
+If a lot of people want to use it then that's fine, there's no secret sauce, but you're on your own.

dean_research_tools-0.1.0/src/dean_research_tools.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+LICENSE
+README.md
+pyproject.toml
+src/dean_research_tools/__init__.py
+src/dean_research_tools/config.py
+src/dean_research_tools/embeddings.py
+src/dean_research_tools/models.py
+src/dean_research_tools/retriever.py
+src/dean_research_tools.egg-info/PKG-INFO
+src/dean_research_tools.egg-info/SOURCES.txt
+src/dean_research_tools.egg-info/dependency_links.txt
+src/dean_research_tools.egg-info/requires.txt
+src/dean_research_tools.egg-info/top_level.txt
+tests/test_func_args_sync.py

dean_research_tools-0.1.0/src/dean_research_tools.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

dean_research_tools-0.1.0/src/dean_research_tools.egg-info/requires.txt

@@ -0,0 +1,11 @@
+cohere>=5.0.0
+langchain-core>=0.3.0
+orjson>=3.9.0
+pgvector>=0.3.0
+psycopg[binary]>=3.2.0
+pydantic>=2.7.0
+
+[dev]
+build>=1.2.2
+pytest>=8.2.0
+twine>=5.1.1

dean_research_tools-0.1.0/src/dean_research_tools.egg-info/top_level.txt

@@ -0,0 +1 @@
+dean_research_tools

dean_research_tools-0.1.0/tests/test_func_args_sync.py

@@ -0,0 +1,141 @@
+from inspect import isfunction, signature
+from typing import Any, cast, get_args
+
+from langchain_core.tools import ArgsSchema
+from pydantic import SecretStr
+
+from dean_research_tools import PGTools
+from dean_research_tools.config import Settings
+from dean_research_tools.retriever import AVAILABLE_TOOLS
+
+
+def _safe_schema(model: ArgsSchema) -> dict[str, Any]:
+    if isinstance(model, dict):
+        return model
+    attr_names = ["model_json_schema", "schema"]
+    for attr_name in attr_names:
+        if hasattr(model, attr_name):
+            return getattr(model, attr_name)()
+
+    raise ValueError(
+        f"Model {model} does not have a schema method or model_json_schema method"
+    )
+
+
+TYPE_MAP = {
+    "string": "str",
+    "null": "None",
+    "integer": "int",
+    "number": "float",
+    "boolean": "bool",
+    "array": "list",
+    "object": "dict",
+}
+
+
+def _parse_model_arg(model_args: dict[str, Any], arg: str) -> set[str]:
+    if arg not in model_args:
+        raise ValueError(f"Argument {arg} not found in model args")
+    arg_info = model_args[arg]
+    return _parse_arg_info(arg_info)
+
+
+def _parse_arg_info(arg_info: dict[str, Any]) -> set[str]:
+    inner_types = ""
+    if "items" in arg_info:
+        inner_types_raw = sorted(list(_parse_arg_info(arg_info["items"])))
+        inner_types = "[" + ",".join(inner_types_raw) + "]"
+    if "type" in arg_info:
+        arg_type = TYPE_MAP[arg_info["type"]]
+        return {arg_type + inner_types}
+    if "anyOf" in arg_info:
+        types = set()
+        for option in arg_info["anyOf"]:
+            types.update(_parse_arg_info(option))
+        return types
+    raise ValueError(f"Argument info {arg_info} does not have a type or anyOf field")
+
+
+def _alphabatise_inner(annotation: str) -> str:
+    open_indx = annotation.find("[")
+    while open_indx > 0:
+        close_indx = annotation.find("]", open_indx)
+        inner = annotation[open_indx + 1 : close_indx]
+        inner_parts = [p.strip() for p in inner.split(",")]
+        inner_parts.sort()
+        inner_sorted = ",".join(inner_parts)
+        annotation = (
+            annotation[: open_indx + 1] + inner_sorted + annotation[close_indx:]
+        )
+        open_indx = annotation.find("[", close_indx)
+    return annotation
+
+
+def _parse_func_arg(annotation: str) -> set[str]:
+    annotation = _alphabatise_inner(annotation)
+    split_chars = [",", "|"]
+    for char in split_chars:
+        if char in annotation:
+            parts = [p.strip() for p in annotation.split(char)]
+            return {p for p in parts}
+
+    return {annotation}
+
+
+blank = Settings(
+    db_host="a",
+    db_port=1,
+    db_name="b",
+    db_user="c",
+    db_password="d",
+    azure_openai_api_key=SecretStr("e"),
+    azure_embedding_endpoint="f",
+    azure_openai_embedding_deployment="g",
+)
+
+
+def test_get_all_tools():
+    """Tests that the functions returned by _get_all_tools match the functions defined in the PGTools class."""
+    funcs_in_all = set(x.name for x in PGTools(settings=blank)._get_all_tools())
+
+    funcs_in_class = set(
+        x for x in dir(PGTools) if x[0] != "_" and isfunction(getattr(PGTools, x))
+    )
+
+    if funcs_in_class != funcs_in_all:
+        missing_from_class = funcs_in_all - funcs_in_class
+        missing_from_function = funcs_in_class - funcs_in_all
+        raise ValueError(
+            f"Mismatch between PGTools class and _get_all_tools. Missing from class: {missing_from_class}. Missing from function: {missing_from_function}"
+        )
+
+
+def test_available_tools():
+    """Tests that the functions defined in the PGTools class match the functions listed in AVAILABLE_TOOLS."""
+    funcs_in_class = set(
+        x for x in dir(PGTools) if x[0] != "_" and isfunction(getattr(PGTools, x))
+    )
+
+    funcs_in_literal = set(get_args(AVAILABLE_TOOLS))
+
+    if funcs_in_class != funcs_in_literal:
+        missing_from_class = funcs_in_literal - funcs_in_class
+        missing_from_literal = funcs_in_class - funcs_in_literal
+        raise ValueError(
+            f"Mismatch between PGTools class and AVAILABLE_TOOLS literal. Missing from class: {missing_from_class}. Missing from literal: {missing_from_literal}"
+        )
+
+
+def test_arg_schema_sync():
+    """Tests that the function arguments defined in the PGTools class match the arguments defined in the ArgsSchema for each tool."""
+    for struct_tool in PGTools(settings=blank)._get_all_tools():
+        func_args = signature(getattr(PGTools, struct_tool.name)).parameters
+        model_args = _safe_schema(struct_tool.args_schema)["properties"]
+        for arg in func_args:
+            if arg == "self":
+                continue
+            func_annot = str(func_args[arg].annotation)
+
+            model_types = _parse_model_arg(model_args, arg)
+            func_types = _parse_func_arg(func_annot)
+            assert func_types == model_types