docsif 0.1.0__py3-none-any.whl

sif/embedding/model.py

@@ -0,0 +1,126 @@
+"""Embedding model interface."""
+
+from abc import ABC, abstractmethod
+from typing import Protocol
+
+from sif.models.embedding import ModelType
+
+
+class EmbeddingModel(ABC):
+    """Abstract base class for embedding models.
+
+    Implements a common interface for different embedding providers,
+    allowing models to be used interchangeably.
+    """
+
+    def __init__(
+        self,
+        model_id: str,
+        embedding_dim: int,
+        max_tokens: int,
+    ) -> None:
+        """Initialize the embedding model."""
+        self._model_id = model_id
+        self._embedding_dim = embedding_dim
+        self._max_tokens = max_tokens
+        self._loaded = False
+
+    @property
+    def model_id(self) -> str:
+        """Get the model ID."""
+        return self._model_id
+
+    @property
+    def embedding_dim(self) -> int:
+        """Get the embedding dimension."""
+        return self._embedding_dim
+
+    @property
+    def max_tokens(self) -> int:
+        """Get the maximum token count."""
+        return self._max_tokens
+
+    @property
+    def loaded(self) -> bool:
+        """Check if the model is loaded."""
+        return self._loaded
+
+    @abstractmethod
+    def load(self) -> None:
+        """Load the model into memory."""
+        ...
+
+    @abstractmethod
+    def unload(self) -> None:
+        """Unload the model from memory."""
+        ...
+
+    @abstractmethod
+    def embed(
+        self,
+        texts: list[str],
+        normalize: bool = True,
+    ) -> list[list[float]]:
+        """Generate embeddings for a list of texts.
+
+        Args:
+            texts: List of texts to embed
+            normalize: Whether to normalize embeddings
+
+        Returns:
+            List of embeddings, one per input text
+        """
+        ...
+
+    @abstractmethod
+    def embed_single(
+        self,
+        text: str,
+        normalize: bool = True,
+    ) -> list[float]:
+        """Generate embedding for a single text.
+
+        Args:
+            text: Text to embed
+            normalize: Whether to normalize the embedding
+
+        Returns:
+            Embedding vector
+        """
+        ...
+
+    def count_tokens(self, text: str) -> int:
+        """Count tokens in a text.
+
+        Default implementation uses a rough estimate.
+        Subclasses can override for more accurate counts.
+        """
+        # Rough estimate: 1 token ~= 4 characters
+        return len(text) // 4
+
+
+class EmbeddingModelFactory(Protocol):
+    """Protocol for embedding model factories.
+
+    Implements the Factory pattern for creating embedding models.
+    """
+
+    def create_model(
+        self,
+        model_type: ModelType,
+        model_path: str | None,
+        model_name: str,
+        **kwargs: dict[str, any],
+    ) -> EmbeddingModel:
+        """Create an embedding model instance.
+
+        Args:
+            model_type: Type of model to create
+            model_path: Path to model file (for local models)
+            model_name: Model name or identifier
+            **kwargs: Additional model-specific parameters
+
+        Returns:
+            Configured embedding model instance
+        """
+        ...

sif/indexing/__init__.py

@@ -0,0 +1,21 @@
+"""Document indexing functionality for SIF.
+
+This package provides document indexing capabilities:
+- File system scanning
+- Markdown parsing and chunking
+- Document metadata extraction
+"""
+
+from sif.indexing.chunker import Chunker, create_chunker
+from sif.indexing.parser import MarkdownParser, ParsedDocument
+from sif.indexing.scanner import FileScanner, ScanResult
+
+
+__all__ = [
+    "Chunker",
+    "FileScanner",
+    "MarkdownParser",
+    "ParsedDocument",
+    "ScanResult",
+    "create_chunker",
+]

sif/indexing/chunker.py

