aurelian 0.3.2__py3-none-any.whl

aurelian/agents/ubergraph/__init__.py

@@ -0,0 +1,40 @@
+"""
+Ubergraph agent package for working with ontologies via the UberGraph endpoint.
+"""
+
+from .ubergraph_agent import (
+    ubergraph_agent, 
+    ASSUMPTIONS,
+    add_ontology_assumptions,
+    add_prefixes,
+)
+from .ubergraph_config import Dependencies, DEFAULT_PREFIXES, get_config
+from .ubergraph_gradio import chat
+from .ubergraph_tools import (
+    query_ubergraph,
+    QueryResults,
+    simplify_value,
+    simplify_results,
+)
+
+__all__ = [
+    # Agent
+    "ubergraph_agent",
+    "ASSUMPTIONS",
+    "add_ontology_assumptions",
+    "add_prefixes",
+    
+    # Config
+    "Dependencies",
+    "DEFAULT_PREFIXES",
+    "get_config",
+    
+    # Tools
+    "query_ubergraph",
+    "QueryResults",
+    "simplify_value",
+    "simplify_results",
+    
+    # Gradio
+    "chat",
+]

aurelian/agents/ubergraph/ubergraph_agent.py

@@ -0,0 +1,71 @@
+"""
+Agent for working with ontologies via UberGraph endpoint.
+"""
+from typing import Dict
+
+from pydantic_ai import Agent, RunContext
+
+from .ubergraph_config import Dependencies, get_config
+from .ubergraph_tools import query_ubergraph
+
+# Assumptions about the UberGraph data model
+ASSUMPTIONS = {
+    "provenance": (
+        "When formulating your response to tool outputs,",
+        " you can extemporize with your own knowledge, but if you do so,"
+        " you must be clear about which statements come from the ontology"
+        " vs your own knowledge.",
+    ),
+    "ids": "include both IDs and labels in responses, unless directed not to do so.",
+    "obo": "Assume OBO style ontology and OBO PURLs (http://purl.obolibrary.org/obo/).",
+    "rg": (
+        "All edges are stored as simple triples, e.g CL:0000080 BFO:0000050 UBERON:0000179"
+        " for 'circulating cell' 'part of' 'haemolymphatic fluid'"
+    ),
+    "ont_graph": (
+        "Direct (asserted) edges are stored in the `renci:ontology` graph." "Use this by default, even for subClassOf."
+    ),
+    "entailed": (
+        "Indirect (entailed) edges (including reflexive) are stored in the `renci:redundant` graph"
+        "Use this for queries that require transitive closure, e.g. rdfs:subClassOf+"
+        "Note however that other triples like rdfs:label are NOT in this graph - use renci:ontology for these."
+    ),
+    "paths": "In general you should NOT use paths like rdfs:subClassOf+, use the entailed graph.",
+    "ro": "RO is used for predicates. Common relations include BFO:0000050 for part-of.",
+    "is_a": "rdfs:subClassOf is used for is_a relationships.",
+    "labels": "rdfs:label used for labels. IDs/URIs are typically OBO-style.",
+    "oboInOwl": "assume obiInOwl for synonyms, e.g. oboInOwl:hasExactSynonym.",
+    "blazegraph": (
+        "Blazegraph is used as the underlying triplestore."
+        "This means you SHOULD do relevance-ranked match queries over CONTAINS. "
+        "E.g. ?c rdfs:label ?v . ?v bds:search 'circulating cell' ; ?v bds:relevance ?score ."
+    ),
+    "def": "IAO:0000115 is used for definitions.",
+    "xref": "assume oboInOwl:hasDbXref for simple cross-references.",
+    "mixed_language": "Do not assume all labels are language tagged.",
+}
+
+# Create the UberGraph agent
+ubergraph_agent = Agent(
+    "openai:gpt-4o",
+    deps_type=Dependencies,
+    result_type=str,
+)
+
+# Register tools
+ubergraph_agent.tool(query_ubergraph)
+
+
+@ubergraph_agent.system_prompt
+def add_ontology_assumptions(ctx: RunContext[Dependencies]) -> str:
+    """Add ontology assumptions to the system prompt."""
+    return "\n\n" + "\n\n".join([f"Assumption: {desc}" for name, desc in ASSUMPTIONS.items()])
+
+
+@ubergraph_agent.system_prompt
+def add_prefixes(ctx: RunContext[Dependencies]) -> str:
+    """Add SPARQL prefixes to the system prompt."""
+    prefixes = ctx.deps.prefixes
+    return "\n\nAssume the following prefixes are auto-included:" + "\n".join(
+        [f"\nPrefix: {prefix}: {expansion}" for prefix, expansion in prefixes.items()]
+    )

