aurelian 0.1.0__py3-none-any.whl

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')