aurelian 0.3.2__py3-none-any.whl

aurelian/agents/biblio/biblio_agent.py

@@ -0,0 +1,94 @@
+"""
+Agent for working with bibliographies and citation data.
+"""
+from pydantic_ai import Agent, RunContext
+
+from .biblio_config import BiblioDependencies
+from .biblio_tools import search_bibliography, lookup_pmid, search_web, retrieve_web_page
+
+
+biblio_agent = Agent(
+    model="openai:gpt-4o",
+    deps_type=BiblioDependencies,
+    result_type=str,
+    system_prompt=(
+        "You are an AI assistant that help organize a bibliography."
+        " You can use different functions to access the store, for example:"
+        "  - `search` to find biblio by text query"
+        "  - `lookup_phenopacket` to retrieve a specific phenopacket by ID"
+        "You can also use `lookup_pmid` to retrieve the text of a PubMed ID, or `search_web` to search the web."
+        "While you are knowledgeable about clinical genetics, you should always use the store and "
+        "functions provided to answer questions, rather than providing your own opinion or knowledge,"
+        " unless explicitly asked. For example, if you are asked to 'review' something then you "
+        "can add your own perspective and understanding. "
+        "You should endeavour to provide answers in narrative form that would be understood "
+        "by a clinical geneticists, but provide backup using assertions from the store."
+        " providing IDs of terms alongside labels is encouraged, unless asked not to."
+        "Stick to markdown, and all prefixed IDs should by hyperlinked with bioregistry,"
+        " i.e https://bioregistry.io/{curie}."
+        "tables are a good way of summarizing or comparing multiple patients, use markdown"
+        " tables for this. Use your judgment in how to roll up tables, and whether values"
+        " should be present/absent, increased/decreased, or more specific."
+    ),
+)
+
+
+@biblio_agent.tool
+async def search_bibliography_tool(ctx: RunContext[BiblioDependencies], query: str):
+    """
+    Performs a retrieval search over the biblio database.
+
+    The query can be any text, such as name of a disease, phenotype, gene, etc.
+
+    The objects returned are "biblio" which is a structured representation
+    of a patient. Each is uniquely identified by a phenopacket ID (essentially
+    the patient ID).
+
+    The objects returned are summaries of biblio; omit some details such
+    as phenotypes. Use `lookup_biblio` to retrieve full details of a phenopacket.
+
+    Note that the phenopacket store may not be complete, and the retrieval
+    method may be imperfect
+    """
+    return await search_bibliography(ctx, query)
+
+
+@biblio_agent.tool
+async def lookup_pmid_tool(ctx: RunContext[BiblioDependencies], pmid: str):
+    """
+    Lookup the text of a PubMed ID, using its PMID.
+
+    A PMID should be of the form "PMID:nnnnnnn" (no underscores).
+
+    NOTE: Phenopacket IDs are typically of the form PMID_nnn_PatientNumber,
+    but this should be be assumed. To reliably get PMIDs for a phenopacket,
+    use `lookup_phenopacket` to retrieve examine the `externalReferences`
+    field.
+
+    Returns: full text if available, otherwise abstract
+    """
+    return await lookup_pmid(ctx, pmid)
+
+
+@biblio_agent.tool
+async def search_web_tool(ctx: RunContext[BiblioDependencies], query: str):
+    """
+    Search the web using a text query.
+
+    Note, this will not retrieve the full content, for that you
+    should use `retrieve_web_page`.
+
+    Returns: matching web pages plus summaries
+    """
+    return await search_web(ctx, query)
+
+
+@biblio_agent.tool
+async def retrieve_web_page_tool(ctx: RunContext[BiblioDependencies], url: str):
+    """
+    Fetch the contents of a web page.
+
+    Returns:
+        The contents of the web page.
+    """
+    return await retrieve_web_page(ctx, url)

aurelian/agents/biblio/biblio_config.py

