rag-python 0.1.0__py3-none-any.whl

rag_python/providers/openai_provider.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from openai import OpenAI
+
+from ..config import OPENAI_API_KEY, LLM_MODEL, EMBEDDING_MODEL
+
+
+class OpenAIProvider:
+    def __init__(self, *, api_key: str | None = None) -> None:
+        api_key = api_key or OPENAI_API_KEY
+        if not api_key:
+            raise RuntimeError("OpenAI API key missing. Set OPENAI_API_KEY or pass api_key=...")
+        self._client = OpenAI(api_key=api_key)
+
+    def generate(
+        self,
+        *,
+        user: str,
+        system: str | None = None,
+        model: str | None = None,
+        temperature: float = 0.2,
+        max_tokens: int = 1024,
+    ) -> str:
+        r = self._client.chat.completions.create(
+            model=model or LLM_MODEL,
+            messages=[
+                {"role": "system", "content": system or "You are a helpful assistant."},
+                {"role": "user", "content": user},
+            ],
+            temperature=temperature,
+            max_tokens=max_tokens,
+        )
+        return (r.choices[0].message.content or "").strip()
+
+    def embed(self, texts: list[str], *, model: str | None = None) -> list[list[float]]:
+        if not texts:
+            return []
+        out: list[list[float]] = []
+        batch_size = 100
+        for i in range(0, len(texts), batch_size):
+            batch = [t[:8000] for t in texts[i : i + batch_size]]
+            r = self._client.embeddings.create(input=batch, model=model or EMBEDDING_MODEL)
+            for e in r.data:
+                out.append(e.embedding)
+        return out
+

rag_python/query_rewriting.py

@@ -0,0 +1,65 @@
+"""Query understanding & rewriting: expand/rewrite user query for better retrieval."""
+from .config import LLM_MODEL
+from .providers import LLMProvider, make_llm_provider
+
+
+def rewrite_for_retrieval(
+    query: str,
+    *,
+    n_queries: int = 3,
+    llm: LLMProvider | None = None,
+    llm_model: str | None = None,
+) -> list[str]:
+    """Generate n alternative queries (multi-query) for better recall."""
+    if n_queries <= 1:
+        return [query]
+    llm = llm or make_llm_provider("openai")
+    prompt = (
+        "You are a search expert. Given the user question below, generate "
+        f"{n_queries} different search queries that would help find relevant documents. "
+        "Output ONLY the queries, one per line, no numbering or bullets."
+    )
+    try:
+        text = llm.generate(
+            system=prompt,
+            user=query,
+            model=llm_model or LLM_MODEL,
+            temperature=0.3,
+            max_tokens=200,
+        )
+        lines = [line.strip() for line in text.splitlines() if line.strip()]
+        cleaned = []
+        for line in lines[:n_queries]:
+            for sep in (". ", ") ", "- ", " ", ""):
+                if line.startswith(sep):
+                    line = line[len(sep):].strip()
+                if line and not line[0].isdigit():
+                    break
+            if line:
+                cleaned.append(line)
+        if not cleaned:
+            return [query]
+        return cleaned[:n_queries]
+    except Exception:
+        return [query]
+
+
+def understand_query(query: str, *, llm: LLMProvider | None = None, llm_model: str | None = None) -> str:
+    """Optional: rewrite single query for clarity (e.g. spell correction, intent)."""
+    llm = llm or make_llm_provider("openai")
+    try:
+        out = llm.generate(
+            system=(
+                "Rewrite the user's question to be a clear, standalone search query. "
+                "Preserve intent. Output only the rewritten query."
+            ),
+            user=query,
+            model=llm_model or LLM_MODEL,
+            temperature=0,
+            max_tokens=150,
+        )
+        out = (out or "").strip().strip('"')
+        return out if out else query
+    except Exception:
+        return query
+

rag_python/rag_pipeline.py

