aurelian 0.3.2__py3-none-any.whl

aurelian/agents/robot/robot_ontology_agent.py

@@ -0,0 +1,139 @@
+"""
+Agent for creating ROBOT templates and compiling to ontologies.
+"""
+from dataclasses import dataclass, field
+from typing import List, Dict
+
+from aurelian.agents.filesystem.filesystem_tools import inspect_file, download_url_as_markdown, list_files
+from aurelian.agents.robot.robot_config import RobotDependencies
+from aurelian.agents.robot.robot_tools import write_and_compile_template, fetch_documentation
+from aurelian.utils.async_utils import run_sync
+from aurelian.utils.search_utils import web_search
+from pydantic_ai import Agent, RunContext, Tool
+
+from aurelian.dependencies.workdir import WorkDir, HasWorkdir
+
+SYSTEM = """
+Background:
+
+Your job is to iteratively build an ontology via *robot templates*,
+These are tabular data (CSV syntax) with a special header that compiles to OWL.
+
+For example, if the request is for an animals ontology, you could start with a CSV with columns Name, ParentTaxon, Eats, 
+with rows filled out with some example animals.
+
+The main tool you will use is `write_and_compile_template` which writes the the template content to
+a local file after compiling to OWL. This also takes a list of ontologies to import, which
+should also be on the file system.
+
+Sometimes you may need to work with multiple dependent ontologies. For example, if you have a vehicle class
+hierarchy in `vehicles.csv` and parts in `parts.csv`, and vehicles depends on parts, you would first iterate
+on `parts.csv` (e.g. calling `write_and_compile_template`, with no imports), then write vehicles using
+`write_and_compile_template` with `['parts.csv']` as the dependencies/imports.
+
+## Robot template CSV structure:
+
+Robot template files have an additional metadata row below the header row. This is called the "template row". It specifies how each column maps to OWL. Typical values will be "ID" for the unique identifier, LABEL for the name, "SC %" for the parent class. Consult the docs for details. Note that this is always beneath the main header row. This can seem a bit duplicative, but that's OK. An example might be:
+
+identifier,name,parent,synonyms
+ID,LABEL,SC %,A oboInOwl:hasExactSynonym
+ANIMAL:1,chicken,aves,Gallus gallus|chick
+
+The first row is a normal header with human-friendly columns. The 2nd is the robot template row. After that are the usual data rows.
+
+Here "A oboInOwl:hasExactSynonym" in the template row for "synonyms" indicate this column should be interpreted as an owl annotation using that property. Generally the value for annotations is literals/text.
+
+Another common piece of metadata is definitions. For OBO ontologies, IAO must be used here. For non-OBO ontologies people may want to use skos
+
+Some ontologies may need to use other relationships. For part-of parents, use "'part of' some %" (this means that the class indicated by the ID is part-of some X, where X is the value in the part-of column). Use other relationships as appropriate. If you are unclear about the semantics, then consult the docs. You can also work through the docs with the user.
+
+Note that any terms referenced as parents or in logical axioms such as part-of should be in the ontology, so make sure they have rows in the CSV.  It's OK to use the label. For example:
+
+identifier,primary_name,parent,madeOf
+ID,LABEL,SC %,SC 'made of' some %
+VON:1,vehicle,,
+VON:2.car,vehicle,wheel|chassis
+VON:3,wheel,car part,
+VON:4,chassis,car part,
+
+If in doubt, use "A <propertyName>" for metadata and "SC '<relationName>' some %" for logical relationships / graph edges.
+
+If your working dir doesn't contain any object or annotation properties you can make them in a seperate
+imported ontology, TYPE is useful for determining the OWL type, for example:
+
+```
+ID,Label,Type,Definition
+ID,LABEL,TYPE,A IAO:0000115
+IAO:0000115,definition,owl:AnnotationProperty
+BFO:0000050,part_of,owl:ObjectProperty
+```
+
+If you need any more detailed documentation, you can fetch it with `fetch_documentation`
+
+You can look at files with `inspect_file`
+
+Use scientific language as far as possible. For IDs, these should be numeric curies unless the user requests otherwise. If the user wants to substitute actual ontology term IDs for these, use lookup_curies_get_lookup_get 
+"""
+
+
+
+
+robot_ontology_agent = Agent(
+    model="openai:gpt-4o",
+    deps_type=RobotDependencies,
+    system_prompt=SYSTEM,
+    tools=[
+        Tool(write_and_compile_template, max_retries=2),
+        Tool(fetch_documentation),
+        Tool(inspect_file),
+        Tool(list_files),
+        Tool(download_url_as_markdown),
+    ]
+)
+
+
+@robot_ontology_agent.system_prompt
+def include_templates_in_prompt(ctx: RunContext[RobotDependencies]) -> str:
+    files_names = ctx.deps.workdir.list_file_names()
+    s = "Working directory files/templates:"
+    if files_names:
+        for f in files_names:
+            s += f"{f}\n---"
+            s += ctx.deps.workdir.read_file(f)
+            s += "\n"
+    return s
+
+
+@robot_ontology_agent.system_prompt
+def include_prefixes_in_prompt(ctx: RunContext[RobotDependencies]) -> str:
+    pmap = ctx.deps.prefix_map
+    return f"Prefixes: {pmap}"
+
+
+
+
+
+
+def chat(workdir: str, **kwargs):
+    import gradio as gr
+    deps = RobotDependencies()
+    deps.workdir.location = workdir
+
+    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: robot_ontology_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="robot AI Assistant",
+        examples=[
+            ["Create an ontology of snacks"],
+        ]
+    )

