aurelian 0.3.2__py3-none-any.whl

aurelian/agents/chemistry/chemistry_config.py

@@ -0,0 +1,71 @@
+"""
+Configuration classes for the chemistry agent.
+"""
+from dataclasses import dataclass, field
+from typing import Optional
+
+from pydantic import BaseModel
+from aurelian.dependencies.workdir import HasWorkdir
+
+
+class ChemicalStructure(BaseModel):
+    """
+    Model for representing chemical structures.
+    """
+    chebi_id: Optional[str] = None
+    smiles: Optional[str] = None
+    inchi: Optional[str] = None
+    name: Optional[str] = None
+
+    @property
+    def chebi_local_id(self) -> Optional[str]:
+        if self.chebi_id:
+            return self.chebi_id.split(":")[1]
+        return None
+
+    @property
+    def chebi_image_url(self) -> str:
+        local_id = self.chebi_local_id
+        if local_id:
+            return f"https://www.ebi.ac.uk/chebi/displayImage.do?defaultImage=true&imageIndex=0&chebiId={local_id}"
+        return ""
+
+    @classmethod
+    def from_id(cls, id: str) -> 'ChemicalStructure':
+        if ":" in id:
+            prefix, local_id = id.split(":")
+            if prefix.lower() != "chebi":
+                raise ValueError(f"Invalid prefix: {prefix}")
+            id = "CHEBI:" + local_id
+        else:
+            id = "CHEBI:" + id
+        return cls(chebi_id=id)
+
+    @classmethod
+    def from_anything(cls, id: str) -> 'ChemicalStructure':
+        if ":" in id:
+            return cls.from_id(id)
+        # check if valid smiles
+        from rdkit import Chem
+        mol = Chem.MolFromSmiles(id)
+        if mol:
+            return cls(smiles=id)
+        raise ValueError(f"Invalid identifier: {id}")
+
+
+@dataclass
+class ChemistryDependencies(HasWorkdir):
+    """
+    Configuration for the chemistry agent.
+    """
+    max_search_results: int = 30
+
+
+def get_config() -> ChemistryDependencies:
+    """
+    Get the Chemistry agent configuration.
+    
+    Returns:
+        ChemistryDependencies: The chemistry dependencies
+    """
+    return ChemistryDependencies()

aurelian/agents/chemistry/chemistry_evals.py