@@ -0,0 +1,241 @@
+"""Full RAG pipeline: Query → Understanding/Rewrite → Retrieval (multi-query) → Rerank → LLM → Guardrails → Eval/Retry."""
+from dataclasses import dataclass
+from pathlib import Path
+
+from .cleaning import clean_document
+from .chunking import chunk_text, Chunk
+from .document_loaders import load_file, load_directory, LoadedDocument
+from .vector_store import ingest_chunks, delete_all
+from .retrieval import retrieve as rag_retrieve
+from .generation import generate
+from .guardrails import check_prompt_injection, check_hallucination
+from .evaluation import evaluate_rag, should_retry, self_correct
+from .providers import LLMProvider, EmbeddingProvider, make_llm_provider, make_embedding_provider
+from .config import DATA_DIR, CHUNK_SIZE, CHUNK_OVERLAP, CHUNK_STRATEGY
+from .options import QueryConfig, SearchConfig
+
+
+@dataclass
+class RAGResponse:
+    answer: str
+    sources: list[dict]
+    evaluation: dict
+    retried: bool
+
+
+def _embed_fn_for_semantic(embedder: EmbeddingProvider, embedding_model: str | None):
+    def _fn(texts: list[str]):
+        return embedder.embed(texts, model=embedding_model)
+
+    return _fn
+
+
+def _load_documents(
+    paths: list[Path] | None = None,
+    data_path: Path | None = None,
+    *,
+    extensions: tuple[str, ...] = (".txt", ".md", ".pdf", ".docx"),
+) -> list[LoadedDocument]:
+    """Load documents from explicit paths and/or a data directory."""
+    docs: list[LoadedDocument] = []
+    seen_sources: set[str] = set()
+
+    if paths:
+        for raw in paths:
+            path = Path(raw)
+            if path.is_file():
+                doc = load_file(path)
+                if doc and doc.content.strip() and doc.source not in seen_sources:
+                    seen_sources.add(doc.source)
+                    docs.append(doc)
+            elif path.is_dir():
+                for doc in load_directory(path, extensions=extensions):
+                    if doc.source not in seen_sources:
+                        seen_sources.add(doc.source)
+                        docs.append(doc)
+
+    if data_path is not None:
+        root = Path(data_path)
+        if root.is_file():
+            doc = load_file(root)
+            if doc and doc.content.strip() and doc.source not in seen_sources:
+                docs.append(doc)
+        elif root.is_dir():
+            for doc in load_directory(root, extensions=extensions):
+                if doc.source not in seen_sources:
+                    seen_sources.add(doc.source)
+                    docs.append(doc)
+
+    return docs
+
+
+def _ingest_documents(
+    docs: list[LoadedDocument],
+    *,
+    clean: bool,
+    chunk_strategy: str,
+    chunk_size: int,
+    chunk_overlap: int,
+    embedding_model: str | None,
+    embedder: EmbeddingProvider,
+) -> int:
+    all_chunks: list[Chunk] = []
+    for doc in docs:
+        text = doc.content
+        if clean:
+            text = clean_document(text)
+        if chunk_strategy == "semantic":
+            chunks = chunk_text(
+                text,
+                strategy="semantic",
+                chunk_size=chunk_size,
+                overlap=chunk_overlap,
+                metadata=doc.metadata,
+                embed_fn=_embed_fn_for_semantic(embedder, embedding_model),
+            )
+        else:
+            chunks = chunk_text(
+                text,
+                strategy=chunk_strategy,
+                chunk_size=chunk_size,
+                overlap=chunk_overlap,
+                metadata=doc.metadata,
+            )
+        all_chunks.extend(chunks)
+
+    if not all_chunks:
+        return 0
+
+    ids = [f"chunk_{i}" for i in range(len(all_chunks))]
+    texts = [c.text for c in all_chunks]
+    metadatas = [c.metadata for c in all_chunks]
+    ingest_chunks(ids, texts, metadatas, embedder=embedder, embed_model=embedding_model)
+    return len(all_chunks)
+
+
+def ingest(
+    data_path: Path | str | None = None,
+    *,
+    paths: list[Path | str] | None = None,
+    clean: bool = True,
+    chunk_strategy: str | None = None,
+    chunk_size: int | None = None,
+    chunk_overlap: int | None = None,
+    extensions: tuple[str, ...] | None = None,
+    reindex: bool = False,
+    embedding_model: str | None = None,
+    embedder: EmbeddingProvider | None = None,
+) -> int:
+    """
+    Load documents → clean → chunk → embed → vector store.
+    Pass ``paths`` for multiple files/folders, or ``data_path`` for a single root.
+    Returns number of chunks ingested.
+    """
+    if reindex:
+        delete_all()
+    strategy = chunk_strategy or CHUNK_STRATEGY
+    size = chunk_size or CHUNK_SIZE
+    overlap = chunk_overlap or CHUNK_OVERLAP
+    ext = extensions or (".txt", ".md", ".pdf", ".docx")
+    embedder = embedder or make_embedding_provider("openai")
+
+    path_list = [Path(p) for p in paths] if paths else None
+    root = Path(data_path) if data_path else (None if path_list else Path(DATA_DIR))
+    docs = _load_documents(path_list, root, extensions=ext)
+    return _ingest_documents(
+        docs,
+        clean=clean,
+        chunk_strategy=strategy,
+        chunk_size=size,
+        chunk_overlap=overlap,
+        embedding_model=embedding_model,
+        embedder=embedder,
+    )
+
+
+def query(
+    question: str,
+    *,
+    search: SearchConfig | None = None,
+    query_config: QueryConfig | None = None,
+    multi_query: bool | None = None,
+    use_guardrails: bool | None = None,
+    use_retry: bool | None = None,
+    eval_threshold: float | None = None,
+    max_retries: int | None = None,
+    llm_model: str | None = None,
+    embedding_model: str | None = None,
+    llm: LLMProvider | None = None,
+    embedder: EmbeddingProvider | None = None,
+) -> RAGResponse:
+    """
+    Run full pipeline: input guardrails → retrieval → rerank → generate →
+    output guardrails → evaluation → retry/self-correction if needed.
+    """
+    search_cfg = search or SearchConfig()
+    q_cfg = query_config or QueryConfig()
+    use_guardrails = q_cfg.use_guardrails if use_guardrails is None else use_guardrails
+    use_retry = q_cfg.use_retry if use_retry is None else use_retry
+    eval_threshold = q_cfg.eval_threshold if eval_threshold is None else eval_threshold
+    max_retries = q_cfg.max_retries if max_retries is None else max_retries
+
+    llm = llm or make_llm_provider("openai")
+    embedder = embedder or make_embedding_provider("openai")
+
+    if use_guardrails:
+        safe, msg = check_prompt_injection(question, llm=llm, llm_model=llm_model)
+        if not safe:
+            return RAGResponse(
+                answer=f"Request blocked: {msg}",
+                sources=[],
+                evaluation={},
+                retried=False,
+            )
+
+    q_understood = question
+
+    hits = rag_retrieve(
+        q_understood,
+        retriever=search_cfg.retriever,
+        multi_query=multi_query,
+        n_queries=search_cfg.multi_query_n,
+        top_k_retrieve=search_cfg.top_k_retrieve,
+        top_k_rerank=search_cfg.top_k_rerank,
+        rerank_enabled=search_cfg.rerank_enabled,
+        embedder=embedder,
+        embedding_model=embedding_model,
+        llm=llm,
+        llm_model=llm_model,
+    )
+    context_chunks = [h[0] for h in hits]
+    sources = [{"text": h[0][:200], "metadata": h[1], "score": h[2]} for h in hits]
+    context_str = "\n\n".join(context_chunks)
+
+    if not context_chunks:
+        return RAGResponse(
+            answer="No relevant documents found for your question.",
+            sources=[],
+            evaluation={"faithfulness": 1.0, "relevance": 0.0},
+            retried=False,
+        )
+
+    answer = generate(q_understood, context_chunks, model=llm_model, llm=llm)
+    retried = False
+
+    if use_guardrails:
+        grounded, msg = check_hallucination(answer, context_str, llm=llm, llm_model=llm_model)
+        if not grounded:
+            answer = answer + f"\n\n[Note: {msg}]"
+
+    eval_scores = evaluate_rag(q_understood, answer, context_str, llm=llm, llm_model=llm_model)
+
+    if use_retry and should_retry(eval_scores["faithfulness"], eval_scores["relevance"], eval_threshold):
+        for _ in range(max(0, max_retries - 1)):
+            answer = self_correct(q_understood, answer, context_str, llm=llm, llm_model=llm_model)
+            eval_scores = evaluate_rag(q_understood, answer, context_str, llm=llm, llm_model=llm_model)
+            retried = True
+            if not should_retry(eval_scores["faithfulness"], eval_scores["relevance"], eval_threshold):
+                break
+
+    return RAGResponse(answer=answer, sources=sources, evaluation=eval_scores, retried=retried)
+

