piragi 0.1.0__py3-none-any.whl

ragi/embeddings.py

@@ -0,0 +1,150 @@
+"""Embedding generation using local or remote models."""
+
+import os
+from typing import List, Optional
+
+from .types import Chunk
+
+
+class EmbeddingGenerator:
+    """Generate embeddings using local sentence-transformers or remote API."""
+
+    def __init__(
+        self,
+        model: str = "nvidia/llama-embed-nemotron-8b",
+        device: str | None = None,
+        base_url: str | None = None,
+        api_key: str | None = None,
+    ) -> None:
+        """
+        Initialize the embedding generator.
+
+        Args:
+            model: Embedding model to use (default: nvidia/llama-embed-nemotron-8b)
+            device: Device to run on ('cuda', 'cpu', or None for auto-detect) - only for local models
+            base_url: Optional API base URL for remote embeddings (e.g., https://api.openai.com/v1)
+            api_key: Optional API key for remote embeddings
+        """
+        self.model_name = model
+        self.base_url = base_url
+        self.api_key = api_key
+        self.use_remote = base_url is not None
+
+        if self.use_remote:
+            # Use OpenAI-compatible API client
+            from openai import OpenAI
+
+            if self.api_key is None:
+                self.api_key = os.getenv("EMBEDDING_API_KEY", "not-needed")
+
+            self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)
+            self.model = None
+        else:
+            # Use local sentence-transformers
+            from sentence_transformers import SentenceTransformer
+
+            self.model = SentenceTransformer(
+                model,
+                trust_remote_code=True,
+                device=device,
+            )
+            self.client = None
+
+    def embed_chunks(self, chunks: List[Chunk]) -> List[Chunk]:
+        """
+        Generate embeddings for a list of chunks.
+
+        Args:
+            chunks: List of chunks to embed
+
+        Returns:
+            Chunks with embeddings added
+        """
+        if not chunks:
+            return chunks
+
+        # Extract texts
+        texts = [chunk.text for chunk in chunks]
+
+        # Generate embeddings in batch (documents don't need instruction prefix)
+        embeddings = self._generate_embeddings(texts)
+
+        # Add embeddings to chunks
+        for chunk, embedding in zip(chunks, embeddings):
+            chunk.embedding = embedding.tolist()
+
+        return chunks
+
+    def _generate_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """
+        Generate embeddings for a list of texts.
+
+        Args:
+            texts: List of text strings
+
+        Returns:
+            List of embedding vectors
+        """
+        try:
+            if self.use_remote:
+                # Use OpenAI-compatible API
+                response = self.client.embeddings.create(
+                    input=texts,
+                    model=self.model_name,
+                )
+                return [item.embedding for item in response.data]
+            else:
+                # Use local sentence-transformers
+                # Use encode_document for document chunks if available
+                if hasattr(self.model, "encode_document"):
+                    embeddings = self.model.encode_document(texts)
+                else:
+                    embeddings = self.model.encode(texts)
+                return embeddings
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate embeddings: {e}")
+
+    def embed_query(self, query: str, task_instruction: str | None = None) -> List[float]:
+        """
+        Generate embedding for a single query.
+
+        Args:
+            query: Query text
+            task_instruction: Optional task instruction for query
+                (e.g., "Retrieve relevant documents for this question")
+
+        Returns:
+            Embedding vector
+        """
+        try:
+            if self.use_remote:
+                # Use OpenAI-compatible API
+                query_text = query
+                if task_instruction:
+                    query_text = f"{task_instruction}\n{query}"
+
+                response = self.client.embeddings.create(
+                    input=query_text,
+                    model=self.model_name,
+                )
+                return response.data[0].embedding
+            else:
+                # Use local sentence-transformers
+                # Use encode_query for search queries if available
+                if hasattr(self.model, "encode_query"):
+                    if task_instruction:
+                        query_with_instruction = f"{task_instruction}\n{query}"
+                        embedding = self.model.encode_query(query_with_instruction)
+                    else:
+                        embedding = self.model.encode_query(query)
+                else:
+                    query_text = query
+                    if task_instruction:
+                        query_text = f"{task_instruction}\n{query}"
+                    embedding = self.model.encode(query_text)
+
+                return embedding.tolist()
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate query embedding: {e}")

ragi/loader.py

