athenaeum-kb 0.1.0__py3-none-any.whl

athenaeum/__init__.py

@@ -0,0 +1,28 @@
+"""Athenaeum - Tools for intelligent interaction with knowledge bases."""
+
+from athenaeum.athenaeum import Athenaeum
+from athenaeum.config import AthenaeumConfig
+from athenaeum.models import (
+    ChunkMetadata,
+    ContentSearchHit,
+    Document,
+    Excerpt,
+    Metadata,
+    SearchHit,
+    TOCEntry,
+)
+from athenaeum.ocr import OCRProvider, get_ocr_provider
+
+__all__ = [
+    "Athenaeum",
+    "AthenaeumConfig",
+    "ChunkMetadata",
+    "ContentSearchHit",
+    "Document",
+    "Excerpt",
+    "Metadata",
+    "OCRProvider",
+    "SearchHit",
+    "TOCEntry",
+    "get_ocr_provider",
+]

athenaeum/athenaeum.py

@@ -0,0 +1,286 @@
+"""Main orchestrator class for Athenaeum."""
+
+from __future__ import annotations
+
+import shutil
+import uuid
+from pathlib import Path
+from typing import Literal
+
+from langchain_core.embeddings import Embeddings
+
+from athenaeum.chunker import chunk_markdown
+from athenaeum.config import AthenaeumConfig
+from athenaeum.document_store import DocumentStore
+from athenaeum.models import ChunkMetadata, ContentSearchHit, Document, Excerpt, SearchHit
+from athenaeum.ocr import OCRProvider, get_ocr_provider
+from athenaeum.search.bm25 import BM25Index
+from athenaeum.search.hybrid import reciprocal_rank_fusion
+from athenaeum.search.vector import VectorIndex
+from athenaeum.storage import StorageManager
+from athenaeum.toc import extract_toc
+
+
+class Athenaeum:
+    """Main entry point for the Athenaeum knowledge base."""
+
+    def __init__(
+        self,
+        embeddings: Embeddings,
+        config: AthenaeumConfig | None = None,
+        ocr_provider: OCRProvider | None = None,
+    ) -> None:
+        self._config = config or AthenaeumConfig()
+        self._storage = StorageManager(self._config.storage_dir)
+        self._doc_store = DocumentStore(self._storage)
+        self._ocr = ocr_provider or get_ocr_provider("markitdown")
+        self._bm25 = BM25Index()
+        self._vector = VectorIndex(
+            embeddings=embeddings,
+            persist_directory=self._storage.ensure_chroma_dir(),
+            collection_name="athenaeum",
+        )
+        self._reindex_bm25()
+
+    def _reindex_bm25(self) -> None:
+        """Rebuild BM25 index from all stored documents."""
+        for doc in self._doc_store.list_all():
+            md_path = Path(doc.path_to_md)
+            if md_path.exists():
+                text = md_path.read_text()
+                chunks = chunk_markdown(
+                    text, doc.id, self._config.chunk_size, self._config.chunk_overlap
+                )
+                self._bm25.add_chunks(chunks)
+
+    def load_doc(self, path: str) -> str:
+        """Load a document into the knowledge base.
+
+        Args:
+            path: Path to the document file.
+
+        Returns:
+            The document ID.
+        """
+        file_path = Path(path).resolve()
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+
+        ext = file_path.suffix.lower()
+        supported = self._ocr.supported_extensions()
+        if ".*" not in supported and ext not in supported:
+            raise ValueError(
+                f"Unsupported file type: {ext}. Supported: {sorted(supported)}"
+            )
+
+        doc_id = uuid.uuid4().hex[:12]
+
+        # Copy raw file
+        raw_dest = self._storage.raw_path(doc_id, ext)
+        shutil.copy2(file_path, raw_dest)
+
+        # Convert to markdown
+        markdown = self._ocr.convert(file_path)
+        md_dest = self._storage.content_md_path(doc_id)
+        md_dest.write_text(markdown)
+
+        # Extract metadata
+        toc = extract_toc(markdown)
+        lines = markdown.split("\n")
+
+        doc = Document(
+            id=doc_id,
+            name=file_path.name,
+            path_to_raw=str(raw_dest),
+            path_to_md=str(md_dest),
+            num_lines=len(lines),
+            table_of_contents=toc,
+            file_size=file_path.stat().st_size,
+            file_type=ext,
+        )
+        self._doc_store.add(doc)
+
+        # Index
+        chunks = chunk_markdown(
+            markdown, doc_id, self._config.chunk_size, self._config.chunk_overlap
+        )
+        self._bm25.add_chunks(chunks)
+        self._vector.add_chunks(chunks)
+
+        return doc_id
+
+    def list_docs(self) -> list[SearchHit]:
+        """List all documents in the knowledge base."""
+        results = []
+        for doc in self._doc_store.list_all():
+            results.append(
+                SearchHit(
+                    id=doc.id,
+                    name=doc.name,
+                    num_lines=doc.num_lines,
+                    table_of_contents=doc.format_toc(),
+                )
+            )
+        return results
+
+    def search_docs(
+        self,
+        query: str,
+        top_k: int = 10,
+        scope: Literal["names", "contents"] = "contents",
+        strategy: Literal["hybrid", "bm25", "vector"] = "hybrid",
+    ) -> list[SearchHit]:
+        """Search across all documents.
+
+        Args:
+            query: Search query text.
+            top_k: Maximum number of results.
+            scope: ``"contents"`` to search within documents, ``"names"`` to search names only.
+            strategy: Search strategy (only for ``scope="contents"``).
+
+        Returns:
+            Ranked list of matching documents.
+        """
+        if scope == "names":
+            return self._search_by_name(query, top_k)
+
+        chunks = self._search_chunks(query, top_k=top_k * 3, strategy=strategy)
+
+        # Aggregate chunks by document
+        doc_scores: dict[str, float] = {}
+        doc_snippets: dict[str, str] = {}
+        for chunk, score in chunks:
+            if chunk.doc_id not in doc_scores or score > doc_scores[chunk.doc_id]:
+                doc_scores[chunk.doc_id] = score
+                doc_snippets[chunk.doc_id] = chunk.text[:200]
+
+        results = []
+        for doc_id, score in sorted(doc_scores.items(), key=lambda x: x[1], reverse=True):
+            doc = self._doc_store.get(doc_id)
+            if doc is None:
+                continue
+            results.append(
+                SearchHit(
+                    id=doc.id,
+                    name=doc.name,
+                    num_lines=doc.num_lines,
+                    table_of_contents=doc.format_toc(),
+                    score=score,
+                    snippet=doc_snippets.get(doc_id, ""),
+                )
+            )
+
+        return results[:top_k]
+
+    def search_doc_contents(
+        self,
+        doc_id: str,
+        query: str,
+        top_k: int = 5,
+        strategy: Literal["hybrid", "bm25", "vector"] = "hybrid",
+    ) -> list[ContentSearchHit]:
+        """Search within a specific document.
+
+        Args:
+            doc_id: Document identifier.
+            query: Search query text.
+            top_k: Maximum number of results.
+            strategy: Search strategy.
+
+        Returns:
+            List of matching content fragments.
+        """
+        doc = self._doc_store.get(doc_id)
+        if doc is None:
+            raise ValueError(f"Document not found: {doc_id}")
+
+        chunks = self._search_chunks(query, top_k=top_k, strategy=strategy, doc_id=doc_id)
+
+        return [
+            ContentSearchHit(
+                doc_id=chunk.doc_id,
+                line_range=(chunk.start_line, chunk.end_line),
+                text=chunk.text,
+                score=score,
+            )
+            for chunk, score in chunks
+        ]
+
+    def read_doc(
+        self,
+        doc_id: str,
+        start_line: int = 1,
+        end_line: int = 100,
+    ) -> Excerpt:
+        """Read a range of lines from a document.
+
+        Args:
+            doc_id: Document identifier.
+            start_line: Starting line number (1-indexed).
+            end_line: Ending line number (1-indexed, inclusive).
+
+        Returns:
+            An ``Excerpt`` with the requested lines.
+        """
+        doc = self._doc_store.get(doc_id)
+        if doc is None:
+            raise ValueError(f"Document not found: {doc_id}")
+
+        md_path = Path(doc.path_to_md)
+        lines = md_path.read_text().split("\n")
+
+        start_idx = max(0, start_line - 1)
+        end_idx = min(len(lines), end_line)
+        selected = lines[start_idx:end_idx]
+
+        return Excerpt(
+            doc_id=doc_id,
+            line_range=(start_idx + 1, end_idx),
+            text="\n".join(selected),
+            total_lines=len(lines),
+        )
+
+    def _search_by_name(self, query: str, top_k: int) -> list[SearchHit]:
+        query_lower = query.lower()
+        scored = []
+        for doc in self._doc_store.list_all():
+            name_lower = doc.name.lower()
+            if query_lower in name_lower:
+                # Simple relevance: exact match > contains
+                score = 1.0 if query_lower == name_lower else 0.5
+                scored.append((doc, score))
+
+        scored.sort(key=lambda x: x[1], reverse=True)
+        return [
+            SearchHit(
+                id=doc.id,
+                name=doc.name,
+                num_lines=doc.num_lines,
+                table_of_contents=doc.format_toc(),
+                score=score,
+            )
+            for doc, score in scored[:top_k]
+        ]
+
+    def _search_chunks(
+        self,
+        query: str,
+        top_k: int,
+        strategy: Literal["hybrid", "bm25", "vector"] = "hybrid",
+        doc_id: str | None = None,
+    ) -> list[tuple[ChunkMetadata, float]]:
+        """Internal: run search using the given strategy."""
+        if strategy == "bm25":
+            return self._bm25.search(query, top_k=top_k, doc_id=doc_id)
+
+        if strategy == "vector":
+            return self._vector.search(query, top_k=top_k, doc_id=doc_id)
+
+        # hybrid
+        bm25_results = self._bm25.search(query, top_k=top_k, doc_id=doc_id)
+        vector_results = self._vector.search(query, top_k=top_k, doc_id=doc_id)
+        return reciprocal_rank_fusion(
+            [bm25_results, vector_results],
+            k=self._config.rrf_k,
+            top_k=top_k,
+        )