@@ -0,0 +1,79 @@
+"""
+Evaluation module for the Chemistry agent.
+
+This module implements evaluations for the Chemistry 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.chemistry.chemistry_agent import chemistry_agent
+from aurelian.agents.chemistry.chemistry_config import ChemistryDependencies
+
+class ChemistryMetadata(Dict[str, Any]):
+    """Simple metadata dictionary for Chemistry evaluations."""
+    pass
+
+# Define individual evaluation cases
+case1 = Case(
+    name="caffeine_structure",
+    inputs="Explain the structure of caffeine (CHEBI:27732)",
+    expected_output="methylxanthine",  # Should mention methylxanthine structure
+    metadata=metadata("medium", "structure_explanation")
+)
+
+case2 = Case(
+    name="aspirin_properties",
+    inputs="What does the structure of aspirin (CHEBI:15365) tell us about its properties?",
+    expected_output="acetyl",  # Should mention acetyl group
+    metadata=metadata("medium", "structure_property_relationship"),
+    evaluators=[
+        LLMJudge(
+            rubric="Answer should explain how the acetyl group affects aspirin's properties and mention its action as a COX inhibitor",
+            include_input=True
+        )
+    ]
+)
+
+case3 = Case(
+    name="smiles_interpretation",
+    inputs="Interpret this SMILES: CC(=O)OC1=CC=CC=C1C(=O)O",
+    expected_output="aspirin",  # This is aspirin
+    metadata=metadata("hard", "smiles_interpretation")
+)
+
+case4 = Case(
+    name="functional_groups",
+    inputs="Identify all functional groups in paracetamol (CHEBI:46195)",
+    expected_output="amide",  # Should identify the amide group
+    metadata=metadata("medium", "functional_group_identification")
+)
+
+def create_eval_dataset() -> Dataset[str, str, MetadataDict]:
+    """
+    Create a dataset for evaluating the Chemistry agent.
+    
+    Returns:
+        Dataset of Chemistry evaluation cases with appropriate evaluators
+    """
+    # Collect all cases
+    cases = [case1, case2, case3, case4]
+    
+    # Dataset-level evaluators
+    evaluators = [
+        SubstringEvaluator(),
+        LLMJudge(
+            rubric="Answer should be scientifically accurate and use proper chemistry terminology",
+            model="anthropic:claude-3-7-sonnet-latest"
+        )
+    ]
+    
+    return Dataset(
+        cases=cases,
+        evaluators=evaluators
+    )

aurelian/agents/chemistry/chemistry_gradio.py

@@ -0,0 +1,50 @@
+"""
+Gradio UI for the chemistry agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.agents.chemistry.chemistry_agent import chemistry_agent
+from aurelian.agents.chemistry.chemistry_config import ChemistryDependencies
+from aurelian.utils.async_utils import run_sync
+
+
+def chat(deps: Optional[ChemistryDependencies] = None, workdir: str = None, **kwargs):
+    """
+    Initialize a chat interface for the chemistry 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 = ChemistryDependencies()
+        
+    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: chemistry_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="Chemistry AI Assistant",
+        examples=[
+            ["Explain the structure of caffeine (CHEBI:27732)"],
+            ["What does the structure of aspirin (CHEBI:15365) tell us about its properties?"],
+            ["Interpret this SMILES: CC(=O)OC1=CC=CC=C1C(=O)O"]
+        ]
+    )

aurelian/agents/chemistry/chemistry_mcp.py

@@ -0,0 +1,120 @@
+"""
+MCP tools for working with chemical structures.
+"""
+import os
+from typing import Dict, List
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.chemistry.chemistry_tools as ct
+import aurelian.agents.filesystem.filesystem_tools as fst
+from aurelian.agents.chemistry.chemistry_agent import SYSTEM
+from aurelian.agents.chemistry.chemistry_config import ChemistryDependencies
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("chemistry", instructions=SYSTEM)
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> ChemistryDependencies:
+    deps = ChemistryDependencies()
+    # Set the location from environment variable or default
+    loc = os.getenv("AURELIAN_WORKDIR", "/tmp/aurelian")
+    deps.workdir = WorkDir(loc)
+    return deps
+
+def ctx() -> RunContext[ChemistryDependencies]:
+    rc: RunContext[ChemistryDependencies] = RunContext[ChemistryDependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def draw_structure_and_interpret(identifier: str, question: str) -> str:
+    """
+    Draw a chemical structure and analyze it.
+
+    Args:
+        identifier: A ChEBI ID (e.g., CHEBI:16236) or SMILES string
+        question: A specific question about the structure
+        
+    Returns:
+        Analysis of the structure in response to the question
+    """
+    return await ct.draw_structure_and_interpret(ctx(), identifier, question)
+
+
+@mcp.tool()
+async def chebi_search_terms(query: str) -> List[Dict]:
+    """
+    Search ChEBI for a term.
+
+    Args:
+        query: The search text
+        
+    Returns:
+        A list of matching ChEBI terms
+    """
+    return await ct.chebi_search_terms(ctx(), query)
+
+
+@mcp.tool()
+async def search_web_for_chemistry(query: str) -> str:
+    """
+    Search the web for chemistry information.
+
+    Args:
+        query: The search query
+        
+    Returns:
+        Search results with summaries
+    """
+    return await ct.search_web_for_chemistry(ctx(), query)
+
+
+@mcp.tool()
+async def retrieve_chemistry_web_page(url: str) -> str:
+    """
+    Fetch the contents of a web page related to chemistry.
+
+    Args:
+        url: The URL to fetch
+        
+    Returns:
+        The contents of the web page
+    """
+    return await ct.retrieve_chemistry_web_page(ctx(), url)
+
+
+@mcp.tool()
+async def inspect_file(data_file: str) -> str:
+    """
+    Inspect a file in the working directory.
+
+    Args:
+        data_file: name of file
+        
+    Returns:
+        Contents of the file
+    """
+    return await fst.inspect_file(ctx(), data_file)
+
+
+@mcp.tool()
+async def list_files() -> str:
+    """
+    List files in the working directory.
+
+    Returns:
+        List of files in the working directory
+    """
+    return await fst.list_files(ctx())
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')

aurelian/agents/chemistry/chemistry_tools.py

@@ -0,0 +1,121 @@
+"""
+Tools for the chemistry agent.
+"""
+import io
+import httpx
+from functools import lru_cache
+from typing import List, Dict, Optional
+
+from oaklib import get_adapter
+from pydantic_ai import RunContext, BinaryContent, ModelRetry
+
+from aurelian.agents.chemistry.chemistry_config import ChemistryDependencies, ChemicalStructure
+from aurelian.utils.ontology_utils import search_ontology
+from aurelian.utils.search_utils import web_search, retrieve_web_page
+
+
+@lru_cache
+def get_chebi_adapter():
+    """Get the ChEBI adapter from oaklib."""
+    return get_adapter(f"sqlite:obo:chebi")
+
+
+def smiles_to_image(smiles: str) -> bytes:
+    """
+    Convert a SMILES string to an image.
+    
+    Args:
+        smiles: The SMILES representation of a molecule
+        
+    Returns:
+        bytes: PNG image of the molecular structure
+        
+    Raises:
+        ValueError: If the SMILES string is invalid
+    """
+    from rdkit import Chem
+    from rdkit.Chem import Draw
+    mol = Chem.MolFromSmiles(smiles)
+    if not mol:
+        raise ValueError(f"Invalid SMILES: {smiles}")
+    img = Draw.MolToImage(mol)
+    img_bytes = io.BytesIO()
+    img.save(img_bytes, format='PNG')
+    return img_bytes.getvalue()
+
+
+async def draw_structure_and_interpret(ctx: RunContext[ChemistryDependencies], identifier: str, question: str) -> str:
+    """
+    Draw a chemical structure and analyze it.
+
+    Args:
+        ctx: The run context
+        identifier: CHEBI ID (e.g. CHEBI:12345) or a SMILES string
+        question: Question about the structure to be answered
+        
+    Returns:
+        str: Analysis of the chemical structure
+    """
+    print(f"Draw Structure: {identifier}, then: {question}")
+    structure = ChemicalStructure.from_anything(identifier)
+    image_url = structure.chebi_image_url
+    img = None
+    
+    if image_url:
+        image_response = httpx.get(image_url)
+        img = BinaryContent(data=image_response.content, media_type='image/png')
+    else:
+        if structure.smiles:
+            img = BinaryContent(data=smiles_to_image(structure.smiles), media_type='image/png')
+    
+    if not img:
+        raise ModelRetry("Could not find image for structure")
+        
+    from aurelian.agents.chemistry.image_agent import structure_image_agent
+    result = await structure_image_agent.run(
+        [question, img],
+        deps=ctx.deps)
+    return result.data
+
+
+async def chebi_search_terms(ctx: RunContext[ChemistryDependencies], query: str) -> List[Dict]:
+    """
+    Finds similar ontology terms to the search query in ChEBI.
+
+    Args:
+        ctx: The run context 
+        query: The search query
+        
+    Returns:
+        List[Dict]: List of matching ChEBI terms
+    """
+    print(f"ChEBI Term Search: {query}")
+    return search_ontology(get_chebi_adapter(), query, limit=ctx.deps.max_search_results)
+
+
+async def search_web_for_chemistry(query: str) -> str:
+    """
+    Search the web using a text query.
+
+    Args:
+        query: The search query
+        
+    Returns:
+        str: Matching web pages plus summaries
+    """
+    print(f"Web Search: {query}")
+    return web_search(query)
+
+
+async def retrieve_chemistry_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
+    """
+    print(f"Fetch URL: {url}")
+    return retrieve_web_page(url)