@@ -0,0 +1,131 @@
+"""Document loading using markitdown."""
+
+import glob
+import os
+from pathlib import Path
+from typing import List, Union
+from urllib.parse import urlparse
+
+from markitdown import MarkItDown
+
+from .types import Document
+
+
+class DocumentLoader:
+    """Load documents from various sources using markitdown."""
+
+    def __init__(self) -> None:
+        """Initialize the document loader."""
+        self.converter = MarkItDown()
+
+    def load(self, source: Union[str, List[str]]) -> List[Document]:
+        """
+        Load documents from file paths, URLs, or glob patterns.
+
+        Args:
+            source: Single path/URL, list of paths/URLs, or glob pattern
+
+        Returns:
+            List of loaded documents
+        """
+        if isinstance(source, str):
+            sources = [source]
+        else:
+            sources = source
+
+        documents = []
+        for src in sources:
+            documents.extend(self._load_single(src))
+
+        return documents
+
+    def _load_single(self, source: str) -> List[Document]:
+        """Load from a single source (file, URL, or glob pattern)."""
+        # Check if it's a URL
+        if self._is_url(source):
+            return [self._load_url(source)]
+
+        # Check if it's a glob pattern
+        if any(char in source for char in ["*", "?", "[", "]"]):
+            return self._load_glob(source)
+
+        # Single file
+        if os.path.isfile(source):
+            return [self._load_file(source)]
+
+        # Directory - load all files
+        if os.path.isdir(source):
+            return self._load_directory(source)
+
+        raise ValueError(f"Invalid source: {source}")
+
+    def _is_url(self, source: str) -> bool:
+        """Check if source is a URL."""
+        try:
+            result = urlparse(source)
+            return all([result.scheme, result.netloc])
+        except Exception:
+            return False
+
+    def _load_file(self, file_path: str) -> Document:
+        """Load a single file."""
+        try:
+            result = self.converter.convert(file_path)
+            content = result.text_content
+
+            # Extract metadata
+            metadata = {
+                "filename": os.path.basename(file_path),
+                "file_type": Path(file_path).suffix.lstrip("."),
+                "file_path": os.path.abspath(file_path),
+            }
+
+            return Document(content=content, source=file_path, metadata=metadata)
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to load file {file_path}: {e}")
+
+    def _load_url(self, url: str) -> Document:
+        """Load content from a URL."""
+        try:
+            result = self.converter.convert(url)
+            content = result.text_content
+
+            metadata = {
+                "url": url,
+                "source_type": "url",
+            }
+
+            return Document(content=content, source=url, metadata=metadata)
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to load URL {url}: {e}")
+
+    def _load_glob(self, pattern: str) -> List[Document]:
+        """Load files matching a glob pattern."""
+        files = glob.glob(pattern, recursive=True)
+        files = [f for f in files if os.path.isfile(f)]
+
+        if not files:
+            raise ValueError(f"No files found matching pattern: {pattern}")
+
+        return [self._load_file(f) for f in files]
+
+    def _load_directory(self, directory: str) -> List[Document]:
+        """Load all files from a directory recursively."""
+        pattern = os.path.join(directory, "**", "*")
+        files = glob.glob(pattern, recursive=True)
+        files = [f for f in files if os.path.isfile(f)]
+
+        if not files:
+            raise ValueError(f"No files found in directory: {directory}")
+
+        documents = []
+        for f in files:
+            try:
+                documents.append(self._load_file(f))
+            except Exception:
+                # Skip files that can't be processed
+                continue
+
+        return documents

ragi/retrieval.py

