aurelian 0.3.2__py3-none-any.whl

aurelian/__init__.py

@@ -0,0 +1,9 @@
+"""ubergraph-agent package."""
+
+import importlib_metadata
+
+try:
+    __version__ = importlib_metadata.version(__name__)
+except importlib_metadata.PackageNotFoundError:
+    # package is not installed
+    __version__ = "0.0.0"  # pragma: no cover

aurelian/agents/amigo/__init__.py

@@ -0,0 +1,3 @@
+"""
+AmiGO agent module for interacting with GO KnowledgeBase via AmiGO solr endpoint.
+"""

aurelian/agents/amigo/amigo_agent.py

@@ -0,0 +1,77 @@
+"""
+Agent for interacting with GO KnowledgeBase via AmiGO solr endpoint.
+"""
+from aurelian.agents.amigo.amigo_config import AmiGODependencies
+from aurelian.agents.amigo.amigo_tools import (
+    find_gene_associations,
+    find_gene_associations_for_pmid,
+    lookup_uniprot_entry,
+    uniprot_mapping
+)
+from aurelian.agents.literature.literature_tools import (
+    lookup_pmid as literature_lookup_pmid,
+    search_literature_web,
+    retrieve_literature_page
+)
+from pydantic_ai import Agent, Tool
+
+SYSTEM = """
+You are a biocurator that can answer questions using the Gene Ontology knowledge base via the AmiGO API.
+
+Do not assume the knowledge base is complete or always correct. Your job is to help curators find mistakes
+or missing information. A particular pervasive issue in GO is over-annotation based on phenotypes - a gene
+should only be annotated to a process if it is involved in that process, i.e., if the activity of the
+gene process is an identifiable step in the pathway.
+
+You can help with:
+- Finding gene associations for specific genes or gene products
+- Finding gene associations cited in specific publications by PMID
+- Looking up protein information via UniProt
+- Mapping UniProt accessions to other databases
+- Analyzing gene function and involvement in biological processes
+- Identifying potential over-annotations or missing annotations
+
+You can also use your general knowledge of genes and biological processes, and do additional searches
+when needed to provide context or verification.
+"""
+
+amigo_agent = Agent(
+    model="openai:gpt-4o",
+    deps_type=AmiGODependencies,
+    system_prompt=SYSTEM,
+    tools=[
+        Tool(find_gene_associations),
+        Tool(find_gene_associations_for_pmid),
+        Tool(lookup_uniprot_entry),
+        Tool(uniprot_mapping),
+        Tool(literature_lookup_pmid, name="lookup_pmid",
+             description="""Lookup the text of a PubMed article by its PMID.
+
+A PMID should be of the form "PMID:nnnnnnn" (no underscores).
+
+This is useful for retrieving the full text of papers referenced in GO annotations
+to verify the evidence for gene annotations or identify potential over-annotations.
+
+Args:
+    pmid: The PubMed ID to look up
+        
+Returns:
+    str: Full text if available, otherwise abstract"""),
+        Tool(search_literature_web, name="search_web",
+             description="""Search the web using a text query.
+    
+Args:
+    query: The search query
+        
+Returns:
+    str: Search results with summaries"""),
+        Tool(retrieve_literature_page, name="retrieve_web_page",
+             description="""Fetch the contents of a web page.
+    
+Args:
+    url: The URL to fetch
+        
+Returns:
+    str: The contents of the web page""")
+    ]
+)

aurelian/agents/amigo/amigo_config.py