aurelian/agents/ubergraph/ubergraph_config.py

@@ -0,0 +1,79 @@
+"""
+Configuration for the Ubergraph agent.
+"""
+from dataclasses import dataclass, field
+import os
+from typing import Dict, Optional
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+
+# Default UberGraph endpoint
+UBERGRAPH_ENDPOINT = "https://ubergraph.apps.renci.org/sparql"
+
+# Default SPARQL prefixes
+DEFAULT_PREFIXES = {
+    "owl": "http://www.w3.org/2002/07/owl#",
+    "rdf": "http://www.w3.org/1999/02/22-rdf-syntax-ns#",
+    "rdfs": "http://www.w3.org/2000/01/rdf-schema#",
+    "schema": "http://schema.org/",
+    "obo": "http://purl.obolibrary.org/obo/",
+    "xsd": "http://www.w3.org/2001/XMLSchema#",
+    "renci": "http://reasoner.renci.org/",
+    "oboInOwl": "http://www.geneontology.org/formats/oboInOwl#",
+    "BFO": "http://purl.obolibrary.org/obo/BFO_",
+    "RO": "http://purl.obolibrary.org/obo/RO_",
+    "GO": "http://purl.obolibrary.org/obo/GO_",
+    "SO": "http://purl.obolibrary.org/obo/SO_",
+    "CHEBI": "http://purl.obolibrary.org/obo/CHEBI_",
+    "CL": "http://purl.obolibrary.org/obo/CL_",
+    "UBERON": "http://purl.obolibrary.org/obo/UBERON_",
+    "IAO": "http://purl.obolibrary.org/obo/IAO_",
+    "OBI": "http://purl.obolibrary.org/obo/OBI_",
+    "biolink": "https://w3id.org/biolink/vocab/",
+    "bds": "http://www.bigdata.com/rdf/search#",
+}
+
+
+@dataclass
+class Dependencies(HasWorkdir):
+    """Configuration for the UberGraph agent."""
+    
+    # SPARQL endpoint
+    endpoint: str = UBERGRAPH_ENDPOINT
+    
+    # Prefixes for SPARQL queries
+    prefixes: Dict[str, str] = field(default_factory=lambda: DEFAULT_PREFIXES)
+    
+    # Maximum number of results to return
+    max_results: int = 20
+    
+    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(
+    endpoint: Optional[str] = None, 
+    prefixes: Optional[Dict[str, str]] = None,
+    max_results: Optional[int] = None,
+) -> Dependencies:
+    """Get the UberGraph configuration from environment variables or defaults."""
+    # Initialize from environment or defaults
+    config_endpoint = endpoint or os.environ.get("UBERGRAPH_ENDPOINT", UBERGRAPH_ENDPOINT)
+    config_max_results = max_results or int(os.environ.get("MAX_RESULTS", "20"))
+    
+    # Get workdir from environment if specified
+    workdir_path = os.environ.get("AURELIAN_WORKDIR", None)
+    workdir = WorkDir(location=workdir_path) if workdir_path else None
+    
+    # Create config with the specified values
+    config = Dependencies(
+        endpoint=config_endpoint,
+        prefixes=prefixes or DEFAULT_PREFIXES,
+        max_results=config_max_results,
+        workdir=workdir,
+    )
+    
+    return config

aurelian/agents/ubergraph/ubergraph_gradio.py