athenaeum/chunker.py

@@ -0,0 +1,77 @@
+"""Line-aware markdown chunking with heading-boundary snapping."""
+
+from __future__ import annotations
+
+import re
+
+from athenaeum.models import ChunkMetadata
+
+_HEADING_RE = re.compile(r"^#{1,6}\s+")
+
+
+def _find_heading_lines(lines: list[str]) -> set[int]:
+    """Return 0-indexed line numbers that are markdown headings."""
+    return {i for i, line in enumerate(lines) if _HEADING_RE.match(line.strip())}
+
+
+def chunk_markdown(
+    markdown: str,
+    doc_id: str,
+    chunk_size: int = 80,
+    chunk_overlap: int = 20,
+) -> list[ChunkMetadata]:
+    """Split markdown into overlapping, line-based chunks.
+
+    Chunks are snapped to heading boundaries when a heading falls within the
+    overlap region, so sections start cleanly.
+
+    Args:
+        markdown: Full markdown text.
+        doc_id: Parent document ID.
+        chunk_size: Target number of lines per chunk.
+        chunk_overlap: Number of overlapping lines between consecutive chunks.
+
+    Returns:
+        List of ``ChunkMetadata`` instances.
+    """
+    if not markdown:
+        return []
+
+    lines = markdown.split("\n")
+    total = len(lines)
+
+    heading_lines = _find_heading_lines(lines)
+    chunks: list[ChunkMetadata] = []
+    start = 0
+    chunk_index = 0
+
+    while start < total:
+        end = min(start + chunk_size, total)
+        text = "\n".join(lines[start:end])
+
+        chunks.append(
+            ChunkMetadata(
+                doc_id=doc_id,
+                chunk_index=chunk_index,
+                start_line=start + 1,  # 1-indexed
+                end_line=end,  # 1-indexed inclusive
+                text=text,
+            )
+        )
+        chunk_index += 1
+
+        if end >= total:
+            break
+
+        # Compute next start with overlap
+        next_start = end - chunk_overlap
+
+        # Snap to heading if one exists in the overlap zone [next_start, end)
+        for line_idx in range(next_start, end):
+            if line_idx in heading_lines:
+                next_start = line_idx
+                break
+
+        start = next_start
+
+    return chunks