@@ -0,0 +1,85 @@
+"""
+Configuration classes for the AmiGO agent.
+"""
+from dataclasses import dataclass, field
+
+from bioservices import UniProt
+from oaklib import get_adapter
+from oaklib.implementations import AmiGOImplementation
+
+from aurelian.dependencies.workdir import HasWorkdir
+from aurelian.agents.uniprot.uniprot_tools import normalize_uniprot_id
+
+# Initialize UniProt service
+uniprot_service = UniProt(verbose=False)
+
+
+@dataclass
+class AmiGODependencies(HasWorkdir):
+    """
+    Configuration for the AmiGO agent.
+    
+    Args:
+        taxon: NCBI Taxonomy ID, defaults to human (9606)
+    """
+    taxon: str = field(default="9606")
+    
+    def get_uniprot_service(self) -> UniProt:
+        """
+        Get the UniProt service for protein lookups.
+        
+        Returns:
+            UniProt: The UniProt service
+        """
+        return uniprot_service
+    
+    def get_amigo_adapter(self) -> AmiGOImplementation:
+        """
+        Get the AmiGO adapter for the specified taxon.
+        
+        Returns:
+            AmiGOImplementation: The OAK AmiGO adapter
+        """
+        return get_adapter(f"amigo:NCBITaxon:{self.taxon}")
+        
+    def get_gene_id(self, gene_term: str) -> str:
+        """
+        Normalize a gene identifier.
+        
+        Args:
+            gene_term: The gene identifier
+            
+        Returns:
+            str: The normalized gene identifier
+        """
+        return gene_term
+
+
+def normalize_pmid(pmid: str) -> str:
+    """
+    Normalize a PubMed ID to the format PMID:nnnnnnn.
+    
+    Args:
+        pmid: The PubMed ID
+        
+    Returns:
+        str: The normalized PubMed ID
+    """
+    if ":" in pmid:
+        pmid = pmid.split(":", 1)[1]
+    if not pmid.startswith("PMID:"):
+        pmid = f"PMID:{pmid}"
+    return pmid
+
+
+def get_config(taxon: str = "9606") -> AmiGODependencies:
+    """
+    Get the AmiGO configuration.
+    
+    Args:
+        taxon: NCBI Taxonomy ID, defaults to human (9606)
+        
+    Returns:
+        AmiGODependencies: The AmiGO dependencies
+    """
+    return AmiGODependencies(taxon=taxon)

aurelian/agents/amigo/amigo_evals.py

@@ -0,0 +1,73 @@
+"""
+Evaluation module for the AmiGO agent.
+
+This module implements evaluations for the AmiGO 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.amigo.amigo_agent import amigo_agent
+from aurelian.agents.amigo.amigo_config import AmiGODependencies
+
+class AmiGOMetadata(Dict[str, Any]):
+    """Simple metadata dictionary for AmiGO evaluations."""
+    pass
+
+# Define individual evaluation cases
+case1 = Case(
+    name="uniprot_annotations",
+    inputs="What are some annotations for the protein UniProtKB:Q9UMS5",
+    expected_output="GO:",  # Should contain GO terms
+    metadata=metadata("medium", "annotation_retrieval")
+)
+
+case2 = Case(
+    name="paper_overannotation",
+    inputs="Check PMID:19661248 for over-annotation",
+    expected_output="annotation",  # Should evaluate the paper's annotations
+    metadata=metadata("hard", "literature_assessment")
+)
+
+case3 = Case(
+    name="pathway_genes",
+    inputs="What genes are involved in the ribosome biogenesis pathway?",
+    expected_output="ribosom",  # Should mention ribosome-related genes
+    metadata=metadata("medium", "pathway_analysis")
+)
+
+case4 = Case(
+    name="database_mapping",
+    inputs="Map UniProtKB:P04637 to KEGG database",
+    expected_output="KEGG",  # Should contain KEGG IDs
+    metadata=metadata("easy", "database_mapping")
+)
+
+case5 = Case(
+    name="dna_repair_genes",
+    inputs="Search for genes involved in DNA repair and show me their annotations",
+    expected_output="repair",  # Should mention DNA repair annotations
+    metadata=metadata("medium", "gene_function_search")
+)
+
+def create_eval_dataset() -> Dataset[str, str, MetadataDict]:
+    """
+    Create a dataset for evaluating the AmiGO agent.
+    
+    Returns:
+        Dataset of AmiGO evaluation cases with appropriate evaluators
+    """
+    # Collect all cases
+    cases = [case1, case2, case3, case4, case5]
+    
+    # Dataset-level evaluators
+    evaluators = [SubstringEvaluator()]
+    
+    return Dataset(
+        cases=cases,
+        evaluators=evaluators
+    )

aurelian/agents/amigo/amigo_gradio.py

@@ -0,0 +1,52 @@
+"""
+Gradio UI for the AmiGO agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.agents.amigo.amigo_agent import amigo_agent
+from aurelian.agents.amigo.amigo_config import AmiGODependencies
+from aurelian.utils.async_utils import run_sync
+
+
+def chat(deps: Optional[AmiGODependencies] = None, taxon: Optional[str] = None, **kwargs):
+    """
+    Initialize a chat interface for the AmiGO agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        taxon: Optional NCBI Taxonomy ID, defaults to human (9606)
+        **kwargs: Additional arguments to pass to the agent
+        
+    Returns:
+        A Gradio chat interface
+    """
+    if deps is None:
+        deps = AmiGODependencies()
+        
+    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: amigo_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="AmiGO AI Assistant",
+        examples=[
+            ["What are some annotations for the protein UniProtKB:Q9UMS5"],
+            ["Check PMID:19661248 for over-annotation"],
+            ["What genes are involved in the ribosome biogenesis pathway?"],
+            ["Map UniProtKB:P04637 to KEGG database"],
+            ["Search for genes involved in DNA repair and show me their annotations"]
+        ]
+    )

