aurelian 0.3.2__py3-none-any.whl

aurelian/mcp/mcp_discovery.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python
+"""
+Script for discovering and testing MCP implementations.
+"""
+
+import argparse
+import importlib
+import inspect
+import sys
+from pathlib import Path
+from typing import List, Optional
+
+
+def list_mcp_tools(module_path: str) -> List[str]:
+    """
+    List all MCP tools in a given module.
+    
+    Args:
+        module_path: Dot-separated path to the module
+        
+    Returns:
+        List of tool names
+    """
+    try:
+        # Import the module
+        module = importlib.import_module(module_path)
+        
+        # Find the MCP server instance
+        mcp = getattr(module, "mcp", None)
+        if not mcp:
+            print(f"No 'mcp' instance found in {module_path}")
+            return []
+        
+        # Get all functions with the MCP tool decorator
+        tools = []
+        for name, func in inspect.getmembers(module, inspect.isfunction):
+            # Check if this function is an MCP tool
+            if hasattr(func, "__mcp_tool__") and func.__mcp_tool__ is True:
+                tools.append(name)
+                
+        return tools
+    except Exception as e:
+        print(f"Error importing {module_path}: {e}")
+        return []
+
+
+def main():
+    """Command line interface."""
+    parser = argparse.ArgumentParser(description="Discover and test MCP implementations")
+    parser.add_argument("--agent", "-a", help="Agent type to test (e.g., 'diagnosis')")
+    parser.add_argument("--list", "-l", action="store_true", help="List available MCP modules")
+    args = parser.parse_args()
+    
+    # Base path for agent modules
+    base_path = Path(__file__).parent.parent
+    
+    if args.list:
+        # List all MCP modules
+        agents_dir = base_path / "agents"
+        print("Available MCP modules:")
+        for agent_dir in agents_dir.iterdir():
+            if agent_dir.is_dir():
+                mcp_file = agent_dir / f"{agent_dir.name}_mcp.py"
+                if mcp_file.exists():
+                    print(f"  - {agent_dir.name}")
+        return
+        
+    if not args.agent:
+        print("Error: Please specify an agent type with --agent")
+        sys.exit(1)
+        
+    # Check if the MCP module exists
+    agent_type = args.agent
+    module_path = f"aurelian.agents.{agent_type}.{agent_type}_mcp"
+    
+    # List the tools
+    tools = list_mcp_tools(module_path)
+    if tools:
+        print(f"MCP tools for {agent_type}:")
+        for tool in tools:
+            print(f"  - {tool}")
+    else:
+        print(f"No MCP tools found for {agent_type}")
+
+
+if __name__ == "__main__":
+    main()

aurelian/mcp/mcp_test.py

@@ -0,0 +1,31 @@
+"""
+MCP tools for creating LinkML schemas and example datasets
+"""
+import sys
+
+from mcp.server.fastmcp import FastMCP
+
+
+# Initialize FastMCP server
+mcp = FastMCP("test")
+
+@mcp.tool()
+async def add_two_numbers(n1: int, n2: int) -> int:
+    """
+    Add two numbers together.
+
+    Args:
+        n1: first number
+        n2: second number
+
+    Returns:
+        sum of the two numbers
+    """
+    return n1 + n2
+
+
+if __name__ == "__main__":
+    print("Running the LinkML MCP tools")
+    print("Use Ctrl-C to exit", file=sys.stderr)
+    # Initialize and run the server
+    mcp.run(transport='stdio')

aurelian/mcp/phenopackets_mcp_test.py