rag_python/reranker.py

@@ -0,0 +1,64 @@
+"""Reranking: cross-encoder to score query–document relevance."""
+from typing import Any
+
+from .config import RERANKER_MODEL, RERANK_ENABLED
+
+_reranker = None
+
+
+def _get_reranker():
+    global _reranker
+    if _reranker is None:
+        try:
+            from sentence_transformers import CrossEncoder
+        except ImportError as e:
+            raise ImportError(
+                "Reranking requires optional dependencies. Install with: pip install rag-python[rerank]"
+            ) from e
+        _reranker = CrossEncoder(RERANKER_MODEL)
+    return _reranker
+
+
+def rerank(
+    query: str,
+    documents: list[str],
+    top_k: int = 5,
+    *,
+    rerank_enabled: bool | None = None,
+) -> list[tuple[str, float, int]]:
+    """Return (document, score, original_index) sorted by relevance."""
+    if not documents:
+        return []
+    use_rerank = RERANK_ENABLED if rerank_enabled is None else rerank_enabled
+    if not use_rerank:
+        return [(documents[i], 0.0, i) for i in range(min(top_k, len(documents)))]
+
+    try:
+        pairs = [(query, d) for d in documents]
+        model = _get_reranker()
+        scores = model.predict(pairs)
+        indexed = [(documents[i], float(scores[i]), i) for i in range(len(documents))]
+        indexed.sort(key=lambda x: -x[1])
+        return indexed[: top_k if top_k < len(indexed) else len(indexed)]
+    except ImportError:
+        # Graceful fallback: keep retrieval order
+        return [(documents[i], 0.0, i) for i in range(min(top_k, len(documents)))]
+
+
+def rerank_with_metadata(
+    query: str,
+    doc_meta_list: list[tuple[str, dict[str, Any]]],
+    top_k: int = 5,
+    *,
+    rerank_enabled: bool | None = None,
+) -> list[tuple[str, dict[str, Any], float]]:
+    """Rerank (document, metadata) pairs; return (doc, meta, score)."""
+    if not doc_meta_list:
+        return []
+    docs = [d for d, _ in doc_meta_list]
+    metas = [m for _, m in doc_meta_list]
+    results = rerank(query, docs, top_k=min(top_k, len(docs)), rerank_enabled=rerank_enabled)
+    out = []
+    for doc, score, idx in results:
+        out.append((doc, metas[idx], score))
+    return out