@@ -0,0 +1,226 @@
+"""Document chunking strategies."""
+
+from __future__ import annotations
+
+import re
+from abc import ABC, abstractmethod
+
+from sif.core.models import DocumentChunk
+
+
+class Chunker(ABC):
+    """Base class for document chunkers."""
+
+    @abstractmethod
+    def chunk(self, text: str) -> list[DocumentChunk]:
+        """Split text into chunks.
+
+        Args:
+            text: Text to chunk
+
+        Returns:
+            List of chunks
+        """
+        pass
+
+
+class FixedSizeChunker(Chunker):
+    """Chunk text into fixed-size pieces with overlap."""
+
+    def __init__(self, chunk_size: int = 1000, overlap: int = 100) -> None:
+        """Initialize fixed-size chunker.
+
+        Args:
+            chunk_size: Target chunk size in characters
+            overlap: Overlap between chunks in characters
+        """
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+
+    def chunk(self, text: str) -> list[DocumentChunk]:
+        """Split text into fixed-size chunks."""
+        chunks = []
+        start = 0
+        sequence = 0
+
+        while start < len(text):
+            end = start + self.chunk_size
+
+            # Try to find a good break point
+            if end < len(text):
+                # Look for newline or space
+                for i in range(min(end + 100, len(text)) - 1, end - 1, -1):
+                    if text[i] in "\n ":
+                        end = i + 1
+                        break
+
+            chunk_text = text[start:end].strip()
+            if chunk_text:
+                chunks.append(
+                    DocumentChunk(
+                        content=chunk_text,
+                        sequence=sequence,
+                        start_pos=start,
+                        end_pos=end,
+                    )
+                )
+                sequence += 1
+
+            # Move start with overlap
+            start = end - self.overlap
+            start = min(end, start)
+
+        return chunks
+
+
+class MarkdownChunker(Chunker):
+    """Chunk markdown documents by headings."""
+
+    def __init__(self, max_chunk_size: int = 2000) -> None:
+        """Initialize markdown chunker.
+
+        Args:
+            max_chunk_size: Maximum chunk size in characters
+        """
+        self.max_chunk_size = max_chunk_size
+        self.heading_pattern = re.compile(r"^(#{1,6}\s+.+)$", re.MULTILINE)
+
+    def chunk(self, text: str) -> list[DocumentChunk]:
+        """Split markdown by headings."""
+        # Find all headings
+        headings = list(self.heading_pattern.finditer(text))
+
+        if not headings:
+            # No headings, use fixed-size chunking
+            return FixedSizeChunker(self.max_chunk_size).chunk(text)
+
+        chunks = []
+        sequence = 0
+
+        for i, match in enumerate(headings):
+            start = match.start()
+            end = headings[i + 1].start() if i + 1 < len(headings) else len(text)
+
+            chunk_text = text[start:end].strip()
+
+            # If chunk is too large, split it further
+            if len(chunk_text) > self.max_chunk_size:
+                sub_chunks = FixedSizeChunker(self.max_chunk_size, 100).chunk(chunk_text)
+                for sub_chunk in sub_chunks:
+                    sub_chunk.sequence = sequence
+                    chunks.append(sub_chunk)
+                    sequence += 1
+            else:
+                chunks.append(
+                    DocumentChunk(
+                        content=chunk_text,
+                        sequence=sequence,
+                        start_pos=start,
+                        end_pos=end,
+                    )
+                )
+                sequence += 1
+
+        return chunks
+
+
+class CodeChunker(Chunker):
+    """Chunk code files by functions/classes."""
+
+    def __init__(self, max_chunk_size: int = 2000, language: str = "python") -> None:
+        """Initialize code chunker.
+
+        Args:
+            max_chunk_size: Maximum chunk size in characters
+            language: Programming language
+        """
+        self.max_chunk_size = max_chunk_size
+        self.language = language
+
+    def chunk(self, text: str) -> list[DocumentChunk]:
+        """Split code by functions/classes."""
+        # Simple regex-based approach for Python
+        if self.language == "python":
+            pattern = re.compile(r"^(def\s+|class\s+)", re.MULTILINE)
+        else:
+            # Fallback to fixed-size for other languages
+            return FixedSizeChunker(self.max_chunk_size).chunk(text)
+
+        matches = list(pattern.finditer(text))
+
+        if not matches:
+            return FixedSizeChunker(self.max_chunk_size).chunk(text)
+
+        chunks = []
+        sequence = 0
+
+        for i, match in enumerate(matches):
+            start = match.start()
+            end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
+
+            chunk_text = text[start:end].strip()
+
+            if len(chunk_text) > self.max_chunk_size:
+                sub_chunks = FixedSizeChunker(self.max_chunk_size, 100).chunk(chunk_text)
+                for sub_chunk in sub_chunks:
+                    sub_chunk.sequence = sequence
+                    chunks.append(sub_chunk)
+                    sequence += 1
+            else:
+                chunks.append(
+                    DocumentChunk(
+                        content=chunk_text,
+                        sequence=sequence,
+                        start_pos=start,
+                        end_pos=end,
+                    )
+                )
+                sequence += 1
+
+        return chunks
+
+
+class AutoChunker(Chunker):
+    """Automatically select chunking strategy based on content."""
+
+    def __init__(self, max_chunk_size: int = 2000) -> None:
+        """Initialize auto chunker.
+
+        Args:
+            max_chunk_size: Maximum chunk size in characters
+        """
+        self.max_chunk_size = max_chunk_size
+
+    def chunk(self, text: str) -> list[DocumentChunk]:
+        """Automatically select and apply chunking strategy."""
+        # Check if markdown
+        if "#" in text and re.search(r"^#+\s+", text, re.MULTILINE):
+            return MarkdownChunker(self.max_chunk_size).chunk(text)
+
+        # Check if code (has function/class definitions)
+        if re.search(r"^(def\s+|class\s+)", text, re.MULTILINE):
+            return CodeChunker(self.max_chunk_size).chunk(text)
+
+        # Default to fixed-size
+        return FixedSizeChunker(self.max_chunk_size).chunk(text)
+
+
+def create_chunker(strategy: str, **kwargs) -> Chunker:
+    """Create a chunker by strategy name.
+
+    Args:
+        strategy: Chunking strategy ("fixed", "markdown", "code", "auto")
+        **kwargs: Additional arguments for the chunker
+
+    Returns:
+        Chunker instance
+    """
+    if strategy == "fixed":
+        return FixedSizeChunker(**kwargs)
+    if strategy == "markdown":
+        return MarkdownChunker(**kwargs)
+    if strategy == "code":
+        return CodeChunker(**kwargs)
+    if strategy == "auto":
+        return AutoChunker(**kwargs)
+    raise ValueError(f"Unknown chunking strategy: {strategy}")