@@ -0,0 +1,40 @@
+"""
+Configuration for the Biblio agent.
+"""
+from dataclasses import dataclass, field
+from typing import Optional
+
+from linkml_store import Client
+from linkml_store.api import Collection
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+from . import HANDLE, DB_NAME, COLLECTION_NAME
+
+
+@dataclass
+class BiblioDependencies(HasWorkdir):
+    """Configuration for the Biblio agent."""
+    
+    max_results: int = field(default=10)
+    _collection: Optional[Collection] = None
+    
+    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()
+
+    @property
+    def collection(self) -> Collection:
+        """Get the database collection, initializing it if needed."""
+        if self._collection is None:
+            client = Client()
+            client.attach_database(HANDLE, alias=DB_NAME)
+            db = client.databases[DB_NAME]
+            self._collection = db.get_collection(COLLECTION_NAME)
+        return self._collection
+
+
+def get_config() -> BiblioDependencies:
+    """Get the Biblio configuration with default settings."""
+    return BiblioDependencies()

aurelian/agents/biblio/biblio_gradio.py

@@ -0,0 +1,67 @@
+"""
+Gradio interface for the Biblio agent.
+"""
+from typing import List, Optional
+
+import gradio as gr
+
+from .biblio_agent import biblio_agent
+from .biblio_config import BiblioDependencies, get_config
+
+
+async def get_info(query: str, history: List[str], deps: BiblioDependencies) -> str:
+    """
+    Process a query using the biblio agent.
+    
+    Args:
+        query: The user query
+        history: The conversation history
+        deps: The dependencies for the agent
+        
+    Returns:
+        The agent's response
+    """
+    print(f"QUERY: {query}")
+    print(f"HISTORY: {history}")
+    
+    # Add history to the query if available
+    if history:
+        query += "## History"
+        for h in history:
+            query += f"\n{h}"
+    
+    # Run the agent
+    result = await biblio_agent.run(query, deps=deps)
+    return result.data
+
+
+def chat(deps: Optional[BiblioDependencies] = None, **kwargs):
+    """
+    Create a Gradio chat interface for the Biblio agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        kwargs: Additional keyword arguments for the agent
+        
+    Returns:
+        A Gradio ChatInterface
+    """
+    if deps is None:
+        deps = get_config()
+        
+    def get_info_wrapper(query: str, history: List[str]) -> str:
+        # Use run_sync to handle the async function
+        from aurelian.utils.async_utils import run_sync
+        return run_sync(lambda: get_info(query, history, deps))
+    
+    return gr.ChatInterface(
+        fn=get_info_wrapper,
+        type="messages",
+        title="Biblio AI Assistant",
+        examples=[
+            ["What patients have liver disease?"],
+            ["What biblio involve genes from metabolic pathways"],
+            ["How does the type of variant affect phenotype in peroxisomal disorders?"],
+            ["Examine biblio for skeletal dysplasias, check them against publications"],
+        ],
+    )

aurelian/agents/biblio/biblio_mcp.py