aurelian/agents/robot/robot_tools.py

@@ -0,0 +1,50 @@
+from typing import Optional, List
+
+from pydantic_ai import RunContext, ModelRetry
+
+from aurelian.agents.robot.assets import ROBOT_ONTOLOGY_AGENT_CONTENTS_DIR
+from aurelian.agents.robot.robot_config import RobotDependencies
+from aurelian.utils.robot_ontology_utils import run_robot_template_command
+
+
+async def write_and_compile_template(ctx: RunContext[RobotDependencies], template: str, save_to_file: str= "core.csv", import_ontologies: Optional[List[str]] = None) -> str:
+    """
+    Adds a template to the file system and compile it to OWL
+
+    Args:
+        ctx: context
+        template: robot template as string. Do not truncate, always pass the whole template, including header.
+        save_to_file: file name to save the templates to. Defaults to core.csv. Only written if file compiles to OWL
+        import_ontologies: list of ontologies to import. These should be files in the working directory.
+
+    Returns:
+
+    """
+    print(f"Validating template: {template}")
+    try:
+        ctx.deps.workdir.write_file(save_to_file, template)
+        output_path = run_robot_template_command(
+            ctx.deps.workdir,
+            save_to_file,
+            import_ontologies=import_ontologies,
+            prefix_map=ctx.deps.prefix_map,
+            output_path=None,
+        ),
+        if save_to_file and template:
+            ctx.deps.workdir.write_file(save_to_file, template)
+    except Exception as e:
+        raise ModelRetry(f"Template does not compile: {e}")
+    return f"Template compiled to {output_path}"
+
+
+async def fetch_documentation(ctx: RunContext[RobotDependencies]) -> str:
+    """
+    Fetch the documentation for the robot ontology agent.
+
+    Args:
+        ctx: context
+
+    Returns:
+        str: documentation
+    """
+    return open(ROBOT_ONTOLOGY_AGENT_CONTENTS_DIR / "template.md").read()

aurelian/agents/talisman/__init__.py

@@ -0,0 +1,3 @@
+"""
+Talisman agent package for advanced gene analysis using UniProt and NCBI Entrez.
+"""

aurelian/agents/talisman/talisman_agent.py

