ragup 0.1.0__py3-none-any.whl

ragup/embeddings/models.py

@@ -0,0 +1,22 @@
+"""
+Embedding model registry for RagUp.
+"""
+
+DEFAULT_EMBEDDING_MODEL = "all-MiniLM-L6-v2"
+
+SUPPORTED_EMBEDDING_MODELS = {
+    "all-MiniLM-L6-v2": "sentence-transformers/all-MiniLM-L6-v2",
+    "bge-small-en": "BAAI/bge-small-en-v1.5",
+    "bge-base-en": "BAAI/bge-base-en-v1.5",
+}
+
+
+def get_model_name(model_name: str) -> str:
+    """
+    Resolve a model alias to its Hugging Face identifier.
+    """
+
+    return SUPPORTED_EMBEDDING_MODELS.get(
+        model_name,
+        model_name,
+    )

ragup/exceptions.py

@@ -0,0 +1,57 @@
+"""
+Custom exceptions used throughout RagUp.
+"""
+
+
+class RagUpError(Exception):
+    """Base exception for all RagUp errors."""
+
+    pass
+
+
+class UnsupportedFileTypeError(RagUpError):
+    """Raised when the provided file type is not supported."""
+
+    pass
+
+
+class FileNotIndexedError(RagUpError):
+    """Raised when ask() or search() is called before ragup()."""
+
+    pass
+
+
+class InvalidAPIKeyError(RagUpError):
+    """Raised when an invalid API key is provided."""
+
+    pass
+
+
+class VectorDBError(RagUpError):
+    """Raised when a vector database operation fails."""
+
+    pass
+
+
+class EmbeddingError(RagUpError):
+    """Raised when embedding generation fails."""
+
+    pass
+
+
+class LLMError(RagUpError):
+    """Raised when the LLM provider returns an error."""
+
+    pass
+
+
+class LoaderError(RagUpError):
+    """Raised when a document cannot be loaded."""
+
+    pass
+
+
+class ChunkingError(RagUpError):
+    """Raised when chunk generation fails."""
+
+    pass

ragup/llm/__init__.py

@@ -0,0 +1,9 @@
+"""
+LLM module.
+"""
+
+from .answer import generate_answer
+
+__all__ = [
+    "generate_answer",
+]

ragup/llm/answer.py

@@ -0,0 +1,54 @@
+"""
+Gemini answer generation for RagUp.
+"""
+
+from __future__ import annotations
+
+from litellm import completion
+
+from .prompt import build_prompt
+from ..exceptions import LLMError
+
+
+DEFAULT_GEMINI_MODEL = "gemini/gemini-2.5-flash"
+
+
+def generate_answer(
+    question: str,
+    context: list[dict],
+    api_key: str,
+) -> str:
+    """
+    Generate an answer using Gemini.
+    """
+
+    if not api_key:
+        raise LLMError(
+            "A Gemini API key is required."
+        )
+
+    prompt = build_prompt(
+        question=question,
+        context=context,
+    )
+
+    try:
+
+        response = completion(
+            model=DEFAULT_GEMINI_MODEL,
+            api_key=api_key,
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt,
+                }
+            ],
+            temperature=0.0,
+        )
+
+        return response.choices[0].message.content
+
+    except Exception as exc:
+        raise LLMError(
+            "Failed to generate an answer from Gemini."
+        ) from exc

ragup/llm/prompt.py

@@ -0,0 +1,36 @@
+"""
+Prompt builder for RagUp.
+"""
+
+from __future__ import annotations
+
+
+def build_prompt(
+    question: str,
+    context: list[dict],
+) -> str:
+    """
+    Build the prompt sent to Gemini.
+    """
+
+    context_text = "\n\n".join(
+        item["text"] for item in context
+    )
+
+    return f"""
+You are a helpful AI assistant.
+
+Answer the user's question ONLY using the provided context.
+
+If the answer cannot be found in the context, reply:
+
+"I could not find the answer in the provided documents."
+
+Context:
+{context_text}
+
+Question:
+{question}
+
+Answer:
+""".strip()

ragup/loaders/__init__.py

@@ -0,0 +1,62 @@
+"""
+Loader factory for RagUp.
+"""
+
+from pathlib import Path
+
+from .base import BaseLoader
+from .pdf_loader import PDFLoader
+from .txt_loader import TXTLoader
+from .json_loader import JSONLoader
+
+from ..exceptions import UnsupportedFileTypeError
+
+
+_LOADER_MAP = {
+    ".pdf": PDFLoader,
+    ".txt": TXTLoader,
+    ".json": JSONLoader,
+}
+
+
+def get_loader(file_path: str | Path) -> BaseLoader:
+    """
+    Return the appropriate loader for the given file.
+
+    Parameters
+    ----------
+    file_path : str | Path
+        Path to the document.
+
+    Returns
+    -------
+    BaseLoader
+        Loader instance.
+
+    Raises
+    ------
+    UnsupportedFileTypeError
+        If the file extension is not supported.
+    """
+
+    path = Path(file_path)
+
+    extension = path.suffix.lower()
+
+    loader = _LOADER_MAP.get(extension)
+
+    if loader is None:
+        raise UnsupportedFileTypeError(
+            f"Unsupported file type: {extension}"
+        )
+
+    return loader(path)
+
+
+__all__ = [
+    "BaseLoader",
+    "PDFLoader",
+    "TXTLoader",
+    "JSONLoader",
+    "get_loader",
+]

