aurelian 0.3.2__py3-none-any.whl

aurelian/agents/uniprot/uniprot_gradio.py

@@ -0,0 +1,48 @@
+"""
+Gradio interface for the UniProt agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.agents.uniprot.uniprot_agent import uniprot_agent
+from aurelian.agents.uniprot.uniprot_config import UniprotConfig
+from aurelian.utils.async_utils import run_sync
+
+
+def chat(deps: Optional[UniprotConfig] = None, **kwargs):
+    """
+    Initialize a chat interface for the UniProt agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        **kwargs: Additional arguments to pass to the agent
+        
+    Returns:
+        A Gradio chat interface
+    """
+    if deps is None:
+        deps = UniprotConfig()
+
+    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: uniprot_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="UniProt AI Assistant",
+        examples=[
+            ["Search for human insulin protein"],
+            ["Look up UniProt entry P01308"],
+            ["Map UniProt IDs P01308,P01009 to PDB database"],
+            ["What domains are present in UniProt entry P53_HUMAN?"],
+            ["Find all proteins related to Alzheimer's disease"]
+        ]
+    )

aurelian/agents/uniprot/uniprot_mcp.py

@@ -0,0 +1,168 @@
+"""
+MCP integration for the UniProt 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 .uniprot_agent import uniprot_agent, UNIPROT_SYSTEM_PROMPT
+from .uniprot_config import UniprotConfig, 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_uniprot_mcp_tools() -> List[Dict[str, Any]]:
+    """Get the MCP tools for the UniProt agent.
+
+    Returns:
+        A list of tools in the MCP format
+    """
+    return tools_to_anthropic_tools(uniprot_agent)
+
+
+def get_uniprot_mcp_messages(
+    messages: Optional[List[Dict[str, Any]]] = None,
+    system: Optional[str] = None,
+) -> List[Dict[str, Any]]:
+    """Get MCP messages for the UniProt agent.
+
+    Args:
+        messages: Previous messages
+        system: System message override
+
+    Returns:
+        List of MCP messages
+    """
+    if messages is None:
+        messages = []
+    
+    if system is None:
+        system = UNIPROT_SYSTEM_PROMPT
+    
+    if not messages or messages[0].get("role") != "system":
+        messages = [{"role": "system", "content": system}] + messages
+    
+    return messages
+
+
+def handle_uniprot_mcp_request(
+    messages: List[Dict[str, Any]],
+    config: Optional[UniprotConfig] = 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 UniProt agent.
+
+    Args:
+        messages: The messages from the client
+        config: The UniProt 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_uniprot_mcp_messages(messages)
+    mcp_tools = get_uniprot_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,
+    )

aurelian/agents/uniprot/uniprot_tools.py

