aurelian 0.3.2__py3-none-any.whl

aurelian/agents/diagnosis/diagnosis_tools.py

@@ -0,0 +1,204 @@
+"""
+Tools for the Diagnosis agent.
+"""
+import asyncio
+from functools import lru_cache
+from typing import Dict, List, Optional
+
+from oaklib import get_adapter
+from pydantic_ai import RunContext, ModelRetry
+
+from aurelian.utils.data_utils import obj_to_dict
+from aurelian.utils.ontology_utils import search_ontology
+from aurelian.utils.search_utils import web_search, retrieve_web_page as fetch_web_page
+from .diagnosis_config import DiagnosisDependencies, get_config, HAS_PHENOTYPE
+
+
+@lru_cache
+def get_mondo_adapter():
+    """
+    Get the MONDO ontology adapter.
+    
+    Returns:
+        The MONDO ontology adapter from OAK
+    """
+    return get_adapter("sqlite:obo:mondo")
+
+
+async def find_disease_id(
+    ctx: RunContext[DiagnosisDependencies], 
+    query: str
+) -> List[Dict]:
+    """
+    Finds the disease ID for a given search term.
+
+    OAK search term syntax is used; the default strategy is to match
+    labels:
+
+        ```
+        find_disease_id("Dravet syndrome")
+        ```
+
+    You can use OAK expressions, e.g, all labels
+    that start with "Peroxisomal biogenesis disorder":
+
+        ```
+        find_disease_id("l^Peroxisomal biogenesis disorder")
+        ```
+
+    Args:
+        ctx: The run context
+        query: The label search term to use
+        
+    Returns:
+        List of matching disease IDs and names
+    """
+    print(f"Disease Search: {query}")
+    
+    try:
+        config = ctx.deps or get_config()
+        adapter = get_mondo_adapter()
+        
+        # Execute the potentially blocking operation in a thread pool
+        results = await asyncio.to_thread(
+            search_ontology, 
+            adapter, 
+            query, 
+            limit=config.max_search_results
+        )
+        
+        if not results:
+            print(f"No results for query: {query} using {adapter}")
+            raise ModelRetry(
+                f"No disease IDs found for query: {query}. Try a different search term."
+            )
+        print(f"Got {len(results)} results for {query}")
+            
+        return results
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error finding disease ID: {str(e)}")
+
+
+async def find_disease_phenotypes(
+    ctx: RunContext[DiagnosisDependencies], 
+    query: str
+) -> List[Dict]:
+    """
+    Finds the phenotypes for a disease ID.
+
+    Example:
+
+        ```
+        find_disease_phenotypes("MONDO:0007947")
+        ```
+
+    Args:
+        ctx: The run context
+        query: The disease ID to search for (should be an ID but can be a name)
+        
+    Returns:
+        List of phenotypes for the disease
+    """
+    print(f"Phenotype query: {query}")
+    
+    try:
+        config = ctx.deps or get_config()
+        
+        # Determine if we have a disease ID or need to search for one
+        if ":" in query:
+            query_ids = [query]
+        else:
+            # Find the disease ID from the name
+            disease_results = await find_disease_id(ctx, query)
+            if not disease_results:
+                raise ModelRetry(f"Could not find disease for query: {query}")
+                
+            # Extract just the IDs from the results
+            query_ids = [result.get("id") for result in disease_results if "id" in result]
+            if not query_ids:
+                raise ModelRetry(f"Could not find valid disease IDs for query: {query}")
+        
+        # Get the phenotype associations
+        monarch_adapter = config.monarch_adapter
+        
+        # Execute the potentially blocking operation in a thread pool
+        assocs = await asyncio.to_thread(
+            monarch_adapter.associations, 
+            subjects=query_ids, 
+            predicates=[HAS_PHENOTYPE]
+        )
+        
+        # Convert to dictionaries
+        results = []
+        for assoc in assocs:
+            dict_assoc = await asyncio.to_thread(obj_to_dict, assoc)
+            results.append(dict_assoc)
+        
+        if not results:
+            disease_label = query
+            if query_ids and query_ids[0] != query:
+                disease_label = f"{query} ({query_ids[0]})"
+            raise ModelRetry(f"No phenotypes found for disease: {disease_label}")
+            
+        print(f"Results[{query_ids}]: {results}")
+        return results
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error finding disease phenotypes: {str(e)}")
+
+
+async def search_web(query: str) -> str:
+    """
+    Search the web using a text query.
+
+    Note: This will not retrieve the full content. For that, use `retrieve_web_page`.
+
+    Args:
+        query: The search query
+        
+    Returns:
+        Matching web pages plus summaries
+    """
+    print(f"Web Search: {query}")
+    
+    try:
+        # Execute the potentially blocking operation in a thread pool
+        results = await asyncio.to_thread(web_search, query)
+        
+        if not results or results.strip() == "":
+            raise ModelRetry(f"No web search results found for query: {query}")
+            
+        return results
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error searching the web: {str(e)}")
+
+
+async def retrieve_web_page(url: str) -> str:
+    """
+    Fetch the contents of a web page.
+
+    Args:
+        url: The URL of the web page to retrieve
+        
+    Returns:
+        The contents of the web page
+    """
+    print(f"Fetch URL: {url}")
+    
+    try:
+        # Execute the potentially blocking operation in a thread pool
+        content = await asyncio.to_thread(fetch_web_page, url)
+        
+        if not content or content.strip() == "":
+            raise ModelRetry(f"No content found at URL: {url}")
+            
+        return content
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error retrieving web page: {str(e)}")

