Part 1: Getting Started with Document Q&A Agents

from gaia.agents.base.agent import Agent
from gaia.agents.chat.tools import RAGToolsMixin, FileToolsMixin

class MyAgent(Agent, RAGToolsMixin, FileToolsMixin):
    def _register_tools(self):
        self.register_rag_tools()    # Gets all RAG tools
        self.register_file_tools()   # Gets all file tools
        # Agent now has ~15 tools available

from gaia.agents.base.agent import Agent
from gaia.agents.base.console import AgentConsole
from gaia.agents.base.tools import tool
from gaia.rag.sdk import RAGSDK, RAGConfig
from fnmatch import fnmatch
from pathlib import Path

class SmartDocAgent(Agent):
    """Agent with smart document discovery."""

    def __init__(self, **kwargs):
        self.rag = RAGSDK(RAGConfig())
        self.indexed_files = set()
        super().__init__(**kwargs)

    def _get_system_prompt(self) -> str:
        indexed_docs = "\n".join(f"- {Path(f).name}" for f in self.indexed_files)

        return f"""You are an intelligent document assistant.

Indexed documents:
{indexed_docs or "None yet"}

**Smart Discovery Workflow:**
When user asks about something (e.g., "oil & gas manual"):
1. Use search_files to find it
2. Index it automatically
3. Then query to answer their question

This creates a more natural user experience where file paths don't need to be specified."""

    def _create_console(self):
        return AgentConsole()

    def _register_tools(self):
        @tool
        def search_files(pattern: str) -> dict:
            """Find files matching a pattern (searches common locations)."""
            # Search common directories
            search_paths = [
                Path.home() / "Documents",
                Path.home() / "Downloads",
                Path.home() / "Desktop",
                Path.cwd(),
            ]

            # Support "*" wildcards (agent sends patterns with them) and fall back to
            # substring matching by wrapping non-wildcard patterns with "*".
            normalized_pattern = pattern.lower()
            if "*" not in normalized_pattern and "?" not in normalized_pattern:
                normalized_pattern = f"*{normalized_pattern}*"

            found_files = []
            for search_path in search_paths:
                if search_path.exists():
                    # Find PDF files matching pattern
                    for pdf in search_path.rglob("*.pdf"):
                        if fnmatch(pdf.name.lower(), normalized_pattern):
                            found_files.append(str(pdf))

            return {
                "files": found_files,
                "count": len(found_files),
                "message": f"Found {len(found_files)} file(s)"
            }

        @tool
        def index_document(file_path: str) -> dict:
            """Index a document for searching."""
            if not Path(file_path).exists():
                return {"error": f"File not found: {file_path}"}

            result = self.rag.index_document(file_path)
            if result.get("success"):
                self.indexed_files.add(file_path)
                return {
                    "status": "success",
                    "chunks": result.get("num_chunks"),
                    "file": Path(file_path).name
                }
            return {"error": result.get("error")}

        @tool
        def query_documents(query: str) -> dict:
            """Search all indexed documents."""
            if not self.indexed_files:
                return {"error": "No documents indexed"}

            response = self.rag.query(query)
            return {
                "answer": response.text,
                "sources": [Path(f).name for f in response.source_files],
                "scores": response.chunk_scores
            }

# Use it
agent = SmartDocAgent()

# User can now ask naturally!
result = agent.process_query("Find and search the user manual for installation steps")
# Agent will:
# 1. Call search_files("user manual")
# 2. Call index_document(found_file)
# 3. Call query_documents("installation steps")
# 4. Return the answer!