aurelian 0.3.2__py3-none-any.whl

aurelian/agents/gocam/gocam_curator_agent.py

@@ -0,0 +1,46 @@
+from pydantic_ai import Agent, Tool, RunContext
+
+from aurelian.agents.gocam import GOCAMDependencies, validate_gocam_model
+from aurelian.agents.gocam.gocam_agent import PREDICATES_INFO, SYSTEM_CURIES, core_tools, \
+    get_documents_for_prompt
+from aurelian.agents.gocam.gocam_tools import lookup_gocam_local
+
+CURATOR_SYSTEM = f"""
+You are a GO-CAM curator in charge of curating GO-CAMs.
+
+When curating a GO-CAM, you should first review all relevant information:
+
+- read the relevant papers provided to you `literature_lookup_pmid`
+- find any additional references or papers of relevance `lookup_uniprot_entry`
+- check the relevant GO-CAM best practices (see `fetch_document`)
+- find the right uniprot IDs, and check what information is known `lookup_uniprot_entry`
+- find the relevant GO and ontology terms `lookup_go_term`
+
+{PREDICATES_INFO}
+
+{SYSTEM_CURIES}
+"""
+gocam_curator_agent = Agent(
+    model="anthropic:claude-3-7-sonnet-latest",
+    deps_type=GOCAMDependencies,
+    system_prompt=CURATOR_SYSTEM,
+    tools=[
+        *core_tools,
+        Tool(lookup_gocam_local),
+        Tool(validate_gocam_model),
+    ],
+)
+
+
+@gocam_curator_agent.system_prompt
+def add_documents(ctx: RunContext[GOCAMDependencies]) -> str:
+    """
+    Add available GO-CAM documents to the system prompt.
+
+    Args:
+        ctx: The run context
+
+    Returns:
+        A string containing the list of available GO-CAM documents
+    """
+    return get_documents_for_prompt()

aurelian/agents/gocam/gocam_evals.py