aurelian/agents/diagnosis_agent.py

@@ -0,0 +1,28 @@
+"""
+Agent for performing diagnoses, validated against Monarch KG - DEPRECATED
+
+This module is maintained for backward compatibility.
+Please use aurelian.agents.diagnosis.diagnosis_agent instead.
+"""
+
+from aurelian.agents.diagnosis.diagnosis_agent import diagnosis_agent
+from aurelian.agents.diagnosis.diagnosis_config import DiagnosisDependencies, get_config
+from aurelian.agents.diagnosis.diagnosis_gradio import chat
+from aurelian.agents.diagnosis.diagnosis_tools import (
+    find_disease_id,
+    find_disease_phenotypes,
+    search_web,
+    retrieve_web_page,
+    get_mondo_adapter,
+)
+
+__all__ = [
+    "diagnosis_agent",
+    "DiagnosisDependencies",
+    "chat",
+    "find_disease_id",
+    "find_disease_phenotypes",
+    "search_web",
+    "retrieve_web_page",
+    "get_mondo_adapter",
+]

aurelian/agents/draw/__init__.py

@@ -0,0 +1,3 @@
+"""
+Draw Agent module for creating SVG drawings.
+"""

aurelian/agents/draw/draw_agent.py

@@ -0,0 +1,39 @@
+"""
+Agent for creating SVG drawings based on text descriptions.
+"""
+from pydantic import BaseModel
+
+from aurelian.agents.draw.draw_config import DrawDependencies
+from aurelian.agents.draw.draw_tools import (
+    judge_drawing, DrawingFeedback
+)
+from pydantic_ai import Agent, Tool
+
+SYSTEM = """
+You are an expert scientific artist specializing in creating clear illustrations and figures.
+
+When creating SVG drawings:
+1. Focus on clarity and simplicity over excessive detail
+2. Use appropriate shapes, lines, and basic colors
+3. Ensure the drawing is recognizable and representative of the description
+4. Use valid SVG syntax with width and height attributes
+
+ALWAYS `judge_drawing` to get feedback on your drawings, iterate on them, and improve clarity.
+Even if you think you have the correct drawing, you MUST call this AT LEAST once to get a second
+opinion and to make sure it renders OK.
+"""
+
+class SVGDrawing(BaseModel):
+    svg_content: str
+    legend: str
+    feedback: DrawingFeedback
+
+draw_agent = Agent(
+    model="openai:gpt-4o",
+    deps_type=DrawDependencies,
+    system_prompt=SYSTEM,
+    result_type=SVGDrawing,
+    tools=[
+        Tool(judge_drawing),
+    ]
+)

aurelian/agents/draw/draw_config.py

