aurelian 0.3.2__py3-none-any.whl

aurelian/agents/phenopackets/phenopackets_tools.py

@@ -0,0 +1,127 @@
+"""
+Tools for the phenopackets agent.
+"""
+from typing import List, Dict, Optional
+
+from pydantic_ai import RunContext, ModelRetry
+
+from aurelian.agents.phenopackets.phenopackets_config import PhenopacketsDependencies
+from aurelian.utils.data_utils import flatten
+from aurelian.agents.literature.literature_tools import (
+    lookup_pmid as literature_lookup_pmid,
+    search_literature_web,
+    retrieve_literature_page
+)
+
+
+async def search_phenopackets(ctx: RunContext[PhenopacketsDependencies], query: str) -> List[Dict]:
+    """
+    Performs a retrieval search over the Phenopackets database.
+
+    The query can be any text, such as name of a disease, phenotype, gene, etc.
+
+    The objects returned are "Phenopackets" which is a structured representation
+    of a patient. Each is uniquely identified by a phenopacket ID (essentially
+    the patient ID).
+
+    The objects returned are summaries of Phenopackets; some details such
+    as phenotypes are omitted. Use `lookup_phenopacket` to retrieve full details.
+
+    Args:
+        ctx: The run context
+        query: The search query text
+
+    Returns:
+        List[Dict]: List of phenopackets matching the query
+    """
+    print(f"SEARCH PHENOPACKETS: {query} // {ctx.deps}")
+    try:
+        qr = ctx.deps.collection.search(query, index_name="llm", limit=ctx.deps.max_results)
+        objs = []
+        for score, row in qr.ranked_rows:
+            obj = flatten(row, preserve_keys=["interpretations", "diseases"])
+            obj["relevancy_score"] = score
+            objs.append(obj)
+            print(f"RESULT: {obj}")
+        
+        if not objs:
+            raise ModelRetry(f"No phenopackets found matching the query: {query}. Try a different search term.")
+        
+        return objs
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error searching phenopackets: {str(e)}")
+
+
+async def lookup_phenopacket(ctx: RunContext[PhenopacketsDependencies], phenopacket_id: str) -> Dict:
+    """
+    Performs a lookup of an individual Phenopacket by its ID.
+
+    IDs are typically of the form PMID_nnn_PatientNumber, but this should not be assumed.
+
+    Args:
+        ctx: The run context
+        phenopacket_id: The ID of the Phenopacket to look up
+
+    Returns:
+        Dict: The phenopacket data
+    """
+    print(f"LOOKUP PHENOPACKET: {phenopacket_id}")
+    try:
+        qr = ctx.deps.collection.find({"id": phenopacket_id})
+        if not qr.rows:
+            raise ModelRetry(f"Could not find phenopacket with ID {phenopacket_id}. The ID may be incorrect.")
+        return qr.rows[0]
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error looking up phenopacket {phenopacket_id}: {str(e)}")
+
+
+async def lookup_pmid(pmid: str) -> str:
+    """
+    Lookup the text of a PubMed article by its PMID.
+    
+    A PMID should be of the form "PMID:nnnnnnn" (no underscores).
+    
+    NOTE: Phenopacket IDs are typically of the form PMID_nnn_PatientNumber,
+    but this should not be assumed. To reliably get PMIDs for a phenopacket,
+    use `lookup_phenopacket` to retrieve and examine the `externalReferences` field.
+    
+    Args:
+        pmid: The PubMed ID to look up
+        
+    Returns:
+        str: Full text if available, otherwise abstract
+    """
+    print(f"LOOKUP PMID FOR PHENOPACKET: {pmid}")
+    return await literature_lookup_pmid(pmid)
+
+
+async def search_web(query: str) -> str:
+    """
+    Search the web using a text query.
+    
+    Args:
+        query: The search query
+        
+    Returns:
+        str: Search results with summaries
+    """
+    print(f"PHENOPACKET WEB SEARCH: {query}")
+    return await search_literature_web(query)
+
+
+async def retrieve_web_page(url: str) -> str:
+    """
+    Fetch the contents of a web page.
+    
+    Args:
+        url: The URL to fetch
+        
+    Returns:
+        str: The contents of the web page
+    """
+    print(f"FETCH WEB PAGE FOR PHENOPACKET: {url}")
+    return await retrieve_literature_page(url)