rag_python/retrieval.py

@@ -0,0 +1,61 @@
+"""Retrieval: multi-query retrieval + reranking."""
+from typing import Any
+
+from .vector_store import retrieve as chroma_retrieve
+from .query_rewriting import rewrite_for_retrieval
+from .reranker import rerank_with_metadata
+from .providers import EmbeddingProvider, LLMProvider
+from .options import RetrieverStrategy
+from .config import TOP_K_RETRIEVE, TOP_K_RERANK, MULTI_QUERY_N
+
+
+def retrieve(
+    query: str,
+    *,
+    embedder: EmbeddingProvider,
+    embedding_model: str | None = None,
+    retriever: RetrieverStrategy = "multi_query",
+    multi_query: bool | None = None,
+    n_queries: int | None = None,
+    top_k_retrieve: int | None = None,
+    top_k_rerank: int | None = None,
+    rerank_enabled: bool | None = None,
+    llm: LLMProvider | None = None,
+    llm_model: str | None = None,
+) -> list[tuple[str, dict[str, Any], float]]:
+    """
+    Retrieve relevant chunks using vector or multi-query search, then rerank.
+    Returns list of (document_text, metadata, rerank_score).
+    """
+    top_k_retrieve = top_k_retrieve or TOP_K_RETRIEVE
+    top_k_rerank = top_k_rerank or TOP_K_RERANK
+    n_queries = n_queries or MULTI_QUERY_N
+    use_multi_query = retriever == "multi_query" if multi_query is None else multi_query
+
+    queries = [query]
+    if use_multi_query and n_queries > 1:
+        rewritten = rewrite_for_retrieval(query, n_queries=n_queries, llm=llm, llm_model=llm_model)
+        if rewritten:
+            queries = rewritten
+
+    seen_docs: set[str] = set()
+    all_candidates: list[tuple[str, dict, float]] = []
+    for q in queries:
+        emb = embedder.embed([q], model=embedding_model)[0]
+        hits = chroma_retrieve(emb, top_k=top_k_retrieve)
+        for doc, meta, dist in hits:
+            key = (doc[:200], meta.get("source", ""))
+            if key in seen_docs:
+                continue
+            seen_docs.add(key)
+            all_candidates.append((doc, meta, -dist))
+
+    if not all_candidates:
+        return []
+    docs = [c[0] for c in all_candidates]
+    metas = [c[1] for c in all_candidates]
+    reranked = rerank_with_metadata(
+        query, list(zip(docs, metas)), top_k=top_k_rerank, rerank_enabled=rerank_enabled
+    )
+    return reranked
+

