document-rag-mcp 0.1.0__py3-none-any.whl

document_rag_mcp/__init__.py

@@ -0,0 +1,2 @@
+# document-rag-mcp package
+__version__ = "0.1.0"

document_rag_mcp/cli.py

@@ -0,0 +1,136 @@
+import asyncio
+import os
+from pathlib import Path
+import click
+
+
+@click.group()
+@click.option(
+    "--config",
+    "-c",
+    type=click.Path(exists=True, path_type=Path),
+    envvar="DOCRAG_CONFIG",
+    help="Path to the YAML configuration file.",
+)
+@click.option(
+    "--chunking-model",
+    envvar="DOCRAG_CHUNKING_MODEL",
+    help="Override the local model used for semantic chunking boundary detection.",
+)
+@click.pass_context
+def main(ctx: click.Context, config: Path | None, chunking_model: str | None) -> None:
+    """RAG MCP Server CLI — Manage document indexing and semantic search server."""
+    if config:
+        os.environ["DOCRAG_CONFIG"] = str(config.resolve())
+    if chunking_model:
+        os.environ["DOCRAG_CHUNKING_MODEL"] = chunking_model
+
+
+@main.command()
+@click.option(
+    "--transport",
+    type=click.Choice(["stdio", "http"]),
+    default="stdio",
+    help="Transport protocol to run the server on (stdio or http/sse).",
+)
+@click.option(
+    "--host",
+    default="127.0.0.1",
+    help="Host to bind the HTTP server to (http transport only).",
+)
+@click.option(
+    "--port",
+    default=8000,
+    type=int,
+    help="Port to run the HTTP server on (http transport only).",
+)
+def serve(transport: str, host: str, port: int) -> None:
+    """Start the Model Context Protocol (MCP) server."""
+    # Lazy import to ensure CLI overrides are set in environment first
+    import document_rag_mcp.server as mcp_server
+
+    transport_mcp = "sse" if transport == "http" else "stdio"
+    click.echo(f"Starting document-rag-mcp server on {transport} transport...")
+    if transport == "http":
+        click.echo(f"Binding to http://{host}:{port}")
+
+    mcp_server.mcp.run(transport=transport_mcp, host=host, port=port)
+
+
+@main.command()
+@click.argument("query")
+@click.option(
+    "--collection",
+    "-c",
+    default=None,
+    help="Filter search results by a specific collection.",
+)
+@click.option(
+    "--top-k",
+    "-k",
+    default=5,
+    type=int,
+    help="Number of results to return.",
+)
+def search(query: str, collection: str | None, top_k: int) -> None:
+    """Run a semantic search against the indexed collections."""
+    import document_rag_mcp.server as mcp_server
+
+    async def run_search():
+        click.echo(f"Searching for: '{query}'...")
+        results = await mcp_server.search_engine.search(
+            query=query, collection_name=collection, top_k=top_k
+        )
+        if not results:
+            click.echo("No matching documents found.")
+            return
+
+        for i, r in enumerate(results):
+            click.echo(f"\n[{i+1}] Score: {r.score:.4f} | File: {r.metadata.file_path}")
+            click.echo(f"Collection: {r.metadata.collection} | Page: {r.metadata.page_number or 1}")
+            click.echo(f"Title: {r.metadata.title or 'N/A'} | Section: {r.metadata.section or 'N/A'}")
+            click.echo("-" * 40)
+            click.echo(r.text.strip())
+            click.echo("=" * 60)
+
+    asyncio.run(run_search())
+
+
+@main.command()
+@click.option(
+    "--collection",
+    "-c",
+    default=None,
+    help="Limit ingestion to a specific collection.",
+)
+def ingest(collection: str | None) -> None:
+    """Recursively scan folders and index new/changed files immediately."""
+    import document_rag_mcp.server as mcp_server
+
+    async def run_ingest():
+        click.echo("Starting one-shot ingestion scan...")
+        res = await mcp_server.ingest_now(collection=collection)
+        click.echo(res)
+
+    asyncio.run(run_ingest())
+
+
+@main.command()
+def collections() -> None:
+    """List all configured collections and their indexing status."""
+    import document_rag_mcp.server as mcp_server
+
+    click.echo("Configured Collections:")
+    click.echo("=" * 60)
+    for coll in mcp_server.config.collections:
+        stats = mcp_server.vector_store.collection_stats(coll.name)
+        paths_str = ", ".join(str(p) for p in coll.paths)
+        click.echo(f"Name:          {coll.name}")
+        click.echo(f"Paths:         {paths_str}")
+        click.echo(f"Patterns:      {', '.join(coll.file_patterns)}")
+        click.echo(f"Indexed Chunks: {stats['count']}")
+        click.echo("-" * 60)
+
+
+if __name__ == "__main__":
+    main()