aurelian/agents/rag/__init__.py

@@ -0,0 +1,40 @@
+"""
+RAG agent package for retrieval-augmented generation against document collections.
+"""
+
+# Constants
+COLLECTION_NAME = "main"
+
+# isort: skip_file
+from .rag_agent import rag_agent  # noqa: E402
+from .rag_config import RagDependencies, get_config  # noqa: E402
+from .rag_gradio import chat  # noqa: E402
+from .rag_tools import (  # noqa: E402
+    search_documents,
+    inspect_document,
+    lookup_pmid,
+    search_web,
+    retrieve_web_page,
+)
+
+__all__ = [
+    # Constants
+    "COLLECTION_NAME",
+    
+    # Agent
+    "rag_agent",
+    
+    # Config
+    "RagDependencies",
+    "get_config",
+    
+    # Tools
+    "search_documents",
+    "inspect_document",
+    "lookup_pmid",
+    "search_web",
+    "retrieve_web_page",
+    
+    # Gradio
+    "chat",
+]

aurelian/agents/rag/rag_agent.py

@@ -0,0 +1,83 @@
+"""
+Agent for retrieval-augmented generation (RAG) against document collections.
+"""
+from pydantic_ai import Agent, RunContext
+
+from .rag_config import RagDependencies
+from .rag_tools import search_documents, inspect_document, lookup_pmid, search_web, retrieve_web_page
+
+
+rag_agent = Agent(
+    model="openai:gpt-4o",
+    deps_type=RagDependencies,
+    result_type=str,
+    system_prompt=(
+        "You are an AI assistant that help explore a literature collection via RAG."
+        " You can use different functions to access the store, for example:"
+        "  - `search_documents` to find documents by text query"
+        "  - `inspect_document` to retrieve a specific document (by title/name)"
+        "You can also use `lookup_pmid` to retrieve the text of a PubMed ID, or `search_web` to search the web."
+    ),
+)
+
+
+@rag_agent.tool
+async def search_documents_tool(ctx: RunContext[RagDependencies], query: str):
+    """
+    Performs a retrieval search over the RAG database.
+    
+    The query can be any text, such as name of a disease, phenotype, gene, etc.
+    """
+    return await search_documents(ctx, query)
+
+
+@rag_agent.tool
+async def inspect_document_tool(ctx: RunContext[RagDependencies], query: str):
+    """
+    Returns the content of the document.
+    
+    Args:
+        query: E.g. title
+    """
+    return await inspect_document(ctx, query)
+
+
+@rag_agent.tool
+async def lookup_pmid_tool(ctx: RunContext[RagDependencies], pmid: str):
+    """
+    Lookup the text of a PubMed ID, using its PMID.
+    
+    A PMID should be of the form "PMID:nnnnnnn" (no underscores).
+    
+    NOTE: Phenopacket IDs are typically of the form PMID_nnn_PatientNumber,
+    but this should be be assumed. To reliably get PMIDs for a phenopacket,
+    use `lookup_phenopacket` to retrieve examine the `externalReferences`
+    field.
+    
+    Returns: full text if available, otherwise abstract
+    """
+    return await lookup_pmid(ctx, pmid)
+
+
+@rag_agent.tool
+async def search_web_tool(ctx: RunContext[RagDependencies], query: str):
+    """
+    Search the web using a text query.
+    
+    Note, this will not retrieve the full content, for that you
+    should use `retrieve_web_page`.
+    
+    Returns: matching web pages plus summaries
+    """
+    return await search_web(ctx, query)
+
+
+@rag_agent.tool
+async def retrieve_web_page_tool(ctx: RunContext[RagDependencies], url: str):
+    """
+    Fetch the contents of a web page.
+    
+    Returns:
+        The contents of the web page.
+    """
+    return await retrieve_web_page(ctx, url)

