aurelian 0.3.2__py3-none-any.whl

aurelian/agents/linkml/linkml_tools.py

@@ -0,0 +1,102 @@
+from typing import Optional
+
+import yaml
+from linkml.generators import JsonSchemaGenerator
+from linkml_runtime.linkml_model import SchemaDefinition
+from linkml_runtime.loaders import yaml_loader
+from pydantic import BaseModel
+from pydantic_ai import RunContext, ModelRetry
+
+from aurelian.agents.linkml.linkml_config import LinkMLDependencies
+from aurelian.dependencies.workdir import WorkDir, HasWorkdir
+
+
+class LinkMLError(ModelRetry):
+    pass
+
+class SchemaValidationError(LinkMLError):
+    """Base exception for all schema validation errors."""
+    def __init__(self, message="Schema validation failed", details=None):
+        self.details = details or {}
+        super().__init__(message)
+
+
+class ValidationResult(BaseModel):
+    valid: bool
+    info_messages: Optional[list[str]] = None
+
+
+async def validate_then_save_schema(ctx: RunContext[HasWorkdir], schema_as_str: str, save_to_file: str= "schema.yaml") -> ValidationResult:
+    """
+    Validate a LinkML schema.
+
+    Args:
+        ctx: context
+        schema_as_str: linkml schema (as yaml) to validate. Do not truncate, always pass the whole schema.
+        save_to_file: file name to save the schema to. Defaults to schema.yaml
+
+    Returns:
+
+    """
+    print(f"Validating schema: {schema_as_str}")
+    msgs = []
+    try:
+        schema_dict = yaml.safe_load(schema_as_str)
+        print("YAML is valid")
+    except Exception as e:
+        raise SchemaValidationError(f"Schema is not valid yaml: {e}")
+    if "id" not in schema_dict:
+        raise SchemaValidationError("Schema does not have a top level id")
+    if "name" not in schema_dict:
+        raise SchemaValidationError("Schema does not have a top level name")
+    try:
+        schema_obj = yaml_loader.loads(schema_as_str, target_class=SchemaDefinition)
+    except Exception as e:
+        raise ModelRetry(f"Schema does not validate: {e} // {schema_as_str}")
+    try:
+        gen = JsonSchemaGenerator(schema_obj)
+        gen.serialize()
+    except Exception as e:
+        raise ModelRetry(f"Schema does not convert to JSON-Schema: {e} // {schema_as_str}")
+    try:
+        if save_to_file and schema_as_str:
+            msgs.append(f"Writing schema to {save_to_file}")
+            workdir = ctx.deps.workdir
+            workdir.write_file(save_to_file, schema_as_str)
+    except Exception as e:
+        raise ModelRetry(f"Schema does not validate: {e} // {schema_as_str}")
+    return ValidationResult(valid=True, info_messages=msgs)
+
+
+async def validate_data(ctx: RunContext[LinkMLDependencies], schema: str, data_file: str) -> str:
+    """
+    Validate data file against a schema.
+
+    This assumes the data file is present in the working directory.
+    You can write data to the working directory using the `write_to_file` tool.
+
+    Args:
+        ctx:
+        schema: the schema (as a YAML string)
+        data_file: the name of the data file in the working directory
+
+    Returns:
+
+    """
+    print(f"Validating data file: {data_file} using schema: {schema}")
+    try:
+        schema = yaml_loader.loads(schema, target_class=SchemaDefinition)
+    except Exception as e:
+        return f"Schema does not validate: {e}"
+    try:
+        from linkml.validator import validate
+        instances = ctx.deps.parse_objects_from_file(data_file)
+        for instance in instances:
+            print(f"Validating {instance}")
+            rpt = validate(instance, schema)
+            print(f"Validation report: {rpt}")
+            if rpt.results:
+                return f"Data does not validate:\n{rpt.results}"
+        return f"{len(instances)} instances all validate successfully"
+    except Exception as e:
+        raise ModelRetry(f"Data does not validate: {e}")

aurelian/agents/literature/__init__.py

@@ -0,0 +1,3 @@
+"""
+Literature agent module for working with scientific literature, publications, and references.
+"""

