RAG SDK

@dataclass
class RAGConfig:
    """Configuration for RAG SDK."""

    model: str = "Qwen3.5-35B-A3B-GGUF"
    max_tokens: int = 1024
    chunk_size: int = 500  # max useful; chunks are capped at EMBED_MAX_CHARS (2000)
    chunk_overlap: int = 100
    max_chunks: int = 5
    embedding_model: str = "user.embeddinggemma-300m-GGUF"
    cache_dir: str = field(default_factory=default_rag_cache_dir)  # ~/.gaia/cache/rag
    show_stats: bool = False
    use_local_llm: bool = True
    base_url: Optional[str] = None  # LEMONADE_BASE_URL, else GAIA's own Lemonade

    # Memory management
    max_indexed_files: int = 100
    max_total_chunks: int = 10000
    enable_lru_eviction: bool = True

    # File size limits
    max_file_size_mb: int = 100
    warn_file_size_mb: int = 50

    # Chunking strategy
    use_llm_chunking: bool = False

    # VLM settings
    vlm_model: str = "Qwen3-VL-4B-Instruct-GGUF"
    use_vlm: bool = True  # False: never read PDF/PPTX images, even if the VLM is installed

    # Security
    allowed_paths: Optional[List[str]] = None

@dataclass
class RAGResponse:
    """Response from RAG operations with enhanced metadata."""

    text: str
    chunks: Optional[List[str]] = None
    chunk_scores: Optional[List[float]] = None
    stats: Optional[Dict[str, Any]] = None
    source_files: Optional[List[str]] = None
    chunk_metadata: Optional[List[Dict[str, Any]]] = None
    query_metadata: Optional[Dict[str, Any]] = None

class RAGSDK:
    """
    Simple RAG SDK for PDF document Q&A following GAIA patterns.

    Supports:
    - Documents: PDF, TXT, MD, CSV, JSON
    - Backend Code: Python, Java, C/C++, Go, Rust, Ruby, PHP, Swift, Kotlin, Scala
    - Web Code: JavaScript/TypeScript, HTML, CSS/SCSS/SASS/LESS, Vue, Svelte, Astro
    - Config: YAML, XML, TOML, INI, ENV, Properties
    - Build: Gradle, CMake, Makefiles
    - Database: SQL
    """

    def __init__(self, config: Optional[RAGConfig] = None):
        """Initialize RAG SDK."""
        ...

    def index_document(self, file_path: str) -> Dict[str, Any]:
        """
        Index a document for retrieval.

        Args:
            file_path: Path to document or code file

        Returns:
            Dict with indexing results and statistics:
            {
                "success": bool,
                "file_name": str,
                "file_type": str,
                "file_size_mb": float,
                "num_pages": int (for PDFs),
                "num_chunks": int,
                "total_indexed_files": int,
                "total_chunks": int,
                "vlm_pages": int (pages enhanced with VLM),
                "total_images": int (images processed),
                "error": str (if failed),
                "error_type": str (some failures, see below)
            }

        error_type, when set, names a failure a caller can act on:
            "file_locked": Windows refused the read; another program holds
                the document exclusively or has locked part of it (or the user
                lacks read permission, which Windows reports the same way).
                Any format. Index it again once the program lets go.
            "cache_busy": another GAIA process holds this document's cache.
            "embedding_unavailable": the embedder or Lemonade could not run.
                Two booleans come with it: "lemonade_unreachable" (Lemonade
                Server could not be reached at all) and "embedder_busy" (the
                server timed out, busy with another model or app).
        """
        ...

    def remove_document(self, file_path: str) -> bool:
        """Remove a document from the index."""
        ...

    def reindex_document(self, file_path: str) -> Dict[str, Any]:
        """Reindex a document (remove old chunks and add new ones)."""
        ...

    def query(self, question: str, include_metadata: bool = True) -> RAGResponse:
        """
        Query the indexed documents with enhanced metadata tracking.

        Args:
            question: Question to ask about the documents
            include_metadata: Whether to include detailed metadata in response

        Returns:
            RAGResponse with answer, retrieved chunks, and metadata
        """
        ...

    def clear_cache(self):
        """Delete only files the cache wrote (hash-named *.json / *.json.sig
        and *_extracted.md). Raises ValueError if cache_dir is $HOME or ~/.gaia."""
        ...

    def get_status(self) -> Dict[str, Any]:
        """Get RAG system status."""
        ...

    # Private methods for internal use

    def _extract_text_from_pdf(self, pdf_path: str) -> tuple:
        """
        Extract text from PDF file with VLM for images.

        Returns:
            (text, num_pages, metadata) tuple
        """
        ...

    def _split_text_into_chunks(self, text: str) -> List[str]:
        """
        Split text into semantic chunks.

        Uses intelligent splitting that:
        - Respects natural document boundaries
        - Keeps semantic units together
        - Maintains context with overlap
        - Protects VLM content blocks
        """
        ...

    def _create_vector_index(self, chunks: List[str]) -> tuple:
        """Create FAISS vector index from chunks."""
        ...

    def _retrieve_chunks(self, query: str) -> tuple:
        """Retrieve relevant chunks for query."""
        ...

    def _retrieve_chunks_from_file(self, query: str, file_path: str) -> tuple:
        """
        Retrieve relevant chunks from a specific file using cached per-file index.

        Much faster than global search because:
        1. Uses pre-computed embeddings
        2. Searches smaller, file-specific FAISS index
        3. No need to rebuild index on each query
        """
        ...

