aurelian 0.3.2__py3-none-any.whl

aurelian/agents/monarch/monarch_config.py

@@ -0,0 +1,45 @@
+"""
+Configuration for the Monarch agent.
+"""
+from dataclasses import dataclass, field
+import os
+from typing import Dict, List, Optional
+
+from oaklib import get_adapter
+from oaklib.interfaces import OboGraphInterface
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+
+
+@dataclass
+class MonarchDependencies(HasWorkdir):
+    """Configuration for the Monarch agent."""
+    
+    # Default taxon ID for humans
+    taxon: str = "9606"
+    
+    def __post_init__(self):
+        """Initialize the config with default values."""
+        # HasWorkdir doesn't have a __post_init__ method, so we don't call super()
+        if self.workdir is None:
+            self.workdir = WorkDir()
+            
+    def get_monarch_adapter(self) -> OboGraphInterface:
+        """Get a configured Monarch adapter."""
+        return get_adapter("monarch:")
+        
+    def get_mondo_adapter(self) -> OboGraphInterface:
+        """Get a configured Mondo adapter."""
+        return get_adapter("sqlite:obo:mondo")
+
+
+def get_config() -> MonarchDependencies:
+    """Get the Monarch configuration from environment variables or defaults."""
+    workdir_path = os.environ.get("AURELIAN_WORKDIR", None)
+    workdir = WorkDir(location=workdir_path) if workdir_path else None
+    taxon = os.environ.get("MONARCH_TAXON", "9606")
+    
+    return MonarchDependencies(
+        workdir=workdir,
+        taxon=taxon
+    )

aurelian/agents/monarch/monarch_gradio.py