@@ -0,0 +1,115 @@
+"""
+MCP tools for working with bibliographies and citation data.
+"""
+import os
+from typing import Dict, List
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.biblio.biblio_tools as bt
+from aurelian.agents.biblio.biblio_agent import biblio_agent
+from aurelian.agents.biblio.biblio_config import BiblioDependencies
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("biblio", instructions=biblio_agent.system_prompt)
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> BiblioDependencies:
+    deps = BiblioDependencies()
+    # Set the location from environment variable or default
+    loc = os.getenv("AURELIAN_WORKDIR", "/tmp/aurelian")
+    deps.workdir = WorkDir(loc)
+    return deps
+
+def ctx() -> RunContext[BiblioDependencies]:
+    rc: RunContext[BiblioDependencies] = RunContext[BiblioDependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def search_bibliography(query: str) -> List[Dict]:
+    """
+    Performs a retrieval search over the biblio database.
+
+    Args:
+        query: The search query (disease, phenotype, gene, etc.)
+        
+    Returns:
+        A list of biblio objects matching the query
+
+    The query can be any text, such as name of a disease, phenotype, gene, etc.
+
+    The objects returned are "biblio" which is a structured representation
+    of a patient. Each is uniquely identified by a phenopacket ID (essentially
+    the patient ID).
+
+    The objects returned are summaries of biblio; omit some details such
+    as phenotypes. Use `lookup_biblio` to retrieve full details of a phenopacket.
+
+    Note that the phenopacket store may not be complete, and the retrieval
+    method may be imperfect
+    """
+    return await bt.search_bibliography(ctx(), query)
+
+
+@mcp.tool()
+async def lookup_pmid(pmid: str) -> str:
+    """
+    Lookup the text of a PubMed ID, using its PMID.
+
+    Args:
+        pmid: The PubMed ID to look up (format: "PMID:nnnnnnn")
+        
+    Returns:
+        The full text if available, otherwise abstract
+
+    A PMID should be of the form "PMID:nnnnnnn" (no underscores).
+
+    NOTE: Phenopacket IDs are typically of the form PMID_nnn_PatientNumber,
+    but this should be assumed. To reliably get PMIDs for a phenopacket,
+    use `lookup_phenopacket` to retrieve examine the `externalReferences`
+    field.
+    """
+    return await bt.lookup_pmid(ctx(), pmid)
+
+
+@mcp.tool()
+async def search_web(query: str) -> str:
+    """
+    Search the web using a text query.
+
+    Args:
+        query: The search query
+        
+    Returns:
+        Matching web pages plus summaries
+
+    Note, this will not retrieve the full content, for that you
+    should use `retrieve_web_page`.
+    """
+    return await bt.search_web(ctx(), query)
+
+
+@mcp.tool()
+async def retrieve_web_page(url: str) -> str:
+    """
+    Fetch the contents of a web page.
+
+    Args:
+        url: The URL to fetch
+        
+    Returns:
+        The contents of the web page
+    """
+    return await bt.retrieve_web_page(ctx(), url)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')

aurelian/agents/biblio/biblio_tools.py

@@ -0,0 +1,164 @@
+"""
+Tools for the Biblio agent for working with bibliographic data.
+"""
+import asyncio
+from typing import Dict, List
+
+from pydantic_ai import RunContext, ModelRetry
+
+from aurelian.utils.data_utils import flatten
+from aurelian.utils.pubmed_utils import get_pmid_text
+from aurelian.utils.search_utils import web_search, retrieve_web_page as fetch_web_page
+from .biblio_config import BiblioDependencies
+
+
+async def search_bibliography(
+    ctx: RunContext[BiblioDependencies], 
+    query: str
+) -> List[Dict]:
+    """
+    Performs a retrieval search over the biblio database.
+
+    Args:
+        ctx: The run context
+        query: The search query (disease, phenotype, gene, etc.)
+
+    Returns:
+        A list of biblio objects matching the query
+
+    The query can be any text, such as name of a disease, phenotype, gene, etc.
+
+    The objects returned are "biblio" which is a structured representation
+    of a patient. Each is uniquely identified by a phenopacket ID (essentially
+    the patient ID).
+
+    The objects returned are summaries of biblio; omit some details such
+    as phenotypes. Use `lookup_biblio` to retrieve full details of a phenopacket.
+
+    Note that the phenopacket store may not be complete, and the retrieval
+    method may be imperfect
+    """
+    try:
+        print(f"SEARCH: {query}")
+        
+        # Execute the potentially blocking operation in a thread pool
+        def _search():
+            qr = ctx.deps.collection.search(query, index_name="llm", limit=ctx.deps.max_results)
+            objs = []
+            for score, row in qr.ranked_rows:
+                obj = flatten(row, preserve_keys=["interpretations", "diseases"])
+                obj["relevancy_score"] = score
+                objs.append(obj)
+                print(f"RESULT: {obj}")
+            return objs
+            
+        objs = await asyncio.to_thread(_search)
+        
+        if not objs:
+            raise ModelRetry(f"No results found for query: {query}")
+            
+        return objs
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error searching bibliography: {str(e)}")
+
+
+async def lookup_pmid(
+    ctx: RunContext[BiblioDependencies], 
+    pmid: str
+) -> str:
+    """
+    Lookup the text of a PubMed ID, using its PMID.
+
+    Args:
+        ctx: The run context
+        pmid: The PubMed ID to look up (format: "PMID:nnnnnnn")
+
+    Returns:
+        The full text if available, otherwise abstract
+
+    A PMID should be of the form "PMID:nnnnnnn" (no underscores).
+
+    NOTE: Phenopacket IDs are typically of the form PMID_nnn_PatientNumber,
+    but this should be be assumed. To reliably get PMIDs for a phenopacket,
+    use `lookup_phenopacket` to retrieve examine the `externalReferences`
+    field.
+    """
+    try:
+        print(f"LOOKUP PMID: {pmid}")
+        
+        # Execute the potentially blocking operation in a thread pool
+        text = await asyncio.to_thread(get_pmid_text, pmid)
+        
+        if not text or text.strip() == "":
+            raise ModelRetry(f"No text found for PMID: {pmid}")
+            
+        return text
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error retrieving text from PMID: {str(e)}")
+
+
+async def search_web(
+    ctx: RunContext[BiblioDependencies], 
+    query: str
+) -> str:
+    """
+    Search the web using a text query.
+
+    Args:
+        ctx: The run context
+        query: The search query
+
+    Returns:
+        Matching web pages plus summaries
+
+    Note, this will not retrieve the full content, for that you
+    should use `retrieve_web_page`.
+    """
+    try:
+        print(f"Web Search: {query}")
+        
+        # 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 web: {str(e)}")
+
+
+async def retrieve_web_page(
+    ctx: RunContext[BiblioDependencies], 
+    url: str
+) -> str:
+    """
+    Fetch the contents of a web page.
+
+    Args:
+        ctx: The run context
+        url: The URL to fetch
+
+    Returns:
+        The contents of the web page
+    """
+    try:
+        print(f"Fetch URL: {url}")
+        
+        # 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 for 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/biblio_agent.py

@@ -0,0 +1,46 @@
+"""
+Agent for working with bibliographies
+
+This module re-exports components from the biblio/ package for backward compatibility.
+"""
+from typing import Dict, List
+
+# Re-export from biblio package
+from aurelian.agents.biblio import (
+    biblio_agent,
+    BiblioDependencies,
+    get_config,
+    search_bibliography,
+    lookup_pmid,
+    search_web,
+    retrieve_web_page,
+    chat,
+)
+
+# Re-export the older synchronous versions of functions for compatibility
+@biblio_agent.tool
+def search_bibliography_sync(ctx, query: str) -> List[Dict]:
+    """Legacy synchronous version of search_bibliography"""
+    import asyncio
+    return asyncio.run(search_bibliography(ctx, query))
+
+
+@biblio_agent.tool
+def lookup_pmid_sync(ctx, pmid: str) -> str:
+    """Legacy synchronous version of lookup_pmid"""
+    import asyncio
+    return asyncio.run(lookup_pmid(ctx, pmid))
+
+
+@biblio_agent.tool
+def search_web_sync(ctx, query: str) -> str:
+    """Legacy synchronous version of search_web"""
+    import asyncio
+    return asyncio.run(search_web(ctx, query))
+
+
+@biblio_agent.tool
+def retrieve_web_page_sync(ctx, url: str) -> str:
+    """Legacy synchronous version of retrieve_web_page"""
+    import asyncio
+    return asyncio.run(retrieve_web_page(ctx, url))

aurelian/agents/checklist/__init__.py

@@ -0,0 +1,44 @@
+"""
+Checklist agent package for validating papers against checklists (e.g., STREAMS).
+"""
+from pathlib import Path
+
+THIS_DIR = Path(__file__).parent
+CONTENT_DIR = THIS_DIR / "content"
+CONTENT_METADATA_PATH = CONTENT_DIR / "checklists.yaml"
+
+# These imports must be after constants are defined
+# isort: skip_file
+from .checklist_agent import checklist_agent, add_checklists  # noqa: E402
+from .checklist_config import ChecklistDependencies, get_config  # noqa: E402
+from .checklist_gradio import chat  # noqa: E402
+from .checklist_tools import (  # noqa: E402
+    all_checklists,
+    retrieve_text_from_pmid,
+    retrieve_text_from_doi,
+    fetch_checklist,
+)
+
+__all__ = [
+    # Constants
+    "THIS_DIR",
+    "CONTENT_DIR",
+    "CONTENT_METADATA_PATH",
+    
+    # Agent
+    "checklist_agent",
+    "add_checklists",
+    
+    # Config
+    "ChecklistDependencies",
+    "get_config",
+    
+    # Tools
+    "all_checklists",
+    "retrieve_text_from_pmid",
+    "retrieve_text_from_doi",
+    "fetch_checklist",
+    
+    # Gradio
+    "chat",
+]

aurelian/agents/checklist/checklist_agent.py

@@ -0,0 +1,85 @@
+"""
+Agent for validating papers against checklists, e.g STREAMS
+"""
+from pydantic_ai import Agent, RunContext
+
+from .checklist_config import ChecklistDependencies
+from .checklist_tools import all_checklists, retrieve_text_from_pmid, retrieve_text_from_doi, fetch_checklist
+
+
+checklist_agent = Agent(
+    model="openai:gpt-4o",
+    system_prompt=(
+        "Your role is to evaluate papers to ensure they conform to relevant checklists."
+        "When asked to look at or review a paper, you should first select the "
+        "appropriate checklist from the list of available checklists. Retrieve the checklist."
+        " evaluate the paper according to the checklist, and return results that include both"
+        " complete evaluation for each checklist item, and a general summary."
+        " if a particular checklist item succeeds, say PASS and then any relevant details."
+        " Include examples if relevant. If a particular checklist item fails, say FAIL and provide"
+        " Explanation. If unclear state OTHER and provide an explanation."
+        " Return this as a markdown table."
+        "\nThe available checklists are:"
+    ),
+    deps_type=ChecklistDependencies,
+)
+
+
+@checklist_agent.system_prompt
+def add_checklists(ctx: RunContext[ChecklistDependencies]) -> str:
+    """
+    Add available checklists to the system prompt.
+    
+    Args:
+        ctx: The run context
+        
+    Returns:
+        A string containing the list of available checklists
+    """
+    meta = all_checklists()
+    return "\n".join([f"- {c['id']}: {c['title']}" for c in meta["checklists"]])
+
+
+@checklist_agent.tool
+async def retrieve_text_from_pmid_tool(ctx: RunContext[ChecklistDependencies], pmid: str) -> str:
+    """
+    Lookup the text of a PubMed ID, using its PMID.
+
+    Args:
+        ctx: The run context
+        pmid: The PubMed ID to look up
+    
+    Returns: 
+        Full text if available, otherwise abstract
+    """
+    return await retrieve_text_from_pmid(ctx, pmid)
+
+
+@checklist_agent.tool
+async def retrieve_text_from_doi_tool(ctx: RunContext[ChecklistDependencies], doi: str) -> str:
+    """
+    Lookup the text of a DOI.
+
+    Args:
+        ctx: The run context
+        doi: The DOI to look up
+    
+    Returns: 
+        Full text if available, otherwise abstract
+    """
+    return await retrieve_text_from_doi(ctx, doi)
+
+
+@checklist_agent.tool
+async def fetch_checklist_tool(ctx: RunContext[ChecklistDependencies], checklist_id: str) -> str:
+    """
+    Lookup the checklist entry for a given checklist accession number.
+
+    Args:
+        ctx: The run context
+        checklist_id: The checklist ID (e.g. STREAM, STORMS, ARRIVE)
+    
+    Returns:
+        The content of the checklist
+    """
+    return await fetch_checklist(ctx, checklist_id)

aurelian/agents/checklist/checklist_config.py

@@ -0,0 +1,28 @@
+"""
+Configuration for the Checklist agent.
+"""
+from dataclasses import dataclass
+import os
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+
+
+@dataclass
+class ChecklistDependencies(HasWorkdir):
+    """Configuration for the Checklist agent."""
+    
+    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() -> ChecklistDependencies:
+    """Get the Checklist 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 ChecklistDependencies(
+        workdir=workdir,
+    )