aurelian 0.3.2__py3-none-any.whl

aurelian/agents/ontology_mapper/ontology_mapper_mcp.py

@@ -0,0 +1,81 @@
+"""
+MCP tools for creating ontology mappings.
+"""
+import os
+from typing import Dict, List, Optional
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.ontology_mapper.ontology_mapper_tools as omt
+from aurelian.agents.ontology_mapper.ontology_mapper_agent import ONTOLOGY_MAPPER_SYSTEM_PROMPT
+from aurelian.agents.ontology_mapper.ontology_mapper_config import OntologyMapperDependencies, get_config
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("ontology_mapper", instructions=ONTOLOGY_MAPPER_SYSTEM_PROMPT)
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> OntologyMapperDependencies:
+    deps = get_config()
+    # Set the location from environment variable or default
+    loc = os.getenv("AURELIAN_WORKDIR", "/tmp/aurelian")
+    deps.workdir = WorkDir(loc)
+    return deps
+
+def ctx() -> RunContext[OntologyMapperDependencies]:
+    rc: RunContext[OntologyMapperDependencies] = RunContext[OntologyMapperDependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def search_terms(query: str, ont: Optional[str] = None, limit: int = 10) -> List[Dict]:
+    """
+    Search for ontology terms matching a query.
+
+    Args:
+        query: The search query text
+        ont: Optional ontology ID to search in (e.g., 'cl', 'go', 'uberon')
+        limit: Maximum number of results to return
+        
+    Returns:
+        List of matching ontology terms with their details
+    """
+    return await omt.search_terms(ctx(), query, ont, limit)
+
+
+@mcp.tool()
+async def search_web(query: str) -> str:
+    """
+    Search the web for ontology-related information.
+
+    Args:
+        query: The search query
+        
+    Returns:
+        Search results with summaries
+    """
+    return await omt.search_web(ctx(), query)
+
+
+@mcp.tool()
+async def retrieve_web_page(url: str) -> str:
+    """
+    Fetch the contents of a web page related to ontologies.
+
+    Args:
+        url: The URL to fetch
+        
+    Returns:
+        The contents of the web page
+    """
+    return await omt.retrieve_web_page(ctx(), url)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')

aurelian/agents/ontology_mapper/ontology_mapper_tools.py

@@ -0,0 +1,147 @@
+"""
+Tools for the Ontology Mapper agent.
+"""
+import asyncio
+from functools import lru_cache
+from typing import Dict, List, Optional
+
+from oaklib import get_adapter
+from pydantic_ai import RunContext, ModelRetry
+
+from aurelian.utils.ontology_utils import search_ontology
+from aurelian.utils.search_utils import web_search, retrieve_web_page as fetch_web_page
+from .ontology_mapper_config import OntologyMapperDependencies, get_config
+
+
+@lru_cache
+def get_ontology_adapter(ont: str):
+    """
+    Get an adapter for the specified ontology.
+
+    Args:
+        ont: The ontology ID to get an adapter for (e.g. cl, go, uberon)
+
+    Returns:
+        An OAK adapter for the specified ontology
+    """
+    ont = ont.lower()
+    return get_adapter(f"sqlite:obo:{ont}")
+
+
+async def search_terms(
+    ctx: RunContext[OntologyMapperDependencies], 
+    ontology_id: str, 
+    query: str
+) -> List[Dict]:
+    """
+    Finds similar ontology terms to the search query.
+
+    For example:
+
+        ```
+        search_terms("go", "cycle cycle and related processes")
+        ```
+
+    Relevancy ranking is used, with semantic similarity, which means
+    queries need only be close in semantic space. E.g. while GO does not
+    deal with diseases, this may return relevant pathways or structures:
+
+        ```
+        search_terms("go", "terms most relevant to Parkinson disease")
+        ```
+
+    Args:
+        ctx: The run context
+        ontology_id: The ontology ID to search in (e.g. cl, go, uberon)
+        query: The search query
+        
+    Returns:
+        A list of matching ontology terms
+    """
+    print(f"Term Search: {ontology_id} {query}")
+    
+    try:
+        if " " in ontology_id:
+            raise ModelRetry(
+                "Invalid ontology ID, use an OBO style ID like cl, mondo, chebi, etc."
+            )
+            
+        config = ctx.deps or get_config()
+        if ontology_id.lower() not in [ont.lower() for ont in config.ontologies]:
+            allowed_onts = ", ".join(config.ontologies)
+            raise ModelRetry(
+                f"Ontology '{ontology_id}' not in allowed list: {allowed_onts}"
+            )
+            
+        adapter = get_ontology_adapter(ontology_id)
+        # Execute the potentially blocking operation in a thread pool
+        results = await asyncio.to_thread(
+            search_ontology, 
+            adapter, 
+            query, 
+            limit=config.max_search_results
+        )
+        
+        if not results:
+            raise ModelRetry(f"No results found for query '{query}' in ontology '{ontology_id}'")
+            
+        return results
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error searching ontology: {str(e)}")
+
+
+async def search_web(query: str) -> str:
+    """
+    Search the web using a text query.
+
+    Note, this will not retrieve the full content, for that you
+    should use `retrieve_web_page`.
+
+    Args:
+        query: The search query
+        
+    Returns: 
+        Matching web pages plus summaries
+    """
+    print(f"Web Search: {query}")
+    
+    try:
+        # 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 the web: {str(e)}")
+
+
+async def retrieve_web_page(url: str) -> str:
+    """
+    Fetch the contents of a web page.
+
+    Args:
+        url: The URL of the web page to retrieve
+        
+    Returns:
+        The contents of the web page
+    """
+    print(f"Fetch URL: {url}")
+    
+    try:
+        # 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 at 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)}")

aurelian/agents/phenopackets/__init__.py

@@ -0,0 +1,3 @@
+"""
+Phenopackets agent module for working with phenopacket databases.
+"""

aurelian/agents/phenopackets/phenopackets_agent.py

@@ -0,0 +1,58 @@
+"""
+Agent for working with phenopacket databases.
+"""
+from aurelian.agents.phenopackets.phenopackets_config import PhenopacketsDependencies
+from aurelian.agents.phenopackets.phenopackets_tools import (
+    search_phenopackets,
+    lookup_phenopacket,
+    lookup_pmid,
+    search_web,
+    retrieve_web_page
+)
+from aurelian.agents.filesystem.filesystem_tools import inspect_file, list_files
+from pydantic_ai import Agent, Tool
+
+SYSTEM = """
+You are an AI assistant that can answer questions using the Phenopacket database.
+
+Phenopackets are standardized data structures for representing phenotypic and genetic information 
+about patients with rare diseases or genetic disorders.
+
+You can help with:
+- Searching for phenopackets by disease, phenotype, gene, etc.
+- Looking up specific phenopackets by ID
+- Analyzing and comparing information from multiple phenopackets
+- Finding correlations between phenotypes, genes, and variants
+- Retrieving literature related to phenopackets via PubMed
+
+You can use different functions to access the database:
+- `search_phenopackets` to find phenopackets by text query
+- `lookup_phenopacket` to retrieve a specific phenopacket by ID
+- `lookup_pmid` to retrieve the text of a PubMed article 
+- `search_web` and `retrieve_web_page` for additional information
+
+Always use the database and functions provided to answer questions, rather than providing 
+your own knowledge, unless explicitly asked. Provide answers in a narrative form 
+understandable by clinical geneticists, with supporting evidence from the database.
+
+When presenting terms, include IDs alongside labels when available (e.g., HP:0001234). 
+All prefixed IDs should be hyperlinked with Bioregistry, i.e., https://bioregistry.io/{curie}.
+
+Use markdown tables for summarizing or comparing multiple patients, with appropriate 
+column headers and clear organization of information.
+"""
+
+phenopackets_agent = Agent(
+    model="openai:gpt-4o",
+    deps_type=PhenopacketsDependencies,
+    system_prompt=SYSTEM,
+    tools=[
+        Tool(search_phenopackets),
+        Tool(lookup_phenopacket),
+        Tool(lookup_pmid),
+        Tool(search_web),
+        Tool(retrieve_web_page),
+        Tool(inspect_file),
+        Tool(list_files),
+    ]
+)

aurelian/agents/phenopackets/phenopackets_config.py

@@ -0,0 +1,72 @@
+"""
+Configuration classes for the phenopackets 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
+
+HANDLE = "mongodb://localhost:27017/phenopackets"
+DB_NAME = "phenopackets"
+COLLECTION_NAME = "main"
+
+
+@dataclass
+class PhenopacketsDependencies(HasWorkdir):
+    """
+    Configuration for the phenopackets agent.
+    """
+    max_results: int = field(default=10)
+    db_path: str = field(default=HANDLE)
+    db_name: str = field(default=DB_NAME)
+    collection_name: str = field(default=COLLECTION_NAME)
+    _collection: Optional[Collection] = None
+
+    def __post_init__(self):
+        """Initialize the config with default values."""
+        # Initialize workdir if not provided
+        if self.workdir is None:
+            self.workdir = WorkDir()
+
+    @property
+    def collection(self) -> Collection:
+        """
+        Get the phenopackets collection, initializing the connection if needed.
+        
+        Returns:
+            Collection: The phenopackets collection
+        """
+        if self._collection is None:
+            client = Client()
+            print(f"Attaching to database: {self.db_path} with alias: {self.db_name}")
+            client.attach_database(self.db_path, alias=self.db_name)
+            db = client.databases[self.db_name]
+            self._collection = db.get_collection(self.collection_name)
+        return self._collection
+
+
+def get_config() -> PhenopacketsDependencies:
+    """
+    Get the Phenopackets configuration from environment variables or defaults.
+    
+    Returns:
+        PhenopacketsDependencies: The phenopackets dependencies
+    """
+    workdir_path = os.environ.get("AURELIAN_WORKDIR", None)
+    workdir = WorkDir(location=workdir_path) if workdir_path else None
+    
+    # Get any environment-specific settings
+    db_path = os.environ.get("PHENOPACKETS_DB_PATH", HANDLE)
+    db_name = os.environ.get("PHENOPACKETS_DB_NAME", DB_NAME)
+    collection_name = os.environ.get("PHENOPACKETS_COLLECTION", COLLECTION_NAME)
+    
+    return PhenopacketsDependencies(
+        workdir=workdir,
+        db_path=db_path,
+        db_name=db_name,
+        collection_name=collection_name
+    )

aurelian/agents/phenopackets/phenopackets_evals.py

@@ -0,0 +1,99 @@
+"""
+Evaluation module for the Phenopackets agent.
+
+This module implements evaluations for the Phenopackets agent using the pydantic-ai-evals framework.
+"""
+import asyncio
+import sys
+from typing import Optional, Any, Dict, Callable, Awaitable
+
+from aurelian.evaluators.model import MetadataDict, metadata
+from aurelian.evaluators.substring_evaluator import SubstringEvaluator
+from pydantic_evals import Case, Dataset
+from pydantic_evals.evaluators import LLMJudge
+
+from aurelian.agents.phenopackets.phenopackets_agent import phenopackets_agent
+from aurelian.agents.phenopackets.phenopackets_config import PhenopacketsDependencies
+
+class PhenopacketsMetadata(Dict[str, Any]):
+    """Simple metadata dictionary for Phenopackets evaluations."""
+    pass
+
+# Define individual evaluation cases
+case1 = Case(
+    name="liver_disease_patients",
+    inputs="What patients have liver disease?",
+    expected_output="hepat",  # Should mention hepatic/liver terms
+    metadata=metadata("medium", "phenotype_query")
+)
+
+case2 = Case(
+    name="metabolic_pathway_genes",
+    inputs="What phenopackets involve genes from metabolic pathways?",
+    expected_output="metabol",  # Should mention metabolic genes/pathways
+    metadata=metadata("hard", "gene_pathway_query"),
+    evaluators=[
+        LLMJudge(
+            rubric="""
+            Answer should:
+            1. Identify phenopackets containing genes involved in metabolic pathways
+            2. Link the metabolic genes to their corresponding phenotypes
+            3. Explain how these genes relate to metabolic pathways
+            4. Provide patient/case IDs where applicable
+            """,
+            include_input=True
+        )
+    ]
+)
+
+case3 = Case(
+    name="variant_effect_peroxisomal",
+    inputs="How does the type of variant affect phenotype in peroxisomal disorders?",
+    expected_output="peroxisom",  # Should discuss peroxisomal disorders
+    metadata=metadata("hard", "variant_phenotype_correlation")
+)
+
+case4 = Case(
+    name="skeletal_dysplasia_comparison",
+    inputs="Examine phenopackets for skeletal dysplasias and compare their phenotypes",
+    expected_output="skeletal",  # Should discuss skeletal terms
+    metadata=metadata("medium", "comparative_phenotype_analysis")
+)
+
+case5 = Case(
+    name="pnpla6_mutations",
+    inputs="Look up any patients with mutations in the PNPLA6 gene",
+    expected_output="PNPLA6",  # Should mention the PNPLA6 gene
+    metadata=metadata("easy", "gene_mutation_query")
+)
+
+def create_eval_dataset() -> Dataset[str, str, MetadataDict]:
+    """
+    Create a dataset for evaluating the Phenopackets agent.
+    
+    Returns:
+        Dataset of Phenopackets evaluation cases with appropriate evaluators
+    """
+    # Collect all cases
+    cases = [case1, case2, case3, case4, case5]
+    
+    # Dataset-level evaluators
+    evaluators = [
+        SubstringEvaluator(),
+        LLMJudge(
+            rubric="""
+            Evaluate the answer based on:
+            1. Accuracy in identifying relevant phenopackets based on the query
+            2. Correct interpretation of phenotype-genotype relationships
+            3. Proper use of HPO terms and gene identifiers
+            4. Comprehensive analysis of phenotypic data when requested
+            5. Clear presentation of results including patient/case identifiers when appropriate
+            """,
+            model="anthropic:claude-3-7-sonnet-latest"
+        )
+    ]
+    
+    return Dataset(
+        cases=cases,
+        evaluators=evaluators
+    )

aurelian/agents/phenopackets/phenopackets_gradio.py

@@ -0,0 +1,55 @@
+"""
+Gradio UI for the phenopackets agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.agents.phenopackets.phenopackets_agent import phenopackets_agent
+from aurelian.agents.phenopackets.phenopackets_config import PhenopacketsDependencies
+from aurelian.utils.async_utils import run_sync
+
+
+def chat(deps: Optional[PhenopacketsDependencies] = None, db_path: Optional[str] = None, collection_name: Optional[str] = None, **kwargs):
+    """
+    Initialize a chat interface for the phenopackets agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        db_path: Optional database path, defaults to MongoDB localhost
+        collection_name: Optional collection name, defaults to "main"
+        **kwargs: Additional arguments to pass to the agent
+        
+    Returns:
+        A Gradio chat interface
+    """
+    if deps is None:
+        deps = PhenopacketsDependencies()
+        
+    if db_path:
+        deps.db_path = db_path
+    if collection_name:
+        deps.collection_name = collection_name
+
+    def get_info(query: str, history: List[str]) -> str:
+        print(f"QUERY: {query}")
+        print(f"HISTORY: {history}")
+        if history:
+            query += "## History"
+            for h in history:
+                query += f"\n{h}"
+        result = run_sync(lambda: phenopackets_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="Phenopackets AI Assistant",
+        examples=[
+            ["What patients have liver disease?"],
+            ["What phenopackets involve genes from metabolic pathways?"],
+            ["How does the type of variant affect phenotype in peroxisomal disorders?"],
+            ["Examine phenopackets for skeletal dysplasias and compare their phenotypes"],
+            ["Look up any patients with mutations in the PNPLA6 gene"]
+        ]
+    )

aurelian/agents/phenopackets/phenopackets_mcp.py

@@ -0,0 +1,178 @@
+"""
+MCP tools for working with phenopacket databases.
+"""
+import os
+from typing import Dict, List
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.filesystem.filesystem_tools as fst
+from aurelian.agents.phenopackets.phenopackets_agent import SYSTEM
+import aurelian.agents.phenopackets.phenopackets_tools as pt
+from aurelian.agents.phenopackets.phenopackets_config import PhenopacketsDependencies
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("phenopackets", instructions=SYSTEM)
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> PhenopacketsDependencies:
+    deps = PhenopacketsDependencies()
+    # Set the location from environment variable or default
+    loc = os.getenv("AURELIAN_WORKDIR", "/tmp/aurelian")
+    deps.workdir = WorkDir(loc)
+    
+    # Get database connection parameters from environment if available
+    db_path = os.getenv("PHENOPACKETS_DB_PATH")
+    db_name = os.getenv("PHENOPACKETS_DB_NAME") 
+    collection_name = os.getenv("PHENOPACKETS_COLLECTION_NAME")
+    
+    if db_path:
+        deps.db_path = db_path
+    if db_name:
+        deps.db_name = db_name
+    if collection_name:
+        deps.collection_name = collection_name
+        
+    return deps
+
+def ctx() -> RunContext[PhenopacketsDependencies]:
+    rc: RunContext[PhenopacketsDependencies] = RunContext[PhenopacketsDependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def search_phenopackets(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:
+        query: The search query text
+
+    Returns:
+        List[Dict]: List of phenopackets matching the query
+    """
+    return await pt.search_phenopackets(ctx(), query)
+
+
+@mcp.tool()
+async def lookup_phenopacket(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:
+        phenopacket_id: The ID of the Phenopacket to look up
+
+    Returns:
+        Dict: The phenopacket data
+    """
+    return await pt.lookup_phenopacket(ctx(), phenopacket_id)
+
+
+@mcp.tool()
+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
+    """
+    return await pt.lookup_pmid(pmid)
+
+
+@mcp.tool()
+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
+    """
+    return await pt.search_web(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:
+        str: The contents of the web page
+    """
+    return await pt.retrieve_web_page(url)
+
+
+@mcp.tool()
+async def inspect_file(data_file: str) -> str:
+    """
+    Inspect a file in the working directory.
+
+    Args:
+        data_file: name of file
+
+    Returns:
+        str: Contents of the file
+    """
+    return await fst.inspect_file(ctx(), data_file)
+
+
+@mcp.tool()
+async def list_files() -> str:
+    """
+    List files in the working directory.
+
+    Returns:
+        str: List of files in the working directory
+    """
+    return await fst.list_files(ctx())
+
+
+@mcp.tool()
+async def write_to_file(file_name: str, data: str) -> str:
+    """
+    Write data to a file in the working directory.
+
+    Args:
+        file_name: Name of the file to write
+        data: Data to write to the file
+
+    Returns:
+        str: Confirmation message
+    """
+    return await fst.write_to_file(ctx(), file_name, data)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')