@@ -0,0 +1,51 @@
+"""
+Gradio interface for the Monarch agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.utils.async_utils import run_sync
+from .monarch_agent import monarch_agent
+from .monarch_config import MonarchDependencies
+
+
+def chat(deps: Optional[MonarchDependencies] = None, taxon: Optional[str] = None, **kwargs):
+    """
+    Initialize a chat interface for the Monarch agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        taxon: Optional taxon ID to use, defaults to human (9606)
+        **kwargs: Additional arguments to pass to the agent
+        
+    Returns:
+        A Gradio chat interface
+    """
+    if deps is None:
+        deps = MonarchDependencies()
+        
+    if taxon:
+        deps.taxon = taxon
+
+    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: monarch_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="Monarch Knowledge Base AI Assistant",
+        examples=[
+            ["Find associations for gene BRCA1"],
+            ["What diseases are associated with the APOE gene?"],
+            ["Find information about disease MONDO:0007254"],
+            ["What genes are associated with Alzheimer's disease?"]
+        ]
+    )

aurelian/agents/monarch/monarch_mcp.py

@@ -0,0 +1,65 @@
+"""
+MCP tools for interacting with the Monarch Knowledge Base.
+"""
+import os
+from typing import Dict, List, Optional
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.monarch.monarch_tools as mt
+from aurelian.agents.monarch.monarch_agent import MONARCH_SYSTEM_PROMPT
+from aurelian.agents.monarch.monarch_config import MonarchDependencies, get_config
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("monarch", instructions=MONARCH_SYSTEM_PROMPT)
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> MonarchDependencies:
+    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[MonarchDependencies]:
+    rc: RunContext[MonarchDependencies] = RunContext[MonarchDependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def find_gene_associations(gene_id: str) -> List[Dict]:
+    """
+    Find associations for a given gene or gene product.
+
+    Args:
+        gene_id: Gene identifier (e.g., HGNC symbol like "BRCA1")
+        
+    Returns:
+        List of association objects containing subject, predicate, object details
+    """
+    return await mt.find_gene_associations(ctx(), gene_id)
+
+
+@mcp.tool()
+async def find_disease_associations(disease_id: str) -> List[Dict]:
+    """
+    Find associations for a given disease.
+
+    Args:
+        disease_id: Disease identifier (e.g., MONDO:0007254)
+        
+    Returns:
+        List of association objects containing subject, predicate, object details
+    """
+    return await mt.find_disease_associations(ctx(), disease_id)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')

aurelian/agents/monarch/monarch_tools.py

@@ -0,0 +1,113 @@
+"""
+Tools for interacting with the Monarch Knowledge Base.
+"""
+import asyncio
+from typing import Dict, List, Optional
+
+from pydantic_ai import RunContext, ModelRetry
+
+from aurelian.utils.data_utils import obj_to_dict
+from aurelian.utils.async_utils import run_sync
+from .monarch_config import MonarchDependencies, get_config
+
+
+def get_gene_id(ctx: RunContext[MonarchDependencies], gene_term: str) -> str:
+    """
+    Normalize a gene identifier.
+    
+    Args:
+        ctx: Run context with dependencies
+        gene_term: The gene term to normalize
+        
+    Returns:
+        Normalized gene ID
+    """
+    # Currently just a pass-through, but could be enhanced with ID normalization
+    return gene_term
+
+
+async def find_gene_associations(ctx: RunContext[MonarchDependencies], gene_id: str) -> List[Dict]:
+    """
+    Find associations for a given gene ID.
+    
+    Args:
+        ctx: Run context with dependencies
+        gene_id: The gene ID to find associations for
+        
+    Returns:
+        List of associations for the gene
+    """
+    config = ctx.deps or get_config()
+    adapter = config.get_monarch_adapter()
+    
+    try:
+        normalized_gene_id = get_gene_id(ctx, gene_id)
+        
+        # Execute the potentially blocking operation in a thread pool
+        associations = await asyncio.to_thread(adapter.associations, [normalized_gene_id])
+        
+        if not associations:
+            raise ModelRetry(f"No gene associations found for {gene_id}. Try a different gene identifier.")
+        
+        # Convert associations to dictionaries asynchronously 
+        result = []
+        for assoc in associations:
+            dict_assoc = await asyncio.to_thread(obj_to_dict, assoc)
+            result.append(dict_assoc)
+            
+        return result
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error retrieving gene associations: {str(e)}")
+
+
+def get_disease_id(ctx: RunContext[MonarchDependencies], term: str) -> str:
+    """
+    Normalize a disease identifier.
+    
+    Args:
+        ctx: Run context with dependencies
+        term: The disease term to normalize
+        
+    Returns:
+        Normalized disease ID
+    """
+    # Currently just a pass-through, but could be enhanced with ID normalization
+    return term
+
+
+async def find_disease_associations(ctx: RunContext[MonarchDependencies], disease_id: str) -> List[Dict]:
+    """
+    Find associations for a given disease ID.
+    
+    Args:
+        ctx: Run context with dependencies
+        disease_id: The disease ID to find associations for
+        
+    Returns:
+        List of associations for the disease
+    """
+    config = ctx.deps or get_config()
+    adapter = config.get_monarch_adapter()
+    
+    try:
+        normalized_disease_id = get_disease_id(ctx, disease_id)
+        
+        # Execute the potentially blocking operation in a thread pool
+        associations = await asyncio.to_thread(adapter.associations, [normalized_disease_id])
+        
+        if not associations:
+            raise ModelRetry(f"No disease associations found for {disease_id}. Try a different disease identifier.")
+            
+        # Convert associations to dictionaries asynchronously
+        result = []
+        for assoc in associations:
+            dict_assoc = await asyncio.to_thread(obj_to_dict, assoc)
+            result.append(dict_assoc)
+            
+        return result
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error retrieving disease associations: {str(e)}")

aurelian/agents/oak/oak_config.py

@@ -0,0 +1,27 @@
+"""
+Configuration for the OAK agent.
+"""
+import os
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+
+
+class OakDependencies(HasWorkdir):
+    """
+    OAK agent dependencies that include a working directory.
+    
+    This allows the agent to maintain state and access files on the local filesystem.
+    """
+    pass
+
+
+def get_config() -> OakDependencies:
+    """Get a default configuration for the OAK agent.
+    
+    Returns:
+        OakDependencies object with default settings
+    """
+    workdir_path = os.environ.get("AURELIAN_WORKDIR", None)
+    workdir = WorkDir(location=workdir_path) if workdir_path else None
+
+    return OakDependencies(workdir=workdir)

aurelian/agents/oak/oak_gradio.py

@@ -0,0 +1,57 @@
+"""
+Gradio UI for the OAK agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.dependencies.workdir import HasWorkdir
+from aurelian.utils.async_utils import run_sync
+from pydantic_ai import Agent, Tool
+
+# Create an Agent
+oak_agent = Agent(
+    model="openai:gpt-4o",
+    deps_type=HasWorkdir,
+    system_prompt="""
+    You are an expert OAK (Ontology Access Kit) assistant. You can help users interact with ontology databases,
+    search ontologies, and perform common ontology operations.
+    Always provide clear explanations of what you're doing.
+    """,
+    tools=[]  # OAK tools would be added here
+)
+
+def chat(deps: Optional[HasWorkdir] = None, **kwargs):
+    """
+    Initialize a chat interface for the OAK agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        **kwargs: Additional arguments to pass to the agent
+        
+    Returns:
+        A Gradio chat interface
+    """
+    if deps is None:
+        deps = HasWorkdir()
+
+    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: oak_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="OAK AI Assistant",
+        examples=[
+            ["Search for 'diabetes' in the Human Phenotype Ontology"],
+            ["Find terms related to 'heart' in GO"],
+            ["Get information about the term 'HP:0001250'"]
+        ]
+    )