aurelian/agents/rag/rag_config.py

@@ -0,0 +1,80 @@
+"""
+Configuration for the RAG agent.
+"""
+from dataclasses import dataclass, field
+import os
+from typing import Optional
+
+from linkml_store import Client
+from linkml_store.api import Collection
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+from . import COLLECTION_NAME
+
+
+@dataclass
+class RagDependencies:
+    """Configuration for the RAG agent."""
+    
+    # Required fields
+    db_path: str
+    
+    # Optional fields with defaults
+    collection_name: str = COLLECTION_NAME
+    max_results: int = 10
+    max_content_len: int = 5000
+    workdir: Optional[WorkDir] = None
+    _collection: Optional[Collection] = None
+    
+    def __post_init__(self):
+        """Initialize the config with default values."""
+        if self.workdir is None:
+            self.workdir = WorkDir()
+
+    @property
+    def collection(self) -> Collection:
+        """Get the database collection, initializing it if needed."""
+        if self._collection is None:
+            client = Client()
+            db_path = self.db_path
+            client.attach_database(db_path)
+            db = client.databases[db_path]
+            self._collection = db.get_collection(self.collection_name)
+        return self._collection
+
+
+def get_config(db_path: Optional[str] = None, collection_name: Optional[str] = None) -> RagDependencies:
+    """
+    Get the RAG configuration from environment variables or defaults.
+    
+    Args:
+        db_path: The database path to use (overrides environment variable)
+        collection_name: The collection name to use (overrides environment variable)
+        
+    Returns:
+        A RagDependencies instance
+    """
+    # Try to get from environment, then use provided values or defaults
+    env_db_path = os.environ.get("AURELIAN_RAG_DB_PATH", None)
+    env_collection = os.environ.get("AURELIAN_RAG_COLLECTION", COLLECTION_NAME)
+    
+    # Use provided values first, then environment, then defaults
+    final_db_path = db_path or env_db_path
+    final_collection = collection_name or env_collection
+    
+    # For testing purposes, if no DB path is provided, use a default one
+    # This is only used for running basic smoke tests
+    if not final_db_path:
+        if os.environ.get("TESTING", "0") == "1":
+            final_db_path = "memory://test"
+        else:
+            raise ValueError("Database path must be provided either as parameter or via AURELIAN_RAG_DB_PATH environment variable")
+    
+    workdir_path = os.environ.get("AURELIAN_WORKDIR", None)
+    workdir = WorkDir(location=workdir_path) if workdir_path else None
+    
+    return RagDependencies(
+        db_path=final_db_path,
+        collection_name=final_collection,
+        workdir=workdir,
+    )

aurelian/agents/rag/rag_gradio.py