document_rag_mcp/config.py

@@ -0,0 +1,94 @@
+import os
+from pathlib import Path
+from typing import Any
+import yaml
+from pydantic import BaseModel, Field
+from pydantic_settings import (
+    BaseSettings,
+    SettingsConfigDict,
+    PydanticBaseSettingsSource,
+)
+
+class CollectionConfig(BaseModel):
+    name: str                           # e.g., "project-docs"
+    paths: list[Path]                   # folders to watch/scan
+    file_patterns: list[str] = ["*.txt", "*.md", "*.pdf"]
+
+
+class EmbeddingConfig(BaseModel):
+    base_url: str = "http://localhost:8080/v1"   # lemonade default
+    api_key: str = "unused"
+    model: str = "embed-gemma-300m-FLM"
+    dimensions: int = 768               # gemma embedding dims
+    batch_size: int = 32
+
+
+class VisionConfig(BaseModel):
+    enabled: bool = False
+    base_url: str = "http://localhost:8080/v1"
+    api_key: str = "unused"
+    model: str = "gpt-4o"
+
+
+class ChunkingConfig(BaseModel):
+    max_chunk_size: int = 512          # tokens
+    similarity_threshold: float = 0.5  # for semantic boundary detection
+    local_model: str = "all-MiniLM-L6-v2"  # overridable via --chunking-model CLI flag
+
+
+class StorageConfig(BaseModel):
+    data_dir: Path = Path("./data")    # ChromaDB + SQLite storage
+
+
+class ServerConfig(BaseModel):
+    host: str = "127.0.0.1"
+    port: int = 8000
+
+
+class AppConfig(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_prefix="DOCRAG_",
+        env_nested_delimiter="__",
+        extra="ignore",
+    )
+
+    collections: list[CollectionConfig] = Field(default_factory=list)
+    embedding: EmbeddingConfig = EmbeddingConfig()
+    vision: VisionConfig = VisionConfig()
+    chunking: ChunkingConfig = ChunkingConfig()
+    storage: StorageConfig = StorageConfig()
+    server: ServerConfig = ServerConfig()
+
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        # Prioritize environment variables over initialization arguments (YAML data)
+        return env_settings, init_settings, dotenv_settings, file_secret_settings
+
+
+def load_config(config_path: Path | str | None = None) -> AppConfig:
+    """Loads application configuration from YAML file and overrides with environment variables.
+
+    Environment variables must be prefixed with DOCRAG_ and nested with double underscores.
+    For example: DOCRAG_EMBEDDING__MODEL="text-embedding-3-small"
+    """
+    if not config_path:
+        config_path = os.environ.get("DOCRAG_CONFIG")
+
+    yaml_data: dict[str, Any] = {}
+    if config_path:
+        path = Path(config_path)
+        if path.exists():
+            with open(path, "r", encoding="utf-8") as f:
+                loaded = yaml.safe_load(f)
+                if isinstance(loaded, dict):
+                    yaml_data = loaded
+
+    # Pydantic BaseSettings automatically prioritizes env vars over values passed in init.
+    return AppConfig(**yaml_data)