athenaeum/config.py

@@ -0,0 +1,16 @@
+"""Configuration for Athenaeum."""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal
+
+
+@dataclass
+class AthenaeumConfig:
+    """Configuration for an Athenaeum instance."""
+
+    storage_dir: Path = field(default_factory=lambda: Path.home() / ".athenaeum")
+    chunk_size: int = 80
+    chunk_overlap: int = 20
+    rrf_k: int = 60
+    default_strategy: Literal["hybrid", "bm25", "vector"] = "hybrid"

athenaeum/document_store.py

@@ -0,0 +1,53 @@
+"""JSON-backed document registry."""
+
+from __future__ import annotations
+
+from athenaeum.models import Document
+from athenaeum.storage import StorageManager
+
+
+class DocumentStore:
+    """Manages the document registry backed by ``metadata.json``."""
+
+    def __init__(self, storage: StorageManager) -> None:
+        self._storage = storage
+        self._docs: dict[str, Document] = {}
+        self._load()
+
+    def _load(self) -> None:
+        raw = self._storage.load_metadata()
+        docs_raw = raw.get("documents", {})
+        for doc_id, data in docs_raw.items():
+            self._docs[doc_id] = Document.model_validate(data)
+
+    def _save(self) -> None:
+        data = {
+            "documents": {
+                doc_id: doc.model_dump(mode="json") for doc_id, doc in self._docs.items()
+            }
+        }
+        self._storage.save_metadata(data)
+
+    def add(self, doc: Document) -> None:
+        """Add or update a document in the registry."""
+        self._docs[doc.id] = doc
+        self._save()
+
+    def get(self, doc_id: str) -> Document | None:
+        """Get a document by ID, or None if not found."""
+        return self._docs.get(doc_id)
+
+    def list_all(self) -> list[Document]:
+        """Return all documents."""
+        return list(self._docs.values())
+
+    def remove(self, doc_id: str) -> Document | None:
+        """Remove a document from the registry. Returns the removed doc or None."""
+        doc = self._docs.pop(doc_id, None)
+        if doc is not None:
+            self._save()
+        return doc
+
+    @property
+    def count(self) -> int:
+        return len(self._docs)

athenaeum/models.py

@@ -0,0 +1,88 @@
+"""Pydantic data models for Athenaeum."""
+
+from datetime import UTC, datetime
+
+from pydantic import BaseModel, Field
+
+
+class Metadata(BaseModel):
+    """Basic document metadata."""
+
+    id: str = Field(..., description="Unique document identifier")
+    name: str = Field(..., description="Document display name")
+
+
+class TOCEntry(BaseModel):
+    """Table of contents entry with line range."""
+
+    title: str = Field(..., description="Section title")
+    level: int = Field(..., description="Header level (1 for h1, 2 for h2)")
+    start_line: int = Field(..., description="Starting line number (1-indexed)")
+    end_line: int | None = Field(None, description="Ending line number (1-indexed, inclusive)")
+
+
+class Document(BaseModel):
+    """Full document record stored in the knowledge base."""
+
+    id: str = Field(..., description="Unique document identifier (UUID)")
+    name: str = Field(..., description="Original filename")
+    path_to_raw: str = Field(..., description="Path to original file")
+    path_to_md: str = Field(..., description="Path to converted markdown")
+    num_lines: int = Field(..., description="Total number of lines in markdown")
+    table_of_contents: list[TOCEntry] = Field(
+        default_factory=list, description="Parsed table of contents"
+    )
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    file_size: int = Field(0, description="Original file size in bytes")
+    file_type: str = Field("", description="Original file extension")
+
+    def format_toc(self) -> str:
+        """Format table of contents as a readable string."""
+        if not self.table_of_contents:
+            return "No table of contents available"
+
+        lines = []
+        for entry in self.table_of_contents:
+            indent = "  " * (entry.level - 1)
+            line_info = f"[lines {entry.start_line}-{entry.end_line or '?'}]"
+            lines.append(f"{indent}- {entry.title} {line_info}")
+        return "\n".join(lines)
+
+
+class SearchHit(BaseModel):
+    """Search result for document-level search."""
+
+    id: str = Field(..., description="Document identifier")
+    name: str = Field(..., description="Document name")
+    num_lines: int = Field(..., description="Total lines in document")
+    table_of_contents: str = Field(..., description="Formatted table of contents")
+    score: float = Field(default=0.0, description="Search relevance score")
+    snippet: str = Field(default="", description="Relevant text snippet")
+
+
+class Excerpt(BaseModel):
+    """Text excerpt from a document."""
+
+    doc_id: str = Field(..., description="Document identifier")
+    line_range: tuple[int, int] = Field(..., description="Line range (start, end), 1-indexed")
+    text: str = Field(..., description="Extracted text content")
+    total_lines: int = Field(0, description="Total lines in document")
+
+
+class ContentSearchHit(BaseModel):
+    """Search result for within-document content search."""
+
+    doc_id: str = Field(..., description="Document identifier")
+    line_range: tuple[int, int] = Field(..., description="Line range of the match")
+    text: str = Field(..., description="Matching text content")
+    score: float = Field(0.0, description="Search relevance score")
+
+
+class ChunkMetadata(BaseModel):
+    """Metadata stored with each chunk in the vector store."""
+
+    doc_id: str = Field(..., description="Parent document ID")
+    chunk_index: int = Field(..., description="Chunk index within document")
+    start_line: int = Field(..., description="Starting line number")
+    end_line: int = Field(..., description="Ending line number")
+    text: str = Field(..., description="Chunk text content")