@@ -0,0 +1,103 @@
+"""
+Test for phenopackets MCP functionality
+"""
+import os
+import tempfile
+from typing import List, Optional, Dict
+
+# try to import, don't die if import fails
+try:
+    from mcp import Client
+except ImportError:
+    print("mcp package not found. Please install it to run this test.")
+    Client = None
+from pydantic import BaseModel
+
+
+class Message(BaseModel):
+    role: str
+    content: str
+
+
+async def test_phenopackets_mcp():
+    """Test the phenopackets MCP agent."""
+    client = Client("/tmp/phenopackets-mcp", exec_args=["python", "-m", "aurelian.agents.phenopackets.phenopackets_mcp"])
+
+    import time
+    time.sleep(1)  # Give the server time to start
+
+    # Set up a temporary working directory for the test
+    tempdir = os.path.join(tempfile.gettempdir(), "test_phenopackets")
+    os.makedirs(tempdir, exist_ok=True)
+    os.environ["AURELIAN_WORKDIR"] = tempdir
+
+    # For testing without a real database - would normally come from configuration
+    os.environ["PHENOPACKETS_DB_PATH"] = "mongodb://localhost:27017/test_phenopackets"
+    os.environ["PHENOPACKETS_DB_NAME"] = "test_phenopackets"
+    os.environ["PHENOPACKETS_COLLECTION_NAME"] = "test_main"
+
+    convo: List[Message] = []
+
+    def add_to_convo(role: str, content: str) -> Message:
+        msg = Message(role=role, content=content)
+        convo.append(msg)
+        return msg
+
+    # Make a query
+    add_to_convo("user", "What tools are available for working with phenopackets?")
+
+    message = convo[-1].content
+
+    response = await client.chat(messages=message)
+    print(f"Got response: {response}")
+
+    # Get available tools
+    tool_choices = await client.get_tool_choice(messages=convo)
+    print(f"Available tools: {[t['id'] for t in tool_choices]}")
+
+    # Test the list_files tool
+    list_files_tool = next((t for t in tool_choices if t["id"] == "list_files"), None)
+    if list_files_tool:
+        print(f"Testing list_files tool...")
+        tool_input_schema = await client.get_tool_input_schema(tool_id=list_files_tool["id"])
+        tool_result = await client.execute_tool(tool_id=list_files_tool["id"], tool_input='{}')
+        print(f"list_files result: {tool_result}")
+
+    # Create a test file
+    write_file_tool = next((t for t in tool_choices if t["id"] == "write_to_file"), None)
+    if write_file_tool:
+        print(f"Testing write_to_file tool...")
+        tool_input = '{"file_name": "test.txt", "data": "This is a test file for phenopackets MCP"}'
+        tool_result = await client.execute_tool(tool_id=write_file_tool["id"], tool_input=tool_input)
+        print(f"write_to_file result: {tool_result}")
+
+    # Check if the file was created
+    if list_files_tool:
+        tool_result = await client.execute_tool(tool_id=list_files_tool["id"], tool_input='{}')
+        print(f"list_files after writing: {tool_result}")
+
+    # Read the file
+    inspect_file_tool = next((t for t in tool_choices if t["id"] == "inspect_file"), None)
+    if inspect_file_tool:
+        print(f"Testing inspect_file tool...")
+        tool_input = '{"data_file": "test.txt"}'
+        tool_result = await client.execute_tool(tool_id=inspect_file_tool["id"], tool_input=tool_input)
+        print(f"inspect_file result: {tool_result}")
+
+    # Try a web search
+    search_web_tool = next((t for t in tool_choices if t["id"] == "search_web"), None)
+    if search_web_tool:
+        print(f"Testing search_web tool...")
+        tool_input = '{"query": "phenopackets specification"}'
+        try:
+            tool_result = await client.execute_tool(tool_id=search_web_tool["id"], tool_input=tool_input)
+            print(f"search_web result: {tool_result[:200]}...")  # Just show first 200 chars
+        except Exception as e:
+            print(f"search_web failed (expected in test): {e}")
+
+    await client.close()
+
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(test_phenopackets_mcp())

aurelian/tools/web/url_download.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass
+
+from pydantic import TypeAdapter
+from pydantic_ai.tools import Tool
+from typing_extensions import TypedDict
+
+from aurelian.dependencies.workdir import WorkDir
+
+
+__all__ = ["download_url_tool"]
+
+class URLDownloadResult(TypedDict):
+    """The result of downloading a URL."""
+
+    size: int
+    """The size of the downloaded content."""
+    location: str
+    """The location of the downloaded content."""
+
+ta = TypeAdapter(list[URLDownloadResult])
+
+@dataclass
+class URLDownloadTool:
+    """A tool for downloading a URL."""
+    workdir: WorkDir
+
+    async def __call__(self, url:str, local_file_name: str) -> list[dict]:
+        """Download the contents of a URL.
+
+        Args:
+            url: URL of the web page
+            local_file_name: Name of the local file to save
+
+        Returns:
+            The contents of the web page.
+        """
+        import asyncio
+        import aurelian.utils.search_utils as su
+        data = await asyncio.to_thread(su.retrieve_web_page, url)
+        self.workdir.write_file(local_file_name, data)
+        return ta.validate_python([{"location": local_file_name, "size": len(data)}])
+
+def download_url_tool(workdir: WorkDir):
+    """Create a URL download tool."""
+    udt = URLDownloadTool(workdir=workdir)
+    c = udt.__call__
+    return Tool(
+        udt.__call__,
+        name="download_web_page",
+        description="Download the contents of a URL.",
+    )