aurelian/agents/literature/literature_agent.py

@@ -0,0 +1,55 @@
+"""
+Agent for working with scientific literature and publications.
+"""
+from aurelian.agents.literature.literature_config import LiteratureDependencies
+from aurelian.agents.literature.literature_tools import (
+    lookup_pmid,
+    lookup_doi,
+    convert_pmid_to_doi,
+    convert_doi_to_pmid,
+    get_article_abstract,
+    extract_text_from_pdf_url,
+    search_literature_web,
+    retrieve_literature_page
+)
+from aurelian.agents.filesystem.filesystem_tools import inspect_file, list_files
+from pydantic_ai import Agent, Tool
+
+SYSTEM = """
+You are an expert scientific literature assistant that helps users access and analyze scientific publications.
+
+You can help with:
+- Finding and retrieving full text of articles using PubMed IDs or DOIs
+- Converting between PubMed IDs and DOIs
+- Extracting text from PDF articles
+- Searching for scientific literature on specific topics
+- Analyzing and summarizing scientific papers
+
+Always provide accurate citations for any scientific information, including:
+- Article titles
+- Authors
+- Journal names
+- Publication dates
+- DOIs and/or PubMed IDs
+
+When quoting or referencing a specific part of a paper, always indicate which section it comes from
+(e.g., abstract, methods, results, discussion).
+"""
+
+literature_agent = Agent(
+    model="openai:gpt-4o",
+    deps_type=LiteratureDependencies,
+    system_prompt=SYSTEM,
+    tools=[
+        Tool(lookup_pmid),
+        Tool(lookup_doi),
+        Tool(convert_pmid_to_doi),
+        Tool(convert_doi_to_pmid),
+        Tool(get_article_abstract),
+        Tool(extract_text_from_pdf_url),
+        Tool(search_literature_web),
+        Tool(retrieve_literature_page),
+        Tool(inspect_file),
+        Tool(list_files),
+    ]
+)

aurelian/agents/literature/literature_config.py

@@ -0,0 +1,35 @@
+"""
+Configuration classes for the literature agent.
+"""
+from dataclasses import dataclass
+import os
+from typing import Optional
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+
+
+@dataclass
+class LiteratureDependencies(HasWorkdir):
+    """
+    Configuration for the literature agent.
+    """
+    max_results: int = 10
+    
+    def __post_init__(self):
+        """Initialize the config with default values."""
+        # Initialize workdir if not provided
+        if self.workdir is None:
+            self.workdir = WorkDir()
+
+
+def get_config() -> LiteratureDependencies:
+    """
+    Get the Literature agent configuration from environment variables or defaults.
+    
+    Returns:
+        LiteratureDependencies: The literature dependencies
+    """
+    workdir_path = os.environ.get("AURELIAN_WORKDIR", None)
+    workdir = WorkDir(location=workdir_path) if workdir_path else None
+    
+    return LiteratureDependencies(workdir=workdir)

aurelian/agents/literature/literature_gradio.py

@@ -0,0 +1,52 @@
+"""
+Gradio UI for the literature agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from aurelian.agents.literature.literature_agent import literature_agent
+from aurelian.agents.literature.literature_config import LiteratureDependencies
+from aurelian.utils.async_utils import run_sync
+
+
+def chat(deps: Optional[LiteratureDependencies] = None, workdir: str = None, **kwargs):
+    """
+    Initialize a chat interface for the literature 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 = LiteratureDependencies()
+        
+    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: literature_agent.run_sync(query, deps=deps, **kwargs))
+        return result.data
+
+    return gr.ChatInterface(
+        fn=get_info,
+        type="messages",
+        title="Scientific Literature Assistant",
+        examples=[
+            ["Look up this article: PMID:31653696"],
+            ["Find information about Alzheimer's disease genetics in recent papers"],
+            ["What is the DOI for PMID:27629041?"],
+            ["Get the abstract of PMID:30478089"],
+            ["Convert this DOI to a PMID: 10.1038/nature12373"]
+        ]
+    )

aurelian/agents/literature/literature_mcp.py