athenaeum/ocr/__init__.py

@@ -0,0 +1,48 @@
+"""OCR provider registry and factory."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from athenaeum.ocr.base import OCRProvider
+from athenaeum.ocr.custom import CustomOCR
+from athenaeum.ocr.markitdown import MarkitdownOCR
+
+OCRBackend = Literal["markitdown", "docling", "mistral", "lighton"]
+
+
+def get_ocr_provider(backend: OCRBackend = "markitdown", **kwargs: object) -> OCRProvider:
+    """Factory to create an OCR provider by name.
+
+    Args:
+        backend: One of ``"markitdown"``, ``"docling"``, ``"mistral"``, ``"lighton"``.
+        **kwargs: Passed to the provider constructor.
+
+    Returns:
+        An ``OCRProvider`` instance.
+    """
+    if backend == "markitdown":
+        return MarkitdownOCR()
+    if backend == "docling":
+        from athenaeum.ocr.docling import DoclingOCR
+
+        return DoclingOCR()
+    if backend == "mistral":
+        from athenaeum.ocr.mistral import MistralOCR
+
+        return MistralOCR(**kwargs)  # type: ignore[arg-type]
+    if backend == "lighton":
+        from athenaeum.ocr.lighton import LightOnOCR
+
+        return LightOnOCR(**kwargs)  # type: ignore[arg-type]
+
+    raise ValueError(f"Unknown OCR backend: {backend!r}")
+
+
+__all__ = [
+    "CustomOCR",
+    "MarkitdownOCR",
+    "OCRBackend",
+    "OCRProvider",
+    "get_ocr_provider",
+]

athenaeum/ocr/base.py

@@ -0,0 +1,28 @@
+"""Abstract base class for OCR providers."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+
+class OCRProvider(ABC):
+    """Base class for OCR/document-to-markdown converters."""
+
+    @abstractmethod
+    def convert(self, file_path: Path) -> str:
+        """Convert a file to markdown text.
+
+        Args:
+            file_path: Path to the source file.
+
+        Returns:
+            Markdown string of the file contents.
+        """
+
+    @abstractmethod
+    def supported_extensions(self) -> set[str]:
+        """Return the set of file extensions this provider supports.
+
+        Extensions should include the leading dot, e.g. ``{".pdf", ".docx"}``.
+        """

athenaeum/ocr/custom.py

@@ -0,0 +1,26 @@
+"""Custom OCR provider wrapping a user-supplied callable."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+
+from athenaeum.ocr.base import OCRProvider
+
+
+class CustomOCR(OCRProvider):
+    """Wrap an arbitrary ``(Path) -> str`` callable as an OCR provider."""
+
+    def __init__(
+        self,
+        fn: Callable[[Path], str],
+        extensions: set[str] | None = None,
+    ) -> None:
+        self._fn = fn
+        self._extensions = extensions or {".*"}
+
+    def convert(self, file_path: Path) -> str:
+        return self._fn(file_path)
+
+    def supported_extensions(self) -> set[str]:
+        return self._extensions

athenaeum/ocr/docling.py

@@ -0,0 +1,30 @@
+"""OCR provider using Docling."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from athenaeum.ocr.base import OCRProvider
+
+
+class DoclingOCR(OCRProvider):
+    """Convert documents to markdown using Docling."""
+
+    _EXTENSIONS = {".pdf", ".pptx", ".docx", ".xlsx", ".html", ".md"}
+
+    def __init__(self) -> None:
+        try:
+            from docling.document_converter import DocumentConverter
+        except ImportError as e:
+            raise ImportError(
+                "Docling is not installed. Install with: pip install 'athenaeum-kb[docling]'"
+            ) from e
+        self._converter = DocumentConverter()
+
+    def convert(self, file_path: Path) -> str:
+        result = self._converter.convert(str(file_path))
+        md: str = result.document.export_to_markdown()
+        return md
+
+    def supported_extensions(self) -> set[str]:
+        return self._EXTENSIONS

athenaeum/ocr/lighton.py