sif/indexing/indexer.py

@@ -0,0 +1,226 @@
+"""Document indexer for orchestrating the indexing process."""
+
+from dataclasses import dataclass
+from enum import Enum, auto
+from pathlib import Path
+
+from sif.core.models import Collection, Document, DocumentChunk
+from sif.database.repositories import DocumentRepository
+from sif.embedding.manager import EmbeddingManager
+from sif.indexing.chunker import Chunker
+from sif.indexing.parser import MarkdownParser
+from sif.indexing.scanner import FileScanner
+from sif.utils.logging import get_logger
+from sif.utils.progress import ProgressTracker, RichProgressTracker
+
+
+logger = get_logger(__name__)
+
+
+class IndexStatus(Enum):
+    """Status of indexing operation."""
+
+    PENDING = auto()
+    INDEXING = auto()
+    COMPLETED = auto()
+    FAILED = auto()
+    SKIPPED = auto()
+
+
+@dataclass
+class IndexResult:
+    """Result of indexing operation."""
+
+    files_processed: int
+    files_added: int
+    files_updated: int
+    files_skipped: int
+    files_failed: int
+    chunks_created: int
+    errors: list[str]
+
+
+class DocumentIndexer:
+    """Orchestrates the document indexing process.
+
+    Coordinates file scanning, parsing, chunking, embedding generation,
+    and database storage for document indexing.
+    """
+
+    def __init__(  # noqa: PLR0913
+        self,
+        scanner: FileScanner,
+        parser: MarkdownParser,
+        chunker: Chunker,
+        embedding_manager: EmbeddingManager,
+        repository: DocumentRepository,
+        progress_tracker: ProgressTracker | None = None,
+    ) -> None:
+        """Initialize document indexer.
+
+        Args:
+            scanner: File scanner
+            parser: Markdown parser
+            chunker: Document chunker
+            embedding_manager: Embedding manager
+            repository: Document repository
+            progress_tracker: Progress tracker
+        """
+        self._scanner = scanner
+        self._parser = parser
+        self._chunker = chunker
+        self._embedding_manager = embedding_manager
+        self._repository = repository
+        self._progress = progress_tracker or RichProgressTracker()
+
+    def index_collection(self, collection: Collection) -> IndexResult:
+        """Index all documents in a collection.
+
+        Args:
+            collection: Collection to index
+
+        Returns:
+            Indexing result
+        """
+        logger.info(f"Indexing collection: {collection.name}")
+
+        # Scan for files
+        scan_result = self._scanner.scan(
+            Path(collection.path),
+            pattern=collection.pattern,
+            ignore_patterns=collection.ignore_patterns,
+        )
+
+        logger.info(f"Found {scan_result.file_count} files to index")
+
+        # Index files
+        self._progress.start("Indexing documents", total=scan_result.file_count)
+
+        result = IndexResult(
+            files_processed=0,
+            files_added=0,
+            files_updated=0,
+            files_skipped=0,
+            files_failed=0,
+            chunks_created=0,
+            errors=[],
+        )
+
+        for file_path in scan_result.files:
+            try:
+                file_result = self._index_file(file_path, collection.id)
+
+                result.files_processed += 1
+                if file_result == IndexStatus.ADDED:
+                    result.files_added += 1
+                elif file_result == IndexStatus.COMPLETED:
+                    result.files_updated += 1
+                elif file_result == IndexStatus.SKIPPED:
+                    result.files_skipped += 1
+
+                self._progress.update()
+
+            except Exception as e:
+                logger.error(f"Error indexing {file_path}: {e}")
+                result.files_failed += 1
+                result.errors.append(f"{file_path}: {e!s}")
+
+        self._progress.finish()
+
+        # Update collection stats
+        collection.document_count = self._repository.list_by_collection(collection.id).__len__()
+        collection.mark_indexed()
+
+        logger.info(
+            f"Indexing complete: {result.files_added} added, "
+            f"{result.files_updated} updated, {result.files_failed} failed"
+        )
+
+        return result
+
+    def _index_file(
+        self,
+        file_path: Path,
+        collection_id: str,
+    ) -> IndexStatus:
+        """Index a single file.
+
+        Args:
+            file_path: Path to file
+            collection_id: Collection ID
+
+        Returns:
+            Index status
+        """
+        import hashlib  # noqa: PLC0415
+        import uuid  # noqa: PLC0415
+
+        # Calculate checksum
+        with open(file_path, "rb") as f:
+            checksum = hashlib.sha256(f.read()).hexdigest()
+
+        # Check if file already indexed
+        existing = self._repository.get_by_path(str(file_path), collection_id)
+        if existing:
+            if existing.checksum == checksum:
+                logger.debug(f"File unchanged, skipping: {file_path}")
+                return IndexStatus.SKIPPED
+
+        # Parse file
+        parse_result = self._parser.parse(file_path)
+
+        # Create document
+        document = Document(
+            id=str(uuid.uuid4()) if not existing else existing.id,
+            collection_id=collection_id,
+            path=str(file_path),
+            content=parse_result.content,
+            title=parse_result.title or file_path.name,
+            checksum=checksum,
+            file_size=file_path.stat().st_size,
+            mtime=file_path.stat().st_mtime,
+            metadata=parse_result.metadata,
+        )
+
+        # Chunk document
+        chunks = self._chunker.chunk(parse_result.content)
+
+        # Generate embeddings for chunks
+        if chunks:
+            chunk_texts = [c.content for c in chunks]
+            embedding_response = self._embedding_manager.embed(chunk_texts)
+
+            # Create document chunks with embeddings
+            for _i, (chunk, embedding) in enumerate(zip(chunks, embedding_response.embeddings)):
+                doc_chunk = DocumentChunk(
+                    id=str(uuid.uuid4()),
+                    document_id=document.id,
+                    content=chunk.content,
+                    start_pos=chunk.start_pos,
+                    end_pos=chunk.end_pos,
+                    token_count=chunk.token_count,
+                    embedding=embedding,
+                )
+                document.add_chunk(doc_chunk)
+
+        # Save to repository
+        if existing:
+            self._repository.update(document)
+            return IndexStatus.COMPLETED
+        self._repository.create(document)
+        return IndexStatus.ADDED
+
+    def reindex_collection(self, collection: Collection) -> IndexResult:
+        """Reindex all documents in a collection.
+
+        Args:
+            collection: Collection to reindex
+
+        Returns:
+            Indexing result
+        """
+        # Clear existing documents
+        self._repository.delete_by_collection(collection.id)
+
+        # Reindex
+        return self.index_collection(collection)