@@ -0,0 +1,126 @@
+"""
+Agent for working with gene information using the UniProt API and NCBI Entrez.
+"""
+from pydantic_ai import Agent
+
+from .talisman_config import TalismanConfig, get_config
+from .talisman_tools import (
+    get_gene_description, 
+    get_gene_descriptions, 
+    get_genes_from_list, 
+    analyze_gene_set
+)
+
+# System prompt for the Talisman agent
+TALISMAN_SYSTEM_PROMPT = """
+You are a helpful assistant that specializes in providing gene information using both UniProt and NCBI Entrez.
+You can retrieve gene descriptions for single gene IDs or lists of multiple gene IDs, and analyze gene sets to identify functional relationships.
+
+You can:
+- Get a description for a single gene using its identifier
+- Get descriptions for multiple genes using a list of identifiers
+- Parse a string containing multiple gene identifiers in various formats
+- Analyze sets of genes to identify biological relationships and shared properties
+
+Gene identifiers can be:
+- UniProt accession numbers (e.g., P12345, Q934N0)
+- Gene symbols (e.g., INS, TP53, mmoX)
+- Gene names
+- Ensembl IDs (e.g., ENSG00000139618)
+- NCBI protein IDs (e.g., 8YJT_C2)
+
+IMPORTANT: The system handles both gene symbols and UniProt identifiers intelligently:
+1. When a gene symbol is provided, the system will:
+   - First attempt to look up the corresponding UniProt protein accession
+   - Search UniProt for detailed information
+   - Fall back to NCBI Entrez if the gene is not found in UniProt
+
+2. When a UniProt ID is provided, the system will:
+   - Directly retrieve the information from UniProt
+   - Fall back to NCBI Entrez if needed
+
+3. When a protein ID is provided, the system will:
+   - Search the protein database in NCBI Entrez
+   - Return detailed protein information
+
+This multi-database approach ensures comprehensive coverage of gene and protein information.
+
+When returning gene information, present it in a clear, organized manner with:
+- Gene name and symbol
+- Description of gene function
+- Organism information when available
+- Disease associations if relevant
+- Source database (UniProt, NCBI, or both)
+
+For multiple genes, organize the results in a table format for easy readability.
+
+For gene set analysis, you will receive a detailed summary of:
+- Shared biological processes these genes may participate in
+- Potential protein-protein interactions or functional relationships
+- Common cellular localization patterns
+- Involvement in similar pathways
+- Coordinated activities or cooperative functions
+- Any disease associations that multiple genes in this set share
+
+The analysis will cover multiple types of relationships:
+- Functional relationships
+- Pathway relationships
+- Regulatory relationships
+- Localization patterns
+- Physical interactions
+- Genetic interactions
+
+IMPORTANT: For gene set analysis, ALWAYS include a distinct section titled "## Terms" 
+that contains a semicolon-delimited list of functional terms relevant to the gene set, 
+ordered by relevance. These terms should include:
+- Gene Ontology biological process terms (e.g., DNA repair, oxidative phosphorylation, signal transduction)
+- Molecular function terms (e.g., kinase activity, DNA binding, transporter activity)
+- Cellular component/localization terms (e.g., nucleus, plasma membrane, mitochondria)
+- Pathway names (e.g., glycolysis, TCA cycle, MAPK signaling)
+- Co-regulation terms (e.g., stress response regulon, heat shock response)
+- Interaction networks (e.g., protein complex formation, signaling cascade)
+- Metabolic process terms (e.g., fatty acid synthesis, amino acid metabolism)
+- Regulatory mechanisms (e.g., transcriptional regulation, post-translational modification)
+- Disease associations (if relevant, e.g., virulence, pathogenesis, antibiotic resistance)
+- Structural and functional domains/motifs (e.g., helix-turn-helix, zinc finger)
+
+Example of Terms section:
+## Terms
+DNA damage response; p53 signaling pathway; apoptosis; cell cycle regulation; tumor suppression; DNA repair; protein ubiquitination; transcriptional regulation; nuclear localization; cancer predisposition
+
+IMPORTANT: After the Terms section, ALWAYS include a "## Gene Summary Table" with a markdown table 
+summarizing the genes analyzed, with the following columns in this exact order:
+- ID: The gene identifier (same as Gene Symbol)
+- Annotation: Genomic coordinates or accession with position information
+- Genomic Context: Information about the genomic location (chromosome, plasmid, etc.)
+- Organism: The organism the gene belongs to
+- Description: The protein/gene function description
+
+Example of Gene Summary Table:
+## Gene Summary Table
+| ID | Annotation | Genomic Context | Organism | Description |
+|-------------|-------------|----------|----------------|------------|
+| BRCA1 | NC_000017.11 (43044295..43125483) | Chromosome 17 | Homo sapiens | Breast cancer type 1 susceptibility protein |
+| TP53 | NC_000017.11 (7668402..7687550) | Chromosome 17 | Homo sapiens | Tumor suppressor protein |
+
+For bacterial genes, the table format would be:
+| ID | Annotation | Genomic Context | Organism | Description |
+|-------------|-------------|----------|----------------|------------|
+| invA | NC_003197.2 (3038407..3040471, complement) | Chromosome | Salmonella enterica | Invasion protein |
+| DVUA0001 | NC_005863.1 (699..872, complement) | Plasmid pDV | Desulfovibrio vulgaris str. Hildenborough | Hypothetical protein |
+
+Use this information to help researchers understand the potential functional relationships between genes.
+"""
+
+# Create the agent with the system prompt
+talisman_agent = Agent(
+    model="openai:gpt-4o",
+    system_prompt=TALISMAN_SYSTEM_PROMPT,
+    deps_type=TalismanConfig,
+)
+
+# Register the tools with the agent
+talisman_agent.tool(get_gene_description)
+talisman_agent.tool(get_gene_descriptions)
+talisman_agent.tool(get_genes_from_list)
+talisman_agent.tool(analyze_gene_set)