@@ -0,0 +1,40 @@
+"""OCR provider using LightOnOCR-2-1B."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from athenaeum.ocr.base import OCRProvider
+
+
+class LightOnOCR(OCRProvider):
+    """Convert documents to markdown using LightOnOCR-2-1B (local model)."""
+
+    _EXTENSIONS = {".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"}
+
+    def __init__(self, device: str = "cpu") -> None:
+        try:
+            from transformers import pipeline
+        except ImportError as e:
+            raise ImportError(
+                "transformers/torch not installed. "
+                "Install with: pip install 'athenaeum-kb[lighton]'"
+            ) from e
+        self._pipe = pipeline(
+            "image-text-to-text",
+            model="lightonai/LightOnOCR-2-1B",
+            device=device,
+        )
+
+    def convert(self, file_path: Path) -> str:
+        from PIL import Image
+
+        image = Image.open(file_path)
+        result = self._pipe(image)
+        if isinstance(result, list) and len(result) > 0:
+            text: str = result[0].get("generated_text", "")
+            return text
+        return str(result)
+
+    def supported_extensions(self) -> set[str]:
+        return self._EXTENSIONS

athenaeum/ocr/markitdown.py

@@ -0,0 +1,29 @@
+"""OCR provider using Microsoft markitdown."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from markitdown import MarkItDown
+
+from athenaeum.ocr.base import OCRProvider
+
+
+class MarkitdownOCR(OCRProvider):
+    """Convert documents to markdown using markitdown."""
+
+    _EXTENSIONS = {
+        ".pdf", ".pptx", ".docx", ".xlsx",
+        ".json", ".csv", ".txt", ".md",
+        ".html", ".xml", ".rtf", ".epub",
+    }
+
+    def __init__(self) -> None:
+        self._converter = MarkItDown()
+
+    def convert(self, file_path: Path) -> str:
+        result = self._converter.convert(str(file_path))
+        return result.text_content
+
+    def supported_extensions(self) -> set[str]:
+        return self._EXTENSIONS

athenaeum/ocr/mistral.py

@@ -0,0 +1,44 @@
+"""OCR provider using Mistral OCR API."""
+
+from __future__ import annotations
+
+import base64
+import os
+from pathlib import Path
+
+from athenaeum.ocr.base import OCRProvider
+
+
+class MistralOCR(OCRProvider):
+    """Convert documents to markdown using Mistral's OCR API."""
+
+    _EXTENSIONS = {".pdf"}
+
+    def __init__(self, api_key: str | None = None) -> None:
+        try:
+            from mistralai import Mistral
+        except ImportError as e:
+            raise ImportError(
+                "Mistral SDK is not installed. Install with: pip install 'athenaeum-kb[mistral]'"
+            ) from e
+        key = api_key or os.environ.get("MISTRAL_API_KEY", "")
+        if not key:
+            raise ValueError("MISTRAL_API_KEY must be set or passed explicitly.")
+        self._client = Mistral(api_key=key)
+
+    def convert(self, file_path: Path) -> str:
+        from mistralai import DocumentURLChunk, OCRRequest
+
+        encoded = base64.standard_b64encode(file_path.read_bytes()).decode()
+        data_uri = f"data:application/pdf;base64,{encoded}"
+        response = self._client.ocr.process(
+            request=OCRRequest(
+                document=DocumentURLChunk(document_url=data_uri),
+                model="mistral-ocr-latest",
+            )
+        )
+        pages = [page.markdown for page in response.pages]
+        return "\n\n".join(pages)
+
+    def supported_extensions(self) -> set[str]:
+        return self._EXTENSIONS

athenaeum/search/__init__.py

@@ -0,0 +1,7 @@
+"""Search subsystem for Athenaeum."""
+
+from athenaeum.search.bm25 import BM25Index
+from athenaeum.search.hybrid import reciprocal_rank_fusion
+from athenaeum.search.vector import VectorIndex
+
+__all__ = ["BM25Index", "VectorIndex", "reciprocal_rank_fusion"]

athenaeum/search/bm25.py