aurelian/agents/ontology_mapper/__init__.py

@@ -0,0 +1,31 @@
+"""
+Ontology Mapper agent package for working with ontology term mappings.
+"""
+
+from .ontology_mapper_agent import ontology_mapper_agent, add_ontologies
+from .ontology_mapper_config import OntologyMapperDependencies
+from .ontology_mapper_gradio import chat
+from .ontology_mapper_tools import (
+    search_terms,
+    search_web,
+    retrieve_web_page,
+    get_ontology_adapter,
+)
+
+__all__ = [
+    # Agent
+    "ontology_mapper_agent",
+    "add_ontologies",
+    
+    # Config
+    "OntologyMapperDependencies",
+    
+    # Tools
+    "search_terms",
+    "search_web",
+    "retrieve_web_page",
+    "get_ontology_adapter",
+    
+    # Gradio
+    "chat",
+]

aurelian/agents/ontology_mapper/ontology_mapper_agent.py

@@ -0,0 +1,56 @@
+"""
+Agent for creating ontology mappings.
+"""
+from pydantic_ai import Agent, RunContext
+
+from .ontology_mapper_config import OntologyMapperDependencies, get_config
+from .ontology_mapper_tools import (
+    search_terms,
+    search_web,
+    retrieve_web_page,
+)
+
+# System prompt for the ontology mapper agent
+ONTOLOGY_MAPPER_SYSTEM_PROMPT = (
+    "You are an expert on OBO ontologies."
+    " Your task is to assist the user in finding the relevant ontology terms,"
+    " given inputs such as search queries, lists of terms to map, alternate ontology classes, etc."
+    " You have access to a limited set of ontologies, which you can search using the `search_terms` function."
+    " This uses embedding-based search, so you can use partial terms, alternate names, etc."
+    " You can also expand the users search terms as appropriate, making use of any context provided."
+    " You should show your reasoning, and your candidate list (as many as appropriate)."
+    " Do not be completely literal in the task of matching ontology terms. If something seems out of scope"
+    " for an ontology, give the appropriate response and recommendation. "
+    " If a term is in scope and can't be found, suggest a term request."
+    " Give detailed provenance chains in <details> tags."
+    " Show ontology term IDs together with labels whenever possible."
+    " IMPORTANT: precision is important. If a user makes a query for a term then you should only return terms"
+    " that represent the SAME CONCEPT. Sometimes this will not be possible, and only close concepts can be found."
+    " Here you can report the close terms, but make it clear these are NOT THE SAME. Before doing this, you should"
+    " try strategies like varying your search term, based on your knowledge of that ontology"
+    " You must NEVER guess ontology term IDs, the query results should always be the source of truth."
+    "Stick to markdown, and all prefixed IDs should by hyperlinked with bioregistry,"
+    " i.e https://bioregistry.io/{curie}."
+)
+
+# Create the agent with the system prompt
+ontology_mapper_agent = Agent(
+    model="openai:gpt-4o",
+    deps_type=OntologyMapperDependencies,
+    result_type=str,
+    system_prompt=ONTOLOGY_MAPPER_SYSTEM_PROMPT,
+)
+
+# Register the tools with the agent
+ontology_mapper_agent.tool(search_terms)
+ontology_mapper_agent.tool_plain(search_web)
+ontology_mapper_agent.tool_plain(retrieve_web_page)
+
+
+@ontology_mapper_agent.system_prompt
+def add_ontologies(ctx: RunContext[OntologyMapperDependencies]) -> str:
+    """Add the list of allowed ontologies to the system prompt."""
+    allowed_ontologies = ctx.deps.ontologies
+    if allowed_ontologies:
+        return f"Allowed ontologies: {allowed_ontologies}"
+    return "Use any ontology (ideally in OBO repository)"

