ragforge-ml 0.1.0__py3-none-any.whl

ragforge/__init__.py

@@ -0,0 +1,12 @@
+"""RAGforge — local-first RAG pipeline with an eval harness."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from ragforge.pipeline import Answer, Pipeline, Source
+
+try:
+    __version__ = version("ragforge-ml")
+except PackageNotFoundError:
+    __version__ = "0.0.0+local"
+
+__all__ = ["Answer", "Pipeline", "Source", "__version__"]

ragforge/cli.py

@@ -0,0 +1,205 @@
+"""Command-line interface — ``ragforge`` / ``rf``.
+
+Subcommands:
+    ingest    load PDF / Markdown documents into a collection
+    ask       run the full RAG pipeline (retrieve, rerank, generate)
+    eval      run the eval harness on a JSONL of QA samples
+    serve     start the FastAPI server
+    info      print the current collection stats
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Annotated
+
+import typer
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+
+app = typer.Typer(
+    name="ragforge",
+    help="Local-first RAG pipeline with an eval harness.",
+    no_args_is_help=True,
+    rich_markup_mode="rich",
+)
+console = Console()
+
+
+@app.command()
+def ingest(
+    paths: Annotated[list[Path], typer.Argument(help="Files or directories to index")],
+    collection: Annotated[str, typer.Option(help="Qdrant collection name")] = "ragforge",
+    store_path: Annotated[Path, typer.Option(help="Qdrant embedded storage dir")] = Path(
+        "qdrant_storage"
+    ),
+    chunk_size: Annotated[int, typer.Option()] = 1024,
+    chunk_overlap: Annotated[int, typer.Option()] = 128,
+) -> None:
+    """Load documents into the vector store."""
+    from ragforge import Pipeline
+
+    console.print(Panel.fit(f"[bold purple]ragforge ingest[/]  [dim]{collection}[/]"))
+    rag = Pipeline.from_defaults(collection=collection, store_path=store_path, use_reranker=False)
+    rag.chunk_size = chunk_size
+    rag.chunk_overlap = chunk_overlap
+    n = rag.ingest([str(p) for p in paths])
+    console.print(
+        f"[green]ok[/]  indexed [bold]{n}[/] chunks   total in collection: {rag.store.count()}"
+    )
+
+
+@app.command()
+def ask(
+    question: Annotated[str, typer.Argument()],
+    collection: Annotated[str, typer.Option()] = "ragforge",
+    store_path: Annotated[Path, typer.Option()] = Path("qdrant_storage"),
+    model_id: Annotated[
+        str, typer.Option(help="HF causal LM to generate the answer")
+    ] = "Qwen/Qwen2.5-3B-Instruct",
+    k: Annotated[int, typer.Option(help="Top-k after rerank")] = 5,
+    max_new_tokens: Annotated[int, typer.Option()] = 256,
+    quantize: Annotated[
+        str | None, typer.Option(help="Optional turboquant-ml method, e.g. bnb-nf4")
+    ] = None,
+) -> None:
+    """Retrieve, rerank, and generate an answer."""
+    from ragforge import Pipeline
+
+    console.print(Panel.fit(f"[bold purple]ragforge ask[/]  [dim]{question}[/]"))
+
+    if quantize:
+        from ragforge.llm import QuantizedHFLLM
+
+        llm = QuantizedHFLLM(model_id, method=quantize)
+    else:
+        from ragforge.llm import HFLLM
+
+        llm = HFLLM(model_id)
+
+    rag = Pipeline.from_defaults(collection=collection, store_path=store_path, llm=llm)
+    ans = rag.ask(question, top_k=k, max_new_tokens=max_new_tokens)
+    console.print(f"\n[bold]Answer[/]  [dim]({ans.latency_ms:.0f} ms)[/]\n{ans.text}\n")
+
+    table = Table(title="Sources", show_header=True)
+    table.add_column("#", style="dim", width=3)
+    table.add_column("score", justify="right")
+    table.add_column("source")
+    for i, s in enumerate(ans.sources, 1):
+        loc = s.metadata.get("path", s.id)
+        page = s.metadata.get("page")
+        loc = f"{loc}  p.{page}" if page else loc
+        table.add_row(str(i), f"{s.score:.3f}", loc)
+    console.print(table)
+
+
+@app.command()
+def eval(
+    dataset: Annotated[Path, typer.Argument(help="JSONL with {question, ground_truth} per line")],
+    collection: Annotated[str, typer.Option()] = "ragforge",
+    store_path: Annotated[Path, typer.Option()] = Path("qdrant_storage"),
+    model_id: Annotated[str, typer.Option()] = "Qwen/Qwen2.5-3B-Instruct",
+    metrics: Annotated[
+        str, typer.Option(help="Comma-separated")
+    ] = "context_recall,answer_relevance,faithfulness",
+    out: Annotated[Path | None, typer.Option(help="Save JSON report here")] = None,
+    limit: Annotated[int | None, typer.Option(help="Cap the dataset for quick runs")] = None,
+) -> None:
+    """Run the eval harness over a JSONL dataset."""
+    import time
+
+    from ragforge import Pipeline
+    from ragforge.eval import evaluate
+    from ragforge.eval.report import EvalReport
+    from ragforge.llm import HFLLM
+
+    samples_in = _read_jsonl(dataset, limit=limit)
+    console.print(
+        Panel.fit(f"[bold purple]ragforge eval[/]  n={len(samples_in)}  metrics={metrics}")
+    )
+
+    llm = HFLLM(model_id)
+    rag = Pipeline.from_defaults(collection=collection, store_path=store_path, llm=llm)
+
+    samples_out: list[dict] = []
+    latencies: list[float] = []
+    for s in samples_in:
+        t0 = time.perf_counter()
+        ans = rag.ask(s["question"])
+        latencies.append((time.perf_counter() - t0) * 1000)
+        samples_out.append(
+            {
+                "question": s["question"],
+                "answer": ans.text,
+                "contexts": [src.text for src in ans.sources],
+                "ground_truth": s.get("ground_truth"),
+            }
+        )
+
+    metric_list = [m.strip() for m in metrics.split(",") if m.strip()]
+    res = evaluate(samples_out, encoder=rag.encoder, metrics=metric_list)
+    report = EvalReport(
+        n=res["n"],
+        means=res["means"],
+        per_sample=res["per_sample"],
+        latencies_ms=latencies,
+    )
+    console.print(report.as_table())
+    if out:
+        report.save(out)
+        console.print(f"[green]ok[/]  saved {out}")
+
+
+@app.command()
+def serve(
+    collection: Annotated[str, typer.Option()] = "ragforge",
+    store_path: Annotated[Path, typer.Option()] = Path("qdrant_storage"),
+    model_id: Annotated[str, typer.Option()] = "Qwen/Qwen2.5-3B-Instruct",
+    host: Annotated[str, typer.Option()] = "127.0.0.1",
+    port: Annotated[int, typer.Option()] = 8000,
+) -> None:
+    """Start the FastAPI server."""
+    import uvicorn
+
+    from ragforge import Pipeline
+    from ragforge.llm import HFLLM
+    from ragforge.serve import build_app
+
+    llm = HFLLM(model_id)
+    rag = Pipeline.from_defaults(collection=collection, store_path=store_path, llm=llm)
+    fastapi_app = build_app(rag)
+    uvicorn.run(fastapi_app, host=host, port=port, log_level="info")
+
+
+@app.command()
+def info(
+    collection: Annotated[str, typer.Option()] = "ragforge",
+    store_path: Annotated[Path, typer.Option()] = Path("qdrant_storage"),
+) -> None:
+    """Print the current collection stats."""
+    from ragforge import Pipeline
+
+    rag = Pipeline.from_defaults(collection=collection, store_path=store_path, use_reranker=False)
+    console.print(f"collection : [bold]{collection}[/]")
+    console.print(f"store      : [bold]{store_path}[/]")
+    console.print(f"vectors    : [bold]{rag.store.count()}[/]")
+    console.print(f"encoder    : [bold]{rag.encoder.model_id}[/] dim={rag.encoder.dim}")
+
+
+def _read_jsonl(path: Path, *, limit: int | None) -> list[dict]:
+    rows: list[dict] = []
+    with path.open(encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            rows.append(json.loads(line))
+            if limit and len(rows) >= limit:
+                break
+    return rows
+
+
+if __name__ == "__main__":
+    app()

ragforge/embed/__init__.py

@@ -0,0 +1,5 @@
+"""Embedding models — thin wrapper around sentence-transformers."""
+
+from ragforge.embed.encoder import Encoder, SentenceTransformerEncoder
+
+__all__ = ["Encoder", "SentenceTransformerEncoder"]

ragforge/embed/encoder.py

@@ -0,0 +1,62 @@
+"""Encoder protocol + sentence-transformers backend.
+
+The protocol is intentionally small — :meth:`encode` is the only required
+method — so swapping the backend (e.g. for an in-house ONNX encoder, an API
+client, or a cached encoder) is cheap.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+import numpy as np
+
+
+class Encoder(Protocol):
+    """Anything that can map a batch of texts to fixed-size vectors."""
+
+    dim: int
+
+    def encode(
+        self, texts: list[str], *, batch_size: int = 32, normalize: bool = True
+    ) -> np.ndarray: ...
+
+
+class SentenceTransformerEncoder:
+    """Default encoder.
+
+    ``BAAI/bge-small-en-v1.5`` (33 M params, 384-dim) is the default because it
+    sits at the Pareto frontier of "small enough to run anywhere" / "good
+    enough on MTEB". For multilingual corpora, pass
+    ``intfloat/multilingual-e5-small`` instead.
+    """
+
+    def __init__(
+        self,
+        model_id: str = "BAAI/bge-small-en-v1.5",
+        *,
+        device: str | None = None,
+        cache_folder: str | None = None,
+    ) -> None:
+        from sentence_transformers import SentenceTransformer
+
+        self.model_id = model_id
+        self._model = SentenceTransformer(model_id, device=device, cache_folder=cache_folder)
+        self.dim = self._model.get_sentence_embedding_dimension()
+
+    def encode(
+        self, texts: list[str], *, batch_size: int = 32, normalize: bool = True
+    ) -> np.ndarray:
+        if not texts:
+            return np.zeros((0, self.dim), dtype=np.float32)
+        vecs = self._model.encode(
+            texts,
+            batch_size=batch_size,
+            normalize_embeddings=normalize,
+            convert_to_numpy=True,
+            show_progress_bar=False,
+        )
+        return vecs.astype(np.float32, copy=False)
+
+    def __repr__(self) -> str:
+        return f"SentenceTransformerEncoder({self.model_id!r}, dim={self.dim})"

ragforge/eval/__init__.py

@@ -0,0 +1,17 @@
+"""Evaluation harness — measure RAG answer quality without an external API."""
+
+from ragforge.eval.metrics import (
+    answer_relevance,
+    context_recall,
+    evaluate,
+    faithfulness,
+)
+from ragforge.eval.report import EvalReport
+
+__all__ = [
+    "EvalReport",
+    "answer_relevance",
+    "context_recall",
+    "evaluate",
+    "faithfulness",
+]

ragforge/eval/metrics.py

@@ -0,0 +1,178 @@
+"""RAG quality metrics — embedding-based, pure-Python, no external API.
+
+All metrics return a float in ``[0, 1]`` (higher is better) and accept the
+same arguments so the orchestrator can iterate them uniformly:
+
+    metric(question, answer, contexts, ground_truth=None, *, encoder) -> float
+
+The three metrics implemented here are RAGAS-compatible interpretations:
+
+- ``context_recall``: of the n-grams in the ground-truth answer, what fraction
+  appear in any retrieved context block?
+- ``answer_relevance``: cosine similarity between the question and the
+  answer's embedding, with a length-penalty for empty/non-answers.
+- ``faithfulness``: of the claims (sentences) in the answer, what fraction
+  have a cosine similarity > τ with at least one context block?
+
+These are deliberately *embedding-based* rather than NLI-based: they need no
+extra model, run on CPU, and correlate strongly with judgment scores on
+short-answer QA. For long-form evaluation, plug in a stronger judge model.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from collections.abc import Iterable
+
+import numpy as np
+
+from ragforge.embed import Encoder
+
+logger = logging.getLogger("ragforge.eval")
+
+_WORD = re.compile(r"\w+", re.UNICODE)
+_SENT = re.compile(r"(?<=[.!?])\s+")
+
+
+def context_recall(
+    question: str,
+    answer: str,
+    contexts: list[str],
+    ground_truth: str | None,
+    *,
+    encoder: Encoder | None = None,
+    ngram: int = 1,
+) -> float:
+    """Fraction of ground-truth ngrams found in any retrieved context."""
+    if not ground_truth or not contexts:
+        return 0.0
+    gold = _ngrams(ground_truth.lower(), ngram)
+    if not gold:
+        return 0.0
+    joined = " ".join(contexts).lower()
+    present = sum(1 for g in gold if " ".join(g) in joined)
+    return present / len(gold)
+
+
+def answer_relevance(
+    question: str,
+    answer: str,
+    contexts: list[str],
+    ground_truth: str | None = None,
+    *,
+    encoder: Encoder,
+) -> float:
+    """Cosine similarity between question and answer embeddings.
+
+    Clamped to ``[0, 1]`` (negative cosine = unrelated, treated as 0).
+    "I don't know"-style answers are hard-capped at 0.2 even if they happen
+    to be lexically similar to the question.
+    """
+    if not answer.strip():
+        return 0.0
+    score = max(0.0, _cos_pair(encoder, question, answer))
+    if _looks_like_dont_know(answer):
+        return min(score, 0.2)
+    return min(score, 1.0)
+
+
+def faithfulness(
+    question: str,
+    answer: str,
+    contexts: list[str],
+    ground_truth: str | None = None,
+    *,
+    encoder: Encoder,
+    tau: float = 0.55,
+) -> float:
+    """Fraction of answer sentences supported by at least one context block.
+
+    A sentence is "supported" if its cosine similarity with at least one
+    retrieved context block crosses ``tau``. Tuned for ``bge-small`` on
+    factual short-answer QA; raise for stricter eval.
+    """
+    if not answer.strip() or not contexts:
+        return 0.0
+    sentences = [s.strip() for s in _SENT.split(answer.strip()) if s.strip()]
+    if not sentences:
+        return 0.0
+    sent_vecs = encoder.encode(sentences)
+    ctx_vecs = encoder.encode(contexts)
+    # cosine on normalized vectors == dot product
+    sims = sent_vecs @ ctx_vecs.T
+    supported = (sims.max(axis=1) >= tau).sum()
+    return float(supported) / len(sentences)
+
+
+# ---------------------------------------------------------------- orchestrator
+
+
+def evaluate(
+    samples: list[dict],
+    *,
+    encoder: Encoder,
+    metrics: Iterable[str] = ("context_recall", "answer_relevance", "faithfulness"),
+) -> dict:
+    """Run multiple metrics over a list of QA samples and return per-metric means.
+
+    Each sample is a dict with at least:
+        question: str
+        answer: str
+        contexts: list[str]
+        ground_truth: str | None (required for context_recall)
+    """
+    registry = {
+        "context_recall": context_recall,
+        "answer_relevance": answer_relevance,
+        "faithfulness": faithfulness,
+    }
+    by_metric: dict[str, list[float]] = {m: [] for m in metrics}
+
+    for s in samples:
+        q = s.get("question", "")
+        a = s.get("answer", "")
+        ctxs = s.get("contexts", [])
+        gt = s.get("ground_truth")
+        for m in metrics:
+            fn = registry[m]
+            score = fn(q, a, ctxs, gt, encoder=encoder)
+            by_metric[m].append(float(score))
+
+    means = {m: (sum(v) / len(v) if v else 0.0) for m, v in by_metric.items()}
+    return {
+        "n": len(samples),
+        "means": means,
+        "per_sample": by_metric,
+    }
+
+
+# ------------------------------------------------------------------ utilities
+
+
+def _ngrams(text: str, n: int) -> set[tuple[str, ...]]:
+    words = _WORD.findall(text)
+    if len(words) < n:
+        return set()
+    return {tuple(words[i : i + n]) for i in range(len(words) - n + 1)}
+
+
+def _cos_pair(encoder: Encoder, a: str, b: str) -> float:
+    vecs = encoder.encode([a, b])
+    return float(np.dot(vecs[0], vecs[1]))
+
+
+_DK_PATTERNS = (
+    "i don't know",
+    "i do not know",
+    "not enough information",
+    "cannot determine",
+    "no answer",
+    "n/a",
+    "based on the provided context",
+)
+
+
+def _looks_like_dont_know(text: str) -> bool:
+    t = text.lower()
+    return any(p in t for p in _DK_PATTERNS)

ragforge/eval/report.py

@@ -0,0 +1,57 @@
+"""Pretty-printable eval report."""
+
+from __future__ import annotations
+
+import json
+import statistics
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass
+class EvalReport:
+    n: int
+    means: dict[str, float]
+    per_sample: dict[str, list[float]]
+    latencies_ms: list[float] = field(default_factory=list)
+    extras: dict = field(default_factory=dict)
+
+    @property
+    def p50_ms(self) -> float:
+        return statistics.median(self.latencies_ms) if self.latencies_ms else 0.0
+
+    @property
+    def p95_ms(self) -> float:
+        if not self.latencies_ms:
+            return 0.0
+        s = sorted(self.latencies_ms)
+        return s[int(0.95 * (len(s) - 1))]
+
+    def as_table(self) -> str:
+        lines = [
+            "+----------------+--------+",
+            "|  metric        |  mean  |",
+            "+----------------+--------+",
+        ]
+        for name, val in self.means.items():
+            lines.append(f"| {name:<14s} |  {val:.3f} |")
+        lines.append("+----------------+--------+")
+        if self.latencies_ms:
+            lines.append(f"n={self.n}  ·  p50={self.p50_ms:.0f}ms  ·  p95={self.p95_ms:.0f}ms")
+        else:
+            lines.append(f"n={self.n}")
+        return "\n".join(lines)
+
+    def save(self, path: str | Path) -> None:
+        Path(path).write_text(
+            json.dumps(
+                {
+                    "n": self.n,
+                    "means": self.means,
+                    "per_sample": self.per_sample,
+                    "latencies_ms": self.latencies_ms,
+                    "extras": self.extras,
+                },
+                indent=2,
+            )
+        )

ragforge/ingest/__init__.py

@@ -0,0 +1,41 @@
+"""Document ingestion — load files into a stream of (text, metadata) pairs."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any
+
+from ragforge.ingest.markdown import load_markdown
+from ragforge.ingest.pdf import load_pdf
+
+Loader = dict[str, Any]
+
+
+def iter_documents(paths: list[str | Path]) -> Iterator[tuple[str, dict]]:
+    """Dispatch each path to the right loader based on extension.
+
+    Directories are walked recursively. Unsupported file types are skipped
+    silently — extending the registry is a one-line change.
+    """
+    for raw in paths:
+        p = Path(raw)
+        if p.is_dir():
+            for child in sorted(p.rglob("*")):
+                if child.is_file():
+                    yield from _load_one(child)
+        elif p.is_file():
+            yield from _load_one(p)
+
+
+def _load_one(path: Path) -> Iterator[tuple[str, dict]]:
+    suffix = path.suffix.lower()
+    if suffix == ".pdf":
+        yield from load_pdf(path)
+    elif suffix in {".md", ".markdown"}:
+        yield from load_markdown(path)
+    elif suffix in {".txt", ".rst"}:
+        yield path.read_text(encoding="utf-8", errors="ignore"), {"path": str(path)}
+
+
+__all__ = ["iter_documents", "load_markdown", "load_pdf"]

ragforge/ingest/chunking.py

@@ -0,0 +1,102 @@
+"""Recursive character splitter with overlap — a faithful, dependency-free port
+of the recipe that LangChain popularized.
+
+Why character-based rather than token-based?
+    Embedding models truncate at a token budget, but tokenizers vary. A
+    well-tuned character size (≈4 chars/token for English) is portable across
+    encoders and avoids re-tokenizing during chunking. For mixed-language
+    corpora, plug in a token-aware splitter; the interface is one function.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+_DEFAULT_SEPARATORS = ("\n\n", "\n", ". ", " ", "")
+
+
+@dataclass
+class Chunk:
+    text: str
+    metadata: dict
+
+
+def split(
+    text: str,
+    metadata: dict | None = None,
+    *,
+    size: int = 1024,
+    overlap: int = 128,
+    separators: Iterable[str] = _DEFAULT_SEPARATORS,
+) -> list[Chunk]:
+    """Recursively split ``text`` into windows of ~``size`` chars.
+
+    Each chunk overlaps the previous one by ``overlap`` chars so a query that
+    falls on the boundary still hits a chunk with full context. The recursion
+    tries successively finer separators (paragraph → line → sentence → word →
+    char) so we cut at a semantically reasonable boundary when possible.
+    """
+    metadata = metadata or {}
+    seps = list(separators)
+    raw_chunks = _recursive_split(text, seps, size)
+    merged = _merge_with_overlap(raw_chunks, size=size, overlap=overlap)
+    return [
+        Chunk(text=c, metadata={**metadata, "chunk": i, "n_chunks": len(merged)})
+        for i, c in enumerate(merged)
+    ]
+
+
+def split_documents(
+    docs: Iterable[tuple[str, dict]],
+    *,
+    size: int = 1024,
+    overlap: int = 128,
+) -> list[Chunk]:
+    out: list[Chunk] = []
+    for text, meta in docs:
+        out.extend(split(text, meta, size=size, overlap=overlap))
+    return out
+
+
+def _recursive_split(text: str, separators: list[str], size: int) -> list[str]:
+    if len(text) <= size:
+        return [text]
+    sep = separators[0]
+    rest = separators[1:] or [""]
+    if sep == "":
+        return [text[i : i + size] for i in range(0, len(text), size)]
+    parts = text.split(sep)
+    out: list[str] = []
+    buf = ""
+    for part in parts:
+        candidate = (buf + sep + part) if buf else part
+        if len(candidate) <= size:
+            buf = candidate
+            continue
+        if buf:
+            out.append(buf)
+        if len(part) > size:
+            out.extend(_recursive_split(part, rest, size))
+            buf = ""
+        else:
+            buf = part
+    if buf:
+        out.append(buf)
+    return out
+
+
+def _merge_with_overlap(chunks: list[str], *, size: int, overlap: int) -> list[str]:
+    """Sliding-window merge: ensures successive chunks share ``overlap`` chars."""
+    if overlap <= 0 or not chunks:
+        return chunks
+    out: list[str] = []
+    for c in chunks:
+        if out and len(out[-1]) + len(c) <= size + overlap:
+            out[-1] = out[-1] + c
+        else:
+            if out:
+                tail = out[-1][-overlap:]
+                c = tail + c if not c.startswith(tail) else c
+            out.append(c)
+    return out