aurelian 0.3.3__py3-none-any.whl → 0.4.1__py3-none-any.whl

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)}")

aurelian/agents/paperqa/paperqa_config.py

@@ -0,0 +1,142 @@
+"""
+Configuration for the PaperQA agent.
+"""
+from dataclasses import dataclass, field
+import os
+from typing import Optional, List
+
+from paperqa import Settings as PQASettings
+from paperqa.settings import (
+    AnswerSettings,
+    ParsingSettings,
+    PromptSettings,
+    AgentSettings,
+    IndexSettings, Settings,
+)
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+
+
+@dataclass
+class PaperQADependencies(HasWorkdir):
+    """Configuration for the PaperQA agent."""
+
+    paper_directory: str = field(
+        default_factory=lambda: os.getcwd(),
+        metadata={"description": "Directory containing papers to be searched"}
+    )
+    index_name: Optional[str] = field(
+        default=None,
+        metadata={"description": "Optional name for the search index. If None, it will be generated based on settings."}
+    )
+
+    llm: str = field(
+        default="gpt-4.1-2025-04-14",
+        metadata={"description": "LLM to use for queries and answer generation. Default is gpt-4.1-2025-04-14."}
+    )
+    summary_llm: str = field(
+        default="gpt-4.1-2025-04-14",
+        metadata={"description": "LLM to use for summarization. Default is gpt-4.1-2025-04-14."}
+    )
+    embedding: str = field(
+        default="text-embedding-3-small",
+        metadata={"description": "Embedding model to use. Default is text-embedding-3-small."}
+    )
+    temperature: float = field(
+        default=0.1,
+        metadata={"description": "Temperature for LLM generation. Default is 0.1."}
+    )
+
+    search_count: int = field(
+        default=8,
+        metadata={"description": "Number of papers to retrieve in searches. Default is 8."}
+    )
+
+    evidence_k: int = field(
+        default=10,
+        metadata={"description": "Number of evidence pieces to retrieve. Default is 10."}
+    )
+    answer_max_sources: int = field(
+        default=5,
+        metadata={"description": "Maximum number of sources to use in answers. Default is 5."}
+    )
+    max_concurrent_requests: int = field(
+        default=4,
+        metadata={"description": "Maximum number of concurrent requests to LLMs. Default is 4."}
+    )
+
+    chunk_size: int = field(
+        default=5000,
+        metadata={"description": "Size of document chunks for embedding. Default is 5000."}
+    )
+    overlap: int = field(
+        default=250,
+        metadata={"description": "Overlap between chunks. Default is 250."}
+    )
+
+    workdir: Optional[WorkDir] = None
+
+    def __post_init__(self):
+        """Initialize the config with default values."""
+        if self.workdir is None:
+            self.workdir = WorkDir()
+
+    def set_paperqa_settings(self) -> PQASettings:
+        """
+        Convert to PaperQA Settings object.
+
+        This allows users to customize all PaperQA settings through the dependencies object.
+        Any changes to the dependencies will be reflected in the returned Settings object.
+        """
+        return PQASettings(
+            llm=self.llm,
+            summary_llm=self.summary_llm,
+            embedding=self.embedding,
+            temperature=self.temperature,
+
+            answer=AnswerSettings(
+                evidence_k=self.evidence_k,
+                answer_max_sources=self.answer_max_sources,
+                max_concurrent_requests=self.max_concurrent_requests,
+            ),
+
+            parsing=ParsingSettings(
+                chunk_size=self.chunk_size,
+                overlap=self.overlap,
+            ),
+
+            agent=AgentSettings(
+                agent_llm=self.llm,
+                search_count=self.search_count,
+                index=IndexSettings(
+                    name=self.index_name,
+                    paper_directory=self.paper_directory,
+                    recurse_subdirectories=False,
+                ),
+            ),
+        )
+
+
+def get_config() -> PaperQADependencies:
+    """
+    Get the PaperQA configuration from environment variables or defaults.
+
+    Returns:
+        A PaperQADependencies instance with default settings.
+
+    Note:
+        Users can modify the returned object to customize settings.
+        Example:
+            ```python
+            deps = get_config()
+            deps.llm = "claude-3-sonnet-20240229"  # Use Claude instead of default GPT-4
+            deps.temperature = 0.5  # Increase temperature
+            deps.evidence_k = 15  # Retrieve more evidence
+            ```
+    """
+    workdir_path = os.environ.get("AURELIAN_WORKDIR", None)
+    workdir = WorkDir(location=workdir_path) if workdir_path else None
+
+    return PaperQADependencies(
+        workdir=workdir,
+    )