aurelian/agents/ontology_mapper/ontology_mapper_config.py

@@ -0,0 +1,50 @@
+"""
+Configuration for the Ontology Mapper agent.
+"""
+from dataclasses import dataclass, field
+import os
+from typing import List, Optional
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+
+
+@dataclass
+class OntologyMapperDependencies(HasWorkdir):
+    """
+    Configuration for the ontology mapper agent.
+
+    We include a default set of ontologies because the initial text embedding index is slow..
+    this can easily be changed e.g. in command line
+    """
+    max_search_results: int = 30
+    ontologies: List[str] = field(
+        default_factory=lambda: ["mondo", "hp", "go", "uberon", "cl", "mp", "envo", "obi"]
+    )
+    
+    def __post_init__(self):
+        """Initialize the config with default values."""
+        # HasWorkdir doesn't have a __post_init__ method, so we don't call super()
+        if self.workdir is None:
+            self.workdir = WorkDir()
+
+
+def get_config(ontologies: Optional[List[str]] = None) -> OntologyMapperDependencies:
+    """Get the Ontology Mapper configuration from environment variables or defaults."""
+    workdir_path = os.environ.get("AURELIAN_WORKDIR", None)
+    workdir = WorkDir(location=workdir_path) if workdir_path else None
+    
+    config = OntologyMapperDependencies(workdir=workdir)
+    
+    # Set ontologies if specified
+    if ontologies:
+        config.ontologies = ontologies
+    
+    # Allow environment variable to override max search results
+    max_results_env = os.environ.get("MAX_SEARCH_RESULTS")
+    if max_results_env:
+        try:
+            config.max_search_results = int(max_results_env)
+        except ValueError:
+            pass
+    
+    return config

aurelian/agents/ontology_mapper/ontology_mapper_evals.py