@@ -0,0 +1,136 @@
+"""
+Tools for interacting with the UniProt database.
+"""
+from typing import Dict, List, Optional
+
+from pydantic_ai import RunContext, ModelRetry
+
+from .uniprot_config import UniprotConfig, get_config
+
+
+def normalize_uniprot_id(uniprot_id: str) -> str:
+    """Normalize a Uniprot ID by removing any version number.
+
+    Args:
+        uniprot_id: The Uniprot ID
+
+    Returns:
+        The normalized Uniprot ID
+    """
+    if ":" in uniprot_id:
+        return uniprot_id.split(":")[-1]
+    return uniprot_id
+
+
+def search(ctx: RunContext[UniprotConfig], query: str) -> str:
+    """Search UniProt with a query string.
+
+    Args:
+        ctx: The run context with access to the config
+        query: The search query
+
+    Returns:
+        The search results in TSV format
+    """
+    config = ctx.deps or get_config()
+    u = config.get_uniprot_client()
+    
+    try:
+        results = u.search(query, frmt="tsv", columns="accession,id,gene_names")
+        if not results or results.strip() == "":
+            raise ModelRetry(f"No results found for query: {query}")
+        return results
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error searching UniProt: {str(e)}")
+
+
+def lookup_uniprot_entry(ctx: RunContext[UniprotConfig], uniprot_acc: str) -> str:
+    """Lookup the Uniprot entry for a given Uniprot accession number.
+
+    Args:
+        ctx: The run context with access to the config
+        uniprot_acc: The Uniprot accession
+
+    Returns:
+        The UniProt entry in text format
+    """
+    config = ctx.deps or get_config()
+    u = config.get_uniprot_client()
+    
+    try:
+        uniprot_acc = normalize_uniprot_id(uniprot_acc)
+        result = u.retrieve(uniprot_acc, frmt="txt")
+        if not result or result.strip() == "":
+            raise ModelRetry(f"No entry found for accession: {uniprot_acc}")
+        return result
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error looking up UniProt entry: {str(e)}")
+
+
+def uniprot_mapping(
+    ctx: RunContext[UniprotConfig], 
+    target_database: str, 
+    uniprot_accs: List[str]
+) -> Dict:
+    """Perform a mapping of Uniprot accessions to another database.
+
+    Args:
+        ctx: The run context with access to the config
+        target_database: The target database (e.g., KEGG, PDB)
+        uniprot_accs: The Uniprot accessions
+
+    Returns:
+        A dictionary mapping UniProt accessions to entries in the target database
+    """
+    config = ctx.deps or get_config()
+    u = config.get_uniprot_client()
+    
+    try:
+        if not uniprot_accs:
+            raise ModelRetry("No UniProt accessions provided for mapping")
+            
+        normalized_accs = [normalize_uniprot_id(x) for x in uniprot_accs]
+        result = u.mapping("UniProtKB_AC-ID", target_database, ",".join(normalized_accs))
+        
+        if not result:
+            raise ModelRetry(f"No mappings found for accessions to {target_database}")
+        return result
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error mapping UniProt entries: {str(e)}")
+
+def map_to_uniprot(ctx: RunContext[UniprotConfig], external_ids: List[str]) -> Dict:
+    """Map Uniprot accessions to UniProt IDs.
+
+    Args:
+        ctx: The run context with access to the config
+        external_ids: The external IDs, as prefixed IDs
+
+    Returns:
+        A dictionary mapping external IDs to UniProt IDs
+    """
+    config = ctx.deps or get_config()
+    u = config.get_uniprot_client()
+    results = {}
+    # split the external IDs by prefix
+    ids_by_prefix = {}  
+    for external_id in external_ids:
+        prefix, id = external_id.split(":")
+        if prefix not in ids_by_prefix:
+            ids_by_prefix[prefix] = []
+        ids_by_prefix[prefix].append(id)
+    for prefix, ids in ids_by_prefix.items():
+        result = u.mapping(prefix, "UniProtKB_AC-ID", ",".join(ids))
+        if 'results' in result:
+            print(result)
+            uniprot_ids = [entry['to'] for entry in result['results']]
+            results[prefix] = uniprot_ids
+        else:
+            print("No mapping found.")
+    return results
+

aurelian/agents/web/web_config.py

@@ -0,0 +1,27 @@
+"""
+Configuration for the Web agent.
+"""
+import os
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+
+
+class WebDependencies(HasWorkdir):
+    """
+    Web 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() -> WebDependencies:
+    """Get a default configuration for the Web agent.
+    
+    Returns:
+        WebDependencies object with default settings
+    """
+    workdir_path = os.environ.get("AURELIAN_WORKDIR", None)
+    workdir = WorkDir(location=workdir_path) if workdir_path else None
+
+    return WebDependencies(workdir=workdir)

aurelian/agents/web/web_gradio.py

@@ -0,0 +1,48 @@
+"""
+Gradio UI for the Web agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.dependencies.workdir import HasWorkdir
+from aurelian.agents.web.web_mcp import mcp
+from aurelian.utils.async_utils import run_sync
+
+
+def chat(deps: Optional[HasWorkdir] = None, **kwargs):
+    """
+    Initialize a chat interface for the Web agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        **kwargs: Additional arguments to pass to the agent
+        
+    Returns:
+        A Gradio chat interface
+    """
+    if deps is None:
+        from aurelian.agents.web.web_mcp import deps as get_deps
+        deps = get_deps()
+
+    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}"
+        # Use MCP for the agent
+        result = run_sync(lambda: mcp.query(query))
+        return result
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="Web Search AI Assistant",
+        examples=[
+            ["Search for information about gene ontology"],
+            ["Find recent research papers on CRISPR"],
+            ["Look up information about linkML data modeling"]
+        ]
+    )

aurelian/agents/web/web_mcp.py