aurelian/utils/async_utils.py

@@ -0,0 +1,15 @@
+import asyncio
+from typing import Callable
+
+
+def run_sync(f: Callable):
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    result = f()
+
+    # Ensure it's a coroutine before running it
+    if asyncio.iscoroutine(result):
+        result = loop.run_until_complete(result)
+
+    loop.close()
+    return result

aurelian/utils/data_utils.py

@@ -0,0 +1,32 @@
+from typing import Dict, List, Optional, Union
+
+from linkml_runtime.dumpers import json_dumper
+from linkml_runtime.utils.yamlutils import YAMLRoot
+from pydantic import BaseModel
+
+
+def flatten(d: Dict, preserve_keys: Optional[List] = None) -> Dict:
+    """Flatten a dictionary"""
+    out = {}
+    for k, v in d.items():
+        if isinstance(v, list):
+            if preserve_keys and k in preserve_keys:
+                out[k] = [flatten(x, preserve_keys=preserve_keys) for x in v]
+            else:
+                out[f"{k}_count"] = len(v)
+        elif isinstance(v, dict):
+            out[k] = flatten(v, preserve_keys=preserve_keys)
+        else:
+            out[k] = v
+    return out
+
+
+def obj_to_dict(obj: Union[object, YAMLRoot, BaseModel, Dict]) -> Dict:
+    if isinstance(obj, YAMLRoot):
+        return json_dumper.to_dict(obj)
+    elif isinstance(obj, BaseModel):
+        return obj.model_dump()
+    elif isinstance(obj, dict):
+        return obj
+    else:
+        raise ValueError(f"Cannot convert object of type {type(obj)} to dict")

aurelian/utils/documentation_manager.py

@@ -0,0 +1,59 @@
+from pydantic import BaseModel
+from pathlib import Path
+from typing import Dict, Optional, List
+from dataclasses import dataclass
+
+class Document(BaseModel):
+    """
+    A document is a file in the documentation directory.
+    """
+    id: str
+    title: str
+    path: str
+    metadata: Optional[Dict] = None
+
+@dataclass
+class DocumentationManager:
+    """
+    A manager for the documentation directory.
+    """
+    
+    documents_dir: Path
+    collection_name: Optional[str] = None
+
+    def all_documents(self) -> List[Document]:
+        """
+        Get all available documents.
+        """
+        return [Document(
+            id=file_path.stem,
+            title=file_path.stem.replace("_", " "),
+            path=str(file_path),
+            metadata=None
+        ) for file_path in self.documents_dir.glob("*.md")]
+
+    def get_documents_for_prompt(self, extra_text: str = "") -> str:
+        """
+        Get the documents for the system prompt.
+
+        Returns:
+            A string containing the list of available GO annotation best practice documents
+        """
+        docs = self.all_documents()
+        if not docs:
+            raise AssertionError("No best practice documents are available")
+
+        docs_text = "\n\nThe following documents are available:\n"
+        docs_text += "\n".join([f"- {d.title}" for d in docs])
+        docs_text += "\n\n" + extra_text
+        return docs_text
+
+    def fetch_document(self, id_or_title: str) -> Document:
+        """
+        Fetch a document by its ID or title.
+        """
+        for document in self.all_documents():
+            if document.id == id_or_title or document.title == id_or_title:
+                return document
+        raise KeyError(f"Document with ID or title '{id_or_title}' not found")
+

aurelian/utils/doi_fetcher.py