aurelian/agents/chemistry/image_agent.py

@@ -0,0 +1,15 @@
+"""
+Agent specifically for interpreting chemical structure images.
+"""
+from pydantic_ai import Agent
+
+# Separate agent for image interpretation to avoid circular imports
+structure_image_agent = Agent(
+    model='openai:gpt-4o',
+    system_prompt="""You are an expert chemist, able to interpret
+    chemical structure diagrams and answer questions on them.
+    Use the information in the provided chemical structure image to
+    answer questions about molecular properties, functional groups,
+    potential reactivity, or other chemical characteristics.
+    """
+)

aurelian/agents/d4d/__init__.py

@@ -0,0 +1,30 @@
+"""
+D4D (Datasheets for Datasets) agent package for extracting dataset metadata.
+"""
+
+# isort: skip_file
+from .d4d_agent import data_sheets_agent  # noqa: E402
+from .d4d_config import D4DConfig, get_config  # noqa: E402
+from .d4d_gradio import chat  # noqa: E402
+from .d4d_tools import (  # noqa: E402
+    get_full_schema,
+    process_website_or_pdf,
+    extract_text_from_pdf,
+)
+
+__all__ = [
+    # Agent
+    "data_sheets_agent",
+    
+    # Config
+    "D4DConfig",
+    "get_config",
+    
+    # Tools
+    "get_full_schema",
+    "process_website_or_pdf",
+    "extract_text_from_pdf",
+    
+    # Gradio
+    "chat",
+]

aurelian/agents/d4d/d4d_agent.py