@@ -0,0 +1,125 @@
+"""Retrieval and answer generation using OpenAI-compatible APIs."""
+
+import os
+from typing import List, Optional
+
+from openai import OpenAI
+
+from .types import Answer, Citation
+
+
+class Retriever:
+    """Generate answers from retrieved chunks using OpenAI-compatible LLM APIs."""
+
+    def __init__(
+        self,
+        model: str = "llama3.2",
+        api_key: str | None = None,
+        base_url: str | None = None,
+    ) -> None:
+        """
+        Initialize the retriever.
+
+        Args:
+            model: Model name to use (default: llama3.2 for Ollama)
+            api_key: API key (optional for local models like Ollama)
+            base_url: Base URL for OpenAI-compatible API (e.g., http://localhost:11434/v1 for Ollama)
+        """
+        self.model = model
+
+        # Default to Ollama if no base_url provided
+        if base_url is None:
+            base_url = os.getenv("LLM_BASE_URL", "http://localhost:11434/v1")
+
+        # API key is optional for local models
+        if api_key is None:
+            api_key = os.getenv("LLM_API_KEY", "not-needed")
+
+        self.client = OpenAI(
+            api_key=api_key,
+            base_url=base_url,
+        )
+
+    def generate_answer(
+        self,
+        query: str,
+        citations: List[Citation],
+        system_prompt: Optional[str] = None,
+    ) -> Answer:
+        """
+        Generate an answer from retrieved citations.
+
+        Args:
+            query: User's question
+            citations: Retrieved citations
+            system_prompt: Optional custom system prompt
+
+        Returns:
+            Answer with citations
+        """
+        if not citations:
+            return Answer(
+                text="I couldn't find any relevant information to answer your question.",
+                citations=[],
+                query=query,
+            )
+
+        # Build context from citations
+        context = self._build_context(citations)
+
+        # Generate answer
+        answer_text = self._generate_with_llm(query, context, system_prompt)
+
+        return Answer(
+            text=answer_text,
+            citations=citations,
+            query=query,
+        )
+
+    def _build_context(self, citations: List[Citation]) -> str:
+        """Build context string from citations."""
+        context_parts = []
+
+        for i, citation in enumerate(citations, 1):
+            source_info = f"Source {i} ({citation.source}):"
+            context_parts.append(f"{source_info}\n{citation.chunk}\n")
+
+        return "\n".join(context_parts)
+
+    def _generate_with_llm(
+        self,
+        query: str,
+        context: str,
+        system_prompt: Optional[str] = None,
+    ) -> str:
+        """Generate answer using OpenAI-compatible API."""
+        if system_prompt is None:
+            system_prompt = (
+                "You are a helpful assistant that answers questions based on the provided context. "
+                "Always cite your sources by mentioning which source number you're referring to. "
+                "If the context doesn't contain enough information to answer the question, say so. "
+                "Be concise and accurate."
+            )
+
+        user_prompt = f"""Context from documents:
+
+{context}
+
+Question: {query}
+
+Please answer the question based on the context provided above. Cite your sources."""
+
+        try:
+            response = self.client.chat.completions.create(
+                model=self.model,
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt},
+                ],
+                temperature=0.3,
+            )
+
+            return response.choices[0].message.content or ""
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate answer: {e}")

ragi/store.py