@@ -0,0 +1,238 @@
+import os
+import re
+from tempfile import NamedTemporaryFile
+from typing import Any, Dict, List, Optional
+
+import logfire
+import requests
+import requests_cache
+from bs4 import BeautifulSoup
+from markitdown import MarkItDown
+from openai import BaseModel
+from pydantic import Field
+
+
+class FullTextInfo(BaseModel):
+    """Data model for full text information."""
+
+    success: bool = True
+    abstract: Optional[str] = Field(None, description="Abstract of the article")
+    text: Optional[str] = Field(None, description="Full text of the article")
+    source: Optional[str] = Field(None, description="Source of the full text")
+    metadata: Optional[Dict[str, Any]] = Field(None, description="Metadata of the article")
+    pdf_url: Optional[str] = Field(None, description="URL to the PDF version of the article")
+
+
+class DOIFetcher:
+    """Fetch metadata and full text for a DOI using various APIs."""
+
+    def __init__(self, email: Optional[str] = None, url_prefixes: Optional[List[str]] = None):
+        """Initialize the DOI fetcher with a contact email (required by some APIs).
+
+        Args:
+            email (str): Contact email for API access
+            url_prefixes (List[str]): List of URL prefixes to check for full text
+
+        """
+        self.email = email or os.getenv("EMAIL") or "test@example.com"
+        self.url_prefixes = url_prefixes or os.getenv("DOI_FULL_TEXT_URLS", "").split(",")
+        self.headers = {"User-Agent": f"DOIFetcher/1.0 (mailto:{email})", "Accept": "application/json"}
+
+    def clean_text(self, text: str) -> str:
+        """Clean extracted text by removing extra whitespace and normalized characters.
+
+        Args:
+            text:
+
+        Returns:
+            str: The cleaned text
+
+        """
+        # Remove extra whitespace
+        text = re.sub(r"\s+", " ", text)
+        # Remove non-printable characters
+        text = "".join(char for char in text if char.isprintable())
+        return text.strip()
+
+    def get_metadata(self, doi: str, strict=False) -> Optional[Dict[str, Any]]:
+        """Fetch metadata for a DOI using the Crossref API.
+
+        Args:
+            doi (str): The DOI to look up
+            strict (bool): Raise exceptions if API call fails
+
+        Returns:
+            Optional[Dict[str, Any]]: Metadata dictionary if successful, None otherwise
+
+        """
+        base_url = "https://api.crossref.org/works/"
+        try:
+            response = requests.get(f"{base_url}{doi}", headers=self.headers)
+            response.raise_for_status()
+            return response.json()["message"]
+        except Exception as e:
+            if strict:
+                raise e
+            logfire.warn(f"Error fetching metadata: {e}")
+            return None
+
+    def get_unpaywall_info(self, doi: str, strict=False) -> Optional[Dict[str, Any]]:
+        """Check Unpaywall for open access versions.
+
+        Example:
+            >>> fetcher = DOIFetcher()
+            >>> doi = "10.1038/nature12373"
+            >>> unpaywall_data = fetcher.get_unpaywall_info(doi)
+            >>> assert unpaywall_data["doi"] == doi
+            >>> unpaywall_data["best_oa_location"]["url_for_pdf"]
+            'https://europepmc.org/articles/pmc4221854?pdf=render'
+
+        Args:
+            doi (str): The DOI to look up
+            strict (bool): Raise exceptions if API call fails
+
+        Returns:
+            Optional[Dict[str, Any]]: Unpaywall data if successful, None otherwise
+
+        """
+        base_url = f"https://api.unpaywall.org/v2/{doi}"
+        try:
+            response = requests.get(f"{base_url}?email={self.email}")
+            response.raise_for_status()
+            return response.json()
+        except Exception as e:
+            if strict:
+                raise e
+            logfire.warn(f"Error fetching Unpaywall data: {e}")
+            return None
+
+    def get_full_text(self, doi: str, fallback_to_abstract=True) -> Optional[str]:
+        """Get the full text of a paper using various methods.
+
+        Example:
+            >>> fetcher = DOIFetcher()
+            >>> doi = "10.1128/msystems.00045-18"
+            >>> full_text = fetcher.get_full_text(doi)
+            >>> assert "Populus Microbiome" in full_text
+
+        Args:
+            doi:
+            fallback_to_abstract:
+
+        Returns:
+            str: The full text if available, else abstract text if fallback_to_abstract,
+              else None
+
+        """
+        info = self.get_full_text_info(doi)
+        if not info:
+            return None
+        text = info.text
+        if text:
+            return self.clean_text(text)
+        if info.pdf_url:
+            text = self.text_from_pdf_url(info.pdf_url)
+            if text:
+                return self.clean_text(text)
+        message = "FULL TEXT NOT AVAILABLE"
+        if fallback_to_abstract:
+            metadata = info.metadata or {}
+            abstract = metadata.get("abstract")
+            if abstract:
+                return self.clean_text(abstract) + f"\n\n{message}"
+        return message
+
+    def get_full_text_info(self, doi: str) -> Optional[FullTextInfo]:
+        """Attempt to get the full text of a paper using various methods.
+
+            >>> fetcher = DOIFetcher()
+            >>> doi = "10.1128/msystems.00045-18"
+            >>> info = fetcher.get_full_text_info(doi)
+            >>> metadata = info.metadata
+            >>> metadata["type"]
+            'journal-article'
+            >>> metadata["title"][0][0:20]
+            'Exploration of the B'
+            >>> assert info.pdf_url is not None
+            >>> info.pdf_url
+            'https://europepmc.org/articles/pmc6172771?pdf=render'
+
+        Args:
+            doi (str): The DOI to fetch
+
+        Returns:
+            FullTextInfo: Full text information
+
+        """
+        # Get metadata
+        metadata = self.get_metadata(doi)
+
+        # Check Unpaywall
+        unpaywall_data = self.get_unpaywall_info(doi)
+        if unpaywall_data and unpaywall_data.get("is_oa"):
+            locations = unpaywall_data.get("oa_locations", [])
+            if unpaywall_data.get("best_oa_location"):
+                best_oa_location = unpaywall_data.get("best_oa_location")
+                locations = [best_oa_location] + locations
+
+            # Find best open access location
+            for location in locations:
+                pdf_url = location.get("url_for_pdf")
+                if pdf_url:
+                    return FullTextInfo(text=None, pdf_url=pdf_url, source="unpaywall", metadata=metadata)
+
+        # Fallback
+        url_prefixes = os.getenv("DOI_FULL_TEXT_URLS", "").split(",")
+
+        for url_prefix in url_prefixes:
+            url_prefix.rstrip("/")
+            url = f"{url_prefix}/{doi}"
+            try:
+                response = requests.get(url)
+                if response.status_code == 200:
+                    soup = BeautifulSoup(response.text, "html.parser")
+                    pdf_embed = soup.find("embed", id="pdf")
+                    if pdf_embed and pdf_embed.get("src"):
+                        pdf_url = pdf_embed["src"]
+                        # Remove any URL parameters after #
+                        pdf_url = pdf_url.split("#")[0]
+                        if not pdf_url.startswith("http"):
+                            pdf_url = "https:" + pdf_url
+                        return FullTextInfo(
+                            pdf_url=pdf_url,
+                            source=url,
+                            metadata=metadata,
+                        )
+            except Exception:
+                continue
+
+    def text_from_pdf_url(self, pdf_url: str, raise_for_status=False) -> Optional[str]:
+        """Extract text from a PDF URL.
+
+        Example:
+            >>> fetcher = DOIFetcher()
+            >>> pdf_url = "https://ceur-ws.org/Vol-1747/IT201_ICBO2016.pdf"
+            >>> text = fetcher.text_from_pdf_url(pdf_url)
+            >>> assert "biosphere" in text
+
+        Args:
+            pdf_url:
+            raise_for_status:
+
+        Returns:
+
+        """
+        session = requests_cache.CachedSession("pdf_cache")
+        # Download the PDF
+        response = session.get(pdf_url)
+        if raise_for_status:
+            response.raise_for_status()
+        if response.status_code != 200:
+            return None
+        with NamedTemporaryFile(delete=False) as tmpf:
+            tmpf.write(response.content)
+            tmp_name = tmpf.name
+            with open(tmp_name, "wb") as f:
+                f.write(response.content)
+            md = MarkItDown()
+            return md.convert(tmpf.name).text_content