ezrag-engine 0.1.0__py3-none-any.whl

ezrag/__init__.py

@@ -0,0 +1,7 @@
+"""ezrag — a 3-line RAG micro-framework for Hugging Face Spaces."""
+
+from .engine import EZRAG, Chunk
+
+__all__ = ["EZRAG", "Chunk"]
+__version__ = "0.1.0"
+

ezrag/engine.py

@@ -0,0 +1,349 @@
+"""ezrag.engine — hidden orchestration of the EZRAG pipeline.
+
+Everything heavy lives here and stays invisible to developers:
+text chunking, embedding generation, in-memory similarity search,
+context bundling and LLM query execution. Consumers only touch the
+public surface exposed by ``ezrag.EZRAG``.
+"""
+
+from __future__ import annotations
+
+import logging
+import os as _os
+import re
+from dataclasses import dataclass, field
+from pathlib import Path as _Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+try:
+    import numpy as np
+except ImportError:  # pragma: no cover
+    np = None  # type: ignore
+
+try:
+    from huggingface_hub import InferenceClient
+except ImportError:  # pragma: no cover
+    InferenceClient = None  # type: ignore
+
+try:
+    import litellm
+except ImportError:  # pragma: no cover
+    litellm = None  # type: ignore
+
+logger = logging.getLogger("ezrag")
+
+__all__ = ["EZRAG", "Chunk"]
+
+DEFAULT_EMBEDDING_MODEL = "all-MiniLM-L6-v2"
+DEFAULT_LLM = _os.getenv("EZRAG_LLM", "mistralai/Mistral-7B-Instruct-v0.2")
+
+_SENTENCE_BOUNDARY = re.compile(r"(?<=[.!?。！？])\s+")
+_WHITESPACE = re.compile(r"\s+")
+
+_SYSTEM_PROMPT = (
+    "You are a precise, grounded assistant. Answer ONLY from the provided "
+    "CONTEXT delimited by <context></context>. Cite every fact with source "
+    "tags like [1]. If the CONTEXT is empty or clearly cannot answer the "
+    "question, reply exactly with 'I do not have enough information to "
+    "answer.' Never invent facts."
+)
+
+
+def _normalize_text(text: str) -> str:
+    return _WHITESPACE.sub(" ", text).strip()
+
+
+@dataclass(frozen=True)
+class Chunk:
+    """A retrievable unit of knowledge plus provenance and relevance score."""
+
+    text: str
+    score: float = 0.0
+    metadata: Dict[str, Any] = field(default_factory=dict, compare=False)
+
+
+class EZRAG:
+    """Lazy, 3-line RAG: ``EZRAG().load(source).ask(question)``.
+
+    Chunking, embeddings, in-memory retrieval and LLM calls are handled
+    transparently. Nothing touches the network until data is loaded or a
+    question is asked.
+    """
+
+    def __init__(
+        self,
+        llm: Optional[str] = None,
+        embedding_model: str = DEFAULT_EMBEDDING_MODEL,
+        provider: str = "hf",
+        hf_token: Optional[str] = None,
+        top_k: int = 4,
+        min_score: float = 0.10,
+        chunk_size: int = 900,
+        chunk_overlap: int = 120,
+        max_new_tokens: int = 400,
+        temperature: float = 0.3,
+    ) -> None:
+        self.llm = llm or _os.getenv("EZRAG_LLM") or DEFAULT_LLM
+        self.embedding_model = embedding_model
+        self.provider = provider
+        self.hf_token = hf_token or _os.getenv("HF_TOKEN")
+        self.top_k = top_k
+        self.min_score = min_score
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+        self.max_new_tokens = max_new_tokens
+        self.temperature = temperature
+        self._encoder = None
+        self._vectors: List[np.ndarray] = []
+        self._chunks: List[Chunk] = []
+        self._cache: Dict[str, np.ndarray] = {}
+
+    # ------------------------------------------------------------------ #
+    # Public API
+    # ------------------------------------------------------------------ #
+
+    def load(self, source: Any) -> "EZRAG":
+        """Index a raw text, a file path, or a ``{name: text}`` mapping."""
+        if isinstance(source, dict):
+            for name, payload in source.items():
+                self._read(payload, {"source": str(name)})
+        elif isinstance(source, (list, tuple)):
+            for item in source:
+                self.load(item)
+        else:
+            self._read(source, {})
+        return self
+
+    def load_files(self, *paths: Union[str, _Path]) -> "EZRAG":
+        """Index one or more ``.txt``/``.md`` files from disk."""
+        for path in paths:
+            self.load(path)
+        return self
+
+    def search(self, question: str, top_k: Optional[int] = None) -> List[Chunk]:
+        """Return the most relevant chunks for a question, best first."""
+        if not self._vectors:
+            raise RuntimeError("no knowledge loaded yet - call load() first")
+        query_vec = self._encode([Chunk(text=question)])[0]
+        query_vec = query_vec / (float(np.linalg.norm(query_vec)) + 1e-9)
+        k = top_k or self.top_k
+        candidates: List[Chunk] = []
+        offset = 0
+        for block in self._vectors:
+            scores = block @ query_vec
+            for local in np.argsort(scores)[::-1][:k]:
+                score = float(scores[local])
+                if score < self.min_score:
+                    continue
+                chunk = self._chunks[offset + int(local)]
+                candidates.append(
+                    Chunk(text=chunk.text, score=score, metadata=dict(chunk.metadata))
+                )
+            offset += block.shape[0]
+        candidates.sort(key=lambda c: c.score, reverse=True)
+        return candidates[:k]
+
+    def ask(self, question: str, **overrides: Any) -> str:
+        """Retrieve evidence and answer using the configured LLM."""
+        top_k = overrides.pop("top_k", self.top_k)
+        hits = self.search(question, top_k=top_k)
+        if not hits:
+            return "I do not have enough information in the loaded documents to answer this."
+        context = self._bundle(hits)
+        prompt = self._build_prompt(question, context)
+        logger.info("asking %s with %d retrieved chunks", self.llm, len(hits))
+        return self._call_llm(prompt, **overrides)
+
+    def query(self, question: str, **overrides: Any) -> str:
+        """Alias for :meth:`ask`."""
+        return self.ask(question, **overrides)
+
+    def stats(self) -> Dict[str, Any]:
+        """Diagnostic summary of the indexed knowledge base."""
+        return {
+            "chunks": len(self._chunks),
+            "vectors": sum(int(b.shape[0]) for b in self._vectors),
+            "embedding_model": self.embedding_model,
+            "llm": self.llm,
+            "provider": self.provider,
+        }
+
+    def __len__(self) -> int:
+        return len(self._chunks)
+
+    def __repr__(self) -> str:
+        return (
+            f"EZRAG(embedding={self.embedding_model!r}, llm={self.llm!r}, "
+            f"provider={self.provider!r}, chunks={len(self._chunks)})"
+        )
+
+    # ------------------------------------------------------------------ #
+    # Ingestion
+    # ------------------------------------------------------------------ #
+
+    @staticmethod
+    def _is_path(value: str) -> bool:
+        try:
+            return _Path(value).is_file()
+        except (OSError, ValueError):
+            return False
+
+    def _read(self, source: Any, metadata: Dict[str, Any]) -> List[Chunk]:
+        try:
+            if hasattr(source, "read"):
+                payload = source.read()
+            elif isinstance(source, _Path):
+                payload = source.read_text(encoding="utf-8")
+            elif isinstance(source, str) and self._is_path(source):
+                payload = _Path(source).read_text(encoding="utf-8")
+            elif isinstance(source, str):
+                payload = source
+            else:
+                raise TypeError(f"unsupported source type: {type(source).__name__}")
+        except (UnicodeDecodeError, PermissionError) as exc:
+            raise ValueError(f"cannot read source {source!r}: {exc}") from exc
+        return self._ingest(str(payload), metadata)
+
+    def _ingest(self, payload: str, metadata: Dict[str, Any]) -> List[Chunk]:
+        chunks = self._chunk_text(payload)
+        logger.info("ingesting %d chunks (%d chars)", len(chunks), len(payload))
+        for start in range(0, len(chunks), 128):
+            batch = chunks[start:start + 128]
+            vectors = self._encode(batch)
+            self._vectors.append(vectors)
+            for i, chunk in enumerate(batch):
+                chunk_meta = dict(metadata)
+                chunk_meta["chunk_index"] = len(self._chunks) + i
+                self._chunks.append(
+                    Chunk(text=chunk.text, metadata=chunk_meta)
+                )
+        return chunks
+
+    def _chunk_text(self, text: str) -> List[Chunk]:
+        """Split text into sentence-aware chunks with a sliding overlap."""
+        text = _normalize_text(text)
+        if not text:
+            return []
+        if len(text) <= self.chunk_size:
+            return [Chunk(text=text)]
+        chunks: List[Chunk] = []
+        current = ""
+        for sentence in _SENTENCE_BOUNDARY.split(text):
+            candidate = f"{current} {sentence}".strip() if current else sentence
+            if len(candidate) > self.chunk_size and current:
+                chunks.append(Chunk(text=current))
+                overlap = (
+                    current[-(self.chunk_overlap or 0):] if self.chunk_overlap else ""
+                )
+                current = f"{overlap} {sentence}".strip()
+            else:
+                current = candidate
+        if current:
+            chunks.append(Chunk(text=current))
+        return chunks
+
+    # ------------------------------------------------------------------ #
+    # Embeddings
+    # ------------------------------------------------------------------ #
+
+    def _build_encoder(self) -> None:
+        if np is None:
+            raise RuntimeError("missing dependency: pip install numpy")
+        try:
+            from sentence_transformers import SentenceTransformer
+        except ImportError as exc:
+            raise RuntimeError(
+                "missing dependency: pip install sentence-transformers"
+            ) from exc
+        logger.info("loading embedding model %s", self.embedding_model)
+        self._encoder = SentenceTransformer(self.embedding_model)
+
+    def _encode(self, chunks: List[Chunk]) -> np.ndarray:
+        """Encode chunk texts into normalised float32 vectors (batched)."""
+        texts = [c.text for c in chunks]
+        vectors: Dict[int, np.ndarray] = {}
+        remaining: List[Tuple[int, str]] = []
+        for i, text in enumerate(texts):
+            hit = self._cache.get(text)
+            if hit is None:
+                remaining.append((i, text))
+            else:
+                vectors[i] = hit
+        if remaining:
+            if self._encoder is None:
+                self._build_encoder()
+            fresh = self._encoder.encode(
+                [text for _, text in remaining],
+                normalize_embeddings=True,
+                show_progress_bar=False,
+            )
+            for (i, text), vec in zip(remaining, np.asarray(fresh, dtype=np.float32)):
+                vectors[i] = vec
+                if len(self._cache) < 20_000:
+                    self._cache[text] = vec
+        return np.stack([vectors[i] for i in range(len(texts))])
+
+    # ------------------------------------------------------------------ #
+    # Context bundling + LLM execution
+    # ------------------------------------------------------------------ #
+
+    @staticmethod
+    def _bundle(hits: List[Chunk]) -> str:
+        lines = ["<context>"]
+        for i, hit in enumerate(hits, start=1):
+            src = hit.metadata.get("source", "document")
+            idx = hit.metadata.get("chunk_index", i)
+            lines.append(f"[{i}] (source: {src}, chunk #{idx}) {hit.text}")
+        lines.append(f"</context>\nBest match score: {hits[0].score:.2f}")
+        return "\n".join(lines)
+
+    @staticmethod
+    def _build_prompt(question: str, context: str) -> str:
+        return (
+            f"{_SYSTEM_PROMPT}\n\n"
+            f"QUESTION: {question}\n\n"
+            f"{context}\n\n"
+            "ANSWER:"
+        )
+
+    def _call_llm(self, prompt: str, **overrides: Any) -> str:
+        model = overrides.pop("llm", None) or self.llm
+        provider = overrides.pop("provider", self.provider)
+        max_new_tokens = overrides.pop("max_new_tokens", self.max_new_tokens)
+        temperature = overrides.pop("temperature", self.temperature)
+
+        if provider == "hf":
+            if InferenceClient is None:
+                raise RuntimeError(
+                    "the 'hf' provider requires huggingface_hub: pip install huggingface_hub"
+                )
+            client = InferenceClient(model=model, token=self.hf_token)
+            output = client.text_generation(
+                prompt,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=0.9,
+                stop_sequences=["</s>", "<|endoftext|>"],
+            )
+            if hasattr(output, "generated_text"):
+                return str(output.generated_text).strip()
+            return str(output).strip()
+
+        if provider == "litellm":
+            if litellm is None:
+                raise RuntimeError(
+                    "the 'litellm' provider requires litellm: pip install litellm"
+                )
+            response = litellm.completion(
+                model=model,
+                messages=[{"role": "user", "content": prompt}],
+                temperature=temperature,
+                max_tokens=max_new_tokens,
+            )
+            return str(response["choices"][0]["message"]["content"]).strip()
+
+        raise ValueError(
+            f"unknown provider {provider!r}; use 'hf' (Hugging Face Serverless) "
+            "or 'litellm' (any OpenAI-compatible gateway)"
+        )
+