aurelian/agents/paperqa/paperqa_gradio.py

@@ -0,0 +1,90 @@
+"""
+Gradio interface for the PaperQA agent.
+"""
+import os
+import logging
+from typing import List, Optional, Any
+
+import gradio as gr
+
+from aurelian.utils.async_utils import run_sync
+from .paperqa_agent import paperqa_agent
+from .paperqa_config import PaperQADependencies, get_config
+
+logger = logging.getLogger(__name__)
+
+
+async def get_info(query: str, history: List[str], deps: PaperQADependencies, **kwargs) -> str:
+    """
+    Process a query using the PaperQA agent.
+    
+    Args:
+        query: The user query
+        history: The conversation history
+        deps: The dependencies for the agent
+        model: Optional model override
+        
+    Returns:
+        The agent's response
+    """
+    logger.info(f"QUERY: {query}")
+    logger.debug(f"HISTORY: {history}")
+    
+    if history:
+        query += "\n\n## Previous Conversation:\n"
+        for h in history:
+            query += f"\n{h}"
+    
+    result = await paperqa_agent.run(query, deps=deps, **kwargs)
+    return result.data
+
+
+def chat(deps: Optional[PaperQADependencies] = None, model=None, **kwargs):
+    """
+    Create a Gradio chat interface for the PaperQA agent.
+
+    Args:
+        deps: Optional dependencies configuration
+        model: Optional model override
+        kwargs: Additional keyword arguments for dependencies
+
+    Returns:
+        A Gradio ChatInterface
+    """
+    if deps is None:
+        deps = get_config()
+
+        for key, value in kwargs.items():
+            if hasattr(deps, key):
+                setattr(deps, key, value)
+
+    paper_dir = os.path.join(os.getcwd(), "pqa_source")
+    os.makedirs(paper_dir, exist_ok=True)
+    deps.paper_directory = paper_dir
+    os.environ["PQA_HOME"] = paper_dir
+    print(f"Using dedicated papers directory at: {paper_dir}")
+
+    def get_info_wrapper(query: str, history: List[str]) -> str:
+        """Wrapper for the async get_info function."""
+        import asyncio
+        return asyncio.run(get_info(query, history, deps, **kwargs))
+
+    return gr.ChatInterface(
+        fn=get_info_wrapper,
+        type="messages",
+        title="PaperQA AI Assistant",
+        description="""This assistant helps you search and analyze scientific papers. You can:
+        - Search for papers on a topic
+        - Ask questions about the papers in the repository
+        - Add specific papers by path or URL:  (if paths, use absolute paths!)"
+        - Add multiple papers from a directory: (use absolute paths!)"
+        - List all papers in the collection
+        
+        Supported document types: PDF, TXT, HTML, and Markdown files""",
+        examples=[
+            ["Search for papers on CRISPR gene editing"],
+            ["What are the main challenges in CRISPR gene editing?"],
+            ["What is the relationship between CRISPR and Cas9?"],
+            ["List all the indexed papers"],
+        ],
+    )

aurelian/agents/paperqa/paperqa_mcp.py