document_rag_mcp/embedding/__init__.py

@@ -0,0 +1 @@
+# Embedding module

document_rag_mcp/embedding/client.py

@@ -0,0 +1,47 @@
+from openai import AsyncOpenAI
+from ..config import EmbeddingConfig
+
+
+class EmbeddingClient:
+    def __init__(self, config: EmbeddingConfig):
+        self.client = AsyncOpenAI(base_url=config.base_url, api_key=config.api_key)
+        self.model = config.model
+        self.dimensions = config.dimensions
+        self.batch_size = config.batch_size
+
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        """Generates embeddings for a list of texts in batches."""
+        if not texts:
+            return []
+
+        embeddings: list[list[float]] = []
+
+        # Process in batches
+        for i in range(0, len(texts), self.batch_size):
+            batch = texts[i : i + self.batch_size]
+            try:
+                # Try sending with dimensions parameter
+                response = await self.client.embeddings.create(
+                    input=batch,
+                    model=self.model,
+                    dimensions=self.dimensions,
+                )
+            except Exception:
+                # Fallback for APIs that don't support the dimensions parameter
+                response = await self.client.embeddings.create(
+                    input=batch,
+                    model=self.model,
+                )
+
+            # Extract embeddings in correct order
+            batch_embeddings = [data.embedding for data in response.data]
+            embeddings.extend(batch_embeddings)
+
+        return embeddings
+
+    async def embed_query(self, query: str) -> list[float]:
+        """Generates embedding for a single search query."""
+        results = await self.embed([query])
+        if not results:
+            raise ValueError("Failed to generate embedding for query.")
+        return results[0]

document_rag_mcp/ingestion/__init__.py

@@ -0,0 +1 @@
+# Ingestion module

document_rag_mcp/ingestion/chunker.py