@@ -0,0 +1,108 @@
+"""
+Evaluation module for the Ontology Mapper agent.
+
+This module implements evaluations for the Ontology Mapper 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.ontology_mapper.ontology_mapper_agent import ontology_mapper_agent
+from aurelian.agents.ontology_mapper.ontology_mapper_config import OntologyMapperDependencies, get_config
+
+class OntologyMapperMetadata(Dict[str, Any]):
+    """Simple metadata dictionary for Ontology Mapper evaluations."""
+    pass
+
+# Define individual evaluation cases
+case1 = Case(
+    name="cell_ontology_neuron",
+    inputs="Find the term in the cell ontology for neuron",
+    expected_output="CL:0000540",  # Neuron ID in Cell Ontology
+    metadata=metadata("easy", "basic_term_lookup")
+)
+
+case2 = Case(
+    name="middle_fingers_terms",
+    inputs="Best terms to use for the middle 3 fingers",
+    expected_output="digit",  # Should mention digit terms
+    metadata=metadata("medium", "anatomical_term_suggestion"),
+    evaluators=[
+        LLMJudge(
+            rubric="""
+            Answer should:
+            1. Identify appropriate anatomical terms for index, middle, and ring fingers
+            2. Include ontology IDs where possible
+            3. Explain why these terms are appropriate
+            """,
+            include_input=True
+        )
+    ]
+)
+
+case3 = Case(
+    name="go_cell_division",
+    inputs="What is the term for the process of cell division in GO?",
+    expected_output="GO:0051301",  # Cell division GO term
+    metadata=metadata("easy", "go_term_lookup")
+)
+
+case4 = Case(
+    name="mp_term_mapping",
+    inputs="""
+    Find good MP terms for the following. If no matches can be found, suggest appropriate action
+
+    * CA1 TBS Reduced
+    * CA1 TBS Increased
+    * Surface righting Reduced
+    * Contextual fear conditioning (shock context/context shock) Reduced
+    * Morris water maze Reduced
+    * Rotarod Increased
+    * Rotarod Reduced
+    """,
+    expected_output="MP:",  # Should provide MP IDs
+    metadata=metadata("hard", "complex_term_mapping")
+)
+
+case5 = Case(
+    name="cross_ontology_mapping",
+    inputs="Map the Human Phenotype Ontology term for 'Seizure' (HP:0001250) to Mouse Phenotype Ontology",
+    expected_output="MP:0002064",  # Seizure in MP
+    metadata=metadata("medium", "cross_ontology_mapping")
+)
+
+def create_eval_dataset() -> Dataset[str, str, MetadataDict]:
+    """
+    Create a dataset for evaluating the Ontology Mapper agent.
+    
+    Returns:
+        Dataset of Ontology Mapper 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 of ontology term identification
+            2. Correctness of ontology IDs provided
+            3. Appropriateness of term selection for the given context
+            4. Clear explanation of term selection logic and relationships
+            5. Suitable handling of cases where no exact match exists
+            """,
+            model="anthropic:claude-3-7-sonnet-latest"
+        )
+    ]
+    
+    return Dataset(
+        cases=cases,
+        evaluators=evaluators
+    )

aurelian/agents/ontology_mapper/ontology_mapper_gradio.py

@@ -0,0 +1,58 @@
+"""
+Gradio interface for the Ontology Mapper agent.
+"""
+import os
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.utils.async_utils import run_sync
+from .ontology_mapper_agent import ontology_mapper_agent
+from .ontology_mapper_config import OntologyMapperDependencies, get_config
+
+
+def chat(deps: Optional[OntologyMapperDependencies] = None, **kwargs):
+    """
+    Initialize a chat interface for the Ontology Mapper agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        **kwargs: Additional arguments to pass to the agent
+        
+    Returns:
+        A Gradio chat interface
+    """
+    if deps is None:
+        deps = get_config()
+
+    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: ontology_mapper_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="Ontology Mapper AI Assistant",
+        examples=[
+            "Find the term in the cell ontology for neuron",
+            "Best terms to use for the middle 3 fingers",
+            "What is the term for the process of cell division in GO?",
+            """
+            Find good MP terms for the following. If no matches can be found, suggest appropriate action
+
+            * CA1 TBS Reduced
+            * CA1 TBS Increased
+            * Surface righting Reduced
+            * Contextual fear conditioning (shock context/context shock) Reduced
+            * Morris water maze Reduced
+            * Rotarod Increased
+            * Rotarod Reduced
+            """,
+        ]
+    )