@@ -0,0 +1,67 @@
+"""
+Evaluation module for the GO-CAM agent.
+
+This module implements evaluations for the GO-CAM 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 aurelian.agents.gocam.gocam_agent import gocam_agent
+from aurelian.agents.gocam.gocam_config import GOCAMDependencies
+
+class GOCAMMetadata(Dict[str, Any]):
+    """Simple metadata dictionary for GO-CAM evaluations."""
+    pass
+
+# Define individual evaluation cases
+case1 = Case(
+    name="apoptosis_genes",
+    inputs="Find a model relating to apoptosis and list their genes",
+    expected_output="CASP",
+    metadata=metadata("medium", "information_retrieval")
+)
+
+case2 = Case(
+    name="count_gene_products",
+    inputs="How many distinct gene products in 62b4ffe300001804? Answer with a number, e.g. 7.",
+    expected_output="4",
+    metadata=metadata("easy", "counting")
+)
+
+case3 = Case(
+    name="nonexistent_model",
+    inputs="Find a model with ID gomodel:1234 and summarize it",
+    expected_output=None,  # Just checking the agent doesn't error out
+    metadata=metadata("medium", "error_handling")
+)
+
+case4 = Case(
+    name="signaling_receptor_id",
+    inputs="When curating GO-CAMs, the activity unit for a ligand of a signaling receptor should use which GO MF ID for the activity unit?",
+    expected_output="0048018",
+    metadata=metadata("hard", "knowledge_retrieval")
+)
+
+def create_eval_dataset() -> Dataset[str, str, MetadataDict]:
+    """
+    Create a dataset for evaluating the GO-CAM agent.
+    
+    Returns:
+        Dataset of GO-CAM evaluation cases with appropriate evaluators
+    """
+    # Collect all cases
+    cases = [case1, case2, case3, case4]
+    
+    # Dataset-level evaluators
+    evaluators = [SubstringEvaluator()]
+    
+    return Dataset(
+        cases=cases,
+        evaluators=evaluators
+    )
+

aurelian/agents/gocam/gocam_gradio.py

@@ -0,0 +1,89 @@
+"""
+Gradio UI for the GOCAM agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.agents.gocam.gocam_agent import gocam_agent
+from aurelian.agents.gocam.gocam_config import GOCAMDependencies
+from aurelian.utils.async_utils import run_sync
+
+
+def ui():
+    """
+    Initialize a basic interface for the GOCAM agent (non-chat).
+    
+    Returns:
+        A Gradio interface
+    """
+    deps = GOCAMDependencies()
+
+    def get_info(query: str):
+        print(f"QUERY: {query}")
+        result = run_sync(lambda: gocam_agent.run_sync(query, deps=deps))
+        return result.data
+
+    demo = gr.Interface(
+        fn=get_info,
+        inputs=gr.Textbox(
+            label="Ask about any GO-CAMs", placeholder="What is the function of caspase genes in apoptosis pathways?"
+        ),
+        outputs=gr.Textbox(label="GO-CAM Information"),
+        title="GO-CAM AI Assistant",
+        description="Ask me anything about GO-CAMs and I will try to provide you with the information you need.",
+    )
+    return demo
+
+
+def chat(deps: Optional[GOCAMDependencies] = None, db_path: Optional[str] = None, collection_name: Optional[str] = None, **kwargs):
+    """
+    Initialize a chat interface for the GOCAM 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 = GOCAMDependencies()
+        
+    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: gocam_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="GO-CAM AI Assistant",
+        examples=[
+            ["What is the function of caspase genes in apoptosis pathways?"],
+            ["What models involve autophagy?"],
+            [
+                (
+                    "find the wikipedia article on the integrated stress response pathway,"
+                    " download it, and summarize the genes and what they do."
+                    " then find similar GO-CAMs, look up their details,"
+                    " and compare them to the reviews"
+                )
+            ],
+            ["Find models involving the NLRP3 inflammasome. Compare the GO-CAM model with information available from uniprot"],
+            ["Examine models for antimicrobial resistance, look for commonalities in genes"],
+            ["When curating GO-CAMs, the activity unit for a ligand of a signaling receptor should use which GO MF ID for the activity unit?"],
+        ],
+    )

aurelian/agents/gocam/gocam_mcp.py

@@ -0,0 +1,224 @@
+"""
+MCP tools for working with Gene Ontology Causal Activity Models.
+"""
+import os
+from typing import Dict, List
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.filesystem.filesystem_tools as fst
+from aurelian.agents.gocam.gocam_agent import SYSTEM
+import aurelian.agents.gocam.gocam_tools as gt
+from aurelian.agents.literature.literature_tools import lookup_pmid as lit_lookup_pmid, search_literature_web, retrieve_literature_page
+from aurelian.agents.gocam.gocam_config import GOCAMDependencies
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("gocam", instructions=SYSTEM)
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> GOCAMDependencies:
+    deps = GOCAMDependencies()
+    # 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("GOCAM_DB_PATH")
+    db_name = os.getenv("GOCAM_DB_NAME") 
+    collection_name = os.getenv("GOCAM_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[GOCAMDependencies]:
+    rc: RunContext[GOCAMDependencies] = RunContext[GOCAMDependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def search_gocams(query: str) -> List[Dict]:
+    """
+    Performs a retrieval search over the GO-CAM database.
+
+    The query can be any text, such as name of a pathway, genes, or
+    a complex sentence.
+
+    The objects returned are summaries of GO-CAM models; they do not contain full
+    details. Use `lookup_gocam` to retrieve full details of a model.
+
+    This tool uses a retrieval method that is not guaranteed to always return
+    complete results, and some results may be less relevant than others.
+    You MAY use your judgment in filtering these.
+    
+    Args:
+        query: The search query text
+        
+    Returns:
+        List[Dict]: List of GOCAM models matching the query
+    """
+    return await gt.search_gocams(ctx(), query)
+
+
+@mcp.tool()
+async def lookup_gocam(model_id: str) -> Dict:
+    """
+    Performs a lookup of a GO-CAM model by its ID, and returns the model.
+    
+    Args:
+        model_id: The ID of the GO-CAM model to look up
+        
+    Returns:
+        Dict: The GO-CAM model data
+    """
+    return await gt.lookup_gocam(ctx(), model_id)
+
+
+@mcp.tool()
+async def lookup_uniprot_entry(uniprot_acc: str) -> str:
+    """
+    Lookup the Uniprot entry for a given Uniprot accession number.
+
+    This can be used to obtain further information about a protein in
+    a GO-CAM.
+
+    Args:
+        uniprot_acc: The Uniprot accession
+        
+    Returns:
+        str: Detailed functional and other info about the protein
+    """
+    return await gt.lookup_uniprot_entry(ctx(), uniprot_acc)
+
+
+@mcp.tool()
+async def lookup_pmid(pmid: str) -> str:
+    """
+    Lookup the text of a PubMed article by its PMID.
+
+    Note that assertions in GO-CAMs may reference PMIDs, so this tool
+    is useful for validating assertions. A common task is to align
+    the text of a PMID with the text of an assertion, or extracting text
+    snippets from the publication that support the assertion.
+    
+    Args:
+        pmid: The PubMed ID to look up
+        
+    Returns:
+        str: Full text if available, otherwise abstract
+    """
+    return await lit_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 search_literature_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 retrieve_literature_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)
+
+
+@mcp.tool()
+async def fetch_document(name: str, format: str = "md") -> str:
+    """
+    Lookup the GO-CAM document by name.
+
+    Args:
+        name: The document name (e.g. "How_to_annotate_complexes_in_GO-CAM")
+        format: The format of the document (defaults to "md")
+    
+    Returns:
+        The content of the document
+    """
+    return await gt.fetch_document(ctx(), name, format)
+
+
+@mcp.tool()
+async def validate_gocam_model(model_data: str, format: str = "json") -> Dict:
+    """
+    Validate a GO-CAM model against the pydantic schema.
+    
+    Args:
+        model_data: The model data as a JSON/YAML string
+        format: The format of the input data (json or yaml)
+    
+    Returns:
+        Dict with validation results, including success status and errors if any
+    """
+    return await gt.validate_gocam_model(ctx(), model_data, format)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')