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

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,65 @@
+"""
+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,
+)
+
+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)}")

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"],
+        ],
+    )