aurelian/agents/talisman/talisman_config.py

@@ -0,0 +1,66 @@
+"""
+Configuration for the Talisman agent.
+"""
+from dataclasses import dataclass, field
+import os
+from typing import Any, Dict, Optional
+
+from bioservices import UniProt
+from bioservices.eutils import EUtils as NCBI
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+
+
+@dataclass
+class TalismanConfig(HasWorkdir):
+    """Configuration for the Talisman agent."""
+
+    # Options for the bioservices UniProt client
+    uniprot_client_options: Dict[str, Any] = field(default_factory=dict)
+    
+    # Options for the bioservices NCBI client
+    ncbi_client_options: Dict[str, Any] = field(default_factory=dict)
+    
+    # OpenAI API key for LLM-based analysis
+    openai_api_key: Optional[str] = None
+    
+    # Model to use for gene set analysis
+    model_name: str = "gpt-4o"
+
+    def __post_init__(self):
+        """Initialize the config with default values."""
+        # Initialize with default options if none provided
+        if self.uniprot_client_options is None or len(self.uniprot_client_options) == 0:
+            self.uniprot_client_options = {"verbose": False}
+            
+        if self.ncbi_client_options is None or len(self.ncbi_client_options) == 0:
+            self.ncbi_client_options = {"verbose": False, "email": "MJoachimiak@lbl.gov"}
+        
+        # Initialize the workdir if not already set
+        if self.workdir is None:
+            self.workdir = WorkDir()
+        
+        # Try to get OpenAI API key from environment if not provided
+        if self.openai_api_key is None:
+            import os
+            self.openai_api_key = os.environ.get("OPENAI_API_KEY")
+
+    def get_uniprot_client(self) -> UniProt:
+        """Get a configured UniProt client."""
+        return UniProt(**self.uniprot_client_options)
+        
+    def get_ncbi_client(self) -> NCBI:
+        """Get a configured NCBI client."""
+        return NCBI(**self.ncbi_client_options)
+
+
+def get_config() -> TalismanConfig:
+    """Get the Talisman 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 TalismanConfig(
+        workdir=workdir,
+        uniprot_client_options={"verbose": False},
+        ncbi_client_options={"verbose": False}
+    )

aurelian/agents/talisman/talisman_gradio.py

@@ -0,0 +1,50 @@
+"""
+Gradio interface for the Talisman agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.agents.talisman.talisman_agent import talisman_agent
+from aurelian.agents.talisman.talisman_config import TalismanConfig
+from aurelian.utils.async_utils import run_sync
+
+
+def chat(deps: Optional[TalismanConfig] = None, **kwargs):
+    """
+    Initialize a chat interface for the Talisman agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        **kwargs: Additional arguments to pass to the agent
+        
+    Returns:
+        A Gradio chat interface
+    """
+    if deps is None:
+        deps = TalismanConfig()
+
+    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: talisman_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="Talisman Gene Analysis Assistant",
+        examples=[
+            ["Get description for TP53"],
+            ["Get information about the BRCA1 gene"],
+            ["Get descriptions for multiple genes: INS, ALB, APOE"],
+            ["What is the function of KRAS?"],
+            ["Analyze the relationship between TP53 and MDM2"],
+            ["Analyze these genes and their functional relationships: BRCA1, BRCA2, ATM, PARP1"],
+            ["Get descriptions for ENSG00000139618, ENSG00000141510"]
+        ]
+    )

aurelian/agents/talisman/talisman_mcp.py

