aurelian 0.1.0__py3-none-any.whl

aurelian/agents/ontology_mapper/ontology_mapper_mcp.py

@@ -0,0 +1,81 @@
+"""
+MCP tools for creating ontology mappings.
+"""
+import os
+from typing import Dict, List, Optional
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.ontology_mapper.ontology_mapper_tools as omt
+from aurelian.agents.ontology_mapper.ontology_mapper_agent import ONTOLOGY_MAPPER_SYSTEM_PROMPT
+from aurelian.agents.ontology_mapper.ontology_mapper_config import OntologyMapperDependencies, get_config
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("ontology_mapper", instructions=ONTOLOGY_MAPPER_SYSTEM_PROMPT)
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> OntologyMapperDependencies:
+    deps = get_config()
+    # Set the location from environment variable or default
+    loc = os.getenv("AURELIAN_WORKDIR", "/tmp/aurelian")
+    deps.workdir = WorkDir(loc)
+    return deps
+
+def ctx() -> RunContext[OntologyMapperDependencies]:
+    rc: RunContext[OntologyMapperDependencies] = RunContext[OntologyMapperDependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def search_terms(query: str, ont: Optional[str] = None, limit: int = 10) -> List[Dict]:
+    """
+    Search for ontology terms matching a query.
+
+    Args:
+        query: The search query text
+        ont: Optional ontology ID to search in (e.g., 'cl', 'go', 'uberon')
+        limit: Maximum number of results to return
+        
+    Returns:
+        List of matching ontology terms with their details
+    """
+    return await omt.search_terms(ctx(), query, ont, limit)
+
+
+@mcp.tool()
+async def search_web(query: str) -> str:
+    """
+    Search the web for ontology-related information.
+
+    Args:
+        query: The search query
+        
+    Returns:
+        Search results with summaries
+    """
+    return await omt.search_web(ctx(), query)
+
+
+@mcp.tool()
+async def retrieve_web_page(url: str) -> str:
+    """
+    Fetch the contents of a web page related to ontologies.
+
+    Args:
+        url: The URL to fetch
+        
+    Returns:
+        The contents of the web page
+    """
+    return await omt.retrieve_web_page(ctx(), url)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')

aurelian/agents/ontology_mapper/ontology_mapper_tools.py

@@ -0,0 +1,147 @@
+"""
+Tools for the Ontology Mapper agent.
+"""
+import asyncio
+from functools import lru_cache
+from typing import Dict, List, Optional
+
+from oaklib import get_adapter
+from pydantic_ai import RunContext, ModelRetry
+
+from aurelian.utils.ontology_utils import search_ontology
+from aurelian.utils.search_utils import web_search, retrieve_web_page as fetch_web_page
+from .ontology_mapper_config import OntologyMapperDependencies, get_config
+
+
+@lru_cache
+def get_ontology_adapter(ont: str):
+    """
+    Get an adapter for the specified ontology.
+
+    Args:
+        ont: The ontology ID to get an adapter for (e.g. cl, go, uberon)
+
+    Returns:
+        An OAK adapter for the specified ontology
+    """
+    ont = ont.lower()
+    return get_adapter(f"sqlite:obo:{ont}")
+
+
+async def search_terms(
+    ctx: RunContext[OntologyMapperDependencies], 
+    ontology_id: str, 
+    query: str
+) -> List[Dict]:
+    """
+    Finds similar ontology terms to the search query.
+
+    For example:
+
+        ```
+        search_terms("go", "cycle cycle and related processes")
+        ```
+
+    Relevancy ranking is used, with semantic similarity, which means
+    queries need only be close in semantic space. E.g. while GO does not
+    deal with diseases, this may return relevant pathways or structures:
+
+        ```
+        search_terms("go", "terms most relevant to Parkinson disease")
+        ```
+
+    Args:
+        ctx: The run context
+        ontology_id: The ontology ID to search in (e.g. cl, go, uberon)
+        query: The search query
+        
+    Returns:
+        A list of matching ontology terms
+    """
+    print(f"Term Search: {ontology_id} {query}")
+    
+    try:
+        if " " in ontology_id:
+            raise ModelRetry(
+                "Invalid ontology ID, use an OBO style ID like cl, mondo, chebi, etc."
+            )
+            
+        config = ctx.deps or get_config()
+        if ontology_id.lower() not in [ont.lower() for ont in config.ontologies]:
+            allowed_onts = ", ".join(config.ontologies)
+            raise ModelRetry(
+                f"Ontology '{ontology_id}' not in allowed list: {allowed_onts}"
+            )
+            
+        adapter = get_ontology_adapter(ontology_id)
+        # Execute the potentially blocking operation in a thread pool
+        results = await asyncio.to_thread(
+            search_ontology, 
+            adapter, 
+            query, 
+            limit=config.max_search_results
+        )
+        
+        if not results:
+            raise ModelRetry(f"No results found for query '{query}' in ontology '{ontology_id}'")
+            
+        return results
+    except Exception as e:
+        if "ModelRetry" in str(type(e)):
+            raise e
+        raise ModelRetry(f"Error searching ontology: {str(e)}")
+
+
+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: The search query
+        
+    Returns: 
+        Matching web pages plus summaries
+    """
+    print(f"Web Search: {query}")
+    
+    try:
+        # 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 the web: {str(e)}")
+
+
+async def retrieve_web_page(url: str) -> str:
+    """
+    Fetch the contents of a web page.
+
+    Args:
+        url: The URL of the web page to retrieve
+        
+    Returns:
+        The contents of the web page
+    """
+    print(f"Fetch URL: {url}")
+    
+    try:
+        # 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 at 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/paperqa/__init__.py

@@ -0,0 +1,27 @@
+"""
+PaperQA agent package for scientific literature search and analysis.
+"""
+
+# isort: skip_file
+from .paperqa_agent import paperqa_agent  # noqa: E402
+from .paperqa_config import PaperQADependencies, get_config  # noqa: E402
+from .paperqa_gradio import chat  # noqa: E402
+from .paperqa_tools import (  # noqa: E402
+    search_papers,
+    query_papers,
+    add_paper,
+    add_papers,
+    list_papers,
+)
+
+__all__ = [
+    "paperqa_agent",
+    "PaperQADependencies",
+    "get_config",
+    "search_papers",
+    "query_papers",
+    "add_paper",
+    "add_papers",
+    "list_papers",
+    "chat",
+]

aurelian/agents/paperqa/paperqa_agent.py

@@ -0,0 +1,66 @@
+"""
+Agent for PaperQA integration with Aurelian.
+"""
+import logging
+from pydantic_ai import Agent
+
+paperqa_logger = logging.getLogger("aurelian.agents.paperqa")
+paperqa_logger.setLevel(logging.INFO)
+
+for handler in list(paperqa_logger.handlers):
+    paperqa_logger.removeHandler(handler)
+
+console = logging.StreamHandler()
+console.setLevel(logging.INFO)
+formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+console.setFormatter(formatter)
+paperqa_logger.addHandler(console)
+
+paperqa_logger.propagate = False
+
+from .paperqa_config import PaperQADependencies
+from .paperqa_tools import (
+    search_papers,
+    query_papers,
+    add_paper,
+    add_papers,
+    list_papers,
+    build_index
+)
+
+PAPERQA_SYSTEM_PROMPT = """
+        You are an AI assistant that helps explore scientific literature using PaperQA.
+        You can use different functions to search for papers and analyze them:
+          - `search_papers` to find papers by topic or keyword from outside this repository.
+          - `query_papers` to ask questions about the papers in the repository
+          - `add_paper` to add a specific paper by file path or URL (with auto_index=True by default)
+          - `add_papers` to add multiple papers from a directory (with auto_index=True by default)
+          - `list_papers` to see all papers in the collection
+          - `build_index` to manually rebuild the search index
+        
+        When adding papers with `add_paper` or `add_papers`:
+        - For `add_paper`, the URL must be a direct link to a PDF (e.g., "https://example.com/paper.pdf")
+        - For `add_paper`, you can provide a citation string to attribute the source
+        - For `add_papers`, you provide a directory containing papers and an optional citation format
+        - By default, auto_index=True, which automatically rebuilds the index after adding papers
+        - You can set auto_index=False if you want to add multiple papers before indexing
+        - After adding papers with auto_index=False, use `build_index()` to make them searchable
+        
+        When showing paper information, format using Markdown for readability.
+        When papers have been successfully retrieved, proceed to analyzing them.
+                        """
+
+paperqa_agent = Agent(
+    model="openai:gpt-4o-2024-11-20",
+    deps_type=PaperQADependencies,
+    result_type=str,
+    system_prompt=PAPERQA_SYSTEM_PROMPT,
+    defer_model_check=True,
+)
+
+paperqa_agent.tool(search_papers)
+paperqa_agent.tool(query_papers)
+paperqa_agent.tool(add_paper)
+paperqa_agent.tool(add_papers)
+paperqa_agent.tool(list_papers)
+paperqa_agent.tool(build_index)

aurelian/agents/paperqa/paperqa_cli.py

@@ -0,0 +1,305 @@
+"""
+CLI commands for the PaperQA agent.
+"""
+import os
+import asyncio
+import logging
+import sys
+from pathlib import Path
+import click
+from paperqa import agent_query
+
+from aurelian.agents.paperqa.paperqa_config import get_config
+from paperqa.agents.search import get_directory_index
+from paperqa.settings import IndexSettings
+
+logger = logging.getLogger(__name__)
+
+def setup_logging():
+    """Set up logging for the PaperQA CLI."""
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        handlers=[logging.StreamHandler(sys.stdout)]
+    )
+
+
+def check_api_key():
+    """Check if the OpenAI API key is set.
+
+    Returns:
+        bool: True if key is set, False otherwise
+    """
+    if not os.environ.get("OPENAI_API_KEY"):
+        logger.error("OPENAI_API_KEY environment variable must be set.")
+        click.echo("Error: OPENAI_API_KEY environment variable must be set.")
+        return False
+    return True
+
+
+def setup_and_configure_paper_directory(directory):
+    """
+    Setup and configure a paper directory with proper paths.
+
+    Args:
+        directory: Input directory path (can be relative)
+
+    Returns:
+        tuple: (resolved_path, settings, config) tuple with properly configured settings
+    """
+    directory = str(Path(directory).resolve())
+
+    config = get_config()
+    config.paper_directory = directory
+
+    os.environ["PQA_HOME"] = directory
+
+    if not os.path.exists(directory):
+        logger.info(f"Creating paper directory: {directory}")
+        os.makedirs(directory, exist_ok=True)
+
+    settings = config.set_paperqa_settings()
+    settings.agent.index = IndexSettings(
+        name=config.index_name,
+        paper_directory=directory,
+        recurse_subdirectories=False
+    )
+
+    return directory, settings, config
+
+
+def get_document_files(directory):
+    """
+    Get all indexable document files in the given directory.
+
+    Args:
+        directory: Directory to search for document files
+
+    Returns:
+        dict: Dictionary with file lists by type and a combined list
+    """
+    document_extensions = ['.pdf', '.txt', '.html', '.md']
+    all_files = [f for f in os.listdir(directory)
+                if any(f.lower().endswith(ext) for ext in document_extensions)]
+
+    return {
+        'all': all_files,
+        'pdf': [f for f in all_files if f.lower().endswith('.pdf')],
+        'txt': [f for f in all_files if f.lower().endswith('.txt')],
+        'html': [f for f in all_files if f.lower().endswith('.html')],
+        'md': [f for f in all_files if f.lower().endswith('.md')],
+    }
+
+
+@click.group(name="paperqa")
+@click.option("-v", "--verbose", count=True, help="Increase verbosity level (-v for INFO, -vv for DEBUG)")
+@click.option("-q", "--quiet", is_flag=True, help="Suppress non-error output")
+def paperqa_cli(verbose, quiet):
+    """PaperQA management commands for indexing and querying documents.
+    
+    PaperQA supports PDF, TXT, HTML, and Markdown files in all operations.
+    
+    Examples:
+        # Index documents in a directory
+        aurelian paperqa index -d /path/to/papers
+        
+        # Ask a question about indexed papers
+        aurelian paperqa ask "What is the role of tau protein in Alzheimer's?" -d /path/to/papers
+        
+        # List indexed papers
+        aurelian paperqa list -d /path/to/papers
+        
+        # Run with increased verbosity
+        aurelian paperqa --verbose index -d /path/to/papers
+        
+        # Add documents through the agent
+        # (Using these commands in chat modes like Gradio or MCP)
+        "Add the paper from /path/to/paper.pdf"
+        "Add all papers from the directory /path/to/papers/"
+    """
+    setup_logging()
+
+    if verbose >= 2:
+        logging.getLogger("aurelian.agents.paperqa").setLevel(logging.DEBUG)
+    elif verbose == 1:
+        logging.getLogger("aurelian.agents.paperqa").setLevel(logging.INFO)
+    else:
+        logging.getLogger("aurelian.agents.paperqa").setLevel(logging.WARNING)
+
+    if quiet:
+        logging.getLogger("aurelian.agents.paperqa").setLevel(logging.ERROR)
+
+
+@paperqa_cli.command()
+@click.option(
+    "--directory", "-d", 
+    required=True,
+    help="Paper directory containing PDF, TXT, HTML, and MD files to index",
+)
+def index(directory):
+    """Index documents for search and querying.
+    
+    This command scans the specified directory for documents (PDF, TXT, HTML, MD)
+    and creates a searchable index for them. The index is stored in the .pqa
+    subdirectory of the specified directory.
+    
+    Example:
+        aurelian paperqa index -d ~/research/papers
+    """
+    if not check_api_key():
+        return
+
+    paper_dir, settings, _ = setup_and_configure_paper_directory(directory)
+
+    docs = get_document_files(paper_dir)
+
+    if not docs['all']:
+        logger.warning(f"No indexable documents found in {paper_dir}")
+        click.echo(f"No indexable documents found in {paper_dir}")
+        return
+
+    # detailed breakdown
+    logger.info(f"Found {len(docs['all'])} documents in {paper_dir}:")
+    if docs['pdf']: logger.info(f"  - {len(docs['pdf'])} PDF files")
+    if docs['txt']: logger.info(f"  - {len(docs['txt'])} text files")
+    if docs['html']: logger.info(f"  - {len(docs['html'])} HTML files")
+    if docs['md']: logger.info(f"  - {len(docs['md'])} Markdown files")
+    logger.info(f"Index will be stored in: {paper_dir}/.pqa")
+    logger.info("Indexing papers... (this may take a while)")
+
+    async def run_index():
+        try:
+            index = await get_directory_index(
+                settings=settings,
+                build=True,
+            )
+            index_files = await index.index_files
+            logger.info(f"Success! Indexed {len(index_files)} document chunks from your PDF files.")
+        except Exception as e:
+            logger.error(f"Error indexing papers: {str(e)}")
+
+    try:
+        asyncio.run(run_index())
+    except Exception as e:
+        logger.error(f"Error: {str(e)}")
+
+
+@paperqa_cli.command()
+@click.argument("query", required=True)
+@click.option(
+"--directory", "-d",
+required=True,
+help="Paper directory containing indexed documents",
+)
+def ask(query, directory):
+    """Ask a question about the indexed documents.
+    
+    This command searches the indexed documents for information relevant to the
+    provided query and generates an AI-powered answer with references. Make sure
+    to run the 'index' command first to create an index.
+    
+    Example:
+        aurelian paperqa ask "What are the key findings on tau proteins?" -d ~/research/papers
+    """
+    if not check_api_key():
+        return
+
+    paper_dir, settings, _ = setup_and_configure_paper_directory(directory)
+
+    async def run_query():
+        try:
+            docs = get_document_files(paper_dir)
+
+            if not docs['all']:
+                logger.warning(f"No indexable documents found in {paper_dir}")
+                logger.info(f"Add documents (PDF, TXT, HTML, MD) to the directory and then run 'aurelian paperqa index -d {paper_dir}'")
+                return
+
+            try:
+                index = await get_directory_index(settings=settings, build=False)
+                index_files = await index.index_files
+
+                if not index_files:
+                    logger.warning(f"No indexed papers found. PDF files exist but haven't been indexed.")
+                    logger.info(f"Run 'aurelian paperqa index -d {paper_dir}' to index the papers.")
+                    return
+            except Exception as e:
+                if "was empty, please rebuild it" in str(e):
+                    logger.warning(f"Index is empty. Run 'aurelian paperqa index -d {paper_dir}' to index papers.")
+                    return
+                raise
+
+            logger.info(f"Querying {len(index_files)} papers about: {query}")
+            logger.info("This may take a moment...")
+
+            response = await agent_query(
+                query=query,
+                settings=settings
+            )
+
+            click.echo(f"Answer: {response.session.answer}" +
+                       f"\n\nReferences: {response.session.references}")
+
+        except Exception as e:
+            logger.error(f"Error querying papers: {str(e)}")
+
+    loop = asyncio.get_event_loop()
+    try:
+        loop.run_until_complete(run_query())
+    except Exception as e:
+        logger.error(f"Error: {str(e)}")
+
+
+@paperqa_cli.command()
+@click.option(
+    "--directory", "-d", 
+    required=True,
+    help="Paper directory containing documents",
+)
+def list(directory):
+    """List documents in the directory and their indexing status.
+    
+    This command displays all documents (PDF, TXT, HTML, MD) in the specified
+    directory and shows which ones have been indexed. Use this to verify that
+    your documents are properly recognized and indexed.
+    
+    Example:
+        aurelian paperqa list -d ~/research/papers
+    """
+    if not check_api_key():
+        return
+    
+    paper_dir, settings, _ = setup_and_configure_paper_directory(directory)
+
+    docs = get_document_files(paper_dir)
+
+    logger.info(f"Documents in directory {paper_dir}:")
+    for doc in docs['all']:
+        if doc.lower().endswith('.pdf'):
+            logger.info(f"  - {doc} [PDF]")
+        elif doc.lower().endswith('.txt'):
+            logger.info(f"  - {doc} [TXT]")
+        elif doc.lower().endswith('.html'):
+            logger.info(f"  - {doc} [HTML]")
+        elif doc.lower().endswith('.md'):
+            logger.info(f"  - {doc} [MD]")
+    
+    async def list_indexed():
+        try:
+            index = await get_directory_index(settings=settings, build=False)
+            index_files = await index.index_files
+            if index_files:
+                logger.info(f"Indexed papers ({len(index_files)}):")
+                for file in index_files:
+                    logger.info(f"  - {file}")
+            else:
+                logger.warning(f"No indexed papers found. Run 'aurelian paperqa index -d {paper_dir}' to index papers.")
+        except Exception as e:
+            logger.error(f"Error accessing index: {str(e)}")
+            logger.info(f"Run 'aurelian paperqa index -d {paper_dir}' to create or rebuild the index.")
+    
+    loop = asyncio.get_event_loop()
+    try:
+        loop.run_until_complete(list_indexed())
+    except Exception as e:
+        logger.error(f"Error: {str(e)}")