@@ -0,0 +1,72 @@
+"""BM25 keyword search index."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from rank_bm25 import BM25Okapi
+
+from athenaeum.models import ChunkMetadata
+
+
+@dataclass
+class _Entry:
+    chunk: ChunkMetadata
+    tokens: list[str]
+
+
+class BM25Index:
+    """In-memory BM25 index over document chunks."""
+
+    def __init__(self) -> None:
+        self._entries: list[_Entry] = []
+        self._bm25: BM25Okapi | None = None
+
+    def _rebuild(self) -> None:
+        if self._entries:
+            corpus = [e.tokens for e in self._entries]
+            self._bm25 = BM25Okapi(corpus)
+        else:
+            self._bm25 = None
+
+    @staticmethod
+    def _tokenize(text: str) -> list[str]:
+        return text.lower().split()
+
+    def add_chunks(self, chunks: list[ChunkMetadata]) -> None:
+        """Add chunks to the index and rebuild."""
+        for chunk in chunks:
+            self._entries.append(_Entry(chunk=chunk, tokens=self._tokenize(chunk.text)))
+        self._rebuild()
+
+    def remove_document(self, doc_id: str) -> None:
+        """Remove all chunks for a document and rebuild."""
+        self._entries = [e for e in self._entries if e.chunk.doc_id != doc_id]
+        self._rebuild()
+
+    def search(
+        self,
+        query: str,
+        top_k: int = 10,
+        doc_id: str | None = None,
+    ) -> list[tuple[ChunkMetadata, float]]:
+        """Search the index, returning (chunk, score) pairs sorted by score descending."""
+        if not self._bm25 or not self._entries:
+            return []
+
+        tokens = self._tokenize(query)
+        scores = self._bm25.get_scores(tokens)
+
+        results: list[tuple[ChunkMetadata, float]] = []
+        for i, score in enumerate(scores):
+            entry = self._entries[i]
+            if doc_id is not None and entry.chunk.doc_id != doc_id:
+                continue
+            results.append((entry.chunk, float(score)))
+
+        results.sort(key=lambda x: x[1], reverse=True)
+        return results[:top_k]
+
+    @property
+    def size(self) -> int:
+        return len(self._entries)

athenaeum/search/hybrid.py

@@ -0,0 +1,37 @@
+"""Reciprocal Rank Fusion for combining ranked search results."""
+
+from __future__ import annotations
+
+from athenaeum.models import ChunkMetadata
+
+
+def reciprocal_rank_fusion(
+    ranked_lists: list[list[tuple[ChunkMetadata, float]]],
+    k: int = 60,
+    top_k: int = 10,
+) -> list[tuple[ChunkMetadata, float]]:
+    """Combine multiple ranked lists using Reciprocal Rank Fusion.
+
+    RRF score for a chunk = sum over lists of 1 / (k + rank), where rank is
+    1-indexed. Higher scores indicate better combined relevance.
+
+    Args:
+        ranked_lists: List of ranked result lists, each containing (chunk, score) pairs.
+        k: RRF constant (default 60).
+        top_k: Number of results to return.
+
+    Returns:
+        Merged and re-ranked list of (chunk, rrf_score) pairs.
+    """
+    scores: dict[str, float] = {}
+    chunks: dict[str, ChunkMetadata] = {}
+
+    for ranked_list in ranked_lists:
+        for rank, (chunk, _score) in enumerate(ranked_list, start=1):
+            key = f"{chunk.doc_id}:{chunk.chunk_index}"
+            scores[key] = scores.get(key, 0.0) + 1.0 / (k + rank)
+            chunks[key] = chunk
+
+    merged = [(chunks[key], score) for key, score in scores.items()]
+    merged.sort(key=lambda x: x[1], reverse=True)
+    return merged[:top_k]

athenaeum/search/vector.py

@@ -0,0 +1,72 @@
+"""Vector similarity search index using LangChain + Chroma."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from langchain_chroma import Chroma
+from langchain_core.embeddings import Embeddings
+
+from athenaeum.models import ChunkMetadata
+
+
+class VectorIndex:
+    """Vector similarity search backed by Chroma."""
+
+    def __init__(
+        self,
+        embeddings: Embeddings,
+        persist_directory: str | Path | None = None,
+        collection_name: str = "athenaeum",
+    ) -> None:
+        kwargs: dict[str, object] = {
+            "embedding_function": embeddings,
+            "collection_name": collection_name,
+        }
+        if persist_directory is not None:
+            kwargs["persist_directory"] = str(persist_directory)
+        self._store = Chroma(**kwargs)  # type: ignore[arg-type]
+        self._embeddings = embeddings
+
+    def add_chunks(self, chunks: list[ChunkMetadata]) -> None:
+        """Add chunks to the vector store."""
+        if not chunks:
+            return
+        texts = [c.text for c in chunks]
+        metadatas = [c.model_dump() for c in chunks]
+        ids = [f"{c.doc_id}:{c.chunk_index}" for c in chunks]
+        self._store.add_texts(texts=texts, metadatas=metadatas, ids=ids)
+
+    def remove_document(self, doc_id: str) -> None:
+        """Remove all chunks for a document from the vector store."""
+        self._store._collection.delete(where={"doc_id": doc_id})
+
+    def search(
+        self,
+        query: str,
+        top_k: int = 10,
+        doc_id: str | None = None,
+    ) -> list[tuple[ChunkMetadata, float]]:
+        """Search for similar chunks, returning (chunk, score) pairs.
+
+        Scores are similarity scores (higher = more similar).
+        """
+        kwargs: dict[str, object] = {"k": top_k}
+        if doc_id is not None:
+            kwargs["filter"] = {"doc_id": doc_id}
+
+        results = self._store.similarity_search_with_relevance_scores(query, **kwargs)  # type: ignore[arg-type]
+
+        output: list[tuple[ChunkMetadata, float]] = []
+        for doc, score in results:
+            meta = doc.metadata
+            chunk = ChunkMetadata(
+                doc_id=meta["doc_id"],
+                chunk_index=meta["chunk_index"],
+                start_line=meta["start_line"],
+                end_line=meta["end_line"],
+                text=meta["text"],
+            )
+            output.append((chunk, float(score)))
+
+        return output

athenaeum/storage.py

@@ -0,0 +1,68 @@
+"""Storage layout manager for Athenaeum."""
+
+from __future__ import annotations
+
+import json
+import shutil
+from pathlib import Path
+from typing import Any
+
+
+class StorageManager:
+    """Manages the on-disk layout under the storage root.
+
+    Layout::
+
+        <root>/
+            docs/<doc_id>/raw.*       # original file
+            docs/<doc_id>/content.md  # converted markdown
+            index/chroma/             # Chroma persistent directory
+            metadata.json             # document registry
+    """
+
+    def __init__(self, root: Path) -> None:
+        self.root = root
+        self.root.mkdir(parents=True, exist_ok=True)
+
+    @property
+    def docs_dir(self) -> Path:
+        return self.root / "docs"
+
+    @property
+    def chroma_dir(self) -> Path:
+        return self.root / "index" / "chroma"
+
+    @property
+    def metadata_path(self) -> Path:
+        return self.root / "metadata.json"
+
+    def doc_dir(self, doc_id: str) -> Path:
+        d = self.docs_dir / doc_id
+        d.mkdir(parents=True, exist_ok=True)
+        return d
+
+    def raw_path(self, doc_id: str, suffix: str) -> Path:
+        """Return path for storing the original file."""
+        return self.doc_dir(doc_id) / f"raw{suffix}"
+
+    def content_md_path(self, doc_id: str) -> Path:
+        """Return path for the converted markdown."""
+        return self.doc_dir(doc_id) / "content.md"
+
+    def remove_doc(self, doc_id: str) -> None:
+        """Remove a document's directory."""
+        d = self.docs_dir / doc_id
+        if d.exists():
+            shutil.rmtree(d)
+
+    def ensure_chroma_dir(self) -> Path:
+        self.chroma_dir.mkdir(parents=True, exist_ok=True)
+        return self.chroma_dir
+
+    def load_metadata(self) -> dict[str, Any]:
+        if self.metadata_path.exists():
+            return json.loads(self.metadata_path.read_text())  # type: ignore[no-any-return]
+        return {}
+
+    def save_metadata(self, data: dict[str, Any]) -> None:
+        self.metadata_path.write_text(json.dumps(data, indent=2, default=str))