@@ -0,0 +1,168 @@
+"""
+MCP integration for the Talisman agent.
+"""
+import json
+import os
+from typing import Dict, Any, List, Optional
+
+from anthropic import Anthropic
+from pydantic_ai import Agent
+
+# Import directly from the MCP module
+from .talisman_agent import talisman_agent, TALISMAN_SYSTEM_PROMPT
+from .talisman_config import TalismanConfig, get_config
+
+
+def tools_to_anthropic_tools(agent: Agent) -> List[Dict[str, Any]]:
+    """Convert pydantic-ai Agent tools to Anthropic tools format.
+    
+    Args:
+        agent: The pydantic-ai Agent with tools
+        
+    Returns:
+        A list of tools in the Anthropic format
+    """
+    anthropic_tools = []
+    for tool in agent.tools:
+        anthropic_tool = {
+            "name": tool.name,
+            "description": tool.description,
+            "input_schema": {
+                "type": "object",
+                "properties": {},
+                "required": []
+            }
+        }
+        
+        for param_name, param in tool.params.items():
+            if param_name == "ctx":
+                continue
+                
+            param_schema = {"type": "string"}
+            if param.annotation.__origin__ == list:
+                param_schema = {
+                    "type": "array",
+                    "items": {"type": "string"}
+                }
+                
+            anthropic_tool["input_schema"]["properties"][param_name] = param_schema
+            if param.default == param.empty:
+                anthropic_tool["input_schema"]["required"].append(param_name)
+                
+        anthropic_tools.append(anthropic_tool)
+        
+    return anthropic_tools
+
+
+def create_bedrock_compatible_resp(
+    response_id: str, 
+    model: str, 
+    role: str, 
+    content: List[Dict[str, Any]]
+) -> Dict[str, Any]:
+    """Create a response compatible with the Bedrock format.
+    
+    Args:
+        response_id: The response ID
+        model: The model name
+        role: The role (usually 'assistant')
+        content: The content from the Anthropic response
+        
+    Returns:
+        A response in the Bedrock format
+    """
+    return {
+        "id": response_id,
+        "model": model,
+        "choices": [
+            {
+                "message": {
+                    "role": role,
+                    "content": content
+                }
+            }
+        ]
+    }
+
+
+def get_talisman_mcp_tools() -> List[Dict[str, Any]]:
+    """Get the MCP tools for the Talisman agent.
+
+    Returns:
+        A list of tools in the MCP format
+    """
+    return tools_to_anthropic_tools(talisman_agent)
+
+
+def get_talisman_mcp_messages(
+    messages: Optional[List[Dict[str, Any]]] = None,
+    system: Optional[str] = None,
+) -> List[Dict[str, Any]]:
+    """Get MCP messages for the Talisman agent.
+
+    Args:
+        messages: Previous messages
+        system: System message override
+
+    Returns:
+        List of MCP messages
+    """
+    if messages is None:
+        messages = []
+    
+    if system is None:
+        system = TALISMAN_SYSTEM_PROMPT
+    
+    if not messages or messages[0].get("role") != "system":
+        messages = [{"role": "system", "content": system}] + messages
+    
+    return messages
+
+
+def handle_talisman_mcp_request(
+    messages: List[Dict[str, Any]],
+    config: Optional[TalismanConfig] = None,
+    model: str = "claude-3-5-sonnet-20240620",
+    max_tokens: int = 4096,
+    temperature: float = 0.0,
+) -> Dict[str, Any]:
+    """Handle an MCP request for the Talisman agent.
+
+    Args:
+        messages: The messages from the client
+        config: The Talisman configuration (optional)
+        model: The model to use
+        max_tokens: The maximum number of tokens to generate
+        temperature: The temperature to use for generation
+
+    Returns:
+        An MCP response
+    """
+    config = config or get_config()
+    
+    api_key = os.environ.get("ANTHROPIC_API_KEY")
+    if not api_key:
+        raise ValueError("ANTHROPIC_API_KEY environment variable is required for MCP")
+
+    client = Anthropic(api_key=api_key)
+    
+    # Prepare the messages and tools
+    mcp_messages = get_talisman_mcp_messages(messages)
+    mcp_tools = get_talisman_mcp_tools()
+    
+    # Send the request to Anthropic
+    response = client.messages.create(
+        model=model,
+        messages=mcp_messages,
+        tools=mcp_tools,
+        max_tokens=max_tokens,
+        temperature=temperature,
+    )
+
+    # Return the response in a format compatible with the MCP protocol
+    return create_bedrock_compatible_resp(
+        response_id=response.id,
+        model=model,
+        role="assistant",
+        content=response.content,
+    )