@@ -0,0 +1,125 @@
+from datetime import datetime, timezone
+import hashlib
+import os
+from ..config import ChunkingConfig
+from ..models import ChunkMetadata, DocumentChunk
+from .extractor import ExtractedPage
+
+
+class DocumentChunker:
+    def __init__(self, config: ChunkingConfig):
+        self.config = config
+        self._init_chunkers()
+
+    def _init_chunkers(self) -> None:
+        import chonkie
+
+        # 1. Initialize semantic chunker (for TXT / PDF)
+        try:
+            self.semantic_chunker = chonkie.SemanticChunker(
+                embedding_model=self.config.local_model,
+                chunk_size=self.config.max_chunk_size,
+                threshold=self.config.similarity_threshold,
+            )
+        except Exception as e:
+            # Fallback to SentenceChunker/TokenChunker if sentence-transformers is missing or fails
+            print(
+                f"Warning: Failed to load SemanticChunker with '{self.config.local_model}': {e}"
+            )
+            print("Falling back to TokenChunker for semantic text.")
+            self.semantic_chunker = chonkie.TokenChunker(chunk_size=self.config.max_chunk_size)
+
+        # 2. Initialize recursive chunker for Markdown (uses split rules rather than semantic embeddings)
+        self.markdown_chunker = chonkie.RecursiveChunker(
+            chunk_size=self.config.max_chunk_size,
+        )
+
+    def chunk_pages(
+        self,
+        pages: list[ExtractedPage],
+        file_path: str,
+        collection_name: str,
+        file_hash: str,
+        file_type: str,
+        last_modified: datetime,
+    ) -> list[DocumentChunk]:
+        """Transforms a list of ExtractedPages into a flat list of metadata-rich DocumentChunks."""
+        file_name = os.path.basename(file_path)
+
+        # Resolve document title once for the entire document
+        document_title = None
+        if pages and pages[0].metadata:
+            document_title = pages[0].metadata.get("title")
+        if not document_title and pages and pages[0].headings:
+            for heading_text, level in pages[0].headings:
+                if level == 1:
+                    document_title = heading_text
+                    break
+        if not document_title:
+            document_title = os.path.splitext(file_name)[0]
+
+        raw_chunks: list[tuple[str, int | None, list[tuple[str, int]]]] = []
+
+        # Generate chunks per page
+        for page in pages:
+            text_to_chunk = page.text.strip()
+            if not text_to_chunk:
+                continue
+
+            if file_type == "md":
+                chunks_out = self.markdown_chunker.chunk(text_to_chunk)
+            else:
+                chunks_out = self.semantic_chunker.chunk(text_to_chunk)
+
+            for c in chunks_out:
+                raw_chunks.append((c.text, page.page_number, page.headings))
+
+        total_chunks = len(raw_chunks)
+        document_chunks: list[DocumentChunk] = []
+
+        for idx, (chunk_text, page_num, headings) in enumerate(raw_chunks):
+            # Resolve section / heading
+            section = None
+            if headings:
+                # Check headings in reverse order (closest heading before/within the chunk)
+                for heading_text, _ in reversed(headings):
+                    if heading_text in chunk_text:
+                        section = heading_text
+                        break
+                if not section:
+                    section = headings[0][0]
+
+            # Compute chunk-level hash
+            chunk_hash = hashlib.sha256(chunk_text.encode("utf-8")).hexdigest()
+
+            # Deterministic chunk ID
+            chunk_id_seed = f"{file_path}_{idx}"
+            chunk_id = hashlib.sha256(chunk_id_seed.encode("utf-8")).hexdigest()
+
+            metadata = ChunkMetadata(
+                file_path=file_path,
+                file_name=file_name,
+                collection=collection_name,
+                file_hash=file_hash,
+                chunk_hash=chunk_hash,
+                chunk_index=idx,
+                total_chunks=total_chunks,
+                title=document_title,
+                section=section,
+                file_type=file_type,
+                page_number=page_num,
+                last_modified=last_modified,
+                ingested_at=datetime.now(timezone.utc),
+                vision_processed=False,
+            )
+
+            document_chunks.append(
+                DocumentChunk(
+                    id=chunk_id,
+                    text=chunk_text,
+                    metadata=metadata,
+                    embedding=None,
+                )
+            )
+
+        return document_chunks

document_rag_mcp/ingestion/extractor.py