@@ -0,0 +1,48 @@
+"""
+Gradio interface for the UberGraph agent.
+"""
+import os
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.utils.async_utils import run_sync
+from .ubergraph_agent import ubergraph_agent
+from .ubergraph_config import Dependencies, get_config
+
+
+def chat(deps: Optional[Dependencies] = None, **kwargs):
+    """
+    Initialize a chat interface for the UberGraph 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: ubergraph_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="UberGraph SPARQL Assistant",
+        examples=[
+            "Find all cell types that are part of the heart",
+            "What is the definition of CL:0000746?",
+            "What genes are expressed in neurons?",
+            "What are the subclasses of skeletal muscle tissue?",
+        ]
+    )

aurelian/agents/ubergraph/ubergraph_mcp.py

@@ -0,0 +1,69 @@
+"""
+MCP tools for working with ontologies via UberGraph endpoint.
+"""
+import os
+from typing import Dict, Optional
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.ubergraph.ubergraph_tools as ut
+from aurelian.agents.ubergraph.ubergraph_config import Dependencies, get_config
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server with combined system prompt
+SYSTEM_PROMPT = """
+You are an expert ontologist with access to the UberGraph SPARQL endpoint.
+
+UberGraph is a knowledge graph built from multiple OBO ontologies, including GO, Uberon, CL, ChEBI, and more.
+You can help users explore ontology terms, relationships, and hierarchies through SPARQL queries.
+
+IMPORTANT ASSUMPTIONS:
+- When formulating your response to tool outputs, you can extemporize with your own knowledge,
+  but if you do so, you must be clear about which statements come from the ontology vs your own knowledge.
+- Include both IDs and labels in responses, unless directed not to do so.
+- Assume OBO style ontology and OBO PURLs (http://purl.obolibrary.org/obo/).
+- All edges are stored as simple triples, e.g CL:0000080 BFO:0000050 UBERON:0000179 for 'circulating cell'
+  'part of' 'haemolymphatic fluid'
+- Direct (asserted) edges are stored in the 'renci:ontology' graph. Use this by default, even for subClassOf.
+- Indirect (entailed) edges (including reflexive) are stored in the 'renci:redundant' graph. Use this for
+  queries that require transitive closure, e.g. rdfs:subClassOf+
+"""
+
+mcp = FastMCP("ubergraph", instructions=SYSTEM_PROMPT)
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> Dependencies:
+    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[Dependencies]:
+    rc: RunContext[Dependencies] = RunContext[Dependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def query_ubergraph(query: str, format: Optional[str] = "text") -> Dict:
+    """
+    Execute a SPARQL query against the UberGraph endpoint.
+    
+    Args:
+        query: The SPARQL query to execute
+        format: Output format (text or json)
+        
+    Returns:
+        The query results
+    """
+    return await ut.query_ubergraph(ctx(), query, format)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')

aurelian/agents/ubergraph/ubergraph_tools.py

@@ -0,0 +1,118 @@
+"""
+Tools for interacting with the UberGraph SPARQL endpoint.
+"""
+import asyncio
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel
+from pydantic_ai import RunContext, ModelRetry
+from SPARQLWrapper import JSON, SPARQLWrapper
+
+from .ubergraph_config import Dependencies, get_config
+
+
+class QueryResults(BaseModel):
+    """Results of a SPARQL query."""
+    results: List[Dict] = []
+
+
+def simplify_value(v: Dict, prefixes=None) -> Any:
+    """
+    Simplify a SPARQL query result value.
+    
+    Args:
+        v: The value to simplify
+        prefixes: Optional mapping of prefixes to expansions
+        
+    Returns:
+        The simplified value
+    """
+    if prefixes and v["type"] == "uri":
+        for prefix, expansion in prefixes.items():
+            if v["value"].startswith(expansion):
+                return f"{prefix}:{v['value'][len(expansion):]}"
+    return v["value"]
+
+
+def simplify_results(results: Dict, prefixes=None, limit=20) -> List[Dict]:
+    """
+    Simplify SPARQL query results.
+    
+    Args:
+        results: The query results to simplify
+        prefixes: Optional mapping of prefixes to expansions
+        limit: Maximum number of results to return
+        
+    Returns:
+        A list of simplified results
+    """
+    rows = []
+    n = 0
+    for r in results["results"]["bindings"]:
+        n += 1
+        if n > limit:
+            break
+        row = {}
+        for k, v in r.items():
+            row[k] = simplify_value(v, prefixes)
+        rows.append(row)
+    return rows
+
+
+async def query_ubergraph(ctx: RunContext[Dependencies], query: str) -> QueryResults:
+    """
+    Performs a SPARQL query over Ubergraph then returns the results as triples.
+
+    Ubergraph is a triplestore that contains many OBO ontologies and precomputed
+    relation graph edges.
+    
+    Args:
+        ctx: The run context
+        query: The SPARQL query to execute
+        
+    Returns:
+        The query results
+    """
+    config = ctx.deps or get_config()
+    prefixes = config.prefixes
+    endpoint = config.endpoint
+    
+    # Add prefixes to query
+    prefixed_query = ""
+    for k, v in prefixes.items():
+        prefixed_query += f"PREFIX {k}: <{v}>\n"
+    prefixed_query += query
+    
+    print("## Query")
+    print(prefixed_query)
+    print("##")
+    
+    try:
+        # Create SPARQL wrapper
+        sw = SPARQLWrapper(endpoint)
+        sw.setQuery(prefixed_query)
+        sw.setReturnFormat(JSON)
+        
+        # Execute the query in a thread pool
+        ret = await asyncio.to_thread(sw.queryAndConvert)
+        
+        # Process the results
+        results = simplify_results(ret, prefixes, limit=config.max_results)
+        print("num results=", len(results))
+        print("results=", results)
+        
+        if not results:
+            raise ModelRetry(f"No results found for SPARQL query. Try refining your query.")
+            
+        return QueryResults(results=results)
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        
+        # Handle specific SPARQL errors
+        if "syntax error" in str(e).lower():
+            raise ModelRetry(f"SPARQL syntax error: {str(e)}")
+        elif "time" in str(e).lower() and "out" in str(e).lower():
+            raise ModelRetry("Query timed out. Try to simplify your query or reduce its scope.")
+        else:
+            raise ModelRetry(f"Error executing SPARQL query: {str(e)}")

aurelian/agents/uniprot/__init__.py

@@ -0,0 +1,37 @@
+"""
+UniProt agent package for interacting with the UniProt database.
+"""
+
+from .uniprot_agent import uniprot_agent, UNIPROT_SYSTEM_PROMPT
+from .uniprot_config import UniprotConfig, get_config
+from .uniprot_gradio import chat
+from .uniprot_mcp import (
+    get_uniprot_mcp_tools,
+    get_uniprot_mcp_messages,
+    handle_uniprot_mcp_request,
+)
+from .uniprot_tools import lookup_uniprot_entry, search, uniprot_mapping, normalize_uniprot_id
+
+__all__ = [
+    # Agent
+    "uniprot_agent",
+    "UNIPROT_SYSTEM_PROMPT",
+    
+    # Config
+    "UniprotConfig",
+    "get_config",
+    
+    # Tools
+    "lookup_uniprot_entry",
+    "search",
+    "uniprot_mapping",
+    "normalize_uniprot_id",
+    
+    # Gradio
+    "chat",
+    
+    # MCP
+    "get_uniprot_mcp_tools",
+    "get_uniprot_mcp_messages",
+    "handle_uniprot_mcp_request",
+]

aurelian/agents/uniprot/uniprot_agent.py

@@ -0,0 +1,43 @@
+"""
+Agent for working with the UniProt database and API.
+"""
+from pydantic_ai import Agent
+
+from .uniprot_config import UniprotConfig, get_config
+from .uniprot_tools import lookup_uniprot_entry, search, uniprot_mapping
+
+# System prompt for the UniProt agent
+UNIPROT_SYSTEM_PROMPT = """
+You are a helpful assistant that specializes in accessing and interpreting information from the UniProt database.
+UniProt is a comprehensive, high-quality resource of protein sequence and functional information.
+
+You can:
+- Search UniProt with queries
+- Look up detailed information about specific proteins using UniProt accession numbers
+- Map UniProt accessions to entries in other databases
+
+When using protein IDs:
+- UniProt accession numbers (e.g., P12345) are stable identifiers for protein entries
+- Some proteins may be referenced by their entry name (e.g., ALBU_HUMAN)
+- UniProt IDs may sometimes include version numbers (e.g., P12345.2) which can be normalized
+
+When returning information about proteins, present it in a clear, organized manner with:
+- Key protein attributes like name, gene, organism, and length
+- Functional information including catalytic activity and pathways
+- Structural information if available
+- Disease associations if relevant
+
+For search results, summarize the key findings and highlight the most relevant matches.
+"""
+
+# Create the agent with the system prompt
+uniprot_agent = Agent(
+    model="openai:gpt-4o",
+    system_prompt=UNIPROT_SYSTEM_PROMPT,
+    deps_type=UniprotConfig,
+)
+
+# Register the tools with the agent
+uniprot_agent.tool(search)
+uniprot_agent.tool(lookup_uniprot_entry)
+uniprot_agent.tool(uniprot_mapping)

aurelian/agents/uniprot/uniprot_config.py

@@ -0,0 +1,43 @@
+"""
+Configuration for the UniProt agent.
+"""
+from dataclasses import dataclass, field
+import os
+from typing import Any, Dict, Optional
+
+from bioservices import UniProt
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+
+
+@dataclass
+class UniprotConfig(HasWorkdir):
+    """Configuration for the UniProt agent."""
+
+    # Options for the bioservices UniProt client
+    uniprot_client_options: Dict[str, Any] = field(default_factory=dict)
+
+    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.uniprot_client_options is None or len(self.uniprot_client_options) == 0:
+            self.uniprot_client_options = {"verbose": False}
+        
+        # Initialize the workdir if not already set
+        if self.workdir is None:
+            self.workdir = WorkDir()
+
+    def get_uniprot_client(self) -> UniProt:
+        """Get a configured UniProt client."""
+        return UniProt(**self.uniprot_client_options)
+
+
+def get_config() -> UniprotConfig:
+    """Get the UniProt configuration from environment variables or defaults."""
+    workdir_path = os.environ.get("AURELIAN_WORKDIR", None)
+    workdir = WorkDir(location=workdir_path) if workdir_path else None
+    
+    return UniprotConfig(
+        workdir=workdir,
+        uniprot_client_options={"verbose": False}
+    )

aurelian/agents/uniprot/uniprot_evals.py

@@ -0,0 +1,99 @@
+"""
+Evaluation module for the UniProt agent.
+
+This module implements evaluations for the UniProt 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.uniprot.uniprot_agent import uniprot_agent
+from aurelian.agents.uniprot.uniprot_config import UniprotConfig
+
+class UniprotMetadata(Dict[str, Any]):
+    """Simple metadata dictionary for UniProt evaluations."""
+    pass
+
+# Define individual evaluation cases
+case1 = Case(
+    name="human_insulin",
+    inputs="Search for human insulin protein",
+    expected_output="P01308",  # Human insulin UniProt ID
+    metadata=metadata("easy", "protein_search")
+)
+
+case2 = Case(
+    name="uniprot_entry_lookup",
+    inputs="Look up UniProt entry P01308",
+    expected_output="insulin",  # Should identify this as insulin
+    metadata=metadata("easy", "id_lookup"),
+    evaluators=[
+        LLMJudge(
+            rubric="""
+            Answer should:
+            1. Correctly identify P01308 as human insulin
+            2. Include key information about the protein's function
+            3. Mention its role in glucose homeostasis
+            4. Provide information about protein structure
+            """,
+            include_input=True
+        )
+    ]
+)
+
+case3 = Case(
+    name="id_mapping",
+    inputs="Map UniProt IDs P01308,P01009 to PDB database",
+    expected_output="PDB",  # Should return PDB IDs
+    metadata=metadata("medium", "database_mapping")
+)
+
+case4 = Case(
+    name="domain_identification",
+    inputs="What domains are present in UniProt entry P53_HUMAN?",
+    expected_output="domain",  # Should discuss protein domains
+    metadata=metadata("medium", "protein_domain_analysis")
+)
+
+case5 = Case(
+    name="disease_association",
+    inputs="Find all proteins related to Alzheimer's disease",
+    expected_output="amyloid",  # Should mention amyloid proteins
+    metadata=metadata("hard", "disease_association_query")
+)
+
+def create_eval_dataset() -> Dataset[str, str, MetadataDict]:
+    """
+    Create a dataset for evaluating the UniProt agent.
+    
+    Returns:
+        Dataset of UniProt 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 protein information provided
+            2. Correct identification of UniProt IDs and cross-references
+            3. Comprehensive coverage of protein structure and function
+            4. Proper description of protein domains and modifications
+            5. Accurate representation of protein-disease associations
+            """,
+            model="anthropic:claude-3-7-sonnet-latest"
+        )
+    ]
+    
+    return Dataset(
+        cases=cases,
+        evaluators=evaluators
+    )