diff --git a/.devcontainer/rag-lab/devcontainer.json b/.devcontainer/rag-lab/devcontainer.json new file mode 100644 index 0000000..c08d0f7 --- /dev/null +++ b/.devcontainer/rag-lab/devcontainer.json @@ -0,0 +1,16 @@ +{ + "name": "labagent RAG lab", + "build": { "dockerfile": "../../Dockerfile.lab", "context": "../.." }, + "remoteUser": "coder", + "workspaceFolder": "/home/coder/lab", + "forwardPorts": [8501], + "portsAttributes": { "8501": { "label": "RAG playground", "onAutoForward": "openBrowser" } }, + "customizations": { + "vscode": { + "extensions": ["ms-python.python", "ms-toolsai.jupyter"], + "settings": { "python.defaultInterpreterPath": "/usr/local/bin/python" } + } + }, + "postCreateCommand": "test -f .env || labagent rag init --dir .", + "containerEnv": { "HF_HUB_OFFLINE": "1" } +} diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..f90e795 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,9 @@ +.git +.venv +dist +*.egg-info +**/__pycache__ +**/.terraform +tests +docs +.pytest_cache diff --git a/Dockerfile.lab b/Dockerfile.lab new file mode 100644 index 0000000..514f51f --- /dev/null +++ b/Dockerfile.lab @@ -0,0 +1,47 @@ +# Web IDE + RAG lab image: VS Code in the browser (code-server) with Jupyter and every RAG dependency pre-installed. +# Build from the repo root: docker build -f Dockerfile.lab -t labagent-rag-lab:local . +# Run: labagent rag ide (or: docker run -p 127.0.0.1:8080:8080 -e PASSWORD=... labagent-rag-lab:local) +FROM python:3.11-bookworm + +ARG TARGETARCH +ARG CODE_SERVER_VERSION=4.141.0 +ARG CODE_SERVER_SHA256_AMD64=f384beb4f3ba207e0bca27d0dbd36dd825d5aa909d78c6d4baa8703eaf6b77ba +ARG CODE_SERVER_SHA256_ARM64=0bca33a9b01180ad88f0b8969a8e34ae1c76417df471697e7cd65c20ef5dd60d + +RUN apt-get update && apt-get install -y --no-install-recommends curl git sudo ca-certificates \ + && rm -rf /var/lib/apt/lists/* + +# code-server pinned by version and SHA-256: the build fails if the download changes. +RUN set -eux; \ + case "${TARGETARCH:-$(dpkg --print-architecture)}" in \ + amd64) sha="$CODE_SERVER_SHA256_AMD64"; arch=amd64 ;; \ + arm64) sha="$CODE_SERVER_SHA256_ARM64"; arch=arm64 ;; \ + *) echo "unsupported architecture"; exit 1 ;; \ + esac; \ + curl -fsSL -o /tmp/cs.deb "https://github.com/coder/code-server/releases/download/v${CODE_SERVER_VERSION}/code-server_${CODE_SERVER_VERSION}_${arch}.deb"; \ + echo "${sha} /tmp/cs.deb" | sha256sum -c -; \ + dpkg -i /tmp/cs.deb; rm /tmp/cs.deb + +RUN useradd -m -s /bin/bash -u 1000 coder && echo "coder ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/coder +ENV HF_HOME=/opt/hf-cache SENTENCE_TRANSFORMERS_HOME=/opt/hf-cache +# CPU-only torch keeps the image far smaller than the default CUDA wheels. +RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu + +# labagent from this checkout, with the lab extras (chromadb, sentence-transformers, groq, langchain, qdrant-client, jupyterlab ...). +COPY pyproject.toml README.md LICENSE MANIFEST.in /src/ +COPY labagent /src/labagent +RUN pip install --no-cache-dir "/src[rag-lab]" + +# Pre-download the embedding model so the first ingest has no wait. Learners can run offline afterwards. +RUN python -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('all-MiniLM-L6-v2')" \ + && chmod -R a+rX /opt/hf-cache + +USER coder +WORKDIR /home/coder/lab +RUN labagent rag init --dir /home/coder/lab && python -m ipykernel install --user --name python3 +# Extensions come from the Open VSX registry; a network failure here should not fail the whole build. +RUN code-server --install-extension ms-python.python --install-extension ms-toolsai.jupyter || echo "extensions not installed" + +EXPOSE 8080 +# PASSWORD comes from the environment (labagent rag ide sets a random one). Bound on all interfaces INSIDE the container only. +CMD ["code-server", "--bind-addr", "0.0.0.0:8080", "--auth", "password", "/home/coder/lab"] diff --git a/examples/01_rag_starter_lab.py b/examples/01_rag_starter_lab.py new file mode 100644 index 0000000..744e512 --- /dev/null +++ b/examples/01_rag_starter_lab.py @@ -0,0 +1,32 @@ +"""End-to-end local RAG: ingest, hybrid retrieval, answer. No cloud account needed. + + pip install 'labagent[rag-local]' + labagent rag init --dir rag-lab && cd rag-lab + python ../examples/01_rag_starter_lab.py + +Set GROQ_API_KEY in .env for real answers; without any key you get a clearly marked mock answer. +This is the laptop version of the labagent Databricks packs (`labagent launch rag-sec-filings`). +""" +from pathlib import Path + +from labagent.ai.rag import RAGLabClient +from labagent.ai.rag.config import load_dotenv + +load_dotenv(Path(".env")) +client = RAGLabClient() + +result = client.ingest_documents(["data"]) +print(f"indexed {result.chunks} chunks from {result.documents} documents ({result.embedding_model}, {result.vector_store})") +for w in result.warnings: + print("warning:", w) + +for question in ("Why does hybrid search help with product codes?", "How big should chunks be?"): + resp = client.query(question, top_k=3) + print(f"\nQ: {question}\nA: {resp.answer}") + for c in resp.chunks: + print(f" #{c.rank} {c.doc_id} similarity={c.similarity} keywords={c.keyword_score}") + m = resp.metrics + print(f" {m.provider}/{m.model} retrieval {m.retrieval_ms} ms generation {m.generation_ms} ms tokens {m.total_tokens}") + +print("\nvector inspection:", client.inspect_vectors(result.doc_ids[0])["items"][0]["embedding_head"]) +client.reset() diff --git a/labagent/ai/__init__.py b/labagent/ai/__init__.py new file mode 100644 index 0000000..73656bd --- /dev/null +++ b/labagent/ai/__init__.py @@ -0,0 +1 @@ +"""AI lab helpers that run on a laptop at zero cost (extra: labagent[rag-local]).""" diff --git a/labagent/ai/rag/__init__.py b/labagent/ai/rag/__init__.py new file mode 100644 index 0000000..b9c5cc1 --- /dev/null +++ b/labagent/ai/rag/__init__.py @@ -0,0 +1,15 @@ +"""Local, zero-cost RAG lab SDK. Install with: pip install 'labagent[rag-local]'. + +Exports load lazily so `labagent rag init` and the base CLI work without pydantic installed. +""" +import importlib + +_EXPORTS = {"RAGLabClient": ".client", "chunk_text": ".client", "RAGSettings": ".config", + "RAGResponse": ".models", "EvalResult": ".models", "IngestionResult": ".models"} +__all__ = sorted(_EXPORTS) + + +def __getattr__(name): + if name in _EXPORTS: + return getattr(importlib.import_module(_EXPORTS[name], __name__), name) + raise AttributeError(name) diff --git a/labagent/ai/rag/adapters.py b/labagent/ai/rag/adapters.py new file mode 100644 index 0000000..ec405f2 --- /dev/null +++ b/labagent/ai/rag/adapters.py @@ -0,0 +1,236 @@ +"""Vector store adapters. Add a store by subclassing BaseVectorAdapter and registering it in make_adapter(). + +Shipped: ChromaAdapter (default, embedded, free), MemoryAdapter (no dependencies), QdrantAdapter (remote or local). +Pinecone and pgvector fit the same four methods. For managed pgvector on Lakebase use the pack instead: +`labagent launch rag-sec-filings`. +""" +import math +import uuid +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional + +from .config import RAGSettings +from .models import VectorHit + + +class VectorStoreError(RuntimeError): + pass + + +Where = Optional[Dict[str, Any]] + + +def norm_meta(meta: Optional[Dict[str, Any]]) -> Dict[str, str]: + """Stores keep metadata as strings so filters behave the same everywhere (fiscal_year 2025 == "2025").""" + return {str(k): str(v) for k, v in (meta or {}).items() if v is not None} + + +def matches(meta: Dict[str, Any], where: Where) -> bool: + """where = {"ticker": "AAPL"} or {"ticker": ["AAPL", "WMT"]} (any of). Several keys must all match.""" + for k, v in (where or {}).items(): + allowed = [str(x) for x in v] if isinstance(v, (list, tuple, set)) else [str(v)] + if str(meta.get(k)) not in allowed: + return False + return True + + +class BaseVectorAdapter(ABC): + name = "base" + + @abstractmethod + def upsert(self, ids: List[str], embeddings: List[List[float]], documents: List[str], doc_ids: List[str], + metadatas: Optional[List[Dict[str, Any]]] = None, titles: Optional[List[Optional[str]]] = None) -> None: ... + + @abstractmethod + def query(self, embedding: List[float], top_k: int, where: Where = None) -> List[VectorHit]: + """Nearest chunks by cosine similarity, best first. `where` keeps only chunks whose metadata matches.""" + + @abstractmethod + def get_doc(self, doc_id: str) -> List[VectorHit]: + """Every chunk of one document, with embeddings.""" + + @abstractmethod + def reset(self) -> None: ... + + @abstractmethod + def count(self) -> int: ... + + +def _cosine(a: List[float], b: List[float]) -> float: + dot = sum(x * y for x, y in zip(a, b)) + na, nb = math.sqrt(sum(x * x for x in a)), math.sqrt(sum(y * y for y in b)) + return dot / (na * nb) if na and nb else 0.0 + + +class MemoryAdapter(BaseVectorAdapter): + """Pure Python store. Fine for a few thousand chunks; use it when chromadb is not installed.""" + name = "memory" + + def __init__(self) -> None: + self._rows: Dict[str, VectorHit] = {} + + def upsert(self, ids, embeddings, documents, doc_ids, metadatas=None, titles=None): + metadatas = metadatas or [{}] * len(ids) + titles = titles or [None] * len(ids) + for i, e, d, doc, m, t in zip(ids, embeddings, documents, doc_ids, metadatas, titles): + self._rows[i] = VectorHit(chunk_id=i, doc_id=doc, text=d, similarity=0.0, embedding=e, title=t, metadata=norm_meta(m)) + + def query(self, embedding, top_k, where=None): + scored = [r.model_copy(update={"similarity": _cosine(embedding, r.embedding or []), "embedding": None}) + for r in self._rows.values() if matches(r.metadata, where)] + return sorted(scored, key=lambda r: r.similarity, reverse=True)[:top_k] + + def get_doc(self, doc_id): + return [r for r in self._rows.values() if r.doc_id == doc_id] + + def reset(self): + self._rows.clear() + + def count(self): + return len(self._rows) + + +class ChromaAdapter(BaseVectorAdapter): + """Embedded ChromaDB. In memory unless settings.chroma_path is set.""" + name = "chroma" + + def __init__(self, settings: RAGSettings): + try: + import chromadb + except ImportError as e: + raise VectorStoreError("chromadb is not installed. Fix: pip install 'labagent[rag-local]'") from e + self._client = chromadb.PersistentClient(path=settings.chroma_path) if settings.chroma_path else chromadb.EphemeralClient() + self._collection_name = settings.collection + self._col = self._open() + + def _open(self): + return self._client.get_or_create_collection(self._collection_name, metadata={"hnsw:space": "cosine"}) + + @staticmethod + def _where(where: Where): + conds = [{k: ({"$in": [str(x) for x in v]} if isinstance(v, (list, tuple, set)) else str(v))} for k, v in (where or {}).items()] + return None if not conds else conds[0] if len(conds) == 1 else {"$and": conds} + + @staticmethod + def _split(meta: Dict[str, Any]): + meta = dict(meta) + return meta.pop("doc_id"), meta.pop("_title", None), meta + + def upsert(self, ids, embeddings, documents, doc_ids, metadatas=None, titles=None): + metadatas = metadatas or [{}] * len(ids) + titles = titles or [None] * len(ids) + metas = [{**norm_meta(m), "doc_id": d, **({"_title": t} if t else {})} for m, d, t in zip(metadatas, doc_ids, titles)] + self._col.upsert(ids=ids, embeddings=embeddings, documents=documents, metadatas=metas) + + def query(self, embedding, top_k, where=None): + n = self._col.count() + if n == 0: + return [] + r = self._col.query(query_embeddings=[embedding], n_results=min(top_k, n), include=["documents", "metadatas", "distances"], where=self._where(where)) + out = [] + for i, d, m, dist in zip(r["ids"][0], r["documents"][0], r["metadatas"][0], r["distances"][0]): + doc, title, meta = self._split(m) + out.append(VectorHit(chunk_id=i, doc_id=doc, text=d, similarity=1.0 - dist, title=title, metadata=meta)) + return out + + def get_doc(self, doc_id): + r = self._col.get(where={"doc_id": doc_id}, include=["documents", "embeddings", "metadatas"]) + embeddings = r["embeddings"] if r["embeddings"] is not None else [] + out = [] + for i, d, e, m in zip(r["ids"], r["documents"], embeddings, r["metadatas"]): + _, title, meta = self._split(m) + out.append(VectorHit(chunk_id=i, doc_id=doc_id, text=d, similarity=1.0, embedding=[float(x) for x in e], title=title, metadata=meta)) + return out + + def reset(self): + self._client.delete_collection(self._collection_name) + self._col = self._open() + + def count(self): + return self._col.count() + + +class QdrantAdapter(BaseVectorAdapter): + """Qdrant over HTTP (Qdrant Cloud or a local container). With no URL it uses an in-process instance. + + Not exercised against a live server in labagent's tests: treat the first run as a smoke test. + """ + name = "qdrant" + + def __init__(self, settings: RAGSettings, dim: int): + try: + from qdrant_client import QdrantClient + except ImportError as e: + raise VectorStoreError("qdrant-client is not installed. Fix: pip install qdrant-client") from e + self._client = QdrantClient(url=settings.qdrant_url, api_key=settings.qdrant_api_key) if settings.qdrant_url else QdrantClient(":memory:") + self._collection = settings.collection + self._dim = dim + self._ensure() + + def _ensure(self): + from qdrant_client.models import Distance, VectorParams + existing = {c.name for c in self._client.get_collections().collections} + if self._collection not in existing: + self._client.create_collection(self._collection, vectors_config=VectorParams(size=self._dim, distance=Distance.COSINE)) + + @staticmethod + def _pid(chunk_id: str) -> str: + return str(uuid.uuid5(uuid.NAMESPACE_URL, chunk_id)) # Qdrant point ids must be UUIDs or ints + + def upsert(self, ids, embeddings, documents, doc_ids, metadatas=None, titles=None): + from qdrant_client.models import PointStruct + metadatas = metadatas or [{}] * len(ids) + titles = titles or [None] * len(ids) + points = [PointStruct(id=self._pid(i), vector=e, payload={"chunk_id": i, "doc_id": d, "text": t, "title": ti, "meta": norm_meta(m)}) + for i, e, t, d, m, ti in zip(ids, embeddings, documents, doc_ids, metadatas, titles)] + self._client.upsert(self._collection, points=points) + + @staticmethod + def _filter(where: Where): + from qdrant_client.models import FieldCondition, Filter, MatchAny, MatchValue + if not where: + return None + return Filter(must=[FieldCondition(key=f"meta.{k}", match=(MatchAny(any=[str(x) for x in v]) if isinstance(v, (list, tuple, set)) else MatchValue(value=str(v)))) + for k, v in where.items()]) + + @staticmethod + def _hit(p, similarity=None, embedding=None): + return VectorHit(chunk_id=p.payload["chunk_id"], doc_id=p.payload["doc_id"], text=p.payload["text"], title=p.payload.get("title"), + metadata=p.payload.get("meta") or {}, similarity=float(p.score) if similarity is None else similarity, embedding=embedding) + + def query(self, embedding, top_k, where=None): + flt = self._filter(where) + if hasattr(self._client, "query_points"): + res = self._client.query_points(self._collection, query=embedding, limit=top_k, with_payload=True, query_filter=flt).points + else: # older qdrant-client + res = self._client.search(self._collection, query_vector=embedding, limit=top_k, with_payload=True, query_filter=flt) + return [self._hit(p) for p in res] + + def get_doc(self, doc_id): + from qdrant_client.models import FieldCondition, Filter, MatchValue + points, _ = self._client.scroll(self._collection, scroll_filter=Filter(must=[FieldCondition(key="doc_id", match=MatchValue(value=doc_id))]), + limit=10000, with_payload=True, with_vectors=True) + return [self._hit(p, similarity=1.0, embedding=list(p.vector)) for p in points] + + def reset(self): + self._client.delete_collection(self._collection) + self._ensure() + + def count(self): + return self._client.count(self._collection, exact=True).count + + +def make_adapter(settings: RAGSettings, dim: int, warnings: Optional[List[str]] = None) -> BaseVectorAdapter: + kind = settings.vector_db_type.lower() + if kind == "memory": + return MemoryAdapter() + if kind == "qdrant": + return QdrantAdapter(settings, dim) + if kind == "chroma": + try: + return ChromaAdapter(settings) + except VectorStoreError as e: + if warnings is not None: + warnings.append(f"{e} Using the in-memory store for now; data is lost when the process exits.") + return MemoryAdapter() + raise VectorStoreError(f"unknown VECTOR_DB_TYPE '{settings.vector_db_type}': use chroma, qdrant or memory") diff --git a/labagent/ai/rag/cli.py b/labagent/ai/rag/cli.py new file mode 100644 index 0000000..9a59b07 --- /dev/null +++ b/labagent/ai/rag/cli.py @@ -0,0 +1,124 @@ +"""`labagent rag init | ide | playground`. Wired into the main parser by labagent.cli (argparse, like launch/destroy).""" +import os +import secrets +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Dict, List, Optional + +from . import samples + +IMAGE = "labagent-rag-lab:local" +CONTAINER_LAB = "/home/coder/lab" + + +def add_parser(sub) -> None: + rp = sub.add_parser("rag", help="Local, zero-cost RAG lab: init a workspace, open the web IDE (ide), or the visual playground") + rs = rp.add_subparsers(dest="rag_cmd", required=True) + i = rs.add_parser("init", help="Create .env, data/sample_docs, notebooks/ and src/main.py") + i.add_argument("--dir", default=".", help="Target folder (default: current)") + i.add_argument("--force", action="store_true", help="Overwrite files that already exist") + d = rs.add_parser("ide", help="Start a browser VS Code (code-server) on the workspace") + d.add_argument("--dir", default=".", help="Workspace folder (initialised if it has no .env)") + d.add_argument("--port", type=int, default=8080) + d.add_argument("--mode", choices=("auto", "local", "docker"), default="auto", + help="local runs the code-server on your PATH; docker runs the lab image; auto prefers docker") + d.add_argument("--image", default=IMAGE, help="Docker image to run (built from Dockerfile.lab in a labagent checkout if missing)") + d.add_argument("--print-command", action="store_true", help="Show the command and exit") + s = rs.add_parser("playground", aliases=["serve"], help="Visual playground on port 8501: chunks, scores and answers side by side") + s.add_argument("--dir", default=".", help="Project folder with .env and data/") + s.add_argument("--port", type=int, default=8501) + + +def init(target: Path, force: bool = False) -> int: + files = {target / rel: content for rel, content in samples.workspace().items()} + clash = [str(p) for p in files if p.exists()] + if clash and not force: + print("error: would overwrite " + ", ".join(clash) + " (use --force)", file=sys.stderr) + return 2 + for p, content in files.items(): + p.parent.mkdir(parents=True, exist_ok=True) + p.write_bytes(content if isinstance(content, bytes) else content.encode()) + (target / ".env").chmod(0o600) + print(f"created {len(files)} files in {target.resolve()}") + print(f"next: labagent rag ide --dir {target}" if str(target) != "." else "next: labagent rag ide") + return 0 + + +def _find_dockerfile() -> Optional[Path]: + """Dockerfile.lab sits at the root of a labagent checkout (three levels above this file). Not shipped in the wheel.""" + root = Path(__file__).resolve().parents[3] + return root / "Dockerfile.lab" if (root / "Dockerfile.lab").is_file() else None + + +def ide_command(mode: str, workspace: Path, port: int, password: str, image: str = IMAGE) -> List[str]: + """The command that serves the web IDE. The port is published on 127.0.0.1 only, and a password is always required.""" + if mode == "local": + return ["code-server", "--bind-addr", f"127.0.0.1:{port}", "--auth", "password", str(workspace)] + return ["docker", "run", "--rm", "-p", f"127.0.0.1:{port}:8080", "-e", "PASSWORD", "-v", f"{workspace}:{CONTAINER_LAB}", image] + + +def _pick_mode(mode: str) -> str: + if mode != "auto": + return mode + return "docker" if shutil.which("docker") else "local" + + +def _docker_ready(image: str) -> Optional[str]: + """None when docker can run `image` (building it if needed), else an error message.""" + if subprocess.run(["docker", "info"], capture_output=True).returncode != 0: + return "docker is installed but its daemon is not running. Start Docker Desktop and retry." + if subprocess.run(["docker", "image", "inspect", image], capture_output=True).returncode == 0: + return None + dockerfile = _find_dockerfile() + if dockerfile is None: + return (f"image '{image}' not found and no Dockerfile.lab is nearby. Run this from a labagent checkout " + "(git clone https://github.com/becloudready/labagent) or pass --image with a prebuilt image.") + print(f"building {image} from {dockerfile} (first time takes several minutes) ...") + rc = subprocess.call(["docker", "build", "-f", str(dockerfile), "-t", image, str(dockerfile.parent)]) + return None if rc == 0 else "docker build failed" + + +def ide(workspace: Path, port: int, mode: str, image: str, print_only: bool = False) -> int: + mode = _pick_mode(mode) + if mode == "local" and not shutil.which("code-server"): + print("error: code-server is not on PATH. Install it (https://coder.com/docs/code-server/install) or use --mode docker.", file=sys.stderr) + return 2 + if not (workspace / ".env").exists(): + rc = init(workspace) + if rc: + return rc + workspace = workspace.resolve() + password = os.environ.get("CODE_SERVER_PASSWORD") or secrets.token_urlsafe(12) + cmd = ide_command(mode, workspace, port, password, image) + if print_only: + print(" ".join(cmd)) + return 0 + if mode == "docker": + err = _docker_ready(image) + if err: + print(f"error: {err}", file=sys.stderr) + return 2 + print(f"web IDE: http://127.0.0.1:{port}\npassword: {password}\nnotebook: notebooks/01_build_your_first_rag.ipynb (Ctrl+C to stop)", flush=True) + return subprocess.call(cmd, env={**os.environ, "PASSWORD": password}) + + +def playground(target: Path, port: int) -> int: + try: + import streamlit # noqa: F401 + except ImportError: + print("error: streamlit is not installed. Fix: pip install 'labagent[rag-local]'", file=sys.stderr) + return 2 + script = Path(__file__).with_name("playground.py") + cmd = [sys.executable, "-m", "streamlit", "run", str(script), "--server.port", str(port), "--server.headless", "true", "--", "--dir", str(target.resolve())] + print(f"playground: http://localhost:{port}") + return subprocess.call(cmd) + + +def run(args) -> int: + if args.rag_cmd == "init": + return init(Path(args.dir), args.force) + if args.rag_cmd == "ide": + return ide(Path(args.dir), args.port, args.mode, args.image, args.print_command) + return playground(Path(args.dir), args.port) diff --git a/labagent/ai/rag/client.py b/labagent/ai/rag/client.py new file mode 100644 index 0000000..281c909 --- /dev/null +++ b/labagent/ai/rag/client.py @@ -0,0 +1,233 @@ +"""RAGLabClient: the high-level SDK for the local RAG lab. + + client = RAGLabClient() + client.ingest_documents(["data/"]) + print(client.query("What is a vector database?").answer) + +How it connects to labagent: `labagent rag init` scaffolds a project that uses this client, `labagent rag serve` shows +it in a browser, and the Databricks packs (`labagent launch rag-sec-filings`) deploy the same pipeline to Lakebase. +""" +import asyncio +import re +import time +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union + +from .adapters import BaseVectorAdapter, make_adapter +from .config import RAGSettings +from .embeddings import Embedder, make_embedder +from .llm import LLMGateway +from .adapters import Where, matches +from .models import EvalResult, IngestionResult, Metrics, RAGResponse, RetrievedChunk + +TEXT_SUFFIXES = {".txt", ".md", ".rst"} +PDF_SUFFIX = ".pdf" +_WORD = re.compile(r"[a-z0-9]+") +_STOP = frozenset("a an the of to in on for and or is are was be what how why when which who does do it its with as by at from this that".split()) +EMBED_BATCH = 64 + + +def chunk_text(text: str, size: int = 800, overlap: int = 100) -> List[str]: + """Paragraph-aware chunks of about `size` characters; a tail of `overlap` characters carries into the next chunk.""" + paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()] + chunks: List[str] = [] + cur = "" + for p in paragraphs: + while len(p) > size: # a single huge paragraph: cut at a space near the limit + cut = max(p.rfind(" ", 0, size), size // 2) + chunks.append((cur + " " + p[:cut + 1]).strip() if cur else p[:cut + 1].strip()) + cur, p = "", p[cut + 1:].lstrip() + if cur and len(cur) + len(p) + 2 > size: + chunks.append(cur) + cur = cur[-overlap:] if overlap else "" + cur = (cur + "\n\n" + p).strip() if cur else p + if cur: + chunks.append(cur) + return [c for c in chunks if c] + + +def _slug(path: Path) -> str: + return re.sub(r"[^a-z0-9]+", "-", path.stem.lower()).strip("-") or "doc" + + +def _keywords(text: str) -> List[str]: + return [w for w in _WORD.findall(text.lower()) if w not in _STOP] + + +def _pdf_text(path: Path) -> str: + from pypdf import PdfReader # lazy: only needed for PDFs + return "\n\n".join((page.extract_text() or "") for page in PdfReader(str(path)).pages) + + +def _has_pypdf() -> bool: + try: + import pypdf # noqa: F401 + return True + except ImportError: + return False + + +def _expand(sources: Sequence[Union[str, Path]]) -> Tuple[List[Path], Dict[str, str]]: + files: List[Path] = [] + skipped: Dict[str, str] = {} + pdf_ok = _has_pypdf() + ok = TEXT_SUFFIXES | ({PDF_SUFFIX} if pdf_ok else set()) + + def note_pdf(f: Path) -> None: + if f.suffix.lower() == PDF_SUFFIX and not pdf_ok: + skipped[str(f)] = "PDF support needs pypdf: pip install pypdf" + + for s in sources: + p = Path(s).expanduser() + if p.is_dir(): + for f in sorted(x for x in p.rglob("*") if x.is_file()): + if f.suffix.lower() in ok: + files.append(f) + else: + note_pdf(f) + elif p.is_file(): + if p.suffix.lower() in ok: + files.append(p) + else: + skipped[str(p)] = (f"unsupported type '{p.suffix}' (supported: {', '.join(sorted(ok))})" if p.suffix.lower() != PDF_SUFFIX + else "PDF support needs pypdf: pip install pypdf") + else: + skipped[str(p)] = "not found" + return files, skipped + + +class RAGLabClient: + def __init__(self, vector_db: Optional[str] = None, llm_provider: Optional[str] = None, settings: Optional[RAGSettings] = None, + embedder: Optional[Embedder] = None, store: Optional[BaseVectorAdapter] = None, llm: Optional[LLMGateway] = None): + """vector_db ("chroma" | "qdrant" | "memory") and llm_provider ("groq" | "ollama" | "openai" | "anthropic" | "mock") + override the environment. Anything left out comes from RAGSettings.from_env(); a missing API key never raises.""" + self.settings = settings or RAGSettings.from_env() + override = {k: v for k, v in (("vector_db_type", vector_db), ("llm_provider", llm_provider)) if v} + if override: + self.settings = self.settings.model_copy(update=override) + self.warnings: List[str] = [] + self.embedder = embedder or make_embedder(self.settings, self.warnings) + self.store = store or make_adapter(self.settings, self.embedder.dim, self.warnings) + self.llm = llm or LLMGateway(self.settings) + + # ---- ingestion ------------------------------------------------------------------------------------------------- + def _read_chunks(self, path: Path) -> Tuple[str, List[str]]: + text = _pdf_text(path) if path.suffix.lower() == PDF_SUFFIX else path.read_text(errors="replace") + return _slug(path), chunk_text(text, self.settings.chunk_chars, self.settings.chunk_overlap) + + def _index(self, per_doc: List[Tuple[str, List[str]]]) -> Tuple[int, List[str]]: + ids: List[str] = [] + texts: List[str] = [] + docs: List[str] = [] + seen: Dict[str, int] = {} + for doc_id, chunks in per_doc: + seen[doc_id] = seen.get(doc_id, 0) + 1 + if seen[doc_id] > 1: # two files with the same stem + doc_id = f"{doc_id}-{seen[doc_id]}" + for n, c in enumerate(chunks): + ids.append(f"{doc_id}::{n}") + texts.append(c) + docs.append(doc_id) + for i in range(0, len(texts), EMBED_BATCH): + b = slice(i, i + EMBED_BATCH) + self.store.upsert(ids[b], self.embedder.embed(texts[b]), texts[b], docs[b]) + return len(texts), sorted(set(docs)) + + def _result(self, t0: float, n_docs: int, n_chunks: int, skipped: Dict[str, str], doc_ids: List[str]) -> IngestionResult: + return IngestionResult(documents=n_docs, chunks=n_chunks, skipped=skipped, elapsed_ms=round((time.perf_counter() - t0) * 1000, 1), + embedding_model=self.embedder.name, vector_store=self.store.name, warnings=list(self.warnings), doc_ids=doc_ids) + + def ingest(self, sources: Union[str, Path, Sequence[Union[str, Path]]]) -> IngestionResult: + """Short form of ingest_documents(): accepts one path or a list.""" + return self.ingest_documents([sources] if isinstance(sources, (str, Path)) else sources) + + def ingest_documents(self, sources: Sequence[Union[str, Path]]) -> IngestionResult: + """Read .txt/.md/.rst (and .pdf when pypdf is installed) files or folders, chunk, embed and index them. Re-ingesting a file overwrites its chunks; call reset() to start clean.""" + t0 = time.perf_counter() + files, skipped = _expand(sources) + n, doc_ids = self._index([self._read_chunks(f) for f in files]) + return self._result(t0, len(files), n, skipped, doc_ids) + + async def aingest_documents(self, sources: Sequence[Union[str, Path]]) -> IngestionResult: + """Async variant: files are read and chunked concurrently in worker threads; embedding runs off the event loop.""" + t0 = time.perf_counter() + files, skipped = _expand(sources) + per_doc = await asyncio.gather(*(asyncio.to_thread(self._read_chunks, f) for f in files)) + n, doc_ids = await asyncio.to_thread(self._index, list(per_doc)) + return self._result(t0, len(files), n, skipped, doc_ids) + + def ingest_chunks(self, rows: Sequence[Dict[str, Any]], progress: Optional[Callable[[int, int], None]] = None) -> IngestionResult: + """Index pre-chunked rows: {"id", "content" (or "text"), "title"?, "meta"?, "doc_id"?}. Use this for datasets such as the + SEC 10-K snapshot (labagent.datasets.fetch_snapshot). `meta` becomes filterable. doc_id defaults to the id minus a trailing -N.""" + t0 = time.perf_counter() + ids, texts, docs, metas, titles = [], [], [], [], [] + for r in rows: + text = r.get("content") or r.get("text") + if not text: + continue + ids.append(str(r["id"])) + texts.append(text) + docs.append(str(r.get("doc_id") or re.sub(r"-\d+$", "", str(r["id"])))) + metas.append(r.get("meta") or {}) + titles.append(r.get("title")) + for i in range(0, len(texts), EMBED_BATCH): + b = slice(i, i + EMBED_BATCH) + self.store.upsert(ids[b], self.embedder.embed(texts[b]), texts[b], docs[b], metas[b], titles[b]) + if progress: + progress(min(i + EMBED_BATCH, len(texts)), len(texts)) + return self._result(t0, len(set(docs)), len(texts), {}, sorted(set(docs))) + + # ---- query ----------------------------------------------------------------------------------------------------- + def retrieve(self, prompt: str, top_k: int = 3, where: Where = None, keyword_weight: float = 0.3, + candidates: Optional[int] = None) -> List[RetrievedChunk]: + """Hybrid retrieval. Fetch `candidates` nearest chunks by vector (default 4x top_k), keep those whose metadata matches `where`, + then rank by (1 - keyword_weight) * similarity + keyword_weight * keyword overlap. keyword_weight=0 is pure vector search.""" + hits = self.store.query(self.embedder.embed([prompt])[0], candidates or max(top_k * 4, top_k), where) + words = set(_keywords(prompt)) + scored = [] + for h in hits: + kw = len(words & set(_keywords(h.text + " " + (h.title or "")))) / len(words) if words else 0.0 + scored.append(((1 - keyword_weight) * h.similarity + keyword_weight * kw, kw, h)) + scored.sort(key=lambda t: t[0], reverse=True) + return [RetrievedChunk(doc_id=h.doc_id, chunk_id=h.chunk_id, text=h.text, similarity=round(h.similarity, 4), distance=round(1 - h.similarity, 4), + keyword_score=round(kw, 3), rank=i + 1, title=h.title, metadata=h.metadata) for i, (_, kw, h) in enumerate(scored[:top_k])] + + def evaluate(self, cases: Sequence[Dict[str, Any]], top_k: int = 5, min_hit_rate: Optional[float] = None, **retrieve_kw: Any) -> EvalResult: + """Retrieval quality check. A case is {"question", "expect": {metadata subset}} and hits when any top_k chunk matches `expect`. + Extra keyword arguments (where=, keyword_weight=) go to retrieve(), so you can compare settings on the same questions.""" + hits, misses = 0, [] + for c in cases: + got = self.retrieve(c["question"], top_k, **retrieve_kw) + if any(matches(g.metadata, c["expect"]) for g in got): + hits += 1 + else: + misses.append({"question": c["question"], "expected": c["expect"], "got": [g.chunk_id for g in got]}) + rate = hits / len(cases) if cases else 0.0 + return EvalResult(cases=len(cases), hits=hits, hit_rate=round(rate, 3), misses=misses, passed=None if min_hit_rate is None else rate >= min_hit_rate) + + def query(self, prompt: str, top_k: int = 3, where: Where = None, keyword_weight: float = 0.3) -> RAGResponse: + if self.store.count() == 0: + raise ValueError("nothing is indexed yet: call ingest_documents() first (or run `labagent rag init` for sample data)") + t0 = time.perf_counter() + chunks = self.retrieve(prompt, top_k, where, keyword_weight) + t1 = time.perf_counter() + context = "\n\n".join(f"[{c.doc_id}] {c.title + chr(10) if c.title else ''}{c.text}" for c in chunks) + res = self.llm.complete(prompt, context) + t2 = time.perf_counter() + m = Metrics(retrieval_ms=round((t1 - t0) * 1000, 1), generation_ms=round((t2 - t1) * 1000, 1), prompt_tokens=res.prompt_tokens, + completion_tokens=res.completion_tokens, total_tokens=res.prompt_tokens + res.completion_tokens, tokens_estimated=res.estimated, + provider=res.provider, model=res.model) + return RAGResponse(answer=res.text, chunks=chunks, metrics=m, warnings=self.warnings + res.warnings) + + # ---- inspection ------------------------------------------------------------------------------------------------ + def inspect_vectors(self, doc_id: str) -> dict: + """What the store holds for one document: chunk ids, text and the first 8 numbers of each embedding.""" + rows = self.store.get_doc(doc_id) + if not rows: + raise KeyError(f"no document '{doc_id}' (indexed: use result.doc_ids from ingest_documents)") + return {"doc_id": doc_id, "chunks": len(rows), "embedding_dim": len(rows[0].embedding or []), "embedding_model": self.embedder.name, + "vector_store": self.store.name, + "items": [{"chunk_id": r.chunk_id, "text": r.text, "embedding_head": [round(x, 4) for x in (r.embedding or [])[:8]]} for r in rows]} + + def reset(self) -> None: + self.store.reset() diff --git a/labagent/ai/rag/config.py b/labagent/ai/rag/config.py new file mode 100644 index 0000000..086fbeb --- /dev/null +++ b/labagent/ai/rag/config.py @@ -0,0 +1,81 @@ +"""Settings for the local RAG lab, read from environment variables (and an optional .env file). + +Connects to the labagent ecosystem: the Databricks packs (`labagent launch rag-*`) deploy RAG to a managed Lakebase +database. This module is the free, laptop-only sibling: same ideas (chunk, embed, retrieve, answer) with an embedded +vector store, so a learner can experiment before touching any cloud account. +""" +import os +from pathlib import Path +from typing import Dict, Optional + +from pydantic import BaseModel, Field + +DEFAULT_EMBEDDING_MODEL = "all-MiniLM-L6-v2" +GROQ_URL = "https://api.groq.com/openai/v1" +OPENAI_URL = "https://api.openai.com/v1" +ANTHROPIC_URL = "https://api.anthropic.com/v1" +OLLAMA_URL = "http://localhost:11434/v1" + + +def load_dotenv(path: Path) -> Dict[str, str]: + """Load KEY=VALUE lines into os.environ without overriding variables that are already set. Returns what was set.""" + set_now: Dict[str, str] = {} + if not path.is_file(): + return set_now + for line in path.read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + key, value = key.strip(), value.strip() + if key.startswith("export "): + key = key[len("export "):].strip() + if value[:1] in ("'", '"') and value[-1:] == value[:1] and len(value) > 1: + value = value[1:-1] + else: + value = value.split(" #", 1)[0].strip() # inline comment on an unquoted value + if key and value and not os.environ.get(key): # an existing empty variable counts as unset + os.environ[key] = value + set_now[key] = value + return set_now + + +class RAGSettings(BaseModel): + """All knobs for RAGLabClient. Build with RAGSettings.from_env(); API keys never appear in repr or logs.""" + + vector_db_type: str = Field("chroma", description="chroma | qdrant | memory") + embedding_model_name: str = DEFAULT_EMBEDDING_MODEL + chroma_path: Optional[str] = Field(None, description="Persist Chroma here. Unset keeps everything in memory.") + qdrant_url: Optional[str] = None + qdrant_api_key: Optional[str] = Field(None, repr=False) + collection: str = "labagent_rag" + chunk_chars: int = 800 + chunk_overlap: int = 100 + llm_provider: Optional[str] = Field(None, description="groq | ollama | openai | anthropic | mock. Unset picks the first usable.") + groq_api_key: Optional[str] = Field(None, repr=False) + groq_model: str = "openai/gpt-oss-20b" + openai_api_key: Optional[str] = Field(None, repr=False) + openai_model: str = "gpt-4o-mini" + anthropic_api_key: Optional[str] = Field(None, repr=False) + anthropic_model: str = "claude-haiku-4-5-20251001" + ollama_url: str = OLLAMA_URL + ollama_model: str = "qwen2.5-coder:0.5b" + llm_timeout_s: float = 60.0 + + @classmethod + def from_env(cls, env: Optional[Dict[str, str]] = None) -> "RAGSettings": + e = os.environ if env is None else env + + def g(name: str) -> Optional[str]: + return e.get(name) or None + + values = { + "vector_db_type": g("VECTOR_DB_TYPE"), "embedding_model_name": g("EMBEDDING_MODEL_NAME"), + "chroma_path": g("CHROMA_PATH"), "qdrant_url": g("QDRANT_URL"), "qdrant_api_key": g("QDRANT_API_KEY"), + "collection": g("RAG_COLLECTION"), "llm_provider": g("LLM_PROVIDER"), + "groq_api_key": g("GROQ_API_KEY"), "groq_model": g("GROQ_MODEL"), + "openai_api_key": g("OPENAI_API_KEY"), "openai_model": g("OPENAI_MODEL"), + "anthropic_api_key": g("ANTHROPIC_API_KEY"), "anthropic_model": g("ANTHROPIC_MODEL"), + "ollama_url": g("OLLAMA_URL"), "ollama_model": g("OLLAMA_MODEL"), + } + return cls(**{k: v for k, v in values.items() if v is not None}) diff --git a/labagent/ai/rag/embeddings.py b/labagent/ai/rag/embeddings.py new file mode 100644 index 0000000..d4bea1b --- /dev/null +++ b/labagent/ai/rag/embeddings.py @@ -0,0 +1,63 @@ +"""Embedding providers. Default is a local CPU model: no per-token cost, no API key.""" +import hashlib +import math +import re +from typing import List, Optional, Protocol + +from .config import RAGSettings + +_WORD = re.compile(r"[a-z0-9]+") + + +class Embedder(Protocol): + name: str + dim: int + + def embed(self, texts: List[str]) -> List[List[float]]: ... + + +class SentenceTransformerEmbedder: + """sentence-transformers on CPU (all-MiniLM-L6-v2, 384 dims). The model downloads once, then runs offline.""" + + def __init__(self, model_name: str): + from sentence_transformers import SentenceTransformer # lazy: heavy import + self.name = model_name + self._model = SentenceTransformer(model_name, device="cpu") + dim_fn = getattr(self._model, "get_embedding_dimension", None) or self._model.get_sentence_embedding_dimension # renamed in newer releases + self.dim = int(dim_fn()) + + def embed(self, texts: List[str]) -> List[List[float]]: + return self._model.encode(texts, normalize_embeddings=True, show_progress_bar=False).tolist() + + +class HashingEmbedder: + """Dependency-free fallback: hashed bag of words. Matches shared words, NOT meaning ("car" will not match "automobile"). + + Used when sentence-transformers is not installed, and in tests. Good enough to run the lab end to end, not to judge quality. + """ + + def __init__(self, dim: int = 384): + self.name = f"hashing-{dim}" + self.dim = dim + + def embed(self, texts: List[str]) -> List[List[float]]: + out = [] + for text in texts: + vec = [0.0] * self.dim + for w in _WORD.findall(text.lower()): + vec[int(hashlib.md5(w.encode()).hexdigest(), 16) % self.dim] += 1.0 + norm = math.sqrt(sum(v * v for v in vec)) or 1.0 + out.append([v / norm for v in vec]) + return out + + +def make_embedder(settings: RAGSettings, warnings: Optional[List[str]] = None) -> Embedder: + if settings.embedding_model_name == "hashing": + return HashingEmbedder() + try: + return SentenceTransformerEmbedder(settings.embedding_model_name) + except ImportError: + if warnings is not None: + warnings.append("sentence-transformers is not installed: using a keyword-only hashing embedder, so results will " + "not reflect meaning. Fix: pip install 'labagent[rag-local]'") + return HashingEmbedder() diff --git a/labagent/ai/rag/llm.py b/labagent/ai/rag/llm.py new file mode 100644 index 0000000..fb8c644 --- /dev/null +++ b/labagent/ai/rag/llm.py @@ -0,0 +1,107 @@ +"""LLM gateway: Groq, Ollama, OpenAI and Anthropic over plain HTTPS, with a graceful fallback to a mock answer. + +Order when LLM_PROVIDER is unset: Groq (if GROQ_API_KEY), OpenAI, Anthropic, local Ollama (if it answers), then mock. +A provider that fails at call time falls through to the next one and the reason is returned as a warning. +API keys are only ever sent to their own provider and are never logged. +""" +import json +import urllib.error +import urllib.request +from typing import Callable, List, Optional + +from pydantic import BaseModel + +from .config import ANTHROPIC_URL, GROQ_URL, OPENAI_URL, RAGSettings + +SYSTEM_PROMPT = ("Answer the question using only the context. Cite sources as [doc_id]. " + "If the context does not contain the answer, say you do not know.") + +Transport = Callable[[str, dict, dict, float], dict] + + +class LLMResult(BaseModel): + text: str + provider: str + model: str + prompt_tokens: int + completion_tokens: int + estimated: bool = False + warnings: List[str] = [] + + +def http_json(url: str, headers: dict, body: dict, timeout: float) -> dict: + # An explicit User-Agent is required: Groq's gateway answers Python's default one with 403 "error code: 1010". + req = urllib.request.Request(url, data=json.dumps(body).encode(), + headers={"Content-Type": "application/json", "User-Agent": "labagent-rag/0.2 (+https://github.com/becloudready/labagent)", **headers}) + with urllib.request.urlopen(req, timeout=timeout) as r: # noqa: S310 (https endpoints or the local Ollama URL) + return json.loads(r.read()) + + +def _est(text: str) -> int: + return max(1, len(text) // 4) + + +class LLMGateway: + def __init__(self, settings: RAGSettings, transport: Optional[Transport] = None): + self.s = settings + self._post = transport or http_json + + def _candidates(self) -> List[str]: + s = self.s + if s.llm_provider: + return [s.llm_provider.lower(), "mock"] + order = [name for name, key in (("groq", s.groq_api_key), ("openai", s.openai_api_key), ("anthropic", s.anthropic_api_key)) if key] + return order + ["ollama", "mock"] + + def complete(self, question: str, context: str) -> LLMResult: + warnings: List[str] = [] + user = f"Context:\n{context}\n\nQuestion: {question}" + for provider in self._candidates(): + try: + res = self._call(provider, user, question, context) + res.warnings = warnings + return res + except _Unavailable as e: + warnings.append(str(e)) + except urllib.error.HTTPError as e: + hint = {401: "check the API key", 403: "request blocked", 404: "model not available to this key: set GROQ_MODEL, list yours at /openai/v1/models", 429: "rate limit reached"}.get(e.code, "") + warnings.append(f"{provider} failed: HTTP {e.code}{' (' + hint + ')' if hint else ''}; trying the next provider") + except (urllib.error.URLError, OSError, KeyError, ValueError) as e: + warnings.append(f"{provider} failed ({type(e).__name__}); trying the next provider") + raise RuntimeError("unreachable: mock never fails") + + def _call(self, provider: str, user: str, question: str, context: str) -> LLMResult: + s = self.s + if provider == "mock": + first = context.split("\n\n")[0] if context else "no context retrieved" + text = ("[mock answer: no LLM provider is configured] Most relevant passage:\n" + first + + "\n\nFix: set GROQ_API_KEY (free tier at console.groq.com) or run `ollama serve`.") + return LLMResult(text=text, provider="mock", model="none", prompt_tokens=_est(user), completion_tokens=_est(text), estimated=True) + if provider == "anthropic": + if not s.anthropic_api_key: + raise _Unavailable("anthropic selected but ANTHROPIC_API_KEY is not set") + d = self._post(f"{ANTHROPIC_URL}/messages", {"x-api-key": s.anthropic_api_key, "anthropic-version": "2023-06-01"}, + {"model": s.anthropic_model, "max_tokens": 600, "system": SYSTEM_PROMPT, "messages": [{"role": "user", "content": user}]}, s.llm_timeout_s) + text = "".join(b.get("text", "") for b in d["content"]) + u = d.get("usage") or {} + return LLMResult(text=text, provider=provider, model=s.anthropic_model, prompt_tokens=u.get("input_tokens") or _est(user), + completion_tokens=u.get("output_tokens") or _est(text), estimated=not u) + url, key, model = { + "groq": (GROQ_URL, s.groq_api_key, s.groq_model), + "openai": (OPENAI_URL, s.openai_api_key, s.openai_model), + "ollama": (s.ollama_url.rstrip("/"), "ollama", s.ollama_model), + }.get(provider, (None, None, None)) + if url is None: + raise _Unavailable(f"unknown LLM_PROVIDER '{provider}': use groq, ollama, openai, anthropic or mock") + if not key: + raise _Unavailable(f"{provider} selected but its API key is not set") + d = self._post(f"{url}/chat/completions", {"Authorization": f"Bearer {key}"}, + {"model": model, "temperature": 0.1, "messages": [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user}]}, s.llm_timeout_s) + text = d["choices"][0]["message"]["content"] + u = d.get("usage") or {} + return LLMResult(text=text, provider=provider, model=model, prompt_tokens=u.get("prompt_tokens") or _est(user), + completion_tokens=u.get("completion_tokens") or _est(text), estimated=not u) + + +class _Unavailable(Exception): + pass diff --git a/labagent/ai/rag/models.py b/labagent/ai/rag/models.py new file mode 100644 index 0000000..e93fe71 --- /dev/null +++ b/labagent/ai/rag/models.py @@ -0,0 +1,73 @@ +"""Structured results returned by RAGLabClient.""" +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, Field + + +class RetrievedChunk(BaseModel): + doc_id: str + chunk_id: str + text: str + similarity: float = Field(description="Cosine similarity, higher is closer (1.0 identical).") + distance: float = Field(description="1 - similarity.") + keyword_score: float = Field(0.0, description="Share of query words found in the chunk (hybrid rerank signal).") + rank: int + title: Optional[str] = None + metadata: Dict[str, Any] = {} + + @property + def url(self) -> Optional[str]: + """Source link when the chunk carries one (SEC filings do).""" + return self.metadata.get("url") + + +class Metrics(BaseModel): + retrieval_ms: float = 0.0 + generation_ms: float = 0.0 + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + tokens_estimated: bool = Field(False, description="True when the provider returned no usage and counts are len/4.") + provider: str = "mock" + model: str = "" + + +class RAGResponse(BaseModel): + answer: str + chunks: List[RetrievedChunk] + metrics: Metrics + warnings: List[str] = [] + + @property + def retrieved_chunks(self) -> List[RetrievedChunk]: + """Alias of `chunks`, for readability in lab code: print(response.retrieved_chunks).""" + return self.chunks + + +class IngestionResult(BaseModel): + documents: int + chunks: int + skipped: Dict[str, str] = {} + elapsed_ms: float + embedding_model: str + vector_store: str + warnings: List[str] = [] + doc_ids: List[str] = [] + + +class VectorHit(BaseModel): + chunk_id: str + doc_id: str + text: str + similarity: float + embedding: Optional[List[float]] = None + title: Optional[str] = None + metadata: Dict[str, Any] = {} + + +class EvalResult(BaseModel): + cases: int + hits: int + hit_rate: float + misses: List[Dict[str, Any]] = [] + passed: Optional[bool] = None diff --git a/labagent/ai/rag/playground.py b/labagent/ai/rag/playground.py new file mode 100644 index 0000000..7fac3fe --- /dev/null +++ b/labagent/ai/rag/playground.py @@ -0,0 +1,65 @@ +"""Streamlit playground: retrieval, top-K chunks and the LLM answer side by side. + +Started by `labagent rag serve` (default port 8501). It uses the same RAGLabClient as the Python SDK and the notebook, +so what you see here is exactly what your code gets. +""" +import argparse +from pathlib import Path + +import streamlit as st + +from labagent.ai.rag import RAGLabClient +from labagent.ai.rag.config import load_dotenv + + +def _args() -> Path: + p = argparse.ArgumentParser() + p.add_argument("--dir", default=".") + return Path(p.parse_known_args()[0].dir) + + +project = _args() +load_dotenv(project / ".env") +st.set_page_config(page_title="labagent RAG playground", layout="wide") + + +@st.cache_resource(show_spinner="Loading the embedding model and indexing data/ ...") +def build(path: str): + client = RAGLabClient() + result = client.ingest_documents([Path(path) / "data"]) + return client, result + + +client, ingested = build(str(project)) +st.title("RAG playground") +st.caption(f"{ingested.documents} documents, {ingested.chunks} chunks | embeddings: {ingested.embedding_model} | store: {ingested.vector_store}") +for w in ingested.warnings: + st.warning(w) +if ingested.chunks == 0: + st.error(f"No documents found in {project / 'data'}. Run `labagent rag init` first or add .md/.txt files.") + st.stop() + +with st.sidebar: + top_k = st.slider("Top K chunks", 1, 10, 3) + doc = st.selectbox("Inspect a document's vectors", [""] + ingested.doc_ids) + +question = st.text_input("Ask a question about your documents", "What is hybrid search and when should I use it?") +if st.button("Run", type="primary") and question: + resp = client.query(question, top_k=top_k) + left, right = st.columns(2) + with left: + st.subheader("Retrieved chunks") + for c in resp.chunks: + with st.expander(f"#{c.rank} {c.doc_id} similarity {c.similarity:.3f} keywords {c.keyword_score:.2f}", expanded=c.rank == 1): + st.write(c.text) + with right: + st.subheader("Answer") + st.write(resp.answer) + m = resp.metrics + st.caption(f"{m.provider} {m.model} | retrieval {m.retrieval_ms} ms | generation {m.generation_ms} ms | " + f"tokens {m.total_tokens}{' (estimated)' if m.tokens_estimated else ''}") + for w in resp.warnings: + st.warning(w) +if doc: + st.subheader(f"Vectors for {doc}") + st.json(client.inspect_vectors(doc)) diff --git a/labagent/ai/rag/samples.py b/labagent/ai/rag/samples.py new file mode 100644 index 0000000..dbedad8 --- /dev/null +++ b/labagent/ai/rag/samples.py @@ -0,0 +1,280 @@ +"""Workspace content written by `labagent rag init` (and baked into the lab image). Kept as Python so nothing needs packaging. + +One source of truth: the Dockerfile and the devcontainer both call `labagent rag init`, so the files a learner sees in the +web IDE are exactly these. +""" +import json +from typing import Dict, List, Union + +ENV_TEMPLATE = """# labagent rag lab settings. Never commit real keys. +# Vector store: chroma (embedded, free) | qdrant | memory +VECTOR_DB_TYPE=chroma +EMBEDDING_MODEL_NAME=all-MiniLM-L6-v2 + +# Uncomment to keep the index between runs +# CHROMA_PATH=./.chroma + +# LLM: leave every key empty to get a clearly marked mock answer, or set one of these. +# Free tier key: console.groq.com +GROQ_API_KEY= +# Groq model (default openai/gpt-oss-20b). If you get HTTP 404, your account lacks that model: list yours with GET /openai/v1/models +# GROQ_MODEL= +# Force a provider: groq | ollama | openai | anthropic | mock +# LLM_PROVIDER=ollama +# OLLAMA_MODEL=qwen2.5-coder:0.5b +# OPENAI_API_KEY= +# ANTHROPIC_API_KEY= +# Remote Qdrant (Qdrant Cloud): both values come from its dashboard +# QDRANT_URL= +# QDRANT_API_KEY= +""" + +SECURITY_POLICY = """# Acme Corp Enterprise Security Policy + +## Access control + +Every employee gets a single identity managed by the company directory. Multi-factor authentication is required for all +systems. Shared accounts are not allowed. Access is granted by role and reviewed every quarter by the system owner. + +## Passwords and secrets + +Passwords must be at least 14 characters. Use the company password manager. Secrets such as API keys and database +credentials must never be stored in source code or chat tools. Store them in the secrets manager and rotate them every +90 days or immediately after a suspected leak. + +## Devices + +Company laptops use full disk encryption and install operating system updates within 14 days of release. Lost or stolen +devices must be reported to the security desk within one hour. +""" + +DATA_HANDLING = """# Acme Corp Data Handling Standard + +## Classification + +Data is Public, Internal, Confidential or Restricted. Customer records, payment data and employee records are Restricted. +When unsure, treat the data as Confidential until the data owner decides. + +## Storage and sharing + +Restricted data is encrypted at rest and in transit. It may only be stored in approved systems and never on personal +devices or personal cloud storage. Share Confidential data only with named people, and use expiring links. + +## Retention + +Customer records are kept for seven years after the account closes, then deleted. Logs that contain personal data are kept +for 90 days. The data owner approves any exception in writing. + +## Using AI tools + +Do not paste Restricted data into public AI assistants. Approved internal AI tools may process Confidential data. Every +AI-generated answer used in a customer-facing document must be reviewed by a person. +""" + +INCIDENT_LINES: List[str] = [ + "Acme Corp Incident Response Runbook", + "", + "1. Detect and report", + "Anyone who suspects a security incident reports it to the security desk at once.", + "The desk opens a ticket and assigns a severity from 1 (critical) to 4 (low).", + "", + "2. Contain", + "For severity 1 and 2 the on-call engineer isolates the affected systems within 30 minutes.", + "Do not delete evidence. Disable compromised accounts and rotate exposed secrets.", + "", + "3. Communicate", + "The incident lead updates leadership every hour during a severity 1 incident.", + "Legal decides whether customers or regulators must be notified. Notification of a", + "personal data breach is made within 72 hours of discovery.", + "", + "4. Recover and review", + "Restore service from clean backups. Hold a blameless review within five business days", + "and track every follow-up action to completion.", +] + + +def _esc(t: str) -> str: + return t.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)") + + +def make_pdf(lines: List[str]) -> bytes: + """A minimal one-page text PDF (Helvetica 11pt). Enough for the lab's PDF-ingestion exercise, no PDF library needed.""" + content = "BT /F1 11 Tf 14 TL 56 780 Td\n" + "\n".join(f"({_esc(l)}) Tj T*" for l in lines) + "\nET" + objs = ["<< /Type /Catalog /Pages 2 0 R >>", "<< /Type /Pages /Kids [3 0 R] /Count 1 >>", + "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 842] /Contents 4 0 R /Resources << /Font << /F1 5 0 R >> >> >>", + f"<< /Length {len(content.encode('latin-1'))} >>\nstream\n{content}\nendstream", "<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>"] + out = b"%PDF-1.4\n" + offsets = [] + for i, o in enumerate(objs, 1): + offsets.append(len(out)) + out += f"{i} 0 obj\n{o}\nendobj\n".encode("latin-1") + xref = len(out) + out += f"xref\n0 {len(objs) + 1}\n0000000000 65535 f \n".encode() + out += b"".join(f"{o:010d} 00000 n \n".encode() for o in offsets) + out += f"trailer\n<< /Size {len(objs) + 1} /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF\n".encode() + return out + + +MAIN_PY = '''"""Command-line RAG over data/sample_docs. Run from the lab folder: + + python src/main.py "What is our enterprise security policy?" + python src/main.py --top-k 5 "How long do we keep customer records?" + +Put your own .md, .txt or .pdf files in data/ and pass --data data to index them instead. +""" +import argparse +from pathlib import Path + +from labagent.ai.rag import RAGLabClient +from labagent.ai.rag.config import load_dotenv + +ROOT = Path(__file__).resolve().parent.parent + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("question") + p.add_argument("--data", default=str(ROOT / "data" / "sample_docs")) + p.add_argument("--top-k", type=int, default=3) + args = p.parse_args() + + load_dotenv(ROOT / ".env") # GROQ_API_KEY etc.; a missing key falls back to Ollama or a mock answer + rag = RAGLabClient() # vector_db="chroma", llm_provider picked from your keys + result = rag.ingest(args.data) + print(f"indexed {result.chunks} chunks from {result.documents} documents") + for w in result.warnings: + print("warning:", w) + + # TODO 1: print the answer. + # TODO 2: print each retrieved chunk with its doc_id and similarity. + # TODO 3: try --top-k 1 and --top-k 8. What changes in the answer? + response = rag.query(args.question, top_k=args.top_k) + print(response.answer) + + +if __name__ == "__main__": + main() +''' + + +def notebook() -> str: + def md(src): return {"cell_type": "markdown", "metadata": {}, "source": src.splitlines(True)} + def code(src): return {"cell_type": "code", "metadata": {}, "execution_count": None, "outputs": [], "source": src.splitlines(True)} + cells = [ + md("# Build your first RAG\n\nYou will chunk documents, turn the chunks into vectors, search them, and let a language model answer from what you found.\n\n" + "Time: about 30 minutes. Cost: zero. Embeddings run on this machine. Answers come from Groq (free key in `.env`), local Ollama, or a mock answer if neither is set.\n"), + md("## 1. Read the documents\n\nThe `data/sample_docs` folder holds three company documents: two Markdown files and one PDF.\n"), + code("from pathlib import Path\n\nfor f in sorted(Path('../data/sample_docs').iterdir()):\n print(f.name, f.stat().st_size, 'bytes')\n"), + md("## 2. Chunk\n\nModels read short passages, so long documents are split. Try different sizes and watch the chunk count.\n"), + code("from labagent.ai.rag import chunk_text\n\ntext = Path('../data/sample_docs/security_policy.md').read_text()\nfor size in (200, 400, 800):\n chunks = chunk_text(text, size=size, overlap=50)\n print(size, '->', len(chunks), 'chunks')\nprint('\\n--- first chunk at 400 ---\\n', chunk_text(text, size=400, overlap=50)[0])\n"), + md("**Your turn:** which size would you choose for policy documents, and why? Write one sentence here before moving on.\n"), + md("## 3. Embed\n\nAn embedding is a list of numbers that captures meaning. Similar texts get similar vectors.\n"), + code("from labagent.ai.rag import RAGLabClient\nfrom labagent.ai.rag.config import load_dotenv\n\nload_dotenv(Path('../.env'))\nrag = RAGLabClient() # Chroma + MiniLM on CPU, no cost\nvecs = rag.embedder.embed(['Passwords need 14 characters', 'Use a long password', 'The cafeteria opens at noon'])\nprint(len(vecs[0]), 'numbers per vector')\n\ndef cos(a, b): return sum(x * y for x, y in zip(a, b))\nprint('password vs long password :', round(cos(vecs[0], vecs[1]), 3))\nprint('password vs cafeteria :', round(cos(vecs[0], vecs[2]), 3))\n"), + md("The first pair should score much higher. If both are near equal, check the warnings above: without `sentence-transformers` the client falls back to a word-matching embedder.\n"), + md("## 4. Index\n"), + code("result = rag.ingest('../data/sample_docs')\nprint(result.documents, 'documents,', result.chunks, 'chunks,', result.embedding_model, 'in', result.vector_store)\nprint('skipped:', result.skipped)\nrag.inspect_vectors(result.doc_ids[0])['items'][0]\n"), + md("## 5. Retrieve\n\nEach result has a similarity score and a keyword score. Retrieval reranks by `0.7 * similarity + 0.3 * keyword overlap`.\n"), + code("for c in rag.retrieve('How often must secrets be rotated?', top_k=3):\n print(c.rank, c.doc_id, 'similarity', c.similarity, 'keywords', c.keyword_score)\n print(' ', c.text[:120].replace('\\n', ' '))\n"), + md("**Your turn:** ask a question that only the PDF can answer (look at its text first). Does the PDF show up at rank 1?\n"), + code("# YOUR QUESTION HERE\nquestion = 'What is our enterprise security policy?'\nfor c in rag.retrieve(question, top_k=3):\n print(c.rank, c.doc_id, c.similarity)\n"), + md("## 6. Generate\n\nThe retrieved chunks go into the prompt. The model is told to answer only from them and to say when it does not know.\n"), + code("resp = rag.query(question, top_k=3)\nprint(resp.answer)\nprint()\nprint(resp.metrics)\nfor w in resp.warnings:\n print('warning:', w)\n"), + md("## 7. Break it on purpose\n\n1. Ask something the documents do not cover, such as `What is the lunch menu?`. A good system says it does not know.\n" + "2. Set `top_k=1` and ask a question that needs two documents.\n3. Call `rag.reset()`, re-ingest with `chunk_chars=100` in `RAGSettings`, and compare answers.\n"), + code("print(rag.query('What is the lunch menu?').answer)\n"), + md("## Next\n\nOpen the visual playground in a terminal with `labagent rag playground` and compare scores side by side. " + "For a managed version on Databricks Lakebase, see `labagent launch rag-sec-filings`.\n"), + ] + for i, c in enumerate(cells): + c["id"] = f"cell{i:02d}" # nbformat 4.5 requires cell ids + nb = {"cells": cells, "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python"}}, + "nbformat": 4, "nbformat_minor": 5} + return json.dumps(nb, indent=1) + "\n" + + +def sec_notebook() -> str: + """Tutorial notebook: RAG over SEC 10-K filings (segments match the video outline).""" + def md(src): return {"cell_type": "markdown", "metadata": {}, "source": src.splitlines(True)} + def code(src): return {"cell_type": "code", "metadata": {}, "execution_count": None, "outputs": [], "source": src.splitlines(True)} + cells = [ + md("# RAG over SEC 10-K filings\n\nYou will build a research assistant that answers questions from real annual reports (Form 10-K) of Apple, Pfizer, " + "Walmart, JPMorgan and Chevron, with citations, filters and a quality check.\n\n" + "**You need:** this web IDE and a free Groq key in `.env` (`GROQ_API_KEY=`). Without a key everything works except the written answers.\n\n" + "**Time:** about 25 minutes.\n"), + code("import os\nfrom pathlib import Path\nfrom labagent.ai.rag import RAGLabClient\nfrom labagent.ai.rag.config import load_dotenv\n\n" + "load_dotenv(Path('../.env'))\nrag = RAGLabClient(vector_db='chroma', llm_provider='groq') # Chroma in memory, MiniLM embeddings on CPU\n\n" + "def show(chunks):\n for c in chunks:\n m = c.metadata\n print(f\"#{c.rank} {m.get('ticker')} {m.get('fiscal_year')} Item {m.get('item')} ({m.get('section')}) \"\n f\"similarity {c.similarity:.3f} keywords {c.keyword_score:.2f}\")\n"), + md("## 1. Why RAG for filings\n\nA 10-K runs to hundreds of pages. Pasting it into a prompt is slow, costly and often too long for the model. " + "RAG keeps the documents outside the model, finds the few passages that matter for each question, and gives only those to the model.\n\n" + "The loop: **load, chunk, embed, store, retrieve, generate, evaluate.** We do each step below.\n"), + md("## 2. Load the data\n\nThe filings were already split into section-aware chunks and published as a versioned dataset. " + "`fetch_snapshot` downloads it once and checks its SHA-256 every time, so everyone gets identical data.\n"), + code("from collections import Counter\nfrom labagent.datasets import fetch_snapshot\n\nrows = fetch_snapshot('sec-10k-2026.10')\nprint(len(rows), 'chunks')\n" + "print(Counter(r['meta']['ticker'] for r in rows))\nprint(Counter(r['meta']['section'] for r in rows))\nprint(sorted(Counter((r['meta']['ticker'], r['meta']['fiscal_year']) for r in rows)))\nprint()\nprint(rows[0]['title'])\nprint(rows[0]['meta'])\nprint(rows[0]['content'][:300], '...')\n"), + md("Every chunk carries metadata: company, fiscal year, section (Item 1 Business, Item 1A Risk Factors, Item 7 MD&A) and the URL of the filing. " + "Each company has **two fiscal years**, so the same question can match two near-identical passages. You will use the metadata for filters and citations.\n\n**Known gap:** the JPMorgan and Chevron MD&A sections are not in this dataset. We will see that show up in step 8.\n"), + md("## 3. Embed and index\n\nAn embedding turns text into a list of numbers so that similar meanings are close together. " + "`all-MiniLM-L6-v2` runs on the CPU here, so it costs nothing per call.\n"), + code("import time\nt = time.time()\nresult = rag.ingest_chunks(rows, progress=lambda d, n: print(f'\\r{d}/{n}', end=''))\n" + "print(f\"\\nindexed {result.chunks} chunks in {time.time() - t:.0f}s with {result.embedding_model} into {result.vector_store}\")\n"), + md("## 4. Retrieve with vectors only\n\nThe question becomes a vector, and the store returns the closest chunks. `keyword_weight=0` means pure vector search.\n"), + code("q = 'What patent expirations and generic competition risks does Pfizer disclose?'\nshow(rag.retrieve(q, top_k=5, keyword_weight=0))\n"), + md("## 5. Break it: when vectors alone are not enough\n\nVectors capture meaning, not exact wording. Ask for a precise phrase and watch the ranking.\n"), + code("q = '52- or 53-week fiscal year'\nprint('vectors only:')\nshow(rag.retrieve(q, top_k=3, keyword_weight=0))\n" + "print('\\nvectors + keyword rerank:')\nshow(rag.retrieve(q, top_k=3, keyword_weight=0.3))\n"), + md("The first run ranks a Walmart chunk first. The rerank adds a keyword-overlap score and puts the Apple chunks, which contain that exact phrase, on top. " + "This is hybrid retrieval.\n\n**Your turn:** try `'iPhone 17 Pro'` and `'Permian Basin'`. Does the keyword score change the order? " + "Sometimes it does not, and that is worth knowing.\n"), + md("## 6. Filter with metadata\n\nSometimes you already know the company. A filter is exact, so no other company can appear.\n"), + code("q = 'What are the main supply chain risks?'\nprint('no filter:')\nshow(rag.retrieve(q, top_k=5))\n" + "print('\\nonly Walmart:')\nshow(rag.retrieve(q, top_k=5, where={'ticker': 'WMT'}))\n"), + md("Each company has two years of filings. To compare only the latest year, filter on `fiscal_year` too.\n"), + code("q = 'What patent expirations and generic competition risks does Pfizer disclose?'\nprint('both years:')\nshow(rag.retrieve(q, top_k=4, where={'ticker': 'PFE'}))\n" + "print('\\n2025 only:')\nshow(rag.retrieve(q, top_k=4, where={'ticker': 'PFE', 'fiscal_year': '2025'}))\n"), + md("**Your turn:** combine filters with a list, for example `where={'ticker': ['AAPL', 'WMT'], 'item': '1A'}`.\n"), + md("## 7. Generate an answer with citations\n\nThe retrieved chunks go into the prompt. The model must answer only from them and say so when it cannot.\n"), + code("resp = rag.query(\"What are Apple's main supply chain and manufacturing risks?\", top_k=5, where={'ticker': 'AAPL', 'fiscal_year': '2025'})\nprint(resp.answer)\n" + "print('\\nprovider:', resp.metrics.provider, resp.metrics.model, '|', resp.metrics.generation_ms, 'ms |', resp.metrics.total_tokens, 'tokens')\n" + "for c in resp.retrieved_chunks:\n print('-', c.title, '->', c.url)\nfor w in resp.warnings:\n print('warning:', w)\n"), + code("# A question the filings cannot answer. A good system admits it.\nprint(rag.query('What is the recipe for Walmart\\'s store-brand cookies?', top_k=3).answer)\n"), + md("## 8. Evaluate before you trust it\n\nWrite questions with the answer you expect, then measure how often retrieval finds a chunk from the right company and section. " + "A score with a pass mark is how you know a change helped or hurt.\n"), + code("cases = [\n" + " {'question': \"What are Apple's main supply chain and manufacturing risks?\", 'expect': {'ticker': 'AAPL', 'item': '1A'}},\n" + " {'question': 'How does Apple describe its products and services?', 'expect': {'ticker': 'AAPL', 'item': '1'}},\n" + " {'question': 'What credit and interest rate risks does JPMorgan Chase face?', 'expect': {'ticker': 'JPM', 'item': '1A'}},\n" + " {'question': \"How did JPMorgan Chase's net income and revenue change this year?\", 'expect': {'ticker': 'JPM', 'item': '7'}},\n" + " {'question': 'What patent expirations and generic competition risks does Pfizer disclose?', 'expect': {'ticker': 'PFE', 'item': '1A'}},\n" + " {'question': \"What are Pfizer's main product lines and business segments?\", 'expect': {'ticker': 'PFE', 'item': '1'}},\n" + " {'question': \"How do oil and gas price changes affect Chevron's earnings?\", 'expect': {'ticker': 'CVX', 'item': '7'}},\n" + " {'question': 'What climate and regulatory risks does Chevron identify?', 'expect': {'ticker': 'CVX', 'item': '1A'}},\n" + " {'question': \"How did Walmart's comparable sales and operating income perform?\", 'expect': {'ticker': 'WMT', 'item': '7'}},\n" + " {'question': 'What competition and supply risks does Walmart face?', 'expect': {'ticker': 'WMT', 'item': '1A'}},\n]\n\n" + "ev = rag.evaluate(cases, top_k=5, min_hit_rate=0.6)\nprint(f'hit rate {ev.hit_rate} ({ev.hits}/{ev.cases}) pass mark 0.6 -> {\"PASS\" if ev.passed else \"FAIL\"}')\n" + "for m in ev.misses:\n print('MISS:', m['question'])\n print(' expected', m['expected'], '| got', m['got'])\n"), + md("Two misses, and both are the data gap from step 2: there is no JPMorgan or Chevron MD&A in the dataset, so no retrieval setting can find it. " + "Evaluation tells you whether to fix the **retrieval** or the **data**.\n\n**Your turn:** add three of your own questions to `cases`, then run `rag.evaluate(cases, keyword_weight=0)` " + "and `rag.evaluate(cases, keyword_weight=0.3)` and compare.\n"), + md("## 9. What it cost, and where next\n\nEmbeddings ran on this machine, so they were free. Answers cost fractions of a cent on Groq's free tier. " + "The compute you are using is billed per second only while this IDE runs and stops when it goes idle.\n\n" + "**Next:** the same pipeline runs managed on Databricks Lakebase (Postgres with pgvector) and Databricks Apps:\n\n" + "```\npip install 'labagent[rag]'\nlabagent launch rag-sec-filings --name demo\n```\n"), + ] + for i, c in enumerate(cells): + c["id"] = f"cell{i:02d}" # nbformat 4.5 requires cell ids + nb = {"cells": cells, "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python"}}, + "nbformat": 4, "nbformat_minor": 5} + return json.dumps(nb, indent=1) + "\n" + + +def workspace() -> Dict[str, Union[str, bytes]]: + """Relative path -> content for a fresh lab workspace.""" + return { + ".env": ENV_TEMPLATE, + "data/sample_docs/security_policy.md": SECURITY_POLICY, + "data/sample_docs/data_handling.md": DATA_HANDLING, + "data/sample_docs/incident_response.pdf": make_pdf(INCIDENT_LINES), + "notebooks/01_build_your_first_rag.ipynb": notebook(), + "notebooks/02_sec_filings_rag.ipynb": sec_notebook(), + "src/main.py": MAIN_PY, + } diff --git a/labagent/cli.py b/labagent/cli.py index e9fe7f7..66587dc 100644 --- a/labagent/cli.py +++ b/labagent/cli.py @@ -46,7 +46,7 @@ def _packs_main(args) -> int: return env.exit_code -KNOWN = {"run", "packs", "dataset", "launch", "destroy", "bootstrap-path", "commands"} +KNOWN = {"run", "packs", "rag", "dataset", "launch", "destroy", "bootstrap-path", "commands"} def _nl_main(argv) -> int: @@ -102,6 +102,8 @@ def main(argv=None) -> int: r.add_argument("--output", default="result.json", help="Result envelope path, or - for stdout") r.add_argument("--ruleset", help="Path to a ruleset JSON for audit/cleanup (merged into the job)") sub.add_parser("packs", help="List available lab packs") + from .ai.rag import cli as rag_cli + rag_cli.add_parser(sub) dp = sub.add_parser("dataset", help="Package a chunk file as a versioned, checksummed snapshot (maintainers)") dp.add_argument("--chunks", required=True, help="JSONL of chunks, e.g. ~/.labagent/cache/rag-sec-filings/.jsonl") dp.add_argument("--name", required=True) @@ -130,6 +132,8 @@ def main(argv=None) -> int: return 2 print(json.dumps({k: m[k] for k in ("file", "chunks", "bytes", "sha256")}, indent=2)) return 0 + if args.cmd == "rag": + return rag_cli.run(args) if args.cmd in ("packs", "launch", "destroy"): return _packs_main(args) if args.cmd == "commands": diff --git a/labagent/packs/rag_ide/__init__.py b/labagent/packs/rag_ide/__init__.py new file mode 100644 index 0000000..64bacb5 --- /dev/null +++ b/labagent/packs/rag_ide/__init__.py @@ -0,0 +1 @@ +"""rag-ide: a browser VS Code per learner (code-server) with the RAG lab preinstalled, on Modal.""" diff --git a/labagent/packs/rag_ide/modal_gateway.py b/labagent/packs/rag_ide/modal_gateway.py new file mode 100644 index 0000000..58a5294 --- /dev/null +++ b/labagent/packs/rag_ide/modal_gateway.py @@ -0,0 +1,132 @@ +"""Everything that touches Modal lives here, behind a small interface, so the pack logic is testable with a fake +and a second platform can be added later by writing another gateway with the same methods. + +Why Sandboxes (not Functions, as in the first spike): each learner gets their own volume, their own tag, and can be stopped on +their own. The code-server process runs as an unprivileged user, never root. +""" +from __future__ import annotations + +import os +import re +from pathlib import Path +from typing import Any, Dict, List, Optional + +from ..base import PackError + +CODE_SERVER_VERSION = "4.141.0" +CODE_SERVER_SHA256_AMD64 = "f384beb4f3ba207e0bca27d0dbd36dd825d5aa909d78c6d4baa8703eaf6b77ba" # Modal workers are x86_64 +MODEL = "all-MiniLM-L6-v2" +LAB_DIR = "/home/coder/lab" +WORK_DIR = LAB_DIR + "/work" # the learner's persistent volume +PORT = 8080 + +# Pip spec for the labagent that goes into the image: a checkout path ("local") or the PyPI release that contains labagent.ai.rag. +PYPI_SPEC = "labagent[rag-lab]>=0.3.0" +IGNORE = [".git", ".venv", "venv", "dist", "*.egg-info", "**/__pycache__", "**/.terraform", "tests", "docs", "spikes", ".pytest_cache"] + + +def slug(text: str, limit: int = 40) -> str: + return re.sub(r"[^a-z0-9-]+", "-", text.lower()).strip("-")[:limit] or "x" + + +def volume_name(cohort: str, learner: str) -> str: + return f"labagent-rag-ide-{slug(cohort, 20)}-{slug(learner, 20)}" + + +def app_name(cohort: str) -> str: + return f"labagent-rag-ide-{slug(cohort, 30)}" + + +def start_script(idle_seconds: int) -> str: + """Runs as root only long enough to hand the volume to the learner's user, then replaces itself with code-server as that user.""" + # WORK_DIR is a symlink into the volume: chown the target, not the link. + return (f"chown -R coder:coder \"$(readlink -f {WORK_DIR})\" 2>/dev/null; " + f"exec runuser -u coder -- env HOME=/home/coder code-server --bind-addr 0.0.0.0:{PORT} --auth password " + f"--idle-timeout-seconds {idle_seconds} {LAB_DIR}") + + +class ModalGateway: + """Real Modal calls. Imported lazily so labagent works without the `modal` extra.""" + + def __init__(self) -> None: + try: + import modal + except ImportError as e: + raise PackError("the 'modal' package is not installed. Fix: pip install 'labagent[modal]' then run: modal token new") from e + self.modal = modal + + # -- auth + def check_auth(self) -> str: + """Returns a label for the active login, or raises with the one command that fixes it.""" + try: + from modal.config import config + if not (config.get("token_id") and config.get("token_secret")): + raise PackError("no Modal login found. Run: modal token new") + return os.environ.get("MODAL_PROFILE") or "your Modal account" + except PackError: + raise + except Exception as e: + raise PackError(f"could not read the Modal login ({type(e).__name__}). Run: modal token new") from e + + # -- image + def image(self, source: str) -> Any: + """source is 'pypi' or a path to a labagent checkout. Built once per change and cached by Modal.""" + m = self.modal + img = (m.Image.debian_slim(python_version="3.11") + .apt_install("curl", "git", "ca-certificates") + .run_commands( + "set -eux; curl -fsSL -o /tmp/cs.deb https://github.com/coder/code-server/releases/download/" + f"v{CODE_SERVER_VERSION}/code-server_{CODE_SERVER_VERSION}_amd64.deb; " + f"echo '{CODE_SERVER_SHA256_AMD64} /tmp/cs.deb' | sha256sum -c -; dpkg -i /tmp/cs.deb; rm /tmp/cs.deb", + "useradd -m -s /bin/bash -u 1000 coder") + .pip_install("torch", index_url="https://download.pytorch.org/whl/cpu")) + if source == "pypi": + img = img.pip_install(PYPI_SPEC) + else: + img = img.add_local_dir(source, "/src", copy=True, ignore=IGNORE).run_commands('pip install "/src[rag-lab]"') + return (img + .env({"HF_HOME": "/opt/hf-cache", "SENTENCE_TRANSFORMERS_HOME": "/opt/hf-cache", "HF_HUB_OFFLINE": "0"}) + .run_commands(f"python -c \"from sentence_transformers import SentenceTransformer; SentenceTransformer('{MODEL}')\"", "chmod -R a+rX /opt/hf-cache") + .env({"HF_HUB_OFFLINE": "1"}) + .run_commands(f"labagent rag init --dir {LAB_DIR}", f"mkdir -p {WORK_DIR}", "chown -R coder:coder /home/coder", + "python -m ipykernel install --sys-prefix --name python3", + "runuser -u coder -- env HOME=/home/coder code-server --install-extension ms-python.python " + "--install-extension ms-toolsai.jupyter || echo 'extensions not installed'")) + + def app(self, name: str) -> Any: + return self.modal.App.lookup(name, create_if_missing=True) + + def build(self, app: Any, image: Any) -> None: + image.build(app) + + # -- learners + def launch(self, app: Any, cohort: str, learner: str, password: str, image: Any, cpu: float, memory_mib: int, + timeout_s: int, idle_s: int, env: Dict[str, str]) -> Dict[str, str]: + m = self.modal + vol = m.Volume.from_name(volume_name(cohort, learner), create_if_missing=True) + sb = m.Sandbox.create("sh", "-c", start_script(idle_s), app=app, image=image, cpu=cpu, memory=memory_mib, timeout=timeout_s, + volumes={WORK_DIR: vol}, encrypted_ports=[PORT], env={"PASSWORD": password, **env}) + sb.set_tags({"labagent": "rag-ide", "cohort": slug(cohort), "learner": slug(learner)}) + tunnel = sb.tunnels(timeout=120)[PORT] + return {"sandbox_id": sb.object_id, "url": tunnel.url, "volume": volume_name(cohort, learner)} + + def is_alive(self, sandbox_id: str) -> bool: + try: + return self.modal.Sandbox.from_id(sandbox_id).poll() is None + except Exception: + return False + + def stop(self, sandbox_id: str) -> None: + try: + self.modal.Sandbox.from_id(sandbox_id).terminate() + except Exception: + pass # already gone + + def find(self, cohort: str) -> List[str]: + return [sb.object_id for sb in self.modal.Sandbox.list(tags={"labagent": "rag-ide", "cohort": slug(cohort)})] + + def delete_volume(self, name: str) -> None: + try: + self.modal.Volume.delete(name) + except Exception: + pass diff --git a/labagent/packs/rag_ide/pack.py b/labagent/packs/rag_ide/pack.py new file mode 100644 index 0000000..420ddad --- /dev/null +++ b/labagent/packs/rag_ide/pack.py @@ -0,0 +1,191 @@ +"""rag-ide: give each learner a private browser VS Code with the RAG lab, on serverless Modal. No learner sign-up, nothing on their laptop. + +The instructor's Modal login stays on the instructor's machine. Each learner gets a URL and a random password (written to a +0600 credentials file, never to state or logs), a persistent volume at work/, and an idle timeout. See docs for the cost model. +""" +from __future__ import annotations + +import csv +import os +import re +import secrets +from pathlib import Path +from typing import Any, Dict, List + +from ..base import Pack, PackContext, PackError, Param, Step, state_dir +from .modal_gateway import ModalGateway, app_name, slug + +MAX_LEARNERS = 60 + + +def learner_names(spec: Any) -> List[str]: + """'3' -> learner-01..03; 'ana,ben' -> ana, ben. Names become resource names, so they are slugged and must be unique.""" + text = str(spec).strip() + if re.fullmatch(r"\d+", text): + n = int(text) + names = [f"learner-{i:02d}" for i in range(1, n + 1)] + else: + names = [slug(p) for p in text.split(",") if p.strip()] + if not names: + raise PackError("learners must be a count (e.g. 5) or a comma list of names (e.g. ana,ben)") + if len(set(names)) != len(names): + raise PackError("learner names must be unique after normalising") + if len(names) > MAX_LEARNERS: + raise PackError(f"at most {MAX_LEARNERS} learners per cohort in one launch") + return names + + +def checkout_root() -> str: + root = Path(__file__).resolve().parents[3] + return str(root) if (root / "pyproject.toml").is_file() and (root / "labagent" / "ai" / "rag").is_dir() else "" + + +def creds_path(name: str) -> Path: + return state_dir() / "rag-ide" / f"{name}-credentials.csv" + + +class _LazyModal: + """Creates the real gateway on first use, inside a step, so a missing `modal` package is a clean step failure + (and `--dry-run` works without it) instead of a crash while the pack is being set up.""" + + def __init__(self) -> None: + self._gw = None + + def __getattr__(self, name: str): + if self._gw is None: + self._gw = ModalGateway() + return getattr(self._gw, name) + + +class RagIde(Pack): + id = "rag-ide" + title = "Browser VS Code lab per learner (Modal)" + summary = "Serverless web IDE with the RAG lab for each learner: one Modal login, nothing for learners to install" + aliases = ("ide", "web-ide", "rag-web-ide") + default_name = "cohort" + params = [ + Param("learners", "1", "Count (5) or comma list of names (ana,ben)"), + Param("hours", 4, "Hard cap per session, 1 to 24"), + Param("idle_minutes", 15, "Shut a session down after this long with no browser connected"), + Param("cpu", 2.0, "CPU cores per learner"), + Param("memory_mib", 4096, "Memory per learner"), + Param("image", "auto", "auto | pypi | path to a labagent checkout"), + Param("llm", "none", "none (mock answers) | groq-shared (copies GROQ_API_KEY from your shell into each learner's IDE)"), + Param("groq_model", "openai/gpt-oss-20b", "Used when llm=groq-shared"), + Param("show_passwords", "auto", "auto (print them for 3 learners or fewer) | yes | no. They are always in the credentials file"), + Param("keep_work", False, "destroy: keep the learners' volumes"), + ] + + def services(self, profile): + return {"modal": _LazyModal()} + + def _validate(self, ctx: PackContext) -> None: + p = ctx.params + names = learner_names(p["learners"]) + if not 1 <= int(p["hours"]) <= 24: + raise PackError("hours must be between 1 and 24") + if p["show_passwords"] not in ("auto", "yes", "no"): + raise PackError("show_passwords must be auto, yes or no") + if p["llm"] not in ("none", "groq-shared"): + raise PackError("llm must be 'none' or 'groq-shared'") + if p["llm"] == "groq-shared" and not os.environ.get("GROQ_API_KEY"): + raise PackError("llm=groq-shared needs GROQ_API_KEY in your shell. The key is visible to every learner: use a throwaway key with a low limit.") + ctx.mem["names"] = names + + def _preflight(self, ctx: PackContext): + self._validate(ctx) + workspace = ctx.services["modal"].check_auth() + src = ctx.params["image"] + if src == "auto": + src = checkout_root() or "pypi" + elif src != "pypi" and not (Path(src) / "pyproject.toml").is_file(): + raise PackError(f"image path '{src}' is not a labagent checkout") + ctx.log(f" Modal workspace: {workspace} | image source: {src} | learners: {len(ctx.mem['names'])}") + return {"workspace": workspace, "image_source": src, "app": app_name(ctx.name)} + + def _build(self, ctx: PackContext): + gw = ctx.services["modal"] + app = gw.app(ctx.state["app"]) + ctx.mem["app"] = app + ctx.mem["image"] = gw.image(ctx.state["image_source"]) + gw.build(app, ctx.mem["image"]) # first build takes several minutes; later runs hit Modal's cache + return {"image_built": True} + + def _learners(self, ctx: PackContext): + gw = ctx.services["modal"] + p = ctx.params + names = ctx.mem.get("names") or learner_names(p["learners"]) + learners: Dict[str, Dict[str, str]] = dict(ctx.state.get("learners", {})) + passwords: Dict[str, str] = {} + env = {"GROQ_API_KEY": os.environ["GROQ_API_KEY"], "GROQ_MODEL": p["groq_model"], "LLM_PROVIDER": "groq"} if p["llm"] == "groq-shared" else {} + app = ctx.mem.get("app") or gw.app(ctx.state["app"]) + image = ctx.mem.get("image") or gw.image(ctx.state["image_source"]) + for n in names: + if n in learners and gw.is_alive(learners[n]["sandbox_id"]): + ctx.log(f" {n}: already running") + continue + pw = secrets.token_urlsafe(12) + ctx.log(f" starting {n} ...") + learners[n] = gw.launch(app, ctx.name, n, pw, image, float(p["cpu"]), int(p["memory_mib"]), int(p["hours"]) * 3600, int(p["idle_minutes"]) * 60, env) + passwords[n] = pw + ctx.state["learners"] = learners # persist progress so a failure part-way resumes + ctx.mem["passwords"] = passwords + return {"learners": learners} + + def _credentials(self, ctx: PackContext): + path = creds_path(ctx.name) + old: Dict[str, str] = {} + if path.exists(): + old = {r["learner"]: r["password"] for r in csv.DictReader(path.open())} + pw = {**old, **ctx.mem.get("passwords", {})} + path.parent.mkdir(parents=True, exist_ok=True) + with open(os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600), "w", newline="") as f: + w = csv.writer(f) + w.writerow(["learner", "url", "password"]) + for n, info in sorted(ctx.state["learners"].items()): + w.writerow([n, info["url"], pw.get(n, "(lost: destroy and relaunch this learner)")]) + learners = sorted(ctx.state["learners"].items()) + sp = ctx.params["show_passwords"] + show = sp == "yes" or (sp == "auto" and len(learners) <= 3) + ctx.log("\n Open each URL in a browser and sign in with the password:") + for n, info in learners: + ctx.log(f" {n}: {info['url']}" + (f" password: {pw.get(n, '(lost: destroy and relaunch)')}" if show else "")) + ctx.log(f"\n Passwords are in {path} (0600, one row per learner)" + ("" if show else f"\n View them: cat {path}")) + return {"credentials_file": str(path)} + + def deploy_steps(self, ctx: PackContext) -> List[Step]: + return [ + Step("preflight", "Check the Modal login and settings", self._preflight), + Step("image", "Build the lab image (first time: several minutes)", self._build), + Step("learners", "Start one private IDE per learner", self._learners), + Step("credentials", "Write the credentials file", self._credentials), + ] + + def _stop_all(self, ctx: PackContext): + gw = ctx.services["modal"] + ids = {i["sandbox_id"] for i in ctx.state.get("learners", {}).values()} | set(gw.find(ctx.name)) + for sid in ids: + gw.stop(sid) + return {"learners": {}, "stopped": len(ids)} + + def _volumes(self, ctx: PackContext): + if ctx.params["keep_work"]: + return {} + gw = ctx.services["modal"] + for vol in ctx.state.get("volumes", []): + gw.delete_volume(vol) + return {"volumes": []} + + def _forget(self, ctx: PackContext): + creds_path(ctx.name).unlink(missing_ok=True) + return {"credentials_file": None} + + def destroy_steps(self, ctx: PackContext) -> List[Step]: + # remember volume names before the learner records are cleared + ctx.state["volumes"] = sorted({i["volume"] for i in ctx.state.get("learners", {}).values() if i.get("volume")} | set(ctx.state.get("volumes", []))) + return [ + Step("preflight", "Check the Modal login", lambda c: {"workspace": c.services["modal"].check_auth()}), + Step("stop", "Stop every learner's IDE", self._stop_all), + Step("volumes", "Delete learner volumes (unless keep_work)", self._volumes), + Step("credentials", "Remove the credentials file", self._forget), + ] diff --git a/labagent/packs/registry.py b/labagent/packs/registry.py index 9e697fe..79676fc 100644 --- a/labagent/packs/registry.py +++ b/labagent/packs/registry.py @@ -9,7 +9,8 @@ def builtin() -> Dict[str, Pack]: from .rag.pack import RagLakebase, RagSecFilings, RagSecFilingsLive from .db_agent.pack import DbAgent - return {p.id: p for p in (RagLakebase(), RagSecFilings(), RagSecFilingsLive(), DbAgent())} + from .rag_ide.pack import RagIde + return {p.id: p for p in (RagLakebase(), RagSecFilings(), RagSecFilingsLive(), DbAgent(), RagIde())} def all_packs() -> Dict[str, Pack]: diff --git a/pyproject.toml b/pyproject.toml index 4597e6a..f642d36 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,10 @@ dependencies = ["boto3>=1.34", "cryptography>=42"] [project.optional-dependencies] rag = ["psycopg[binary]>=3.1"] -dev = ["pytest>=8", "moto[all]>=5", "build>=1.2", "twine>=5"] +rag-local = ["pydantic>=2", "chromadb>=0.5", "sentence-transformers>=3", "streamlit>=1.35", "pypdf>=4"] +modal = ["modal>=1.0"] +rag-lab = ["pydantic>=2", "chromadb>=0.5", "sentence-transformers>=3", "streamlit>=1.35", "pypdf>=4", "groq>=0.9", "langchain>=0.3", "qdrant-client>=1.9", "jupyterlab>=4", "ipykernel>=6"] +dev = ["pytest>=8", "pydantic>=2", "moto[all]>=5", "build>=1.2", "twine>=5"] [project.urls] Homepage = "https://github.com/becloudready/labagent" diff --git a/tests/test_rag_ide.py b/tests/test_rag_ide.py new file mode 100644 index 0000000..361fec0 --- /dev/null +++ b/tests/test_rag_ide.py @@ -0,0 +1,205 @@ +import csv +import os + +import pytest + +from labagent.packs import registry +from labagent.packs.base import load_state, run_pack +from labagent.packs.rag_ide import modal_gateway as mg +from labagent.packs.rag_ide.pack import RagIde, creds_path, learner_names + + +class FakeModal: + def __init__(self, logged_in=True): + self.logged_in = logged_in + self.sandboxes = {} + self.volumes = set() + self.launches = [] + self.builds = 0 + self.deleted = [] + + def check_auth(self): + if not self.logged_in: + from labagent.packs.base import PackError + raise PackError("no Modal login found. Run: modal token new") + return "acme-workspace" + + def app(self, name): return {"app": name} + def image(self, source): return {"image": source} + def build(self, app, image): self.builds += 1 + + def launch(self, app, cohort, learner, password, image, cpu, memory_mib, timeout_s, idle_s, env): + sid = f"sb-{len(self.launches) + 1}" + self.launches.append({"learner": learner, "password": password, "timeout_s": timeout_s, "idle_s": idle_s, "env": env, "cpu": cpu}) + self.sandboxes[sid] = True + vol = mg.volume_name(cohort, learner) + self.volumes.add(vol) + return {"sandbox_id": sid, "url": f"https://{learner}.example.modal.host", "volume": vol} + + def is_alive(self, sid): return self.sandboxes.get(sid, False) + def stop(self, sid): self.sandboxes[sid] = False + def find(self, cohort): return [s for s, up in self.sandboxes.items() if up] + def delete_volume(self, name): self.deleted.append(name); self.volumes.discard(name) + + +@pytest.fixture(autouse=True) +def home(tmp_path, monkeypatch): + monkeypatch.setenv("LABAGENT_HOME", str(tmp_path)) + + +def launch(fake, **params): + return run_pack(RagIde(), "launch", "c1", {"learners": "2", "image": "pypi", **params}, services={"modal": fake}) + + +def test_registered_and_resolvable_by_alias(): + assert registry.resolve("rag-ide").id == "rag-ide" and registry.resolve("web-ide").id == "rag-ide" + + +def test_learner_names_count_list_and_limits(): + assert learner_names("3") == ["learner-01", "learner-02", "learner-03"] + assert learner_names("Ana, Ben Q") == ["ana", "ben-q"] + for bad in ("", "a,A", "99"): + with pytest.raises(Exception): + learner_names(bad) + + +def test_launch_starts_each_learner_with_caps_and_writes_private_credentials(tmp_path): + fake = FakeModal() + env = launch(fake, hours=3, idle_minutes=10, cpu=1.5) + assert env.exit_code == 0 and len(fake.launches) == 2 and fake.builds == 1 + assert fake.launches[0]["timeout_s"] == 3 * 3600 and fake.launches[0]["idle_s"] == 600 and fake.launches[0]["cpu"] == 1.5 + assert fake.launches[0]["env"] == {} # no LLM key unless asked + path = creds_path("c1") + assert path.stat().st_mode & 0o777 == 0o600 + rows = list(csv.DictReader(path.open())) + assert [r["learner"] for r in rows] == ["learner-01", "learner-02"] and all(len(r["password"]) >= 12 for r in rows) + assert len({r["password"] for r in rows}) == 2 + # passwords never reach persisted state or the envelope + blob = str(load_state("rag-ide", "c1")) + str(env.__dict__) + assert not any(r["password"] in blob for r in rows) + + +def test_rerun_resumes_and_only_restarts_dead_learners(): + fake = FakeModal() + launch(fake) + first = {r["learner"]: r["password"] for r in csv.DictReader(creds_path("c1").open())} + fake.sandboxes["sb-1"] = False # learner-01 timed out + launch(fake) + assert len(fake.launches) == 3 and fake.launches[2]["learner"] == "learner-01" + after = {r["learner"]: r["password"] for r in csv.DictReader(creds_path("c1").open())} + assert after["learner-02"] == first["learner-02"] and after["learner-01"] != first["learner-01"] + + +def test_groq_shared_requires_a_key_and_passes_it_only_when_asked(monkeypatch): + monkeypatch.delenv("GROQ_API_KEY", raising=False) + env = launch(FakeModal(), llm="groq-shared") + assert env.exit_code != 0 and "GROQ_API_KEY" in str(env.__dict__) + monkeypatch.setenv("GROQ_API_KEY", "gsk_test") + fake = FakeModal() + assert launch(fake, llm="groq-shared").exit_code == 0 + assert fake.launches[0]["env"]["GROQ_API_KEY"] == "gsk_test" and fake.launches[0]["env"]["LLM_PROVIDER"] == "groq" + assert "gsk_test" not in str(load_state("rag-ide", "c1")) and "gsk_test" not in creds_path("c1").read_text() + + +def test_not_logged_in_fails_with_the_fix(): + env = launch(FakeModal(logged_in=False)) + assert env.exit_code != 0 and "modal token new" in str(env.__dict__) + + +def test_bad_inputs_are_rejected_before_anything_runs(): + fake = FakeModal() + for bad in ({"hours": 0}, {"hours": 25}, {"llm": "other"}, {"learners": "x,x"}): + assert launch(fake, **bad).exit_code != 0 + assert fake.launches == [] and fake.builds == 0 + + +def test_destroy_stops_everyone_deletes_volumes_and_credentials(): + fake = FakeModal() + launch(fake) + env = run_pack(RagIde(), "destroy", "c1", {}, services={"modal": fake}) + assert env.exit_code == 0 and not any(fake.sandboxes.values()) + assert sorted(fake.deleted) == sorted([mg.volume_name("c1", "learner-01"), mg.volume_name("c1", "learner-02")]) and not creds_path("c1").exists() + + +def test_destroy_keep_work_keeps_volumes(): + fake = FakeModal() + launch(fake) + run_pack(RagIde(), "destroy", "c1", {"keep_work": True}, services={"modal": fake}) + assert fake.deleted == [] and not any(fake.sandboxes.values()) + + +def test_destroy_also_stops_orphans_found_by_tag(): + fake = FakeModal() + fake.sandboxes["orphan"] = True # started earlier, state file lost + run_pack(RagIde(), "destroy", "c1", {}, services={"modal": fake}) + assert fake.sandboxes["orphan"] is False + + +def test_start_script_drops_root_and_sets_idle_shutdown(): + s = mg.start_script(900) + assert "runuser -u coder" in s and "--idle-timeout-seconds 900" in s and "--auth password" in s and s.index("chown") < s.index("exec") and "readlink -f" in s + + +def test_volume_and_app_names_are_safe(): + assert mg.volume_name("My Cohort!", "Ana B.") == "labagent-rag-ide-my-cohort-ana-b" + assert len(mg.volume_name("x" * 90, "y" * 90)) <= 64 + + +def test_real_modal_config_is_read_correctly(monkeypatch): + pytest.importorskip("modal") + import modal.config + from labagent.packs.base import PackError + gw = mg.ModalGateway() + monkeypatch.setenv("MODAL_TOKEN_ID", "ak-test") + monkeypatch.setenv("MODAL_TOKEN_SECRET", "as-test") + assert gw.check_auth() # real config.get("token_id") path, the one that once raised KeyError + monkeypatch.setattr(modal.config.config, "get", lambda key, *a, **k: None) + with pytest.raises(PackError, match="modal token new"): + gw.check_auth() + + +def run_logged(fake, **params): + lines = [] + env = run_pack(RagIde(), "launch", "c1", {"image": "pypi", **params}, services={"modal": fake}, log=lines.append) + return env, "\n".join(lines) + + +def test_urls_are_printed_and_small_cohorts_also_see_passwords(): + fake = FakeModal() + env, out = run_logged(fake, learners="2") + rows = list(csv.DictReader(creds_path("c1").open())) + assert env.exit_code == 0 and all(r["url"] in out for r in rows) and all(r["password"] in out for r in rows) + + +def test_large_cohorts_print_urls_but_not_passwords_unless_asked(): + env, out = run_logged(FakeModal(), learners="5") + rows = list(csv.DictReader(creds_path("c1").open())) + assert all(r["url"] in out for r in rows) and not any(r["password"] in out for r in rows) and "cat " in out + _, shown = run_logged(FakeModal(), learners="5", show_passwords="yes") + assert all(r["password"] in shown for r in csv.DictReader(creds_path("c1").open())) + _, hidden = run_logged(FakeModal(), learners="2", show_passwords="no") + assert not any(r["password"] in hidden for r in csv.DictReader(creds_path("c1").open())) + + +def test_rerun_prints_urls_again_with_the_original_passwords(): + fake = FakeModal() + launch(fake) + _, out = run_logged(fake, learners="2") + rows = list(csv.DictReader(creds_path("c1").open())) + assert all(r["url"] in out and r["password"] in out for r in rows) + + +def test_missing_modal_package_is_a_clean_failure_and_dry_run_needs_no_modal(monkeypatch): + import builtins + real_import = builtins.__import__ + + def no_modal(name, *a, **k): + if name == "modal" or name.startswith("modal."): + raise ImportError("No module named 'modal'") + return real_import(name, *a, **k) + + monkeypatch.setattr(builtins, "__import__", no_modal) + dry = run_pack(RagIde(), "launch", "c1", {"learners": "2"}, dry_run=True) + assert dry.exit_code == 0 and [i["step"] for i in dry.items] == ["preflight", "image", "learners", "credentials"] + env = run_pack(RagIde(), "launch", "c1", {"learners": "2"}) # real services, no modal installed + assert env.exit_code != 0 and "labagent[modal]" in str(env.__dict__) diff --git a/tests/test_rag_local.py b/tests/test_rag_local.py new file mode 100644 index 0000000..7517fea --- /dev/null +++ b/tests/test_rag_local.py @@ -0,0 +1,342 @@ +import asyncio +import json +import urllib.error + +import pytest + +pytest.importorskip("pydantic") + +from labagent.ai.rag import RAGLabClient, RAGSettings, chunk_text # noqa: E402 +from labagent.ai.rag import cli as rag_cli # noqa: E402 +from labagent.ai.rag.adapters import MemoryAdapter, VectorStoreError, make_adapter # noqa: E402 +from labagent.ai.rag.config import load_dotenv # noqa: E402 +from labagent.ai.rag.llm import LLMGateway # noqa: E402 + +S = dict(vector_db_type="memory", embedding_model_name="hashing") + + +def offline(url, headers, body, timeout): + raise urllib.error.URLError("offline") + + +def make(tmp_path, **kw): + rag_cli.init(tmp_path, force=True) + s = RAGSettings(**{**S, **kw}) + return RAGLabClient(settings=s, llm=LLMGateway(s, transport=offline)) + + +def test_init_creates_project_and_refuses_overwrite(tmp_path, capsys): + assert rag_cli.init(tmp_path) == 0 + assert (tmp_path / ".env").stat().st_mode & 0o777 == 0o600 + nb = json.loads((tmp_path / "notebooks" / "01_build_your_first_rag.ipynb").read_text()) + assert nb["nbformat"] == 4 and any("RAGLabClient" in "".join(c["source"]) for c in nb["cells"]) + assert len(list((tmp_path / "data" / "sample_docs").iterdir())) == 3 + assert (tmp_path / "src" / "main.py").is_file() and (tmp_path / "notebooks" / "01_build_your_first_rag.ipynb").is_file() + assert rag_cli.init(tmp_path) == 2 and "--force" in capsys.readouterr().err + + +def test_ingest_query_inspect_reset(tmp_path): + c = make(tmp_path) + r = c.ingest_documents([tmp_path / "data", tmp_path / "missing.md", tmp_path / ".env"]) + assert r.documents == 2 + (1 if pypdf_available() else 0) and r.chunks >= 2 and set(r.skipped) >= {str(tmp_path / "missing.md"), str(tmp_path / ".env")} + resp = c.query("how often must secrets be rotated", top_k=2) + assert resp.chunks[0].doc_id == "security-policy" and len(resp.chunks) == 2 and resp.retrieved_chunks is resp.chunks + assert 0 <= resp.chunks[0].similarity <= 1 and resp.chunks[0].distance == pytest.approx(1 - resp.chunks[0].similarity, abs=1e-3) + assert resp.metrics.provider == "mock" and resp.metrics.total_tokens > 0 and resp.metrics.tokens_estimated + assert "mock answer" in resp.answer + info = c.inspect_vectors("security-policy") + assert info["embedding_dim"] == 384 and info["chunks"] >= 1 and len(info["items"][0]["embedding_head"]) == 8 + with pytest.raises(KeyError): + c.inspect_vectors("nope") + c.reset() + with pytest.raises(ValueError, match="ingest_documents"): + c.query("anything") + + +def pypdf_available(): + try: + import pypdf # noqa: F401 + return True + except ImportError: + return False + + +def test_sdk_accepts_vector_db_and_llm_provider_and_single_path(tmp_path, monkeypatch): + for k in ("GROQ_API_KEY", "VECTOR_DB_TYPE", "LLM_PROVIDER"): + monkeypatch.delenv(k, raising=False) + rag_cli.init(tmp_path) + rag = RAGLabClient(vector_db="memory", llm_provider="groq", settings=RAGSettings(embedding_model_name="hashing")) + assert rag.settings.vector_db_type == "memory" and rag.settings.llm_provider == "groq" + assert rag.ingest(str(tmp_path / "data" / "sample_docs")).chunks > 0 + rag.llm = LLMGateway(rag.settings, transport=offline) # no GROQ_API_KEY and no network: must degrade, not raise + r = rag.query("What is our enterprise security policy?") + assert r.metrics.provider == "mock" and any("GROQ" in w or "API key" in w for w in r.warnings) + + +def test_generated_pdf_is_readable(): + pytest.importorskip("pypdf") + import io + from pypdf import PdfReader + from labagent.ai.rag import samples + text = PdfReader(io.BytesIO(samples.make_pdf(samples.INCIDENT_LINES))).pages[0].extract_text() + assert "Incident Response Runbook" in text and "72 hours" in text + + +def test_pdf_without_pypdf_is_skipped_with_fix(tmp_path, monkeypatch): + from labagent.ai.rag import client as cl + monkeypatch.setattr(cl, "_has_pypdf", lambda: False) + rag_cli.init(tmp_path) + r = make(tmp_path).ingest(tmp_path / "data") + assert any("pip install pypdf" in v for v in r.skipped.values()) and r.documents == 2 + + +def test_ide_command_binds_loopback_and_requires_password(tmp_path): + d = rag_cli.ide_command("docker", tmp_path, 9000, "pw") + assert "127.0.0.1:9000:8080" in d and "pw" not in " ".join(d) and "PASSWORD" in d + loc = rag_cli.ide_command("local", tmp_path, 9000, "pw") + assert loc[:3] == ["code-server", "--bind-addr", "127.0.0.1:9000"] and "password" in loc + + +def test_ide_print_command_initialises_workspace(tmp_path, capsys, monkeypatch): + monkeypatch.setattr(rag_cli.shutil, "which", lambda n: "/usr/bin/" + n) + assert rag_cli.ide(tmp_path / "ws", 8080, "local", rag_cli.IMAGE, print_only=True) == 0 + assert (tmp_path / "ws" / ".env").exists() and "code-server" in capsys.readouterr().out + + +def test_ide_missing_code_server_and_docker_fail_clearly(tmp_path, capsys, monkeypatch): + monkeypatch.setattr(rag_cli.shutil, "which", lambda n: None) + assert rag_cli.ide(tmp_path, 8080, "local", rag_cli.IMAGE) == 2 and "code-server" in capsys.readouterr().err + assert rag_cli._pick_mode("auto") == "local" + + +def test_image_and_devcontainer_pin_what_they_download(): + from pathlib import Path + root = Path(__file__).parent.parent + df = (root / "Dockerfile.lab").read_text() + assert "sha256sum -c" in df and "all-MiniLM-L6-v2" in df and "--auth", "password" in df + dc = json.loads((root / ".devcontainer" / "rag-lab" / "devcontainer.json").read_text()) + assert (root / ".devcontainer" / "rag-lab" / dc["build"]["dockerfile"]).resolve() == (root / "Dockerfile.lab") + + +def test_async_ingest_matches_sync(tmp_path): + a = make(tmp_path) + r = asyncio.run(a.aingest_documents([tmp_path / "data"])) + b = make(tmp_path) + assert r.chunks == b.ingest_documents([tmp_path / "data"]).chunks and a.store.count() == b.store.count() + + +def test_chunking_bounds_and_overlap(): + text = "\n\n".join(f"Paragraph {i}. " + "word " * 60 for i in range(10)) + chunks = chunk_text(text, size=400, overlap=50) + assert len(chunks) > 3 and all(len(c) <= 460 for c in chunks) + assert chunk_text("x " * 2000, size=300, overlap=0) + assert chunk_text("") == [] + + +def test_provider_chain_falls_back_with_warning(): + calls = [] + + def transport(url, headers, body, timeout): + calls.append(url) + if "groq" in url: + raise urllib.error.URLError("down") + return {"choices": [{"message": {"content": "from ollama"}}], "usage": {"prompt_tokens": 7, "completion_tokens": 3}} + + g = LLMGateway(RAGSettings(groq_api_key="k"), transport=transport) + r = g.complete("q", "ctx") + assert r.provider == "ollama" and r.text == "from ollama" and (r.prompt_tokens, r.completion_tokens, r.estimated) == (7, 3, False) + assert any("groq failed" in w for w in r.warnings) and "groq" in calls[0] + + +def test_no_keys_and_no_ollama_gives_mock_not_an_exception(): + r = LLMGateway(RAGSettings(), transport=offline).complete("q", "ctx") + assert r.provider == "mock" and "GROQ_API_KEY" in r.text + + +def test_forced_provider_without_key_warns(): + r = LLMGateway(RAGSettings(llm_provider="openai"), transport=offline).complete("q", "ctx") + assert r.provider == "mock" and any("OPENAI_API_KEY" in w or "API key" in w for w in r.warnings) + + +def test_anthropic_payload_and_usage(): + seen = {} + + def transport(url, headers, body, timeout): + seen.update(url=url, headers=headers, body=body) + return {"content": [{"type": "text", "text": "hi"}], "usage": {"input_tokens": 5, "output_tokens": 2}} + + r = LLMGateway(RAGSettings(llm_provider="anthropic", anthropic_api_key="k"), transport=transport).complete("q", "ctx") + assert r.text == "hi" and r.prompt_tokens == 5 and seen["url"].endswith("/messages") and seen["headers"]["x-api-key"] == "k" + + +def test_keys_never_in_repr(): + assert "secret" not in repr(RAGSettings(groq_api_key="secret", qdrant_api_key="secret")) + + +def test_settings_from_env_and_dotenv(tmp_path, monkeypatch): + s = RAGSettings.from_env({"VECTOR_DB_TYPE": "qdrant", "GROQ_API_KEY": "g", "EMBEDDING_MODEL_NAME": "x"}) + assert (s.vector_db_type, s.groq_api_key, s.embedding_model_name) == ("qdrant", "g", "x") + env = tmp_path / ".env" + env.write_text("# c\nRAG_T_A=1\nRAG_T_B=\nRAG_T_C='quoted'\n") + monkeypatch.setenv("RAG_T_A", "keep") + monkeypatch.delenv("RAG_T_C", raising=False) + assert load_dotenv(env) == {"RAG_T_C": "quoted"} + + +def test_generated_env_template_parses_to_clean_values(tmp_path, monkeypatch): + for k in ("VECTOR_DB_TYPE", "EMBEDDING_MODEL_NAME", "GROQ_API_KEY"): + monkeypatch.delenv(k, raising=False) + rag_cli.init(tmp_path) + load_dotenv(tmp_path / ".env") + s = RAGSettings.from_env() + assert (s.vector_db_type, s.embedding_model_name, s.groq_api_key) == ("chroma", "all-MiniLM-L6-v2", None) + monkeypatch.delenv("VECTOR_DB_TYPE") + monkeypatch.delenv("EMBEDDING_MODEL_NAME") + + +def test_unknown_store_and_missing_chroma_degrade_clearly(): + with pytest.raises(VectorStoreError, match="unknown VECTOR_DB_TYPE"): + make_adapter(RAGSettings(vector_db_type="nope"), 8) + try: + import chromadb # noqa: F401 + except ImportError: + w = [] + assert isinstance(make_adapter(RAGSettings(vector_db_type="chroma"), 8, w), MemoryAdapter) and "labagent[rag-local]" in w[0] + + +def test_cli_is_wired(tmp_path, capsys): + from labagent.cli import main + assert main(["rag", "init", "--dir", str(tmp_path / "p")]) == 0 + assert (tmp_path / "p" / "data" / "sample_docs" / "security_policy.md").exists() + + +def test_http_client_sends_an_explicit_user_agent(monkeypatch): + from labagent.ai.rag import llm + seen = {} + + class R: + def __enter__(self): return self + def __exit__(self, *a): return False + def read(self): return b"{}" + + def fake(req, timeout): + seen.update({k.lower(): v for k, v in req.header_items()}) + return R() + + monkeypatch.setattr(llm.urllib.request, "urlopen", fake) + llm.http_json("https://x.example/v1", {"Authorization": "Bearer k"}, {}, 5) + assert seen["user-agent"].startswith("labagent-rag/") and "python-urllib" not in seen["user-agent"].lower() + + +def test_http_status_is_reported_not_hidden(): + def blocked(url, headers, body, timeout): + raise urllib.error.HTTPError(url, 403, "Forbidden", {}, None) + + r = LLMGateway(RAGSettings(llm_provider="groq", groq_api_key="k"), transport=blocked).complete("q", "ctx") + assert r.provider == "mock" and any("HTTP 403" in w for w in r.warnings) + + +def test_dotenv_accepts_export_prefix_and_empty_existing_var(tmp_path, monkeypatch): + env = tmp_path / ".env" + env.write_text("export RAG_T_X=abc\nRAG_T_Y=def\n") + monkeypatch.delenv("RAG_T_X", raising=False) + monkeypatch.setenv("RAG_T_Y", "") + assert load_dotenv(env) == {"RAG_T_X": "abc", "RAG_T_Y": "def"} + monkeypatch.delenv("RAG_T_X"); monkeypatch.delenv("RAG_T_Y") + + +def test_unavailable_model_gives_an_actionable_warning(): + def nope(url, headers, body, timeout): + raise urllib.error.HTTPError(url, 404, "Not Found", {}, None) + + r = LLMGateway(RAGSettings(llm_provider="groq", groq_api_key="k"), transport=nope).complete("q", "ctx") + assert any("GROQ_MODEL" in w for w in r.warnings) and RAGSettings().groq_model == "openai/gpt-oss-20b" + + +ROWS = [ + {"id": "AAPL-2025-1A-1", "title": "Apple Risk", "content": "Apple depends on suppliers in Asia and faces tariff risk on components.", + "meta": {"ticker": "AAPL", "item": "1A", "fiscal_year": 2025, "url": "https://example.com/aapl"}}, + {"id": "AAPL-2025-7-1", "title": "Apple MD&A", "content": "Apple net sales increased driven by services revenue growth.", + "meta": {"ticker": "AAPL", "item": "7", "fiscal_year": 2025}}, + {"id": "WMT-2025-1A-1", "title": "Walmart Risk", "content": "Walmart faces competition and supplier risk including tariff exposure.", + "meta": {"ticker": "WMT", "item": "1A", "fiscal_year": 2025}}, + {"id": "WMT-2025-7-1", "title": "Walmart MD&A", "content": "Walmart comparable sales rose and operating income grew this year.", + "meta": {"ticker": "WMT", "item": "7", "fiscal_year": 2025}}, +] + + +def stores(): + out = [lambda: MemoryAdapter()] + try: + import chromadb # noqa: F401 + from labagent.ai.rag.adapters import ChromaAdapter + out.append(lambda: ChromaAdapter(RAGSettings(collection="t" + __import__("uuid").uuid4().hex[:8]))) + except ImportError: + pass + try: + import qdrant_client # noqa: F401 + from labagent.ai.rag.adapters import QdrantAdapter + out.append(lambda: QdrantAdapter(RAGSettings(collection="t" + __import__("uuid").uuid4().hex[:8]), 384)) + except ImportError: + pass + return out + + +@pytest.mark.parametrize("make_store", stores()) +def test_chunks_filters_titles_and_citations_work_on_every_store(make_store): + s = RAGSettings(**S) + rag = RAGLabClient(settings=s, store=make_store(), llm=LLMGateway(s, transport=offline)) + seen = [] + res = rag.ingest_chunks(ROWS, progress=lambda d, t: seen.append((d, t))) + assert res.chunks == 4 and res.doc_ids == ["AAPL-2025-1A", "AAPL-2025-7", "WMT-2025-1A", "WMT-2025-7"] and seen[-1] == (4, 4) + only_wmt = rag.retrieve("tariff risk", top_k=4, where={"ticker": "WMT"}) + assert {c.metadata["ticker"] for c in only_wmt} == {"WMT"} and len(only_wmt) == 2 + both = rag.retrieve("tariff risk", top_k=4, where={"ticker": ["AAPL", "WMT"], "item": "1A"}) + assert {c.chunk_id for c in both} == {"AAPL-2025-1A-1", "WMT-2025-1A-1"} + top = rag.retrieve("supplier tariff risk", top_k=1, where={"ticker": "AAPL", "fiscal_year": 2025})[0] + assert top.title == "Apple Risk" and top.url == "https://example.com/aapl" # int filter value matches the stored "2025" + assert rag.retrieve("anything", where={"ticker": "NOPE"}) == [] + r = rag.query("What tariff risks does Apple have?", where={"ticker": "AAPL"}) + assert all(c.metadata["ticker"] == "AAPL" for c in r.retrieved_chunks) + assert rag.inspect_vectors("AAPL-2025-1A")["items"][0]["chunk_id"] == "AAPL-2025-1A-1" + + +def test_keyword_weight_zero_is_pure_vector_and_changes_ranking(): + s = RAGSettings(**S) + rag = RAGLabClient(settings=s, llm=LLMGateway(s, transport=offline)) + rag.ingest_chunks(ROWS) + pure = rag.retrieve("tariff", top_k=4, keyword_weight=0.0) + assert all(abs(c.similarity - (1 - c.distance)) < 1e-3 for c in pure) + mixed = rag.retrieve("tariff", top_k=4, keyword_weight=0.9) + assert mixed[0].keyword_score >= pure[-1].keyword_score and mixed[0].keyword_score > 0 + + +def test_evaluate_reports_hit_rate_misses_and_gate(): + s = RAGSettings(**S) + rag = RAGLabClient(settings=s, llm=LLMGateway(s, transport=offline)) + rag.ingest_chunks(ROWS) + cases = [{"question": "Apple supplier tariff risk", "expect": {"ticker": "AAPL", "item": "1A"}}, + {"question": "Walmart comparable sales operating income", "expect": {"ticker": "WMT", "item": "7"}}, + {"question": "Chevron oil prices", "expect": {"ticker": "CVX"}}] + r = rag.evaluate(cases, top_k=2, min_hit_rate=0.6) + assert (r.cases, r.hits, r.passed) == (3, 2, True) and r.misses[0]["expected"] == {"ticker": "CVX"} + assert rag.evaluate(cases, top_k=2, min_hit_rate=0.9).passed is False + assert rag.evaluate(cases, top_k=2, where={"ticker": "AAPL"}).hits >= 1 + + +def test_ingest_chunks_skips_empty_rows_and_accepts_text_key(): + s = RAGSettings(**S) + rag = RAGLabClient(settings=s, llm=LLMGateway(s, transport=offline)) + res = rag.ingest_chunks([{"id": "a-1", "text": "alpha beta"}, {"id": "b-1", "content": ""}, {"id": "c", "content": "gamma", "doc_id": "custom"}]) + assert res.chunks == 2 and set(res.doc_ids) == {"a", "custom"} + + +def test_sec_tutorial_notebook_is_valid_and_its_code_compiles(): + import ast + from labagent.ai.rag import samples + nb = json.loads(samples.sec_notebook()) + code = ["".join(c["source"]) for c in nb["cells"] if c["cell_type"] == "code"] + assert len(code) >= 8 + for src in code: + ast.parse(src) + assert "notebooks/02_sec_filings_rag.ipynb" in samples.workspace()