@@ -0,0 +1,181 @@
+from collections import Counter
+from dataclasses import dataclass, field
+from pathlib import Path
+import fitz  # PyMuPDF
+import yaml
+
+
+@dataclass
+class ExtractedPage:
+    text: str
+    page_number: int | None = None
+    headings: list[tuple[str, int]] = field(default_factory=list)  # list of (heading_text, heading_level)
+    image_bytes: bytes | None = None
+    metadata: dict[str, any] = field(default_factory=dict)
+
+
+class DocumentExtractor:
+    def __init__(self, vision_enabled: bool = False):
+        self.vision_enabled = vision_enabled
+
+    def extract(self, file_path: Path | str) -> list[ExtractedPage]:
+        """Extracts content from a file depending on its extension."""
+        path = Path(file_path)
+        ext = path.suffix.lower()
+
+        if ext == ".txt":
+            return self._extract_txt(path)
+        elif ext == ".md":
+            return self._extract_md(path)
+        elif ext == ".pdf":
+            return self._extract_pdf(path)
+        else:
+            raise ValueError(f"Unsupported file extension: {ext}")
+
+    def _extract_txt(self, path: Path) -> list[ExtractedPage]:
+        with open(path, "r", encoding="utf-8", errors="ignore") as f:
+            content = f.read()
+
+        headings = []
+        lines = content.splitlines()
+        if lines:
+            first_line = lines[0].strip()
+            # If first line looks like a title, mark it as heading level 1
+            if first_line and len(first_line) < 80 and not first_line.endswith((".", ",", ";")):
+                headings.append((first_line, 1))
+
+        return [
+            ExtractedPage(
+                text=content,
+                page_number=1,
+                headings=headings,
+                metadata={"title": path.stem},
+            )
+        ]
+
+    def _extract_md(self, path: Path) -> list[ExtractedPage]:
+        with open(path, "r", encoding="utf-8", errors="ignore") as f:
+            content = f.read()
+
+        frontmatter = {}
+        body = content
+
+        # Parse frontmatter
+        if content.startswith("---"):
+            parts = content.split("---", 2)
+            if len(parts) >= 3:
+                try:
+                    frontmatter = yaml.safe_load(parts[1]) or {}
+                    body = parts[2]
+                except Exception:
+                    pass
+
+        headings = []
+        for line in body.splitlines():
+            line_strip = line.strip()
+            if line_strip.startswith("#"):
+                # Count the heading level
+                level = 0
+                while level < len(line_strip) and line_strip[level] == "#":
+                    level += 1
+                if level > 0 and level < len(line_strip) and line_strip[level] == " ":
+                    heading_text = line_strip[level:].strip()
+                    if heading_text:
+                        headings.append((heading_text, level))
+
+        # Default title is the frontmatter title, or file stem
+        title = frontmatter.get("title", path.stem)
+        meta = {"title": title, "frontmatter": frontmatter}
+
+        return [
+            ExtractedPage(
+                text=body,
+                page_number=1,
+                headings=headings,
+                metadata=meta,
+            )
+        ]
+
+    def _extract_pdf(self, path: Path) -> list[ExtractedPage]:
+        doc = fitz.open(path)
+        pages = []
+
+        for page_idx, page in enumerate(doc):
+            page_num = page_idx + 1
+            text = page.get_text("text")
+
+            # Scanned check (very little text)
+            non_space_chars = len("".join(text.split()))
+            is_empty_or_scanned = non_space_chars < 40
+
+            image_bytes = None
+            if is_empty_or_scanned and self.vision_enabled:
+                # Render to high-quality PNG (144 DPI)
+                pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
+                image_bytes = pix.tobytes("png")
+
+            # Typography-aware heading detection using get_text("dict")
+            headings = []
+            try:
+                blocks = page.get_text("dict").get("blocks", [])
+                sizes = []
+                for b in blocks:
+                    if b.get("type") == 0:  # text block
+                        for line in b.get("lines", []):
+                            for span in line.get("spans", []):
+                                txt = span.get("text", "").strip()
+                                if txt:
+                                    sizes.append(span.get("size", 10.0))
+
+                # Compute body text font size (most common size)
+                body_size = 10.0
+                if sizes:
+                    body_size = Counter(sizes).most_common(1)[0][0]
+
+                # Detect headings
+                for b in blocks:
+                    if b.get("type") == 0:
+                        for line in b.get("lines", []):
+                            spans = line.get("spans", [])
+                            if not spans:
+                                continue
+                            line_text = "".join(s.get("text", "") for s in spans).strip()
+                            if not line_text:
+                                continue
+
+                            first_span = spans[0]
+                            size = first_span.get("size", 10.0)
+                            font = first_span.get("font", "").lower()
+
+                            is_bold = "bold" in font or "black" in font or "heavy" in font
+                            is_large = size > body_size * 1.2
+
+                            # Short line, doesn't end with typical sentence punctuation
+                            if (
+                                (is_large or is_bold)
+                                and len(line_text) < 120
+                                and not line_text.endswith((".", ":", ";", ","))
+                            ):
+                                if size > body_size * 1.5:
+                                    level = 1
+                                elif size > body_size * 1.3:
+                                    level = 2
+                                else:
+                                    level = 3
+                                headings.append((line_text, level))
+            except Exception:
+                # Fallback: if dict parsing fails, don't break extraction
+                pass
+
+            pages.append(
+                ExtractedPage(
+                    text=text,
+                    page_number=page_num,
+                    headings=headings,
+                    image_bytes=image_bytes,
+                    metadata={"title": path.stem},
+                )
+            )
+
+        doc.close()
+        return pages