@@ -0,0 +1,72 @@
+"""
+Agent for extracting dataset metadata following the datasheets for datasets schema.
+"""
+from pydantic_ai import Agent, RunContext
+
+from .d4d_config import D4DConfig
+from .d4d_tools import get_full_schema, process_website_or_pdf
+
+
+# Create the agent, the full schema will be loaded when needed
+data_sheets_agent = Agent(
+    model="openai:gpt-4o",
+    deps_type=D4DConfig,
+    system_prompt="""
+Below is the complete datasheets for datasets schema:
+
+{schema}
+
+When provided with a URL to a webpage or PDF describing a dataset, your task is to fetch the 
+content, extract all the relevant metadata, and output a YAML document that exactly 
+conforms to the above schema. The output must be valid YAML with all required fields 
+filled in, following the schema exactly.
+""",
+)
+
+
+@data_sheets_agent.system_prompt
+async def add_schema(ctx: RunContext[D4DConfig]) -> str:
+    """
+    Add the full schema to the system prompt.
+    
+    Args:
+        ctx: The run context
+        
+    Returns:
+        The schema to be inserted into the system prompt
+    """
+    schema = await get_full_schema(ctx)
+    return schema
+
+
+@data_sheets_agent.tool
+async def extract_metadata(ctx: RunContext[D4DConfig], url: str) -> str:
+    """
+    Extract metadata from a dataset description document or webpage.
+    
+    Args:
+        ctx: The run context
+        url: The URL of the dataset description (webpage or PDF)
+        
+    Returns:
+        YAML formatted metadata following the datasheets for datasets schema
+    """
+    # Retrieve the content
+    content = await process_website_or_pdf(ctx, url)
+    
+    # Prepare a prompt to extract metadata
+    prompt = f"""
+The following is the content of a document describing a dataset:
+
+{content}
+
+Using the complete datasheets for datasets schema provided above, extract all the metadata 
+from the document and generate a YAML document that exactly conforms to that schema. 
+Ensure that all required fields are present and the output is valid YAML. 
+The dataset URL is: {url}
+
+Generate only the YAML document.
+"""
+    
+    # The prompt will be used as the user message
+    return prompt

aurelian/agents/d4d/d4d_config.py

@@ -0,0 +1,46 @@
+"""
+Configuration for the D4D (Datasheets for Datasets) agent.
+"""
+from dataclasses import dataclass
+import os
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+
+
+@dataclass
+class D4DConfig(HasWorkdir):
+    """Configuration for the D4D agent."""
+    
+    schema_url: str = "https://raw.githubusercontent.com/monarch-initiative/ontogpt/main/src/ontogpt/templates/data_sheets_schema.yaml"
+    
+    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(schema_url: str = None) -> D4DConfig:
+    """
+    Get the D4D configuration from environment variables or defaults.
+    
+    Args:
+        schema_url: The URL to the schema YAML (overrides environment variable)
+        
+    Returns:
+        A D4DConfig instance
+    """
+    # Try to get from environment, then use provided values or defaults
+    env_schema_url = os.environ.get("AURELIAN_D4D_SCHEMA_URL", None)
+    
+    # Use provided values first, then environment, then defaults
+    final_schema_url = schema_url or env_schema_url
+    
+    workdir_path = os.environ.get("AURELIAN_WORKDIR", None)
+    workdir = WorkDir(location=workdir_path) if workdir_path else None
+    
+    config = D4DConfig(workdir=workdir)
+    if final_schema_url:
+        config.schema_url = final_schema_url
+        
+    return config

aurelian/agents/d4d/d4d_gradio.py

@@ -0,0 +1,58 @@
+"""
+Gradio interface for the D4D (Datasheets for Datasets) agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from .d4d_agent import data_sheets_agent
+from .d4d_config import D4DConfig, get_config
+
+
+async def process_url(url: str, history: List[str], config: D4DConfig) -> str:
+    """
+    Process a URL and generate metadata in YAML format.
+    
+    Args:
+        url: The URL to process (webpage or PDF)
+        history: Conversation history
+        config: The agent configuration
+        
+    Returns:
+        YAML formatted metadata
+    """
+    # Run the agent with the URL
+    result = await data_sheets_agent.run(url, deps=config)
+    return result.data
+
+
+def chat(deps: Optional[D4DConfig] = None, **kwargs):
+    """
+    Create a Gradio chat interface for the D4D agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        kwargs: Additional keyword arguments for configuration
+        
+    Returns:
+        A Gradio ChatInterface
+    """
+    # Initialize dependencies if needed
+    if deps is None:
+        deps = get_config(**kwargs)
+    
+    def get_info(url: str, history: List[str]) -> str:
+        """Wrapper for the async process_url function."""
+        import asyncio
+        return asyncio.run(process_url(url, history, deps))
+    
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="Datasheets for Datasets Agent",
+        description="Enter a URL to a webpage or PDF describing a dataset. The agent will generate metadata in YAML format according to the complete datasheets for datasets schema.",
+        examples=[
+            "https://fairhub.io/datasets/2",
+            "https://data.chhs.ca.gov/dataset/99bc1fea-c55c-4377-bad8-f00832fd195d/resource/5a6d5fe9-36e6-4aca-ba4c-bf6edc682cf5/download/hci_crime_752-narrative_examples-10-30-15-ada.pdf"
+        ]
+    )