diff --git a/README.md b/README.md index ec4fd41..1590f41 100644 --- a/README.md +++ b/README.md @@ -14,27 +14,25 @@ A lab is a **pack**: ordered steps that are safe to re-run. `launch` builds it and resumes after a failure. `destroy` removes only what the pack created. -- **Databricks labs.** RAG over real SEC 10-K filings, and a text-to-SQL agent, in your own Free Edition workspace. +- **Databricks labs.** A step-by-step RAG tutorial over real SEC 10-K filings (notebooks you run), and a text-to-SQL agent, in your own Free Edition workspace. - **Web IDE per learner.** Private browser VS Code with the RAG lab preinstalled, on Modal, EC2 or local Docker. - **AWS sandboxes.** Cohort IAM with Terraform, plus plan-then-apply cleanup of shared accounts. -- **Plain English.** `labagent "deploy rag-sec"` shows the command it understood and asks first (local Ollama, any OpenAI-compatible endpoint). +- **Plain English.** `labagent "deploy rag-tutorial"` shows the command it understood and asks first (local Ollama, any OpenAI-compatible endpoint). ## Quick start ```bash pip install 'labagent[rag]' # Python 3.9+; sign in with: databricks auth login labagent packs # list the labs -labagent launch rag-sec-filings --name demo -labagent destroy rag-sec-filings --name demo +labagent launch rag-sec-tutorial --name demo # uploads the notebooks to your workspace; you run them +labagent destroy rag-sec-tutorial --name demo # removes the uploaded folder ``` ## Labs | Pack | What you get | Runs on | |---|---|---| -| `rag-sec-filings` | RAG over a pinned SEC 10-K dataset: hybrid search, filters, citations, a quality gate before deploy | Databricks Free Edition | -| `rag-sec-filings-live` | Same, fetching fresh filings from EDGAR | Databricks Free Edition | -| `rag-lakebase` | Small demo corpus to smoke-test your setup | Databricks Free Edition | +| `rag-sec-tutorial` | Nine notebooks that build RAG on a pinned SEC 10-K dataset one step at a time: embeddings, pgvector, hybrid search, filters, citations, an evaluation gate, and the deployed app. Each ends with a check cell | Databricks Free Edition | | `db-agent` | [db-agent](https://github.com/db-agent/db-agent) text-to-SQL over a Databricks sample dataset, pinned commit | Databricks Free Edition | | `rag-ide` | A private browser VS Code per learner with the RAG lab and notebooks | Modal, EC2 or local Docker | diff --git a/docs/job-contract.md b/docs/job-contract.md index 133c448..61702d2 100644 --- a/docs/job-contract.md +++ b/docs/job-contract.md @@ -8,7 +8,7 @@ | Field | Type | Notes | |---|---|---| | `run_id` | string | Echoed back in the result | -| `command` | `provision` `inventory` `sweep` `reset` `offboard` `audit` `cleanup` | Anything else is refused | +| `command` | `provision` `inventory` `sweep` `reset` `offboard` `audit` `cleanup` `launch` `destroy` | Anything else is refused | | `mode` | `plan` `apply` `dry-run` | `sweep` needs `plan` or `apply`; `provision` supports `dry-run` | | `cohort` | `{id, created_date}` | `created_date` is `dd-mmm-yyyy` and matches the `date` tag | | `learners` | `[{learner_id, username, active?}]` | Roster. Held in memory and on the runner's ephemeral disk only | @@ -16,6 +16,7 @@ | `approved_hash` | string | The plan hash a human approved. Apply is refused on mismatch | | `secrets_pubkey_pem` | PEM | RSA public key. Required for `provision` and `reset` (passwords are RSA-OAEP SHA-256 encrypted to it) | | `ruleset` | object | For `audit` and `cleanup` (plan). See docs/account-cleanup.md | +| `pack`, `name`, `params` | string, string, object | For `launch` and `destroy`: a pack id or alias, the instance name (`[a-z0-9-]`, 31 characters at most) and its parameters. `image` and `profile` are refused. These commands need no AWS account allow list: the pack brings its own credentials | | `options` | object | `all_cohorts`, `force`, `learner_id` (reset), `only_learner_id` (sweep a single learner), `password_reset_required` (provision/reset; default from config), `regions` (audit/cleanup; `["*"]` = all enabled regions) | ## Result (labagent to orchestrator) @@ -33,6 +34,8 @@ } ``` +For `launch`, `secrets[]` holds one `ide_password` per learner, encrypted to `secrets_pubkey_pem` (also when the launch fails part way, so passwords already handed out are not lost). `summary.outputs` carries URLs and hosts only. + Exit code: 0 ok, 1 partial (something is `pending` or `failed`: re-run), 2 failed (nothing was done or the input was refused). ## Rules the engine enforces diff --git a/docs/packs.md b/docs/packs.md index 913d3b8..9298cf4 100644 --- a/docs/packs.md +++ b/docs/packs.md @@ -5,9 +5,10 @@ credentials, or in your CI. Add a lab by writing a `Pack` subclass; the framewor ```bash labagent packs # list -labagent launch rag-lakebase --name demo --dry-run -labagent launch rag-lakebase --name demo # re-run to resume after a failure -labagent destroy rag-lakebase --name demo # removes only what this pack created +labagent packs --json # catalog for dashboards: params, platforms, cost tier +labagent launch rag-sec-tutorial --name demo --dry-run +labagent launch rag-sec-tutorial --name demo # re-run to resume after a failure +labagent destroy rag-sec-tutorial --name demo # removes only what this pack created ``` ## Writing a pack @@ -27,45 +28,50 @@ Ship it in your own package and register it: `[project.entry-points."labagent.pa Rules: `check` is a read-only "already done" test; step outputs are non-secret; tokens stay in memory (`ctx.mem`). State is a 0600 file at `~/.labagent/packs//.json`. -## RAG packs (Databricks Lakebase) +## RAG tutorial (Databricks Lakebase) -One reusable pipeline, pluggable data sources. Each stage is idempotent, and `evaluate` gates the deploy: +`rag-sec-tutorial` builds nothing itself. It uploads nine notebooks and the app code into your workspace, and you run them in order. +Each notebook is short, safe to re-run, and ends with a check cell that stops with a clear message if the step did not work. -`preflight, project, locate, build (fetch, parse, chunk), schema, load (embed), index, evaluate, app` - -| Pack | Source | Use | -|---|---|---| -| `rag-sec-filings` | Pinned snapshot of real SEC 10-Ks ([dataset release](https://github.com/becloudready/labagent-datasets/releases/tag/sec-10k-2026.10)), split into Items. No SEC access or contact needed | Financial research | -| `rag-sec-filings-live` | The same corpus downloaded fresh from EDGAR at launch (needs a contact for the SEC) | Newer filings, other tickers | -| `rag-lakebase` | 25 invented passages | Smoke test only | +| Notebook | You | +|---|---| +| `00_setup` | Check your model endpoints, create the Lakebase project | +| `01_data` | Download the pinned SEC 10-K dataset (checked against a SHA-256) and explore it | +| `02_embeddings` | Call the embedding model, compare similarity, time batches | +| `03_pgvector` | Create the table, load the filings in two stages (everyone but Apple, then Apple) to see why embedding decides what search can find | +| `04_search` | Add HNSW and full-text indexes, run vector, keyword and hybrid (reciprocal rank fusion) search, add filters | +| `05_answer_and_eval` | Generate cited answers, try a question with no answer, score retrieval against a 0.6 gate | +| `06_deploy_app` | Create the Databricks App, give its own identity a read-only role, deploy | +| `99_cleanup` | Remove the app, the schema and the database project | ```bash pip install 'labagent[rag]' # needs the Databricks CLI, then: databricks auth login -labagent launch rag-sec-filings --name sec --param "user_agent=Full Name name@example.com" \ - --param tickers=AAPL,JPM,PFE --param filings=2 +labagent launch rag-sec-tutorial --name demo ``` -What makes it more than a toy: section-aware chunking, hybrid retrieval (pgvector + full-text, rank-fused), metadata filters -(`{"ticker": "AAPL"}`), numbered citations with the filing URL, and an evaluation set that must reach `min_hit_rate` -(default 0.6) before the app is deployed. The snapshot is downloaded from a GitHub release and verified against a SHA-256 committed in -`labagent/packs/rag/sources/datasets.json`; a mismatch is refused. The live pack follows EDGAR rules (declared contact, under 10 requests -per second, cached locally). Filings are never stored in this repo. +It prints a link to the uploaded folder. Open it, start with `00_setup`, and finish with `99_cleanup`. The notebooks run on serverless +compute, need no cluster, and download the dataset themselves. Free Edition allows one Lakebase project and three apps, and stops apps +after 24 hours. `labagent destroy rag-sec-tutorial --name demo` removes only the uploaded folder: run `99_cleanup` first to remove what the +notebooks created. -Free Edition allows one Lakebase project: pass `--param project=` to share it. Each source uses its own schema, and -`destroy` drops only that schema (and the project only if labagent created it). Apps stop after 24 hours. +The one-command versions of this lab (`rag-lakebase`, `rag-sec-filings`, `rag-sec-filings-live`) were removed in favour of the tutorial. +Anyone who types one of those names gets a message pointing here. They are still in the `v0.2.1` release and git tag. -Add a source by subclassing `labagent.packs.rag.sources.Source` (`chunks()` and `eval_set()`) and registering it under the -`labagent.rag_sources` entry point. +Add a data source by subclassing `labagent.packs.rag.sources.Source` and registering it under the `labagent.rag_sources` entry point. ## Publishing a dataset snapshot (maintainers) ```bash -labagent dataset --chunks ~/.labagent/cache/rag-sec-filings-live/.jsonl --name sec-10k --version 2026.11 --out out/ +labagent dataset --chunks chunks.jsonl --name sec-10k --version 2026.11 --out out/ gh release create sec-10k-2026.11 out/* --repo becloudready/labagent-datasets ``` -The build is byte-for-byte reproducible. Then add the release URL and the printed `sha256` to `datasets.json` in a pull request. -Data limitations (for example, no Item 7 for JPM and CVX) are listed in the dataset README. +`chunks.jsonl` has one JSON object per line: `id`, `title`, `content`, `meta`. The build is byte-for-byte reproducible. Then add the release +URL and the printed `sha256` to `labagent/packs/rag/sources/datasets.json` in a pull request. Data limitations (for example, no Item 7 for +JPM and CVX) are listed in the dataset README. + +The SEC chunks were produced by the `rag-sec-filings-live` pack, which no longer exists. The parser it used is still in the code +(`labagent.packs.rag.sources.sec.SecFilings`) but nothing calls it any more. To rebuild the SEC dataset from EDGAR, use the `v0.2.1` tag. ## db-agent (text-to-SQL on Databricks) @@ -91,8 +97,8 @@ Steps: `preflight` (login, SQL warehouse, sample data, serving endpoints), `fetc ## Plain English ```bash -labagent "Run and deploy rag-sec in databricks" # shows the command, asks, then runs it -labagent "tear down rag-sec" --print-only # just show the command +labagent "Run and deploy rag-tutorial in databricks" # shows the command, asks, then runs it +labagent "tear down rag-tutorial" --print-only # just show the command labagent ask what can you deploy # 'ask' lets you skip the quotes ``` @@ -103,7 +109,7 @@ A remote endpoint is announced before your sentence is sent to it. The model never runs anything, and it is not trusted for details: - The instance name (`called x`) and parameters (`profile DEFAULT`, `tickers AAPL,JPM`) are read from your sentence by exact rules, so they cannot be invented. -- A pack you typed (`rag-sec`, or any id or alias) wins over the model's pick. If the model guessed, the pack needs a distinctive word from your sentence. +- A pack you typed (`rag-tutorial`, or any id or alias) wins over the model's pick. If the model guessed, the pack needs a distinctive word from your sentence. - Requests that ask to create and delete in one sentence are refused. A delete needs a delete word, and you confirm by typing the instance name. - Without a terminal, nothing runs unless you pass `--yes`. diff --git a/examples/01_rag_starter_lab.py b/examples/01_rag_starter_lab.py index 744e512..f72a2e6 100644 --- a/examples/01_rag_starter_lab.py +++ b/examples/01_rag_starter_lab.py @@ -5,7 +5,7 @@ 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`). +This is the laptop version of the labagent Databricks packs (`labagent launch rag-sec-tutorial`). """ from pathlib import Path diff --git a/labagent/ai/rag/adapters.py b/labagent/ai/rag/adapters.py index ec405f2..f0e8132 100644 --- a/labagent/ai/rag/adapters.py +++ b/labagent/ai/rag/adapters.py @@ -2,7 +2,7 @@ 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`. +`labagent launch rag-sec-tutorial`. """ import math import uuid diff --git a/labagent/ai/rag/client.py b/labagent/ai/rag/client.py index 281c909..c6ce3a5 100644 --- a/labagent/ai/rag/client.py +++ b/labagent/ai/rag/client.py @@ -5,7 +5,7 @@ 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. +it in a browser, and the Databricks tutorial (`labagent launch rag-sec-tutorial`) builds the same pipeline on Lakebase. """ import asyncio import re diff --git a/labagent/ai/rag/samples.py b/labagent/ai/rag/samples.py index dbedad8..e55bc8b 100644 --- a/labagent/ai/rag/samples.py +++ b/labagent/ai/rag/samples.py @@ -183,7 +183,7 @@ def code(src): return {"cell_type": "code", "metadata": {}, "execution_count": 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 a managed version on Databricks Lakebase, see `labagent launch rag-sec-tutorial`.\n"), ] for i, c in enumerate(cells): c["id"] = f"cell{i:02d}" # nbformat 4.5 requires cell ids @@ -258,7 +258,7 @@ def code(src): return {"cell_type": "code", "metadata": {}, "execution_count": 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"), + "```\npip install 'labagent[rag]'\nlabagent launch rag-sec-tutorial --name demo\n```\n"), ] for i, c in enumerate(cells): c["id"] = f"cell{i:02d}" # nbformat 4.5 requires cell ids diff --git a/labagent/cli.py b/labagent/cli.py index 66587dc..57d0e99 100644 --- a/labagent/cli.py +++ b/labagent/cli.py @@ -12,7 +12,7 @@ from .config import STATE_BACKEND_DIR, load_config from .context import Ctx, GuardError from .envelope import Envelope -from .run import COMMANDS, run_job +from .run import COMMANDS, PACK_COMMANDS, run_job, run_pack_job def _emit(env: Envelope, output: str) -> None: @@ -28,6 +28,9 @@ def _packs_main(args) -> int: from .packs import registry from .packs.base import run_pack if args.cmd == "packs": + if args.json: + print(json.dumps({"packs": [pk.describe() for _, pk in sorted(registry.all_packs().items())]}, indent=2)) + return 0 for pid, pk in sorted(registry.all_packs().items()): print(f"{pid:16} {pk.summary}") return 0 @@ -50,10 +53,10 @@ def _packs_main(args) -> int: def _nl_main(argv) -> int: - """labagent "deploy rag-sec on databricks": translate with a small model, confirm, then run the real command.""" + """labagent "deploy rag-tutorial on databricks": translate with a small model, confirm, then run the real command.""" from . import nl from .packs import registry - q = argparse.ArgumentParser(prog="labagent", description="Plain-English front end: labagent \"deploy rag-sec on databricks\"") + q = argparse.ArgumentParser(prog="labagent", description="Plain-English front end: labagent \"deploy rag-tutorial on databricks\"") q.add_argument("text", nargs="+") q.add_argument("--endpoint", help=f"OpenAI-compatible base URL (default {nl.DEFAULT_URL}, or LABAGENT_LLM_URL)") q.add_argument("--model", help=f"Model (default {nl.DEFAULT_MODEL}, or LABAGENT_LLM_MODEL)") @@ -101,11 +104,12 @@ def main(argv=None) -> int: r.add_argument("--job", required=True, help="Path to job JSON (contains roster: keep it out of git and logs)") 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") + pk = sub.add_parser("packs", help="List available lab packs") + pk.add_argument("--json", action="store_true", help="Machine-readable catalog with parameters, platforms and cost tier (used by the console)") 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("--chunks", required=True, help="JSONL of chunks: one object per line with id, title, content and meta") dp.add_argument("--name", required=True) dp.add_argument("--version", required=True) dp.add_argument("--out", required=True, help="Output directory") @@ -140,8 +144,12 @@ def main(argv=None) -> int: print("\n".join(sorted(COMMANDS))) return 0 try: - cfg = load_config(args.config) job = json.loads(Path(args.job).read_text()) + if job.get("command") in PACK_COMMANDS: # lab packs bring their own credentials: no AWS account guard here + env = run_pack_job(job) + _emit(env, args.output) + return env.exit_code + cfg = load_config(args.config) session = boto3.Session(profile_name=args.profile, region_name=cfg.region) if getattr(args, "ruleset", None) and "ruleset" not in job: job["ruleset"] = json.loads(Path(args.ruleset).read_text()) diff --git a/labagent/nl.py b/labagent/nl.py index 5e062da..ebd24d1 100644 --- a/labagent/nl.py +++ b/labagent/nl.py @@ -56,12 +56,13 @@ def build_messages(text: str, packs: Dict[str, Any]) -> List[Dict[str, str]]: model's job (they are read from the sentence by exact rules in `extract_args`), so it cannot invent them.""" plist = sorted(packs.values(), key=lambda x: x.id) lines = [f"- {p.id}" + (f" (alias: {', '.join(p.aliases)})" if p.aliases else "") + f": {p.summary.split('.')[0]}" for p in plist] + rag = next((p.id for p in plist if "tutorial" in p.id), plist[0].id) # what a vague "the rag app" request most likely means shots = [] for p in plist: short = p.aliases[0] if p.aliases else p.id shots += [(f"deploy {short}", "launch", p.id), (f"delete {p.id}", "destroy", p.id)] - shots += [("run and start the rag app on databricks", "launch", plist[0].id), ("what can you deploy?", "list", None), - ("show me the available packs", "list", None), ("tear down everything for the sec rag", "destroy", next((p.id for p in plist if "sec" in p.id), plist[0].id))] + shots += [("run and start the rag app on databricks", "launch", rag), ("what can you deploy?", "list", None), + ("show me the available packs", "list", None), ("tear down everything for the sec rag", "destroy", rag)] ex = "\n".join(f"Request: {q}\n" + json.dumps({"action": a, "pack": k}) for q, a, k in shots) system = ("Convert the request to JSON: {\"action\": ..., \"pack\": ...}. Reply with JSON only.\n" "action is launch (run, deploy, start, create, build), destroy (delete, remove, tear down) or list (show what exists). " @@ -190,7 +191,7 @@ def validate(raw: str, text: str, resolve: Callable[[str], Any], packs: Optional action = _canonical_action(str(obj.get("action") or ""), text) if action not in ACTIONS: raise NLError(f"I could not tell what to do (got '{action or 'nothing'}'). Say launch, destroy or list, e.g. " - "\"deploy rag-sec\"") + "\"deploy rag-tutorial\"") wants_destroy, wants_launch = bool(DESTROY_WORDS.search(text)), bool(LAUNCH_WORDS.search(text)) if wants_destroy and wants_launch: raise NLError("that asks for two things at once (create and delete). One action at a time: run them as two commands") @@ -198,7 +199,7 @@ def validate(raw: str, text: str, resolve: Callable[[str], Any], packs: Optional raise NLError("the model chose 'destroy' but your request does not ask to delete anything. Nothing was run") if action == "launch" and wants_destroy: raise NLError("your request sounds like a deletion but the model chose 'launch'. Nothing was run: " - "say it explicitly, e.g. \"destroy rag-sec\"") + "say it explicitly, e.g. \"destroy rag-tutorial\"") if action == "list": return Plan("list") pack, warnings = None, [] @@ -215,10 +216,10 @@ def validate(raw: str, text: str, resolve: Callable[[str], Any], packs: Optional raise NLError(str(e).strip("\"'")) from e words = set(re.split(r"[^a-z0-9]+", text.lower())) if packs and not (_distinctive_tokens(pack, packs) & words): - raise NLError("I could not tell which pack you mean (the model guessed). Name it, e.g. \"deploy rag-sec\", " + raise NLError("I could not tell which pack you mean (the model guessed). Name it, e.g. \"deploy rag-tutorial\", " "or run 'labagent packs'") if pack is None: - raise NLError("which pack? Run 'labagent packs' to see them, e.g. \"deploy rag-sec\"") + raise NLError("which pack? Run 'labagent packs' to see them, e.g. \"deploy rag-tutorial\"") args = extract_args(text, pack) plan = Plan(action, pack.id, args["name"] or pack.default_name or None, args["params"], warnings) if not plan.name: @@ -231,7 +232,7 @@ def translate(text: str, packs: Dict[str, Any], resolve: Callable[[str], Any], u chat: Optional[Callable[..., str]] = None) -> Plan: text = " ".join(text.split()) if not text: - raise NLError("say what you want, e.g. labagent \"deploy rag-sec on databricks\"") + raise NLError("say what you want, e.g. labagent \"deploy rag-tutorial on databricks\"") chat = chat or call_llm # looked up at call time so callers and tests can substitute it return validate(chat(url, model, build_messages(text, packs), api_key), text, resolve, packs) diff --git a/labagent/packs/base.py b/labagent/packs/base.py index ba3c5fb..e0e0c8b 100644 --- a/labagent/packs/base.py +++ b/labagent/packs/base.py @@ -54,6 +54,8 @@ class Pack: summary: str = "" aliases: tuple = () # short names accepted on the command line and by the natural-language translator default_name: str = "" # instance name used when --name is omitted + platforms: tuple = () # where the pack can run, for the console's platform picker (first is the default) + cost_tier: str = "" # free | low | metered. A hint for dashboards, not a price params: List[Param] = [] def deploy_steps(self, ctx: PackContext) -> List[Step]: @@ -62,6 +64,16 @@ def deploy_steps(self, ctx: PackContext) -> List[Step]: def destroy_steps(self, ctx: PackContext) -> List[Step]: raise NotImplementedError + def describe(self) -> Dict[str, Any]: + """Machine-readable description for `labagent packs --json`. No secrets: defaults are literals from the code.""" + return {"id": self.id, "title": self.title, "summary": self.summary, "aliases": list(self.aliases), + "default_name": self.default_name, "platforms": list(self.platforms), "cost_tier": self.cost_tier, + "params": [{"name": q.name, "default": q.default, "help": q.help, "required": q.required} for q in self.params]} + + def outputs(self, state: Dict[str, Any]) -> Dict[str, Any]: + """Non-secret facts from state that go in the envelope (URLs, ids). Override when state has nested or mixed content.""" + return {k: v for k, v in state.items() if k in getattr(self, "output_keys", ())} + def services(self, profile: Optional[str]) -> Dict[str, Any]: """Real service objects for this pack. Tests bypass this by passing their own services.""" return {} @@ -98,8 +110,19 @@ def save_state(pack_id: str, name: str, state: Dict[str, Any]) -> None: p.chmod(0o600) +def _seal_passwords(env: Envelope, ctx: PackContext, pubkey_pem: Optional[str]) -> None: + """Packs that mint per-user credentials leave {user: secret} in ctx.mem['passwords']. With a public key they leave only as ciphertext.""" + secrets_in = ctx.mem.get("passwords") or {} + if not pubkey_pem or not secrets_in: + return + from ..secrets_box import encrypt, load_public_key + key = load_public_key(pubkey_pem) + env.secrets = [{"user": user, "kind": "ide_password", "ciphertext": encrypt(key, pw)} for user, pw in sorted(secrets_in.items()) + if not str(pw).startswith("(")] + + def run_pack(pack: Pack, action: str, name: str, given: Dict[str, Any], services: Optional[Dict[str, Any]] = None, - dry_run: bool = False, log: Optional[Callable[[str], None]] = None) -> Envelope: + dry_run: bool = False, log: Optional[Callable[[str], None]] = None, pubkey_pem: Optional[str] = None) -> Envelope: env = Envelope(command=f"{action}:{pack.id}", mode="dry-run" if dry_run else "apply", cohort=name) try: params = pack.resolve(given) @@ -134,11 +157,12 @@ def run_pack(pack: Pack, action: str, name: str, given: Dict[str, Any], services env.summary = {"steps": len(env.items)} return env env.summary = {s: sum(1 for i in env.items if i["state"] == s) for s in ("done", "skipped", "failed")} + _seal_passwords(env, ctx, pubkey_pem) # also on failure: a half-built lab must not lose the passwords it already handed out failed = [i for i in env.items if i["state"] == "failed"] if failed: env.fail(f"step '{failed[0]['step']}' failed: {failed[0].get('error', '')}") else: - env.summary["outputs"] = {k: v for k, v in ctx.state.items() if k in getattr(pack, "output_keys", ())} + env.summary["outputs"] = pack.outputs(ctx.state) if action == "destroy": state_path(pack.id, name).unlink(missing_ok=True) return env diff --git a/labagent/packs/databricks/cli.py b/labagent/packs/databricks/cli.py index bf1c061..6724a65 100644 --- a/labagent/packs/databricks/cli.py +++ b/labagent/packs/databricks/cli.py @@ -44,19 +44,6 @@ def json(self, *args: str, body: Any = None) -> Any: out = self.text(*args, *extra, "-o", "json") return json.loads(out) if out.strip() else {} - def api(self, method: str, path: str, body: Any = None, retries: int = 6) -> Any: - """Rate limits are normal on Free Edition foundation-model endpoints: back off 2s, 4s, 8s... then give up.""" - extra = ["--json", json.dumps(body)] if body is not None else [] - for attempt in range(retries + 1): - try: - out = self.text("api", method, path, *extra) - break - except PackError as e: - if attempt == retries or not any(k in str(e) for k in ("REQUEST_LIMIT_EXCEEDED", "RATE_LIMIT", "429")): - raise - self._sleep(min(2 ** (attempt + 1), 30)) - return json.loads(out) if out.strip() else {} - def profiles(self) -> list: """Configured CLI profiles with their host and whether the login is still valid.""" saved, self.profile = self.profile, None diff --git a/labagent/packs/databricks/models.py b/labagent/packs/databricks/models.py new file mode 100644 index 0000000..4ec55aa --- /dev/null +++ b/labagent/packs/databricks/models.py @@ -0,0 +1,5 @@ +"""How packs recognise the model serving endpoints they need in a Databricks workspace.""" +import re + +EMBED_RE = re.compile(r"gte|bge|embed", re.I) +CHAT_RE = re.compile(r"llama|gpt|claude|gemma|qwen|instruct|chat", re.I) diff --git a/labagent/packs/db_agent/pack.py b/labagent/packs/db_agent/pack.py index 2b319d2..cd2472d 100644 --- a/labagent/packs/db_agent/pack.py +++ b/labagent/packs/db_agent/pack.py @@ -23,7 +23,7 @@ from ..databricks.auth import login, workspace_host from ..databricks.cli import Dbx from ..gitsrc import clone_pinned -from ..rag.pack import CHAT_RE, EMBED_RE +from ..databricks.models import CHAT_RE, EMBED_RE REPO = "https://github.com/db-agent/db-agent" PINNED = "eef0a308d052b6bdd00fdfec7ee1c9a61c01426f" # reviewed commit; change it deliberately (see docs/packs.md) @@ -77,6 +77,8 @@ def _items(resp: Any, key: str) -> List[Dict[str, Any]]: class DbAgent(Pack): id = "db-agent" aliases = ("dbagent", "db-agent-databricks", "text-to-sql") + platforms = ("databricks",) + cost_tier = "free" default_name = "demo" title = "DB Agent: text-to-SQL on Databricks" summary = "Deploy db-agent (github.com/db-agent/db-agent) to Databricks Apps over a Databricks sample dataset, then smoke-test it." diff --git a/labagent/packs/rag/app/app.py b/labagent/packs/rag/app/app.py index 1e422c6..8c1932d 100644 --- a/labagent/packs/rag/app/app.py +++ b/labagent/packs/rag/app/app.py @@ -23,6 +23,9 @@ KEY = re.compile(r"^[a-z_][a-z0-9_]*$") SHOWN_META = ("ticker", "company", "fiscal_year", "section", "item", "form", "url", "source") +# gpt-oss models spend tokens on reasoning before the answer: 500 truncated answers (and their citations) in testing; the longest seen used 868 +MAX_ANSWER_TOKENS = 1200 + SYSTEM = ("You answer questions using only the numbered context passages. Cite the passages you used like [1] or [2]. " "If the context does not contain the answer, say you do not know. Do not use outside knowledge.") @@ -99,7 +102,7 @@ def ask(a: Ask): return {"answer": "I do not know: no matching passages.", "sources": []} context = "\n\n".join(f"[{i}] {h['title']}\n{h['content']}" for i, h in enumerate(hits, 1)) out = _invoke(os.environ["CHAT_ENDPOINT"], {"messages": [{"role": "system", "content": SYSTEM}, - {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {a.question}"}], "max_tokens": 500}) + {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {a.question}"}], "max_tokens": MAX_ANSWER_TOKENS}) return {"answer": out["choices"][0]["message"]["content"], "sources": [{"n": i, "title": h["title"], "url": h["meta"].get("url"), "score": h["score"], "snippet": h["content"][:600], "meta": {k: v for k, v in h["meta"].items() if k in SHOWN_META}} diff --git a/labagent/packs/rag/pack.py b/labagent/packs/rag/pack.py deleted file mode 100644 index 6cb3ff3..0000000 --- a/labagent/packs/rag/pack.py +++ /dev/null @@ -1,329 +0,0 @@ -"""RAG on Databricks Lakebase: one reusable pipeline, pluggable data sources. - -Workflow (each stage idempotent, resumes on re-run): - preflight -> project -> locate -> build (fetch, parse, chunk) -> schema -> load (embed) -> index -> evaluate -> app - -`evaluate` is a quality gate: the app is deployed only if retrieval hits the expected passages often enough. -It runs the same hybrid-retrieval code the deployed app uses (app/retrieval.py). -Runs on the learner's machine with their own Databricks CLI profile: labagent never sees their credentials. -""" -from __future__ import annotations - -import importlib.util -import json -import os -import re -import shutil -import tempfile -from pathlib import Path -from typing import Any, Dict, List, Optional - -from ..base import Pack, PackContext, PackError, Param, Step, state_dir -from ..databricks import pg as pgmod -from ..databricks.auth import login as login_shared -from ..databricks.cli import Dbx -from . import sources as src_mod -from .chunking import make_chunks # noqa: F401 (re-exported for source authors) - -HERE = Path(__file__).parent -EMBED_RE = re.compile(r"gte|bge|embed", re.I) -CHAT_RE = re.compile(r"llama|gpt|claude|gemma|qwen|instruct|chat", re.I) -BATCH = 8 -EMBED_CHARS = 1800 # gte-large accepts 512 tokens; stay under it - - -def _retrieval(): - """Load app/retrieval.py by path: it ships inside the app bundle, not as an importable package.""" - spec = importlib.util.spec_from_file_location("labagent_rag_retrieval", HERE / "app" / "retrieval.py") - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - return mod - - -def _items(resp: Any, key: str) -> List[Dict[str, Any]]: - if isinstance(resp, list): - return resp - return (resp or {}).get(key) or [] - - -class RagPack(Pack): - source_id = "" - output_keys = ("app_url", "project", "embedding_endpoint", "chat_endpoint", "chunks", "eval_hit_rate", "eval_questions") - base_params = [ - Param("profile", None, "Databricks CLI profile (default profile if omitted)"), - Param("project", None, "Lakebase project id (default: rag-). Free Edition allows ONE project per account"), - Param("embedding_endpoint", None, "Serving endpoint for embeddings (auto-discovered if omitted)"), - Param("chat_endpoint", None, "Serving endpoint for answers (auto-discovered if omitted)"), - Param("top_k", 4, "Passages retrieved per question"), - Param("keep_deleted", False, "On destroy, keep the soft-deleted project (its name stays reserved) instead of purging it"), - Param("min_hit_rate", 0.6, "Deploy only if the evaluation hit rate reaches this (0 disables the gate)"), - ] - - def __init__(self): - self.source = src_mod.get(self.source_id) - self.params = self.base_params + list(self.source.params) - - def services(self, profile: Optional[str]) -> Dict[str, Any]: - return {"dbx": Dbx(profile), "pg": pgmod.connect} - - # -- helpers ------------------------------------------------------------- - @staticmethod - def _project(ctx: PackContext) -> str: - return ctx.params["project"] or f"rag-{ctx.name}" - - @staticmethod - def _embed(ctx: PackContext, texts: List[str]) -> List[List[float]]: - r = ctx.services["dbx"].api("post", f"/serving-endpoints/{ctx.state['embedding_endpoint']}/invocations", - {"input": [t[:EMBED_CHARS] for t in texts]}) - return [d["embedding"] for d in r["data"]] - - @staticmethod - def _conn(ctx: PackContext): - d = ctx.services["dbx"] - tok = pgmod.find(d.json("postgres", "generate-database-credential", ctx.state["lakebase_endpoint"]), "token") - if not tok: - raise PackError("could not obtain a Lakebase credential: run `databricks auth login` and retry") - return ctx.services["pg"](ctx.state["pg_host"], ctx.state["user"], tok) - - def _chunks_file(self, ctx: PackContext) -> Path: - return state_dir().parent / "cache" / self.id / f"{ctx.name}.jsonl" - - def _read_chunks(self, ctx: PackContext) -> List[Dict[str, Any]]: - return [json.loads(l) for l in self._chunks_file(ctx).read_text().splitlines() if l.strip()] - - # -- deploy -------------------------------------------------------------- - def deploy_steps(self, ctx: PackContext) -> List[Step]: - d = lambda: ctx.services["dbx"] # noqa: E731 - project = self._project(ctx) - schema = self.source.schema - - def preflight(c: PackContext): - me = login_shared(c) - user = me.get("userName") - if not user: - raise PackError("not authenticated: run `databricks auth login` and retry") - names = [e.get("name", "") for e in _items(d().json("serving-endpoints", "list"), "endpoints")] - emb = c.params["embedding_endpoint"] or next((n for n in names if EMBED_RE.search(n)), None) - chat = c.params["chat_endpoint"] or next((n for n in names if CHAT_RE.search(n) and not EMBED_RE.search(n)), None) - if not emb or not chat: - raise PackError("no embedding/chat serving endpoint found in this workspace. " - f"Pass --param embedding_endpoint=... --param chat_endpoint=... (available: {', '.join(names[:8]) or 'none'})") - others = [p.get("name", "") for p in _items(d().json("postgres", "list-projects"), "projects")] - if others and f"projects/{project}" not in others: - raise PackError(f"this workspace already has a Lakebase project ({others[0].split('/')[-1]}) and Free Edition allows one. " - f"Reuse it with --param project={others[0].split('/')[-1]} or delete it first") - return {"user": user, "embedding_endpoint": emb, "chat_endpoint": chat, "project": project} - - def project_listed() -> bool: - # list membership, not get-project: a deleted project can still answer get-project for a while - return f"projects/{project}" in [p.get("name") for p in _items(d().json("postgres", "list-projects"), "projects")] - - def mk_project(c: PackContext): - try: - d().json("postgres", "create-project", project) - except PackError as e: - if "soft-deleted" in str(e): - raise PackError(f"the name '{project}' is held by a soft-deleted project. Use another name " - f"(--param project=) or purge it: databricks api delete " - f"'/api/2.0/postgres/projects/{project}?purge=true'") from e - raise - return {"project_created": True} - - def locate(c: PackContext): - branches = _items(d().json("postgres", "list-branches", f"projects/{project}"), "branches") - if not branches: - raise PackError("project has no branch yet; wait a minute and re-run") - br = next((b for b in branches if pgmod.find(b, "default", "is_default") is True), branches[0]) - eps = _items(d().json("postgres", "list-endpoints", br["name"]), "endpoints") - host = pgmod.find(eps[0], "host") if eps else None - if not host: - raise PackError("compute endpoint not ready yet; re-run in a minute") - return {"lakebase_endpoint": eps[0]["name"], "pg_host": host} - - def build(c: PackContext): - chunks = self.source.chunks(c) - f = self._chunks_file(c) - f.parent.mkdir(parents=True, exist_ok=True) - f.write_text("\n".join(json.dumps(x) for x in chunks)) - f.chmod(0o600) - return {"chunks": len(chunks), "corpus": self.source.fingerprint(c)} - - def sch(c: PackContext): - dim = len(self._embed(c, ["dimension probe"])[0]) - with self._conn(c) as cx: - cx.execute("CREATE EXTENSION IF NOT EXISTS vector") - cx.execute(f"CREATE SCHEMA IF NOT EXISTS {schema}") - cx.execute(f"CREATE TABLE IF NOT EXISTS {schema}.chunks (id text PRIMARY KEY, title text, content text, " - f"meta jsonb, embedding vector({int(dim)}), " - f"tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', content)) STORED)") - return {"dim": dim} - - def load(c: PackContext): - chunks = self._read_chunks(c) - if c.state.get("loaded_corpus") != self.source.fingerprint(c): - with self._conn(c) as cx: # different data than what is in the table: replace it, do not mix - cx.execute(f"TRUNCATE {schema}.chunks") - for i in range(0, len(chunks), BATCH): - batch = chunks[i:i + BATCH] - vecs = self._embed(c, [f"{x['title']}\n{x['content']}" for x in batch]) - with self._conn(c) as cx: - for x, v in zip(batch, vecs): - cx.execute(f"INSERT INTO {schema}.chunks (id, title, content, meta, embedding) VALUES (%s,%s,%s,%s::jsonb,%s::vector) " - "ON CONFLICT (id) DO UPDATE SET title=EXCLUDED.title, content=EXCLUDED.content, " - "meta=EXCLUDED.meta, embedding=EXCLUDED.embedding", - (x["id"], x["title"], x["content"], json.dumps(x["meta"]), _retrieval().vec_literal(v))) - c.log(f" embedded {min(i + BATCH, len(chunks))}/{len(chunks)}") - return {"loaded_corpus": self.source.fingerprint(c)} - - def loaded(c: PackContext) -> bool: - if c.state.get("loaded_corpus") != self.source.fingerprint(c): - return False - try: - with self._conn(c) as cx: - return cx.execute(f"SELECT count(*) FROM {schema}.chunks").fetchone()[0] == c.state.get("chunks") - except Exception: - return False - - def index(c: PackContext): - with self._conn(c) as cx: - cx.execute(f"CREATE INDEX IF NOT EXISTS chunks_embedding_hnsw ON {schema}.chunks USING hnsw (embedding vector_cosine_ops)") - cx.execute(f"CREATE INDEX IF NOT EXISTS chunks_tsv ON {schema}.chunks USING gin (tsv)") - cx.execute(f"CREATE INDEX IF NOT EXISTS chunks_meta ON {schema}.chunks USING gin (meta jsonb_path_ops)") - - def evaluate(c: PackContext): - metas = [x["meta"] for x in self._read_chunks(c)] - qs = [q for q in self.source.eval_set() - if any(all(str(m.get(k)) == str(v) for k, v in q["expect"].items()) for m in metas)] - if not qs: - c.log(" no evaluation questions apply to this corpus: gate skipped") - return {"eval_questions": 0, "eval_hit_rate": None} - vecs = [v for i in range(0, len(qs), BATCH) for v in self._embed(c, [q["question"] for q in qs[i:i + BATCH]])] - r = _retrieval() - hits = 0 - with self._conn(c) as cx: - for q, v in zip(qs, vecs): - res = r.search(cx, schema, v, q["question"], int(c.params["top_k"])) - hits += 1 if r.is_hit(res, q["expect"]) else 0 - rate = round(hits / len(qs), 3) - c.log(f" retrieval hit rate {rate} on {len(qs)} questions") - floor = float(c.params["min_hit_rate"]) - if rate < floor: - raise PackError(f"retrieval hit rate {rate} is below min_hit_rate {floor} on {len(qs)} questions: " - "not deploying. Inspect chunking/sections, or lower --param min_hit_rate") - return {"eval_questions": len(qs), "eval_hit_rate": rate} - - def app(c: PackContext): - name = f"rag-{c.name}" - res = [{"name": "embedder", "serving_endpoint": {"name": c.state["embedding_endpoint"], "permission": "CAN_QUERY"}}, - {"name": "chat", "serving_endpoint": {"name": c.state["chat_endpoint"], "permission": "CAN_QUERY"}}] - if not d().exists("apps", "get", name): - d().json("apps", "create", body={"name": name, "description": f"labagent RAG ({self.source.id}) on Lakebase", "resources": res}) - sp = pgmod.find(d().json("apps", "get", name), "service_principal_client_id") - if not sp: - raise PackError("app service principal not reported yet; re-run") - branch = c.state["lakebase_endpoint"].rsplit("/endpoints/", 1)[0] - try: - d().json("postgres", "create-role", branch, "--role-id", re.sub(r"[^a-z0-9-]", "-", name.lower()), - body={"spec": {"postgres_role": sp, "identity_type": "SERVICE_PRINCIPAL", "auth_method": "LAKEBASE_OAUTH_V1"}}) - except PackError as e: - if "already exists" not in str(e).lower(): - raise - with self._conn(c) as cx: - cx.execute(f'GRANT USAGE ON SCHEMA {schema} TO "{sp}"') - cx.execute(f'GRANT SELECT ON {schema}.chunks TO "{sp}"') - with tempfile.TemporaryDirectory() as tmp: - bundle = Path(tmp) / "app" - shutil.copytree(HERE / "app", bundle) - y = (bundle / "app.yaml").read_text() - for k, v in {"PGHOST": c.state["pg_host"], "LAKEBASE_ENDPOINT": c.state["lakebase_endpoint"], - "EMBED_ENDPOINT": c.state["embedding_endpoint"], "CHAT_ENDPOINT": c.state["chat_endpoint"], - "RAG_SCHEMA": schema, "TOP_K": str(int(c.params["top_k"]))}.items(): - y = y.replace(f"__{k}__", v) - (bundle / "app.yaml").write_text(y) - (bundle / "ui.json").write_text(json.dumps(self.source.ui([x["meta"] for x in self._read_chunks(c)]), indent=1)) - path = f"/Users/{c.state['user']}/labagent/{name}" - d().text("workspace", "import-dir", str(bundle), path, "--overwrite") - d().text("apps", "deploy", name, "--source-code-path", f"/Workspace{path}") - return {"app_name": name, "app_path": path, "app_url": pgmod.find(d().json("apps", "get", name), "url")} - - return [ - Step("preflight", "Check login, serving endpoints and Free Edition project limit", preflight), - Step("project", "Create the Lakebase project", mk_project, check=lambda c: project_listed()), - Step("locate", "Find the compute endpoint", locate), - Step("build", f"Fetch, parse and chunk the '{self.source.id}' corpus", build, check=lambda c: self._chunks_file(c).exists() and c.state.get("corpus") == self.source.fingerprint(c)), - Step("schema", "Enable pgvector and create the table", sch), - Step("load", "Embed the chunks and load them", load, check=loaded), - Step("index", "Build vector, full-text and metadata indexes", index), - Step("evaluate", "Score retrieval on the evaluation questions (deploy gate)", evaluate), - Step("app", "Deploy the Databricks App", app), - ] - - # -- destroy ------------------------------------------------------------- - def destroy_steps(self, ctx: PackContext) -> List[Step]: - d = lambda: ctx.services["dbx"] # noqa: E731 - name = f"rag-{ctx.name}" - project = self._project(ctx) - - def rm_schema(c: PackContext): - if c.state.get("project_created"): - return # the project delete below removes it - try: - with self._conn(c) as cx: - cx.execute(f"DROP SCHEMA IF EXISTS {self.source.schema} CASCADE") - except Exception as e: - c.log(f"could not drop schema ({type(e).__name__}); remove it manually if the project still exists") - - def rm_project(c: PackContext): - if not c.state.get("project_created"): - c.log("project was not created by labagent: leaving it in place") - return - d().json("postgres", "delete-project", f"projects/{project}") - if str(c.params["keep_deleted"]).lower() not in ("true", "1", "yes"): - # Lakebase soft-deletes: purge so the name can be reused and no copy of the data lingers - d().api("delete", f"/api/2.0/postgres/projects/{project}?purge=true") - - def rm_app(c: PackContext): - try: - d().text("apps", "delete", name) - except PackError as e: - if "DELETING" not in str(e): # a delete already in flight is success - raise - - def rm_code(c: PackContext): - d().text("workspace", "delete", c.state["app_path"], "--recursive") - c.state.pop("app_path", None) - return {} - - return [ - Step("app", "Delete the Databricks App", rm_app, check=lambda c: not d().exists("apps", "get", name)), - Step("code", "Delete uploaded app code", rm_code, check=lambda c: "app_path" not in c.state), - Step("schema", "Drop this pack's schema (only when the project is reused)", rm_schema), - Step("project", "Delete the Lakebase project (only if labagent created it)", rm_project), - ] - - -class RagLakebase(RagPack): - id = "rag-lakebase" - aliases = ("rag-demo",) - default_name = "demo" - title = "RAG demo on Lakebase (pgvector)" - summary = "Demo corpus (25 invented passages) to smoke-test the RAG pipeline and Databricks setup." - source_id = "builtin" - - -class RagSecFilings(RagPack): - id = "rag-sec-filings" - aliases = ("rag-sec", "sec-filings") - default_name = "sec" - title = "Financial research RAG over SEC 10-K filings" - summary = "RAG over a pinned snapshot of real 10-Ks: hybrid search, metadata filters, citations, eval gate." - source_id = "sec-10k-snapshot" - - -class RagSecFilingsLive(RagPack): - id = "rag-sec-filings-live" - aliases = ("rag-sec-live",) - default_name = "sec-live" - title = "Financial research RAG, fresh from SEC EDGAR" - summary = "Same pipeline, but downloads the filings from EDGAR at launch (needs a contact for the SEC)." - source_id = "sec-10k" diff --git a/labagent/packs/rag/sources/__init__.py b/labagent/packs/rag/sources/__init__.py index 677a3a2..49f45f9 100644 --- a/labagent/packs/rag/sources/__init__.py +++ b/labagent/packs/rag/sources/__init__.py @@ -1,5 +1,5 @@ -"""Data sources for RAG packs. A source turns an upstream corpus into chunks (id, title, content, meta) -plus an evaluation set. Add one by subclassing Source; register built-ins below or via the +"""Data sources for the RAG tutorial and for building dataset snapshots. A source turns an upstream corpus into chunks +(id, title, content, meta) plus an evaluation set. Add one by subclassing Source; register built-ins below or via the `labagent.rag_sources` entry point.""" from __future__ import annotations @@ -13,11 +13,6 @@ class Source: params: List[Param] = [] schema: str = "" # Postgres schema for this source's table: keeps sources side by side in one project - def fingerprint(self, ctx: PackContext) -> str: - """Identity of the corpus this run will produce. When it changes, the build and load stages redo their work - (and the table is replaced) instead of trusting a cache or a matching row count from different data.""" - return self.id - def chunks(self, ctx: PackContext) -> List[Dict[str, Any]]: raise NotImplementedError @@ -32,9 +27,8 @@ def eval_set(self) -> List[Dict[str, Any]]: def _builtin() -> Dict[str, Source]: - from .builtin import Builtin from .sec import SecFilings, SecSnapshot - return {s.id: s for s in (Builtin(), SecFilings(), SecSnapshot())} + return {s.id: s for s in (SecFilings(), SecSnapshot())} def get(source_id: str) -> Source: diff --git a/labagent/packs/rag/sources/builtin.py b/labagent/packs/rag/sources/builtin.py deleted file mode 100644 index 8d0ce81..0000000 --- a/labagent/packs/rag/sources/builtin.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Demo source: 25 invented support-policy passages for a fictional retailer. For smoke tests, not for teaching.""" -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any, Dict, List - -from ...base import PackContext -from . import Source - -HERE = Path(__file__).parent - -EVAL = [ - ("How long do refunds take?", "returns"), ("Can I exchange a size?", "returns"), - ("Is shipping free?", "shipping"), ("How fast is express delivery?", "shipping"), - ("My package says delivered but is missing", "shipping"), ("What does the warranty cover?", "warranty"), - ("How do I file a warranty claim?", "warranty"), ("How many points do I earn per dollar?", "loyalty"), - ("When do loyalty points expire?", "loyalty"), ("I forgot my password", "account"), - ("Why did my card payment fail?", "payments"), ("Can I get a copy of my personal data?", "privacy"), -] - - -class Builtin(Source): - id = "builtin" - schema = "rag_builtin" - - def chunks(self, ctx: PackContext) -> List[Dict[str, Any]]: - rows = [json.loads(l) for l in (HERE / "builtin_corpus.jsonl").read_text().splitlines() if l.strip()] - return [{"id": f"builtin-{r['id']}", "title": r["title"], "content": r["content"], "meta": {"source": r["source"]}} - for r in rows] - - def eval_set(self) -> List[Dict[str, Any]]: - return [{"question": q, "expect": {"source": s}} for q, s in EVAL] - - def ui(self, metas): - cats = sorted({m["source"] for m in metas}) - return {"title": "Retail Support Assistant", "tagline": "Demo knowledge base (invented policies)", - "heading": "Ask the support knowledge base", "sub": "Fictional retailer policies, for testing the pipeline.", - "disclaimer": "Demo data: policies are invented.", - "facets": [{"key": "source", "label": "Topic", "kind": "chips"}], - "suggestions": [{"question": q, "label": s.title(), "filters": {"source": s}} for q, s in EVAL if s in cats]} diff --git a/labagent/packs/rag/sources/builtin_corpus.jsonl b/labagent/packs/rag/sources/builtin_corpus.jsonl deleted file mode 100644 index 8d0d1f1..0000000 --- a/labagent/packs/rag/sources/builtin_corpus.jsonl +++ /dev/null @@ -1,25 +0,0 @@ -{"id": 1, "source": "returns", "title": "Return window", "content": "Items can be returned within 30 days of delivery for a full refund if unused and in original packaging. Sale items follow the same window. Return shipping is free for defective items and costs 6.95 dollars otherwise."} -{"id": 2, "source": "returns", "title": "Refund timing", "content": "Refunds are issued to the original payment method within 5 business days of the warehouse receiving the return. Banks may take another 3 to 5 days to post the credit."} -{"id": 3, "source": "returns", "title": "Exchanges", "content": "Exchanges for a different size or colour are free. Start one from the order page and ship the original item back. The replacement ships once the carrier scans the return label."} -{"id": 4, "source": "returns", "title": "Non-returnable items", "content": "Gift cards, downloadable products, personalised items and opened hygiene products cannot be returned. Damaged or wrong items are always covered, whatever the category."} -{"id": 5, "source": "shipping", "title": "Standard shipping", "content": "Standard shipping takes 3 to 7 business days and is free on orders over 50 dollars. Under that threshold it costs 4.95 dollars."} -{"id": 6, "source": "shipping", "title": "Express shipping", "content": "Express delivery takes 1 to 2 business days and costs 14.95 dollars. Orders placed before 2 pm local time ship the same day. Express is not available to PO boxes."} -{"id": 7, "source": "shipping", "title": "International shipping", "content": "We ship to 38 countries. Delivery takes 7 to 14 business days. Duties and taxes are collected at checkout so there are no surprise fees on arrival."} -{"id": 8, "source": "shipping", "title": "Tracking", "content": "A tracking link is emailed when the order leaves the warehouse. Tracking can take up to 24 hours to show movement. If it stalls for 5 days, contact support and we open a carrier trace."} -{"id": 9, "source": "shipping", "title": "Lost packages", "content": "If a package shows delivered but you did not receive it, wait 48 hours and check with neighbours. Then contact support: we reship or refund after a 3 day carrier investigation."} -{"id": 10, "source": "warranty", "title": "Standard warranty", "content": "All products carry a 2 year warranty against manufacturing defects. Normal wear, accidents and misuse are not covered. Keep your order confirmation as proof of purchase."} -{"id": 11, "source": "warranty", "title": "Making a claim", "content": "Submit a claim with the order number and two photos of the defect. Claims are reviewed in 2 business days. Approved claims are repaired, replaced or refunded at our choice."} -{"id": 12, "source": "warranty", "title": "Extended protection", "content": "Extended protection adds 3 years and covers accidental damage on electronics over 100 dollars. It must be bought within 30 days of the original purchase and is non-refundable after 14 days."} -{"id": 13, "source": "loyalty", "title": "Earning points", "content": "Members earn 1 point per dollar spent. Double points apply during the first week of each month. Points are credited when the return window closes."} -{"id": 14, "source": "loyalty", "title": "Redeeming points", "content": "100 points equal 5 dollars off. Points can cover up to half of an order and cannot be used on gift cards. Points expire after 18 months of account inactivity."} -{"id": 15, "source": "loyalty", "title": "Membership tiers", "content": "Silver starts at 500 points a year, Gold at 1500 and Platinum at 4000. Gold adds free express shipping on one order a month. Platinum adds a dedicated support line."} -{"id": 16, "source": "account", "title": "Resetting a password", "content": "Use the Forgot password link on the sign in page. The reset email expires after 30 minutes. If it does not arrive, check spam, then request a new one after 5 minutes."} -{"id": 17, "source": "account", "title": "Changing an email address", "content": "Change your email under Account settings. We send a confirmation to the new address and the change takes effect once you click it. Order history stays attached to the account."} -{"id": 18, "source": "account", "title": "Deleting an account", "content": "Account deletion removes personal data within 30 days. Invoices are retained for 7 years because tax law requires it. Open orders must be delivered or cancelled first."} -{"id": 19, "source": "payments", "title": "Accepted payments", "content": "We accept major credit cards, PayPal and instalment plans from Klarna. Gift cards can be combined with one other payment method per order."} -{"id": 20, "source": "payments", "title": "Failed payments", "content": "A payment fails most often because of address mismatch or a bank fraud hold. Try again with the billing address exactly as the bank has it, or use another method."} -{"id": 21, "source": "payments", "title": "Instalments", "content": "Instalment plans split orders over 100 dollars into 4 interest free payments, charged every 2 weeks. A late payment fee of 7 dollars applies after a 5 day grace period."} -{"id": 22, "source": "privacy", "title": "Data we collect", "content": "We collect order, delivery and device data to fulfil orders and prevent fraud. We do not sell personal data. Marketing email is opt in and can be turned off in one click."} -{"id": 23, "source": "privacy", "title": "Data requests", "content": "You can request a copy of your data or ask for correction at any time. Requests are answered within 30 days and verified against your account email."} -{"id": 24, "source": "support", "title": "Contacting support", "content": "Chat is open 8 am to 10 pm on weekdays and 9 am to 6 pm on weekends. Email replies arrive within 24 hours. Platinum members have a phone line with a 2 minute target wait."} -{"id": 25, "source": "support", "title": "Order cancellation", "content": "Orders can be cancelled until they enter picking, usually within 1 hour of purchase. After that, refuse delivery or use the return process instead."} diff --git a/labagent/packs/rag/sources/sec.py b/labagent/packs/rag/sources/sec.py index 0ea2386..f4bbc55 100644 --- a/labagent/packs/rag/sources/sec.py +++ b/labagent/packs/rag/sources/sec.py @@ -138,9 +138,6 @@ def _json(self, ctx: PackContext, url: str) -> Any: return json.loads(self._get(ctx, url)) # -- source -------------------------------------------------------------- - def fingerprint(self, ctx: PackContext) -> str: - return f"{self.id}:{ctx.params['tickers']}:{ctx.params['filings']}:{ctx.params['sections']}" - def chunks(self, ctx: PackContext) -> List[Dict[str, Any]]: ua = str(ctx.params["user_agent"] or "") if "@" not in ua: @@ -224,11 +221,6 @@ class SecSnapshot(SecFilings): id = "sec-10k-snapshot" params = [Param("dataset", "sec-10k-2026.10", "Pinned snapshot (see labagent/packs/rag/sources/datasets.json)")] - def fingerprint(self, ctx: PackContext) -> str: - from ....datasets import registry - entry = registry().get(str(ctx.params["dataset"]), {}) - return f"{self.id}:{ctx.params['dataset']}:{entry.get('sha256', '')[:12]}" - def chunks(self, ctx: PackContext) -> List[Dict[str, Any]]: from ....datasets import DatasetError, default_http, fetch_snapshot try: diff --git a/labagent/packs/rag/tutorial/__init__.py b/labagent/packs/rag/tutorial/__init__.py new file mode 100644 index 0000000..44c86f0 --- /dev/null +++ b/labagent/packs/rag/tutorial/__init__.py @@ -0,0 +1 @@ +"""Notebooks and helpers for the step-by-step Databricks RAG tutorial (see labagent.packs.rag.tutorial_pack).""" diff --git a/labagent/packs/rag/tutorial/notebooks/00_setup.py b/labagent/packs/rag/tutorial/notebooks/00_setup.py new file mode 100644 index 0000000..434adcd --- /dev/null +++ b/labagent/packs/rag/tutorial/notebooks/00_setup.py @@ -0,0 +1,109 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC # 00. Set up: the pieces and the database +# MAGIC +# MAGIC **Time:** about 5 minutes. **Cost:** none on Databricks Free Edition. +# MAGIC +# MAGIC You will build a research assistant that answers questions from real SEC annual reports (Form 10-K) and cites its sources. This tutorial builds it one step at a time, so you see what each part does and why it is there. +# MAGIC +# MAGIC ``` +# MAGIC question --> [embedding model] --> vector --> [Lakebase: Postgres + pgvector] --> best passages +# MAGIC | +# MAGIC answer with citations <-- [chat model] <---------------- question + passages <------+ +# MAGIC ``` +# MAGIC +# MAGIC | Notebook | You will | +# MAGIC |---|---| +# MAGIC | 00 Set up | Check your workspace and create the database | +# MAGIC | 01 Data | Load and explore the SEC filings | +# MAGIC | 02 Embeddings | Turn text into vectors | +# MAGIC | 03 pgvector | Store vectors in Lakebase and search them | +# MAGIC | 04 Search | Add indexes, keyword search, hybrid search and filters | +# MAGIC | 05 Answer and evaluate | Generate cited answers and score retrieval | +# MAGIC | 06 Deploy the app | Publish the assistant as a Databricks App | +# MAGIC | 99 Clean up | Remove everything you created | +# MAGIC +# MAGIC Run the notebooks in order. Run each cell with **Shift+Enter**. Every notebook ends with a check cell that says whether the step worked. + +# COMMAND ---------- + +# MAGIC %run ./_common + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 1. Check your workspace +# MAGIC +# MAGIC Two kinds of model are needed. An **embedding model** turns text into numbers. A **chat model** writes the answer. Databricks hosts both as pay-per-token endpoints. + +# COMMAND ---------- + +print("signed in as: ", ME) +print("embedding endpoint:", EMBED_ENDPOINT) +print("chat endpoint: ", CHAT_ENDPOINT) +check(EMBED_ENDPOINT and CHAT_ENDPOINT, "an embedding model and a chat model are available in this workspace") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 2. Create the database +# MAGIC +# MAGIC **Lakebase** is managed Postgres inside Databricks. The **pgvector** extension adds a vector column type and similarity search, so one database can hold your text, its metadata and its vectors. +# MAGIC +# MAGIC Free Edition allows **one** Lakebase project. The cell below creates a project called `rag-tutorial`. It does nothing if the project already exists, and it stops with a message if you already have a different project. + +# COMMAND ---------- + +projects = [p["name"] for p in w.api_client.do("GET", "/api/2.0/postgres/projects").get("projects", [])] +print("existing projects:", projects) + +if f"projects/{PROJECT}" in projects: + print(f"project '{PROJECT}' already exists: nothing to create") +elif projects: + raise RuntimeError(f"this workspace already has the Lakebase project {projects[0]!r} and Free Edition allows one. " + f"Delete it in the Lakebase page, or set PROJECT in the _common notebook to its name.") +else: + w.api_client.do("POST", "/api/2.0/postgres/projects", query={"project_id": PROJECT}, body={}) + print(f"creating project '{PROJECT}' ...") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 3. Wait for the database to start +# MAGIC +# MAGIC A new project takes a minute or two to get its compute endpoint. + +# COMMAND ---------- + +host = None +for _ in range(40): + try: + host = lakebase_host() + except Exception: + host = None + if host: + break + print(" still starting ...") + time.sleep(8) +print("database host:", host) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Check + +# COMMAND ---------- + +check(host, "the Lakebase endpoint has a hostname") +token = w.api_client.do("POST", "/api/2.0/postgres/credentials", body={"endpoint": LAKEBASE_ENDPOINT}).get("token") +check(token, "you can get a short-lived database credential") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## What to notice +# MAGIC +# MAGIC - There is **no password to manage**. You ask Databricks for a token that lasts about an hour, and use your own identity as the database user. +# MAGIC - The database is empty. Everything in the next notebooks creates something inside it, and notebook 99 removes it. +# MAGIC +# MAGIC **Next:** `01_data` diff --git a/labagent/packs/rag/tutorial/notebooks/01_data.py b/labagent/packs/rag/tutorial/notebooks/01_data.py new file mode 100644 index 0000000..dde11f3 --- /dev/null +++ b/labagent/packs/rag/tutorial/notebooks/01_data.py @@ -0,0 +1,76 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC # 01. The data: what are we searching? +# MAGIC +# MAGIC **Time:** about 5 minutes. +# MAGIC +# MAGIC RAG is only as good as the passages it can find. Before any embedding, look at the data. +# MAGIC +# MAGIC The filings were already split into **chunks**, short passages that stay inside one section of a 10-K. Each chunk carries **metadata**: company, year, section and the URL of the original on sec.gov. You will use the metadata later for filters and citations. + +# COMMAND ---------- + +# MAGIC %run ./_common + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 1. Download the dataset +# MAGIC +# MAGIC The download is a versioned release on GitHub. The code checks its SHA-256 against a pinned value, so you and every other student work on identical data. + +# COMMAND ---------- + +rows = load_rows() +print(len(rows), "chunks loaded") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 2. Look at one chunk + +# COMMAND ---------- + +one = rows[0] +print("id: ", one["id"]) +print("title: ", one["title"]) +print("meta: ", json.dumps(one["meta"], indent=2)) +print("content:", one["content"][:500], "...") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 3. What is in the dataset? + +# COMMAND ---------- + +from collections import Counter + +print("by company: ", dict(Counter(r["meta"]["ticker"] for r in rows))) +print("by section: ", dict(Counter(r["meta"]["section"] for r in rows))) +print("by year: ", sorted(Counter((r["meta"]["ticker"], r["meta"]["fiscal_year"]) for r in rows).items())) +lengths = sorted(len(r["content"]) for r in rows) +print(f"chunk length: shortest {lengths[0]}, median {lengths[len(lengths) // 2]}, longest {lengths[-1]} characters") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Check + +# COMMAND ---------- + +check(len(rows) == 1880, "all 1,880 chunks loaded") +check({r["meta"]["ticker"] for r in rows} == {"AAPL", "CVX", "JPM", "PFE", "WMT"}, "five companies: Apple, Chevron, JPMorgan, Pfizer, Walmart") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## What to notice +# MAGIC +# MAGIC - Chunks are **about 1,000 characters**. Too small and a passage loses its context. Too large and one vector blurs several topics. +# MAGIC - Each company has **two fiscal years**, so the same question can match two near-identical passages. +# MAGIC - **Known gap:** the JPMorgan and Chevron *Management's Discussion and Analysis* sections are missing from this dataset. You will see this show up in the evaluation in notebook 05. +# MAGIC +# MAGIC **Your turn:** find the longest chunk and read it. Does it look like a good unit to search? +# MAGIC +# MAGIC **Next:** `02_embeddings` diff --git a/labagent/packs/rag/tutorial/notebooks/02_embeddings.py b/labagent/packs/rag/tutorial/notebooks/02_embeddings.py new file mode 100644 index 0000000..948d08f --- /dev/null +++ b/labagent/packs/rag/tutorial/notebooks/02_embeddings.py @@ -0,0 +1,94 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC # 02. Embeddings: turning meaning into numbers +# MAGIC +# MAGIC **Time:** about 10 minutes. +# MAGIC +# MAGIC An **embedding** is a list of numbers that captures the meaning of a text. Texts with similar meaning get vectors that point in similar directions. Searching means finding the stored vectors closest to the vector of your question. +# MAGIC +# MAGIC In this notebook you call the embedding model directly, so you see the raw request before any helper hides it. + +# COMMAND ---------- + +# MAGIC %run ./_common + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 1. One request to the embedding model +# MAGIC +# MAGIC The model is a serving endpoint. You send text, you get a vector back. + +# COMMAND ---------- + +sentences = ["Passwords must be at least 14 characters", "Use a long password", "The cafeteria opens at noon"] + +resp = w.api_client.do("POST", f"/serving-endpoints/{EMBED_ENDPOINT}/invocations", body={"input": sentences}) +vectors = [d["embedding"] for d in resp["data"]] + +print("endpoint: ", EMBED_ENDPOINT) +print("vectors returned: ", len(vectors)) +print("numbers per vector: ", len(vectors[0])) +print("first 6 numbers: ", [round(x, 4) for x in vectors[0][:6]]) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 2. Which sentences are close? +# MAGIC +# MAGIC **Cosine similarity** measures the angle between two vectors: 1.0 is the same direction, 0 is unrelated. + +# COMMAND ---------- + +def cosine(a, b): + dot = sum(x * y for x, y in zip(a, b)) + return dot / ((sum(x * x for x in a) ** 0.5) * (sum(y * y for y in b) ** 0.5)) + +print("password rule vs long password :", round(cosine(vectors[0], vectors[1]), 3)) +print("password rule vs cafeteria hours :", round(cosine(vectors[0], vectors[2]), 3)) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC The first pair shares **meaning** but almost no words, and it scores far higher than the unrelated pair. That is what keyword search cannot do. +# MAGIC +# MAGIC **Your turn:** add your own sentences to the list above and predict the scores before you run it. + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 3. Embed real chunks, politely +# MAGIC +# MAGIC The endpoint limits how fast you can call it. The `embed()` helper in `_common` sends **8 texts per request** and waits and retries if the endpoint says to slow down. Time a small batch. + +# COMMAND ---------- + +rows = load_rows() +sample = [f"{r['title']}\n{r['content']}" for r in rows[:16]] + +t = time.time() +sample_vectors = embed(sample) +print(f"embedded {len(sample_vectors)} chunks in {time.time() - t:.1f}s ({(time.time() - t) / len(sample):.2f}s each)") +print(f"embedding all {len(rows)} chunks will take roughly {len(rows) * (time.time() - t) / len(sample) / 60:.0f} minutes") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Check + +# COMMAND ---------- + +check(len(vectors[0]) == 1024, "each vector has 1,024 numbers (this model's dimension)") +check(cosine(vectors[0], vectors[1]) > cosine(vectors[0], vectors[2]), "similar meaning scores higher than unrelated text") +check(len(sample_vectors) == 16, "16 chunks embedded") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## What to notice +# MAGIC +# MAGIC - The vector has **1,024 numbers**. The database column in the next notebook must be declared with exactly that size. +# MAGIC - Embedding is **the slow step** of building a RAG system. You pay it once per chunk, and again whenever a chunk changes. +# MAGIC - The model sees at most about 512 tokens, so `embed()` cuts each text at 1,800 characters. +# MAGIC +# MAGIC **Next:** `03_pgvector` diff --git a/labagent/packs/rag/tutorial/notebooks/03_pgvector.py b/labagent/packs/rag/tutorial/notebooks/03_pgvector.py new file mode 100644 index 0000000..b7338a3 --- /dev/null +++ b/labagent/packs/rag/tutorial/notebooks/03_pgvector.py @@ -0,0 +1,193 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC # 03. pgvector: store vectors and search them +# MAGIC +# MAGIC **Time:** about 15 minutes (most of it is embedding the filings). +# MAGIC +# MAGIC You will connect to the Lakebase database, create a table with a vector column, and load the filings in **two stages on purpose**. First every company except Apple. You then ask about Apple and see what search does when the right document was never embedded. Then you load Apple and ask again. + +# COMMAND ---------- + +# MAGIC %pip install "psycopg[binary]>=3.1" --quiet + +# COMMAND ---------- + +dbutils.library.restartPython() + +# COMMAND ---------- + +# MAGIC %run ./_common + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 1. Connect to Lakebase +# MAGIC +# MAGIC There is no stored password. You ask Databricks for a token and use it as the Postgres password. The `pg()` helper in `_common` does exactly the two steps below. + +# COMMAND ---------- + +import psycopg + +token = w.api_client.do("POST", "/api/2.0/postgres/credentials", body={"endpoint": LAKEBASE_ENDPOINT})["token"] +conn = psycopg.connect(host=lakebase_host(), port=5432, dbname="databricks_postgres", user=ME, password=token, + sslmode="require", autocommit=True, connect_timeout=30) +print(conn.execute("select version()").fetchone()[0][:60]) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 2. Create the table +# MAGIC +# MAGIC `CREATE EXTENSION vector` turns on pgvector. The `embedding` column is declared `vector(1024)`, the size you saw in notebook 02. The `tsv` column is a generated full-text index of the content: you will use it for keyword search in notebook 04. + +# COMMAND ---------- + +DIM = len(embed(["dimension probe"])[0]) +print("embedding dimension:", DIM) + +with pg() as cx: + cx.execute("CREATE EXTENSION IF NOT EXISTS vector") + cx.execute(f"CREATE SCHEMA IF NOT EXISTS {SCHEMA}") + cx.execute(f"""CREATE TABLE IF NOT EXISTS {SCHEMA}.chunks ( + id text PRIMARY KEY, title text, content text, meta jsonb, + embedding vector({DIM}), + tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', content)) STORED)""") + print("rows in the table now:", cx.execute(f"SELECT count(*) FROM {SCHEMA}.chunks").fetchone()[0]) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 3. A loader: embed, then insert +# MAGIC +# MAGIC For each batch of chunks, embed their text and insert the chunk with its vector. `ON CONFLICT ... DO UPDATE` makes it safe to run twice: a chunk that is already there is replaced, not duplicated. + +# COMMAND ---------- + +def load_chunks(chunks): + started = time.time() + for i in range(0, len(chunks), BATCH): + batch = chunks[i:i + BATCH] + vectors = embed([f"{c['title']}\n{c['content']}" for c in batch]) + with pg() as cx: + for c, v in zip(batch, vectors): + cx.execute(f"INSERT INTO {SCHEMA}.chunks (id, title, content, meta, embedding) VALUES (%s, %s, %s, %s::jsonb, %s::vector) " + "ON CONFLICT (id) DO UPDATE SET title = EXCLUDED.title, content = EXCLUDED.content, " + "meta = EXCLUDED.meta, embedding = EXCLUDED.embedding", + (c["id"], c["title"], c["content"], json.dumps(c["meta"]), vec(v))) + if (i // BATCH) % 25 == 0: + print(f" {min(i + BATCH, len(chunks))}/{len(chunks)} chunks embedded and stored ({time.time() - started:.0f}s)") + print(f"done: {len(chunks)} chunks in {time.time() - started:.0f}s") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 4. Stage one: load everyone except Apple +# MAGIC +# MAGIC The Apple filings stay in Python memory. They are **not** embedded and **not** in the database. + +# COMMAND ---------- + +rows = load_rows() +first = [r for r in rows if r["meta"]["ticker"] != "AAPL"] +later = [r for r in rows if r["meta"]["ticker"] == "AAPL"] +print(f"loading {len(first)} chunks now; holding back {len(later)} Apple chunks") + +load_chunks(first) + +# COMMAND ---------- + +with pg() as cx: + print(cx.execute(f"SELECT meta->>'ticker', count(*) FROM {SCHEMA}.chunks GROUP BY 1 ORDER BY 1").fetchall()) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 5. Ask about Apple +# MAGIC +# MAGIC This is a plain **vector search**: embed the question, then order the table by distance to it. `<=>` is pgvector's cosine distance, so `1 - distance` is the similarity. `meta @> '{"ticker": "AAPL"}'` is a metadata filter. + +# COMMAND ---------- + +def vector_search(question, ticker=None, k=5): + qv = vec(embed([question])[0]) + where, args = ("WHERE meta @> %s::jsonb", [json.dumps({"ticker": ticker})]) if ticker else ("", []) + with pg() as cx: + return cx.execute(f"SELECT meta->>'ticker', meta->>'fiscal_year', meta->>'section', round((1 - (embedding <=> %s::vector))::numeric, 3) AS similarity " + f"FROM {SCHEMA}.chunks {where} ORDER BY embedding <=> %s::vector LIMIT {int(k)}", (qv, *args, qv)).fetchall() + +question = "What are Apple's main supply chain and manufacturing risks?" + +print("Search only Apple (filter):") +only_apple = vector_search(question, ticker="AAPL") +print(" rows found:", len(only_apple)) + +print("\nSearch everything:") +everything = vector_search(question) +for r in everything: + print(" ", r) +before = {"best": max(float(r[3]) for r in everything), "companies": sorted({r[0] for r in everything})} + +# COMMAND ---------- + +# MAGIC %md +# MAGIC **What to notice** +# MAGIC +# MAGIC 1. The Apple-only search finds **nothing**. There are no Apple vectors to search. +# MAGIC 2. The search over everything still returns rows. They are other companies' chunks, the closest of what exists, with lower similarity. **Search never says "I have nothing."** It returns the nearest thing it has. +# MAGIC +# MAGIC The Apple data is not wrong. It was never embedded. When a RAG system answers badly, check that the document is in the index before blaming the model. + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 6. Stage two: load Apple, ask again + +# COMMAND ---------- + +load_chunks(later) + +# COMMAND ---------- + +print("Search only Apple (filter):") +only_apple = vector_search(question, ticker="AAPL") +print(" rows found:", len(only_apple)) +for r in only_apple: + print(" ", r) + +print("\nSearch everything:") +everything = vector_search(question) +for r in everything: + print(" ", r) + +print("\nBefore and after embedding Apple") +print(f" best similarity {before['best']:.3f} -> {max(float(r[3]) for r in everything):.3f}") +print(f" companies found {before['companies']} -> {sorted({r[0] for r in everything})}") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Check + +# COMMAND ---------- + +with pg() as cx: + total = cx.execute(f"SELECT count(*) FROM {SCHEMA}.chunks").fetchone()[0] + with_vectors = cx.execute(f"SELECT count(*) FROM {SCHEMA}.chunks WHERE embedding IS NOT NULL").fetchone()[0] + +check(total == len(rows), f"all {len(rows)} chunks are in the table") +check(with_vectors == total, "every chunk has a vector") +check(len(only_apple) == 5, "the Apple-only search now returns results") +check({r[0] for r in everything} == {"AAPL"}, "searching everything now returns only Apple chunks for an Apple question") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## What to notice +# MAGIC +# MAGIC - The question, the data and the model are identical before and after. Only the index changed. +# MAGIC - One database holds the text, the metadata and the vectors, so a filter and a similarity search happen in a single SQL query. +# MAGIC +# MAGIC **Your turn:** run `vector_search("How does Walmart describe competition?")` and read the similarity scores. Then try a question that has nothing to do with the filings. +# MAGIC +# MAGIC **Next:** `04_search` diff --git a/labagent/packs/rag/tutorial/notebooks/04_search.py b/labagent/packs/rag/tutorial/notebooks/04_search.py new file mode 100644 index 0000000..55646a1 --- /dev/null +++ b/labagent/packs/rag/tutorial/notebooks/04_search.py @@ -0,0 +1,164 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC # 04. Search: indexes, keywords, hybrid and filters +# MAGIC +# MAGIC **Time:** about 15 minutes. +# MAGIC +# MAGIC Vector search finds passages with similar **meaning**. It is weak at exact wording, such as a product code or a phrase. Keyword search is the opposite. In this notebook you add indexes, run both kinds of search, merge them, and narrow results with filters. + +# COMMAND ---------- + +# MAGIC %pip install "psycopg[binary]>=3.1" --quiet + +# COMMAND ---------- + +dbutils.library.restartPython() + +# COMMAND ---------- + +# MAGIC %run ./_common + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 1. Add three indexes +# MAGIC +# MAGIC Without an index every search reads the whole table. With 1,880 rows that is fast anyway, but a real system has millions. +# MAGIC +# MAGIC | Index | Speeds up | Trade-off | +# MAGIC |---|---|---| +# MAGIC | **HNSW** on `embedding` | Nearest-vector search | Approximate: it may miss a true neighbour. Takes time and memory to build | +# MAGIC | **GIN** on `tsv` | Keyword search | Slower inserts | +# MAGIC | **GIN** on `meta` | Metadata filters (`meta @> ...`) | Slower inserts | + +# COMMAND ---------- + +with pg() as cx: + cx.execute(f"CREATE INDEX IF NOT EXISTS chunks_embedding_hnsw ON {SCHEMA}.chunks USING hnsw (embedding vector_cosine_ops)") + cx.execute(f"CREATE INDEX IF NOT EXISTS chunks_tsv ON {SCHEMA}.chunks USING gin (tsv)") + cx.execute(f"CREATE INDEX IF NOT EXISTS chunks_meta ON {SCHEMA}.chunks USING gin (meta jsonb_path_ops)") + for (name,) in cx.execute("SELECT indexname FROM pg_indexes WHERE schemaname = %s ORDER BY 1", (SCHEMA,)).fetchall(): + print(name) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 2. Two kinds of search, side by side +# MAGIC +# MAGIC The question is an exact phrase from Apple's filing. The meaning is generic, so vectors alone struggle. Run both and compare the first column. + +# COMMAND ---------- + +def show(rows): + for n, r in enumerate(rows, 1): + print(f" #{n} {r[0]:5} {r[1]} {r[2][:45]}") + +phrase = "52- or 53-week fiscal year" +qv = vec(embed([phrase])[0]) + +with pg() as cx: + print("VECTOR search (meaning):") + show(cx.execute(f"SELECT meta->>'ticker', meta->>'fiscal_year', meta->>'section' FROM {SCHEMA}.chunks ORDER BY embedding <=> %s::vector LIMIT 5", (qv,)).fetchall()) + + print("\nKEYWORD search (exact words):") + show(cx.execute(f"SELECT meta->>'ticker', meta->>'fiscal_year', meta->>'section' FROM {SCHEMA}.chunks, websearch_to_tsquery('english', %s) q " + "WHERE tsv @@ q ORDER BY ts_rank_cd(tsv, q) DESC LIMIT 5", (phrase,)).fetchall()) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC Vector search may rank another company first, because the phrase is generic. Keyword search finds the chunks that contain those exact words. Neither is enough alone. + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 3. Merge them: reciprocal rank fusion +# MAGIC +# MAGIC Run both searches, then score each passage by where it ranked in each list: `1 / (60 + rank)`. A passage near the top of **both** lists wins. The constant 60 stops one first place from dominating. + +# COMMAND ---------- + +# a worked example with three passages A, B, C +vector_order = ["A", "B", "C"] +keyword_order = ["C", "A"] +scores = {} +for order in (vector_order, keyword_order): + for rank, item in enumerate(order, 1): + scores[item] = scores.get(item, 0) + 1 / (60 + rank) +print({k: round(v, 5) for k, v in sorted(scores.items(), key=lambda kv: -kv[1])}) +print("A is first in the vector list and second in the keyword list, so it beats C, which is first only in the keyword list.") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC The helper `hybrid_search` in `_common` does this against your table. Use it on the same phrase. + +# COMMAND ---------- + +with pg() as cx: + hybrid = hybrid_search(cx, embed([phrase])[0], phrase, k=5) +for n, r in enumerate(hybrid, 1): + print(f" #{n} {r['meta']['ticker']:5} {r['meta']['fiscal_year']} {r['meta']['section'][:45]} score {r['score']}") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 4. Filters +# MAGIC +# MAGIC When you already know the company, a filter is exact: no other company can appear. The filter is part of the same SQL query as the similarity search. + +# COMMAND ---------- + +question = "What are the main supply chain risks?" +qvec = embed([question])[0] + +with pg() as cx: + print("no filter:") + for r in hybrid_search(cx, qvec, question, k=5): + print(" ", r["meta"]["ticker"], r["meta"]["fiscal_year"], r["meta"]["section"][:40]) + + print("\nonly Walmart:") + for r in hybrid_search(cx, qvec, question, k=5, filters={"ticker": "WMT"}): + print(" ", r["meta"]["ticker"], r["meta"]["fiscal_year"], r["meta"]["section"][:40]) + + print("\nonly Walmart, fiscal year 2025:") + for r in hybrid_search(cx, qvec, question, k=5, filters={"ticker": "WMT", "fiscal_year": "2025"}): + print(" ", r["meta"]["ticker"], r["meta"]["fiscal_year"], r["meta"]["section"][:40]) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 5. Optional: ask the database how it ran the search + +# COMMAND ---------- + +with pg() as cx: + plan = cx.execute(f"EXPLAIN SELECT id FROM {SCHEMA}.chunks ORDER BY embedding <=> %s::vector LIMIT 5", (qv,)).fetchall() +print("\n".join(p[0] for p in plan)) +print("\nOn a table this small Postgres may choose to read every row instead of using the HNSW index. That is normal: the index pays off at scale.") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Check + +# COMMAND ---------- + +with pg() as cx: + names = {r[0] for r in cx.execute("SELECT indexname FROM pg_indexes WHERE schemaname = %s", (SCHEMA,)).fetchall()} +check({"chunks_embedding_hnsw", "chunks_tsv", "chunks_meta"} <= names, "the three indexes exist") +check(any(r["meta"]["ticker"] == "AAPL" for r in hybrid[:3]), "hybrid search puts an Apple chunk in the top 3 for the Apple phrase") +check(all(r["meta"]["ticker"] == "WMT" for r in hybrid_search(pg(), qvec, question, k=5, filters={"ticker": "WMT"})), "a filter returns only the filtered company") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## What to notice +# MAGIC +# MAGIC - **Hybrid search** is the default choice when users ask for exact terms and for ideas. +# MAGIC - **Filters** are the cheapest way to improve precision when the question names a company, a year or a department. +# MAGIC - Every extra index makes writes slower. Add the ones your queries need. +# MAGIC +# MAGIC **Your turn:** try `iPhone 17 Pro` and `Permian Basin`. Does hybrid change the order compared with vectors alone? Sometimes it does not, and that is worth knowing. +# MAGIC +# MAGIC **Next:** `05_answer_and_eval` diff --git a/labagent/packs/rag/tutorial/notebooks/05_answer_and_eval.py b/labagent/packs/rag/tutorial/notebooks/05_answer_and_eval.py new file mode 100644 index 0000000..7bd628e --- /dev/null +++ b/labagent/packs/rag/tutorial/notebooks/05_answer_and_eval.py @@ -0,0 +1,165 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC # 05. Answer with citations, then measure +# MAGIC +# MAGIC **Time:** about 15 minutes. +# MAGIC +# MAGIC Retrieval finds passages. **Generation** gives them to a chat model and asks it to answer using only those passages and to cite them. Then you **evaluate**: write questions with known answers and measure how often retrieval finds the right place. Without a score you cannot tell whether a change helped. + +# COMMAND ---------- + +# MAGIC %pip install "psycopg[binary]>=3.1" --quiet + +# COMMAND ---------- + +dbutils.library.restartPython() + +# COMMAND ---------- + +# MAGIC %run ./_common + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 1. The answer function +# MAGIC +# MAGIC Four steps: embed the question, retrieve passages, number them, ask the chat model. The instruction in `SYSTEM` is what stops the model from answering from memory. + +# COMMAND ---------- + +SYSTEM = ("You answer questions using only the numbered context passages. Cite the passages you used like [1] or [2]. " + "If the context does not contain the answer, say you do not know. Do not use outside knowledge.") + + +def ask(question, filters=None, k=TOP_K): + qvec = embed([question])[0] + with pg() as cx: + hits = hybrid_search(cx, qvec, question, k, filters) + if not hits: + return {"answer": "I do not know: no matching passages.", "sources": []} + context = "\n\n".join(f"[{i}] {h['title']}\n{h['content']}" for i, h in enumerate(hits, 1)) + out = w.api_client.do("POST", f"/serving-endpoints/{CHAT_ENDPOINT}/invocations", body={ + "messages": [{"role": "system", "content": SYSTEM}, {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}], + "max_tokens": 1200}) # reasoning models spend tokens thinking first: 500 cut answers off in testing + choice = out["choices"][0] + return {"answer": choice["message"]["content"], "sources": hits, "cut_off": choice.get("finish_reason") == "length"} + + +def show_answer(result): + print(result["answer"]) + if result.get("cut_off"): + print("\n[the model ran out of tokens: raise max_tokens in ask()]") + print("\ncites passage numbers like [1]:", bool(re.search(r"\[\d+\]", result["answer"]))) + print("\nsources:") + for i, h in enumerate(result["sources"], 1): + print(f" [{i}] {h['title']} -> {h['meta'].get('url')}") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 2. Ask a question the filings can answer + +# COMMAND ---------- + +good = ask("What are Apple's main supply chain and manufacturing risks?", filters={"ticker": "AAPL", "fiscal_year": "2025"}) +show_answer(good) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC Each number like `[1]` maps to a passage and a link to the original filing, so a reader can check the answer. +# MAGIC +# MAGIC **Models do not always cite.** In testing, about 1 answer in 9 left out the numbers even though the instruction asked for them. Run the cell a few times. A production system checks for citations and retries or flags the answer. + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 3. Ask a question the filings cannot answer +# MAGIC +# MAGIC A good system says it does not know. This is the behaviour to test before you trust a RAG system. + +# COMMAND ---------- + +bad = ask("What is the recipe for Walmart's store-brand cookies?") +show_answer(bad) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC Models differ here. If yours declines, that is the good outcome. If it answers anyway, look at the sources: they are real passages about something else, and the model stretched them. + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 4. Evaluate retrieval +# MAGIC +# MAGIC Each case is a question and **where the answer should come from**: a company and a section (`1` Business, `1A` Risk Factors, `7` Management's Discussion). A case **hits** when any retrieved passage comes from that company and section. + +# COMMAND ---------- + +cases = [ + ("What are Apple's main supply chain and manufacturing risks?", {"ticker": "AAPL", "item": "1A"}), + ("How does Apple describe its products and services?", {"ticker": "AAPL", "item": "1"}), + ("What credit and interest rate risks does JPMorgan Chase face?", {"ticker": "JPM", "item": "1A"}), + ("How did JPMorgan Chase's net income and revenue change this year?", {"ticker": "JPM", "item": "7"}), + ("What patent expirations and generic competition risks does Pfizer disclose?", {"ticker": "PFE", "item": "1A"}), + ("What are Pfizer's main product lines and business segments?", {"ticker": "PFE", "item": "1"}), + ("How do oil and gas price changes affect Chevron's earnings?", {"ticker": "CVX", "item": "7"}), + ("What climate and regulatory risks does Chevron identify?", {"ticker": "CVX", "item": "1A"}), + ("How did Walmart's comparable sales and operating income perform?", {"ticker": "WMT", "item": "7"}), + ("What competition and supply risks does Walmart face?", {"ticker": "WMT", "item": "1A"}), +] + +rows = load_rows() +in_data = lambda expect: any(all(str(r["meta"].get(k)) == str(v) for k, v in expect.items()) for r in rows) +applicable = [c for c in cases if in_data(c[1])] +missing = [c for c in cases if not in_data(c[1])] +print(f"{len(cases)} questions; {len(missing)} ask for a section that is not in the dataset:") +for q, expect in missing: + print(" not in the data:", expect, "-", q[:60]) + +# COMMAND ---------- + +qvecs = embed([q for q, _ in applicable]) +hits = 0 +with pg() as cx: + for (q, expect), qv in zip(applicable, qvecs): + results = hybrid_search(cx, qv, q, k=TOP_K) + ok = is_hit(results, expect) + hits += ok + print(("HIT " if ok else "MISS ") + str(expect).ljust(36), q[:70]) +rate = hits / len(applicable) +print(f"\nretrieval hit rate: {rate:.2f} on {len(applicable)} questions that the data can answer") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 5. The quality gate +# MAGIC +# MAGIC A deployment can refuse to publish an app whose hit rate is below **0.6**. That rule is the point: a gate turns "it seems to work" into a number you can enforce. + +# COMMAND ---------- + +GATE = 0.6 +print(f"hit rate {rate:.2f} vs gate {GATE}: {'PASS, safe to deploy' if rate >= GATE else 'FAIL, do not deploy'}") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Check + +# COMMAND ---------- + +check(len(good["sources"]) > 0 and not good["cut_off"], "the answer was complete and came with its sources") +check(all(h["meta"].get("url", "").startswith("https://www.sec.gov/") for h in good["sources"]), "every source links to the original filing on sec.gov") +check(rate >= GATE, f"retrieval hit rate {rate:.2f} is at or above the {GATE} gate") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## What to notice +# MAGIC +# MAGIC - The two questions about JPMorgan and Chevron *Management's Discussion* were set aside because **the data does not contain those sections**. No retrieval setting can find what is not there. Evaluation tells you whether to fix **retrieval** or the **data**. +# MAGIC - Change `TOP_K` or the chunking and re-run section 4. Did the rate go up or down? That is how you improve a RAG system: change one thing, measure. +# MAGIC +# MAGIC **Next:** `06_deploy_app` diff --git a/labagent/packs/rag/tutorial/notebooks/06_deploy_app.py b/labagent/packs/rag/tutorial/notebooks/06_deploy_app.py new file mode 100644 index 0000000..e36aee2 --- /dev/null +++ b/labagent/packs/rag/tutorial/notebooks/06_deploy_app.py @@ -0,0 +1,196 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC # 06. Deploy the assistant as a Databricks App +# MAGIC +# MAGIC **Time:** about 10 minutes. This notebook is optional: notebooks 00 to 05 already built a working RAG system. +# MAGIC +# MAGIC A **Databricks App** is a small web service that Databricks hosts for you, behind your workspace login. The app code (a FastAPI service and a web page) was uploaded next to these notebooks in the `app` folder. You will not write it. You will do the four things a deployment script would automate: +# MAGIC +# MAGIC 1. Create the app and tell it which models it may call. +# MAGIC 2. Give the app's own identity a **read-only** database role. The app never uses your credentials. +# MAGIC 3. Fill in the app's settings (database host, model names). +# MAGIC 4. Deploy, and wait until it runs. + +# COMMAND ---------- + +# MAGIC %pip install "psycopg[binary]>=3.1" --quiet + +# COMMAND ---------- + +dbutils.library.restartPython() + +# COMMAND ---------- + +# MAGIC %run ./_common + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 0. Before you start +# MAGIC +# MAGIC The data must be loaded, and the app code must be where labagent uploaded it. + +# COMMAND ---------- + +import os +import shutil + +with pg() as cx: + chunks = cx.execute(f"SELECT count(*) FROM {SCHEMA}.chunks").fetchone()[0] +print("chunks in the database:", chunks) +print("app code folder:", f"{TUTORIAL_FOLDER}/app", "| exists:", os.path.isdir(f"{TUTORIAL_FOLDER}/app")) +check(chunks > 0, "the database has data (run notebook 03 if this fails)") +check(os.path.isdir(f"{TUTORIAL_FOLDER}/app"), "the app code was uploaded (launch the tutorial with labagent, or upload the app folder yourself)") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 1. Create the app +# MAGIC +# MAGIC The app is given **resources**: permission to query exactly two serving endpoints, the embedding model and the chat model. Nothing else. +# MAGIC +# MAGIC Free Edition allows three apps and stops each one after 24 hours. You can start it again later. + +# COMMAND ---------- + +from databricks.sdk.errors import NotFound + +def get_app(): + try: + return w.api_client.do("GET", f"/api/2.0/apps/{APP_NAME}") + except NotFound: + return None + +if get_app() is None: + w.api_client.do("POST", "/api/2.0/apps", body={ + "name": APP_NAME, "description": "RAG over SEC 10-K filings (labagent tutorial)", + "resources": [{"name": "embedder", "serving_endpoint": {"name": EMBED_ENDPOINT, "permission": "CAN_QUERY"}}, + {"name": "chat", "serving_endpoint": {"name": CHAT_ENDPOINT, "permission": "CAN_QUERY"}}]}) + print("app created") +else: + print("app already exists") + +service_principal = None +for _ in range(30): + service_principal = _find(get_app(), "service_principal_client_id") + if service_principal: + break + time.sleep(5) +print("the app's own identity (service principal):", service_principal) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 2. Give the app a read-only database role +# MAGIC +# MAGIC Two steps. Create a Postgres role for the app's identity (it signs in with an OAuth token, no password). Then grant it `SELECT` on the one table it needs. If the app is ever compromised it can read the filings and nothing else. + +# COMMAND ---------- + +role_id = re.sub(r"[^a-z0-9-]", "-", f"app-{APP_NAME}".lower()) +try: + w.api_client.do("POST", f"/api/2.0/postgres/{BRANCH}/roles", query={"role_id": role_id}, + body={"spec": {"postgres_role": service_principal, "identity_type": "SERVICE_PRINCIPAL", "auth_method": "LAKEBASE_OAUTH_V1"}}) + print("role created") +except Exception as e: + if "already exists" not in str(e).lower(): + raise + print("role already exists") + +with pg() as cx: + cx.execute(f'GRANT USAGE ON SCHEMA {SCHEMA} TO "{service_principal}"') + cx.execute(f'GRANT SELECT ON {SCHEMA}.chunks TO "{service_principal}"') +print("granted: USAGE on the schema and SELECT on the chunks table") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 3. Fill in the app's settings +# MAGIC +# MAGIC `app.yaml` tells Databricks how to start the app and which environment variables to set. The uploaded copy has placeholders. You copy the folder and replace them with the values you created in this tutorial. + +# COMMAND ---------- + +source, target = f"{TUTORIAL_FOLDER}/app", f"{TUTORIAL_FOLDER}/app_deploy" +if os.path.isdir(target): + shutil.rmtree(target) +shutil.copytree(source, target) + +settings = {"PGHOST": lakebase_host(), "LAKEBASE_ENDPOINT": LAKEBASE_ENDPOINT, "EMBED_ENDPOINT": EMBED_ENDPOINT, + "CHAT_ENDPOINT": CHAT_ENDPOINT, "RAG_SCHEMA": SCHEMA, "TOP_K": str(TOP_K)} +yaml_path = f"{target}/app.yaml" +text = open(yaml_path).read() +for key, value in settings.items(): + text = text.replace(f"__{key}__", value) +open(yaml_path, "w").write(text) + +print(text) +check("__" not in text, "every placeholder in app.yaml was replaced") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 4. Deploy +# MAGIC +# MAGIC A deployment snapshots the folder and starts it. A brand-new app first needs its compute to start, which takes a few minutes, so the next cell waits for it. + +# COMMAND ---------- + +def compute_state(): + return (get_app().get("compute_status") or {}).get("state") + +if compute_state() in ("STOPPED", "ERROR"): + w.api_client.do("POST", f"/api/2.0/apps/{APP_NAME}/start") + print("the app's compute was stopped: starting it ...") + +# A new app's compute needs a few minutes to start, and Databricks refuses a deployment until it is ACTIVE. +for _ in range(60): + if compute_state() == "ACTIVE": + break + print(" app compute:", compute_state()) + time.sleep(10) +check(compute_state() == "ACTIVE", "the app's compute is active, so it can be deployed") + +deployment = w.api_client.do("POST", f"/api/2.0/apps/{APP_NAME}/deployments", body={"source_code_path": target, "mode": "SNAPSHOT"}) +print("deployment id:", deployment.get("deployment_id")) + +# COMMAND ---------- + +state, url = None, None +for _ in range(60): + app = get_app() + state = (app.get("app_status") or {}).get("state") + dep = (app.get("active_deployment") or app.get("pending_deployment") or {}).get("status", {}).get("state") + url = app.get("url") + print(f" app: {state} | deployment: {dep}") + if state == "RUNNING" and dep == "SUCCEEDED": + break + if dep == "FAILED": + raise RuntimeError("the deployment failed: open the app in Compute > Apps and read its logs") + time.sleep(10) +print("\napp url:", url) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Check + +# COMMAND ---------- + +check(state == "RUNNING", "the app is running") +check(url and url.startswith("https://"), "the app has a URL") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Try it +# MAGIC +# MAGIC Open the URL above in a new browser tab and sign in with your Databricks account. Pick a suggested question, or ask your own. Try the company filter. Open one of the source links: it goes to the original filing on sec.gov. +# MAGIC +# MAGIC ## What to notice +# MAGIC +# MAGIC - The app uses the **same** search and the **same** prompt you wrote in notebooks 04 and 05, served behind a web page. +# MAGIC - It runs with its **own identity** and a read-only role, not with yours. +# MAGIC - On Free Edition the app stops after 24 hours. Start it again from **Compute > Apps**. +# MAGIC +# MAGIC **Next:** `99_cleanup` when you are finished. diff --git a/labagent/packs/rag/tutorial/notebooks/99_cleanup.py b/labagent/packs/rag/tutorial/notebooks/99_cleanup.py new file mode 100644 index 0000000..d2eeab2 --- /dev/null +++ b/labagent/packs/rag/tutorial/notebooks/99_cleanup.py @@ -0,0 +1,110 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC # 99. Clean up +# MAGIC +# MAGIC Run this when you have finished the tutorial. It removes what you created, in the reverse order. Each step is safe to run twice. +# MAGIC +# MAGIC Free Edition allows **one** Lakebase project and **three** apps, so cleaning up also frees those slots for your next project. + +# COMMAND ---------- + +# MAGIC %pip install "psycopg[binary]>=3.1" --quiet + +# COMMAND ---------- + +dbutils.library.restartPython() + +# COMMAND ---------- + +# MAGIC %run ./_common + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 1. Delete the app + +# COMMAND ---------- + +from databricks.sdk.errors import NotFound + +try: + w.api_client.do("DELETE", f"/api/2.0/apps/{APP_NAME}") + print("app deleted (it may show as DELETING for a minute)") +except NotFound: + print("no app to delete") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 2. Drop the tutorial schema +# MAGIC +# MAGIC This deletes the table, the vectors and the indexes. It leaves the database project itself. + +# COMMAND ---------- + +try: + with pg() as cx: + cx.execute(f"DROP SCHEMA IF EXISTS {SCHEMA} CASCADE") + print(f"schema '{SCHEMA}' dropped") +except Exception as e: + print("nothing to drop:", str(e)[:120]) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 3. Delete the database project +# MAGIC +# MAGIC **This cannot be undone.** Lakebase keeps a deleted project for a while, which blocks its name. The second call purges it so you can reuse the name. +# MAGIC +# MAGIC Run the next cell only if you want to remove the whole database. + +# COMMAND ---------- + +DELETE_PROJECT = True # set to False to keep the database project + +if DELETE_PROJECT: + for suffix in ("", "?purge=true"): + try: + w.api_client.do("DELETE", f"/api/2.0/postgres/projects/{PROJECT}{suffix}") + except NotFound: + pass + except Exception as e: + print("delete request:", str(e)[:140]) + time.sleep(5) + remaining = [p["name"] for p in w.api_client.do("GET", "/api/2.0/postgres/projects").get("projects", [])] + print("projects left:", remaining) +else: + print("kept the project") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Check + +# COMMAND ---------- + +app_gone = False +for _ in range(24): # deletion is asynchronous: give it up to two minutes + try: + app = w.api_client.do("GET", f"/api/2.0/apps/{APP_NAME}") + except NotFound: + app_gone = True + break + if "DELETING" in ((app.get("app_status") or {}).get("state"), (app.get("compute_status") or {}).get("state")): + app_gone = True + break + time.sleep(5) +check(app_gone, "the app is deleted or being deleted") +if DELETE_PROJECT: + check(f"projects/{PROJECT}" not in remaining, "the Lakebase project is gone") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Last step +# MAGIC +# MAGIC The notebooks themselves are still in your workspace folder. Delete the folder from the workspace sidebar, or from your terminal: +# MAGIC +# MAGIC ``` +# MAGIC labagent destroy rag-sec-tutorial --name +# MAGIC ``` diff --git a/labagent/packs/rag/tutorial/notebooks/_common.py b/labagent/packs/rag/tutorial/notebooks/_common.py new file mode 100644 index 0000000..bda7f9a --- /dev/null +++ b/labagent/packs/rag/tutorial/notebooks/_common.py @@ -0,0 +1,159 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC # Shared settings and helpers +# MAGIC +# MAGIC You do not run this notebook yourself. Every other notebook starts with `%run ./_common`, which runs this file first. +# MAGIC +# MAGIC It holds the names used across the tutorial, the dataset download (checked against a SHA-256), and a few small helpers. Read it once so nothing feels like magic, then ignore it. + +# COMMAND ---------- + +import gzip +import hashlib +import json +import re +import time +import urllib.request +from pathlib import Path + +from databricks.sdk import WorkspaceClient + +w = WorkspaceClient() +ME = w.current_user.me().user_name + +# COMMAND ---------- + +# Names used by every notebook. Change PROJECT only if that name is already taken in your workspace. +PROJECT = "rag-tutorial" +BRANCH = f"projects/{PROJECT}/branches/production" +LAKEBASE_ENDPOINT = f"{BRANCH}/endpoints/primary" +SCHEMA = "tutorial" +APP_NAME = "rag-tutorial" +TOP_K = 4 + +# The dataset. The download is verified against this SHA-256, so everyone works on identical data. +SNAPSHOT_URL = "__SNAPSHOT_URL__" +SNAPSHOT_SHA256 = "__SNAPSHOT_SHA256__" + +# Where labagent uploaded this tutorial (notebook 06 reads the app code from here). +TUTORIAL_FOLDER = "__WORKSPACE_FOLDER__" + +BATCH = 8 # texts per embedding request: small batches avoid the endpoint's rate limit +EMBED_CHARS = 1800 # the embedding model accepts about 512 tokens, so cut longer text + +# COMMAND ---------- + +# Model serving endpoints. Every Databricks workspace has pay-per-token foundation models. +_names = [e.name for e in w.serving_endpoints.list()] +_is_embedding = lambda n: re.search(r"gte|bge|embed", n, re.I) +EMBED_ENDPOINT = "databricks-gte-large-en" if "databricks-gte-large-en" in _names else next((n for n in _names if _is_embedding(n)), None) +CHAT_ENDPOINT = next((n for n in _names if re.search(r"llama|gpt|claude|gemma|qwen|instruct|chat", n, re.I) and not _is_embedding(n)), None) + + +def check(condition, message): + """Stop with a clear message when a step did not work. Each notebook ends with checks like this.""" + if not condition: + raise AssertionError(f"CHECK FAILED: {message}") + print(f"OK: {message}") + + +# COMMAND ---------- + +def load_rows(): + """The SEC 10-K chunks, downloaded once to local disk and verified. Returns a list of dicts: id, title, content, meta.""" + path = Path("/tmp/sec-10k.chunks.jsonl.gz") + sha = lambda p: hashlib.sha256(p.read_bytes()).hexdigest() + if not path.exists() or sha(path) != SNAPSHOT_SHA256: + req = urllib.request.Request(SNAPSHOT_URL, headers={"User-Agent": "labagent-tutorial"}) + with urllib.request.urlopen(req, timeout=120) as r: + path.write_bytes(r.read()) + if sha(path) != SNAPSHOT_SHA256: + raise RuntimeError("the downloaded dataset does not match its checksum: do not use it") + return [json.loads(line) for line in gzip.open(path, "rt")] + + +def embed(texts): + """Embedding vectors for a list of texts. Sends small batches and waits when the endpoint says to slow down.""" + out = [] + for i in range(0, len(texts), BATCH): + batch = [t[:EMBED_CHARS] for t in texts[i:i + BATCH]] + for attempt in range(8): + try: + resp = w.api_client.do("POST", f"/serving-endpoints/{EMBED_ENDPOINT}/invocations", body={"input": batch}) + out += [d["embedding"] for d in resp["data"]] + break + except Exception as e: + if not re.search(r"REQUEST_LIMIT_EXCEEDED|429|rate limit|too many", str(e), re.I) or attempt == 7: + raise + time.sleep(min(30, 2 ** attempt)) + return out + + +def vec(v): + """A vector as the text Postgres expects: [0.1,0.2,...]""" + return "[" + ",".join(repr(float(x)) for x in v) + "]" + + +# COMMAND ---------- + +def _find(obj, *keys): + """First value found under any of `keys`, searching nested dicts and lists.""" + if isinstance(obj, dict): + for k in keys: + if obj.get(k) not in (None, ""): + return obj[k] + for v in obj.values(): + r = _find(v, *keys) + if r not in (None, ""): + return r + elif isinstance(obj, list): + for v in obj: + r = _find(v, *keys) + if r not in (None, ""): + return r + return None + + +def lakebase_host(): + """The hostname of the Lakebase compute endpoint, or None while it is still starting.""" + eps = w.api_client.do("GET", f"/api/2.0/postgres/{BRANCH}/endpoints").get("endpoints", []) + return _find(eps[0], "host") if eps else None + + +def pg(): + """A Postgres connection to Lakebase. The password is a short-lived token, so every call asks for a fresh one.""" + import psycopg + token = w.api_client.do("POST", "/api/2.0/postgres/credentials", body={"endpoint": LAKEBASE_ENDPOINT})["token"] + return psycopg.connect(host=lakebase_host(), port=5432, dbname="databricks_postgres", user=ME, password=token, + sslmode="require", autocommit=True, connect_timeout=30) + + +# COMMAND ---------- + +RRF_K = 60 + + +def hybrid_search(conn, qvec, qtext, k=TOP_K, filters=None, pool=20): + """Vector search plus full-text search, merged with reciprocal rank fusion. Notebook 04 builds this up one piece at a time.""" + where, fargs = "", [] + if filters: + where, fargs = " AND meta @> %s::jsonb", [json.dumps({a: str(b) for a, b in filters.items()})] + by_vector = conn.execute(f"SELECT id, title, content, meta FROM {SCHEMA}.chunks WHERE true{where} " + "ORDER BY embedding <=> %s::vector LIMIT %s", (*fargs, vec(qvec), pool)).fetchall() + by_text = conn.execute(f"SELECT id, title, content, meta FROM {SCHEMA}.chunks, websearch_to_tsquery('english', %s) q " + f"WHERE tsv @@ q{where} ORDER BY ts_rank_cd(tsv, q) DESC LIMIT %s", (qtext, *fargs, pool)).fetchall() + score, rows = {}, {} + for ranked in (by_vector, by_text): + for rank, r in enumerate(ranked, 1): + score[r[0]] = score.get(r[0], 0.0) + 1.0 / (RRF_K + rank) + rows[r[0]] = r + return [{"id": i, "title": rows[i][1], "content": rows[i][2], "meta": rows[i][3] if isinstance(rows[i][3], dict) else json.loads(rows[i][3] or "{}"), + "score": round(score[i], 5)} for i in sorted(score, key=score.get, reverse=True)[:k]] + + +def is_hit(results, expect): + """True when any result has all the expected metadata values (for example ticker AAPL and item 1A).""" + return any(all(str(r["meta"].get(k)) == str(v) for k, v in expect.items()) for r in results) + + +print(f"ready: signed in as {ME} | embeddings: {EMBED_ENDPOINT} | chat: {CHAT_ENDPOINT}") diff --git a/labagent/packs/rag/tutorial_pack.py b/labagent/packs/rag/tutorial_pack.py new file mode 100644 index 0000000..cc4e71f --- /dev/null +++ b/labagent/packs/rag/tutorial_pack.py @@ -0,0 +1,127 @@ +"""rag-sec-tutorial: the SEC filings RAG lab, as notebooks a student runs step by step. + +This pack builds nothing itself. It uploads nine notebooks (setup, data, embeddings, pgvector, +search, answer and evaluate, deploy the app, clean up, plus one shared helper) and the app code into your Databricks workspace. +The student then runs them in order and learns what each step does. Notebook 06 deploys the finished app from the code uploaded next to the notebooks. + +Nothing is created in Databricks except the uploaded folder. The notebooks create and remove the database and the app themselves. +""" +from __future__ import annotations + +import json +import re +import shutil +import tempfile +from pathlib import Path +from typing import Any, Dict, List, Optional + +from ... import datasets +from ..base import Pack, PackContext, PackError, Param, Step +from ..databricks.auth import login as login_shared +from ..databricks.auth import workspace_host +from ..databricks.cli import Dbx +from ..databricks.models import CHAT_RE, EMBED_RE +from . import sources as src_mod + +HERE = Path(__file__).parent +NOTEBOOKS = HERE / "tutorial" / "notebooks" +APP = HERE / "app" +DATASET = "sec-10k-2026.10" +SOURCE_ID = "sec-10k-snapshot" + + +def _items(resp: Any, key: str) -> List[Dict[str, Any]]: + if isinstance(resp, list): + return resp + return (resp or {}).get(key) or [] + + +def notebook_text(name: str, replacements: Dict[str, str]) -> str: + text = (NOTEBOOKS / name).read_text() + for key, value in replacements.items(): + text = text.replace(f"__{key}__", value) + return text + + +class RagSecTutorial(Pack): + id = "rag-sec-tutorial" + title = "RAG on Databricks, step by step (notebooks)" + summary = "Uploads nine notebooks and the app code to your workspace so a student builds the SEC filings RAG one step at a time." + aliases = ("rag-tutorial", "sec-tutorial", "rag-notebooks") + platforms = ("databricks",) + cost_tier = "free" + default_name = "demo" + output_keys = ("workspace_path", "url", "notebooks") + params = [ + Param("profile", None, "Databricks CLI profile (default profile if omitted)"), + Param("dataset", DATASET, "Pinned dataset snapshot the notebooks download"), + ] + + def services(self, profile: Optional[str]) -> Dict[str, Any]: + return {"dbx": Dbx(profile)} + + @staticmethod + def folder(ctx: PackContext) -> str: + return f"/Users/{ctx.state['user']}/labagent-rag-tutorial-{ctx.name}" + + def deploy_steps(self, ctx: PackContext) -> List[Step]: + d = lambda: ctx.services["dbx"] # noqa: E731 + + def preflight(c: PackContext): + me = login_shared(c) + user = me.get("userName") + if not user: + raise PackError("not authenticated: run `databricks auth login` and retry") + names = [e.get("name", "") for e in _items(d().json("serving-endpoints", "list"), "endpoints")] + if not any(EMBED_RE.search(n) for n in names) or not any(CHAT_RE.search(n) and not EMBED_RE.search(n) for n in names): + raise PackError("this workspace has no embedding or chat model serving endpoint, which the tutorial needs " + f"(found: {', '.join(names[:8]) or 'none'})") + if c.params["dataset"] not in datasets.registry(): + raise PackError(f"unknown dataset '{c.params['dataset']}'. Known: {', '.join(sorted(datasets.registry()))}") + return {"user": user} + + def upload(c: PackContext): + path = self.folder(c) + entry = datasets.registry()[c.params["dataset"]] + rows = datasets.fetch_snapshot(c.params["dataset"], download=c.services.get("download", datasets.default_http)) + ui = src_mod.get(SOURCE_ID).ui([r["meta"] for r in rows]) + repl = {"SNAPSHOT_URL": entry["url"], "SNAPSHOT_SHA256": entry["sha256"], "WORKSPACE_FOLDER": f"/Workspace{path}"} + with tempfile.TemporaryDirectory() as tmp: + bundle = Path(tmp) / "bundle" + bundle.mkdir() + for nb in sorted(NOTEBOOKS.glob("*.py")): + (bundle / nb.name).write_text(notebook_text(nb.name, repl)) + shutil.copytree(APP, bundle / "app", ignore=shutil.ignore_patterns("__pycache__", "*.pyc")) + (bundle / "app" / "ui.json").write_text(json.dumps(ui, indent=1)) + d().text("workspace", "mkdirs", path) + d().text("workspace", "import-dir", str(bundle), path, "--overwrite") + return {"workspace_path": path, "notebooks": len(list(NOTEBOOKS.glob("*.py")))} + + def verify(c: PackContext): + path = c.state["workspace_path"] + listing = _items(d().json("workspace", "list", path), "objects") + names = {Path(o.get("path", "")).name for o in listing if o.get("object_type") == "NOTEBOOK"} + missing = sorted({p.stem for p in NOTEBOOKS.glob("*.py")} - names) + if missing: + raise PackError(f"these notebooks did not upload: {', '.join(missing)}") + status = d().json("workspace", "get-status", path) + url = f"{workspace_host(c)}/browse/folders/{status.get('object_id')}" + c.log(f"\n Open: {url}\n Start with 00_setup and run the notebooks in order. Finish with 99_cleanup.") + return {"url": url} + + return [Step("preflight", "Check login and the model endpoints the notebooks use", preflight), + Step("upload", "Upload the notebooks and the app code", upload), + Step("verify", "Check that every notebook arrived", verify)] + + def destroy_steps(self, ctx: PackContext) -> List[Step]: + d = lambda: ctx.services["dbx"] # noqa: E731 + + def remove(c: PackContext): + if not c.state.get("workspace_path"): + return {} + c.log(" reminder: this removes only the notebooks. Run 99_cleanup first to remove the database and the app you created.") + d().text("workspace", "delete", c.state["workspace_path"], "--recursive") + return {"workspace_path": None, "url": None} + + return [Step("login", "Check login", lambda c: {"user": login_shared(c).get("userName")}), + Step("folder", "Delete the uploaded tutorial folder", remove)] diff --git a/labagent/packs/rag_ide/pack.py b/labagent/packs/rag_ide/pack.py index 3048ec5..1599cda 100644 --- a/labagent/packs/rag_ide/pack.py +++ b/labagent/packs/rag_ide/pack.py @@ -89,6 +89,8 @@ class RagIde(Pack): 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") + platforms = ("local", "ec2", "modal") + cost_tier = "metered" default_name = "cohort" params = [ Param("platform", "modal", "modal (one serverless sandbox per learner) | ec2 (one shared instance per cohort, much cheaper at scale) | local (Docker on this machine, for demos and tests)"), @@ -114,6 +116,13 @@ class RagIde(Pack): def services(self, profile): return {"modal": _LazyModal(), "ec2": _LazyEc2(), "local": _LazyLocal()} + def outputs(self, state): + """URLs only. State also holds sandbox ids and file paths that a dashboard has no use for.""" + out = {"learners": {n: {"url": i.get("url")} for n, i in sorted((state.get("learners") or {}).items())}} + if state.get("host"): + out["host"] = state["host"] + return out + def _validate(self, ctx: PackContext) -> None: p = ctx.params if p["platform"] not in ("modal", "ec2", "local"): diff --git a/labagent/packs/registry.py b/labagent/packs/registry.py index 79676fc..124e6f1 100644 --- a/labagent/packs/registry.py +++ b/labagent/packs/registry.py @@ -7,10 +7,10 @@ def builtin() -> Dict[str, Pack]: - from .rag.pack import RagLakebase, RagSecFilings, RagSecFilingsLive + from .rag.tutorial_pack import RagSecTutorial from .db_agent.pack import DbAgent from .rag_ide.pack import RagIde - return {p.id: p for p in (RagLakebase(), RagSecFilings(), RagSecFilingsLive(), DbAgent(), RagIde())} + return {p.id: p for p in (RagSecTutorial(), DbAgent(), RagIde())} def all_packs() -> Dict[str, Pack]: @@ -27,6 +27,10 @@ def all_packs() -> Dict[str, Pack]: return found +# Packs that no longer exist, and what replaced them. Without this, the old alias `rag-sec` would silently match the new pack by prefix. +REMOVED = {t: "rag-sec-tutorial" for t in ("rag-lakebase", "rag-demo", "rag-sec-filings", "rag-sec", "sec-filings", "rag-sec-filings-live")} + + def get(pack_id: str) -> Pack: packs = all_packs() if pack_id not in packs: @@ -38,6 +42,9 @@ def resolve(token: str) -> Pack: """Exact id, then alias, then a unique prefix. Ambiguity is an error, never a guess.""" packs = all_packs() t = token.strip().lower() + if t in REMOVED: + raise KeyError(f"the pack {token!r} was removed. Use '{REMOVED[t]}': it uploads notebooks that build the same RAG step by step " + f"(labagent launch {REMOVED[t]} --name demo)") if t in packs: return packs[t] by_alias = [p for p in packs.values() if t in [a.lower() for a in p.aliases]] diff --git a/labagent/run.py b/labagent/run.py index c4e4189..e8d9274 100644 --- a/labagent/run.py +++ b/labagent/run.py @@ -1,6 +1,7 @@ """Single entry point used by the CI runner: execute a job document fetched from the orchestrator.""" from __future__ import annotations +import re from typing import Any, Dict from .context import Ctx @@ -8,7 +9,31 @@ from .envelope import Envelope from . import cleanup as cleanup_mod, iam_ops, sweep -COMMANDS = {"provision", "inventory", "sweep", "reset", "offboard", "audit", "cleanup"} +PACK_COMMANDS = {"launch", "destroy"} +COMMANDS = {"provision", "inventory", "sweep", "reset", "offboard", "audit", "cleanup"} | PACK_COMMANDS +# Parameters a remote job may not set: they point at the runner's own disk or identity, not at the lab. +JOB_PARAM_DENY = {"image", "profile"} +NAME_RE = re.compile(r"^[a-z0-9][a-z0-9-]{0,30}$") + + +def run_pack_job(job: Dict[str, Any]) -> Envelope: + """job = {run_id, command: launch|destroy, pack, name, params{}, mode?: dry-run, secrets_pubkey_pem?}. Needs no AWS context.""" + from .packs import registry + from .packs.base import run_pack + cmd, run_id = job["command"], job.get("run_id") + env = Envelope(command=cmd, run_id=run_id) + name, params = str(job.get("name") or ""), job.get("params") or {} + if not NAME_RE.match(name): + return env.fail("name must be lowercase letters, digits and hyphens (1 to 31 characters)") + if not isinstance(params, dict) or JOB_PARAM_DENY & set(params): + return env.fail("params must be an object and may not set: " + ", ".join(sorted(JOB_PARAM_DENY))) + try: + pack = registry.resolve(str(job.get("pack") or "")) + except (KeyError, ValueError) as e: + return env.fail(str(e.args[0] if e.args else e)[:200]) + out = run_pack(pack, cmd, name, params, dry_run=job.get("mode") == "dry-run", pubkey_pem=job.get("secrets_pubkey_pem")) + out.run_id = run_id + return out def run_job(ctx: Ctx, job: Dict[str, Any]) -> Envelope: @@ -17,6 +42,11 @@ def run_job(ctx: Ctx, job: Dict[str, Any]) -> Envelope: run_id = job.get("run_id") if cmd not in COMMANDS: return Envelope(command=str(cmd)[:40], run_id=run_id).fail("unknown command") + if cmd in PACK_COMMANDS: + try: + return run_pack_job(job) + except Exception as e: + return Envelope(command=cmd, run_id=run_id).fail(f"{type(e).__name__}: {str(e)[:200]}") cohort = job.get("cohort") or {} cid, cdate = cohort.get("id"), cohort.get("created_date") learners = job.get("learners") or [] diff --git a/pyproject.toml b/pyproject.toml index 12837b0..8ec573e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,10 +60,10 @@ labagent = [ "terraform/*/*.tftpl", "terraform/*/.terraform.lock.hcl", "terraform/*/tests/*.tftest.hcl", - "packs/*/sources/*.jsonl", "packs/*/sources/*.json", "packs/*/app/*", "packs/*/app/static/*", + "packs/*/tutorial/notebooks/*.py", ] [tool.pytest.ini_options] diff --git a/tests/test_nl.py b/tests/test_nl.py index d48ad02..d15937a 100644 --- a/tests/test_nl.py +++ b/tests/test_nl.py @@ -15,57 +15,57 @@ def ask(text, reply): def test_the_example_request_becomes_the_right_command(): - plan = ask("Run and deploy rag-sec in databricks", {"action": "launch", "pack": "rag-sec", "name": None, "params": {}}) - assert plan.command_line() == "labagent launch rag-sec-filings --name sec" # alias resolved, default name applied + plan = ask("Run and deploy rag-tutorial in databricks", {"action": "launch", "pack": "rag-tutorial", "name": None, "params": {}}) + assert plan.command_line() == "labagent launch rag-sec-tutorial --name demo" # alias resolved, default name applied def test_tolerates_code_fences_and_chatter_around_the_json(): - raw = 'Sure! ```json\n{"action": "launch", "pack": "rag-sec-filings", "name": null, "params": {}}\n``` hope that helps' - assert ask("deploy rag-sec", raw).pack == "rag-sec-filings" + raw = 'Sure! ```json\n{"action": "launch", "pack": "rag-sec-tutorial", "name": null, "params": {}}\n``` hope that helps' + assert ask("deploy rag-tutorial", raw).pack == "rag-sec-tutorial" def test_name_and_params_are_read_from_the_sentence_not_from_the_model(): - plan = ask("deploy rag-sec called team1 with profile MYPROFILE", - {"action": "launch", "pack": "rag-sec", "name": "prod", "params": {"top_k": "9", "profile": "evil"}}) + plan = ask("deploy rag-tutorial called team1 with profile MYPROFILE", + {"action": "launch", "pack": "rag-tutorial", "name": "prod", "params": {"dataset": "evil", "profile": "evil"}}) assert plan.name == "team1" and plan.params == {"profile": "MYPROFILE"} # the model's name and params are ignored - bare = ask("deploy rag-sec", {"action": "launch", "pack": "rag-sec", "name": "prod", "params": {"profile": "x"}}) - assert bare.name == "sec" and bare.params == {} - live = ask("launch rag-sec-filings-live with tickers AAPL,MSFT", {"action": "launch", "pack": "rag-sec-filings-live"}) - assert live.params == {"tickers": "AAPL,MSFT"} # 'filings' inside the pack id is not a parameter + bare = ask("deploy rag-tutorial", {"action": "launch", "pack": "rag-tutorial", "name": "prod", "params": {"profile": "x"}}) + assert bare.name == "demo" and bare.params == {} + ide = ask("launch rag-ide called cohort1 with learners 5 and platform ec2", {"action": "launch", "pack": "rag-ide"}) + assert ide.name == "cohort1" and ide.params == {"learners": "5", "platform": "ec2"} # several parameters, read by exact rules def test_synonym_actions_are_normalised_and_the_sentence_has_the_last_word(): - assert ask("create rag-sec", {"action": "create", "pack": "rag-sec"}).action == "launch" - assert ask("remove rag-sec", {"action": "remove", "pack": "rag-sec"}).action == "destroy" - assert ask("deploy rag-sec", {"action": "gibberish", "pack": "rag-sec"}).action == "launch" + assert ask("create rag-tutorial", {"action": "create", "pack": "rag-tutorial"}).action == "launch" + assert ask("remove rag-tutorial", {"action": "remove", "pack": "rag-tutorial"}).action == "destroy" + assert ask("deploy rag-tutorial", {"action": "gibberish", "pack": "rag-tutorial"}).action == "launch" def test_a_pack_the_model_guessed_must_be_backed_by_the_sentence(): with pytest.raises(nl.NLError, match="which pack"): - ask("deploy it", {"action": "launch", "pack": "rag-lakebase"}) + ask("deploy it", {"action": "launch", "pack": "rag-sec-tutorial"}) with pytest.raises(nl.NLError, match="which pack"): - ask("run and start the rag app", {"action": "launch", "pack": "rag-lakebase"}) - assert ask("start the demo rag", {"action": "launch", "pack": "rag-lakebase"}).pack == "rag-lakebase" - assert ask("spin up the sec filings rag app", {"action": "launch", "pack": "rag-sec-filings"}).pack == "rag-sec-filings" + ask("run and start the rag app", {"action": "launch", "pack": "rag-sec-tutorial"}) # 'rag' and 'app' are too generic to back a guess + assert ask("start the tutorial notebooks", {"action": "launch", "pack": "rag-sec-tutorial"}).pack == "rag-sec-tutorial" + assert ask("spin up the text to sql agent", {"action": "launch", "pack": "db-agent"}).pack == "db-agent" def test_compound_requests_are_refused_not_half_executed(): with pytest.raises(nl.NLError, match="two things"): - ask("deploy rag-sec and then delete everything", {"action": "launch", "pack": "rag-sec"}) + ask("deploy rag-tutorial and then delete everything", {"action": "launch", "pack": "rag-tutorial"}) def test_a_pack_the_user_typed_beats_the_models_pick(): - plan = ask("deploy rag-sec", {"action": "launch", "pack": "rag-lakebase"}) - assert plan.pack == "rag-sec-filings" and plan.warnings + plan = ask("deploy rag-tutorial", {"action": "launch", "pack": "db-agent"}) + assert plan.pack == "rag-sec-tutorial" and plan.warnings def test_destroy_needs_an_explicit_destructive_request(): with pytest.raises(nl.NLError, match="does not ask to delete"): - ask("deploy rag-sec", {"action": "destroy", "pack": "rag-sec", "name": None, "params": {}}) + ask("deploy rag-tutorial", {"action": "destroy", "pack": "rag-tutorial", "name": None, "params": {}}) with pytest.raises(nl.NLError, match="sounds like a deletion"): - ask("tear down rag-sec", {"action": "launch", "pack": "rag-sec", "name": None, "params": {}}) - ok = ask("delete rag-sec", {"action": "destroy", "pack": "rag-sec", "name": None, "params": {}}) - assert ok.action == "destroy" and ok.command_line() == "labagent destroy rag-sec-filings --name sec" + ask("tear down rag-tutorial", {"action": "launch", "pack": "rag-tutorial", "name": None, "params": {}}) + ok = ask("delete rag-tutorial", {"action": "destroy", "pack": "rag-tutorial", "name": None, "params": {}}) + assert ok.action == "destroy" and ok.command_line() == "labagent destroy rag-sec-tutorial --name demo" def test_unknown_pack_ambiguous_pack_and_junk_are_refused(): @@ -74,14 +74,14 @@ def test_unknown_pack_ambiguous_pack_and_junk_are_refused(): with pytest.raises(nl.NLError, match="ambiguous"): ask("deploy rag", {"action": "launch", "pack": "rag", "name": None, "params": {}}) with pytest.raises(nl.NLError, match="usable command"): - ask("deploy rag-sec", "I cannot do that") + ask("deploy rag-tutorial", "I cannot do that") with pytest.raises(nl.NLError, match="launch, destroy or list"): - ask("rag-sec please", {"action": "rm -rf /", "pack": "rag-sec"}) # no verb in the sentence, nonsense from the model - assert ask("deploy rag-sec", {"action": "rm -rf /", "pack": "rag-sec"}).command_line() == "labagent launch rag-sec-filings --name sec" + ask("rag-tutorial please", {"action": "rm -rf /", "pack": "rag-tutorial"}) # no verb in the sentence, nonsense from the model + assert ask("deploy rag-tutorial", {"action": "rm -rf /", "pack": "rag-tutorial"}).command_line() == "labagent launch rag-sec-tutorial --name demo" def test_shell_metacharacters_are_quoted_in_the_displayed_command(): - plan = ask("deploy rag-sec with profile x&y", {"action": "launch", "pack": "rag-sec"}) + plan = ask("deploy rag-tutorial with profile x&y", {"action": "launch", "pack": "rag-tutorial"}) assert plan.argv()[-1] == "profile=x&y" and "'profile=x&y'" in plan.command_line() @@ -90,12 +90,12 @@ def test_list_needs_no_pack(): def test_prompt_is_small_lists_real_packs_and_asks_only_for_action_and_pack(): - msgs = nl.build_messages("deploy rag-sec", PACKS) + msgs = nl.build_messages("deploy rag-tutorial", PACKS) sys_prompt = msgs[0]["content"] - assert "rag-sec-filings" in sys_prompt and "rag-sec" in sys_prompt and "rag-lakebase" in sys_prompt - assert "top_k" not in sys_prompt and "profile" not in sys_prompt # parameters are not the model's job + assert "rag-sec-tutorial" in sys_prompt and "rag-tutorial" in sys_prompt and "db-agent" in sys_prompt and "rag-ide" in sys_prompt + assert "idle_minutes" not in sys_prompt and "profile" not in sys_prompt # parameters are not the model's job assert len(sys_prompt) < 2500 # short enough for a 0.5b model to follow - assert msgs[1]["content"] == "Request: deploy rag-sec" + assert msgs[1]["content"] == "Request: deploy rag-tutorial" class Resp(io.BytesIO): @@ -152,7 +152,7 @@ def test_settings_precedence_and_locality(monkeypatch): @pytest.fixture def wired(monkeypatch): ran = [] - monkeypatch.setattr(nl, "call_llm", lambda *a, **k: json.dumps({"action": "launch", "pack": "rag-sec", "name": None, "params": {}})) + monkeypatch.setattr(nl, "call_llm", lambda *a, **k: json.dumps({"action": "launch", "pack": "rag-tutorial", "name": None, "params": {}})) real = cli.main def fake_main(argv=None): @@ -165,43 +165,51 @@ def fake_main(argv=None): def test_bare_sentence_goes_to_the_translator_and_yes_runs_it(wired, capsys): - assert cli.main(["Run and deploy rag-sec in databricks", "--yes"]) == 0 - assert wired == [["launch", "rag-sec-filings", "--name", "sec"]] + assert cli.main(["Run and deploy rag-tutorial in databricks", "--yes"]) == 0 + assert wired == [["launch", "rag-sec-tutorial", "--name", "demo"]] assert "Interpreted as" in capsys.readouterr().err def test_print_only_and_non_interactive_do_not_execute(wired, capsys): - assert cli.main(["deploy rag-sec", "--print-only"]) == 0 and wired == [] - assert cli.main(["deploy rag-sec"]) == 2 and wired == [] # no TTY and no --yes + assert cli.main(["deploy rag-tutorial", "--print-only"]) == 0 and wired == [] + assert cli.main(["deploy rag-tutorial"]) == 2 and wired == [] # no TTY and no --yes assert "--yes" in capsys.readouterr().err def test_destroy_requires_typing_the_name(monkeypatch, wired): - monkeypatch.setattr(nl, "call_llm", lambda *a, **k: json.dumps({"action": "destroy", "pack": "rag-sec", "name": None, "params": {}})) + monkeypatch.setattr(nl, "call_llm", lambda *a, **k: json.dumps({"action": "destroy", "pack": "rag-tutorial", "name": None, "params": {}})) monkeypatch.setattr("sys.stdin", type("T", (), {"isatty": lambda s: True})()) monkeypatch.setattr("builtins.input", lambda prompt="": "y") - assert cli.main(["delete rag-sec"]) == 1 and wired == [] # 'y' is not the name - monkeypatch.setattr("builtins.input", lambda prompt="": "sec") - assert cli.main(["delete rag-sec"]) == 0 and wired == [["destroy", "rag-sec-filings", "--name", "sec"]] + assert cli.main(["delete rag-tutorial"]) == 1 and wired == [] # 'y' is not the name + monkeypatch.setattr("builtins.input", lambda prompt="": "demo") + assert cli.main(["delete rag-tutorial"]) == 0 and wired == [["destroy", "rag-sec-tutorial", "--name", "demo"]] def test_known_subcommands_still_work_and_aliases_resolve(capsys): assert cli.main(["packs"]) == 0 - assert cli.main(["launch", "rag-sec", "--dry-run", "--output", "-"]) == 0 # alias + default name, no --name needed - assert '"cohort": "sec"' in capsys.readouterr().out + assert cli.main(["launch", "rag-tutorial", "--dry-run", "--output", "-"]) == 0 # alias + default name, no --name needed + assert '"cohort": "demo"' in capsys.readouterr().out def test_unreachable_model_gives_an_actionable_error(monkeypatch, capsys): def boom(*a, **k): raise nl.NLError("cannot reach the model endpoint at http://localhost:11434/v1. Run: ollama serve") monkeypatch.setattr(nl, "call_llm", boom) - assert cli.main(["deploy rag-sec"]) == 2 and "ollama serve" in capsys.readouterr().err + assert cli.main(["deploy rag-tutorial"]) == 2 and "ollama serve" in capsys.readouterr().err def test_the_guess_guard_stays_correct_as_more_packs_exist(): - """Regression: adding a pack made 'rag' look distinctive. A word shared by 3+ packs never backs a model's guess.""" - from labagent.nl import _distinctive_tokens - d = _distinctive_tokens(PACKS["rag-lakebase"], PACKS) - assert "rag" not in d and "lakebase" in d - assert {"sec", "filings"} <= _distinctive_tokens(PACKS["rag-sec-filings"], PACKS) # a snapshot/live pair still counts + """Regression: adding a pack once made 'rag' look distinctive. A word shared by 3+ packs never backs a model's guess.""" + from labagent.nl import _distinctive_tokens, literal_pack + tut = _distinctive_tokens(PACKS["rag-sec-tutorial"], PACKS) + assert "rag" not in tut and {"tutorial", "notebooks"} <= tut assert {"agent", "databricks"} & _distinctive_tokens(PACKS["db-agent"], PACKS) + assert literal_pack("deploy rag-tutorial on databricks", PACKS).id == "rag-sec-tutorial" # naming the alias outright still wins + assert literal_pack("launch rag-ide for me", PACKS).id == "rag-ide" + assert literal_pack("start the rag-tutorial notebooks", PACKS).id == "rag-sec-tutorial" + + +def test_removed_pack_names_point_to_the_tutorial_instead_of_guessing(): + for old in ("rag-sec", "rag-sec-filings", "rag-lakebase", "rag-sec-filings-live"): + with pytest.raises(KeyError, match="was removed.*rag-sec-tutorial"): + registry.resolve(old) diff --git a/tests/test_pack_jobs.py b/tests/test_pack_jobs.py new file mode 100644 index 0000000..755d7f6 --- /dev/null +++ b/tests/test_pack_jobs.py @@ -0,0 +1,108 @@ +"""launch and destroy as job-contract commands: what the console's runner sends, and what comes back.""" +import base64 +import json + +import pytest +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import padding, rsa + +from labagent.cli import main +from labagent.packs import registry +from labagent.packs.base import Pack, PackError, Param, Step +from labagent.run import run_job + + +@pytest.fixture +def keys(): + priv = rsa.generate_private_key(public_exponent=65537, key_size=2048) + pem = priv.public_key().public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo).decode() + return priv, pem + + +def decrypt(priv, b64): + return priv.decrypt(base64.b64decode(b64), padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)).decode() + + +class Fake(Pack): + id, title, summary = "fake-ide", "Fake", "test" + params = [Param("fail", False), Param("image", "auto"), Param("profile", None)] + + def deploy_steps(self, ctx): + def make(c): + c.mem["passwords"] = {"ana": "pw-ana-1", "ben": "pw-ben-2", "lost": "(lost: relaunch)"} + c.state["learners"] = {"ana": {"url": "https://x/ana", "sandbox_id": "sb-1"}} + if c.params["fail"]: + raise PackError("boom after minting passwords") + return {"host": "h.example"} + return [Step("make", "make", make)] + + def destroy_steps(self, ctx): + return [Step("rm", "remove", lambda c: {})] + + def outputs(self, state): + return {"learners": {n: {"url": i["url"]} for n, i in state.get("learners", {}).items()}} + + +@pytest.fixture(autouse=True) +def fake_pack(monkeypatch, tmp_path): + monkeypatch.setenv("LABAGENT_HOME", str(tmp_path)) + monkeypatch.setattr(registry, "resolve", lambda name: Fake() if name == "fake-ide" else (_ for _ in ()).throw(KeyError(f"unknown pack '{name}'"))) + + +def job(**kw): + return {"run_id": "r1", "command": "launch", "pack": "fake-ide", "name": "demo", "params": {}, **kw} + + +def test_launch_returns_encrypted_passwords_and_urls_only(keys): + priv, pem = keys + env = run_job(None, job(secrets_pubkey_pem=pem)) # no AWS context needed + assert env.status == "ok" and env.run_id == "r1" + assert {s["user"]: decrypt(priv, s["ciphertext"]) for s in env.secrets} == {"ana": "pw-ana-1", "ben": "pw-ben-2"} # '(lost...)' is not a secret + text = json.dumps(env.to_dict()) + assert "pw-ana-1" not in text and "sb-1" not in text and "https://x/ana" in text + + +def test_without_a_public_key_nothing_secret_is_returned(): + env = run_job(None, job()) + assert env.secrets == [] and "pw-ana" not in json.dumps(env.to_dict()) + + +def test_a_failed_launch_still_returns_the_passwords_it_already_minted(keys): + priv, pem = keys + env = run_job(None, job(params={"fail": True}, secrets_pubkey_pem=pem)) + assert env.status == "failed" and "boom" in env.error + assert decrypt(priv, env.secrets[0]["ciphertext"]) == "pw-ana-1" + + +def test_destroy_and_dry_run(): + assert run_job(None, job(command="destroy")).status == "ok" + dry = run_job(None, job(mode="dry-run")) + assert dry.mode == "dry-run" and dry.items[0]["state"] == "planned" + + +@pytest.mark.parametrize("bad", [ + {"name": "Has Space"}, {"name": ""}, {"name": "x" * 40}, {"name": "-lead"}, + {"params": {"image": "/etc"}}, {"params": {"profile": "prod"}}, {"params": ["a"]}, + {"pack": "nope"}, {"params": {"unknown": 1}}, +]) +def test_bad_jobs_are_refused_without_running(bad): + env = run_job(None, job(**bad)) + assert env.status == "failed" and env.error and env.secrets == [] + + +def test_cli_run_executes_pack_jobs_without_aws(tmp_path, keys, capsys): + _, pem = keys + f = tmp_path / "job.json" + f.write_text(json.dumps(job(secrets_pubkey_pem=pem))) + out = tmp_path / "result.json" + assert main(["run", "--job", str(f), "--output", str(out)]) == 0 + res = json.loads(out.read_text()) + assert res["status"] == "ok" and len(res["secrets"]) == 2 + f.write_text(json.dumps(job(params={"fail": True}))) + assert main(["run", "--job", str(f), "--output", str(out)]) == 2 + + +def test_rag_ide_outputs_hide_sandbox_ids(): + from labagent.packs.rag_ide.pack import RagIde + out = RagIde().outputs({"learners": {"ana": {"url": "https://u", "sandbox_id": "sb-9"}}, "host": "1.2.3.4.sslip.io", "credentials_file": "/home/x/c.csv"}) + assert out == {"learners": {"ana": {"url": "https://u"}}, "host": "1.2.3.4.sslip.io"} diff --git a/tests/test_packs.py b/tests/test_packs.py index 9a73f10..f5faf33 100644 --- a/tests/test_packs.py +++ b/tests/test_packs.py @@ -1,124 +1,13 @@ +"""Pack framework, the Databricks CLI wrapper and the shared login handling. Each pack has its own test file.""" import json + import pytest from labagent.cli import main from labagent.packs import base, registry -from labagent.packs.base import Pack, PackError, Param, Step, run_pack +from labagent.packs.base import Pack, PackContext, PackError, Param, Step, run_pack +from labagent.packs.databricks.auth import login from labagent.packs.databricks.cli import Dbx, redact -from labagent.packs.rag.pack import RagLakebase, RagSecFilings -from labagent.packs.rag.sources import get as get_source - -DIM = 8 - - -class FakeDbx: - """Records calls and answers like the Databricks CLI would; state lives in dicts.""" - def __init__(self, projects=None, endpoints=("databricks-gte-large-en", "databricks-meta-llama-3-3-70b-instruct")): - self.projects = dict(projects or {}) - self.apps = {} - self.calls = [] - self.soft_deleted, self.purged = set(), [] - self.ghosts = set() # deleted projects that still answer get-project - self.endpoints = endpoints - - def json(self, *a, body=None): - self.calls.append(a) - if a[:2] == ("current-user", "me"): - return {"userName": "learner@example.com"} - if a[0] == "serving-endpoints": - return [{"name": n} for n in self.endpoints] - if a[:2] == ("postgres", "list-projects"): - return {"projects": [{"name": f"projects/{p}"} for p in self.projects]} - if a[:2] == ("postgres", "create-project"): - if a[2] in self.soft_deleted: - raise PackError("A soft-deleted project currently has this resource name.") - self.projects[a[2]] = True - return {} - if a[:2] == ("postgres", "delete-project"): - n = a[2].split("/")[-1] - self.projects.pop(n, None) - self.soft_deleted.add(n) - return {} - if a[:2] == ("postgres", "list-branches"): - return {"branches": [{"name": f"{a[2]}/branches/main", "status": {"default": True}}]} - if a[:2] == ("postgres", "list-endpoints"): - return {"endpoints": [{"name": f"{a[2]}/endpoints/primary", "status": {"hosts": {"host": "ep.example.invalid"}}}]} - if a[:2] == ("postgres", "generate-database-credential"): - return {"token": "TOKEN-SECRET"} - if a[:2] == ("postgres", "create-role"): - return {} - if a[:2] == ("apps", "create"): - assert len(a) == 2, "apps create takes the name from --json, not a positional" - n = body["name"] - self.apps[n] = {"url": f"https://{n}.example.invalid", "service_principal_client_id": "sp-1234"} - return {} - if a[:2] == ("apps", "get"): - return self.apps[a[2]] - raise AssertionError(f"unexpected json call {a}") - - def text(self, *a, stdin=None): - self.calls.append(a) - if a[:2] == ("apps", "delete"): - self.apps.pop(a[2], None) - return "" - - def exists(self, *a): - self.calls.append(a) - if a[:2] == ("postgres", "get-project"): - return a[2].split("/")[-1] in self.projects or a[2].split("/")[-1] in self.ghosts - if a[:2] == ("apps", "get"): - return a[2] in self.apps - return False - - def api(self, method, path, body=None): - self.calls.append((method, path)) - if method == "delete": - n = path.split("/projects/")[1].split("?")[0] - self.purged.append(n) - self.soft_deleted.discard(n) - return {"done": True} - return {"data": [{"embedding": [0.1] * DIM} for _ in body["input"]]} - - -class FakePg: - def __init__(self): - self.rows = {} - self.sql = [] - self.tokens = [] - - def connect(self, host, user, token, dbname="databricks_postgres"): - self.tokens.append(token) - return FakeConn(self) - - -class FakeConn: - def __init__(self, pg): - self.pg = pg - - def __enter__(self): - return self - - def __exit__(self, *a): - return False - - def execute(self, sql, params=None): - self.pg.sql.append(sql.split("(")[0].strip()[:60]) - rows = self.pg.rows - if sql.startswith("INSERT INTO"): - rows[params[0]] = params - if sql.startswith("TRUNCATE"): - rows.clear() - res = [] - if "SELECT count" in sql: - res = [(len(rows),)] - elif "websearch_to_tsquery" in sql: - res = [] - elif "ORDER BY embedding" in sql: - res = [(p[0], p[1], p[2], json.loads(p[3])) for p in rows.values()] - out = type("R", (), {})() - out.fetchone = lambda: res[0] - out.fetchall = lambda: res - return out @pytest.fixture @@ -127,88 +16,49 @@ def home(tmp_path, monkeypatch): return tmp_path -def services(dbx=None, pg=None): - pg = pg or FakePg() - return {"dbx": dbx or FakeDbx(), "pg": pg.connect}, pg - - -def test_launch_end_to_end_and_idempotent(home): - svc, pg = services() - env = run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=svc) - assert env.status == "ok", env.error - assert [i["state"] for i in env.items] == ["done"] * 9 - assert len(pg.rows) == 25 - out = env.summary["outputs"] - assert out["app_url"].startswith("https://rag-demo") and out["chunks"] == 25 and out["eval_questions"] == 12 - # re-run resumes: project and loaded rows are detected, nothing is re-embedded - svc2, pg2 = services(dbx=svc["dbx"], pg=pg) - again = run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=svc2) - states = {i["step"]: i["state"] for i in again.items} - assert again.status == "ok" and states["project"] == "skipped" and states["load"] == "skipped" and states["build"] == "skipped" - - -def test_secrets_never_in_envelope_or_state(home): - svc, pg = services() - env = run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=svc) - blob = json.dumps(env.to_dict()) + (home / "packs/rag-lakebase/demo.json").read_text() - assert "TOKEN-SECRET" not in blob - assert pg.tokens and set(pg.tokens) == {"TOKEN-SECRET"} - assert oct((home / "packs/rag-lakebase/demo.json").stat().st_mode & 0o777) == "0o600" +class Hello(Pack): + id, title, summary = "hello", "Hello", "demo" + params = [Param("who", "world")] + def deploy_steps(self, ctx): + return [Step("greet", "say hi", lambda c: {"greeting": f"hi {c.params['who']}"})] -def test_free_edition_one_project_limit(home): - svc, _ = services(dbx=FakeDbx(projects={"someone-elses": True})) - env = run_pack(RagLakebase(), "launch", "demo", {}, services=svc) - assert env.status == "failed" and "allows one" in env.error - assert svc["dbx"].projects == {"someone-elses": True} # nothing created - -def test_missing_endpoints_gives_actionable_error(home): - svc, _ = services(dbx=FakeDbx(endpoints=("custom-model",))) - env = run_pack(RagLakebase(), "launch", "demo", {}, services=svc) - assert env.status == "failed" and "--param embedding_endpoint" in env.error +def test_adding_a_pack_needs_no_framework_change(home): + env = run_pack(Hello(), "launch", "x", {"who": "lab"}, services={}) + assert env.status == "ok" and base.load_state("hello", "x")["greeting"] == "hi lab" -def test_destroy_removes_only_what_it_created(home): - svc, _ = services() - run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=svc) - env = run_pack(RagLakebase(), "destroy", "demo", {}, services=svc) - assert env.status == "ok" - assert svc["dbx"].apps == {} and svc["dbx"].projects == {} - assert not (home / "packs/rag-lakebase/demo.json").exists() +def test_unknown_param_is_rejected(home): + env = run_pack(Hello(), "launch", "demo", {"nope": 1}, services={}) + assert env.status == "failed" and "unknown parameter" in env.error -def test_destroy_keeps_a_preexisting_project(home): - dbx = FakeDbx(projects={"rag-demo": True}) - svc, pg = services(dbx=dbx) - run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=svc) - run_pack(RagLakebase(), "destroy", "demo", {}, services=svc) - assert "rag-demo" in dbx.projects - assert any("DROP SCHEMA" in q for q in pg.sql) # a reused project gets only this pack's schema removed +def test_dry_run_touches_nothing(home): + ran = [] + class Probe(Pack): + id = "probe" -def test_dry_run_touches_nothing(home): - svc, pg = services() - env = run_pack(RagLakebase(), "launch", "demo", {}, services=svc, dry_run=True) - assert env.mode == "dry-run" and len(env.items) == 9 and not svc["dbx"].calls and not pg.sql - assert not (home / "packs").exists() + def deploy_steps(self, ctx): + return [Step("a", "first", lambda c: ran.append("a")), Step("b", "second", lambda c: ran.append("b"))] + env = run_pack(Probe(), "launch", "demo", {}, services={}, dry_run=True) + assert env.mode == "dry-run" and len(env.items) == 2 and ran == [] and not (home / "packs").exists() -def test_unknown_param_rejected(home): - env = run_pack(RagLakebase(), "launch", "demo", {"nope": 1}, services=services()[0]) - assert env.status == "failed" and "unknown parameter" in env.error +def test_steps_with_a_check_are_skipped_on_the_second_run(home): + ran = [] -def test_adding_a_pack_needs_no_framework_change(home): - class Hello(Pack): - id, title, summary = "hello", "Hello", "demo" - params = [Param("who", "world")] + class Once(Pack): + id = "once" def deploy_steps(self, ctx): - return [Step("greet", "say hi", lambda c: {"greeting": f"hi {c.params['who']}"})] + return [Step("make", "make it", lambda c: ran.append(1) or {"made": True}, check=lambda c: bool(c.state.get("made")))] - env = run_pack(Hello(), "launch", "x", {"who": "lab"}, services={}) - assert env.status == "ok" and base.load_state("hello", "x")["greeting"] == "hi lab" + run_pack(Once(), "launch", "x", {}, services={}) + again = run_pack(Once(), "launch", "x", {}, services={}) + assert ran == [1] and again.summary["skipped"] == 1 def test_unexpected_exception_reports_class_only(home): @@ -224,12 +74,21 @@ def go(c): assert "abc123" not in json.dumps(env.to_dict()) and "RuntimeError" in env.error -def test_registry_and_cli(capsys, home): - assert "rag-lakebase" in registry.all_packs() - assert main(["packs"]) == 0 and "rag-lakebase" in capsys.readouterr().out +def test_registry_lists_the_current_packs_and_the_cli_refuses_unknown_ones(capsys, home): + assert sorted(registry.all_packs()) == ["db-agent", "rag-ide", "rag-sec-tutorial"] + assert main(["packs"]) == 0 and "rag-sec-tutorial" in capsys.readouterr().out assert main(["launch", "nope", "--name", "x"]) == 2 +def test_removed_packs_exit_with_a_pointer_to_the_tutorial(capsys, home): + for old in ("rag-lakebase", "rag-sec-filings", "rag-sec-filings-live", "rag-sec"): + assert main(["launch", old, "--name", "x"]) == 2 + err = capsys.readouterr().err + assert "was removed" in err and "rag-sec-tutorial" in err + + +# ---- the Databricks CLI wrapper -------------------------------------------------------------------------------------------------- + def test_dbx_wrapper_redacts_tokens(): assert "dapi" not in redact("auth failed dapi" + "a1" * 20) @@ -240,80 +99,20 @@ def run(cmd, **kw): assert "dapi" not in str(e.value) -def test_corpus_is_clean(): - chunks = get_source("builtin").chunks(None) - assert len(chunks) >= 20 and len({c["id"] for c in chunks}) == len(chunks) - assert all("@" not in c["content"] for c in chunks) - - -def test_eval_gate_blocks_deploy_when_retrieval_is_poor(home): - svc, _ = services() - env = run_pack(RagLakebase(), "launch", "demo", {}, services=svc) # default min_hit_rate 0.6; the fake retrieves badly - assert env.status == "failed" and "not deploying" in env.error - assert svc["dbx"].apps == {} - +# ---- shared login handling (used by db-agent and the tutorial) ------------------------------------------------------------------- -def test_dbx_api_retries_rate_limits_then_succeeds(): - calls, naps = [], [] - - def run(cmd, **kw): - calls.append(cmd) - if len(calls) < 3: - return type("P", (), {"returncode": 1, "stdout": "", "stderr": "REQUEST_LIMIT_EXCEEDED: slow down"})() - return type("P", (), {"returncode": 0, "stdout": '{"ok": true}', "stderr": ""})() - d = Dbx(run=run, binary="/bin/databricks", sleep=naps.append) - assert d.api("post", "/x", {"a": 1}) == {"ok": True} and naps == [2, 4] - - -def test_dbx_api_does_not_retry_other_errors(): - def run(cmd, **kw): - return type("P", (), {"returncode": 1, "stdout": "", "stderr": "PERMISSION_DENIED"})() - with pytest.raises(PackError): - Dbx(run=run, binary="/bin/databricks", sleep=lambda s: pytest.fail("slept")).api("post", "/x", {}) - - -def test_destroy_steps_return_no_text_and_tolerate_a_delete_in_flight(home): - """Regression: a step must return a dict or None (the CLI's stdout string once crashed the runner).""" - svc, _ = services() - run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=svc) - dbx = svc["dbx"] - orig = dbx.text - - def text(*a, stdin=None): - if a[:2] == ("apps", "delete"): - dbx.text = orig - raise PackError("databricks apps delete failed: Cannot delete app: state DELETING") - return "some stdout" - dbx.text = text - env = run_pack(RagLakebase(), "destroy", "demo", {}, services=svc) - assert env.status == "ok", env.error - - -def test_changed_corpus_rebuilds_and_replaces_the_table_but_same_corpus_is_skipped(home, monkeypatch): - from labagent.packs.rag.sources.builtin import Builtin - svc, pg = services() - run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=svc) - pg.rows["stale-old-row"] = ("stale-old-row", "t", "c", "{}", "[0.1]") # a row from an earlier, different corpus - monkeypatch.setattr(Builtin, "fingerprint", lambda self, ctx: "builtin:v2") - env = run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=svc) - st = {i["step"]: i["state"] for i in env.items} - assert st["build"] == "done" and st["load"] == "done" # same row count would have looked 'loaded' - assert "stale-old-row" not in pg.rows and len(pg.rows) == 25 - again = run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=svc) - assert {i["step"]: i["state"] for i in again.items}["load"] == "skipped" - - -class StaleDefaultDbx(FakeDbx): +class StaleDefaultDbx: """The CLI's default profile is expired; 'GOOD' is valid for the same workspace.""" - def __init__(self, profiles, **kw): - super().__init__(**kw) + def __init__(self, profiles): self.profile = None self._profiles = profiles def json(self, *a, body=None): - if a[:2] == ("current-user", "me") and self.profile != "GOOD": - raise PackError("databricks current-user me failed: Error: refresh token is invalid") - return super().json(*a, body=body) + if a[:2] == ("current-user", "me"): + if self.profile != "GOOD": + raise PackError("databricks current-user me failed: Error: refresh token is invalid") + return {"userName": "learner@example.com"} + return {} def profiles(self): return self._profiles @@ -322,48 +121,41 @@ def profiles(self): STALE = [{"name": "stale", "host": "https://w1", "valid": False, "default": True}] +def ctx_for(dbx, **params): + return PackContext("x", "demo", params, services={"dbx": dbx}) + + def test_expired_default_profile_falls_back_to_the_valid_one_for_the_same_workspace(home): dbx = StaleDefaultDbx(STALE + [{"name": "GOOD", "host": "https://w1", "valid": True}]) - svc, _ = services(dbx=dbx) - env = run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=svc) - assert env.status == "ok" and dbx.profile == "GOOD" + assert login(ctx_for(dbx))["userName"] == "learner@example.com" and dbx.profile == "GOOD" def test_no_silent_fallback_to_a_different_workspace_and_errors_name_the_valid_profiles(home): dbx = StaleDefaultDbx(STALE + [{"name": "OTHER", "host": "https://w2", "valid": True}]) - env = run_pack(RagLakebase(), "launch", "demo", {}, services=services(dbx=dbx)[0]) - assert env.status == "failed" and "OTHER" in env.error and "--param profile=" in env.error and dbx.profile is None + with pytest.raises(PackError) as e: + login(ctx_for(dbx)) + assert "OTHER" in str(e.value) and "--param profile=" in str(e.value) and dbx.profile is None def test_an_explicit_profile_is_never_overridden(home): dbx = StaleDefaultDbx(STALE + [{"name": "GOOD", "host": "https://w1", "valid": True}]) - env = run_pack(RagLakebase(), "launch", "demo", {"profile": "stale"}, services=services(dbx=dbx)[0]) - assert env.status == "failed" and dbx.profile is None - - -def test_a_deleted_project_that_still_answers_get_is_recreated(home): - """Regression (live): after teardown, get-project still returned the record while list-projects omitted it.""" - dbx = FakeDbx() - dbx.ghosts.add("rag-demo") - env = run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=services(dbx=dbx)[0]) - assert env.status == "ok" and "rag-demo" in dbx.projects - assert {i["step"]: i["state"] for i in env.items}["project"] == "done" - - -def test_destroy_purges_so_the_name_can_be_reused_and_keep_deleted_opts_out(home): - svc, _ = services() - run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=svc) - run_pack(RagLakebase(), "destroy", "demo", {}, services=svc) - dbx = svc["dbx"] - assert dbx.purged == ["rag-demo"] and "rag-demo" not in dbx.soft_deleted - again = run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=svc) # same name works right away - assert again.status == "ok" - run_pack(RagLakebase(), "destroy", "demo", {"keep_deleted": True}, services=svc) - assert dbx.purged == ["rag-demo"] and "rag-demo" in dbx.soft_deleted # opted out: nothing more purged - - -def test_soft_deleted_name_gives_a_plain_explanation(home): - dbx = FakeDbx() - dbx.soft_deleted.add("rag-demo") - env = run_pack(RagLakebase(), "launch", "demo", {"min_hit_rate": 0}, services=services(dbx=dbx)[0]) - assert env.status == "failed" and "soft-deleted" in env.error and "purge=true" in env.error + with pytest.raises(PackError): + login(ctx_for(dbx, profile="stale")) + assert dbx.profile is None + + +# ---- machine-readable catalog (the console seeds its lab list from this) ------------------------------------------------------------ + +def test_packs_json_describes_every_pack_with_platforms_and_params(capsys, home): + assert main(["packs", "--json"]) == 0 + cat = {p["id"]: p for p in json.loads(capsys.readouterr().out)["packs"]} + assert sorted(cat) == ["db-agent", "rag-ide", "rag-sec-tutorial"] + assert cat["rag-ide"]["platforms"] == ["local", "ec2", "modal"] and cat["rag-ide"]["cost_tier"] == "metered" + assert {"name": "hours", "default": 4, "help": "Hard cap per session, 1 to 24", "required": False} in cat["rag-ide"]["params"] + assert all(p["title"] and p["summary"] and p["platforms"] and p["cost_tier"] for p in cat.values()) + + +def test_packs_json_never_contains_a_credential(capsys, home, monkeypatch): + monkeypatch.setenv("GROQ_API_KEY", "gsk_" + "a" * 40) + main(["packs", "--json"]) + assert "gsk_" not in capsys.readouterr().out diff --git a/tests/test_rag_app.py b/tests/test_rag_app.py index 813a4df..b2353a3 100644 --- a/tests/test_rag_app.py +++ b/tests/test_rag_app.py @@ -7,7 +7,6 @@ import pytest -from labagent.packs.rag.sources.builtin import Builtin from labagent.packs.rag.sources.sec import SecFilings APP = Path(__file__).parent.parent / "labagent" / "packs" / "rag" / "app" @@ -26,11 +25,6 @@ def test_sec_ui_suggestions_cover_loaded_companies_with_mixed_topics(): assert all("JPMorgan" not in x["question"] for x in only) # nothing about companies that are not loaded -def test_builtin_ui_only_suggests_loaded_topics(): - ui = Builtin().ui([{"source": "returns"}]) - assert ui["suggestions"] and all(x["filters"] == {"source": "returns"} for x in ui["suggestions"]) - - def test_page_is_self_contained_and_escapes_untrusted_text(): html = (APP / "static" / "index.html").read_text() assert not re.search(r'(src|href)="https?://', html) # no CDN or external assets diff --git a/tests/test_rag_sec.py b/tests/test_rag_sec.py index 4ac5d9a..f96cb93 100644 --- a/tests/test_rag_sec.py +++ b/tests/test_rag_sec.py @@ -1,12 +1,24 @@ +import importlib.util import json +from pathlib import Path import pytest from labagent.packs.base import PackContext, PackError, run_pack from labagent.packs.rag.chunking import chunk_text -from labagent.packs.rag.pack import RagSecFilings, _retrieval from labagent.packs.rag.sources.sec import SecFilings, html_to_text, split_items +APP = Path(__file__).parent.parent / "labagent" / "packs" / "rag" / "app" + + +def _retrieval(): + """The retrieval module ships inside the app bundle, so tests load it by path.""" + spec = importlib.util.spec_from_file_location("labagent_rag_retrieval", APP / "retrieval.py") + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + PARA = "The company faces significant risks from supply chain disruption and component shortages. " * 6 @@ -64,7 +76,7 @@ def http(url, ua): seen.append((url, ua)) return pages[url].encode() params = {"tickers": tickers, "filings": 2, "sections": "1,1A,7", "user_agent": "Test Person test@example.com", **extra} - return PackContext("rag-sec-filings", "t", params, services={"http": http}), seen + return PackContext("sec-source", "t", params, services={"http": http}), seen def test_sec_source_builds_cited_chunks_with_metadata(tmp_path, monkeypatch): @@ -93,13 +105,6 @@ def test_sec_requires_contact_and_limits_volume(tmp_path, monkeypatch): SecFilings().chunks(ctx) -def test_live_pack_requires_user_agent_but_snapshot_pack_does_not(): - from labagent.packs.rag.pack import RagSecFilingsLive - env = run_pack(RagSecFilingsLive(), "launch", "x", {}, services={}) - assert env.status == "failed" and "user_agent" in env.error - assert "user_agent" not in [p.name for p in RagSecFilings().params] - - class Rows: def __init__(self, rows): self.rows = rows diff --git a/tests/test_rag_tutorial.py b/tests/test_rag_tutorial.py new file mode 100644 index 0000000..62e53a5 --- /dev/null +++ b/tests/test_rag_tutorial.py @@ -0,0 +1,244 @@ +import ast +import json +import re +import shutil +import tempfile +from pathlib import Path + +import pytest + +from labagent import datasets +from labagent.packs import registry +from labagent.packs.base import PackError, load_state, run_pack +from labagent.packs.rag import tutorial_pack as tp +from labagent.packs.rag.sources import get as get_source + +NB = tp.NOTEBOOKS +NAMES = ["00_setup", "01_data", "02_embeddings", "03_pgvector", "04_search", "05_answer_and_eval", "06_deploy_app", "99_cleanup", "_common"] + + +def cells(name): + """(kind, text) per cell. kind is 'md', 'magic' or 'code'. MAGIC prefixes are removed.""" + out = [] + for raw in (NB / f"{name}.py").read_text().split("\n# COMMAND ----------\n"): + lines = raw.strip("\n").splitlines() + if lines and lines[0] == "# Databricks notebook source": + lines = lines[1:] + body = [re.sub(r"^# MAGIC ?", "", l) for l in lines] if any(l.startswith("# MAGIC") for l in lines) else lines + text = "\n".join(body).strip("\n") + kind = "md" if text.startswith("%md") else "magic" if text.startswith("%") else "code" + if text: + out.append((kind, text)) + return out + + +def test_every_notebook_is_in_databricks_source_format(): + assert sorted(p.stem for p in NB.glob("*.py")) == sorted(NAMES) + for n in NAMES: + text = (NB / f"{n}.py").read_text() + assert text.startswith("# Databricks notebook source\n"), n + for kind, code in cells(n): + if kind == "code": + ast.parse(code) # every code cell is valid Python + + +def test_notebooks_follow_the_writing_rules(): + for n in NAMES: + text = (NB / f"{n}.py").read_text() + assert not re.search("[—–]", text), f"{n} has an em or en dash" + assert not re.search("[\U0001F300-\U0001FAFF☀-➿]", text), f"{n} has an emoji" + assert "@gmail" not in text and "@example" not in text, n + + +def test_numbered_notebooks_end_with_a_check_and_say_what_is_next(): + for n in NAMES[:-1]: + text = (NB / f"{n}.py").read_text() + assert "check(" in text, f"{n} has no check cell" + if n != "99_cleanup": + assert "**Next:**" in text and "What to notice" in text or n == "00_setup", n + + +def test_psycopg_is_installed_before_it_is_needed_and_python_restarts(): + for n in ("03_pgvector", "04_search", "05_answer_and_eval", "06_deploy_app", "99_cleanup"): + kinds = [(k, t) for k, t in cells(n)] + pip = next(i for i, (k, t) in enumerate(kinds) if t.startswith("%pip install")) + restart = next(i for i, (k, t) in enumerate(kinds) if "restartPython()" in t) + run = next(i for i, (k, t) in enumerate(kinds) if t.startswith("%run ./_common")) + assert pip < restart < run, n + for n in ("00_setup", "01_data", "02_embeddings"): + assert not any(t.startswith("%pip") for _, t in cells(n)), f"{n} should not need a restart" + + +def test_common_is_filled_in_with_the_pinned_dataset_and_no_placeholder_survives(): + entry = datasets.registry()[tp.DATASET] + text = tp.notebook_text("_common.py", {"SNAPSHOT_URL": entry["url"], "SNAPSHOT_SHA256": entry["sha256"], "WORKSPACE_FOLDER": "/Workspace/Users/x/y"}) + assert entry["url"] in text and entry["sha256"] in text and '"/Workspace/Users/x/y"' in text + assert not re.search(r"__[A-Z_]{4,}__", text) + assert "__SNAPSHOT_URL__" in (NB / "_common.py").read_text() # the source file keeps the placeholders + + +def test_the_notebook_search_is_the_same_algorithm_as_the_deployed_apps(): + src = (NB / "_common.py").read_text() + app = (tp.APP / "retrieval.py").read_text() + for fragment in ("websearch_to_tsquery('english', %s)", "ts_rank_cd(tsv, q)", "RRF_K = 60", "ORDER BY embedding <=> %s::vector", "meta @> %s::jsonb", + "1.0 / (RRF_K + rank)"): + assert fragment in src and fragment in app, fragment + + class Conn: # same canned rows for both implementations + def __init__(self): self.n = 0 + def execute(self, sql, params=()): + self.n += 1 + rows = [("a", "A", "text a", {"ticker": "AAPL"}), ("b", "B", "text b", {"ticker": "WMT"}), ("c", "C", "text c", {"ticker": "PFE"})] + return type("R", (), {"fetchall": lambda s, rows=rows if self.n % 2 else rows[::-1]: rows})() + + retrieval = type("M", (), {})() + ns = {} + exec(compile((tp.APP / "retrieval.py").read_text(), "retrieval", "exec"), ns) + mine = {"json": json, "SCHEMA": "tutorial", "TOP_K": 4, "vec": ns["vec_literal"]} + section = src[src.index("RRF_K = 60"):src.index("def is_hit")] + exec(section, mine) + a = [(r["id"], r["score"]) for r in mine["hybrid_search"](Conn(), [0.1, 0.2], "q", 3)] + b = [(r["id"], r["score"]) for r in ns["search"](Conn(), "tutorial", [0.1, 0.2], "q", 3)] + assert a == b and len(a) == 3 + + +def test_notebook_05_evaluation_cases_match_the_packs_own_questions(): + tree = ast.parse("\n\n".join(t for k, t in cells("05_answer_and_eval") if k == "code")) + cases = next(ast.literal_eval(n.value) for n in ast.walk(tree) if isinstance(n, ast.Assign) and getattr(n.targets[0], "id", "") == "cases") + theirs = [(q["question"], q["expect"]) for q in get_source("sec-10k-snapshot").eval_set() if q["expect"]["ticker"] != "XOM"] + assert cases == theirs and len(cases) == 10 + assert "GATE = 0.6" in (NB / "05_answer_and_eval.py").read_text() + + +def test_the_app_notebook_fills_every_placeholder_the_app_yaml_has(): + placeholders = set(re.findall(r"__([A-Z_]+)__", (tp.APP / "app.yaml").read_text())) + nb = (NB / "06_deploy_app.py").read_text() + settings = set(re.findall(r'"([A-Z_]+)":', nb[nb.index("settings = {"):nb.index("yaml_path")])) + assert placeholders == settings, (placeholders, settings) + assert '"CAN_QUERY"' in nb and "SERVICE_PRINCIPAL" in nb and "LAKEBASE_OAUTH_V1" in nb and "GRANT SELECT" in nb + + +def test_staged_apple_demo_is_in_notebook_03_in_order(): + code = [t for k, t in cells("03_pgvector") if k == "code"] + idx = lambda needle, start=0: next(i for i, s in enumerate(code) if i >= start and needle in s) + first = idx("!= \"AAPL\"") + ask1 = idx("vector_search(question, ticker=\"AAPL\")", first) + load2 = idx("load_chunks(later)", ask1) + ask2 = idx("vector_search(question, ticker=\"AAPL\")", load2) + assert first < ask1 < load2 < ask2 and "ON CONFLICT (id) DO UPDATE" in "\n".join(code) + + +# ---- the pack ------------------------------------------------------------------------------------------------------------------ + +class FakeDbx: + def __init__(self, endpoints=("databricks-gte-large-en", "databricks-meta-llama-3-3-70b-instruct")): + self.profile = None + self.endpoints = endpoints + self.calls, self.captured, self.deleted = [], None, [] + + def profiles(self): + return [{"name": "DEFAULT", "host": "https://dbc-test.cloud.databricks.com", "default": True, "valid": True}] + + def json(self, *a, body=None): + self.calls.append(a) + if a[:2] == ("current-user", "me"): + return {"userName": "learner@example.com"} + if a[0] == "serving-endpoints": + return [{"name": n} for n in self.endpoints] + if a[:2] == ("workspace", "list"): + return [{"object_type": "NOTEBOOK", "path": f"{a[2]}/{p.stem}"} for p in tp.NOTEBOOKS.glob("*.py")] + [{"object_type": "DIRECTORY", "path": f"{a[2]}/app"}] + if a[:2] == ("workspace", "get-status"): + return {"object_id": 123456} + return {} + + def text(self, *a, stdin=None): + self.calls.append(a) + if a[:2] == ("workspace", "import-dir"): + self.captured = Path(tempfile.mkdtemp()) / "captured" # outside the pack's own temp dir, which is deleted after the step + shutil.copytree(a[2], self.captured) + self.captured_path = a[3] + if a[:2] == ("workspace", "delete"): + self.deleted.append(a[2]) + return "" + + +@pytest.fixture(autouse=True) +def home(tmp_path, monkeypatch): + monkeypatch.setenv("LABAGENT_HOME", str(tmp_path)) + rows = [{"id": f"AAPL-2025-1A-{i}", "title": "t", "content": "c", "meta": {"ticker": "AAPL", "company": "Apple Inc.", "fiscal_year": "2025", "item": "1A", "section": "Risk Factors"}} + for i in range(3)] + [{"id": "WMT-2025-7-1", "title": "t", "content": "c", "meta": {"ticker": "WMT", "company": "Walmart Inc.", "fiscal_year": "2025", "item": "7", "section": "MD&A"}}] + monkeypatch.setattr(tp.datasets, "fetch_snapshot", lambda *a, **k: rows) + + +def launch(fake, **params): + lines = [] + env = run_pack(tp.RagSecTutorial(), "launch", "demo", params, services={"dbx": fake}, log=lines.append) + return env, "\n".join(lines) + + +def test_registered_with_aliases_and_the_old_one_click_names_are_refused(): + assert registry.resolve("rag-sec-tutorial").id == "rag-sec-tutorial" and registry.resolve("rag-tutorial").id == "rag-sec-tutorial" + with pytest.raises(KeyError, match="was removed"): + registry.resolve("rag-sec") # an old alias must not silently match the new pack by prefix + + +def test_launch_uploads_nine_filled_in_notebooks_and_the_app(tmp_path): + fake = FakeDbx() + env, out = launch(fake) + assert env.exit_code == 0 + assert fake.captured_path == "/Users/learner@example.com/labagent-rag-tutorial-demo" + names = sorted(p.stem for p in fake.captured.glob("*.py")) + assert names == sorted(NAMES) + common = (fake.captured / "_common.py").read_text() + entry = datasets.registry()[tp.DATASET] + assert entry["sha256"] in common and '"/Workspace/Users/learner@example.com/labagent-rag-tutorial-demo"' in common + assert not re.search(r"__[A-Z_]{4,}__", common) + assert {p.name for p in (fake.captured / "app").iterdir()} >= {"app.py", "retrieval.py", "app.yaml", "requirements.txt", "ui.json", "static"} + ui = json.loads((fake.captured / "app" / "ui.json").read_text()) + assert ui and "AAPL" in json.dumps(ui) # the UI menu was built from the dataset + assert not list(fake.captured.rglob("__pycache__")) + assert "browse/folders/123456" in out and "00_setup" in out + assert load_state("rag-sec-tutorial", "demo")["workspace_path"] == fake.captured_path + + +def test_launch_creates_nothing_in_databricks_except_the_folder(): + fake = FakeDbx() + launch(fake) + verbs = {a[:2] for a in fake.calls if a[0] in ("workspace", "postgres", "apps")} + assert verbs <= {("workspace", "mkdirs"), ("workspace", "import-dir"), ("workspace", "list"), ("workspace", "get-status")} + + +def test_missing_model_endpoints_fail_before_uploading_anything(): + fake = FakeDbx(endpoints=("some-other-endpoint",)) + env, _ = launch(fake) + assert env.exit_code != 0 and "embedding or chat" in str(env.__dict__) and fake.captured is None + + +def test_unknown_dataset_is_rejected(): + env, _ = launch(FakeDbx(), dataset="sec-10k-1999.01") + assert env.exit_code != 0 and "unknown dataset" in str(env.__dict__) + + +def test_a_missing_notebook_is_reported_by_the_verify_step(): + fake = FakeDbx() + real = fake.json + fake.json = lambda *a, body=None: ([x for x in real(*a) if not x["path"].endswith("04_search")] if a[:2] == ("workspace", "list") else real(*a, body=body)) + env, _ = launch(fake) + assert env.exit_code != 0 and "04_search" in str(env.__dict__) + + +def test_destroy_deletes_the_folder_and_reminds_about_the_cleanup_notebook(): + fake = FakeDbx() + launch(fake) + lines = [] + env = run_pack(tp.RagSecTutorial(), "destroy", "demo", {}, services={"dbx": fake}, log=lines.append) + assert env.exit_code == 0 and fake.deleted == ["/Users/learner@example.com/labagent-rag-tutorial-demo"] + assert "99_cleanup" in "\n".join(lines) + again = run_pack(tp.RagSecTutorial(), "destroy", "demo", {}, services={"dbx": fake}) + assert again.exit_code == 0 and len(fake.deleted) == 1 # safe to run twice + + +def test_notebooks_ship_in_the_wheel(): + root = Path(__file__).parent.parent + assert "packs/*/tutorial/notebooks/*.py" in (root / "pyproject.toml").read_text() + assert (root / "labagent" / "packs" / "rag" / "tutorial" / "__init__.py").exists()