@@ -0,0 +1,67 @@
+"""
+Gradio interface for the RAG agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from .rag_agent import rag_agent
+from .rag_config import RagDependencies, get_config
+
+
+async def get_info(query: str, history: List[str], deps: RagDependencies, model: str = None) -> str:
+    """
+    Process a query using the RAG agent.
+    
+    Args:
+        query: The user query
+        history: The conversation history
+        deps: The agent dependencies
+        model: Optional model override
+        
+    Returns:
+        The agent's response
+    """
+    print(f"QUERY: {query}")
+    print(f"HISTORY: {history}")
+    
+    # Add history to the query if available
+    if history:
+        query += "## History"
+        for h in history:
+            query += f"\n{h}"
+    
+    # Run the agent
+    result = await rag_agent.run(query, deps=deps, model=model)
+    return result.data
+
+
+def chat(deps: Optional[RagDependencies] = None, model=None, **kwargs):
+    """
+    Create a Gradio chat interface for the RAG agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        model: Optional model override
+        kwargs: Additional keyword arguments for dependencies
+        
+    Returns:
+        A Gradio ChatInterface
+    """
+    # Initialize dependencies if needed
+    if deps is None:
+        deps = get_config(**kwargs) if kwargs else RagDependencies(**kwargs)
+    
+    def get_info_wrapper(query: str, history: List[str]) -> str:
+        """Wrapper for the async get_info function."""
+        import asyncio
+        return asyncio.run(get_info(query, history, deps, model))
+    
+    return gr.ChatInterface(
+        fn=get_info_wrapper,
+        type="messages",
+        title="RAG AI Assistant",
+        examples=[
+            ["What papers in collection are relevant to microbial nitrogen fixation?"],
+        ],
+    )

aurelian/agents/rag/rag_mcp.py

@@ -0,0 +1,107 @@
+"""
+MCP tools for retrieval-augmented generation (RAG) against document collections.
+"""
+import os
+from typing import Dict, List
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.rag.rag_tools as rt
+from aurelian.agents.rag.rag_agent import rag_agent
+from aurelian.agents.rag.rag_config import RagDependencies
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("rag", instructions=rag_agent.system_prompt)
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> RagDependencies:
+    deps = RagDependencies()
+    # Set the location from environment variable or default
+    loc = os.getenv("AURELIAN_WORKDIR", "/tmp/aurelian")
+    deps.workdir = WorkDir(loc)
+    return deps
+
+def ctx() -> RunContext[RagDependencies]:
+    rc: RunContext[RagDependencies] = RunContext[RagDependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def search_documents(query: str) -> List[Dict]:
+    """
+    Performs a retrieval search over the RAG database.
+    
+    Args:
+        query: The search query (any text, such as name of a disease, phenotype, gene, etc.)
+        
+    Returns:
+        A list of document objects matching the query with relevancy scores
+    """
+    return await rt.search_documents(ctx(), query)
+
+
+@mcp.tool()
+async def inspect_document(query: str) -> str:
+    """
+    Returns the content of the document.
+    
+    Args:
+        query: E.g. title
+        
+    Returns:
+        The full content of the document
+    """
+    return await rt.inspect_document(ctx(), query)
+
+
+@mcp.tool()
+async def lookup_pmid(pmid: str) -> str:
+    """
+    Lookup the text of a PubMed ID, using its PMID.
+    
+    Args:
+        pmid: The PubMed ID to look up (format: "PMID:nnnnnnn")
+        
+    Returns:
+        The full text if available, otherwise abstract
+    """
+    return await rt.lookup_pmid(ctx(), pmid)
+
+
+@mcp.tool()
+async def search_web(query: str) -> str:
+    """
+    Search the web using a text query.
+    
+    Args:
+        query: The search query
+        
+    Returns:
+        Matching web pages plus summaries
+    """
+    return await rt.search_web(ctx(), query)
+
+
+@mcp.tool()
+async def retrieve_web_page(url: str) -> str:
+    """
+    Fetch the contents of a web page.
+    
+    Args:
+        url: The URL to fetch
+        
+    Returns:
+        The contents of the web page
+    """
+    return await rt.retrieve_web_page(ctx(), url)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')

aurelian/agents/rag/rag_tools.py

@@ -0,0 +1,189 @@
+"""
+Tools for the RAG agent for retrieval-augmented generation.
+"""
+import asyncio
+from typing import Dict, List
+
+from pydantic_ai import RunContext, ModelRetry
+
+from aurelian.utils.data_utils import flatten
+from aurelian.utils.pubmed_utils import get_pmid_text
+from aurelian.utils.search_utils import web_search, retrieve_web_page as fetch_web_page
+from .rag_config import RagDependencies
+
+
+async def search_documents(
+    ctx: RunContext[RagDependencies], 
+    query: str
+) -> List[Dict]:
+    """
+    Performs a retrieval search over the RAG database.
+    
+    Args:
+        ctx: The run context
+        query: The search query (any text, such as name of a disease, phenotype, gene, etc.)
+        
+    Returns:
+        A list of document objects matching the query with relevancy scores
+    """
+    try:
+        print(f"SEARCH: {query}")
+        
+        # Execute the potentially blocking operation in a thread pool
+        def _search():
+            qr = ctx.deps.collection.search(query, index_name="llm", limit=ctx.deps.max_results)
+            objs = []
+            for score, row in qr.ranked_rows:
+                row["content"] = row["content"][:ctx.deps.max_content_len]
+                obj = flatten(row)
+                obj["relevancy_score"] = score
+                objs.append(obj)
+                print(f"RESULT: {obj}")
+            return objs
+            
+        objs = await asyncio.to_thread(_search)
+        
+        if not objs:
+            raise ModelRetry(f"No results found for query: {query}")
+            
+        return objs
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error searching documents: {str(e)}")
+
+
+async def inspect_document(
+    ctx: RunContext[RagDependencies], 
+    query: str
+) -> List[Dict]:
+    """
+    Returns the content of a document.
+    
+    Args:
+        ctx: The run context
+        query: Identifying information for the document (e.g., title)
+        
+    Returns:
+        The document content
+    """
+    try:
+        print(f"INSPECT DOCUMENT: {query}")
+        
+        # Execute the potentially blocking operation in a thread pool
+        def _inspect():
+            qr = ctx.deps.collection.search(query, index_name="llm", limit=ctx.deps.max_results)
+            for score, row in qr.ranked_rows:
+                return row["content"]
+            return None
+            
+        content = await asyncio.to_thread(_inspect)
+        
+        if not content:
+            raise ModelRetry(f"No document found matching: {query}")
+            
+        return content
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error inspecting document: {str(e)}")
+
+
+async def lookup_pmid(
+    ctx: RunContext[RagDependencies], 
+    pmid: str
+) -> str:
+    """
+    Lookup the text of a PubMed ID, using its PMID.
+    
+    Args:
+        ctx: The run context
+        pmid: The PubMed ID to look up (format: "PMID:nnnnnnn")
+        
+    Returns:
+        The full text if available, otherwise abstract
+        
+    A PMID should be of the form "PMID:nnnnnnn" (no underscores).
+    
+    NOTE: Phenopacket IDs are typically of the form PMID_nnn_PatientNumber,
+    but this should be be assumed. To reliably get PMIDs for a phenopacket,
+    use `lookup_phenopacket` to retrieve examine the `externalReferences`
+    field.
+    """
+    try:
+        print(f"LOOKUP PMID: {pmid}")
+        
+        # Execute the potentially blocking operation in a thread pool
+        text = await asyncio.to_thread(get_pmid_text, pmid)
+        
+        if not text or text.strip() == "":
+            raise ModelRetry(f"No text found for PMID: {pmid}")
+            
+        return text
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error retrieving text from PMID: {str(e)}")
+
+
+async def search_web(
+    ctx: RunContext[RagDependencies], 
+    query: str
+) -> str:
+    """
+    Search the web using a text query.
+    
+    Args:
+        ctx: The run context
+        query: The search query
+        
+    Returns:
+        Matching web pages plus summaries
+        
+    Note, this will not retrieve the full content, for that you
+    should use `retrieve_web_page`.
+    """
+    try:
+        print(f"Web Search: {query}")
+        
+        # Execute the potentially blocking operation in a thread pool
+        results = await asyncio.to_thread(web_search, query)
+        
+        if not results or results.strip() == "":
+            raise ModelRetry(f"No web search results found for query: {query}")
+            
+        return results
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error searching web: {str(e)}")
+
+
+async def retrieve_web_page(
+    ctx: RunContext[RagDependencies], 
+    url: str
+) -> str:
+    """
+    Fetch the contents of a web page.
+    
+    Args:
+        ctx: The run context
+        url: The URL to fetch
+        
+    Returns:
+        The contents of the web page
+    """
+    try:
+        print(f"Fetch URL: {url}")
+        
+        # Execute the potentially blocking operation in a thread pool
+        content = await asyncio.to_thread(fetch_web_page, url)
+        
+        if not content or content.strip() == "":
+            raise ModelRetry(f"No content found for URL: {url}")
+            
+        return content
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error retrieving web page: {str(e)}")