ragup/loaders/base.py

@@ -0,0 +1,31 @@
+"""
+Base loader interface for all document loaders.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+
+class BaseLoader(ABC):
+    """
+    Base class for all document loaders.
+
+    Every loader should return the document as plain text.
+    """
+
+    def __init__(self, file_path: str | Path):
+        self.file_path = Path(file_path)
+
+    @abstractmethod
+    def load(self) -> str:
+        """
+        Load the document and return its text.
+
+        Returns
+        -------
+        str
+            Extracted document text.
+        """
+        raise NotImplementedError

ragup/loaders/json_loader.py

@@ -0,0 +1,44 @@
+"""
+JSON document loader.
+"""
+
+from __future__ import annotations
+
+import json
+
+from .base import BaseLoader
+from ..exceptions import LoaderError
+
+
+class JSONLoader(BaseLoader):
+    """
+    Loader for JSON documents.
+
+    The JSON content is converted into a formatted string
+    suitable for chunking and embedding.
+    """
+
+    def load(self) -> str:
+        """
+        Load a JSON file and return it as formatted text.
+
+        Returns
+        -------
+        str
+            JSON represented as readable text.
+        """
+
+        try:
+            with open(self.file_path, "r", encoding="utf-8") as f:
+                data = json.load(f)
+
+            return json.dumps(
+                data,
+                indent=2,
+                ensure_ascii=False
+            ).strip()
+
+        except Exception as exc:
+            raise LoaderError(
+                f"Failed to load JSON: {self.file_path}"
+            ) from exc

ragup/loaders/pdf_loader.py

@@ -0,0 +1,43 @@
+"""
+PDF document loader.
+"""
+
+from __future__ import annotations
+
+import fitz  # PyMuPDF
+
+from .base import BaseLoader
+from ..exceptions import LoaderError
+
+
+class PDFLoader(BaseLoader):
+    """
+    Loader for PDF documents.
+    """
+
+    def load(self) -> str:
+        """
+        Load a PDF and return its text.
+
+        Returns
+        -------
+        str
+            Extracted text from all pages.
+        """
+
+        try:
+            document = fitz.open(self.file_path)
+
+            pages = []
+
+            for page in document:
+                pages.append(page.get_text())
+
+            document.close()
+
+            return "\n".join(pages).strip()
+
+        except Exception as exc:
+            raise LoaderError(
+                f"Failed to load PDF: {self.file_path}"
+            ) from exc

ragup/loaders/txt_loader.py

@@ -0,0 +1,45 @@
+"""
+TXT document loader.
+"""
+
+from __future__ import annotations
+
+from .base import BaseLoader
+from ..exceptions import LoaderError
+
+
+class TXTLoader(BaseLoader):
+    """
+    Loader for plain text documents.
+    """
+
+    def load(self) -> str:
+        """
+        Load a text file and return its contents.
+
+        Returns
+        -------
+        str
+            Text file contents.
+        """
+
+        try:
+            return self.file_path.read_text(
+                encoding="utf-8"
+            ).strip()
+
+        except UnicodeDecodeError:
+            try:
+                return self.file_path.read_text(
+                    encoding="latin-1"
+                ).strip()
+
+            except Exception as exc:
+                raise LoaderError(
+                    f"Failed to load TXT: {self.file_path}"
+                ) from exc
+
+        except Exception as exc:
+            raise LoaderError(
+                f"Failed to load TXT: {self.file_path}"
+            ) from exc

ragup/retrieval/__init__.py

@@ -0,0 +1,11 @@
+"""
+Retrieval module.
+"""
+
+from .retriever import Retriever
+from .search import format_search_results
+
+__all__ = [
+    "Retriever",
+    "format_search_results",
+]

ragup/retrieval/retriever.py

@@ -0,0 +1,41 @@
+"""
+Semantic retriever for RagUp.
+"""
+
+from __future__ import annotations
+
+from ..embeddings import Embedder
+from ..vectordb import VectorDBManager
+
+
+class Retriever:
+    """
+    Handles semantic retrieval.
+    """
+
+    def __init__(
+        self,
+        embedder: Embedder,
+        vectordb: VectorDBManager,
+    ):
+
+        self.embedder = embedder
+        self.vectordb = vectordb
+
+    def retrieve(
+        self,
+        query: str,
+        top_k: int = 5,
+    ) -> dict:
+        """
+        Retrieve relevant document chunks.
+        """
+
+        query_embedding = self.embedder.embed_query(
+            query
+        )
+
+        return self.vectordb.search(
+            embedding=query_embedding,
+            top_k=top_k,
+        )

ragup/retrieval/search.py

@@ -0,0 +1,55 @@
+"""
+Search result formatter for RagUp.
+"""
+
+from __future__ import annotations
+
+
+def format_search_results(results: dict) -> list[dict]:
+    """
+    Convert the raw ChromaDB response into a clean list of results.
+
+    Parameters
+    ----------
+    results : dict
+        Raw response returned by ChromaDB.
+
+    Returns
+    -------
+    list[dict]
+    """
+
+    if not results:
+        return []
+
+    documents = results.get("documents", [[]])[0]
+    metadatas = results.get("metadatas", [[]])[0]
+    distances = results.get("distances", [[]])[0]
+    ids = results.get("ids", [[]])[0]
+
+    formatted = []
+
+    for idx, document in enumerate(documents):
+
+        metadata = (
+            metadatas[idx]
+            if idx < len(metadatas)
+            else {}
+        )
+
+        distance = (
+            distances[idx]
+            if idx < len(distances)
+            else None
+        )
+
+        result = {
+            "id": ids[idx] if idx < len(ids) else None,
+            "text": document,
+            "metadata": metadata,
+            "distance": distance,
+        }
+
+        formatted.append(result)
+
+    return formatted

ragup/utils/logger.py

@@ -0,0 +1,47 @@
+"""
+Logging utilities for RagUp.
+"""
+
+import logging
+from typing import Optional
+
+
+_LOGGER_NAME = "ragup"
+
+
+def get_logger(verbose: bool = False) -> logging.Logger:
+    """
+    Returns a configured logger instance.
+
+    Parameters
+    ----------
+    verbose : bool
+        Enable DEBUG logging.
+
+    Returns
+    -------
+    logging.Logger
+    """
+
+    logger = logging.getLogger(_LOGGER_NAME)
+
+    if logger.handlers:
+        return logger
+
+    level = logging.DEBUG if verbose else logging.INFO
+
+    logger.setLevel(level)
+
+    handler = logging.StreamHandler()
+
+    formatter = logging.Formatter(
+        "[%(levelname)s] %(message)s"
+    )
+
+    handler.setFormatter(formatter)
+
+    logger.addHandler(handler)
+
+    logger.propagate = False
+
+    return logger

ragup/vectordb/__init__.py

@@ -0,0 +1,9 @@
+"""
+Vector database module.
+"""
+
+from .manager import VectorDBManager
+
+__all__ = [
+    "VectorDBManager",
+]

ragup/vectordb/chroma.py

@@ -0,0 +1,101 @@
+"""
+ChromaDB wrapper for RagUp.
+"""
+
+from __future__ import annotations
+
+import chromadb
+from chromadb.config import Settings
+
+from ..constants import (
+    DEFAULT_COLLECTION_NAME,
+    DEFAULT_VECTOR_DB_DIR,
+)
+from ..exceptions import VectorDBError
+
+
+class ChromaVectorDB:
+    """
+    Wrapper around ChromaDB.
+    """
+
+    def __init__(
+        self,
+        db_path: str = str(DEFAULT_VECTOR_DB_DIR),
+        collection_name: str = DEFAULT_COLLECTION_NAME,
+    ):
+
+        self.client = chromadb.PersistentClient(
+            path=db_path,
+            settings=Settings(
+                anonymized_telemetry=False
+            ),
+        )
+
+        self.collection = self.client.get_or_create_collection(
+            name=collection_name
+        )
+
+    def add(
+        self,
+        ids: list[str],
+        documents: list[str],
+        embeddings: list[list[float]],
+        metadatas: list[dict],
+    ) -> None:
+        """
+        Store document chunks.
+        """
+
+        try:
+
+            self.collection.add(
+                ids=ids,
+                documents=documents,
+                embeddings=embeddings,
+                metadatas=metadatas,
+            )
+
+        except Exception as exc:
+            raise VectorDBError(
+                "Failed to add documents to ChromaDB."
+            ) from exc
+
+    def search(
+        self,
+        embedding: list[float],
+        top_k: int = 5,
+    ) -> dict:
+        """
+        Semantic search.
+        """
+
+        try:
+
+            return self.collection.query(
+                query_embeddings=[embedding],
+                n_results=top_k,
+            )
+
+        except Exception as exc:
+            raise VectorDBError(
+                "Failed to query ChromaDB."
+            ) from exc
+
+    def count(self) -> int:
+
+        return self.collection.count()
+
+    def reset(self) -> None:
+
+        try:
+
+            ids = self.collection.get()["ids"]
+
+            if ids:
+                self.collection.delete(ids=ids)
+
+        except Exception as exc:
+            raise VectorDBError(
+                "Failed to reset ChromaDB."
+            ) from exc