athenaeum/toc.py

@@ -0,0 +1,49 @@
+"""Table of contents extraction from markdown headings."""
+
+from __future__ import annotations
+
+import re
+
+from athenaeum.models import TOCEntry
+
+_HEADING_RE = re.compile(r"^(#{1,6})\s+(.+)$")
+
+
+def extract_toc(markdown: str) -> list[TOCEntry]:
+    """Extract a table of contents from markdown heading lines.
+
+    Each entry records the heading level, title, and the line range it spans.
+    Line numbers are 1-indexed. The ``end_line`` of each entry is set to the
+    line before the next heading at the same or higher level, or to the last
+    line of the document for the final entry.
+    """
+    lines = markdown.split("\n")
+    entries: list[TOCEntry] = []
+
+    for line_no_0, line in enumerate(lines):
+        m = _HEADING_RE.match(line.strip())
+        if m:
+            level = len(m.group(1))
+            title = m.group(2).strip()
+            entries.append(
+                TOCEntry(title=title, level=level, start_line=line_no_0 + 1, end_line=None)
+            )
+
+    # Fill in end_line for each entry
+    total_lines = len(lines)
+    for i, entry in enumerate(entries):
+        # Find next entry at same or higher (lower number) level
+        end = total_lines
+        for j in range(i + 1, len(entries)):
+            if entries[j].level <= entry.level:
+                end = entries[j].start_line - 1
+                break
+        else:
+            # Last entry at this level — ends at document end
+            if i + 1 < len(entries):
+                end = entries[-1].start_line - 1
+                # Actually, just go to end of document
+                end = total_lines
+        entry.end_line = end
+
+    return entries

athenaeum_kb-0.1.0.dist-info/METADATA

@@ -0,0 +1,165 @@
+Metadata-Version: 2.4
+Name: athenaeum-kb
+Version: 0.1.0
+Summary: Tools for intelligent interaction with knowledge bases
+License-Expression: MIT
+Requires-Python: >=3.11
+Requires-Dist: langchain-chroma>=0.2
+Requires-Dist: langchain-core>=0.3
+Requires-Dist: langchain-openai>=0.3
+Requires-Dist: markitdown>=0.1
+Requires-Dist: pydantic>=2.0
+Requires-Dist: rank-bm25>=0.2.2
+Provides-Extra: all-ocr
+Requires-Dist: docling>=2.0; extra == 'all-ocr'
+Requires-Dist: mistralai>=1.0; extra == 'all-ocr'
+Requires-Dist: pillow>=10.0; extra == 'all-ocr'
+Requires-Dist: torch>=2.0; extra == 'all-ocr'
+Requires-Dist: transformers>=4.40; extra == 'all-ocr'
+Provides-Extra: dev
+Requires-Dist: mypy>=1.10; extra == 'dev'
+Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff>=0.4; extra == 'dev'
+Provides-Extra: docling
+Requires-Dist: docling>=2.0; extra == 'docling'
+Provides-Extra: lighton
+Requires-Dist: pillow>=10.0; extra == 'lighton'
+Requires-Dist: torch>=2.0; extra == 'lighton'
+Requires-Dist: transformers>=4.40; extra == 'lighton'
+Provides-Extra: mistral
+Requires-Dist: mistralai>=1.0; extra == 'mistral'
+Description-Content-Type: text/markdown
+
+# Athenaeum 
+
+## Project Scope and Goals
+
+The goal of this project is to build a Python library that equips AI agents with a robust set of tools for intelligent interaction with knowledge bases. The library focuses on document ingestion, semantic search, and structured access to content, making it suitable for agent-based systems, RAG pipelines, and automation workflows.
+
+Once the module is fully tested and validated, the intended outcome is to package and publish it as a reusable Python library on PyPI.
+
+## Tools
+
+### `load_doc`
+Load a document into the knowledge base, automatically extracting content, metadata, and embeddings.
+
+```python
+load_doc(path: str) -> str
+```
+
+**Parameters:**
+- `path`: Path to the document file
+
+**Supported formats:** PDF, PPTX, DOCX, XLSX, JSON, CSV, TXT, MD, HTML, XML, RTF, EPUB
+
+**Returns**: A document identifier (doc_id) that can be used for subsequent operations.
+
+### `list_docs`
+List all documents currently stored in the knowledge base.
+
+```python
+list_docs() -> list[DocSearchHit]
+```
+
+**Returns:** A list of documents, including metadata (id, name, format, etc.) and, when available, a table of contents.
+
+### `search_docs`
+
+Search across all documents in the knowledge base.
+
+```python
+search_docs(
+    query: str,
+    top_k: int = 10,
+    scope: Literal["names", "contents"] = "contents",
+    strategy: Literal["hybrid", "bm25", "vector"] = "hybrid"
+) -> list[DocSearchHit]
+```
+
+**Parameters:**
+- `query`: Search query text
+- `top_k`: Maximum number of results (default: 10)
+- `scope`: Where to search
+  - `"contents"`: Search within document contents (default)
+  - `"names"`: Search only document names
+- `strategy`: Search strategy (only applies when scope is "contents")
+  - `"hybrid"`: Combines vector and BM25 search (default)
+  - `"bm25"`: Keyword-based search only
+  - `"vector"`: Semantic similarity search only
+
+**Returns**: A ranked list of documents matching the query.
+
+### `search_doc_contents`
+
+Search within a specific document.
+
+```python
+search_doc_contents(
+    doc_id: str,
+    query: str,
+    top_k: int = 5,
+    strategy: Literal["hybrid", "bm25", "vector"] = "hybrid"
+) -> list[ContentSearchHit]
+```
+
+**Parameters:**
+- `doc_id`: Document identifier
+- `query`: Search query text
+- `top_k`: Maximum number of results (default: 5)
+- `strategy`: Search strategy
+  - `"hybrid"`: Combines vector and BM25 search (default)
+  - `"bm25"`: Keyword-based search only
+  - `"vector"`: Semantic similarity search only
+
+**Returns**: A list of matching content fragments with relevance scores.
+
+### `read_doc`
+
+Read a specific range of lines from a document.
+
+```python
+read_doc(
+    doc_id: str,
+    start_line: int = 1,
+    end_line: int = 100
+) -> Excerpt
+```
+
+**Parameters:**
+- `doc_id`: Document identifier
+- `start_line`: Starting line number (1-indexed, default: 1)
+- `end_line`: Ending line number (1-indexed, inclusive, default: 100)
+
+**Returns**: A document excerpt containing the requested lines.
+
+## Search Strategies
+
+- **Hybrid Search** (Default): Combines vector similarity search with BM25 keyword search using Reciprocal Rank Fusion (RRF).
+- **Vector Search**: Uses embedding models for semantic similarity search.
+- **BM25 Search**: Traditional keyword-based search using the BM25 algorithm.
+
+## Document Ingestion Workflow
+
+The ingestion pipeline is triggered by the load_doc(path) function and follows these steps:
+1. Validation
+    - Verify that the file exists.
+    - Confirm that the file format is supported.
+2. Content Extraction
+    - Run an OCR or parsing pipeline to convert the raw file into Markdown.
+    - If the document contains images:
+        - Replace them with placeholders in the Markdown.
+        - Store image references for later retrieval.
+3. Pre-processing
+    - Generate document metadata (e.g., id, name, format).
+    - Build a table of contents (TOC) from Markdown headings when possible.
+4. Indexing
+    - Generate vector embeddings using the configured embedding model.
+    - Store embeddings in the vector database for semantic retrieval.
+
+## Data Models
+A preliminary set of domain models is defined in `models.py`. These classes are exploratory and serve as a conceptual starting point for the project.
+
+They are not considered final and may be refactored, renamed, or removed as the overall architecture of Athenaeum evolves and solidifies.
+
+