@@ -0,0 +1,50 @@
+"""
+MCP tools for web search operations.
+"""
+import os
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.web.web_tools as wt
+from aurelian.dependencies.workdir import WorkDir, HasWorkdir
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("web", instructions="Web search operations")
+
+
+def deps() -> HasWorkdir:
+    deps = HasWorkdir()
+    # Set the location from environment variable or default
+    loc = os.getenv("AURELIAN_WORKDIR", "/tmp/aurelian")
+    deps.workdir = WorkDir(loc)
+    return deps
+
+def ctx() -> RunContext[HasWorkdir]:
+    rc: RunContext[HasWorkdir] = RunContext[HasWorkdir](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def search_web(query: str) -> str:
+    """
+    Search the web using a text query.
+
+    Note, this will not retrieve the full content, for that you
+    should use retrieve_web_page tool.
+
+    Args:
+        query: Text query
+
+    Returns: 
+        Matching web pages plus summaries
+    """
+    return await wt.search_web(query)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')

aurelian/agents/web/web_tools.py

@@ -0,0 +1,108 @@
+from typing import List, Optional
+import os
+import xml.etree.ElementTree as ET
+from pydantic import BaseModel, Field
+
+from aurelian.utils.search_utils import web_search
+from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIModel
+
+
+class Citation(BaseModel):
+    """Citation model for query results with citations."""
+    
+    citation_number: str = Field(description="Citation number or identifier")
+    url: str = Field(description="URL source of the citation") 
+    timestamp: Optional[str] = Field(
+        default=None, description="Timestamp of the citation, if available"
+    )
+
+
+class ResultWithCitations(BaseModel):
+    """Generic result model for queries that return content with citations."""
+    
+    content: str = Field(description="Content of the response")
+    citations: List[Citation] = Field(
+        default_factory=list, description="List of citations referenced in the content"
+    )
+
+
+async def search_web(query: str) -> str:
+    """
+    Search the web using a text query.
+
+    Note, this will not retrieve the full content, for that you
+    should use `retrieve_web_page`.
+
+    Args:
+        query: Text query
+
+    Returns: matching web pages plus summaries
+    """
+    print(f"Web Search: {query}")
+    return web_search(query)
+
+
+async def perplexity_query(query: str, model_name: str = "sonar") -> ResultWithCitations:
+    """
+    Query the Perplexity API and return structured results with citations.
+    
+    Args:
+        query: The query to send to Perplexity
+        model_name: The Perplexity model to use (default: "sonar"). Options include "sonar", "sonar-pro", etc.
+        
+    Returns:
+        ResultWithCitations: A structured response with content and citations
+        
+    Raises:
+        ValueError: If the Perplexity API key is not set
+        RuntimeError: If the response parsing fails
+    """
+    perplexity_api_key = os.environ.get("PERPLEXITY_API_KEY")
+    if not perplexity_api_key:
+        raise ValueError("PERPLEXITY_API_KEY environment variable is not set")
+        
+    sonar_model = OpenAIModel(
+        model_name=model_name,
+        base_url='https://api.perplexity.ai',
+        api_key=perplexity_api_key,
+    )
+    
+    agent = Agent(sonar_model,
+                  system_prompt=("Return the response with xml tags <answer> and <citations>,"
+                                 "where citations includes a list with the citation_number, url,"
+                                 " and timestamp for the retrieved citations"))
+    result = agent.run_sync(query)
+    
+    try:
+        # Parse the XML response
+        xml_string = f"<root>{result.data}</root>"
+        root = ET.fromstring(xml_string)
+        
+        # Extract answer content
+        answer_element = root.find("answer")
+        content = answer_element.text.strip() if answer_element is not None and answer_element.text else ""
+        
+        # Extract citations
+        citations_list = []
+        citations_element = root.find("citations")
+        if citations_element is not None:
+            for citation in citations_element.findall(".//citation"):
+                citation_number = citation.find("citation_number")
+                if citation_number is None:
+                    citation_number = citation.find("number")
+                
+                url_element = citation.find("url")
+                timestamp_element = citation.find("timestamp")
+                
+                citations_list.append(
+                    Citation(
+                        citation_number=citation_number.text if citation_number is not None and citation_number.text else "",
+                        url=url_element.text if url_element is not None and url_element.text else "",
+                        timestamp=timestamp_element.text if timestamp_element is not None and timestamp_element.text else None
+                    )
+                )
+        
+        return ResultWithCitations(content=content, citations=citations_list)
+    except ET.ParseError as e:
+        raise RuntimeError(f"Failed to parse Perplexity response: {e}")

aurelian/chat.py

@@ -0,0 +1,23 @@
+from typing import Callable
+
+from aurelian.utils.async_utils import run_sync
+
+
+def get_chatbot(query_func: Callable, **kwargs):
+    import gradio as gr
+
+    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(query_func)
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        **kwargs,
+    )