aurelian/agents/amigo/amigo_mcp.py

@@ -0,0 +1,152 @@
+"""
+MCP tools for interacting with GO KnowledgeBase via AmiGO solr endpoint.
+"""
+import os
+from typing import Dict, List
+
+from mcp.server.fastmcp import FastMCP
+
+from aurelian.agents.amigo.amigo_agent import SYSTEM
+import aurelian.agents.amigo.amigo_tools as at
+from aurelian.agents.amigo.amigo_config import AmiGODependencies
+from aurelian.agents.literature.literature_tools import (
+    lookup_pmid as lit_lookup_pmid, 
+    search_literature_web, 
+    retrieve_literature_page
+)
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("amigo", instructions=SYSTEM)
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> AmiGODependencies:
+    deps = AmiGODependencies()
+    # Set the location from environment variable or default
+    loc = os.getenv("AURELIAN_WORKDIR", "/tmp/aurelian")
+    deps.workdir = WorkDir(loc)
+    
+    # Get taxonomy ID from environment variable if available
+    taxon = os.getenv("AMIGO_TAXON")
+    if taxon:
+        deps.taxon = taxon
+        
+    return deps
+
+def ctx() -> RunContext[AmiGODependencies]:
+    rc: RunContext[AmiGODependencies] = RunContext[AmiGODependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def find_gene_associations(gene_id: str) -> List[Dict]:
+    """
+    Find gene associations for a given gene or gene product.
+    
+    Args:
+        gene_id: Gene or gene product IDs
+        
+    Returns:
+        List[Dict]: List of gene associations
+    """
+    return await at.find_gene_associations(ctx(), gene_id)
+
+
+@mcp.tool()
+async def find_gene_associations_for_pmid(pmid: str) -> List[Dict]:
+    """
+    Find gene associations for a given PubMed ID.
+    
+    Args:
+        pmid: The PubMed ID
+        
+    Returns:
+        List[Dict]: List of gene associations for the PubMed ID
+    """
+    return await at.find_gene_associations_for_pmid(ctx(), pmid)
+
+
+@mcp.tool()
+async def lookup_uniprot_entry(uniprot_acc: str) -> str:
+    """
+    Lookup the Uniprot entry for a given Uniprot accession number.
+    
+    Args:
+        uniprot_acc: The Uniprot accession
+        
+    Returns:
+        str: The Uniprot entry text
+    """
+    return await at.lookup_uniprot_entry(ctx(), uniprot_acc)
+
+
+@mcp.tool()
+async def uniprot_mapping(target_database: str, uniprot_accs: List[str]) -> Dict:
+    """
+    Perform a mapping of Uniprot accessions to another database.
+    
+    Args:
+        target_database: The target database (e.g KEGG, PDB)
+        uniprot_accs: The Uniprot accessions
+        
+    Returns:
+        Dict: Mapping results
+    """
+    return await at.uniprot_mapping(ctx(), target_database, uniprot_accs)
+
+
+@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).
+
+    This is useful for retrieving the full text of papers referenced in GO annotations
+    to verify the evidence for gene annotations or identify potential over-annotations.
+    
+    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)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')