@@ -0,0 +1,174 @@
+"""
+MCP tools for working with scientific literature and publications.
+"""
+import os
+from typing import Optional
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.literature.literature_tools as lt
+import aurelian.agents.filesystem.filesystem_tools as fst
+from aurelian.agents.literature.literature_agent import SYSTEM
+from aurelian.agents.literature.literature_config import LiteratureDependencies
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("literature", instructions=SYSTEM)
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> LiteratureDependencies:
+    deps = LiteratureDependencies()
+    # Set the location from environment variable or default
+    loc = os.getenv("AURELIAN_WORKDIR", "/tmp/aurelian")
+    deps.workdir = WorkDir(loc)
+    return deps
+
+def ctx() -> RunContext[LiteratureDependencies]:
+    rc: RunContext[LiteratureDependencies] = RunContext[LiteratureDependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def lookup_pmid(pmid: str) -> str:
+    """
+    Lookup the text of a PubMed article by its PMID.
+    
+    Args:
+        pmid: The PubMed ID to look up (format: "PMID:nnnnnnn")
+        
+    Returns:
+        Full text if available, otherwise abstract
+    """
+    return await lt.lookup_pmid(ctx(), pmid)
+
+@mcp.tool()
+async def lookup_doi(doi: str) -> str:
+    """
+    Lookup the text of an article by its DOI.
+    
+    Args:
+        doi: The DOI to look up
+        
+    Returns:
+        Full text if available, otherwise abstract
+    """
+    return await lt.lookup_doi(doi)
+
+
+@mcp.tool()
+async def convert_pmid_to_doi(pmid: str) -> str:
+    """
+    Convert a PubMed ID (PMID) to a DOI.
+    
+    Args:
+        pmid: The PubMed ID to convert
+        
+    Returns:
+        The corresponding DOI
+    """
+    return await lt.convert_pmid_to_doi(pmid)
+
+
+@mcp.tool()
+async def convert_doi_to_pmid(doi: str) -> str:
+    """
+    Convert a DOI to a PubMed ID (PMID).
+    
+    Args:
+        doi: The DOI to convert
+        
+    Returns:
+        The corresponding PubMed ID
+    """
+    return await lt.convert_doi_to_pmid(doi)
+
+
+@mcp.tool()
+async def get_article_abstract(identifier: str) -> str:
+    """
+    Get the abstract of an article by its PMID or DOI.
+    
+    Args:
+        identifier: PMID or DOI of the article
+        
+    Returns:
+        The article abstract
+    """
+    return await lt.get_article_abstract(identifier)
+
+
+@mcp.tool()
+async def extract_text_from_pdf_url(url: str) -> str:
+    """
+    Extract text from a PDF at the given URL.
+    
+    Args:
+        url: URL to the PDF file
+        
+    Returns:
+        Extracted text from the PDF
+    """
+    return await lt.extract_text_from_pdf_url(url)
+
+
+@mcp.tool()
+async def search_literature_web(query: str) -> str:
+    """
+    Search the web for scientific literature.
+    
+    Args:
+        query: The search query
+        
+    Returns:
+        Search results with summaries
+    """
+    return await lt.search_literature_web(query)
+
+
+@mcp.tool()
+async def retrieve_literature_page(url: str) -> str:
+    """
+    Fetch the contents of a web page related to scientific literature.
+    
+    Args:
+        url: The URL to fetch
+        
+    Returns:
+        The contents of the web page
+    """
+    return await lt.retrieve_literature_page(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/literature/literature_tools.py

@@ -0,0 +1,182 @@
+"""
+Tools for the literature agent.
+"""
+from typing import Optional, List, Dict
+
+from pydantic_ai import RunContext, ModelRetry
+
+from aurelian.agents.literature.literature_config import LiteratureDependencies
+from aurelian.utils.doi_fetcher import DOIFetcher
+from aurelian.utils.pubmed_utils import (
+    get_pmid_text,
+    get_doi_text,
+    pmid_to_doi,
+    doi_to_pmid,
+    get_abstract_from_pubmed,
+)
+from aurelian.utils.pdf_fetcher import extract_text_from_pdf
+from aurelian.utils.search_utils import web_search, retrieve_web_page
+
+
+async def lookup_pmid(pmid: str) -> str:
+    """
+    Lookup the text of a PubMed article by its PMID.
+    
+    A PMID should be of the form "PMID:nnnnnnn" (no underscores).
+    
+    Args:
+        pmid: The PubMed ID to look up
+        
+    Returns:
+        str: Full text if available, otherwise abstract
+    """
+    print(f"LOOKUP PMID: {pmid}")
+    try:
+        result = get_pmid_text(pmid)
+        print(f"RESULT LENGTH: {len(result)} // {result[:100]}")
+        if not result or "Error" in result:
+            raise ModelRetry(f"Could not retrieve text for PMID: {pmid}. Try using the abstract only or a different identifier.")
+        return result
+    except Exception as e:
+        raise ModelRetry(f"Error retrieving PMID {pmid}: {str(e)}. Try using the abstract only or a different identifier.")
+
+
+async def lookup_doi(doi: str) -> str:
+    """
+    Lookup the text of an article by its DOI.
+    
+    Args:
+        doi: The DOI to look up (e.g., "10.1038/nature12373")
+        
+    Returns:
+        str: Full text if available, otherwise abstract
+    """
+    print(f"LOOKUP DOI: {doi}")
+    try:
+        result = get_doi_text(doi)
+        if not result or "Error" in result or "not available" in result.lower():
+            raise ModelRetry(f"Could not retrieve text for DOI: {doi}. Try using a PubMed ID or a different approach.")
+        return result
+    except Exception as e:
+        raise ModelRetry(f"Error retrieving DOI {doi}: {str(e)}. Try using a PubMed ID or a different approach.")
+
+
+async def convert_pmid_to_doi(pmid: str) -> Optional[str]:
+    """
+    Convert a PubMed ID to a DOI.
+    
+    Args:
+        pmid: The PubMed ID to convert
+        
+    Returns:
+        str: The corresponding DOI, or None if not found
+    """
+    print(f"CONVERT PMID TO DOI: {pmid}")
+    try:
+        result = pmid_to_doi(pmid)
+        if not result:
+            raise ModelRetry(f"Could not convert PMID {pmid} to DOI. This article may not have a DOI assigned.")
+        return result
+    except Exception as e:
+        raise ModelRetry(f"Error converting PMID {pmid} to DOI: {str(e)}")
+
+
+async def convert_doi_to_pmid(doi: str) -> Optional[str]:
+    """
+    Convert a DOI to a PubMed ID.
+    
+    Args:
+        doi: The DOI to convert
+        
+    Returns:
+        str: The corresponding PubMed ID, or None if not found
+    """
+    print(f"CONVERT DOI TO PMID: {doi}")
+    try:
+        result = doi_to_pmid(doi)
+        if not result:
+            raise ModelRetry(f"Could not convert DOI {doi} to PMID. This article may not be indexed in PubMed.")
+        return result
+    except Exception as e:
+        raise ModelRetry(f"Error converting DOI {doi} to PMID: {str(e)}")
+
+
+async def get_article_abstract(pmid: str) -> str:
+    """
+    Get only the abstract of an article by its PubMed ID.
+    
+    Args:
+        pmid: The PubMed ID to look up
+        
+    Returns:
+        str: The article abstract
+    """
+    print(f"GET ABSTRACT: {pmid}")
+    try:
+        result = get_abstract_from_pubmed(pmid)
+        if not result or result.endswith("No abstract available"):
+            raise ModelRetry(f"No abstract available for PMID {pmid}. This article may not have an abstract or the PMID may be incorrect.")
+        return result
+    except Exception as e:
+        raise ModelRetry(f"Error retrieving abstract for PMID {pmid}: {str(e)}")
+
+
+async def extract_text_from_pdf_url(ctx: RunContext[LiteratureDependencies], pdf_url: str) -> str:
+    """
+    Extract text from a PDF at the given URL.
+    
+    Args:
+        ctx: The run context
+        pdf_url: URL to the PDF file
+        
+    Returns:
+        str: The extracted text content
+    """
+    print(f"EXTRACT PDF: {pdf_url}")
+    try:
+        result = extract_text_from_pdf(pdf_url)
+        if not result or "Error" in result:
+            raise ModelRetry(f"Could not extract text from PDF at {pdf_url}. The URL may be invalid or the PDF may be password-protected.")
+        return result
+    except Exception as e:
+        raise ModelRetry(f"Error extracting text from PDF {pdf_url}: {str(e)}")
+
+
+async def search_literature_web(query: str) -> str:
+    """
+    Search the web for scientific literature using a text query.
+    
+    Args:
+        query: The search query (e.g., "alzheimer's disease genetics 2023")
+        
+    Returns:
+        str: Search results with summaries
+    """
+    print(f"LITERATURE WEB SEARCH: {query}")
+    try:
+        result = web_search(query)
+        if not result:
+            raise ModelRetry(f"No search results found for query: {query}. Try using different keywords.")
+        return result
+    except Exception as e:
+        raise ModelRetry(f"Error searching the web for '{query}': {str(e)}")
+
+
+async def retrieve_literature_page(url: str) -> str:
+    """
+    Fetch the contents of a literature webpage.
+    
+    Args:
+        url: The URL to fetch
+        
+    Returns:
+        str: The contents of the webpage
+    """
+    print(f"FETCH LITERATURE URL: {url}")
+    try:
+        result = retrieve_web_page(url)
+        if not result or len(result.strip()) < 20:
+            raise ModelRetry(f"Could not retrieve meaningful content from {url}. The URL may be invalid or require authentication.")
+        return result
+    except Exception as e:
+        raise ModelRetry(f"Error retrieving webpage {url}: {str(e)}")

aurelian/agents/monarch/__init__.py

@@ -0,0 +1,25 @@
+"""
+Monarch agent package for interacting with the Monarch Knowledge Base.
+"""
+
+from .monarch_agent import monarch_agent, MONARCH_SYSTEM_PROMPT
+from .monarch_config import MonarchDependencies, get_config
+from .monarch_gradio import chat
+from .monarch_tools import find_gene_associations, find_disease_associations
+
+__all__ = [
+    # Agent
+    "monarch_agent",
+    "MONARCH_SYSTEM_PROMPT",
+    
+    # Config
+    "MonarchDependencies",
+    "get_config",
+    
+    # Tools
+    "find_gene_associations",
+    "find_disease_associations",
+    
+    # Gradio
+    "chat",
+]

aurelian/agents/monarch/monarch_agent.py

@@ -0,0 +1,44 @@
+"""
+Agent for interacting with the Monarch knowledge base.
+"""
+from pydantic_ai import Agent
+
+from .monarch_config import MonarchDependencies, get_config
+from .monarch_tools import find_gene_associations, find_disease_associations
+
+# System prompt for the Monarch agent
+MONARCH_SYSTEM_PROMPT = """
+You are a helpful assistant specializing in biomedical data from the Monarch Knowledge Base.
+You can help researchers find relationships between genes, diseases, phenotypes, and other biomedical entities.
+
+The Monarch Knowledge Base integrates data from multiple biomedical databases and provides a unified interface
+for querying associations between different biological entities.
+
+You can:
+- Find associations for genes, including what diseases they're linked to
+- Find associations for diseases, including what genes and phenotypes they're linked to
+- Provide information about biological relationships in a structured way
+
+When working with identifiers:
+- Gene symbols should be specified as HGNC or MGI symbols (e.g. "BRCA1")
+- Disease IDs can be specified as MONDO, OMIM, or Orphanet IDs (e.g. "MONDO:0007254")
+- Phenotype IDs can be specified as HP terms (e.g. "HP:0000118")
+
+Present your findings in a clear, organized manner that helps researchers understand the biological significance
+of the associations. Include relevant details about:
+- Source of the associations
+- Strength of evidence
+- Type of relationship (causal, correlative, etc.)
+- Relevant literature references when available
+"""
+
+# Create the agent with the system prompt
+monarch_agent = Agent(
+    model="openai:gpt-4o",
+    system_prompt=MONARCH_SYSTEM_PROMPT,
+    deps_type=MonarchDependencies,
+)
+
+# Register the tools with the agent
+monarch_agent.tool(find_gene_associations)
+monarch_agent.tool(find_disease_associations)