@@ -0,0 +1,155 @@
+"""
+MCP tools for working with PaperQA for scientific literature search and analysis.
+"""
+import os
+from typing import Dict, List, Any, Optional
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.paperqa.paperqa_tools as pt
+from aurelian.agents.paperqa.paperqa_agent import paperqa_agent, PAPERQA_SYSTEM_PROMPT
+from aurelian.agents.paperqa.paperqa_config import PaperQADependencies
+from pydantic_ai import RunContext
+
+mcp = FastMCP("paperqa", instructions=PAPERQA_SYSTEM_PROMPT.strip())
+
+
+from aurelian.dependencies.workdir import WorkDir
+
+def deps() -> PaperQADependencies:
+    deps = PaperQADependencies()
+    loc = os.getenv("AURELIAN_WORKDIR", "/tmp/paperqa")
+    deps.workdir = WorkDir(loc)
+    
+    paper_dir = os.getenv("PAPERQA_PAPER_DIRECTORY", os.path.join(loc, "papers"))
+    deps.paper_directory = paper_dir
+    
+    if os.getenv("PAPERQA_LLM"):
+        deps.llm = os.getenv("PAPERQA_LLM")
+    
+    if os.getenv("PAPERQA_EMBEDDING"):
+        deps.embedding = os.getenv("PAPERQA_EMBEDDING")
+    
+    return deps
+
+def ctx() -> RunContext[PaperQADependencies]:
+    rc: RunContext[PaperQADependencies] = RunContext[PaperQADependencies](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def search_papers(query: str, max_papers: Optional[int] = None) -> Any:
+    """
+    Search for papers relevant to the query using PaperQA.
+    
+    Args:
+        query: The search query for scientific papers
+        max_papers: Maximum number of papers to return (overrides config)
+        
+    Returns:
+        The search results with paper information
+        
+    This searches for scientific papers based on your query. It returns papers 
+    that are most relevant to the topic you're searching for. You can optionally 
+    specify a maximum number of papers to return.
+    """
+    return await pt.search_papers(ctx(), query, max_papers)
+
+
+@mcp.tool()
+async def query_papers(query: str) -> Any:
+    """
+    Query the papers to answer a specific question using PaperQA.
+    
+    Args:
+        query: The question to answer based on the papers
+        
+    Returns:
+        Detailed answer with citations from the papers
+        
+    This tool analyzes the papers in your collection to provide an evidence-based 
+    answer to your question. It extracts relevant information from across papers 
+    and synthesizes a response with citations to the source papers.
+    """
+    return await pt.query_papers(ctx(), query)
+
+
+@mcp.tool()
+async def add_paper(path: str, citation: Optional[str] = None) -> Any:
+    """
+    Add a specific paper to the collection.
+    
+    Args:
+        path: Path to the paper file or URL
+        citation: Optional citation for the paper
+        
+    Returns:
+        Information about the added paper
+        
+    You can add a paper by providing its file path (PDF) or a URL to the paper 
+    (must be accessible). The paper will be added to your collection for searching 
+    and querying. You can optionally provide a citation string.
+    """
+    return await pt.add_paper(ctx(), path, citation)
+
+@mcp.tool()
+async def add_papers(path: str,) -> Any:
+    """
+    Add multiple papers to the collection.
+    Args:
+        path: Path to the paper file or URL
+        citation: Optional citation for the paper
+
+    Returns:
+        Informations about the added papers
+
+    You can add multiple papers by providing its file path (PDF) or a URL to the
+    paper (must be accessible). The paper will be added to your collection for
+    searching and querying.
+    """
+    return await pt.add_papers(ctx(), path)
+
+
+@mcp.tool()
+async def list_papers() -> Any:
+    """
+    List all papers in the current paper directory.
+    
+    Args:
+        None
+        
+    Returns:
+        Information about all papers in the paper directory
+        
+    This lists all papers currently in your collection, showing their file paths and 
+    any other available metadata. Use this to see what papers you have available 
+    for searching and querying.
+    """
+    return await pt.list_papers(ctx())
+
+
+@mcp.tool()
+async def build_index() -> Any:
+    """
+    Rebuild the search index for papers.
+    
+    Args:
+        None
+        
+    Returns:
+        Information about the indexing process
+        
+    This rebuilds the search index for all papers in your paper directory.
+    The index is required for searching and querying papers. You should run this
+    after adding new papers to make them searchable.
+    """
+    return await pt.build_index(ctx())
+
+
+if __name__ == "__main__":
+    print("Running the PaperQA MCP server")
+    print("Use Ctrl-C to exit")
+    mcp.run(transport='stdio')