rag_python/vector_store.py

@@ -0,0 +1,91 @@
+"""Vector store (Chroma) for ingestion and retrieval."""
+from pathlib import Path
+from typing import Any
+
+import chromadb
+from chromadb.config import Settings
+
+from .config import CHROMA_PERSIST_DIR
+from .providers import EmbeddingProvider
+
+COLLECTION_NAME = "rag_docs"
+# chromadb.PersistentClient is a factory function (not a type), so avoid using it in annotations.
+_chroma_client: Any | None = None
+_persist_dir: Path | None = None
+
+
+def set_persist_dir(path: Path | str) -> None:
+    """Override Chroma persistence directory (resets the client)."""
+    global _chroma_client, _persist_dir
+    _chroma_client = None
+    _persist_dir = Path(path)
+
+
+def _get_client() -> Any:
+    global _chroma_client
+    if _chroma_client is None:
+        p = Path(_persist_dir or CHROMA_PERSIST_DIR)
+        p.mkdir(parents=True, exist_ok=True)
+        _chroma_client = chromadb.PersistentClient(path=str(p), settings=Settings(anonymized_telemetry=False))
+    return _chroma_client
+
+
+def get_collection():
+    return _get_client().get_or_create_collection(
+        name=COLLECTION_NAME,
+        metadata={"hnsw:space": "cosine"},
+    )
+
+
+def ingest_chunks(
+    ids: list[str],
+    texts: list[str],
+    metadatas: list[dict[str, Any]],
+    *,
+    embeddings: list[list[float]] | None = None,
+    embedder: EmbeddingProvider | None = None,
+    embed_model: str | None = None,
+) -> None:
+    """Add chunks to Chroma (embeddings required, or provide an embedder)."""
+    if embeddings is None:
+        if not embedder:
+            raise RuntimeError("ingest_chunks requires `embeddings=` or `embedder=`")
+        embeddings = embedder.embed(texts, model=embed_model)
+    coll = get_collection()
+    clean_metas = []
+    for m in metadatas:
+        clean = {}
+        for k, v in m.items():
+            if v is None:
+                continue
+            if isinstance(v, (str, int, float, bool)):
+                clean[k] = v
+            else:
+                clean[k] = str(v)
+        clean_metas.append(clean)
+    coll.add(ids=ids, embeddings=embeddings, documents=texts, metadatas=clean_metas)
+
+
+def retrieve(
+    query_embedding: list[float],
+    top_k: int = 20,
+    where: dict | None = None,
+) -> list[tuple[str, dict, float]]:
+    """Return (document, metadata, distance) for top_k. Lower distance = more similar."""
+    coll = get_collection()
+    res = coll.query(
+        query_embeddings=[query_embedding],
+        n_results=top_k,
+        where=where,
+        include=["documents", "metadatas", "distances"],
+    )
+    docs = res["documents"][0] or []
+    metas = res["metadatas"][0] or []
+    dists = res["distances"][0] or []
+    return list(zip(docs, metas, dists))
+
+
+def delete_all() -> None:
+    """Remove all documents from the collection (for re-ingestion)."""
+    _get_client().delete_collection(COLLECTION_NAME)
+

rag_python-0.1.0.dist-info/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2026
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+