def quick_rag(pdf_path: str, question: str, **kwargs) -> str:
    """
    Convenience function for quick RAG query.

    Args:
        pdf_path: Path to PDF file
        question: Question to ask
        **kwargs: Additional config parameters

    Returns:
        Answer text
    """
    ...

from gaia.rag.sdk import RAGSDK, RAGConfig

# Initialize
config = RAGConfig(show_stats=True)
rag = RAGSDK(config)

# Index document
result = rag.index_document("manual.pdf")
print(f"Indexed {result['num_chunks']} chunks from {result['file_name']}")

# Query
response = rag.query("What are the key features?")
print(response.text)

# View sources
if response.chunk_metadata:
    for meta in response.chunk_metadata:
        print(f"Source: {meta['source_file']} (score: {meta['relevance_score']:.2f})")

# Index multiple documents
for doc in ["manual.pdf", "guide.pdf", "faq.pdf"]:
    rag.index_document(doc)

# Query across all documents
response = rag.query("How do I troubleshoot errors?", include_metadata=True)

# Show which documents were used
print(f"Retrieved from {len(set(response.source_files))} documents")
print(f"Total chunks: {response.query_metadata['num_chunks_retrieved']}")

# Index code files
rag.index_document("src/main.py")
rag.index_document("src/utils.py")
rag.index_document("README.md")

# Ask coding questions
response = rag.query("How does the authentication system work?")
print(response.text)

# View code chunks that were referenced
for i, chunk in enumerate(response.chunks, 1):
    print(f"\n--- Context {i} ---")
    print(chunk[:200] + "...")

def test_rag_initialization():
    """Test RAG SDK initialization."""
    config = RAGConfig(cache_dir=".test_cache")
    rag = RAGSDK(config)
    assert rag is not None

def test_document_indexing():
    """Test document indexing."""
    rag = RAGSDK()
    result = rag.index_document("test.pdf")

    assert result["success"]
    assert result["num_chunks"] > 0
    assert result["file_name"] == "test.pdf"

def test_query_functionality():
    """Test query with indexed documents."""
    rag = RAGSDK()
    rag.index_document("test.pdf")

    response = rag.query("What is the main topic?")

    assert response.text
    assert response.chunks
    assert len(response.chunks) <= rag.config.max_chunks

def test_remove_document():
    """Test document removal."""
    rag = RAGSDK()
    rag.index_document("test.pdf")

    success = rag.remove_document("test.pdf")
    assert success
    assert len(rag.indexed_files) == 0

[project]
dependencies = [
    "pypdf>=3.0",
    "faiss-cpu",  # Vector search
    "numpy",
    "gaia.chat.sdk",
    "gaia.llm.vlm_client",  # For VLM image extraction
]