@@ -0,0 +1,26 @@
+"""
+Configuration classes for the draw agent.
+"""
+from dataclasses import dataclass
+from typing import Optional
+
+from aurelian.dependencies.workdir import HasWorkdir
+
+
+@dataclass
+class DrawDependencies(HasWorkdir):
+    """
+    Configuration for the draw agent.
+    """
+    max_svg_size: int = 1024 * 1024  # 1MB max SVG size
+    judge_feedback: bool = True  # Whether to get judge feedback
+
+
+def get_config() -> DrawDependencies:
+    """
+    Get the Draw agent configuration.
+    
+    Returns:
+        DrawDependencies: The draw dependencies
+    """
+    return DrawDependencies()

aurelian/agents/draw/draw_gradio.py

@@ -0,0 +1,50 @@
+"""
+Gradio UI for the draw agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.agents.draw.draw_agent import draw_agent
+from aurelian.agents.draw.draw_config import DrawDependencies
+from aurelian.utils.async_utils import run_sync
+
+
+def chat(deps: Optional[DrawDependencies] = None, workdir: str = None, **kwargs):
+    """
+    Initialize a chat interface for the draw agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        workdir: Optional working directory path
+        **kwargs: Additional arguments to pass to the agent
+        
+    Returns:
+        A Gradio chat interface
+    """
+    if deps is None:
+        deps = DrawDependencies()
+        
+    if workdir:
+        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: draw_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="Drawing AI Assistant",
+        examples=[
+            ["Draw a simple cat face"],
+            ["Create an SVG of a tree with birds"],
+            ["Draw a basic house with a chimney"]
+        ]
+    )

aurelian/agents/draw/draw_mcp.py

@@ -0,0 +1,94 @@
+"""
+MCP tools for creating SVG drawings.
+"""
+import os
+from typing import Dict, List
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.draw.draw_tools as dt
+from aurelian.agents.draw.draw_agent import SYSTEM
+from aurelian.agents.draw.draw_config import DrawDependencies
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("draw", instructions=SYSTEM)
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> DrawDependencies:
+    deps = DrawDependencies()
+    # Set the location from environment variable or default
+    loc = os.getenv("AURELIAN_WORKDIR", "/tmp/aurelian")
+    deps.workdir = WorkDir(loc)
+    return deps
+
+def ctx() -> RunContext[DrawDependencies]:
+    rc: RunContext[DrawDependencies] = RunContext[DrawDependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def create_svg_drawing(description: str) -> str:
+    """
+    Create an SVG drawing based on a text description.
+
+    Args:
+        description: Detailed description of what to draw
+        
+    Returns:
+        SVG markup of the drawing
+    """
+    return await dt.create_svg_drawing(ctx(), description)
+
+
+@mcp.tool()
+async def convert_svg_to_png(svg_content: str) -> bytes:
+    """
+    Convert SVG content to PNG image.
+
+    Args:
+        svg_content: SVG markup as a string
+        
+    Returns:
+        PNG image data as bytes
+    """
+    return await dt.convert_svg_to_png(ctx(), svg_content)
+
+
+@mcp.tool()
+async def svg_to_data_url(svg_content: str) -> str:
+    """
+    Convert SVG content to a data URL for embedding in HTML.
+
+    Args:
+        svg_content: SVG markup as a string
+        
+    Returns:
+        Data URL representation of the SVG
+    """
+    return await dt.svg_to_data_url(ctx(), svg_content)
+
+
+@mcp.tool()
+async def judge_drawing(svg_content: str, description: str) -> str:
+    """
+    Judge the quality of an SVG drawing based on the description.
+    
+    Args:
+        svg_content: SVG markup as a string
+        description: The original description of what to draw
+        
+    Returns:
+        Feedback on the drawing's clarity and simplicity
+    """
+    return await dt.judge_drawing(ctx(), svg_content, description)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')

aurelian/agents/draw/draw_tools.py

@@ -0,0 +1,100 @@
+"""
+Tools for the draw agent.
+"""
+import base64
+from typing import Optional
+import cairosvg
+from pydantic import BaseModel
+
+from pydantic_ai import RunContext, BinaryContent, ModelRetry
+
+from aurelian.agents.draw.draw_config import DrawDependencies
+
+
+class DrawingFeedback(BaseModel):
+    """
+    Feedback on the drawing's clarity and simplicity.
+    """
+    feedback: str
+    necessary_changes: Optional[str] = None
+    optional_changes: Optional[str] = None
+
+
+
+
+async def convert_svg_to_png(ctx: RunContext[DrawDependencies], svg_content: str) -> bytes:
+    """
+    Convert SVG content to PNG image.
+
+    Args:
+        ctx: The run context
+        svg_content: SVG markup as a string
+        
+    Returns:
+        bytes: PNG image data
+    """
+    print("Converting SVG to PNG")
+    
+    # Check size limits
+    if len(svg_content.encode('utf-8')) > ctx.deps.max_svg_size:
+        raise ValueError(f"SVG content exceeds maximum size of {ctx.deps.max_svg_size} bytes")
+    
+    # Convert SVG to PNG using cairosvg
+    png_bytes = cairosvg.svg2png(bytestring=svg_content.encode('utf-8'))
+    return png_bytes
+
+
+async def svg_to_data_url(ctx: RunContext[DrawDependencies], svg_content: str) -> str:
+    """
+    Convert SVG content to a data URL for embedding in HTML.
+
+    Args:
+        ctx: The run context
+        svg_content: SVG markup as a string
+        
+    Returns:
+        str: Data URL representation of the SVG
+    """
+    print("Converting SVG to data URL")
+    
+    # Check size limits
+    if len(svg_content.encode('utf-8')) > ctx.deps.max_svg_size:
+        raise ValueError(f"SVG content exceeds maximum size of {ctx.deps.max_svg_size} bytes")
+    
+    # Encode as base64 and create data URL
+    b64_svg = base64.b64encode(svg_content.encode('utf-8')).decode('ascii')
+    return f"data:image/svg+xml;base64,{b64_svg}"
+
+
+async def judge_drawing(ctx: RunContext[DrawDependencies], 
+                        svg_content: str, 
+                        description: str,
+                        attempt_number: int = 1) -> DrawingFeedback:
+    """
+    Judge the readability of an SVG drawing based on the description.
+
+    In particular, make sure that text is readable and contained within boxes
+    
+    Args:
+        ctx: The run context
+        svg_content: SVG markup as a string
+        description: Simple natural language narrative summary of the drawing
+        attempt_number: The attempt number for the drawing
+        
+    Returns:
+        DrawingFeedback: Feedback on the drawing's readability
+    """
+    print(f"Judging drawing for: {description}")
+    
+    from aurelian.agents.draw.judge_agent import drawing_judge_agent
+    
+    # Convert to PNG for the judge to see
+    png_bytes = await convert_svg_to_png(ctx, svg_content)
+    img = BinaryContent(data=png_bytes, media_type='image/png')
+    
+    # Get feedback from judge
+    feedback = await drawing_judge_agent.run(
+        [f"Please evaluate this drawing based on this description: {description} (this is attempt #{attempt_number})", img],
+        deps=ctx.deps)
+    
+    return feedback.data

aurelian/agents/draw/judge_agent.py

@@ -0,0 +1,18 @@
+"""
+Agent specifically for judging the quality of drawings.
+"""
+
+from pydantic_ai import Agent
+
+from aurelian.agents.draw.draw_tools import DrawingFeedback
+
+# Separate agent for judging drawings
+drawing_judge_agent = Agent(
+    model='openai:gpt-4o',
+    system_prompt="""You role is to judge the simplicity and clarity of drawings and figures.
+    
+    Specifically, in addition to correctness, you should focus on mistakes made, overlapping
+    or unclear text, etc
+    """,
+    result_type=DrawingFeedback,
+)

aurelian/agents/filesystem/filesystem_config.py

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

aurelian/agents/filesystem/filesystem_gradio.py

@@ -0,0 +1,49 @@
+"""
+Gradio UI for the Filesystem agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.dependencies.workdir import HasWorkdir
+from aurelian.agents.filesystem.filesystem_tools import inspect_file, download_url_as_markdown
+from aurelian.agents.filesystem.filesystem_mcp import mcp
+from aurelian.utils.async_utils import run_sync
+
+
+def chat(deps: Optional[HasWorkdir] = None, **kwargs):
+    """
+    Initialize a chat interface for the Filesystem 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.filesystem.filesystem_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="Filesystem AI Assistant",
+        examples=[
+            ["List all files in the working directory"],
+            ["Download https://example.com and save as example.md"],
+            ["Show me the contents of the file example.md"]
+        ]
+    )