athenaeum_kb-0.1.0.dist-info/RECORD

@@ -0,0 +1,22 @@
+athenaeum/__init__.py,sha256=3oXIneyMISBVERMs4KuJOJ7MrjZj5I7lC19wQaTuzvA,585
+athenaeum/athenaeum.py,sha256=kl-DZnm5qPyY-oVEcJIgZnprobWOumKPziSyz5YlaE8,9330
+athenaeum/chunker.py,sha256=Sk2h5Z3oVhK6BPdXRMWVMONurwIsMgRkRIorbEPl_Po,2050
+athenaeum/config.py,sha256=NcQa5t3sj0XacqmJgWLE88Ci9uHafWy6Z988WE36rqg,441
+athenaeum/document_store.py,sha256=8llflG_k0UFjtr2Xb4YWnDDpFzFvfA_OQZ7WELyP6Qo,1611
+athenaeum/models.py,sha256=ZeVQkPtS5WBPWkdSTYa-Z0ZZ9jtPdqDwnH16803gu54,3647
+athenaeum/storage.py,sha256=4kV-wPtJ-MPS1ApPlr8W66bLYmtYGknnta4-lFHLXuo,2028
+athenaeum/toc.py,sha256=qROrkhUg7JOlLpgHwuM2-JZd8eX2WO4iMsZh0E68zwY,1639
+athenaeum/ocr/__init__.py,sha256=aG2NgzHL8TZXlfw3HDPZI8-t6gNJSTV1iLODIOSLEFE,1304
+athenaeum/ocr/base.py,sha256=qdw--LAUPIloFWuKGaGLD7lPzNttcVL0HOjAJbZBqQk,716
+athenaeum/ocr/custom.py,sha256=pddvXL-3aAR4Fo3nNDxfjVHO0D6xhTmpI6rLPDgazCU,672
+athenaeum/ocr/docling.py,sha256=sPRu6ONiC6L2ZHi7RLYqFHFX6oLjRMsfsSR__9ssJeo,898
+athenaeum/ocr/lighton.py,sha256=oArJ4BvbxByFTYu4HaokKXLngvQdX7_RJIZUlPe6WBc,1197
+athenaeum/ocr/markitdown.py,sha256=M6bCoqsIvm_z_tLQeTq5ie-CaFihPspu0wZUIudrNZ8,727
+athenaeum/ocr/mistral.py,sha256=XMGk1yLhwbdi_glzdsjLdOptbMwmc-86Fate3iSTwUI,1436
+athenaeum/search/__init__.py,sha256=BCmd-4lIZmxVQZTC3GNt5-uz0pI8G9wd35MYmmIWdIA,256
+athenaeum/search/bm25.py,sha256=b-SbtFaXwqt2IzLaytwwSVo8LBm_v0UotJeEeIsEgl8,2069
+athenaeum/search/hybrid.py,sha256=a9kFNq6WBnNZtC6Rocx9ZI5BTe-f53KR8beaVC12cQU,1273
+athenaeum/search/vector.py,sha256=Hu_7gbwiRLYlkpZGUnf0SD9LXWHZ7IxIUEls1w4blFo,2428
+athenaeum_kb-0.1.0.dist-info/METADATA,sha256=5_5bjWHZ3XJLs5Da85lMEgDuNe1IeLB7CmgYFcI1IAo,5549
+athenaeum_kb-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+athenaeum_kb-0.1.0.dist-info/RECORD,,

athenaeum_kb-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any