aurelian/agents/amigo/amigo_tools.py

@@ -0,0 +1,152 @@
+"""
+Tools for the AmiGO agent.
+"""
+from typing import List, Dict
+
+from pydantic_ai import RunContext, ModelRetry
+
+from aurelian.agents.amigo.amigo_config import AmiGODependencies, normalize_pmid
+from aurelian.agents.uniprot.uniprot_tools import normalize_uniprot_id
+from aurelian.utils.data_utils import obj_to_dict
+
+from oaklib.datamodels.association import Association, NegatedAssociation
+from oaklib.implementations.amigo.amigo_implementation import (
+    DEFAULT_SELECT_FIELDS, QUALIFIER, BIOENTITY,
+    BIOENTITY_LABEL, map_predicate, ANNOTATION_CLASS, ANNOTATION_CLASS_LABEL, 
+    REFERENCE, EVIDENCE_TYPE, ASSIGNED_BY,
+    _query as amigo_query,
+    _normalize
+)
+
+
+async def find_gene_associations(ctx: RunContext[AmiGODependencies], gene_id: str) -> List[Dict]:
+    """
+    Find gene associations for a given gene or gene product.
+    
+    Args:
+        ctx: The run context
+        gene_id: Gene or gene product IDs
+        
+    Returns:
+        List[Dict]: List of gene associations
+    """
+    print(f"FIND GENE ASSOCIATIONS: {gene_id}")
+    try:
+        adapter = ctx.deps.get_amigo_adapter()
+        normalized_gene_id = ctx.deps.get_gene_id(gene_id)
+        assocs = [obj_to_dict(a) for a in adapter.associations([normalized_gene_id])]
+        
+        if not assocs:
+            raise ModelRetry(f"No gene associations found for {gene_id}. Try a different gene identifier.")
+            
+        return assocs
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error finding gene associations for {gene_id}: {str(e)}")
+
+
+async def find_gene_associations_for_pmid(ctx: RunContext[AmiGODependencies], pmid: str) -> List[Dict]:
+    """
+    Find gene associations for a given PubMed ID.
+    
+    Args:
+        ctx: The run context
+        pmid: The PubMed ID
+        
+    Returns:
+        List[Dict]: List of gene associations for the PubMed ID
+    """
+    print(f"FIND GENE ASSOCIATIONS FOR PMID: {pmid}")
+    try:
+        normalized_pmid = normalize_pmid(pmid)
+        amigo = ctx.deps.get_amigo_adapter()
+        
+        print(f"Lookup amigo annotations to PMID: {normalized_pmid}")
+        solr = amigo._solr
+        select_fields = DEFAULT_SELECT_FIELDS
+        results = amigo_query(solr, {"reference": normalized_pmid}, select_fields)
+        
+        assocs = []
+        for doc in results:
+            cls = Association
+            quals = set(doc.get(QUALIFIER, []))
+            if "not" in quals:
+                cls = NegatedAssociation
+            assoc = cls(
+                subject=_normalize(doc[BIOENTITY]),
+                subject_label=doc[BIOENTITY_LABEL],
+                predicate=map_predicate(quals),
+                negated=cls == NegatedAssociation,
+                object=doc[ANNOTATION_CLASS],
+                object_label=doc[ANNOTATION_CLASS_LABEL],
+                publications=doc[REFERENCE],
+                evidence_type=doc.get(EVIDENCE_TYPE),
+                primary_knowledge_source=doc[ASSIGNED_BY],
+                aggregator_knowledge_source="infores:go",
+            )
+            assocs.append(obj_to_dict(assoc))
+            
+        if not assocs:
+            raise ModelRetry(f"No gene associations found for PMID {pmid}. Try a different PubMed ID.")
+            
+        return assocs
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error finding gene associations for PMID {pmid}: {str(e)}")
+
+
+async def lookup_uniprot_entry(ctx: RunContext[AmiGODependencies], uniprot_acc: str) -> str:
+    """
+    Lookup the Uniprot entry for a given Uniprot accession number.
+    
+    Args:
+        ctx: The run context
+        uniprot_acc: The Uniprot accession
+        
+    Returns:
+        str: The Uniprot entry text
+    """
+    print(f"LOOKUP UNIPROT: {uniprot_acc}")
+    try:
+        normalized_acc = normalize_uniprot_id(uniprot_acc)
+        uniprot_service = ctx.deps.get_uniprot_service()
+        result = uniprot_service.retrieve(normalized_acc, frmt="txt")
+        
+        if not result or "Error" in result or "Entry not found" in result:
+            raise ModelRetry(f"Could not find UniProt entry for {uniprot_acc}. The accession may be incorrect.")
+            
+        return result
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error retrieving UniProt entry for {uniprot_acc}: {str(e)}")
+
+
+async def uniprot_mapping(ctx: RunContext[AmiGODependencies], target_database: str, uniprot_accs: List[str]) -> Dict:
+    """
+    Perform a mapping of Uniprot accessions to another database.
+    
+    Args:
+        ctx: The run context
+        target_database: The target database (e.g KEGG, PDB)
+        uniprot_accs: The Uniprot accessions
+        
+    Returns:
+        Dict: Mapping results
+    """
+    print(f"UNIPROT MAPPING: {target_database} - {uniprot_accs}")
+    try:
+        uniprot_service = ctx.deps.get_uniprot_service()
+        normalized_accs = [normalize_uniprot_id(x) for x in uniprot_accs]
+        result = uniprot_service.mapping("UniProtKB_AC-ID", target_database, ",".join(normalized_accs))
+        
+        if not result or len(result) == 0:
+            raise ModelRetry(f"No mappings found for {uniprot_accs} to {target_database}. Try a different database or accessions.")
+            
+        return result
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error mapping {uniprot_accs} to {target_database}: {str(e)}")

aurelian/agents/biblio/__init__.py

@@ -0,0 +1,42 @@
+"""
+Biblio agent package for working with bibliographic data and citations.
+"""
+
+# Constants
+HANDLE = "mongodb://localhost:27017/biblio"
+DB_NAME = "biblio"
+COLLECTION_NAME = "main"
+
+# isort: skip_file
+from .biblio_agent import biblio_agent  # noqa: E402
+from .biblio_config import BiblioDependencies, get_config  # noqa: E402
+from .biblio_gradio import chat  # noqa: E402
+from .biblio_tools import (  # noqa: E402
+    search_bibliography,
+    lookup_pmid,
+    search_web,
+    retrieve_web_page,
+)
+
+__all__ = [
+    # Constants
+    "HANDLE",
+    "DB_NAME",
+    "COLLECTION_NAME",
+    
+    # Agent
+    "biblio_agent",
+    
+    # Config
+    "BiblioDependencies",
+    "get_config",
+    
+    # Tools
+    "search_bibliography",
+    "lookup_pmid",
+    "search_web",
+    "retrieve_web_page",
+    
+    # Gradio
+    "chat",
+]