ezrag_engine-0.1.0.dist-info/METADATA

@@ -0,0 +1,117 @@
+Metadata-Version: 2.4
+Name: ezrag-engine
+Version: 0.1.0
+Summary: A 3-line, zero-boilerplate, cloud-ready RAG pipeline for Python and Hugging Face Spaces.
+Author-email: Oguru Vinay Reddy <Vinayreddy.ace@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://huggingface.co/
+Project-URL: Repository, https://github.com/<your-github-username>/ezrag
+Project-URL: Documentation, https://github.com/<your-github-username>/ezrag#readme
+Keywords: rag,llm,retrieval-augmented-generation,huggingface,embeddings,vector-search
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: sentence-transformers>=3.0.0
+Requires-Dist: huggingface_hub>=0.23.0
+Requires-Dist: numpy>=1.26.0
+Requires-Dist: requests>=2.31.0
+Requires-Dist: litellm>=1.40.0
+
+<div align="center">
+
+# **ezrag** — 3-Line, Zero-Boilerplate, Cloud-Ready RAG
+
+**A production-grade Retrieval-Augmented Generation framework for Python and Hugging Face Spaces.**
+
+![PyPI](https://img.shields.io/pypi/v/ezrag.svg)
+![Python](https://img.shields.io/pypi/pyversions/ezrag.svg)
+![PyPI - License](https://img.shields.io/pypi/l/ezrag.svg)
+
+</div>
+
+---
+
+## Why ezrag?
+
+Most RAG stacks take hours to wire up: embedding models, vector stores, chunkers,
+prompts, and LLM clients. ezrag hides all of that behind one class, so you go from
+an empty folder to a **grounded, citation-aware answer** in exactly three lines of
+code.
+
+**What you get out of the box:**
+
+- **Chunking** — sentence-aware splitting with a sliding overlap, zero configuration
+- **Embeddings** — a lightweight `sentence-transformers` model, loaded lazily
+- **Vector DB** — a self-contained in-memory store with fast NumPy cosine search
+- **Context bundling** — evidence is ranked, tagged `[1] [2] [3] ...` and injected into the prompt
+- **LLM execution** — unified calls through the Hugging Face Serverless API or LiteLLM
+- **Weak-evidence guard** — answers truthfully when the documents do not cover the question
+
+## Installation
+
+```bash
+pip install ezrag
+```
+
+## Quick Start
+
+That is the entire experience — import, load, ask:
+
+```python
+from ezrag import EZRAG                     # 1. import
+
+rag = EZRAG().load(SOURCE)                  # 2. initialise + load your source
+print(rag.ask("Your question here"))        # 3. ask — grounded, cited answer
+```
+
+No environment variables to set up, no vector database to spin up, no prompting
+prompts to craft. One object call chain produces a grounded, cited answer.
+
+## Configuration
+
+| Parameter | Default | Purpose |
+| --- | --- | --- |
+| `llm` | `mistralai/Mistral-7B-Instruct-v0.2` | LLM identifier or endpoint |
+| `embedding_model` | `all-MiniLM-L6-v2` | sentence-transformers model |
+| `provider` | `"hf"` | `"hf"` (HF Serverless) or `"litellm"` |
+| `hf_token` | env `HF_TOKEN` | token for protected / gated models |
+| `top_k` | `4` | number of evidence chunks retrieved |
+| `min_score` | `0.10` | retrieval similarity floor |
+| `chunk_size` | `900` | target chunk length (characters) |
+| `chunk_overlap` | `120` | sliding context overlap |
+
+## Providers
+
+- **Hugging Face Serverless** (default) — no custom API wiring. Set the `HF_TOKEN`
+  environment variable only when the model is gated.
+- **LiteLLM** — point `llm` at any OpenAI-compatible endpoint, for example
+  `openai/gpt-4o` or `azure/gpt-4o`.
+
+## How It Works
+
+```
+  source ──▶ chunk ──▶ embed ──▶ in-memory index
+                                       │
+  question ──▶ embed ──▶ top-k search ──┘
+                                       │
+                      context bundle ──▶ LLM ──▶ cited answer
+```
+
+## Hugging Face Spaces
+
+ezrag runs out of the box in a Hugging Face Space — pip pulls in every runtime
+dependency automatically. Store a `HF_TOKEN` Space secret if your chosen model is
+gated.
+
+## License
+
+MIT © Oguru Vinay Reddy

ezrag_engine-0.1.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+ezrag/__init__.py,sha256=Xfk0iZRFiNL4AMn3X9drkFdi4AZqwcmR965hlk6-2Uk,157
+ezrag/engine.py,sha256=V4CEzjOQUJWzkqx_ZkehQ9ABTbEtytHTIH20KEd8tn0,13293
+ezrag_engine-0.1.0.dist-info/METADATA,sha256=0XpaKSri8N2C8i_Z1wSnjJOVIkL9t-8hkCqZU24NDYQ,4395
+ezrag_engine-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+ezrag_engine-0.1.0.dist-info/top_level.txt,sha256=5GepBIq-mzRxEBqwhAyFK3d6EOwjevtdn2pCQQxoIJI,6
+ezrag_engine-0.1.0.dist-info/RECORD,,

ezrag_engine-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

ezrag_engine-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+ezrag