@@ -0,0 +1,177 @@
+"""Vector store using LanceDB."""
+
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import lancedb
+from lancedb.pydantic import LanceModel, Vector
+
+from .types import Chunk, Citation
+
+
+class ChunkModel(LanceModel):
+    """LanceDB schema for chunks."""
+
+    text: str
+    source: str
+    chunk_index: int
+    metadata: Dict[str, Any]
+    vector: Vector(4096)  # nvidia/llama-embed-nemotron-8b dimension
+
+
+class SourceMetadata(LanceModel):
+    """Metadata for tracking source changes."""
+
+    source: str  # File path or URL
+    last_checked: float  # Unix timestamp
+    content_hash: str  # SHA256 hash of content
+    mtime: Optional[float] = None  # File modification time (for files)
+    etag: Optional[str] = None  # HTTP ETag (for URLs)
+    last_modified: Optional[str] = None  # HTTP Last-Modified (for URLs)
+    check_interval: float = 300.0  # Seconds between checks (default: 5 min)
+
+
+class VectorStore:
+    """Vector store using LanceDB."""
+
+    def __init__(self, persist_dir: str = ".ragi") -> None:
+        """
+        Initialize the vector store.
+
+        Args:
+            persist_dir: Directory to persist the vector database
+        """
+        self.persist_dir = persist_dir
+        Path(persist_dir).mkdir(parents=True, exist_ok=True)
+
+        self.db = lancedb.connect(persist_dir)
+        self.table_name = "chunks"
+        self.metadata_table_name = "source_metadata"
+        self.table: Optional[Any] = None
+        self.metadata_table: Optional[Any] = None
+
+        # Initialize tables if they exist
+        if self.table_name in self.db.table_names():
+            self.table = self.db.open_table(self.table_name)
+
+        if self.metadata_table_name in self.db.table_names():
+            self.metadata_table = self.db.open_table(self.metadata_table_name)
+
+    def add_chunks(self, chunks: List[Chunk]) -> None:
+        """
+        Add chunks to the vector store.
+
+        Args:
+            chunks: List of chunks with embeddings
+        """
+        if not chunks:
+            return
+
+        # Validate embeddings
+        for chunk in chunks:
+            if chunk.embedding is None:
+                raise ValueError("All chunks must have embeddings before adding to store")
+
+        # Convert chunks to LanceDB format
+        data = [
+            {
+                "text": chunk.text,
+                "source": chunk.source,
+                "chunk_index": chunk.chunk_index,
+                "metadata": chunk.metadata,
+                "vector": chunk.embedding,
+            }
+            for chunk in chunks
+        ]
+
+        # Create or update table
+        if self.table is None:
+            self.table = self.db.create_table(self.table_name, data=data, mode="overwrite")
+        else:
+            self.table.add(data)
+
+    def search(
+        self,
+        query_embedding: List[float],
+        top_k: int = 5,
+        filters: Optional[Dict[str, Any]] = None,
+    ) -> List[Citation]:
+        """
+        Search for similar chunks.
+
+        Args:
+            query_embedding: Query vector
+            top_k: Number of results to return
+            filters: Optional metadata filters
+
+        Returns:
+            List of citations
+        """
+        if self.table is None:
+            return []
+
+        # Build query
+        query = self.table.search(query_embedding).limit(top_k)
+
+        # Apply filters if provided
+        if filters:
+            filter_conditions = []
+            for key, value in filters.items():
+                # Handle nested metadata filters
+                filter_conditions.append(f"metadata['{key}'] = '{value}'")
+
+            if filter_conditions:
+                query = query.where(" AND ".join(filter_conditions))
+
+        # Execute search
+        results = query.to_list()
+
+        # Convert to citations
+        citations = []
+        for result in results:
+            citation = Citation(
+                source=result["source"],
+                chunk=result["text"],
+                score=1.0 - result["_distance"],  # Convert distance to similarity score
+                metadata=result["metadata"],
+            )
+            citations.append(citation)
+
+        return citations
+
+    def count(self) -> int:
+        """Return the number of chunks in the store."""
+        if self.table is None:
+            return 0
+        return self.table.count_rows()
+
+    def delete_by_source(self, source: str) -> int:
+        """
+        Delete all chunks from a specific source.
+
+        Args:
+            source: Source file path or URL to delete
+
+        Returns:
+            Number of chunks deleted
+        """
+        if self.table is None:
+            return 0
+
+        # Count chunks before deletion
+        count_before = self.table.count_rows()
+
+        # Delete chunks matching the source
+        self.table.delete(f"source = '{source}'")
+
+        # Count chunks after deletion
+        count_after = self.table.count_rows()
+
+        return count_before - count_after
+
+    def clear(self) -> None:
+        """Clear all data from the store."""
+        if self.table_name in self.db.table_names():
+            self.db.drop_table(self.table_name)
+            self.table = None

ragi/types.py

@@ -0,0 +1,54 @@
+"""Type definitions for Ragi."""
+
+from typing import Any, Dict, List, Optional
+from pydantic import BaseModel, Field
+
+
+class Citation(BaseModel):
+    """A single citation with source information."""
+
+    source: str = Field(description="Source file path or URL")
+    chunk: str = Field(description="The actual text chunk")
+    score: float = Field(description="Relevance score (0-1)")
+    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
+
+    @property
+    def preview(self) -> str:
+        """Return a preview of the chunk (first 100 chars)."""
+        return self.chunk[:100] + "..." if len(self.chunk) > 100 else self.chunk
+
+
+class Answer(BaseModel):
+    """Answer with citations from the RAG system."""
+
+    text: str = Field(description="The generated answer")
+    citations: List[Citation] = Field(default_factory=list, description="Source citations")
+    query: str = Field(description="Original query")
+
+    def __str__(self) -> str:
+        """Return the answer text."""
+        return self.text
+
+    def __repr__(self) -> str:
+        """Return detailed representation."""
+        return f"Answer(text='{self.text[:50]}...', citations={len(self.citations)})"
+
+
+class Document(BaseModel):
+    """Internal document representation."""
+
+    content: str = Field(description="Document content in markdown")
+    source: str = Field(description="Source file path or URL")
+    metadata: Dict[str, Any] = Field(default_factory=dict, description="Document metadata")
+    content_hash: Optional[str] = Field(default=None, description="Hash of content for change detection")
+    last_modified: Optional[float] = Field(default=None, description="Last modification timestamp")
+
+
+class Chunk(BaseModel):
+    """A chunk of a document with metadata."""
+
+    text: str = Field(description="Chunk text")
+    source: str = Field(description="Source document")
+    chunk_index: int = Field(description="Index of chunk in document")
+    metadata: Dict[str, Any] = Field(default_factory=dict, description="Chunk metadata")
+    embedding: Optional[List[float]] = Field(default=None, description="Vector embedding")