aurelian 0.1.0__py3-none-any.whl

aurelian/mcp/linkml_mcp_tools.py

@@ -0,0 +1,190 @@
+"""
+MCP tools for creating LinkML schemas and example datasets
+"""
+from linkml.generators import JsonSchemaGenerator
+from mcp.server.fastmcp import FastMCP
+
+from aurelian.agents.linkml.linkml_agent import LinkMLDependencies
+
+# Initialize FastMCP server
+mcp = FastMCP("linkml")
+
+import logfire
+from linkml_runtime.loaders import yaml_loader
+from linkml_runtime.linkml_model import SchemaDefinition
+from linkml.validator import validate
+from pydantic_ai import RunContext, ModelRetry
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> LinkMLDependencies:
+    deps = LinkMLDependencies()
+    deps.workdir = WorkDir("/tmp/linkml")
+    return deps
+
+def ctx() -> RunContext[LinkMLDependencies]:
+    rc: RunContext[LinkMLDependencies] = RunContext[LinkMLDependencies](deps=deps())
+    return rc
+
+
+@mcp.tool()
+async def validate_schema(schema: str, save_to_file: str="schema.yaml") -> str:
+    """
+    Validate a LinkML schema.
+
+    Args:
+        ctx: context
+        schema: schema (as yaml) to validate. Do not truncate, always pass the whole schema.
+        save_to_file: optional file name to save the schema to. Defaults to schema.yaml
+
+    Returns:
+
+    """
+    print(f"Validating schema: {schema}")
+    try:
+        schema = yaml_loader.loads(schema, target_class=SchemaDefinition)
+        gen = JsonSchemaGenerator(schema)
+        gen.serialize()
+        if save_to_file and schema:
+            deps().workdir.write_file(save_to_file, schema)
+    except Exception as e:
+        raise ModelRetry(f"Schema does not validate: {e}")
+    return "VALIDATES"
+
+
+@mcp.tool()
+async def inspect_file(data_file: str) -> str:
+    """
+    Inspect a file in the working directory.
+
+    Args:
+        ctx:
+        data_file: name of file
+
+    Returns:
+
+    """
+    print(f"Inspecting file: {data_file}")
+    return deps().workdir.read_file(data_file)
+
+
+@mcp.tool()
+async def list_files() -> str:
+    """
+    List files in the working directory.
+
+    Args:
+        ctx:
+
+    Returns:
+
+    """
+    return "\n".join(deps().workdir.list_file_names())
+
+@mcp.tool()
+async def write_to_file(data: str, file_name: str) -> str:
+    """
+    Write data to a file in the working directory.
+
+    Args:
+        ctx:
+        data:
+        file_name:
+
+    Returns:
+
+    """
+    print(f"Writing data to file: {file_name}")
+    deps().workdir.write_file(file_name, data)
+    return f"Data written to {file_name}"
+
+@mcp.tool()
+async def validate_data(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:
+
+    """
+    logfire.log(f"Validating data file: {data_file} using schema: {schema}")
+    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:
+        instances = 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:
+        return f"Data does not validate: {e}"
+
+
+@mcp.tool()
+async def search_web(query: str) -> str:
+    """
+    Search the web using a text query.
+
+    Note, this will not retrieve the full content, for that you
+    should use `retrieve_web_page`.
+
+    Args:
+        query: Text query
+
+    Returns: matching web pages plus summaries
+    """
+    print(f"Web Search: {query}")
+    return web_search(query)
+
+@mcp.tool()
+async def retrieve_web_page(url: str) -> str:
+    """
+    Fetch the contents of a web page.
+
+    Args:
+        url: URL of the web page
+
+    Returns:
+        The contents of the web page.
+    """
+    print(f"Fetch URL: {url}")
+    import aurelian.utils.search_utils as su
+    return su.retrieve_web_page(url)
+
+
+@mcp.tool()
+async def download_web_page(url: str, local_file_name: str) -> str:
+    """
+    Download contents of a web page.
+
+    Args:
+        ctx:
+        url: URL of the web page
+        local_file_name: Name of the local file to save the
+
+    Returns:
+        str: message
+    """
+    print(f"Fetch URL: {url}")
+    import aurelian.utils.search_utils as su
+    data = su.retrieve_web_page(url)
+    deps().workdir.write_file(local_file_name, data)
+    return f"Data written to {local_file_name}"
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')

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,18 @@
+import asyncio
+from typing import Callable, Coroutine
+
+
+def run_sync(f: Callable | Coroutine):
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    if isinstance(f, Coroutine):
+        result = f
+    else:
+        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")
+