omarsol commited on
Commit
e73e01f
·
1 Parent(s): 4cca78b

add knowledge base wiki exploration tool

Browse files
.gitignore CHANGED
@@ -164,6 +164,12 @@ cython_debug/
164
  .vscode/
165
  data/chroma-db**/
166
  evaluation_data/chroma-db**/
 
 
 
 
 
 
167
  data/*_md_files/
168
  data/full_stack_ai_engineering/
169
  data/beginner_python_for_ai_engineering/
 
164
  .vscode/
165
  data/chroma-db**/
166
  evaluation_data/chroma-db**/
167
+ # `data/kb/` is built by data.scraping_scripts.build_kb_artifacts and
168
+ # uploaded to HuggingFace by data.scraping_scripts.upload_dbs_to_hf.
169
+ # The runtime downloads it on first start (see scripts/setup.ensure_local_vector_db).
170
+ # Treat it as a build artifact, not a source file. The canonical AGENTS.md
171
+ # content lives in data/scraping_scripts/kb_agents_template.md.
172
+ data/kb/
173
  data/*_md_files/
174
  data/full_stack_ai_engineering/
175
  data/beginner_python_for_ai_engineering/
README.md CHANGED
@@ -181,11 +181,13 @@ Notes:
181
  After updating the JSONL corpus, rebuild the local Chroma index with:
182
 
183
  ```bash
 
 
184
  uv run -m data.scraping_scripts.add_context_to_nodes
185
  uv run -m data.scraping_scripts.create_vector_stores all_sources
186
  ```
187
 
188
- The first command now builds the chunk manifest used by the workflows, and the second command writes dense embeddings into the local Chroma database.
189
  The vector-store build now shows embedding and Chroma upsert progress in the terminal.
190
 
191
  ### Updating Data Sources
 
181
  After updating the JSONL corpus, rebuild the local Chroma index with:
182
 
183
  ```bash
184
+ uv run -m data.scraping_scripts.build_kb_artifacts
185
+ uv run -m data.scraping_scripts.update_kb_wiki --changed-only
186
  uv run -m data.scraping_scripts.add_context_to_nodes
187
  uv run -m data.scraping_scripts.create_vector_stores all_sources
188
  ```
189
 
190
+ The KB commands generate browseable markdown, indexes, and wiki navigation pages for the agent. The context command builds the chunk manifest used by the workflows, and the vector command writes dense embeddings into the local Chroma database.
191
  The vector-store build now shows embedding and Chroma upsert progress in the terminal.
192
 
193
  ### Updating Data Sources
data/.gitattributes ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.lz4 filter=lfs diff=lfs merge=lfs -text
12
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
13
+ *.model filter=lfs diff=lfs merge=lfs -text
14
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
15
+ *.npy filter=lfs diff=lfs merge=lfs -text
16
+ *.npz filter=lfs diff=lfs merge=lfs -text
17
+ *.onnx filter=lfs diff=lfs merge=lfs -text
18
+ *.ot filter=lfs diff=lfs merge=lfs -text
19
+ *.parquet filter=lfs diff=lfs merge=lfs -text
20
+ *.pb filter=lfs diff=lfs merge=lfs -text
21
+ *.pickle filter=lfs diff=lfs merge=lfs -text
22
+ *.pkl filter=lfs diff=lfs merge=lfs -text
23
+ *.pt filter=lfs diff=lfs merge=lfs -text
24
+ *.pth filter=lfs diff=lfs merge=lfs -text
25
+ *.rar filter=lfs diff=lfs merge=lfs -text
26
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
27
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
29
+ *.tar filter=lfs diff=lfs merge=lfs -text
30
+ *.tflite filter=lfs diff=lfs merge=lfs -text
31
+ *.tgz filter=lfs diff=lfs merge=lfs -text
32
+ *.wasm filter=lfs diff=lfs merge=lfs -text
33
+ *.xz filter=lfs diff=lfs merge=lfs -text
34
+ *.zip filter=lfs diff=lfs merge=lfs -text
35
+ *.zst filter=lfs diff=lfs merge=lfs -text
36
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
37
+ # Audio files - uncompressed
38
+ *.pcm filter=lfs diff=lfs merge=lfs -text
39
+ *.sam filter=lfs diff=lfs merge=lfs -text
40
+ *.raw filter=lfs diff=lfs merge=lfs -text
41
+ # Audio files - compressed
42
+ *.aac filter=lfs diff=lfs merge=lfs -text
43
+ *.flac filter=lfs diff=lfs merge=lfs -text
44
+ *.mp3 filter=lfs diff=lfs merge=lfs -text
45
+ *.ogg filter=lfs diff=lfs merge=lfs -text
46
+ *.wav filter=lfs diff=lfs merge=lfs -text
47
+ # Image files - uncompressed
48
+ *.bmp filter=lfs diff=lfs merge=lfs -text
49
+ *.gif filter=lfs diff=lfs merge=lfs -text
50
+ *.png filter=lfs diff=lfs merge=lfs -text
51
+ *.tiff filter=lfs diff=lfs merge=lfs -text
52
+ # Image files - compressed
53
+ *.jpg filter=lfs diff=lfs merge=lfs -text
54
+ *.jpeg filter=lfs diff=lfs merge=lfs -text
55
+ *.webp filter=lfs diff=lfs merge=lfs -text
56
+ ai-tutor-vector-db/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text
57
+ chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text
58
+ ai_for_business_professionals/AI[[:space:]]for[[:space:]]Business[[:space:]]Professionals[[:space:]]190f9b6f42708087863df100e9b4b556/(Video)[[:space:]]AI[[:space:]]Apps[[:space:]]and[[:space:]]AI[[:space:]]Audio[[:space:]]Using[[:space:]]NotebookLM[[:space:]]to[[:space:]]L[[:space:]]1b0f9b6f4270800f8164d5a4a365196a/Screen_Recording_2025-03-14_at_10.05.04.mov filter=lfs diff=lfs merge=lfs -text
59
+ ai_for_business_professionals/AI[[:space:]]for[[:space:]]Business[[:space:]]Professionals[[:space:]]190f9b6f42708087863df100e9b4b556/(Video)[[:space:]]AI[[:space:]]Apps[[:space:]]and[[:space:]]AI[[:space:]]Audio[[:space:]]Using[[:space:]]NotebookLM[[:space:]]to[[:space:]]L[[:space:]]1b0f9b6f4270800f8164d5a4a365196a/Screen_Recording_2025-03-14_at_10.15.13.mov filter=lfs diff=lfs merge=lfs -text
60
+ ai_for_business_professionals/AI[[:space:]]for[[:space:]]Business[[:space:]]Professionals[[:space:]]190f9b6f42708087863df100e9b4b556/(video)[[:space:]]Making[[:space:]]Prompts[[:space:]]Robust[[:space:]]and[[:space:]]Ready[[:space:]]for[[:space:]]Re-Use[[:space:]]190f9b6f42708042bed5e17583e364a9/1.mov filter=lfs diff=lfs merge=lfs -text
61
+ ai_for_business_professionals/AI[[:space:]]for[[:space:]]Business[[:space:]]Professionals[[:space:]]190f9b6f42708087863df100e9b4b556/(video)[[:space:]]Making[[:space:]]Prompts[[:space:]]Robust[[:space:]]and[[:space:]]Ready[[:space:]]for[[:space:]]Re-Use[[:space:]]190f9b6f42708042bed5e17583e364a9/2.mov filter=lfs diff=lfs merge=lfs -text
62
+ ai_for_business_professionals/AI[[:space:]]for[[:space:]]Business[[:space:]]Professionals[[:space:]]190f9b6f42708087863df100e9b4b556/(video)[[:space:]]Making[[:space:]]Prompts[[:space:]]Robust[[:space:]]and[[:space:]]Ready[[:space:]]for[[:space:]]Re-Use[[:space:]]190f9b6f42708042bed5e17583e364a9/3.mov filter=lfs diff=lfs merge=lfs -text
63
+ ai_for_business_professionals/AI[[:space:]]for[[:space:]]Business[[:space:]]Professionals[[:space:]]190f9b6f42708087863df100e9b4b556/(video)[[:space:]]Making[[:space:]]Prompts[[:space:]]Robust[[:space:]]and[[:space:]]Ready[[:space:]]for[[:space:]]Re-Use[[:space:]]190f9b6f42708042bed5e17583e364a9/4.mov filter=lfs diff=lfs merge=lfs -text
64
+ ai_for_business_professionals/AI[[:space:]]for[[:space:]]Business[[:space:]]Professionals[[:space:]]190f9b6f42708087863df100e9b4b556/Using[[:space:]]LLMs[[:space:]]for[[:space:]]Research[[:space:]](Perplexity)[[:space:]]-[[:space:]]Video[[:space:]]190f9b6f42708069954bd75cc797a378/Screen_Recording_2024-11-27_at_17.10.38.mov filter=lfs diff=lfs merge=lfs -text
65
+ chroma-db-all_sources/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text
data/scraping_scripts/README.md CHANGED
@@ -138,6 +138,79 @@ To update library documentation from GitHub repositories:
138
  uv run -m data.scraping_scripts.update_docs_workflow
139
  ```
140
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
141
  This will update all supported documentation sources. You can also specify specific sources:
142
 
143
  ```bash
@@ -225,3 +298,9 @@ leaving the source configured as active for a future rebuild.
225
  6. Vector-store creation:
226
  - `create_vector_stores.py` now shows terminal progress for embedding generation and Chroma upserts
227
  - `all_sources_contextual_nodes.pkl` stores chunk records for embedding; legacy pickles containing LlamaIndex nodes are still accepted by the runtime compatibility helpers
 
 
 
 
 
 
 
138
  uv run -m data.scraping_scripts.update_docs_workflow
139
  ```
140
 
141
+ ### Rebuilding the KB from scratch
142
+
143
+ If you've blown away local KB state (e.g. `rm -rf data/kb data/chroma-db-all_sources`)
144
+ and want everything back, **just run the docs workflow** — it now handles
145
+ fresh-build cases automatically.
146
+
147
+ ```bash
148
+ uv run -m data.scraping_scripts.update_docs_workflow
149
+ ```
150
+
151
+ What happens under the hood when local state is missing:
152
+
153
+ 1. `ensure_required_files_exist` downloads `all_sources_data.jsonl`,
154
+ `all_sources_contextual_nodes.pkl`, and every per-source JSONL (including
155
+ courses) from the private HF dataset `towardsai-tutors/ai-tutor-data`.
156
+ 2. Fresh docs are re-fetched from GitHub / llms.txt indexes.
157
+ 3. `build_kb_artifacts.py` regenerates `data/kb/raw/` and `data/kb/generated/`
158
+ from the JSONL. Course raw pages are reconstructed from the JSONL too —
159
+ you do **not** need the original Notion exports in `data/<course_name>/`.
160
+ 4. `update_kb_wiki.py --changed-only` rebuilds `data/kb/wiki/`. When it
161
+ detects an empty `wiki/` it auto-promotes to a full seed (topic pages,
162
+ recipes/errors index pages), so a clean rebuild produces the full wiki.
163
+ 5. `data/kb/AGENTS.md` is overwritten from
164
+ `data/scraping_scripts/kb_agents_template.md` (the canonical template,
165
+ tracked in git so it survives any `rm -rf data/`).
166
+ 6. Chroma vector store is rebuilt by `create_vector_stores.py`.
167
+
168
+ Prerequisites for a from-scratch rebuild:
169
+
170
+ - `HF_TOKEN` with access to `towardsai-tutors/ai-tutor-data` (so course
171
+ JSONLs and the contextual nodes PKL can be downloaded).
172
+ - `GITHUB_TOKEN` (rate-limit headroom for docs downloads).
173
+ - `GEMINI_API_KEY` or `GOOGLE_API_KEY` (used by `add_context_to_nodes.py`).
174
+ - `COHERE_API_KEY` (used by `create_vector_stores.py`).
175
+
176
+ What this workflow **cannot** restore: the original Notion-exported course
177
+ markdown in `data/<course_name>/`. You don't need it for the runtime — the
178
+ chunked content lives in the JSONL on HF — but if you want to re-run
179
+ `add_course_workflow.py` against fresh course content, you'd need a new
180
+ Notion export.
181
+
182
+ ### Editing KB agent guidance permanently
183
+
184
+ Do **not** edit `data/kb/AGENTS.md` directly. The file is overwritten in two
185
+ places, both reading from the canonical template at
186
+ `data/scraping_scripts/kb_agents_template.md`:
187
+
188
+ - `data.scraping_scripts.update_kb_wiki.write_agents_md` — rewrites it during
189
+ every `update_docs_workflow.py` run, before uploading to HuggingFace.
190
+ - `scripts.setup.ensure_kb_agents_md` — rewrites it on every runtime startup,
191
+ after `ensure_local_vector_db` downloads the snapshot. This catches the
192
+ case where the HF snapshot has a stale AGENTS.md (e.g. uploaded before a
193
+ template change landed in git) and ensures the live file always matches
194
+ the version checked into the current branch.
195
+
196
+ Edit the template, commit it, and either run the docs workflow (to refresh
197
+ the HF snapshot) or just restart the chatbot (the template gets copied into
198
+ place on startup either way).
199
+
200
+ ### Where `data/kb/` actually lives
201
+
202
+ `data/kb/` is gitignored (it's a 100MB+ build artifact). It's:
203
+
204
+ - **Built** by `data.scraping_scripts.build_kb_artifacts` (from
205
+ `data/all_sources_data.jsonl`) and `data.scraping_scripts.update_kb_wiki`.
206
+ - **Uploaded** to `towardsai-tutors/ai-tutor-vector-db` by
207
+ `data.scraping_scripts.upload_dbs_to_hf` (already includes `kb/**`).
208
+ - **Downloaded** by `scripts.setup.ensure_local_vector_db` on the first
209
+ chatbot start (or any start where the local KB is missing).
210
+
211
+ Treat it the same way as `data/chroma-db-all_sources/`: never commit it,
212
+ never edit it by hand, regenerate or re-download when you need it.
213
+
214
  This will update all supported documentation sources. You can also specify specific sources:
215
 
216
  ```bash
 
298
  6. Vector-store creation:
299
  - `create_vector_stores.py` now shows terminal progress for embedding generation and Chroma upserts
300
  - `all_sources_contextual_nodes.pkl` stores chunk records for embedding; legacy pickles containing LlamaIndex nodes are still accepted by the runtime compatibility helpers
301
+
302
+ 7. KB wiki artifact creation:
303
+ - `build_kb_artifacts.py` converts `all_sources_data.jsonl` into generated markdown under `data/kb/raw/`
304
+ - `update_kb_wiki.py --changed-only` refreshes wiki navigation pages under `data/kb/wiki/`
305
+ - The docs and course workflows run these steps automatically unless `--skip-kb` is passed
306
+ - `upload_dbs_to_hf.py` uploads `data/kb/**` with the vector database so the runtime can download one KB bundle
data/scraping_scripts/add_course_workflow.py CHANGED
@@ -381,6 +381,21 @@ def create_vector_stores() -> None:
381
  logger.info("Successfully created vector stores")
382
 
383
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
384
  def upload_to_huggingface(upload_jsonl: bool = False) -> None:
385
  """Upload databases to HuggingFace."""
386
  logger.info("Uploading databases to HuggingFace")
@@ -465,6 +480,11 @@ def main():
465
  parser.add_argument(
466
  "--skip-vectors", action="store_true", help="Skip vector store creation"
467
  )
 
 
 
 
 
468
  parser.add_argument(
469
  "--skip-upload", action="store_true", help="Skip uploading to HuggingFace"
470
  )
@@ -514,6 +534,9 @@ def main():
514
  if not args.skip_merge:
515
  rebuild_all_sources(courses)
516
 
 
 
 
517
  if args.purge_sources:
518
  purge_sources_from_pkl(args.purge_sources)
519
 
 
381
  logger.info("Successfully created vector stores")
382
 
383
 
384
+ def build_kb_artifacts() -> None:
385
+ """Build generated markdown corpus, indexes, and refreshed wiki pages."""
386
+ logger.info("Building KB artifacts")
387
+ result = run_module("data.scraping_scripts.build_kb_artifacts")
388
+ if result.returncode != 0:
389
+ logger.error("Error building KB artifacts - check output above")
390
+ sys.exit(1)
391
+
392
+ result = run_module("data.scraping_scripts.update_kb_wiki", "--changed-only")
393
+ if result.returncode != 0:
394
+ logger.error("Error updating KB wiki - check output above")
395
+ sys.exit(1)
396
+ logger.info("Successfully built KB artifacts")
397
+
398
+
399
  def upload_to_huggingface(upload_jsonl: bool = False) -> None:
400
  """Upload databases to HuggingFace."""
401
  logger.info("Uploading databases to HuggingFace")
 
480
  parser.add_argument(
481
  "--skip-vectors", action="store_true", help="Skip vector store creation"
482
  )
483
+ parser.add_argument(
484
+ "--skip-kb",
485
+ action="store_true",
486
+ help="Skip generated KB markdown/wiki artifact creation",
487
+ )
488
  parser.add_argument(
489
  "--skip-upload", action="store_true", help="Skip uploading to HuggingFace"
490
  )
 
534
  if not args.skip_merge:
535
  rebuild_all_sources(courses)
536
 
537
+ if not args.skip_kb:
538
+ build_kb_artifacts()
539
+
540
  if args.purge_sources:
541
  purge_sources_from_pkl(args.purge_sources)
542
 
data/scraping_scripts/build_kb_artifacts.py ADDED
@@ -0,0 +1,546 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Build agent-browsable KB artifacts from the normalized JSONL corpus."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import argparse
6
+ import csv
7
+ import json
8
+ import logging
9
+ import re
10
+ import shutil
11
+ import uuid
12
+ from dataclasses import dataclass
13
+ from pathlib import Path
14
+ from typing import Any, Iterable
15
+
16
+ try:
17
+ from data.scraping_scripts.process_md_files import (
18
+ clean_document_content,
19
+ content_sha256,
20
+ extract_title,
21
+ generate_url,
22
+ load_source_extension_manifest,
23
+ load_source_url_manifest,
24
+ should_include_file,
25
+ split_frontmatter,
26
+ slugify_identifier,
27
+ stable_doc_id,
28
+ )
29
+ from data.scraping_scripts.source_registry import COURSE_SOURCE_KEYS, SOURCE_CONFIGS
30
+ except ModuleNotFoundError:
31
+ from process_md_files import (
32
+ clean_document_content,
33
+ content_sha256,
34
+ extract_title,
35
+ generate_url,
36
+ load_source_extension_manifest,
37
+ load_source_url_manifest,
38
+ should_include_file,
39
+ split_frontmatter,
40
+ slugify_identifier,
41
+ stable_doc_id,
42
+ )
43
+ from source_registry import COURSE_SOURCE_KEYS, SOURCE_CONFIGS
44
+
45
+
46
+ DEFAULT_INPUT_FILE = Path("data/all_sources_data.jsonl")
47
+ DEFAULT_OUTPUT_DIR = Path("data/kb")
48
+ RAW_DIR_NAME = "raw"
49
+ GENERATED_DIR_NAME = "generated"
50
+ MARKDOWN_EXTENSIONS = {".md", ".mdx"}
51
+ HEADING_RE = re.compile(r"^(#{1,6})\s+(.+?)\s*#*\s*$")
52
+ FENCE_RE = re.compile(r"^\s*(`{3,}|~{3,})")
53
+ INLINE_CODE_RE = re.compile(r"`([^`\n]{2,160})`")
54
+ SYMBOL_RE = re.compile(
55
+ r"\b[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)+\b|"
56
+ r"\b[a-z_]+_[a-z0-9_]+\b|"
57
+ r"\b[A-Z][A-Za-z0-9_]{2,}\b"
58
+ )
59
+ logger = logging.getLogger(__name__)
60
+
61
+
62
+ @dataclass(frozen=True)
63
+ class CorpusRecord:
64
+ doc_id: str
65
+ legacy_doc_id: str
66
+ title: str
67
+ source: str
68
+ source_group: str
69
+ source_path: str
70
+ url: str
71
+ tokens: int
72
+ content_hash: str
73
+ content: str
74
+
75
+
76
+ @dataclass(frozen=True)
77
+ class OriginalMarkdown:
78
+ source_path: str
79
+ path: Path
80
+ url: str
81
+ title: str
82
+ content_hash: str
83
+ content: str
84
+
85
+
86
+ @dataclass
87
+ class OriginalMarkdownIndex:
88
+ by_source_path: dict[str, OriginalMarkdown]
89
+ by_url: dict[str, list[OriginalMarkdown]]
90
+ by_content_hash: dict[str, list[OriginalMarkdown]]
91
+ by_title: dict[str, list[OriginalMarkdown]]
92
+
93
+
94
+ def load_jsonl(path: Path) -> list[dict[str, Any]]:
95
+ rows: list[dict[str, Any]] = []
96
+ with path.open("r", encoding="utf-8") as handle:
97
+ for line in handle:
98
+ if line.strip():
99
+ rows.append(json.loads(line))
100
+ return rows
101
+
102
+
103
+ def looks_like_uuid(value: str) -> bool:
104
+ try:
105
+ uuid.UUID(value)
106
+ except (TypeError, ValueError):
107
+ return False
108
+ return True
109
+
110
+
111
+ def normalize_source_path(row: dict[str, Any]) -> str:
112
+ source_path = str(row.get("source_path") or "").replace("\\", "/").strip("/")
113
+ path = Path(source_path)
114
+ if path.is_absolute() or ".." in path.parts:
115
+ return ""
116
+ return source_path
117
+
118
+
119
+ def normalize_record(row: dict[str, Any]) -> CorpusRecord:
120
+ content = str(row.get("content") or "")
121
+ source = str(row.get("source") or "unknown")
122
+ title = str(row.get("name") or row.get("title") or "Untitled")
123
+ url = str(row.get("url") or "")
124
+ source_path = normalize_source_path(row)
125
+ if source in COURSE_SOURCE_KEYS:
126
+ source_path = ""
127
+ content_hash = str(row.get("content_hash") or content_sha256(content))
128
+ raw_doc_id = str(row.get("doc_id") or "")
129
+ if raw_doc_id and not looks_like_uuid(raw_doc_id) and ":" in raw_doc_id:
130
+ doc_id = raw_doc_id
131
+ else:
132
+ doc_id = stable_doc_id(
133
+ source=source,
134
+ source_path=source_path,
135
+ title=title,
136
+ url=url,
137
+ content_hash=content_hash,
138
+ )
139
+ try:
140
+ tokens = int(row.get("tokens") or 0)
141
+ except (TypeError, ValueError):
142
+ tokens = 0
143
+ return CorpusRecord(
144
+ doc_id=doc_id,
145
+ legacy_doc_id=raw_doc_id if raw_doc_id != doc_id else "",
146
+ title=title,
147
+ source=source,
148
+ source_group="courses" if source in COURSE_SOURCE_KEYS else "docs",
149
+ source_path=source_path,
150
+ url=url,
151
+ tokens=tokens,
152
+ content_hash=content_hash,
153
+ content=content,
154
+ )
155
+
156
+
157
+ def quote_frontmatter(value: str) -> str:
158
+ return json.dumps(value, ensure_ascii=False)
159
+
160
+
161
+ def strip_existing_frontmatter(content: str) -> str:
162
+ frontmatter, body = split_frontmatter(content)
163
+ return body if frontmatter is not None else content
164
+
165
+
166
+ def markdown_with_frontmatter(
167
+ record: CorpusRecord,
168
+ *,
169
+ generated_from: Path,
170
+ content: str,
171
+ source_path: str = "",
172
+ original_path: str = "",
173
+ ) -> str:
174
+ frontmatter_source_path = source_path or record.source_path
175
+ lines = [
176
+ "---",
177
+ f"doc_id: {quote_frontmatter(record.doc_id)}",
178
+ f"legacy_doc_id: {quote_frontmatter(record.legacy_doc_id)}",
179
+ f"source: {quote_frontmatter(record.source)}",
180
+ f"source_group: {quote_frontmatter(record.source_group)}",
181
+ f"title: {quote_frontmatter(record.title)}",
182
+ f"url: {quote_frontmatter(record.url)}",
183
+ f"source_path: {quote_frontmatter(frontmatter_source_path)}",
184
+ f"tokens: {record.tokens}",
185
+ f"content_hash: {quote_frontmatter(record.content_hash)}",
186
+ f"generated_from: {quote_frontmatter(str(generated_from))}",
187
+ ]
188
+ if original_path:
189
+ lines.append(f"original_path: {quote_frontmatter(original_path)}")
190
+ lines.extend(["---", ""])
191
+ body = strip_existing_frontmatter(content).strip()
192
+ return "\n".join(lines) + body + "\n"
193
+
194
+
195
+ def record_slug(record: CorpusRecord) -> str:
196
+ basis = record.source_path or record.title or record.doc_id
197
+ slug = slugify_identifier(basis)
198
+ if slug == "untitled":
199
+ slug = record.content_hash.removeprefix("sha256:")[:12]
200
+ return slug
201
+
202
+
203
+ def safe_relative_path(value: str) -> Path:
204
+ path = Path(value.replace("\\", "/"))
205
+ if path.is_absolute() or ".." in path.parts:
206
+ raise ValueError(f"Unsafe source path: {value}")
207
+ return path
208
+
209
+
210
+ def unique_path(candidate: Path, used_paths: set[Path], content_hash: str) -> Path:
211
+ if candidate not in used_paths:
212
+ used_paths.add(candidate)
213
+ return candidate
214
+
215
+ suffix = content_hash.removeprefix("sha256:")[:8]
216
+ stem = candidate.stem
217
+ candidate = candidate.with_name(f"{stem}-{suffix}{candidate.suffix}")
218
+ counter = 2
219
+ while candidate in used_paths:
220
+ candidate = candidate.with_name(
221
+ f"{stem}-{suffix}-{counter}{candidate.suffix}"
222
+ )
223
+ counter += 1
224
+ used_paths.add(candidate)
225
+ return candidate
226
+
227
+
228
+ def unique_output_path(
229
+ record: CorpusRecord,
230
+ used_paths: set[Path],
231
+ raw_dir: Path,
232
+ ) -> Path:
233
+ slug = record_slug(record)
234
+ base_dir = raw_dir / record.source_group / record.source
235
+ candidate = base_dir / f"{slug}.md"
236
+ if candidate not in used_paths:
237
+ used_paths.add(candidate)
238
+ return candidate
239
+
240
+ suffix = record.content_hash.removeprefix("sha256:")[:8]
241
+ candidate = base_dir / f"{slug}-{suffix}.md"
242
+ counter = 2
243
+ while candidate in used_paths:
244
+ candidate = base_dir / f"{slug}-{suffix}-{counter}.md"
245
+ counter += 1
246
+ used_paths.add(candidate)
247
+ return candidate
248
+
249
+
250
+ def add_index_value(
251
+ index: dict[str, list[OriginalMarkdown]],
252
+ key: str,
253
+ item: OriginalMarkdown,
254
+ ) -> None:
255
+ if key:
256
+ index.setdefault(key, []).append(item)
257
+
258
+
259
+ def build_original_markdown_index(source: str) -> OriginalMarkdownIndex | None:
260
+ config = SOURCE_CONFIGS.get(source)
261
+ if not config:
262
+ return None
263
+
264
+ directory = Path(str(config.get("input_directory") or ""))
265
+ if not directory.exists() or not directory.is_dir():
266
+ return None
267
+
268
+ source_extension_manifest = load_source_extension_manifest(str(directory))
269
+ source_url_manifest = load_source_url_manifest(str(directory))
270
+ by_source_path: dict[str, OriginalMarkdown] = {}
271
+ by_url: dict[str, list[OriginalMarkdown]] = {}
272
+ by_content_hash: dict[str, list[OriginalMarkdown]] = {}
273
+ by_title: dict[str, list[OriginalMarkdown]] = {}
274
+
275
+ for path in sorted(directory.rglob("*")):
276
+ if not path.is_file() or path.suffix not in MARKDOWN_EXTENSIONS:
277
+ continue
278
+ source_path = path.relative_to(directory).as_posix()
279
+ if not should_include_file(source_path, config):
280
+ continue
281
+ raw_content = path.read_text(encoding="utf-8")
282
+ content = clean_document_content(raw_content)
283
+ item = OriginalMarkdown(
284
+ source_path=source_path,
285
+ path=path,
286
+ url=generate_url(
287
+ source_path,
288
+ config,
289
+ source_extension_manifest.get(source_path),
290
+ source_url_manifest.get(source_path),
291
+ ),
292
+ title=extract_title(content) or path.name,
293
+ content_hash=content_sha256(content),
294
+ content=content,
295
+ )
296
+ by_source_path[source_path] = item
297
+ add_index_value(by_url, item.url, item)
298
+ add_index_value(by_content_hash, item.content_hash, item)
299
+ add_index_value(by_title, item.title, item)
300
+
301
+ return OriginalMarkdownIndex(
302
+ by_source_path=by_source_path,
303
+ by_url=by_url,
304
+ by_content_hash=by_content_hash,
305
+ by_title=by_title,
306
+ )
307
+
308
+
309
+ def unique_match(items: list[OriginalMarkdown] | None) -> OriginalMarkdown | None:
310
+ return items[0] if items and len(items) == 1 else None
311
+
312
+
313
+ def find_original_markdown(
314
+ record: CorpusRecord,
315
+ indexes: dict[str, OriginalMarkdownIndex | None],
316
+ ) -> OriginalMarkdown | None:
317
+ if record.source_group != "docs":
318
+ return None
319
+ if record.source not in indexes:
320
+ indexes[record.source] = build_original_markdown_index(record.source)
321
+ index = indexes.get(record.source)
322
+ if index is None:
323
+ return None
324
+
325
+ if record.source_path and (match := index.by_source_path.get(record.source_path)):
326
+ return match
327
+ if match := unique_match(index.by_url.get(record.url)):
328
+ return match
329
+ if match := unique_match(index.by_content_hash.get(record.content_hash)):
330
+ return match
331
+ if match := unique_match(index.by_title.get(record.title)):
332
+ return match
333
+ return None
334
+
335
+
336
+ def manifest_row(
337
+ record: CorpusRecord,
338
+ *,
339
+ output_path: Path,
340
+ source_path: str = "",
341
+ original_path: str = "",
342
+ ) -> dict[str, Any]:
343
+ return {
344
+ "doc_id": record.doc_id,
345
+ "legacy_doc_id": record.legacy_doc_id,
346
+ "title": record.title,
347
+ "source": record.source,
348
+ "source_group": record.source_group,
349
+ "source_path": source_path or record.source_path,
350
+ "url": record.url,
351
+ "tokens": record.tokens,
352
+ "content_hash": record.content_hash,
353
+ "path": output_path.as_posix(),
354
+ "original_path": original_path,
355
+ }
356
+
357
+
358
+ def write_raw_markdown(
359
+ records: list[CorpusRecord],
360
+ *,
361
+ input_file: Path,
362
+ output_dir: Path,
363
+ ) -> list[dict[str, Any]]:
364
+ raw_dir = output_dir / RAW_DIR_NAME
365
+ if raw_dir.exists():
366
+ shutil.rmtree(raw_dir)
367
+ raw_dir.mkdir(parents=True, exist_ok=True)
368
+
369
+ used_paths: set[Path] = set()
370
+ original_indexes: dict[str, OriginalMarkdownIndex | None] = {}
371
+ fallback_counts: dict[str, int] = {}
372
+ manifest: list[dict[str, Any]] = []
373
+ for record in records:
374
+ original = find_original_markdown(record, original_indexes)
375
+ if original:
376
+ output_path = raw_dir / "docs" / record.source / safe_relative_path(
377
+ original.source_path
378
+ )
379
+ output_path = unique_path(output_path, used_paths, record.content_hash)
380
+ source_path = original.source_path
381
+ original_path = original.path.as_posix()
382
+ content = original.content
383
+ else:
384
+ output_path = unique_output_path(record, used_paths, raw_dir)
385
+ source_path = record.source_path
386
+ original_path = ""
387
+ content = record.content
388
+ if record.source_group == "docs":
389
+ fallback_counts[record.source] = fallback_counts.get(record.source, 0) + 1
390
+
391
+ output_path.parent.mkdir(parents=True, exist_ok=True)
392
+ markdown = markdown_with_frontmatter(
393
+ record,
394
+ generated_from=input_file,
395
+ content=content,
396
+ source_path=source_path,
397
+ original_path=original_path,
398
+ )
399
+ output_path.write_text(markdown, encoding="utf-8")
400
+ manifest.append(
401
+ manifest_row(
402
+ record,
403
+ output_path=output_path,
404
+ source_path=source_path,
405
+ original_path=original_path,
406
+ )
407
+ )
408
+ for source, count in sorted(fallback_counts.items()):
409
+ logger.warning(
410
+ "Generated %s docs KB page(s) from JSONL because original markdown was not found.",
411
+ count,
412
+ extra={"source": source},
413
+ )
414
+ return manifest
415
+
416
+
417
+ def iter_markdown_headings(path: Path) -> Iterable[dict[str, Any]]:
418
+ in_fence = False
419
+ fence_char = ""
420
+ for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
421
+ fence_match = FENCE_RE.match(line)
422
+ if fence_match:
423
+ current = fence_match.group(1)[0]
424
+ if not in_fence:
425
+ in_fence = True
426
+ fence_char = current
427
+ elif current == fence_char:
428
+ in_fence = False
429
+ fence_char = ""
430
+ continue
431
+ if in_fence:
432
+ continue
433
+ match = HEADING_RE.match(line)
434
+ if match:
435
+ yield {
436
+ "line": line_number,
437
+ "level": len(match.group(1)),
438
+ "heading": match.group(2).strip(),
439
+ }
440
+
441
+
442
+ def heading_path_for_line(headings: list[dict[str, Any]], line_number: int) -> str:
443
+ stack: list[dict[str, Any]] = []
444
+ for heading in headings:
445
+ if int(heading["line"]) > line_number:
446
+ break
447
+ level = int(heading["level"])
448
+ stack = [item for item in stack if int(item["level"]) < level]
449
+ stack.append(heading)
450
+ return " > ".join(str(item["heading"]) for item in stack)
451
+
452
+
453
+ def extract_symbols(text: str) -> set[str]:
454
+ symbols: set[str] = set()
455
+ for inline_match in INLINE_CODE_RE.finditer(text):
456
+ for match in SYMBOL_RE.finditer(inline_match.group(1)):
457
+ symbols.add(match.group(0))
458
+ for match in SYMBOL_RE.finditer(text):
459
+ symbols.add(match.group(0))
460
+ return {
461
+ symbol
462
+ for symbol in symbols
463
+ if len(symbol) >= 3 and not symbol.lower().startswith(("http", "www"))
464
+ }
465
+
466
+
467
+ def write_generated_indexes(manifest: list[dict[str, Any]], generated_dir: Path) -> None:
468
+ if generated_dir.exists():
469
+ shutil.rmtree(generated_dir)
470
+ generated_dir.mkdir(parents=True, exist_ok=True)
471
+
472
+ manifest_path = generated_dir / "corpus_manifest.jsonl"
473
+ with manifest_path.open("w", encoding="utf-8") as handle:
474
+ for item in manifest:
475
+ handle.write(json.dumps(item, ensure_ascii=False, sort_keys=True) + "\n")
476
+
477
+ heading_rows: list[dict[str, Any]] = []
478
+ symbol_rows: list[dict[str, str]] = []
479
+ for item in manifest:
480
+ path = Path(str(item["path"]))
481
+ text = path.read_text(encoding="utf-8")
482
+ headings = list(iter_markdown_headings(path))
483
+ for heading in headings:
484
+ heading_rows.append(
485
+ {
486
+ "doc_id": item["doc_id"],
487
+ "source": item["source"],
488
+ "title": item["title"],
489
+ "path": item["path"],
490
+ **heading,
491
+ }
492
+ )
493
+ for symbol in sorted(extract_symbols(text)):
494
+ first_line = 1
495
+ for line_number, line in enumerate(text.splitlines(), 1):
496
+ if symbol in line:
497
+ first_line = line_number
498
+ break
499
+ symbol_rows.append(
500
+ {
501
+ "symbol": symbol,
502
+ "source": str(item["source"]),
503
+ "title": str(item["title"]),
504
+ "path": str(item["path"]),
505
+ "heading": heading_path_for_line(headings, first_line),
506
+ "doc_id": str(item["doc_id"]),
507
+ }
508
+ )
509
+
510
+ with (generated_dir / "headings.jsonl").open("w", encoding="utf-8") as handle:
511
+ for row in heading_rows:
512
+ handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")
513
+
514
+ with (generated_dir / "symbols.tsv").open("w", encoding="utf-8", newline="") as handle:
515
+ writer = csv.DictWriter(
516
+ handle,
517
+ fieldnames=["symbol", "source", "title", "path", "heading", "doc_id"],
518
+ delimiter="\t",
519
+ )
520
+ writer.writeheader()
521
+ writer.writerows(symbol_rows)
522
+
523
+
524
+ def build_kb_artifacts(input_file: Path, output_dir: Path) -> dict[str, int]:
525
+ rows = load_jsonl(input_file)
526
+ records = [normalize_record(row) for row in rows if str(row.get("content") or "").strip()]
527
+ manifest = write_raw_markdown(records, input_file=input_file, output_dir=output_dir)
528
+ write_generated_indexes(manifest, output_dir / GENERATED_DIR_NAME)
529
+ return {"documents": len(records), "manifest_rows": len(manifest)}
530
+
531
+
532
+ def main() -> None:
533
+ parser = argparse.ArgumentParser(description="Build local KB markdown artifacts.")
534
+ parser.add_argument("--input-file", default=str(DEFAULT_INPUT_FILE))
535
+ parser.add_argument("--output-dir", default=str(DEFAULT_OUTPUT_DIR))
536
+ args = parser.parse_args()
537
+
538
+ summary = build_kb_artifacts(Path(args.input_file), Path(args.output_dir))
539
+ print(
540
+ "Built KB artifacts: "
541
+ f"{summary['documents']} documents, {summary['manifest_rows']} manifest rows"
542
+ )
543
+
544
+
545
+ if __name__ == "__main__":
546
+ main()
data/scraping_scripts/kb_agents_template.md ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # AI Tutor KB Instructions
2
+
3
+ ## Ground Truth
4
+
5
+ - `raw/` contains generated markdown mirrors of the normalized JSONL corpus. Treat these pages as the local source authority.
6
+ - `wiki/` contains navigation and synthesis pages. Use them to orient, not as final authority.
7
+ - `generated/` contains machine-generated indexes for manifests, headings, and symbols.
8
+
9
+ ## Standard Workflow
10
+
11
+ **First command rule.** Unless the user gives you a specific file path or a
12
+ known symbol/class name to find, your FIRST `run_kb_command` must be one of:
13
+
14
+ - `cat wiki/index.md` — for any "tell me about X" / "what is X" / comparison
15
+ / recipe question where X is a broad topic.
16
+ - `cat wiki/frameworks/<source>.md` or `cat wiki/courses/<source>.md` — when
17
+ the question is already scoped to one known source.
18
+
19
+ Do NOT start with `rg`, `grep`, `find`, or `ls` over `raw/` as the first
20
+ command. The wiki pages already give you the source map, so doing your own
21
+ discovery first wastes budget and pulls noisy hits into context.
22
+
23
+ After the wiki orientation step:
24
+
25
+ 1. Use `rg`, `find`, `sed`, `head`, and `cat` to verify and read raw pages.
26
+ 2. Use generated indexes for exact symbols, headings, source paths, and URLs.
27
+ 3. Read `raw/` pages to verify exact claims, code, and version-sensitive behavior.
28
+ 4. Mention uncertainty when the knowledge base does not cover the answer.
29
+
30
+ ## Exploration Budget
31
+
32
+ - Prefer 3-6 KB commands for normal answers. Stop browsing once you have one clearly relevant source page and enough detail to answer.
33
+ - Use more commands only when the user asks for exhaustive coverage, a comparison across sources, or debugging that requires multiple files.
34
+ - Do not read a long file from top to bottom. Scan headings first, then read only targeted ranges.
35
+ - Avoid broad searches that return thousands of lines. Add `-m 20` or narrow the path before searching `raw/`.
36
+
37
+ ## Efficient Command Patterns
38
+
39
+ - Orientation (always first for broad questions): `cat wiki/index.md`
40
+ - Source map: `cat wiki/courses/<source>.md` or `cat wiki/frameworks/<source>.md`
41
+ - Follow-up search (after orientation): `rg -n -m 20 "query terms" raw/<scoped-path>`
42
+ - Heading scan: `rg -n "^#{1,3} " raw/.../<page>.md`
43
+ - Targeted read: `sed -n 'START,ENDp' raw/.../<page>.md`
44
+ - Exact symbols: `rg -n -m 20 "ClassName|function_name" generated/symbols.tsv raw/`
45
+
46
+ ## Answering From KB Browsing
47
+
48
+ - Default to a concise synthesis rather than rewriting the markdown content verbatim.
49
+ - For "how do I build X?" questions, give the minimal architecture and a small runnable skeleton; expand only if the user asks for a full implementation.
50
+ - Cite the raw/source page inline for key claims. Do not cite every code comment or repeat the same citation after every sentence.
data/scraping_scripts/process_md_files.py CHANGED
@@ -30,12 +30,13 @@ To add, modify, or retire sources, update data/scraping_scripts/source_registry.
30
  """
31
 
32
  import argparse
 
33
  import json
34
  import logging
35
  import os
36
  import re
37
- import uuid
38
  from typing import Dict, List
 
39
 
40
  import tiktoken
41
 
@@ -262,6 +263,34 @@ def clean_document_content(content: str) -> str:
262
  return content.strip()
263
 
264
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
265
  def process_md_files(directory: str, config: Dict) -> List[Dict]:
266
  jsonl_data = []
267
  source_extension_manifest = load_source_extension_manifest(directory)
@@ -280,6 +309,14 @@ def process_md_files(directory: str, config: Dict) -> List[Dict]:
280
  cleaned_content = clean_document_content(content)
281
  title = extract_title(cleaned_content)
282
  token_count = num_tokens_from_string(cleaned_content, "cl100k_base")
 
 
 
 
 
 
 
 
283
 
284
  # Skip very small or extremely large files
285
  if token_count < 100 or token_count > 200_000:
@@ -290,18 +327,19 @@ def process_md_files(directory: str, config: Dict) -> List[Dict]:
290
 
291
  json_object = {
292
  "tokens": token_count,
293
- "doc_id": str(uuid.uuid4()),
294
- "name": (title if title else file),
295
- "url": generate_url(
296
- relative_path,
297
- config,
298
- source_extension_manifest.get(
299
- relative_path.replace("\\", "/")
300
- ),
301
- source_url_manifest.get(relative_path.replace("\\", "/")),
302
  ),
 
 
303
  "retrieve_doc": (token_count <= 8000),
304
  "source": config["source_name"],
 
 
305
  "content": cleaned_content,
306
  }
307
 
 
30
  """
31
 
32
  import argparse
33
+ import hashlib
34
  import json
35
  import logging
36
  import os
37
  import re
 
38
  from typing import Dict, List
39
+ from urllib.parse import urlparse
40
 
41
  import tiktoken
42
 
 
263
  return content.strip()
264
 
265
 
266
+ def content_sha256(content: str) -> str:
267
+ return "sha256:" + hashlib.sha256(content.encode("utf-8")).hexdigest()
268
+
269
+
270
+ def slugify_identifier(value: str) -> str:
271
+ value = os.path.splitext(value.replace("\\", "/"))[0]
272
+ value = value.strip().lower()
273
+ value = re.sub(r"[^a-z0-9]+", "-", value)
274
+ value = value.strip("-")
275
+ return value or "untitled"
276
+
277
+
278
+ def stable_doc_id(
279
+ *,
280
+ source: str,
281
+ source_path: str,
282
+ title: str | None,
283
+ url: str,
284
+ content_hash: str,
285
+ ) -> str:
286
+ if source_path:
287
+ basis = source_path
288
+ else:
289
+ parsed_path = urlparse(url).path.strip("/") if url else ""
290
+ basis = parsed_path or title or content_hash.removeprefix("sha256:")[:12]
291
+ return f"{source}:{slugify_identifier(basis)}"
292
+
293
+
294
  def process_md_files(directory: str, config: Dict) -> List[Dict]:
295
  jsonl_data = []
296
  source_extension_manifest = load_source_extension_manifest(directory)
 
309
  cleaned_content = clean_document_content(content)
310
  title = extract_title(cleaned_content)
311
  token_count = num_tokens_from_string(cleaned_content, "cl100k_base")
312
+ source_path = relative_path.replace("\\", "/")
313
+ url = generate_url(
314
+ relative_path,
315
+ config,
316
+ source_extension_manifest.get(source_path),
317
+ source_url_manifest.get(source_path),
318
+ )
319
+ hash_value = content_sha256(cleaned_content)
320
 
321
  # Skip very small or extremely large files
322
  if token_count < 100 or token_count > 200_000:
 
327
 
328
  json_object = {
329
  "tokens": token_count,
330
+ "doc_id": stable_doc_id(
331
+ source=config["source_name"],
332
+ source_path=source_path,
333
+ title=title,
334
+ url=url,
335
+ content_hash=hash_value,
 
 
 
336
  ),
337
+ "name": (title if title else file),
338
+ "url": url,
339
  "retrieve_doc": (token_count <= 8000),
340
  "source": config["source_name"],
341
+ "source_path": source_path,
342
+ "content_hash": hash_value,
343
  "content": cleaned_content,
344
  }
345
 
data/scraping_scripts/update_docs_workflow.py CHANGED
@@ -374,6 +374,36 @@ def create_vector_stores() -> None:
374
  logger.info("Successfully created vector stores")
375
 
376
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
377
  def upload_to_huggingface(upload_jsonl: bool = False) -> None:
378
  """Upload databases to HuggingFace."""
379
  logger.info("Uploading databases to HuggingFace")
@@ -432,6 +462,11 @@ def main():
432
  parser.add_argument(
433
  "--skip-vectors", action="store_true", help="Skip vector store creation"
434
  )
 
 
 
 
 
435
  parser.add_argument(
436
  "--skip-upload", action="store_true", help="Skip uploading to HuggingFace"
437
  )
@@ -458,6 +493,9 @@ def main():
458
  if not args.skip_process:
459
  process_markdown_files(args.sources)
460
 
 
 
 
461
  if not args.skip_context:
462
  add_context_to_nodes(not args.process_all_context)
463
 
 
374
  logger.info("Successfully created vector stores")
375
 
376
 
377
+ def build_kb_artifacts() -> None:
378
+ """Build generated markdown corpus, indexes, and refreshed wiki pages.
379
+
380
+ Notes for a from-scratch rebuild (e.g. after `rm -rf data/kb`):
381
+
382
+ - `build_kb_artifacts.py` rebuilds `kb/raw/` and `kb/generated/` from
383
+ `data/all_sources_data.jsonl` (which `ensure_required_files_exist`
384
+ downloaded from HuggingFace if it was missing). Course raw pages come
385
+ from the JSONL too, so missing Notion exports don't matter at runtime.
386
+ - `update_kb_wiki.py` is called with `--changed-only` here, which only
387
+ forces overwrites on `index.md`, `log.md`, and per-source pages. The
388
+ topic pages (`wiki/topics/*.md`) and the recipes/errors index pages
389
+ need `--seed-defaults`. The script auto-promotes to `seed_defaults=True`
390
+ when it sees an empty `wiki/` (see `wiki_is_empty` in update_kb_wiki.py),
391
+ so the workflow's `--changed-only` is safe for both fresh and
392
+ incremental rebuilds.
393
+ """
394
+ logger.info("Building KB artifacts")
395
+ result = run_module("data.scraping_scripts.build_kb_artifacts")
396
+ if result.returncode != 0:
397
+ logger.error("Error building KB artifacts - check output above")
398
+ sys.exit(1)
399
+
400
+ result = run_module("data.scraping_scripts.update_kb_wiki", "--changed-only")
401
+ if result.returncode != 0:
402
+ logger.error("Error updating KB wiki - check output above")
403
+ sys.exit(1)
404
+ logger.info("Successfully built KB artifacts")
405
+
406
+
407
  def upload_to_huggingface(upload_jsonl: bool = False) -> None:
408
  """Upload databases to HuggingFace."""
409
  logger.info("Uploading databases to HuggingFace")
 
462
  parser.add_argument(
463
  "--skip-vectors", action="store_true", help="Skip vector store creation"
464
  )
465
+ parser.add_argument(
466
+ "--skip-kb",
467
+ action="store_true",
468
+ help="Skip generated KB markdown/wiki artifact creation",
469
+ )
470
  parser.add_argument(
471
  "--skip-upload", action="store_true", help="Skip uploading to HuggingFace"
472
  )
 
493
  if not args.skip_process:
494
  process_markdown_files(args.sources)
495
 
496
+ if not args.skip_kb:
497
+ build_kb_artifacts()
498
+
499
  if not args.skip_context:
500
  add_context_to_nodes(not args.process_all_context)
501
 
data/scraping_scripts/update_kb_wiki.py ADDED
@@ -0,0 +1,302 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Seed and refresh the agent-facing KB wiki."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import argparse
6
+ import json
7
+ from collections import Counter, defaultdict
8
+ from datetime import UTC, datetime
9
+ from pathlib import Path
10
+ from typing import Any
11
+
12
+ DEFAULT_KB_DIR = Path("data/kb")
13
+ # Template file lives next to this script (inside `data/scraping_scripts/`) so
14
+ # it survives `rm -rf data/kb` and stays in git. Edit the template here when
15
+ # you want a permanent change to the KB agent guidance; the live
16
+ # `data/kb/AGENTS.md` is overwritten from it on every run of this script.
17
+ AGENTS_TEMPLATE_PATH = Path(__file__).resolve().parent / "kb_agents_template.md"
18
+ AUTO_START = "<!-- AUTO-GENERATED:START -->"
19
+ AUTO_END = "<!-- AUTO-GENERATED:END -->"
20
+ TOPIC_KEYWORDS: dict[str, tuple[str, ...]] = {
21
+ "rag": ("rag", "retrieval", "vector", "embedding", "chunk"),
22
+ "agents": ("agent", "tool", "langgraph", "workflow", "mcp"),
23
+ "retrieval": ("retrieval", "retriever", "bm25", "rerank", "search"),
24
+ "lora": ("lora", "qlora", "adapter", "peft", "loraconfig"),
25
+ "quantization": ("quantization", "quantized", "bitsandbytes", "nf4", "4-bit"),
26
+ "fine-tuning": ("fine-tuning", "finetuning", "sft", "trainer", "trl"),
27
+ "model-loading": ("from_pretrained", "automodel", "tokenizer", "device_map"),
28
+ "tool-calling": ("tool calling", "function calling", "tools", "schema"),
29
+ "evaluation": ("evaluation", "eval", "benchmark", "metrics"),
30
+ }
31
+
32
+
33
+ def load_manifest(kb_dir: Path) -> list[dict[str, Any]]:
34
+ manifest_path = kb_dir / "generated" / "corpus_manifest.jsonl"
35
+ if not manifest_path.exists():
36
+ return []
37
+ rows: list[dict[str, Any]] = []
38
+ with manifest_path.open("r", encoding="utf-8") as handle:
39
+ for line in handle:
40
+ if line.strip():
41
+ rows.append(json.loads(line))
42
+ return rows
43
+
44
+
45
+ def wiki_path(kb_dir: Path, *parts: str) -> Path:
46
+ return kb_dir / "wiki" / Path(*parts)
47
+
48
+
49
+ def write_page(path: Path, content: str, *, overwrite: bool) -> bool:
50
+ path.parent.mkdir(parents=True, exist_ok=True)
51
+ if path.exists() and not overwrite:
52
+ return False
53
+ path.write_text(content.rstrip() + "\n", encoding="utf-8")
54
+ return True
55
+
56
+
57
+ def shell_path(row: dict[str, Any], kb_dir: Path) -> str:
58
+ path = str(row.get("path") or "")
59
+ if not path:
60
+ return ""
61
+ try:
62
+ return Path(path).resolve().relative_to(kb_dir.resolve()).as_posix()
63
+ except ValueError:
64
+ pass
65
+ prefix = "data/kb/"
66
+ if path.startswith(prefix):
67
+ return path[len(prefix) :]
68
+ parts = Path(path).parts
69
+ if "raw" in parts:
70
+ return "/".join(parts[parts.index("raw") :])
71
+ return path
72
+
73
+
74
+ def generated_block(content: str) -> str:
75
+ return f"{AUTO_START}\n{content.rstrip()}\n{AUTO_END}"
76
+
77
+
78
+ def write_generated_section(path: Path, header: str, generated: str, *, overwrite: bool) -> None:
79
+ path.parent.mkdir(parents=True, exist_ok=True)
80
+ block = generated_block(generated)
81
+ if not path.exists() or overwrite:
82
+ path.write_text(f"{header.rstrip()}\n\n{block}\n", encoding="utf-8")
83
+ return
84
+
85
+ existing = path.read_text(encoding="utf-8")
86
+ start = existing.find(AUTO_START)
87
+ end = existing.find(AUTO_END, start)
88
+ if start != -1 and end != -1:
89
+ updated = existing[:start] + block + existing[end + len(AUTO_END) :]
90
+ path.write_text(updated.rstrip() + "\n", encoding="utf-8")
91
+ return
92
+ path.write_text(existing.rstrip() + "\n\n" + block + "\n", encoding="utf-8")
93
+
94
+
95
+ def top_sources(manifest: list[dict[str, Any]]) -> list[tuple[str, int]]:
96
+ counts = Counter(str(row.get("source") or "unknown") for row in manifest)
97
+ return sorted(counts.items(), key=lambda item: (-item[1], item[0]))
98
+
99
+
100
+ def seed_index(kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite: bool) -> None:
101
+ group_by_source = {
102
+ str(row.get("source") or "unknown"): str(row.get("source_group") or "docs")
103
+ for row in manifest
104
+ }
105
+ source_lines = []
106
+ for source, count in top_sources(manifest):
107
+ folder = "courses" if group_by_source.get(source) == "courses" else "frameworks"
108
+ source_lines.append(f"- {source}: `wiki/{folder}/{source}.md` - {count} corpus pages")
109
+ topic_lines = [f"- {topic}: `wiki/topics/{topic}.md`" for topic in TOPIC_KEYWORDS]
110
+ content = f"""# AI Tutor KB Index
111
+
112
+ This wiki is a navigation layer over the generated corpus markdown.
113
+
114
+ ## Sources
115
+
116
+ {chr(10).join(source_lines) if source_lines else "- No corpus sources indexed yet."}
117
+
118
+ ## Topics
119
+
120
+ {chr(10).join(topic_lines)}
121
+
122
+ ## Generated Indexes
123
+
124
+ - Corpus manifest: `generated/corpus_manifest.jsonl`
125
+ - Headings: `generated/headings.jsonl`
126
+ - Symbols: `generated/symbols.tsv`
127
+ """
128
+ write_page(wiki_path(kb_dir, "index.md"), content, overwrite=overwrite)
129
+
130
+
131
+ def seed_log(kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite: bool) -> None:
132
+ now = datetime.now(UTC).isoformat(timespec="seconds")
133
+ path = wiki_path(kb_dir, "log.md")
134
+ entry = f"""## {now}
135
+
136
+ - Generated or refreshed the KB wiki scaffold/source maps.
137
+ - Corpus documents indexed: {len(manifest)}
138
+ """
139
+ path.parent.mkdir(parents=True, exist_ok=True)
140
+ if not path.exists():
141
+ path.write_text(f"# KB Log\n\n{entry.rstrip()}\n", encoding="utf-8")
142
+ return
143
+ existing = path.read_text(encoding="utf-8").rstrip()
144
+ path.write_text(f"{existing}\n\n{entry.rstrip()}\n", encoding="utf-8")
145
+
146
+
147
+ def seed_source_pages(kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite: bool) -> None:
148
+ by_source: dict[str, list[dict[str, Any]]] = defaultdict(list)
149
+ for row in manifest:
150
+ by_source[str(row.get("source") or "unknown")].append(row)
151
+
152
+ for source, rows in sorted(by_source.items()):
153
+ rows = sorted(rows, key=lambda item: str(item.get("title") or ""))
154
+ sample = rows[:20]
155
+ page_links = [
156
+ f"- {row.get('title')}: `{shell_path(row, kb_dir)}`"
157
+ for row in sample
158
+ ]
159
+ group = str(rows[0].get("source_group") or "docs")
160
+ folder = "courses" if group == "courses" else "frameworks"
161
+ content = f"""# {source}
162
+
163
+ Use this page to orient inside the `{source}` corpus before reading source pages directly.
164
+
165
+ ## Corpus Pages
166
+
167
+ {chr(10).join(page_links) if page_links else "- No pages indexed."}
168
+
169
+ ## Verification
170
+
171
+ Treat generated corpus markdown as the source of truth for this source. Use exact search for symbols, error strings, and code identifiers.
172
+ """
173
+ write_page(wiki_path(kb_dir, folder, f"{source}.md"), content, overwrite=overwrite)
174
+
175
+
176
+ def matching_topic_rows(
177
+ manifest: list[dict[str, Any]],
178
+ keywords: tuple[str, ...],
179
+ *,
180
+ limit: int = 20,
181
+ ) -> list[dict[str, Any]]:
182
+ scored: list[tuple[int, dict[str, Any]]] = []
183
+ lowered = tuple(keyword.lower() for keyword in keywords)
184
+ for row in manifest:
185
+ title = str(row.get("title") or "")
186
+ path = Path(str(row.get("path") or ""))
187
+ text = title.lower()
188
+ if path.exists():
189
+ text += "\n" + path.read_text(encoding="utf-8").lower()[:20_000]
190
+ score = sum(text.count(keyword) for keyword in lowered)
191
+ if score:
192
+ scored.append((score, row))
193
+ scored.sort(key=lambda item: (-item[0], str(item[1].get("title") or "")))
194
+ return [row for _score, row in scored[:limit]]
195
+
196
+
197
+ def seed_topic_pages(kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite: bool) -> None:
198
+ for topic, keywords in TOPIC_KEYWORDS.items():
199
+ rows = matching_topic_rows(manifest, keywords)
200
+ links = [
201
+ f"- {row.get('title')}: `{shell_path(row, kb_dir)}`"
202
+ for row in rows
203
+ ]
204
+ header = f"""# {topic.replace("-", " ").title()}
205
+
206
+ Use this page as a starting map for questions involving: {", ".join(keywords)}.
207
+ """
208
+ generated = f"""## Starting Corpus Pages
209
+
210
+ {chr(10).join(links) if links else "- No strong corpus matches found yet."}
211
+
212
+ ## Agent Notes
213
+
214
+ - Start here for orientation.
215
+ - Verify exact factual claims in raw corpus markdown pages.
216
+ - Use `rg` over `raw/` and `generated/symbols.tsv` for API/class/function names.
217
+ """
218
+ write_generated_section(
219
+ wiki_path(kb_dir, "topics", f"{topic}.md"),
220
+ header,
221
+ generated,
222
+ overwrite=overwrite,
223
+ )
224
+
225
+
226
+ def seed_index_pages(kb_dir: Path, *, overwrite: bool) -> None:
227
+ for folder, title in (("recipes", "Recipes"), ("errors", "Errors")):
228
+ header = f"""# {title}
229
+
230
+ This index starts empty. Add pages here when repeated recipes or recurring errors emerge from the corpus.
231
+ """
232
+ write_generated_section(
233
+ wiki_path(kb_dir, folder, "index.md"),
234
+ header,
235
+ "- No generated entries yet.",
236
+ overwrite=overwrite,
237
+ )
238
+
239
+
240
+ def write_agents_md(kb_dir: Path) -> None:
241
+ # The canonical AGENTS.md content lives in `kb_agents_template.md` next to
242
+ # this script. Editing the template is the right way to change the agent's
243
+ # KB guidance for good — direct edits to `data/kb/AGENTS.md` would be
244
+ # clobbered on the next workflow run.
245
+ if not AGENTS_TEMPLATE_PATH.exists():
246
+ raise FileNotFoundError(
247
+ f"KB AGENTS.md template missing at {AGENTS_TEMPLATE_PATH}. "
248
+ "Restore it from git (it's tracked alongside this script)."
249
+ )
250
+ content = AGENTS_TEMPLATE_PATH.read_text(encoding="utf-8")
251
+ (kb_dir / "AGENTS.md").parent.mkdir(parents=True, exist_ok=True)
252
+ (kb_dir / "AGENTS.md").write_text(content, encoding="utf-8")
253
+
254
+
255
+ def wiki_is_empty(kb_dir: Path) -> bool:
256
+ """Return True when the wiki has no real content yet (fresh build).
257
+
258
+ Used by both this script and `update_docs_workflow.py` to decide when to
259
+ auto-promote to `seed_defaults=True` (which seeds topic pages and the
260
+ recipes/errors index pages, not just the source pages). Treats a missing
261
+ or empty `wiki/` directory, or one with no `.md` files, as empty.
262
+ """
263
+ wiki_dir = kb_dir / "wiki"
264
+ if not wiki_dir.exists():
265
+ return True
266
+ return not any(wiki_dir.rglob("*.md"))
267
+
268
+
269
+ def update_kb_wiki(kb_dir: Path, *, seed_defaults: bool, changed_only: bool) -> None:
270
+ manifest = load_manifest(kb_dir)
271
+ # Fresh-build safety net: if the wiki dir is empty (e.g. after a full
272
+ # `rm -rf data/kb`), promote to seed_defaults so topic/recipe/error pages
273
+ # are also written. Without this, only `index.md`, `log.md`, and source
274
+ # pages would be created and the wiki would feel half-empty to the agent.
275
+ if not seed_defaults and wiki_is_empty(kb_dir):
276
+ seed_defaults = True
277
+ overwrite = seed_defaults or changed_only
278
+ write_agents_md(kb_dir)
279
+ seed_index(kb_dir, manifest, overwrite=overwrite)
280
+ seed_log(kb_dir, manifest, overwrite=overwrite)
281
+ seed_source_pages(kb_dir, manifest, overwrite=overwrite)
282
+ seed_topic_pages(kb_dir, manifest, overwrite=seed_defaults)
283
+ seed_index_pages(kb_dir, overwrite=seed_defaults)
284
+
285
+
286
+ def main() -> None:
287
+ parser = argparse.ArgumentParser(description="Seed or refresh the KB wiki.")
288
+ parser.add_argument("--kb-dir", default=str(DEFAULT_KB_DIR))
289
+ parser.add_argument("--seed-defaults", action="store_true")
290
+ parser.add_argument("--changed-only", action="store_true")
291
+ args = parser.parse_args()
292
+
293
+ update_kb_wiki(
294
+ Path(args.kb_dir),
295
+ seed_defaults=bool(args.seed_defaults),
296
+ changed_only=bool(args.changed_only),
297
+ )
298
+ print(f"Updated KB wiki in {args.kb_dir}")
299
+
300
+
301
+ if __name__ == "__main__":
302
+ main()
data/scraping_scripts/upload_dbs_to_hf.py CHANGED
@@ -31,8 +31,9 @@ def upload_vector_db(repo_id: str = DEFAULT_REPO_ID) -> None:
31
  allow_patterns=[
32
  "chroma-db-all_sources/**",
33
  "all_sources_contextual_nodes.pkl",
 
34
  ],
35
- ignore_patterns=["*.jsonl", "*.py", "*.txt", "*.ipynb", "*.md", "*.pyc", "*.mdx"],
36
  )
37
 
38
 
 
31
  allow_patterns=[
32
  "chroma-db-all_sources/**",
33
  "all_sources_contextual_nodes.pkl",
34
+ "kb/**",
35
  ],
36
+ ignore_patterns=["*.py", "*.ipynb", "*.pyc", "__pycache__/**"],
37
  )
38
 
39
 
scripts/chat_service.py CHANGED
@@ -2,6 +2,7 @@ from __future__ import annotations
2
 
3
  import json
4
  import os
 
5
  from dataclasses import dataclass
6
  from functools import lru_cache
7
  from threading import Lock
@@ -10,7 +11,12 @@ from uuid import uuid4
10
 
11
  import logfire
12
  from langchain.agents import create_agent
13
- from langchain.agents.middleware import AgentMiddleware
 
 
 
 
 
14
  from langchain.tools import ToolRuntime, tool
15
  from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage, HumanMessage
16
  from langchain_openai import ChatOpenAI
@@ -18,6 +24,18 @@ from langgraph.checkpoint.memory import InMemorySaver
18
 
19
  from .chat_types import ChatEvent, ChatRequest, ChatTurn, SourceMatch
20
  from .chroma_rag import LocalChromaRetriever, format_tool_payload, parse_tool_payload
 
 
 
 
 
 
 
 
 
 
 
 
21
  from .prompts import build_system_prompt
22
  from .setup import (
23
  BM25_INDEX_PATH,
@@ -39,11 +57,36 @@ LEGACY_THOUGHTS_DETAILS_OPEN = "<details><summary>Gemini thoughts</summary>"
39
  LEGACY_THOUGHTS_DETAILS_OPEN_EXPANDED = "<details open><summary>Gemini thoughts</summary>"
40
  CHECKPOINTER = InMemorySaver()
41
  _RETRIEVER_INIT_LOCK = Lock()
 
 
 
 
42
 
43
 
44
  @dataclass(frozen=True)
45
  class AppContext:
46
  allowed_sources: tuple[str, ...]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
47
 
48
 
49
  @lru_cache(maxsize=1)
@@ -100,6 +143,68 @@ def retrieve_tutor_context(query: str, runtime: ToolRuntime[AppContext]) -> str:
100
  return format_tool_payload(query, results)
101
 
102
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
103
  def message_content_to_text(content: Any) -> str:
104
  if isinstance(content, str):
105
  return content
@@ -260,27 +365,114 @@ def sync_thread_with_history(
260
  return branched_thread_id
261
 
262
 
263
- def collect_updated_source_matches(
264
- matches_by_doc_id: dict[str, SourceMatch],
265
- payload: str,
266
- ) -> list[SourceMatch]:
267
- updated: list[SourceMatch] = []
268
  for match in parse_tool_payload(payload):
269
- source_match = SourceMatch(
270
- doc_id=match.doc_id,
271
- title=match.title,
272
- url=match.url,
273
- source_key=match.source,
274
- source_label=SOURCE_KEY_TO_LABEL.get(match.source, match.source),
275
- score=match.score,
276
- group="courses" if match.source in COURSE_SOURCE_KEYS else "docs",
 
 
277
  )
278
- existing = matches_by_doc_id.get(match.doc_id)
279
- if existing and existing.score >= source_match.score:
 
 
 
 
 
 
280
  continue
281
- matches_by_doc_id[match.doc_id] = source_match
282
- updated.append(source_match)
283
- return updated
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
284
 
285
 
286
  def normalize_model_name(model_name: str) -> str:
@@ -335,6 +527,9 @@ def format_tool_args(args: Any) -> str:
335
  query = str(args.get("query", "")).strip()
336
  if query:
337
  return query
 
 
 
338
  if args:
339
  return json.dumps(args, ensure_ascii=False, sort_keys=True)
340
  return ""
@@ -413,8 +608,27 @@ def build_agent(
413
  include_thoughts: bool = False,
414
  ):
415
  model = build_chat_model(model_name, include_thoughts=include_thoughts)
416
- tools: list[Any] = [retrieve_tutor_context]
417
- middleware: list[AgentMiddleware] = []
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
418
  enabled = set(enabled_tools)
419
  if is_google_genai_model(model_name):
420
  if "web_search" in enabled:
@@ -460,7 +674,7 @@ def effective_tool_names(
460
  model_name: str,
461
  enabled_tools: tuple[str, ...],
462
  ) -> tuple[str, ...]:
463
- names = ["retrieve_tutor_context"]
464
  enabled = set(enabled_tools)
465
  if is_google_genai_model(model_name):
466
  if "web_search" in enabled:
@@ -693,7 +907,9 @@ def extract_anthropic_source_matches(
693
 
694
  async def stream_chat(request: ChatRequest) -> AsyncIterator[ChatEvent]:
695
  normalized_history = normalize_history(request.history)
696
- matches_by_doc_id: dict[str, SourceMatch] = {}
 
 
697
  tool_calls_by_id: dict[str, dict[str, Any]] = {}
698
  answer_chunks: list[str] = []
699
  completed_answer = ""
@@ -717,260 +933,234 @@ async def stream_chat(request: ChatRequest) -> AsyncIterator[ChatEvent]:
717
  normalized_history,
718
  )
719
 
720
- yield ChatEvent(
721
- "thread_started",
722
- {"thread_id": active_thread_id},
723
- )
724
- yield ChatEvent(
725
- "message_started",
726
- {"message_id": message_id},
727
- )
728
 
729
- async for chunk in agent.astream(
730
- {"messages": [{"role": "user", "content": request.query}]},
731
- agent_run_config(request, active_thread_id, message_id),
732
- context=AppContext(allowed_sources=request.source_keys),
733
- stream_mode=["messages", "updates"],
734
- version="v2",
735
- ):
736
- if chunk["type"] == "messages":
737
- token, metadata = chunk["data"]
738
- if not isinstance(token, AIMessageChunk):
739
- continue
 
 
 
 
 
740
 
741
- step = str(metadata.get("langgraph_step", ""))
742
- if include_reasoning:
743
- thought_text = "\n\n".join(extract_thought_summaries(token.content))
744
- if thought_text:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
745
  yield ChatEvent(
746
- "reasoning_delta",
747
  {
748
  "message_id": message_id,
749
- "step": step,
750
- "text": thought_text,
 
 
751
  },
752
  )
753
 
754
- for tool_call in getattr(token, "tool_calls", []) or []:
755
- tool_call_id = str(tool_call.get("id") or uuid4().hex)
756
- tool_calls_by_id[tool_call_id] = tool_call
757
- yield ChatEvent(
758
- "tool_call_started",
759
- {
760
- "message_id": message_id,
761
- "call_id": tool_call_id,
762
- "tool_name": str(tool_call.get("name", "tool")),
763
- "args": tool_call.get("args"),
764
- "args_text": format_tool_args(tool_call.get("args")),
765
- },
766
- )
767
 
768
- text_delta = token.text or ""
769
- if not text_delta and token.content:
770
- text_delta = message_content_to_text(token.content)
771
- if text_delta:
772
- answer_chunks.append(text_delta)
773
- yield ChatEvent(
774
- "text_delta",
775
- {
776
- "message_id": message_id,
777
- "text": text_delta,
778
- },
779
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
780
 
781
- token_metadata = getattr(token, "response_metadata", None)
782
- new_queries = [
783
- q
784
- for q in extract_web_search_queries(token_metadata)
785
- if q not in google_search_queries
786
- ]
787
- new_grounding = extract_grounding_source_matches(
788
- token_metadata,
789
- matches_by_doc_id,
790
- )
791
- if (new_queries or new_grounding) and not google_search_call_id:
792
- google_search_call_id = uuid4().hex
793
- yield ChatEvent(
794
- "tool_call_started",
795
- {
796
- "message_id": message_id,
797
- "call_id": google_search_call_id,
798
- "tool_name": GOOGLE_SEARCH_TOOL_NAME,
799
- "args": {"query": "; ".join(new_queries) if new_queries else ""},
800
- "args_text": "; ".join(new_queries),
801
- },
802
- )
803
- if new_queries:
804
- google_search_queries.extend(new_queries)
805
- for source_match in new_grounding:
806
- google_search_match_count += 1
807
- yield ChatEvent(
808
- "source_match",
809
- {
810
- "message_id": message_id,
811
- "doc_id": source_match.doc_id,
812
- "title": source_match.title,
813
- "url": source_match.url,
814
- "source_key": source_match.source_key,
815
- "source_label": source_match.source_label,
816
- "score": source_match.score,
817
- "group": source_match.group,
818
- "call_id": google_search_call_id,
819
- },
820
- )
821
- continue
822
 
823
- if chunk["type"] != "updates":
824
- continue
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
825
 
826
- for step, update in chunk["data"].items():
827
- message = update["messages"][-1]
828
- if (
829
- getattr(message, "type", None) == "tool"
830
- and getattr(message, "name", "") == "retrieve_tutor_context"
831
- ):
832
- payload = message_content_to_text(message.content)
833
- tool_call_id = str(getattr(message, "tool_call_id", "") or uuid4().hex)
834
- for source_match in collect_updated_source_matches(
835
- matches_by_doc_id,
836
- payload,
837
- ):
 
838
  yield ChatEvent(
839
- "source_match",
840
  {
841
  "message_id": message_id,
842
- "doc_id": source_match.doc_id,
843
- "title": source_match.title,
844
- "url": source_match.url,
845
- "source_key": source_match.source_key,
846
- "source_label": source_match.source_label,
847
- "score": source_match.score,
848
- "group": source_match.group,
849
  "call_id": tool_call_id,
 
 
 
 
 
 
 
 
 
850
  },
851
  )
 
852
 
853
- tool_call = tool_calls_by_id.get(
854
- tool_call_id,
855
- {"name": getattr(message, "name", "tool"), "args": None},
856
- )
857
- yield ChatEvent(
858
- "tool_call_completed",
859
- {
860
- "message_id": message_id,
861
- "step": step,
862
- "call_id": tool_call_id,
863
- "tool_name": str(tool_call.get("name", getattr(message, "name", "tool"))),
864
- "args": tool_call.get("args"),
865
- "args_text": format_tool_args(tool_call.get("args")),
866
- "output_text": payload,
867
- },
868
- )
869
- continue
870
-
871
- if step != "model" or getattr(message, "type", None) != "ai":
872
- continue
873
-
874
- message_metadata = getattr(message, "response_metadata", None)
875
- new_queries = [
876
- q
877
- for q in extract_web_search_queries(message_metadata)
878
- if q not in google_search_queries
879
- ]
880
- new_grounding = extract_grounding_source_matches(
881
- message_metadata,
882
- matches_by_doc_id,
883
- )
884
- if (new_queries or new_grounding) and not google_search_call_id:
885
- google_search_call_id = uuid4().hex
886
- yield ChatEvent(
887
- "tool_call_started",
888
- {
889
- "message_id": message_id,
890
- "call_id": google_search_call_id,
891
- "tool_name": GOOGLE_SEARCH_TOOL_NAME,
892
- "args": {"query": "; ".join(new_queries) if new_queries else ""},
893
- "args_text": "; ".join(new_queries),
894
- },
895
- )
896
- if new_queries:
897
- google_search_queries.extend(new_queries)
898
- for source_match in new_grounding:
899
- google_search_match_count += 1
900
- yield ChatEvent(
901
- "source_match",
902
- {
903
- "message_id": message_id,
904
- "doc_id": source_match.doc_id,
905
- "title": source_match.title,
906
- "url": source_match.url,
907
- "source_key": source_match.source_key,
908
- "source_label": source_match.source_label,
909
- "score": source_match.score,
910
- "group": source_match.group,
911
- "call_id": google_search_call_id,
912
- },
913
- )
914
 
915
- if is_anthropic_model(request.model_name):
916
- anthropic_updates, anthropic_tool_uses = extract_anthropic_source_matches(
917
- message.content,
918
- matches_by_doc_id,
 
 
 
 
 
919
  )
920
- for tool_use_id, tool_use in anthropic_tool_uses.items():
921
- if tool_use_id in tool_calls_by_id:
922
- continue
923
- tool_calls_by_id[tool_use_id] = tool_use
924
  yield ChatEvent(
925
  "tool_call_started",
926
  {
927
  "message_id": message_id,
928
- "call_id": tool_use_id,
929
- "tool_name": tool_use["name"],
930
- "args": tool_use.get("args"),
931
- "args_text": format_tool_args(tool_use.get("args")),
 
 
 
 
932
  },
933
  )
934
- for tool_use_id, new_matches in anthropic_updates.items():
935
- for source_match in new_matches:
 
 
 
 
 
 
 
 
 
 
 
 
 
936
  yield ChatEvent(
937
- "source_match",
938
  {
939
  "message_id": message_id,
940
- "doc_id": source_match.doc_id,
941
- "title": source_match.title,
942
- "url": source_match.url,
943
- "source_key": source_match.source_key,
944
- "source_label": source_match.source_label,
945
- "score": source_match.score,
946
- "group": source_match.group,
947
  "call_id": tool_use_id,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
948
  },
949
  )
950
- if not tool_use_id:
951
- continue
952
- tool_call = tool_calls_by_id.get(tool_use_id, {})
953
- tool_name = str(tool_call.get("name") or "web_search")
954
- plural = "" if len(new_matches) == 1 else "s"
955
- output_text = (
956
- f"{tool_name} returned {len(new_matches)} result{plural}."
957
- )
958
- yield ChatEvent(
959
- "tool_call_completed",
960
- {
961
- "message_id": message_id,
962
- "step": step,
963
- "call_id": tool_use_id,
964
- "tool_name": tool_name,
965
- "args": tool_call.get("args"),
966
- "args_text": format_tool_args(tool_call.get("args")),
967
- "output_text": output_text,
968
- },
969
- )
970
 
971
- if getattr(message, "tool_calls", None):
972
- continue
973
- completed_answer = message_content_to_text(message.content)
 
 
974
 
975
  if google_search_call_id:
976
  joined_queries = "; ".join(google_search_queries)
@@ -994,6 +1184,16 @@ async def stream_chat(request: ChatRequest) -> AsyncIterator[ChatEvent]:
994
  )
995
 
996
  answer = "".join(answer_chunks).strip() or completed_answer.strip()
 
 
 
 
 
 
 
 
 
 
997
  yield ChatEvent(
998
  "message_completed",
999
  {
 
2
 
3
  import json
4
  import os
5
+ import shlex
6
  from dataclasses import dataclass
7
  from functools import lru_cache
8
  from threading import Lock
 
11
 
12
  import logfire
13
  from langchain.agents import create_agent
14
+ from langchain.agents.middleware import (
15
+ AgentMiddleware,
16
+ ClearToolUsesEdit,
17
+ ContextEditingMiddleware,
18
+ SummarizationMiddleware,
19
+ )
20
  from langchain.tools import ToolRuntime, tool
21
  from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage, HumanMessage
22
  from langchain_openai import ChatOpenAI
 
24
 
25
  from .chat_types import ChatEvent, ChatRequest, ChatTurn, SourceMatch
26
  from .chroma_rag import LocalChromaRetriever, format_tool_payload, parse_tool_payload
27
+ from .kb_shell import (
28
+ KbCommandError,
29
+ format_command_payload,
30
+ run_kb_command as execute_kb_command,
31
+ )
32
+ from .kb_manifest import (
33
+ extract_raw_paths,
34
+ parse_markdown_citations,
35
+ resolve_manifest_reference,
36
+ source_match_key,
37
+ source_match_payload,
38
+ )
39
  from .prompts import build_system_prompt
40
  from .setup import (
41
  BM25_INDEX_PATH,
 
57
  LEGACY_THOUGHTS_DETAILS_OPEN_EXPANDED = "<details open><summary>Gemini thoughts</summary>"
58
  CHECKPOINTER = InMemorySaver()
59
  _RETRIEVER_INIT_LOCK = Lock()
60
+ KB_TOOL_NAMES = ("run_kb_command",)
61
+ DEFAULT_KB_COMMAND_LIMIT = 20
62
+ _KB_COMMAND_COUNTS: dict[str, int] = {}
63
+ _KB_COMMAND_COUNT_LOCK = Lock()
64
 
65
 
66
  @dataclass(frozen=True)
67
  class AppContext:
68
  allowed_sources: tuple[str, ...]
69
+ kb_session_id: str = ""
70
+ kb_command_limit: int = DEFAULT_KB_COMMAND_LIMIT
71
+
72
+
73
+ def _claim_kb_command_budget(session_id: str, limit: int) -> tuple[bool, int]:
74
+ if not session_id:
75
+ return True, 1
76
+ with _KB_COMMAND_COUNT_LOCK:
77
+ used = _KB_COMMAND_COUNTS.get(session_id, 0)
78
+ if used >= limit:
79
+ return False, used
80
+ used += 1
81
+ _KB_COMMAND_COUNTS[session_id] = used
82
+ return True, used
83
+
84
+
85
+ def _clear_kb_command_budget(session_id: str) -> None:
86
+ if not session_id:
87
+ return
88
+ with _KB_COMMAND_COUNT_LOCK:
89
+ _KB_COMMAND_COUNTS.pop(session_id, None)
90
 
91
 
92
  @lru_cache(maxsize=1)
 
143
  return format_tool_payload(query, results)
144
 
145
 
146
+ RUN_KB_COMMAND_SCHEMA = {
147
+ "title": "run_kb_command",
148
+ "description": "Run a safe, read-only terminal-style command inside the local KB.",
149
+ "type": "object",
150
+ "properties": {
151
+ "command": {
152
+ "type": "string",
153
+ "description": (
154
+ "Single read-only command to run under data/kb. Supported commands: "
155
+ "rg, grep, find, ls, sed, head, cat, wc. Pipes, redirects, command "
156
+ "chaining, network commands, and writes are not allowed."
157
+ ),
158
+ },
159
+ "timeout_seconds": {
160
+ "type": "integer",
161
+ "description": "Command timeout in seconds, capped by the runtime.",
162
+ "default": 8,
163
+ },
164
+ "max_output_chars": {
165
+ "type": "integer",
166
+ "description": "Maximum stdout/stderr characters to return, capped by the runtime.",
167
+ "default": 40000,
168
+ },
169
+ },
170
+ "required": ["command"],
171
+ }
172
+
173
+
174
+ @tool(args_schema=RUN_KB_COMMAND_SCHEMA)
175
+ def run_kb_command(
176
+ command: str,
177
+ runtime: ToolRuntime[AppContext],
178
+ timeout_seconds: int = 8,
179
+ max_output_chars: int = 40000,
180
+ ) -> str:
181
+ """Run a safe, read-only terminal-style command inside the local KB."""
182
+ allowed, used = _claim_kb_command_budget(
183
+ runtime.context.kb_session_id,
184
+ runtime.context.kb_command_limit,
185
+ )
186
+ if not allowed:
187
+ return (
188
+ f"$ {command}\n"
189
+ "error: KB command budget exceeded for this turn "
190
+ f"({used}/{runtime.context.kb_command_limit}). "
191
+ "Use the evidence already collected to answer now."
192
+ )
193
+ try:
194
+ ensure_local_vector_db()
195
+ except Exception as exc: # pragma: no cover - diagnostic only
196
+ logfire.warn("KB artifact download/check failed.", error=str(exc))
197
+ try:
198
+ result = execute_kb_command(
199
+ command,
200
+ timeout_seconds=timeout_seconds,
201
+ max_output_chars=max_output_chars,
202
+ )
203
+ return format_command_payload(result)
204
+ except (KbCommandError, OSError) as exc:
205
+ return f"$ {command}\nerror: {exc}"
206
+
207
+
208
  def message_content_to_text(content: Any) -> str:
209
  if isinstance(content, str):
210
  return content
 
365
  return branched_thread_id
366
 
367
 
368
+ def collect_retrieval_source_matches(payload: str) -> list[SourceMatch]:
369
+ matches: list[SourceMatch] = []
 
 
 
370
  for match in parse_tool_payload(payload):
371
+ matches.append(
372
+ SourceMatch(
373
+ doc_id=match.doc_id,
374
+ title=match.title,
375
+ url=match.url,
376
+ source_key=match.source,
377
+ source_label=SOURCE_KEY_TO_LABEL.get(match.source, match.source),
378
+ score=match.score,
379
+ group="courses" if match.source in COURSE_SOURCE_KEYS else "docs",
380
+ )
381
  )
382
+ return matches
383
+
384
+
385
+ def _record_evidence(target: dict[str, SourceMatch], matches: list[SourceMatch]) -> None:
386
+ for match in matches:
387
+ key = source_match_key(match)
388
+ existing = target.get(key)
389
+ if existing and existing.score >= match.score:
390
  continue
391
+ target[key] = match
392
+
393
+
394
+ def _index_evidence(matches: dict[str, SourceMatch]) -> dict[str, SourceMatch]:
395
+ index: dict[str, SourceMatch] = {}
396
+ for match in matches.values():
397
+ for key in (match.doc_id, match.url, match.title.strip().lower()):
398
+ if key:
399
+ index[key] = match
400
+ return index
401
+
402
+
403
+ def resolve_answer_citations(
404
+ answer: str,
405
+ *,
406
+ retrieval_evidence: dict[str, SourceMatch],
407
+ shell_evidence: dict[str, SourceMatch],
408
+ web_evidence: dict[str, SourceMatch],
409
+ ) -> list[SourceMatch]:
410
+ evidence_indexes = [
411
+ _index_evidence(retrieval_evidence),
412
+ _index_evidence(shell_evidence),
413
+ _index_evidence(web_evidence),
414
+ ]
415
+ resolved: list[SourceMatch] = []
416
+ seen: set[str] = set()
417
+ for label, reference in parse_markdown_citations(answer):
418
+ candidates = [reference, label.strip().lower()]
419
+ match: SourceMatch | None = None
420
+ for index in evidence_indexes:
421
+ for candidate in candidates:
422
+ if candidate and candidate in index:
423
+ match = index[candidate]
424
+ break
425
+ if match:
426
+ break
427
+ if not match:
428
+ manifest_match = resolve_manifest_reference(reference, label=label)
429
+ if manifest_match:
430
+ shell_index = _index_evidence(shell_evidence)
431
+ if (
432
+ manifest_match.doc_id in shell_index
433
+ or manifest_match.url in shell_index
434
+ or manifest_match.title.strip().lower() in shell_index
435
+ ):
436
+ match = shell_index.get(manifest_match.doc_id) or shell_index.get(manifest_match.url) or shell_index.get(manifest_match.title.strip().lower())
437
+ if not match:
438
+ continue
439
+ key = source_match_key(match)
440
+ if key in seen:
441
+ continue
442
+ seen.add(key)
443
+ resolved.append(match)
444
+ return resolved
445
+
446
+
447
+ def extract_shell_source_matches(command: str, output_text: str) -> list[SourceMatch]:
448
+ raw_paths = extract_raw_paths(output_text)
449
+ try:
450
+ tokens = shlex.split(command)
451
+ except ValueError:
452
+ tokens = []
453
+ if tokens:
454
+ executable = tokens[0]
455
+ if executable == "cat":
456
+ raw_paths.extend(tokens[1:])
457
+ elif executable == "sed" and len(tokens) >= 4:
458
+ raw_paths.append(tokens[-1])
459
+ elif executable == "head":
460
+ start = 1
461
+ if len(tokens) > 2 and tokens[1] == "-n":
462
+ start = 3
463
+ raw_paths.extend(tokens[start:])
464
+ matches: list[SourceMatch] = []
465
+ seen: set[str] = set()
466
+ for path in raw_paths:
467
+ match = resolve_manifest_reference(path)
468
+ if not match:
469
+ continue
470
+ key = source_match_key(match)
471
+ if key in seen:
472
+ continue
473
+ seen.add(key)
474
+ matches.append(match)
475
+ return matches
476
 
477
 
478
  def normalize_model_name(model_name: str) -> str:
 
527
  query = str(args.get("query", "")).strip()
528
  if query:
529
  return query
530
+ command = str(args.get("command", "")).strip()
531
+ if command:
532
+ return command
533
  if args:
534
  return json.dumps(args, ensure_ascii=False, sort_keys=True)
535
  return ""
 
608
  include_thoughts: bool = False,
609
  ):
610
  model = build_chat_model(model_name, include_thoughts=include_thoughts)
611
+ tools: list[Any] = [
612
+ retrieve_tutor_context,
613
+ run_kb_command,
614
+ ]
615
+ middleware: list[AgentMiddleware] = [
616
+ ContextEditingMiddleware(
617
+ edits=[
618
+ ClearToolUsesEdit(
619
+ trigger=5_000,
620
+ keep=3,
621
+ placeholder="[tool output cleared to save context]",
622
+ )
623
+ ],
624
+ token_count_method="approximate",
625
+ ),
626
+ SummarizationMiddleware(
627
+ model=model,
628
+ trigger=("tokens", 30_000),
629
+ keep=("messages", 20),
630
+ ),
631
+ ]
632
  enabled = set(enabled_tools)
633
  if is_google_genai_model(model_name):
634
  if "web_search" in enabled:
 
674
  model_name: str,
675
  enabled_tools: tuple[str, ...],
676
  ) -> tuple[str, ...]:
677
+ names = ["retrieve_tutor_context", *KB_TOOL_NAMES]
678
  enabled = set(enabled_tools)
679
  if is_google_genai_model(model_name):
680
  if "web_search" in enabled:
 
907
 
908
  async def stream_chat(request: ChatRequest) -> AsyncIterator[ChatEvent]:
909
  normalized_history = normalize_history(request.history)
910
+ retrieval_evidence: dict[str, SourceMatch] = {}
911
+ shell_evidence: dict[str, SourceMatch] = {}
912
+ web_evidence: dict[str, SourceMatch] = {}
913
  tool_calls_by_id: dict[str, dict[str, Any]] = {}
914
  answer_chunks: list[str] = []
915
  completed_answer = ""
 
933
  normalized_history,
934
  )
935
 
936
+ yield ChatEvent("thread_started", {"thread_id": active_thread_id})
937
+ yield ChatEvent("message_started", {"message_id": message_id})
 
 
 
 
 
 
938
 
939
+ try:
940
+ async for chunk in agent.astream(
941
+ {"messages": [{"role": "user", "content": request.query}]},
942
+ agent_run_config(request, active_thread_id, message_id),
943
+ context=AppContext(
944
+ allowed_sources=request.source_keys,
945
+ kb_session_id=message_id,
946
+ kb_command_limit=DEFAULT_KB_COMMAND_LIMIT,
947
+ ),
948
+ stream_mode=["messages", "updates"],
949
+ version="v2",
950
+ ):
951
+ if chunk["type"] == "messages":
952
+ token, metadata = chunk["data"]
953
+ if not isinstance(token, AIMessageChunk):
954
+ continue
955
 
956
+ step = str(metadata.get("langgraph_step", ""))
957
+ if include_reasoning:
958
+ thought_text = "\n\n".join(
959
+ extract_thought_summaries(token.content)
960
+ )
961
+ if thought_text:
962
+ yield ChatEvent(
963
+ "reasoning_delta",
964
+ {
965
+ "message_id": message_id,
966
+ "step": step,
967
+ "text": thought_text,
968
+ },
969
+ )
970
+
971
+ for tool_call in getattr(token, "tool_calls", []) or []:
972
+ tool_call_id = str(tool_call.get("id") or uuid4().hex)
973
+ tool_calls_by_id[tool_call_id] = tool_call
974
  yield ChatEvent(
975
+ "tool_call_started",
976
  {
977
  "message_id": message_id,
978
+ "call_id": tool_call_id,
979
+ "tool_name": str(tool_call.get("name", "tool")),
980
+ "args": tool_call.get("args"),
981
+ "args_text": format_tool_args(tool_call.get("args")),
982
  },
983
  )
984
 
985
+ text_delta = token.text or ""
986
+ if not text_delta and token.content:
987
+ text_delta = message_content_to_text(token.content)
988
+ if text_delta:
989
+ answer_chunks.append(text_delta)
990
+ yield ChatEvent(
991
+ "text_delta",
992
+ {
993
+ "message_id": message_id,
994
+ "text": text_delta,
995
+ },
996
+ )
 
997
 
998
+ token_metadata = getattr(token, "response_metadata", None)
999
+ new_queries = [
1000
+ q
1001
+ for q in extract_web_search_queries(token_metadata)
1002
+ if q not in google_search_queries
1003
+ ]
1004
+ new_grounding = extract_grounding_source_matches(
1005
+ token_metadata,
1006
+ web_evidence,
 
 
1007
  )
1008
+ if (new_queries or new_grounding) and not google_search_call_id:
1009
+ google_search_call_id = uuid4().hex
1010
+ yield ChatEvent(
1011
+ "tool_call_started",
1012
+ {
1013
+ "message_id": message_id,
1014
+ "call_id": google_search_call_id,
1015
+ "tool_name": GOOGLE_SEARCH_TOOL_NAME,
1016
+ "args": {
1017
+ "query": "; ".join(new_queries)
1018
+ if new_queries
1019
+ else ""
1020
+ },
1021
+ "args_text": "; ".join(new_queries),
1022
+ },
1023
+ )
1024
+ if new_queries:
1025
+ google_search_queries.extend(new_queries)
1026
+ google_search_match_count += len(new_grounding)
1027
+ continue
1028
 
1029
+ if chunk["type"] != "updates":
1030
+ continue
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1031
 
1032
+ for step, update in chunk["data"].items():
1033
+ if not isinstance(update, dict):
1034
+ continue
1035
+ messages = update.get("messages")
1036
+ if not messages:
1037
+ continue
1038
+ message = messages[-1]
1039
+ if getattr(message, "type", None) == "tool":
1040
+ payload = message_content_to_text(message.content)
1041
+ tool_call_id = str(
1042
+ getattr(message, "tool_call_id", "") or uuid4().hex
1043
+ )
1044
+ tool_name = str(getattr(message, "name", "tool"))
1045
+ if tool_name == "retrieve_tutor_context":
1046
+ _record_evidence(
1047
+ retrieval_evidence,
1048
+ collect_retrieval_source_matches(payload),
1049
+ )
1050
 
1051
+ tool_call = tool_calls_by_id.get(
1052
+ tool_call_id,
1053
+ {"name": getattr(message, "name", "tool"), "args": None},
1054
+ )
1055
+ if tool_name == "run_kb_command":
1056
+ args = tool_call.get("args")
1057
+ command = str(
1058
+ args.get("command") if isinstance(args, dict) else ""
1059
+ )
1060
+ _record_evidence(
1061
+ shell_evidence,
1062
+ extract_shell_source_matches(command, payload),
1063
+ )
1064
  yield ChatEvent(
1065
+ "tool_call_completed",
1066
  {
1067
  "message_id": message_id,
1068
+ "step": step,
 
 
 
 
 
 
1069
  "call_id": tool_call_id,
1070
+ "tool_name": str(
1071
+ tool_call.get(
1072
+ "name",
1073
+ getattr(message, "name", "tool"),
1074
+ )
1075
+ ),
1076
+ "args": tool_call.get("args"),
1077
+ "args_text": format_tool_args(tool_call.get("args")),
1078
+ "output_text": payload,
1079
  },
1080
  )
1081
+ continue
1082
 
1083
+ if step != "model" or getattr(message, "type", None) != "ai":
1084
+ continue
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1085
 
1086
+ message_metadata = getattr(message, "response_metadata", None)
1087
+ new_queries = [
1088
+ q
1089
+ for q in extract_web_search_queries(message_metadata)
1090
+ if q not in google_search_queries
1091
+ ]
1092
+ new_grounding = extract_grounding_source_matches(
1093
+ message_metadata,
1094
+ web_evidence,
1095
  )
1096
+ if (new_queries or new_grounding) and not google_search_call_id:
1097
+ google_search_call_id = uuid4().hex
 
 
1098
  yield ChatEvent(
1099
  "tool_call_started",
1100
  {
1101
  "message_id": message_id,
1102
+ "call_id": google_search_call_id,
1103
+ "tool_name": GOOGLE_SEARCH_TOOL_NAME,
1104
+ "args": {
1105
+ "query": "; ".join(new_queries)
1106
+ if new_queries
1107
+ else ""
1108
+ },
1109
+ "args_text": "; ".join(new_queries),
1110
  },
1111
  )
1112
+ if new_queries:
1113
+ google_search_queries.extend(new_queries)
1114
+ google_search_match_count += len(new_grounding)
1115
+
1116
+ if is_anthropic_model(request.model_name):
1117
+ anthropic_updates, anthropic_tool_uses = (
1118
+ extract_anthropic_source_matches(
1119
+ message.content,
1120
+ web_evidence,
1121
+ )
1122
+ )
1123
+ for tool_use_id, tool_use in anthropic_tool_uses.items():
1124
+ if tool_use_id in tool_calls_by_id:
1125
+ continue
1126
+ tool_calls_by_id[tool_use_id] = tool_use
1127
  yield ChatEvent(
1128
+ "tool_call_started",
1129
  {
1130
  "message_id": message_id,
 
 
 
 
 
 
 
1131
  "call_id": tool_use_id,
1132
+ "tool_name": tool_use["name"],
1133
+ "args": tool_use.get("args"),
1134
+ "args_text": format_tool_args(tool_use.get("args")),
1135
+ },
1136
+ )
1137
+ for tool_use_id, new_matches in anthropic_updates.items():
1138
+ if not tool_use_id:
1139
+ continue
1140
+ tool_call = tool_calls_by_id.get(tool_use_id, {})
1141
+ tool_name = str(tool_call.get("name") or "web_search")
1142
+ plural = "" if len(new_matches) == 1 else "s"
1143
+ output_text = (
1144
+ f"{tool_name} returned {len(new_matches)} result{plural}."
1145
+ )
1146
+ yield ChatEvent(
1147
+ "tool_call_completed",
1148
+ {
1149
+ "message_id": message_id,
1150
+ "step": step,
1151
+ "call_id": tool_use_id,
1152
+ "tool_name": tool_name,
1153
+ "args": tool_call.get("args"),
1154
+ "args_text": format_tool_args(tool_call.get("args")),
1155
+ "output_text": output_text,
1156
  },
1157
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1158
 
1159
+ if getattr(message, "tool_calls", None):
1160
+ continue
1161
+ completed_answer = message_content_to_text(message.content)
1162
+ finally:
1163
+ _clear_kb_command_budget(message_id)
1164
 
1165
  if google_search_call_id:
1166
  joined_queries = "; ".join(google_search_queries)
 
1184
  )
1185
 
1186
  answer = "".join(answer_chunks).strip() or completed_answer.strip()
1187
+ for source_match in resolve_answer_citations(
1188
+ answer,
1189
+ retrieval_evidence=retrieval_evidence,
1190
+ shell_evidence=shell_evidence,
1191
+ web_evidence=web_evidence,
1192
+ ):
1193
+ yield ChatEvent(
1194
+ "source_match",
1195
+ source_match_payload(source_match, message_id=message_id),
1196
+ )
1197
  yield ChatEvent(
1198
  "message_completed",
1199
  {
scripts/gradio_presenter.py CHANGED
@@ -71,9 +71,15 @@ def upsert_activity_event(
71
  def format_tool_args(args: Any, args_text: str = "") -> str:
72
  query = ""
73
  url = ""
 
74
  if isinstance(args, dict):
75
  query = str(args.get("query", "")).strip()
76
  url = str(args.get("url", "")).strip()
 
 
 
 
 
77
 
78
  if not query and args_text.strip():
79
  query = args_text.strip()
@@ -115,12 +121,25 @@ def summarize_tool_result(event: ChatEvent) -> str:
115
  match_count = len(matches)
116
 
117
  if tool_name == "retrieve_tutor_context":
 
 
 
 
 
 
 
 
118
  if not matches:
119
  return "_No matching sources found in the selected sources._"
120
  ordered_sources: list[str] = []
121
  seen_sources: set[str] = set()
122
  for match in matches:
123
- source_label = str(match.get("source_label", match.get("source_key", "unknown")))
 
 
 
 
 
124
  if source_label in seen_sources:
125
  continue
126
  seen_sources.add(source_label)
@@ -129,6 +148,23 @@ def summarize_tool_result(event: ChatEvent) -> str:
129
  match_label = "match" if match_count == 1 else "matches"
130
  return f"_Found {match_count} {match_label} from {source_summary}._"
131
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
132
  output_text = str(event.data.get("output_text", "")).strip()
133
  if output_text:
134
  return f"_{output_text}_"
 
71
  def format_tool_args(args: Any, args_text: str = "") -> str:
72
  query = ""
73
  url = ""
74
+ command = ""
75
  if isinstance(args, dict):
76
  query = str(args.get("query", "")).strip()
77
  url = str(args.get("url", "")).strip()
78
+ command = str(args.get("command", "")).strip()
79
+
80
+ if command:
81
+ trimmed = command[:157] + "..." if len(command) > 160 else command
82
+ return f"Command: `{trimmed}`"
83
 
84
  if not query and args_text.strip():
85
  query = args_text.strip()
 
121
  match_count = len(matches)
122
 
123
  if tool_name == "retrieve_tutor_context":
124
+ if not matches:
125
+ try:
126
+ payload = json.loads(str(event.data.get("output_text", "")))
127
+ matches = payload.get("matches", []) if isinstance(payload, dict) else []
128
+ match_count = len(matches)
129
+ except json.JSONDecodeError:
130
+ matches = []
131
+ match_count = 0
132
  if not matches:
133
  return "_No matching sources found in the selected sources._"
134
  ordered_sources: list[str] = []
135
  seen_sources: set[str] = set()
136
  for match in matches:
137
+ source_label = str(
138
+ match.get("source_label")
139
+ or match.get("source_key")
140
+ or match.get("source")
141
+ or "unknown"
142
+ )
143
  if source_label in seen_sources:
144
  continue
145
  seen_sources.add(source_label)
 
148
  match_label = "match" if match_count == 1 else "matches"
149
  return f"_Found {match_count} {match_label} from {source_summary}._"
150
 
151
+ if tool_name == "run_kb_command":
152
+ output_text = str(event.data.get("output_text", "")).strip()
153
+ command = ""
154
+ exit_code = ""
155
+ for line in output_text.splitlines():
156
+ if line.startswith("$ "):
157
+ command = line[2:].strip()
158
+ elif line.startswith("exit_code:"):
159
+ exit_code = line.partition(":")[2].strip()
160
+ if "error:" in output_text:
161
+ return f"_KB command failed: {output_text.split('error:', 1)[1].strip()[:160]}_"
162
+ if command and exit_code:
163
+ return f"_Ran `{command}` with exit code {exit_code}._"
164
+ if command:
165
+ return f"_Ran `{command}`._"
166
+ return "_Ran KB command._"
167
+
168
  output_text = str(event.data.get("output_text", "")).strip()
169
  if output_text:
170
  return f"_{output_text}_"
scripts/kb_manifest.py ADDED
@@ -0,0 +1,172 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import json
4
+ import re
5
+ from dataclasses import dataclass
6
+ from functools import lru_cache
7
+ from pathlib import Path
8
+ from typing import Any
9
+
10
+ from .chat_types import SourceMatch
11
+ from .setup import COURSE_SOURCE_KEYS, KB_DIR, SOURCE_KEY_TO_LABEL
12
+
13
+ KB_DOC_SCHEME_RE = re.compile(r"^kb://doc/(?P<doc_id>[^)\]\s]+)$")
14
+ RAW_PATH_RE = re.compile(r"(?:data/kb/)?raw/[^\s)\]>,:]+?\.(?:mdx|md)")
15
+
16
+
17
+ @dataclass(frozen=True, slots=True)
18
+ class KbManifestEntry:
19
+ doc_id: str
20
+ title: str
21
+ url: str
22
+ source: str
23
+ source_group: str
24
+ path: str
25
+
26
+
27
+ def _normalize_path(value: str) -> str:
28
+ value = value.strip().strip("<>`'\"")
29
+ if value.startswith("./"):
30
+ value = value[2:]
31
+ if value.startswith(f"{KB_DIR}/"):
32
+ return value
33
+ if value.startswith("raw/"):
34
+ return f"{KB_DIR}/{value}"
35
+ return value
36
+
37
+
38
+ @lru_cache(maxsize=4)
39
+ def load_manifest_entries(kb_dir: str = KB_DIR) -> tuple[KbManifestEntry, ...]:
40
+ manifest_path = Path(kb_dir) / "generated" / "corpus_manifest.jsonl"
41
+ if not manifest_path.exists():
42
+ return ()
43
+
44
+ entries: list[KbManifestEntry] = []
45
+ with manifest_path.open("r", encoding="utf-8") as handle:
46
+ for line in handle:
47
+ if not line.strip():
48
+ continue
49
+ row = json.loads(line)
50
+ entries.append(
51
+ KbManifestEntry(
52
+ doc_id=str(row.get("doc_id") or ""),
53
+ title=str(row.get("title") or ""),
54
+ url=str(row.get("url") or ""),
55
+ source=str(row.get("source") or ""),
56
+ source_group=str(row.get("source_group") or ""),
57
+ path=str(row.get("path") or ""),
58
+ )
59
+ )
60
+ return tuple(entries)
61
+
62
+
63
+ def manifest_indexes(
64
+ kb_dir: str = KB_DIR,
65
+ ) -> tuple[dict[str, KbManifestEntry], dict[str, KbManifestEntry], dict[str, KbManifestEntry], dict[str, KbManifestEntry]]:
66
+ by_doc_id: dict[str, KbManifestEntry] = {}
67
+ by_url: dict[str, KbManifestEntry] = {}
68
+ by_path: dict[str, KbManifestEntry] = {}
69
+ by_title: dict[str, KbManifestEntry] = {}
70
+ for entry in load_manifest_entries(kb_dir):
71
+ if entry.doc_id:
72
+ by_doc_id[entry.doc_id] = entry
73
+ if entry.url:
74
+ by_url[entry.url] = entry
75
+ if entry.path:
76
+ by_path[_normalize_path(entry.path)] = entry
77
+ if entry.path.startswith(f"{KB_DIR}/"):
78
+ by_path[entry.path[len(KB_DIR) + 1 :]] = entry
79
+ if entry.title:
80
+ by_title[entry.title.strip().lower()] = entry
81
+ return by_doc_id, by_url, by_path, by_title
82
+
83
+
84
+ def source_match_from_manifest(entry: KbManifestEntry, *, score: float = 1.0) -> SourceMatch:
85
+ return SourceMatch(
86
+ doc_id=entry.doc_id,
87
+ title=entry.title,
88
+ url=entry.url or entry.path,
89
+ source_key=entry.source,
90
+ source_label=SOURCE_KEY_TO_LABEL.get(entry.source, entry.source),
91
+ score=score,
92
+ group="courses" if entry.source in COURSE_SOURCE_KEYS else entry.source_group or "docs",
93
+ )
94
+
95
+
96
+ def resolve_manifest_reference(
97
+ reference: str,
98
+ *,
99
+ label: str = "",
100
+ kb_dir: str = KB_DIR,
101
+ ) -> SourceMatch | None:
102
+ value = reference.strip()
103
+ by_doc_id, by_url, by_path, by_title = manifest_indexes(kb_dir)
104
+
105
+ if match := KB_DOC_SCHEME_RE.match(value):
106
+ entry = by_doc_id.get(match.group("doc_id"))
107
+ return source_match_from_manifest(entry) if entry else None
108
+
109
+ entry = by_url.get(value)
110
+ if entry:
111
+ return source_match_from_manifest(entry)
112
+
113
+ entry = by_path.get(_normalize_path(value))
114
+ if entry:
115
+ return source_match_from_manifest(entry)
116
+
117
+ if label:
118
+ entry = by_title.get(label.strip().lower())
119
+ if entry:
120
+ return source_match_from_manifest(entry)
121
+ return None
122
+
123
+
124
+ def extract_raw_paths(text: str) -> list[str]:
125
+ paths: list[str] = []
126
+ seen: set[str] = set()
127
+ for match in RAW_PATH_RE.finditer(text):
128
+ value = match.group(0).rstrip(":.,;")
129
+ key = _normalize_path(value)
130
+ if key in seen:
131
+ continue
132
+ seen.add(key)
133
+ paths.append(value)
134
+ return paths
135
+
136
+
137
+ def parse_markdown_citations(text: str) -> list[tuple[str, str]]:
138
+ citations: list[tuple[str, str]] = []
139
+ link_re = re.compile(r"\[([^\]]+)\]\(([^)]+)\)")
140
+ for match in link_re.finditer(text):
141
+ citations.append((match.group(1).strip(), match.group(2).strip()))
142
+
143
+ linked_refs = {ref for _label, ref in citations}
144
+ bare_re = re.compile(r"(?<!\()(?P<ref>https?://[^\s<>()]+|kb://doc/[^\s<>()]+)")
145
+ for match in bare_re.finditer(text):
146
+ ref = match.group("ref").rstrip(".,;")
147
+ if ref not in linked_refs:
148
+ citations.append(("", ref))
149
+ for path in extract_raw_paths(text):
150
+ if path not in linked_refs:
151
+ citations.append(("", path))
152
+ return citations
153
+
154
+
155
+ def source_match_key(match: SourceMatch) -> str:
156
+ return match.doc_id or match.url or match.title
157
+
158
+
159
+ def source_match_payload(match: SourceMatch, *, message_id: str, call_id: str = "") -> dict[str, Any]:
160
+ payload: dict[str, Any] = {
161
+ "message_id": message_id,
162
+ "doc_id": match.doc_id,
163
+ "title": match.title,
164
+ "url": match.url,
165
+ "source_key": match.source_key,
166
+ "source_label": match.source_label,
167
+ "score": match.score,
168
+ "group": match.group,
169
+ }
170
+ if call_id:
171
+ payload["call_id"] = call_id
172
+ return payload
scripts/kb_shell.py ADDED
@@ -0,0 +1,394 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import os
4
+ import re
5
+ import shlex
6
+ import shutil
7
+ import subprocess
8
+ from dataclasses import dataclass
9
+ from pathlib import Path
10
+
11
+ DEFAULT_KB_DIR = Path(os.getenv("AI_TUTOR_KB_DIR", "data/kb"))
12
+ DEFAULT_TIMEOUT_SECONDS = 8
13
+ DEFAULT_MAX_OUTPUT_CHARS = 40_000
14
+ SUPPORTED_COMMANDS = frozenset({"rg", "grep", "find", "ls", "sed", "head", "cat", "wc"})
15
+ SHELL_TOKENS = frozenset({"|", "||", "&", "&&", ";", ">", ">>", "<", "<<", "`"})
16
+ RG_VALUE_OPTIONS = frozenset(
17
+ {
18
+ "-g",
19
+ "--glob",
20
+ "-t",
21
+ "--type",
22
+ "-T",
23
+ "--type-not",
24
+ "-m",
25
+ "--max-count",
26
+ }
27
+ )
28
+ RG_FLAG_OPTIONS = frozenset(
29
+ {
30
+ "-n",
31
+ "--line-number",
32
+ "-i",
33
+ "--ignore-case",
34
+ "-S",
35
+ "--smart-case",
36
+ "-F",
37
+ "--fixed-strings",
38
+ "-w",
39
+ "--word-regexp",
40
+ "--hidden",
41
+ "--no-ignore",
42
+ }
43
+ )
44
+ GREP_FLAG_OPTIONS = frozenset({"-i", "--ignore-case", "-w", "--word-regexp", "-F", "-E"})
45
+ LS_FLAG_OPTIONS = frozenset({"-1", "-a", "-l", "-la", "-al"})
46
+ WC_FLAG_OPTIONS = frozenset({"-l", "-w", "-c", "-m"})
47
+
48
+
49
+ @dataclass(frozen=True, slots=True)
50
+ class KbCommandResult:
51
+ command: str
52
+ argv: list[str]
53
+ cwd: str
54
+ exit_code: int
55
+ stdout: str
56
+ stderr: str
57
+ timed_out: bool
58
+ truncated: bool
59
+
60
+
61
+ class KbCommandError(ValueError):
62
+ """Raised when a KB command violates the read-only command policy."""
63
+
64
+
65
+ def _resolve_root(root: Path | None) -> Path:
66
+ resolved = (root or DEFAULT_KB_DIR).resolve()
67
+ if not resolved.exists():
68
+ raise KbCommandError(f"KB root does not exist: {resolved}")
69
+ if not resolved.is_dir():
70
+ raise KbCommandError(f"KB root is not a directory: {resolved}")
71
+ return resolved
72
+
73
+
74
+ def _resolve_inside_root(value: str, root: Path, *, must_exist: bool = True) -> Path:
75
+ path = Path(value)
76
+ if path.is_absolute():
77
+ resolved = path.resolve()
78
+ else:
79
+ resolved = (root / value).resolve()
80
+ if resolved != root and root not in resolved.parents:
81
+ raise KbCommandError("KB command paths must stay inside data/kb")
82
+ if must_exist and not resolved.exists():
83
+ raise KbCommandError(f"KB command path does not exist: {value}")
84
+ return resolved
85
+
86
+
87
+ def _relative_path_arg(value: str, root: Path) -> str:
88
+ resolved = _resolve_inside_root(value, root)
89
+ return "." if resolved == root else resolved.relative_to(root).as_posix()
90
+
91
+
92
+ def _reject_shell_syntax(command: str, tokens: list[str]) -> None:
93
+ if "\n" in command or "\r" in command:
94
+ raise KbCommandError("KB command must be a single command, not a script")
95
+ if "$(" in command or "${" in command:
96
+ raise KbCommandError("Shell expansion is not allowed in KB commands")
97
+ if any(token in SHELL_TOKENS for token in tokens):
98
+ raise KbCommandError("Pipes, redirects, and command chaining are not allowed")
99
+
100
+
101
+ def _numeric_value(token: str, option: str) -> str:
102
+ if not token.isdigit():
103
+ raise KbCommandError(f"{option} expects a positive integer")
104
+ return token
105
+
106
+
107
+ def _safe_pattern_value(token: str, option: str) -> str:
108
+ if "\x00" in token:
109
+ raise KbCommandError(f"{option} contains an invalid value")
110
+ if Path(token).is_absolute() or ".." in Path(token).parts:
111
+ raise KbCommandError(f"{option} values must stay inside data/kb")
112
+ return token
113
+
114
+
115
+ def _is_broad_raw_path(path: str) -> bool:
116
+ normalized = path.rstrip("/")
117
+ return normalized in {".", "raw", "raw/courses", "raw/docs"}
118
+
119
+
120
+ def _reject_unbounded_raw_search(command: str, paths: list[str], has_max_count: bool) -> None:
121
+ if has_max_count:
122
+ return
123
+ if any(_is_broad_raw_path(path) for path in paths):
124
+ raise KbCommandError(
125
+ f"`{command}` over broad raw paths requires -m/--max-count "
126
+ "or a narrower source/file path"
127
+ )
128
+
129
+
130
+ def _build_rg(tokens: list[str], root: Path) -> list[str]:
131
+ if shutil.which("rg") is None:
132
+ raise KbCommandError("`rg` is not installed in this runtime")
133
+ options: list[str] = []
134
+ positionals: list[str] = []
135
+ has_max_count = False
136
+ idx = 1
137
+ parsing_options = True
138
+ while idx < len(tokens):
139
+ token = tokens[idx]
140
+ if parsing_options and token == "--":
141
+ parsing_options = False
142
+ idx += 1
143
+ continue
144
+ if parsing_options and token in RG_FLAG_OPTIONS:
145
+ options.append(token)
146
+ idx += 1
147
+ continue
148
+ if parsing_options and token in RG_VALUE_OPTIONS:
149
+ if idx + 1 >= len(tokens):
150
+ raise KbCommandError(f"{token} requires a value")
151
+ value = tokens[idx + 1]
152
+ if token in {"-m", "--max-count"}:
153
+ value = _numeric_value(value, token)
154
+ has_max_count = True
155
+ else:
156
+ value = _safe_pattern_value(value, token)
157
+ options.extend([token, value])
158
+ idx += 2
159
+ continue
160
+ positionals.append(token)
161
+ parsing_options = False
162
+ idx += 1
163
+
164
+ if not positionals:
165
+ raise KbCommandError("`rg` requires a search pattern")
166
+ pattern = positionals[0]
167
+ paths = [_relative_path_arg(value, root) for value in positionals[1:]] or ["."]
168
+ _reject_unbounded_raw_search("rg", paths, has_max_count)
169
+ return [
170
+ "rg",
171
+ "--color=never",
172
+ "--line-number",
173
+ "--no-heading",
174
+ *options,
175
+ pattern,
176
+ *paths,
177
+ ]
178
+
179
+
180
+ def _build_grep(tokens: list[str], root: Path) -> list[str]:
181
+ options: list[str] = []
182
+ positionals: list[str] = []
183
+ idx = 1
184
+ parsing_options = True
185
+ while idx < len(tokens):
186
+ token = tokens[idx]
187
+ if parsing_options and token == "--":
188
+ parsing_options = False
189
+ idx += 1
190
+ continue
191
+ if parsing_options and token in GREP_FLAG_OPTIONS:
192
+ options.append(token)
193
+ idx += 1
194
+ continue
195
+ positionals.append(token)
196
+ parsing_options = False
197
+ idx += 1
198
+ if not positionals:
199
+ raise KbCommandError("`grep` requires a search pattern")
200
+ pattern = positionals[0]
201
+ paths = [_relative_path_arg(value, root) for value in positionals[1:]] or ["."]
202
+ return ["grep", "-R", "-n", "--color=never", *options, pattern, *paths]
203
+
204
+
205
+ def _build_find(tokens: list[str], root: Path) -> list[str]:
206
+ path = "."
207
+ idx = 1
208
+ if idx < len(tokens) and not tokens[idx].startswith("-"):
209
+ path = _relative_path_arg(tokens[idx], root)
210
+ idx += 1
211
+
212
+ argv = ["find", path]
213
+ while idx < len(tokens):
214
+ token = tokens[idx]
215
+ if token in {"-maxdepth", "-mindepth"}:
216
+ if idx + 1 >= len(tokens):
217
+ raise KbCommandError(f"{token} requires a value")
218
+ argv.extend([token, _numeric_value(tokens[idx + 1], token)])
219
+ idx += 2
220
+ continue
221
+ if token == "-type":
222
+ if idx + 1 >= len(tokens) or tokens[idx + 1] not in {"f", "d"}:
223
+ raise KbCommandError("-type supports only `f` or `d`")
224
+ argv.extend([token, tokens[idx + 1]])
225
+ idx += 2
226
+ continue
227
+ if token in {"-name", "-iname"}:
228
+ if idx + 1 >= len(tokens):
229
+ raise KbCommandError(f"{token} requires a value")
230
+ argv.extend([token, _safe_pattern_value(tokens[idx + 1], token)])
231
+ idx += 2
232
+ continue
233
+ raise KbCommandError("`find` supports only path, -maxdepth, -mindepth, -type, -name, and -iname")
234
+ return argv
235
+
236
+
237
+ def _build_ls(tokens: list[str], root: Path) -> list[str]:
238
+ argv = ["ls"]
239
+ paths: list[str] = []
240
+ for token in tokens[1:]:
241
+ if token in LS_FLAG_OPTIONS:
242
+ argv.append(token)
243
+ continue
244
+ paths.append(_relative_path_arg(token, root))
245
+ return [*argv, *(paths or ["."])]
246
+
247
+
248
+ def _build_sed(tokens: list[str], root: Path) -> list[str]:
249
+ if len(tokens) != 4 or tokens[1] != "-n":
250
+ raise KbCommandError("`sed` supports only: sed -n START,ENDp FILE")
251
+ expression = tokens[2]
252
+ if not re.fullmatch(r"\d+,\d+p", expression):
253
+ raise KbCommandError("`sed` supports only START,ENDp read expressions")
254
+ path = _relative_path_arg(tokens[3], root)
255
+ return ["sed", "-n", expression, path]
256
+
257
+
258
+ def _build_head(tokens: list[str], root: Path) -> list[str]:
259
+ argv = ["head"]
260
+ paths: list[str] = []
261
+ idx = 1
262
+ if idx < len(tokens) and tokens[idx] == "-n":
263
+ if idx + 1 >= len(tokens):
264
+ raise KbCommandError("-n requires a value")
265
+ argv.extend(["-n", _numeric_value(tokens[idx + 1], "-n")])
266
+ idx += 2
267
+ while idx < len(tokens):
268
+ paths.append(_relative_path_arg(tokens[idx], root))
269
+ idx += 1
270
+ if not paths:
271
+ raise KbCommandError("`head` requires at least one file path")
272
+ return [*argv, *paths]
273
+
274
+
275
+ def _build_cat(tokens: list[str], root: Path) -> list[str]:
276
+ if len(tokens) < 2:
277
+ raise KbCommandError("`cat` requires at least one file path")
278
+ return ["cat", *[_relative_path_arg(token, root) for token in tokens[1:]]]
279
+
280
+
281
+ def _build_wc(tokens: list[str], root: Path) -> list[str]:
282
+ argv = ["wc"]
283
+ paths: list[str] = []
284
+ for token in tokens[1:]:
285
+ if token in WC_FLAG_OPTIONS:
286
+ argv.append(token)
287
+ continue
288
+ paths.append(_relative_path_arg(token, root))
289
+ if not paths:
290
+ raise KbCommandError("`wc` requires at least one file path")
291
+ return [*argv, *paths]
292
+
293
+
294
+ def build_kb_command_argv(command: str, *, root: Path | None = None) -> tuple[list[str], Path]:
295
+ resolved_root = _resolve_root(root)
296
+ try:
297
+ tokens = shlex.split(command)
298
+ except ValueError as exc:
299
+ raise KbCommandError(f"Invalid KB command syntax: {exc}") from exc
300
+ if not tokens:
301
+ raise KbCommandError("KB command cannot be empty")
302
+ _reject_shell_syntax(command, tokens)
303
+
304
+ executable = tokens[0]
305
+ if executable not in SUPPORTED_COMMANDS:
306
+ supported = ", ".join(sorted(SUPPORTED_COMMANDS))
307
+ raise KbCommandError(f"Unsupported KB command `{executable}`. Supported: {supported}")
308
+
309
+ builders = {
310
+ "rg": _build_rg,
311
+ "grep": _build_grep,
312
+ "find": _build_find,
313
+ "ls": _build_ls,
314
+ "sed": _build_sed,
315
+ "head": _build_head,
316
+ "cat": _build_cat,
317
+ "wc": _build_wc,
318
+ }
319
+ return builders[executable](tokens, resolved_root), resolved_root
320
+
321
+
322
+ def _truncate(text: str, max_chars: int) -> tuple[str, bool]:
323
+ if len(text) <= max_chars:
324
+ return text, False
325
+ marker = f"\n... output truncated to {max_chars} characters ..."
326
+ return text[: max(0, max_chars - len(marker))] + marker, True
327
+
328
+
329
+ def run_kb_command(
330
+ command: str,
331
+ *,
332
+ root: Path | None = None,
333
+ timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
334
+ max_output_chars: int = DEFAULT_MAX_OUTPUT_CHARS,
335
+ ) -> KbCommandResult:
336
+ argv, resolved_root = build_kb_command_argv(command, root=root)
337
+ timeout = max(1, min(int(timeout_seconds), 30))
338
+ output_limit = max(1_000, min(int(max_output_chars), 80_000))
339
+ env = {
340
+ "HOME": os.getenv("HOME", ""),
341
+ "LANG": os.getenv("LANG", "C.UTF-8"),
342
+ "LC_ALL": os.getenv("LC_ALL", "C.UTF-8"),
343
+ "PATH": os.getenv("PATH", ""),
344
+ }
345
+ timed_out = False
346
+ try:
347
+ completed = subprocess.run(
348
+ argv,
349
+ cwd=resolved_root,
350
+ capture_output=True,
351
+ text=True,
352
+ timeout=timeout,
353
+ env=env,
354
+ check=False,
355
+ )
356
+ exit_code = int(completed.returncode)
357
+ stdout = completed.stdout or ""
358
+ stderr = completed.stderr or ""
359
+ except subprocess.TimeoutExpired as exc:
360
+ timed_out = True
361
+ exit_code = 124
362
+ stdout = exc.stdout if isinstance(exc.stdout, str) else ""
363
+ stderr = exc.stderr if isinstance(exc.stderr, str) else ""
364
+ stderr = (stderr + "\n" if stderr else "") + f"Command timed out after {timeout}s."
365
+
366
+ stdout, stdout_truncated = _truncate(stdout, output_limit)
367
+ stderr, stderr_truncated = _truncate(stderr, min(output_limit, 8_000))
368
+ return KbCommandResult(
369
+ command=command,
370
+ argv=argv,
371
+ cwd=resolved_root.as_posix(),
372
+ exit_code=exit_code,
373
+ stdout=stdout,
374
+ stderr=stderr,
375
+ timed_out=timed_out,
376
+ truncated=stdout_truncated or stderr_truncated,
377
+ )
378
+
379
+
380
+ def format_command_payload(result: KbCommandResult) -> str:
381
+ lines = [
382
+ f"$ {result.command}",
383
+ f"cwd: {result.cwd}",
384
+ f"exit_code: {result.exit_code}",
385
+ ]
386
+ if result.timed_out:
387
+ lines.append("timed_out: true")
388
+ if result.truncated:
389
+ lines.append("truncated: true")
390
+ if result.stdout:
391
+ lines.extend(["stdout:", result.stdout.rstrip()])
392
+ if result.stderr:
393
+ lines.extend(["stderr:", result.stderr.rstrip()])
394
+ return "\n".join(lines)
scripts/probe_gemini_cache.py ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Send the same prefix to Gemini twice and report cached_content_token_count.
2
+
3
+ Run with:
4
+ uv run dotenv -f .env run -- python -m scripts.probe_gemini_cache
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import os
10
+ import time
11
+
12
+ from google import genai
13
+ from google.genai import types
14
+
15
+ MODEL = "gemini-3.5-flash"
16
+
17
+ # Build a stable prefix that comfortably exceeds the 1024-token Flash minimum
18
+ # for implicit caching. Repeat a paragraph until we have ~3k tokens worth.
19
+ SYSTEM = (
20
+ "You are an AI teacher for applied AI, LLM, RAG, and Python topics. "
21
+ "Your job is to answer student questions clearly and accurately. "
22
+ "Ground claims in the provided context and cite sources inline."
23
+ )
24
+ PARAGRAPH = (
25
+ "Retrieval-Augmented Generation (RAG) systems combine a retrieval step "
26
+ "over a corpus with a generation step using a language model. The "
27
+ "retrieval step selects the most relevant passages from a vector index, "
28
+ "and the generation step conditions on those passages to produce a "
29
+ "grounded answer. Common components include a chunker, an embedding "
30
+ "model, a vector store, a retriever, a reranker, and a generation model. "
31
+ "Evaluation typically considers faithfulness, answer relevance, and "
32
+ "context precision and recall.\n\n"
33
+ )
34
+ LONG_CONTEXT = PARAGRAPH * 12 # ~3-4k tokens
35
+
36
+
37
+ def call(client: genai.Client, user_text: str) -> dict:
38
+ start = time.perf_counter()
39
+ response = client.models.generate_content(
40
+ model=MODEL,
41
+ config=types.GenerateContentConfig(system_instruction=SYSTEM),
42
+ contents=[
43
+ {
44
+ "role": "user",
45
+ "parts": [
46
+ {"text": LONG_CONTEXT},
47
+ {"text": user_text},
48
+ ],
49
+ }
50
+ ],
51
+ )
52
+ elapsed = time.perf_counter() - start
53
+ usage = response.usage_metadata
54
+ return {
55
+ "elapsed_s": round(elapsed, 2),
56
+ "prompt_tokens": getattr(usage, "prompt_token_count", None),
57
+ "cached_tokens": getattr(usage, "cached_content_token_count", None),
58
+ "output_tokens": getattr(usage, "candidates_token_count", None),
59
+ "thoughts_tokens": getattr(usage, "thoughts_token_count", None),
60
+ "text_preview": (response.text or "")[:120],
61
+ }
62
+
63
+
64
+ def main() -> None:
65
+ api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
66
+ if not api_key:
67
+ raise SystemExit("Set GEMINI_API_KEY in .env")
68
+
69
+ client = genai.Client(api_key=api_key)
70
+
71
+ print(f"Model: {MODEL}")
72
+ print(f"Stable prefix (system + LONG_CONTEXT) repeats. Only the final user_text differs.")
73
+ print()
74
+
75
+ for i in range(3):
76
+ result = call(client, user_text=f"Question {i+1}: what is RAG in one sentence?")
77
+ print(f"Call {i+1}: {result}")
78
+ time.sleep(0.5)
79
+
80
+
81
+ if __name__ == "__main__":
82
+ main()
scripts/prompts.py CHANGED
@@ -1,15 +1,25 @@
1
  from __future__ import annotations
2
 
 
 
 
3
 
4
  BASE_PROMPT_HEADER = """You are an AI teacher for applied AI, LLM, RAG, and Python topics.
5
 
6
  Your job is to answer student questions clearly and accurately."""
7
 
 
 
8
  RETRIEVAL_TOOL_LINE = (
9
  "- `retrieve_tutor_context` — retrieval over the course and documentation\n"
10
  " corpus. Use this for anything that depends on course content."
11
  )
12
 
 
 
 
 
 
13
  WEB_TOOL_LINES = {
14
  "google-genai": {
15
  "web_search": (
@@ -74,6 +84,21 @@ DO NOT use retrieval for:
74
  - Questions about your own role or capabilities in this app (answer directly).
75
  - Questions you can answer fully from general knowledge with no corpus dependency."""
76
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
77
  RETRIEVAL_CALL_STRATEGY = """## How to call the retrieval tool
78
 
79
  You have two strategies. Pick one per turn:
@@ -97,6 +122,26 @@ You have two strategies. Pick one per turn:
97
  ANSWERING_RULES = """## Answering rules
98
 
99
  - Ground factual claims about the corpus in the retrieved results.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
100
  - If retrieval results are weak or missing, say the topic is not well
101
  covered by the current knowledge base rather than guessing.
102
  - Synthesize retrieved material into a clear teaching explanation. Do not
@@ -122,10 +167,24 @@ def _provider_key(model_name: str) -> str:
122
  return ""
123
 
124
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
125
  def build_system_prompt(model_name: str, enabled_tools: tuple[str, ...]) -> str:
126
  provider = _provider_key(model_name)
127
  enabled = set(enabled_tools)
128
- tool_lines = [RETRIEVAL_TOOL_LINE]
129
  usage_sections: list[str] = []
130
  provider_web_tools = WEB_TOOL_LINES.get(provider, {})
131
  provider_web_usage = WEB_USAGE_SECTIONS.get(provider, {})
@@ -143,6 +202,7 @@ def build_system_prompt(model_name: str, enabled_tools: tuple[str, ...]) -> str:
143
  BASE_PROMPT_HEADER,
144
  f"{intro}\n\n" + "\n".join(tool_lines),
145
  RETRIEVAL_USAGE_SECTION,
 
146
  ]
147
  if usage_sections:
148
  parts.append("## When to use web search / URL reading\n\n" + "\n\n".join(usage_sections))
@@ -151,6 +211,14 @@ def build_system_prompt(model_name: str, enabled_tools: tuple[str, ...]) -> str:
151
  "course material. Combine tools when it helps (e.g. retrieve corpus\n"
152
  "context, then search the web for the latest update)."
153
  )
 
 
 
 
 
 
 
 
154
  parts.append(RETRIEVAL_CALL_STRATEGY)
155
  parts.append(ANSWERING_RULES)
156
  return "\n\n".join(parts) + "\n"
 
1
  from __future__ import annotations
2
 
3
+ import os
4
+ from pathlib import Path
5
+
6
 
7
  BASE_PROMPT_HEADER = """You are an AI teacher for applied AI, LLM, RAG, and Python topics.
8
 
9
  Your job is to answer student questions clearly and accurately."""
10
 
11
+ DEFAULT_KB_AGENTS_PATH = "data/kb/AGENTS.md"
12
+
13
  RETRIEVAL_TOOL_LINE = (
14
  "- `retrieve_tutor_context` — retrieval over the course and documentation\n"
15
  " corpus. Use this for anything that depends on course content."
16
  )
17
 
18
+ KB_TOOL_LINES = (
19
+ "- `run_kb_command` — run safe, read-only terminal-style KB inspection\n"
20
+ " commands such as `rg`, `grep`, `find`, `ls`, `sed`, `head`, `cat`, and `wc`.",
21
+ )
22
+
23
  WEB_TOOL_LINES = {
24
  "google-genai": {
25
  "web_search": (
 
84
  - Questions about your own role or capabilities in this app (answer directly).
85
  - Questions you can answer fully from general knowledge with no corpus dependency."""
86
 
87
+ KB_USAGE_SECTION = """## When to use the KB browsing tools
88
+
89
+ USE KB browsing for:
90
+ - Exact API names, class names, function names, config keys, or error strings.
91
+ - Broad questions where a wiki page can identify the right source pages.
92
+ - Comparisons, recipes, debugging, and verification-heavy answers.
93
+ - Course questions where reading the generated lesson markdown is useful.
94
+ - Targeted terminal-style inspection when `rg`, `grep`, `find`, or `sed`
95
+ can verify a specific string/path faster than top-k retrieval.
96
+
97
+ `run_kb_command` is not a general shell: no pipes, redirects, command chaining,
98
+ network commands, or writes. Follow the Local KB Instructions below for the
99
+ current wiki structure, the first-command rule, and the efficient command
100
+ patterns to use."""
101
+
102
  RETRIEVAL_CALL_STRATEGY = """## How to call the retrieval tool
103
 
104
  You have two strategies. Pick one per turn:
 
122
  ANSWERING_RULES = """## Answering rules
123
 
124
  - Ground factual claims about the corpus in the retrieved results.
125
+ - When using KB browsing, ground factual claims in raw source pages inspected
126
+ with `run_kb_command`, not only in wiki navigation pages.
127
+ - If the user explicitly restricts you to one available local knowledge tool,
128
+ obey that restriction. If the restriction prevents a complete answer, say so.
129
+ - Cite retrieved context using the original source URL or title shown by
130
+ `retrieve_tutor_context`.
131
+ - Cite files explored with `run_kb_command` using the raw file frontmatter URL,
132
+ `kb://doc/<doc_id>`, or a KB-root path like
133
+ `raw/docs/peft/developer_guides/lora.md`.
134
+ - Use inline Markdown links as citations sentence-by-sentence: place a
135
+ `[short source title](url-or-kb-reference)` citation in the same sentence,
136
+ immediately after the claim it supports.
137
+ - Do not use bare bracketed URLs or paths like `[https://...]` or
138
+ `[raw/docs/...]`; the bracket text should be a short source title.
139
+ - Do not place citations inside code blocks. Cite the prose sentence that
140
+ introduces or explains the code instead.
141
+ - Do not put all citations only in a final "Sources" section. A short sources
142
+ recap is okay only if the answer already has inline citations.
143
+ - Do not cite wiki pages unless the wiki page itself is the source of the
144
+ claim. Prefer raw docs and course lessons.
145
  - If retrieval results are weak or missing, say the topic is not well
146
  covered by the current knowledge base rather than guessing.
147
  - Synthesize retrieved material into a clear teaching explanation. Do not
 
167
  return ""
168
 
169
 
170
+ def kb_agents_path() -> Path:
171
+ return Path(os.getenv("AI_TUTOR_KB_AGENTS_PATH", DEFAULT_KB_AGENTS_PATH))
172
+
173
+
174
+ def load_kb_agents_instructions(path: Path | None = None) -> str:
175
+ resolved = path or kb_agents_path()
176
+ try:
177
+ if not resolved.exists() or not resolved.is_file():
178
+ return ""
179
+ return resolved.read_text(encoding="utf-8").strip()
180
+ except OSError:
181
+ return ""
182
+
183
+
184
  def build_system_prompt(model_name: str, enabled_tools: tuple[str, ...]) -> str:
185
  provider = _provider_key(model_name)
186
  enabled = set(enabled_tools)
187
+ tool_lines = [RETRIEVAL_TOOL_LINE, *KB_TOOL_LINES]
188
  usage_sections: list[str] = []
189
  provider_web_tools = WEB_TOOL_LINES.get(provider, {})
190
  provider_web_usage = WEB_USAGE_SECTIONS.get(provider, {})
 
202
  BASE_PROMPT_HEADER,
203
  f"{intro}\n\n" + "\n".join(tool_lines),
204
  RETRIEVAL_USAGE_SECTION,
205
+ KB_USAGE_SECTION,
206
  ]
207
  if usage_sections:
208
  parts.append("## When to use web search / URL reading\n\n" + "\n\n".join(usage_sections))
 
211
  "course material. Combine tools when it helps (e.g. retrieve corpus\n"
212
  "context, then search the web for the latest update)."
213
  )
214
+ kb_agents_instructions = load_kb_agents_instructions()
215
+ if kb_agents_instructions:
216
+ parts.append(
217
+ "## Local KB Instructions\n\n"
218
+ "The following instructions are loaded from `data/kb/AGENTS.md` and "
219
+ "describe the current local KB schema and workflow.\n\n"
220
+ f"{kb_agents_instructions}"
221
+ )
222
  parts.append(RETRIEVAL_CALL_STRATEGY)
223
  parts.append(ANSWERING_RULES)
224
  return "\n\n".join(parts) + "\n"
scripts/setup.py CHANGED
@@ -1,4 +1,5 @@
1
  import os
 
2
 
3
  import logfire
4
  from dotenv import load_dotenv
@@ -33,6 +34,14 @@ VECTOR_DB_DIR = "data/chroma-db-all_sources"
33
  VECTOR_COLLECTION_NAME = "chroma-db-all_sources"
34
  DOCUMENT_DICT_PATH = f"{VECTOR_DB_DIR}/document_dict_all_sources.pkl"
35
  BM25_INDEX_PATH = f"{VECTOR_DB_DIR}/bm25_index_all_sources.pkl"
 
 
 
 
 
 
 
 
36
  DEFAULT_MODEL_NAME = "google-genai:gemini-3.5-flash"
37
 
38
  AVAILABLE_MODELS: tuple[dict[str, str], ...] = (
@@ -44,20 +53,55 @@ CONCURRENCY_COUNT = int(os.getenv("CONCURRENCY_COUNT", 64))
44
  MONGODB_URI = os.getenv("MONGODB_URI")
45
 
46
 
47
- def ensure_local_vector_db() -> None:
48
- if os.path.exists(VECTOR_DB_DIR) and os.path.exists(DOCUMENT_DICT_PATH):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
49
  return
 
 
 
50
 
51
- logfire.warn(
52
- "Vector database does not exist locally, downloading from Hugging Face"
53
- )
54
- from huggingface_hub import snapshot_download
55
 
56
- snapshot_download(
57
- repo_id="towardsai-tutors/ai-tutor-vector-db",
58
- local_dir="data",
59
- repo_type="dataset",
 
 
 
 
 
 
 
60
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
61
 
62
 
63
  mongo_db = (
@@ -77,10 +121,16 @@ __all__ = [
77
  "DEFAULT_MODEL_NAME",
78
  "BM25_INDEX_PATH",
79
  "DOCUMENT_DICT_PATH",
 
 
 
 
 
80
  "SOURCE_KEY_TO_LABEL",
81
  "SOURCE_UI_TO_KEY",
82
  "VECTOR_COLLECTION_NAME",
83
  "VECTOR_DB_DIR",
 
84
  "ensure_local_vector_db",
85
  "mongo_db",
86
  ]
 
1
  import os
2
+ from pathlib import Path
3
 
4
  import logfire
5
  from dotenv import load_dotenv
 
34
  VECTOR_COLLECTION_NAME = "chroma-db-all_sources"
35
  DOCUMENT_DICT_PATH = f"{VECTOR_DB_DIR}/document_dict_all_sources.pkl"
36
  BM25_INDEX_PATH = f"{VECTOR_DB_DIR}/bm25_index_all_sources.pkl"
37
+ KB_DIR = "data/kb"
38
+ KB_MANIFEST_PATH = f"{KB_DIR}/generated/corpus_manifest.jsonl"
39
+ KB_INDEX_PATH = f"{KB_DIR}/wiki/index.md"
40
+ KB_AGENTS_PATH = f"{KB_DIR}/AGENTS.md"
41
+ # Canonical AGENTS.md content for the KB. Tracked in git alongside the
42
+ # scraping scripts so it survives `rm -rf data/kb` and a stale HF snapshot.
43
+ # `ensure_kb_agents_md()` copies it into `data/kb/AGENTS.md` on every startup.
44
+ KB_AGENTS_TEMPLATE_PATH = "data/scraping_scripts/kb_agents_template.md"
45
  DEFAULT_MODEL_NAME = "google-genai:gemini-3.5-flash"
46
 
47
  AVAILABLE_MODELS: tuple[dict[str, str], ...] = (
 
53
  MONGODB_URI = os.getenv("MONGODB_URI")
54
 
55
 
56
+ def ensure_kb_agents_md() -> None:
57
+ """Write `data/kb/AGENTS.md` from the canonical template if it exists.
58
+
59
+ Why: `data/kb/` is gitignored and downloaded from HuggingFace, so the
60
+ `AGENTS.md` that arrives in the snapshot might be stale (whoever uploaded
61
+ last is whoever wrote it). The template at
62
+ `data/scraping_scripts/kb_agents_template.md` is tracked in git and is the
63
+ single source of truth — we overwrite the live file from it on every
64
+ startup so a fresh `git pull` always propagates KB guidance changes to
65
+ the model, regardless of what's in the HF snapshot.
66
+ """
67
+ template_path = Path(KB_AGENTS_TEMPLATE_PATH)
68
+ if not template_path.exists():
69
+ # If the template is missing (partial checkout, etc.), don't clobber
70
+ # whatever AGENTS.md the HF snapshot provided — at least the agent
71
+ # will see *some* guidance.
72
  return
73
+ target = Path(KB_AGENTS_PATH)
74
+ target.parent.mkdir(parents=True, exist_ok=True)
75
+ target.write_text(template_path.read_text(encoding="utf-8"), encoding="utf-8")
76
 
 
 
 
 
77
 
78
+ def ensure_local_vector_db() -> None:
79
+ # On first start (or after `rm -rf data/{kb,chroma-db-all_sources}`),
80
+ # pull the whole bundle from HuggingFace. The repo
81
+ # `towardsai-tutors/ai-tutor-vector-db` contains both
82
+ # `chroma-db-all_sources/` AND `kb/`, both uploaded by
83
+ # `data.scraping_scripts.upload_dbs_to_hf`.
84
+ needs_download = not (
85
+ os.path.exists(VECTOR_DB_DIR)
86
+ and os.path.exists(DOCUMENT_DICT_PATH)
87
+ and os.path.exists(KB_MANIFEST_PATH)
88
+ and os.path.exists(KB_INDEX_PATH)
89
  )
90
+ if needs_download:
91
+ logfire.warn(
92
+ "Vector database does not exist locally, downloading from Hugging Face"
93
+ )
94
+ from huggingface_hub import snapshot_download
95
+
96
+ snapshot_download(
97
+ repo_id="towardsai-tutors/ai-tutor-vector-db",
98
+ local_dir="data",
99
+ repo_type="dataset",
100
+ )
101
+
102
+ # Always refresh AGENTS.md from the local template, regardless of whether
103
+ # we just downloaded or short-circuited. See `ensure_kb_agents_md` docstring.
104
+ ensure_kb_agents_md()
105
 
106
 
107
  mongo_db = (
 
121
  "DEFAULT_MODEL_NAME",
122
  "BM25_INDEX_PATH",
123
  "DOCUMENT_DICT_PATH",
124
+ "KB_AGENTS_PATH",
125
+ "KB_AGENTS_TEMPLATE_PATH",
126
+ "KB_DIR",
127
+ "KB_INDEX_PATH",
128
+ "KB_MANIFEST_PATH",
129
  "SOURCE_KEY_TO_LABEL",
130
  "SOURCE_UI_TO_KEY",
131
  "VECTOR_COLLECTION_NAME",
132
  "VECTOR_DB_DIR",
133
+ "ensure_kb_agents_md",
134
  "ensure_local_vector_db",
135
  "mongo_db",
136
  ]
tests/manual_e2e_langsmith.md ADDED
@@ -0,0 +1,345 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Manual End-to-End LangSmith Runbook
2
+
3
+ This runbook verifies the live chatbot path end to end and then looks up the
4
+ matching LangSmith trace/run from the CLI.
5
+
6
+ Use it when changing agent tools, KB browsing, citation handling, streaming API
7
+ parts, or model/provider settings.
8
+
9
+ ## Prerequisites
10
+
11
+ Run commands from the repository root.
12
+
13
+ Required local artifacts and environment:
14
+
15
+ - `.env` contains `COHERE_API_KEY`
16
+ - `.env` contains `GEMINI_API_KEY` or `GOOGLE_API_KEY`
17
+ - `.env` contains `LANGSMITH_API_KEY`
18
+ - `.env` has `LANGSMITH_TRACING=true`
19
+ - `.env` has `LANGSMITH_PROJECT=ai-tutor-app`
20
+ - `data/chroma-db-all_sources/` exists
21
+ - `data/kb/wiki/index.md` exists
22
+ - `curl`, `jq`, and the `langsmith` CLI are available through `uv run`
23
+
24
+ Quick check:
25
+
26
+ ```bash
27
+ uv run dotenv -f .env run -- python - <<'PY'
28
+ import os
29
+ for key in [
30
+ "COHERE_API_KEY",
31
+ "GEMINI_API_KEY",
32
+ "GOOGLE_API_KEY",
33
+ "LANGSMITH_API_KEY",
34
+ "LANGSMITH_TRACING",
35
+ "LANGSMITH_PROJECT",
36
+ ]:
37
+ value = os.getenv(key)
38
+ print(f"{key}={'set' if value else 'missing'}")
39
+ PY
40
+ ```
41
+
42
+ ## Fast Automated Smoke Tests
43
+
44
+ These tests run a live model-backed request and check the frontend/API stream
45
+ shape.
46
+
47
+ API stream contract:
48
+
49
+ ```bash
50
+ uv run dotenv -f .env run -- env RUN_LIVE_API_E2E=1 \
51
+ pytest tests/test_api.py::test_live_api_stream_exposes_frontend_parts -q
52
+ ```
53
+
54
+ Gradio curl path:
55
+
56
+ ```bash
57
+ uv run dotenv -f .env run -- env RUN_LIVE_GRADIO_CURL_E2E=1 \
58
+ pytest tests/test_gradio_kb_e2e.py::test_live_gradio_curl_can_use_retrieval_and_shell -q
59
+ ```
60
+
61
+ ## Manual API Test
62
+
63
+ Start the FastAPI app in one terminal:
64
+
65
+ ```bash
66
+ uv run dotenv -f .env run -- env \
67
+ AI_TUTOR_API_HOST=127.0.0.1 \
68
+ AI_TUTOR_API_PORT=8000 \
69
+ python -m scripts.api
70
+ ```
71
+
72
+ From another terminal, confirm the server is ready:
73
+
74
+ ```bash
75
+ curl -s http://127.0.0.1:8000/healthz
76
+ ```
77
+
78
+ Create a request payload. This prompt intentionally asks for both structured
79
+ retrieval and KB shell browsing so both citation paths are exercised.
80
+
81
+ ```bash
82
+ cat >/tmp/ai_tutor_e2e_payload.json <<'JSON'
83
+ {
84
+ "messages": [
85
+ {
86
+ "role": "user",
87
+ "parts": [
88
+ {
89
+ "type": "text",
90
+ "text": "Use run_kb_command to answer: can you tell me about codex?"
91
+ }
92
+ ]
93
+ }
94
+ ],
95
+ "sourceKeys": [
96
+ "agentic_ai_engineering",
97
+ "langchain",
98
+ "langgraph",
99
+ "peft",
100
+ "transformers"
101
+ ],
102
+ "enabledTools": [],
103
+ "model": "google-genai:gemini-3.5-flash",
104
+ "includeReasoning": true,
105
+ "threadId": ""
106
+ }
107
+ JSON
108
+ ```
109
+
110
+ Send the request and save the server-sent event stream:
111
+
112
+ ```bash
113
+ curl -sN --max-time 240 http://127.0.0.1:8000/api/chat \
114
+ -H "Content-Type: application/json" \
115
+ -d @/tmp/ai_tutor_e2e_payload.json \
116
+ -o /tmp/ai_tutor_e2e_api.sse
117
+ ```
118
+
119
+ Convert the stream to JSONL parts:
120
+
121
+ ```bash
122
+ grep '^data: ' /tmp/ai_tutor_e2e_api.sse \
123
+ | sed 's/^data: //' \
124
+ | grep -v '^\[DONE\]$' \
125
+ > /tmp/ai_tutor_e2e_api.parts.jsonl
126
+ ```
127
+
128
+ Inspect what the frontend received:
129
+
130
+ ```bash
131
+ jq -r '.type' /tmp/ai_tutor_e2e_api.parts.jsonl | sort | uniq -c
132
+ ```
133
+
134
+ Inspect tool calls:
135
+
136
+ ```bash
137
+ jq -r '
138
+ select(.type == "tool-input-available" or .type == "tool-input-start")
139
+ | [.type, .toolName, ((.input // .inputText // "") | tostring | gsub("\n"; " ") | .[0:220])]
140
+ | @tsv
141
+ ' /tmp/ai_tutor_e2e_api.parts.jsonl
142
+ ```
143
+
144
+ Inspect final answer text:
145
+
146
+ ```bash
147
+ jq -r 'select(.type == "text-delta") | .delta' \
148
+ /tmp/ai_tutor_e2e_api.parts.jsonl \
149
+ | tr -d '\n'
150
+ echo
151
+ ```
152
+
153
+ Inspect source cards:
154
+
155
+ ```bash
156
+ jq -r '
157
+ select(.type == "data-source")
158
+ | [.data.title, .data.url, .data.sourceKey, (.data.score // "")]
159
+ | @tsv
160
+ ' /tmp/ai_tutor_e2e_api.parts.jsonl
161
+ ```
162
+
163
+ Expected result:
164
+
165
+ - The stream includes `tool-input-*`, `tool-output-available`, `text-delta`,
166
+ `source-url`, `source-document`, `data-source`, and `finish` parts.
167
+ - Tool calls include `retrieve_tutor_context` and `run_kb_command`.
168
+ - The answer has inline citations, not only a final sources list.
169
+ - `data-source` parts include the source cards the frontend will render.
170
+
171
+ ## Manual Gradio Test
172
+
173
+ Start the Gradio app in one terminal:
174
+
175
+ ```bash
176
+ uv run dotenv -f .env run -- python -m scripts.main
177
+ ```
178
+
179
+ Submit a Gradio job:
180
+
181
+ ```bash
182
+ cat >/tmp/ai_tutor_gradio_payload.json <<'JSON'
183
+ {
184
+ "data": [
185
+ "Use run_kb_command to answer: can you tell me about codex?",
186
+ [],
187
+ [
188
+ "Agentic AI Engineering",
189
+ "LangChain Docs",
190
+ "LangGraph Docs",
191
+ "PEFT Docs",
192
+ "Transformers Docs"
193
+ ],
194
+ "google-genai:gemini-3.5-flash",
195
+ "",
196
+ false,
197
+ false
198
+ ]
199
+ }
200
+ JSON
201
+
202
+ curl -s http://127.0.0.1:7860/gradio_api/call/chat \
203
+ -H "Content-Type: application/json" \
204
+ -d @/tmp/ai_tutor_gradio_payload.json \
205
+ -o /tmp/ai_tutor_gradio_post.json
206
+ ```
207
+
208
+ Stream the result:
209
+
210
+ ```bash
211
+ EVENT_ID="$(jq -r '.event_id' /tmp/ai_tutor_gradio_post.json)"
212
+
213
+ curl -N --max-time 240 \
214
+ "http://127.0.0.1:7860/gradio_api/call/chat/${EVENT_ID}" \
215
+ -o /tmp/ai_tutor_gradio_stream.sse
216
+ ```
217
+
218
+ Expected result:
219
+
220
+ - The rendered text contains `Using \`retrieve_tutor_context\``.
221
+ - The rendered text contains `Using \`run_kb_command\``.
222
+ - The rendered text contains inline citations and source cards.
223
+
224
+ ## Find The LangSmith Trace
225
+
226
+ List recent traces in the project:
227
+
228
+ ```bash
229
+ uv run dotenv -f .env run -- langsmith trace list \
230
+ --project ai-tutor-app \
231
+ --limit 10 \
232
+ --full \
233
+ --format json \
234
+ -o /tmp/ai_tutor_langsmith_traces.json
235
+ ```
236
+
237
+ Print the likely match. Look for the newest trace whose input contains the
238
+ manual prompt.
239
+
240
+ ```bash
241
+ jq -r '
242
+ .[]
243
+ | [
244
+ (.trace_id // .id),
245
+ (.status // ""),
246
+ (.duration_ms // .latency_ms // .latency // ""),
247
+ ((.inputs.query // .inputs.input // .inputs.messages[-1].content // .inputs // "") | tostring | gsub("\n"; " ") | .[0:220])
248
+ ]
249
+ | @tsv
250
+ ' /tmp/ai_tutor_langsmith_traces.json
251
+ ```
252
+
253
+ Save the trace ID:
254
+
255
+ ```bash
256
+ TRACE_ID="paste-trace-id-here"
257
+ ```
258
+
259
+ Fetch the full trace:
260
+
261
+ ```bash
262
+ uv run dotenv -f .env run -- langsmith trace get "$TRACE_ID" \
263
+ --project ai-tutor-app \
264
+ --full \
265
+ --format json \
266
+ -o "/tmp/ai_tutor_trace_${TRACE_ID}.json"
267
+ ```
268
+
269
+ Fetch the corresponding run. For the root run, the trace ID is usually also
270
+ the run ID. If this fails, open the trace JSON and use the root run ID shown
271
+ there.
272
+
273
+ ```bash
274
+ uv run dotenv -f .env run -- langsmith run get "$TRACE_ID" \
275
+ --project ai-tutor-app \
276
+ --full \
277
+ --format json \
278
+ -o "/tmp/ai_tutor_run_${TRACE_ID}.json"
279
+ ```
280
+
281
+ ## Inspect Trace/Run Details
282
+
283
+ Overall status and latency:
284
+
285
+ ```bash
286
+ jq '{
287
+ id: (.trace_id // .id),
288
+ name,
289
+ status,
290
+ duration_ms: (.duration_ms // .latency_ms // .latency),
291
+ run_count: (.run_count // ((.runs // .child_runs // []) | length))
292
+ }' "/tmp/ai_tutor_trace_${TRACE_ID}.json"
293
+ ```
294
+
295
+ Tool calls:
296
+
297
+ ```bash
298
+ jq -r '
299
+ .. | objects
300
+ | select((.name? // "") | test("retrieve_tutor_context|run_kb_command"))
301
+ | [
302
+ (.name // ""),
303
+ (.status // ""),
304
+ (.duration_ms // .latency_ms // .latency // ""),
305
+ ((.inputs.command // .inputs.query // .inputs.input // "") | tostring | gsub("\n"; " ") | .[0:220])
306
+ ]
307
+ | @tsv
308
+ ' "/tmp/ai_tutor_trace_${TRACE_ID}.json"
309
+ ```
310
+
311
+ LLM calls:
312
+
313
+ ```bash
314
+ jq -r '
315
+ .. | objects
316
+ | select((.run_type? // "") | test("llm|chat_model"))
317
+ | [
318
+ (.name // ""),
319
+ (.status // ""),
320
+ (.duration_ms // .latency_ms // .latency // ""),
321
+ (.prompt_tokens // .usage.prompt_tokens // .outputs.llm_output.token_usage.prompt_tokens // ""),
322
+ (.completion_tokens // .usage.completion_tokens // .outputs.llm_output.token_usage.completion_tokens // "")
323
+ ]
324
+ | @tsv
325
+ ' "/tmp/ai_tutor_trace_${TRACE_ID}.json"
326
+ ```
327
+
328
+ Search the raw trace for the KB shell command sequence:
329
+
330
+ ```bash
331
+ rg -n "run_kb_command|retrieve_tutor_context|sed -n|rg -n|cat wiki" \
332
+ "/tmp/ai_tutor_trace_${TRACE_ID}.json"
333
+ ```
334
+
335
+ For latency debugging, compare these numbers:
336
+
337
+ - Root trace duration.
338
+ - Count of `run_kb_command` calls.
339
+ - Total shell/tool duration.
340
+ - Count of LLM/chat model calls.
341
+ - Longest LLM call duration.
342
+ - Whether the trace ended in `success`, `error`, or cancellation.
343
+
344
+ If shell/tool duration is tiny but total trace duration is large, the bottleneck
345
+ is model inference or repeated model turns, not filesystem access.
tests/test_api.py CHANGED
@@ -1,9 +1,18 @@
1
  from __future__ import annotations
2
 
3
  import json
 
 
 
 
 
 
 
4
  import unittest
5
  from unittest.mock import patch
6
 
 
 
7
  from fastapi.testclient import TestClient
8
 
9
  from scripts.api import app
@@ -18,6 +27,12 @@ def parse_sse_payloads(raw_text: str) -> list[str]:
18
  ]
19
 
20
 
 
 
 
 
 
 
21
  class ApiTestCase(unittest.TestCase):
22
  def test_healthcheck(self) -> None:
23
  with TestClient(app) as client:
@@ -230,5 +245,176 @@ class ApiTestCase(unittest.TestCase):
230
  self.assertLess(part_types.index("tool-input-start"), part_types.index("text-start"))
231
 
232
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
233
  if __name__ == "__main__":
234
  unittest.main()
 
1
  from __future__ import annotations
2
 
3
  import json
4
+ import os
5
+ import re
6
+ import shutil
7
+ import socket
8
+ import subprocess
9
+ import threading
10
+ import time
11
  import unittest
12
  from unittest.mock import patch
13
 
14
+ import pytest
15
+ import uvicorn
16
  from fastapi.testclient import TestClient
17
 
18
  from scripts.api import app
 
27
  ]
28
 
29
 
30
+ def free_port() -> int:
31
+ with socket.socket() as sock:
32
+ sock.bind(("127.0.0.1", 0))
33
+ return int(sock.getsockname()[1])
34
+
35
+
36
  class ApiTestCase(unittest.TestCase):
37
  def test_healthcheck(self) -> None:
38
  with TestClient(app) as client:
 
245
  self.assertLess(part_types.index("tool-input-start"), part_types.index("text-start"))
246
 
247
 
248
+ LIVE_API_E2E = pytest.mark.skipif(
249
+ os.getenv("RUN_LIVE_API_E2E") != "1",
250
+ reason="Set RUN_LIVE_API_E2E=1 to run the live frontend API smoke test.",
251
+ )
252
+
253
+
254
+ def require_live_api_e2e_prereqs() -> None:
255
+ if shutil.which("curl") is None:
256
+ pytest.skip("curl is required for the live API smoke test.")
257
+ if not os.getenv("COHERE_API_KEY"):
258
+ pytest.skip("COHERE_API_KEY is required for live retrieval.")
259
+ if not (os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")):
260
+ pytest.skip("GEMINI_API_KEY or GOOGLE_API_KEY is required for the live model.")
261
+ if not os.path.exists("data/kb/wiki/index.md"):
262
+ pytest.skip("data/kb artifacts must exist before running the live smoke test.")
263
+
264
+
265
+ def run_live_api_chat(prompt: str) -> list[dict]:
266
+ port = free_port()
267
+ config = uvicorn.Config(
268
+ app,
269
+ host="127.0.0.1",
270
+ port=port,
271
+ log_level="warning",
272
+ lifespan="on",
273
+ )
274
+ server = uvicorn.Server(config)
275
+ thread = threading.Thread(target=server.run, daemon=True)
276
+ thread.start()
277
+ try:
278
+ deadline = time.time() + 30
279
+ while time.time() < deadline:
280
+ health = subprocess.run(
281
+ ["curl", "-s", f"http://127.0.0.1:{port}/healthz"],
282
+ capture_output=True,
283
+ text=True,
284
+ check=False,
285
+ timeout=5,
286
+ )
287
+ if health.returncode == 0 and "ok" in health.stdout:
288
+ break
289
+ time.sleep(0.25)
290
+ else:
291
+ pytest.fail("API server did not become healthy")
292
+
293
+ payload = {
294
+ "messages": [
295
+ {
296
+ "role": "user",
297
+ "parts": [
298
+ {
299
+ "type": "text",
300
+ "text": prompt,
301
+ }
302
+ ],
303
+ }
304
+ ],
305
+ "sourceKeys": ["peft", "transformers"],
306
+ "enabledTools": [],
307
+ "model": os.getenv("LIVE_API_E2E_MODEL", "google-genai:gemini-3.5-flash"),
308
+ "includeReasoning": False,
309
+ "threadId": "",
310
+ }
311
+ response = subprocess.run(
312
+ [
313
+ "curl",
314
+ "-sN",
315
+ f"http://127.0.0.1:{port}/api/chat",
316
+ "-H",
317
+ "Content-Type: application/json",
318
+ "-d",
319
+ json.dumps(payload),
320
+ ],
321
+ capture_output=True,
322
+ text=True,
323
+ check=True,
324
+ timeout=240,
325
+ )
326
+ finally:
327
+ server.should_exit = True
328
+ thread.join(timeout=10)
329
+
330
+ payloads = parse_sse_payloads(response.stdout)
331
+ assert payloads[-1] == "[DONE]"
332
+ return [json.loads(item) for item in payloads[:-1]]
333
+
334
+
335
+ def live_part_types(parts: list[dict]) -> list[str]:
336
+ return [part["type"] for part in parts]
337
+
338
+
339
+ def live_tool_names(parts: list[dict]) -> set[str]:
340
+ return {
341
+ part.get("toolName")
342
+ for part in parts
343
+ if part["type"] in {"tool-input-start", "tool-input-available"}
344
+ }
345
+
346
+
347
+ def live_answer_text(parts: list[dict]) -> str:
348
+ return "\n".join(
349
+ part.get("delta", "") for part in parts if part["type"] == "text-delta"
350
+ )
351
+
352
+
353
+ def live_sources(parts: list[dict]) -> list[dict]:
354
+ return [part["data"] for part in parts if part["type"] == "data-source"]
355
+
356
+
357
+ @LIVE_API_E2E
358
+ def test_live_api_stream_exposes_frontend_parts() -> None:
359
+ require_live_api_e2e_prereqs()
360
+
361
+ parts = run_live_api_chat(
362
+ "Use both retrieve_tutor_context and run_kb_command to answer: "
363
+ "how does PEFT configure LoRA with LoraConfig?"
364
+ )
365
+ part_types = [part["type"] for part in parts]
366
+ assert "data-thread" in part_types
367
+ assert "tool-input-start" in part_types
368
+ assert "tool-output-available" in part_types
369
+ assert "text-delta" in part_types
370
+ assert "source-url" in part_types
371
+ assert "source-document" in part_types
372
+ assert "data-source" in part_types
373
+ assert "finish" in part_types
374
+ tool_names = live_tool_names(parts)
375
+ assert {"retrieve_tutor_context", "run_kb_command"} <= tool_names
376
+ final_text = live_answer_text(parts)
377
+ assert "LoRA" in final_text
378
+ assert re.search(r"\[[^\]]+\]\((?:https://|raw/|kb://doc/)", final_text)
379
+ assert "### Sources" not in final_text
380
+ sources = live_sources(parts)
381
+ assert any(source["sourceKey"] in {"peft", "transformers"} for source in sources)
382
+
383
+
384
+ @LIVE_API_E2E
385
+ def test_live_api_stream_obeys_retrieval_only_instruction() -> None:
386
+ require_live_api_e2e_prereqs()
387
+
388
+ parts = run_live_api_chat(
389
+ "Use only the retrieve_tutor_context tool. Do not call run_kb_command. "
390
+ "Answer: how does PEFT configure LoRA with LoraConfig?"
391
+ )
392
+
393
+ assert "data-source" in live_part_types(parts)
394
+ tool_names = live_tool_names(parts)
395
+ assert "retrieve_tutor_context" in tool_names
396
+ assert "run_kb_command" not in tool_names
397
+ assert "LoRA" in live_answer_text(parts)
398
+ assert live_sources(parts)
399
+
400
+
401
+ @LIVE_API_E2E
402
+ def test_live_api_stream_obeys_shell_only_instruction() -> None:
403
+ require_live_api_e2e_prereqs()
404
+
405
+ parts = run_live_api_chat(
406
+ "Use only the run_kb_command tool. Do not call retrieve_tutor_context. "
407
+ "Inspect raw PEFT docs and answer: how does PEFT configure LoRA with "
408
+ "LoraConfig?"
409
+ )
410
+
411
+ assert "data-source" in live_part_types(parts)
412
+ tool_names = live_tool_names(parts)
413
+ assert "run_kb_command" in tool_names
414
+ assert "retrieve_tutor_context" not in tool_names
415
+ assert "LoRA" in live_answer_text(parts)
416
+ assert live_sources(parts)
417
+
418
+
419
  if __name__ == "__main__":
420
  unittest.main()
tests/test_build_kb_artifacts.py ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import json
4
+ from pathlib import Path
5
+
6
+ from data.scraping_scripts import build_kb_artifacts as builder
7
+
8
+
9
+ def write_jsonl(path: Path, rows: list[dict]) -> None:
10
+ with path.open("w", encoding="utf-8") as handle:
11
+ for row in rows:
12
+ handle.write(json.dumps(row) + "\n")
13
+
14
+
15
+ def test_build_kb_artifacts_writes_markdown_and_indexes(
16
+ tmp_path: Path,
17
+ monkeypatch,
18
+ ) -> None:
19
+ input_file = tmp_path / "all_sources_data.jsonl"
20
+ output_dir = tmp_path / "kb"
21
+ docs_dir = tmp_path / "temp_docs_md_files"
22
+ docs_page = docs_dir / "package_reference" / "lora.mdx"
23
+ docs_page.parent.mkdir(parents=True)
24
+ docs_content = """---
25
+ sidebarTitle: LoRA Guide
26
+ ---
27
+ # LoRA Guide
28
+
29
+ Use `LoraConfig` with `get_peft_model`.
30
+ """
31
+ docs_page.write_text(docs_content, encoding="utf-8")
32
+ monkeypatch.setitem(
33
+ builder.SOURCE_CONFIGS,
34
+ "temp_docs",
35
+ {
36
+ "base_url": "https://example.com/docs/",
37
+ "input_directory": str(docs_dir),
38
+ "output_file": str(tmp_path / "temp_docs_data.jsonl"),
39
+ "source_name": "temp_docs",
40
+ "use_include_list": False,
41
+ "included_dirs": [],
42
+ "excluded_dirs": [],
43
+ "excluded_root_files": [],
44
+ "included_root_files": [],
45
+ "url_extension": "",
46
+ },
47
+ )
48
+ write_jsonl(
49
+ input_file,
50
+ [
51
+ {
52
+ "doc_id": "legacy-random-id",
53
+ "name": "LoRA Guide",
54
+ "url": "https://example.com/docs/package_reference/lora",
55
+ "source": "temp_docs",
56
+ "tokens": 400,
57
+ "retrieve_doc": True,
58
+ "content": docs_content,
59
+ },
60
+ {
61
+ "doc_id": "0a4fe6fa-928c-4cbf-951c-d0fd9dce01f3",
62
+ "name": "Lesson 18: Research Loop",
63
+ "url": "https://academy.towardsai.net/courses/take/agent-engineering/multimedia/70289117-lesson-18-the-research-loop",
64
+ "source": "agentic_ai_engineering",
65
+ "tokens": 800,
66
+ "retrieve_doc": True,
67
+ "content": "# Lesson 18: Research Loop\n\nCall `generate_next_queries_tool`.",
68
+ },
69
+ ],
70
+ )
71
+
72
+ summary = builder.build_kb_artifacts(input_file, output_dir)
73
+
74
+ assert summary == {"documents": 2, "manifest_rows": 2}
75
+ manifest_path = output_dir / "generated" / "corpus_manifest.jsonl"
76
+ headings_path = output_dir / "generated" / "headings.jsonl"
77
+ symbols_path = output_dir / "generated" / "symbols.tsv"
78
+ assert manifest_path.exists()
79
+ assert headings_path.exists()
80
+ assert symbols_path.exists()
81
+
82
+ manifest_rows = [
83
+ json.loads(line)
84
+ for line in manifest_path.read_text(encoding="utf-8").splitlines()
85
+ ]
86
+ assert manifest_rows[0]["doc_id"] == "temp_docs:docs-package-reference-lora"
87
+ assert manifest_rows[0]["content_hash"].startswith("sha256:")
88
+ assert manifest_rows[0]["source_path"] == "package_reference/lora.mdx"
89
+ assert manifest_rows[0]["original_path"] == docs_page.as_posix()
90
+
91
+ markdown_path = Path(manifest_rows[0]["path"])
92
+ assert markdown_path == output_dir / "raw" / "docs" / "temp_docs" / "package_reference" / "lora.mdx"
93
+ markdown = markdown_path.read_text(encoding="utf-8")
94
+ assert 'doc_id: "temp_docs:docs-package-reference-lora"' in markdown
95
+ assert 'source_path: "package_reference/lora.mdx"' in markdown
96
+ assert "sidebarTitle" not in markdown
97
+ assert "# LoRA Guide" in markdown
98
+ assert "`LoraConfig`" in markdown
99
+
100
+ course_row = manifest_rows[1]
101
+ assert course_row["source_group"] == "courses"
102
+ assert course_row["source_path"] == ""
103
+
104
+ symbols = symbols_path.read_text(encoding="utf-8")
105
+ assert "LoraConfig" in symbols
106
+ assert "generate_next_queries_tool" in symbols
tests/test_chat_service.py CHANGED
@@ -1,18 +1,23 @@
1
  from __future__ import annotations
2
 
 
3
  import types
4
  import unittest
5
  from unittest.mock import patch
6
 
7
- from langchain_core.messages import AIMessage, HumanMessage
8
 
9
  from scripts.chat_service import (
 
 
10
  agent_run_config,
11
  build_agent,
12
  effective_tool_names,
 
13
  sync_thread_with_history,
 
14
  )
15
- from scripts.chat_types import ChatRequest, ChatTurn
16
 
17
 
18
  class FakeAgent:
@@ -27,6 +32,61 @@ class FakeAgent:
27
  self.updated_states.append((config, payload))
28
 
29
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
30
  class ChatServiceTestCase(unittest.TestCase):
31
  def test_effective_tool_names_follow_provider(self) -> None:
32
  self.assertEqual(
@@ -34,14 +94,24 @@ class ChatServiceTestCase(unittest.TestCase):
34
  "google-genai:gemini-3.5-flash",
35
  ("web_search", "url_context", "web_fetch"),
36
  ),
37
- ("retrieve_tutor_context", "google_search", "url_context"),
 
 
 
 
 
38
  )
39
  self.assertEqual(
40
  effective_tool_names(
41
  "anthropic:claude-haiku-4-5",
42
  ("web_search", "url_context", "web_fetch"),
43
  ),
44
- ("retrieve_tutor_context", "web_search", "web_fetch"),
 
 
 
 
 
45
  )
46
 
47
  def test_agent_run_config_adds_langsmith_metadata(self) -> None:
@@ -59,13 +129,18 @@ class ChatServiceTestCase(unittest.TestCase):
59
  self.assertEqual(config["run_name"], "ai-tutor-agent-turn")
60
  self.assertIn("provider:google-genai", config["tags"])
61
  self.assertIn("tool:retrieve_tutor_context", config["tags"])
 
62
  self.assertIn("tool:google_search", config["tags"])
63
  self.assertEqual(config["metadata"]["thread_id"], "thread_123")
64
  self.assertEqual(config["metadata"]["conversation_id"], "thread_123")
65
  self.assertEqual(config["metadata"]["message_id"], "message_456")
66
  self.assertEqual(
67
  config["metadata"]["available_tools"],
68
- ["retrieve_tutor_context", "google_search"],
 
 
 
 
69
  )
70
  self.assertEqual(
71
  config["metadata"]["source_keys"],
@@ -117,10 +192,117 @@ class ChatServiceTestCase(unittest.TestCase):
117
  disabled_tool_defs = created_agents[1].kwargs["tools"]
118
  self.assertIn({"google_search": {}}, enabled_tool_defs)
119
  self.assertIn({"url_context": {}}, enabled_tool_defs)
120
- self.assertEqual(disabled_tool_defs, [enabled_tool_defs[0]])
 
 
 
 
 
 
 
 
121
  self.assertEqual(len(created_agents[0].kwargs["middleware"]), 1)
122
  self.assertEqual(created_agents[1].kwargs["middleware"], [])
123
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
124
  def test_shorter_history_branches_to_fresh_thread(self) -> None:
125
  agent = FakeAgent(
126
  [
 
1
  from __future__ import annotations
2
 
3
+ import asyncio
4
  import types
5
  import unittest
6
  from unittest.mock import patch
7
 
8
+ from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage
9
 
10
  from scripts.chat_service import (
11
+ _claim_kb_command_budget,
12
+ _clear_kb_command_budget,
13
  agent_run_config,
14
  build_agent,
15
  effective_tool_names,
16
+ resolve_answer_citations,
17
  sync_thread_with_history,
18
+ stream_chat,
19
  )
20
+ from scripts.chat_types import ChatRequest, ChatTurn, SourceMatch
21
 
22
 
23
  class FakeAgent:
 
32
  self.updated_states.append((config, payload))
33
 
34
 
35
+ class FakeStreamingAgent(FakeAgent):
36
+ async def astream(self, *_args, **_kwargs):
37
+ yield {
38
+ "type": "messages",
39
+ "data": (
40
+ AIMessageChunk(
41
+ content="",
42
+ tool_calls=[
43
+ {
44
+ "id": "call_rg",
45
+ "name": "run_kb_command",
46
+ "args": {"command": "rg LoraConfig raw"},
47
+ }
48
+ ],
49
+ ),
50
+ {"langgraph_step": "model"},
51
+ ),
52
+ }
53
+ yield {
54
+ "type": "updates",
55
+ "data": {
56
+ "tools": {
57
+ "messages": [
58
+ ToolMessage(
59
+ content=(
60
+ "$ rg LoraConfig raw\n"
61
+ "cwd: /tmp/kb\n"
62
+ "exit_code: 0\n"
63
+ "stdout:\n"
64
+ "raw/docs/peft/lora.md:3:LoraConfig"
65
+ ),
66
+ name="run_kb_command",
67
+ tool_call_id="call_rg",
68
+ )
69
+ ]
70
+ }
71
+ },
72
+ }
73
+ yield {
74
+ "type": "updates",
75
+ "data": {
76
+ "model": {
77
+ "messages": [
78
+ AIMessage(
79
+ content=(
80
+ "LoraConfig is documented in LoRA "
81
+ "[LoRA](raw/docs/peft/lora.md)."
82
+ )
83
+ )
84
+ ]
85
+ }
86
+ },
87
+ }
88
+
89
+
90
  class ChatServiceTestCase(unittest.TestCase):
91
  def test_effective_tool_names_follow_provider(self) -> None:
92
  self.assertEqual(
 
94
  "google-genai:gemini-3.5-flash",
95
  ("web_search", "url_context", "web_fetch"),
96
  ),
97
+ (
98
+ "retrieve_tutor_context",
99
+ "run_kb_command",
100
+ "google_search",
101
+ "url_context",
102
+ ),
103
  )
104
  self.assertEqual(
105
  effective_tool_names(
106
  "anthropic:claude-haiku-4-5",
107
  ("web_search", "url_context", "web_fetch"),
108
  ),
109
+ (
110
+ "retrieve_tutor_context",
111
+ "run_kb_command",
112
+ "web_search",
113
+ "web_fetch",
114
+ ),
115
  )
116
 
117
  def test_agent_run_config_adds_langsmith_metadata(self) -> None:
 
129
  self.assertEqual(config["run_name"], "ai-tutor-agent-turn")
130
  self.assertIn("provider:google-genai", config["tags"])
131
  self.assertIn("tool:retrieve_tutor_context", config["tags"])
132
+ self.assertIn("tool:run_kb_command", config["tags"])
133
  self.assertIn("tool:google_search", config["tags"])
134
  self.assertEqual(config["metadata"]["thread_id"], "thread_123")
135
  self.assertEqual(config["metadata"]["conversation_id"], "thread_123")
136
  self.assertEqual(config["metadata"]["message_id"], "message_456")
137
  self.assertEqual(
138
  config["metadata"]["available_tools"],
139
+ [
140
+ "retrieve_tutor_context",
141
+ "run_kb_command",
142
+ "google_search",
143
+ ],
144
  )
145
  self.assertEqual(
146
  config["metadata"]["source_keys"],
 
192
  disabled_tool_defs = created_agents[1].kwargs["tools"]
193
  self.assertIn({"google_search": {}}, enabled_tool_defs)
194
  self.assertIn({"url_context": {}}, enabled_tool_defs)
195
+ self.assertEqual(len(disabled_tool_defs), 2)
196
+ disabled_tool_names = {tool.name for tool in disabled_tool_defs}
197
+ self.assertEqual(
198
+ disabled_tool_names,
199
+ {
200
+ "retrieve_tutor_context",
201
+ "run_kb_command",
202
+ },
203
+ )
204
  self.assertEqual(len(created_agents[0].kwargs["middleware"]), 1)
205
  self.assertEqual(created_agents[1].kwargs["middleware"], [])
206
 
207
+ def test_kb_command_budget_blocks_after_limit(self) -> None:
208
+ session_id = "test_budget_session"
209
+ _clear_kb_command_budget(session_id)
210
+ try:
211
+ self.assertEqual(_claim_kb_command_budget(session_id, 2), (True, 1))
212
+ self.assertEqual(_claim_kb_command_budget(session_id, 2), (True, 2))
213
+ self.assertEqual(_claim_kb_command_budget(session_id, 2), (False, 2))
214
+ finally:
215
+ _clear_kb_command_budget(session_id)
216
+
217
+ def test_resolve_answer_citations_uses_current_turn_evidence(self) -> None:
218
+ retrieval = SourceMatch(
219
+ doc_id="peft:lora",
220
+ title="LoRA",
221
+ url="https://example.com/lora",
222
+ source_key="peft",
223
+ source_label="PEFT Docs",
224
+ score=12.0,
225
+ group="docs",
226
+ )
227
+
228
+ resolved = resolve_answer_citations(
229
+ "See [LoRA](https://example.com/lora) and [Other](https://example.com/other).",
230
+ retrieval_evidence={"peft:lora": retrieval},
231
+ shell_evidence={},
232
+ web_evidence={},
233
+ )
234
+
235
+ self.assertEqual(resolved, [retrieval])
236
+
237
+ def test_resolve_answer_citations_ignores_unseen_kb_paths(self) -> None:
238
+ resolved = resolve_answer_citations(
239
+ "See [LoRA](raw/docs/peft/lora.md).",
240
+ retrieval_evidence={},
241
+ shell_evidence={},
242
+ web_evidence={},
243
+ )
244
+
245
+ self.assertEqual(resolved, [])
246
+
247
+ def test_stream_chat_resolves_shell_citation_after_final_answer(self) -> None:
248
+ agent = FakeStreamingAgent([])
249
+ request = ChatRequest(
250
+ query="Use rg to find LoraConfig",
251
+ source_keys=("peft",),
252
+ model_name="google-genai:gemini-3.5-flash",
253
+ include_reasoning=False,
254
+ enabled_tools=(),
255
+ )
256
+
257
+ async def collect_events():
258
+ return [event async for event in stream_chat(request)]
259
+
260
+ shell_match = SourceMatch(
261
+ doc_id="peft:lora",
262
+ title="LoRA",
263
+ url="https://example.com/lora",
264
+ source_key="peft",
265
+ source_label="PEFT Docs",
266
+ score=1.0,
267
+ group="docs",
268
+ )
269
+
270
+ def fake_resolve_manifest_reference(reference, **_kwargs):
271
+ return shell_match if "raw/docs/peft/lora.md" in reference else None
272
+
273
+ with (
274
+ patch("scripts.chat_service.build_agent", return_value=agent),
275
+ patch("scripts.chat_service.new_thread_id", return_value="thread_rg"),
276
+ patch(
277
+ "scripts.chat_service.resolve_manifest_reference",
278
+ side_effect=fake_resolve_manifest_reference,
279
+ ),
280
+ ):
281
+ events = asyncio.run(collect_events())
282
+
283
+ started = [
284
+ event
285
+ for event in events
286
+ if event.type == "tool_call_started"
287
+ and event.data.get("tool_name") == "run_kb_command"
288
+ ]
289
+ completed = [
290
+ event
291
+ for event in events
292
+ if event.type == "tool_call_completed"
293
+ and event.data.get("tool_name") == "run_kb_command"
294
+ ]
295
+ source_matches = [
296
+ event
297
+ for event in events
298
+ if event.type == "source_match"
299
+ ]
300
+
301
+ self.assertEqual(started[0].data["args_text"], "rg LoraConfig raw")
302
+ self.assertIn("rg LoraConfig raw", completed[0].data["output_text"])
303
+ self.assertEqual(source_matches[0].data["source_key"], "peft")
304
+ self.assertNotIn("call_id", source_matches[0].data)
305
+
306
  def test_shorter_history_branches_to_fresh_thread(self) -> None:
307
  agent = FakeAgent(
308
  [
tests/test_gradio_kb_e2e.py ADDED
@@ -0,0 +1,207 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import socket
4
+ import os
5
+
6
+ import pytest
7
+ from gradio_client import Client
8
+
9
+ from scripts.chat_types import ChatEvent
10
+
11
+
12
+ def free_port() -> int:
13
+ with socket.socket() as sock:
14
+ sock.bind(("127.0.0.1", 0))
15
+ return int(sock.getsockname()[1])
16
+
17
+
18
+ def test_gradio_chat_api_surfaces_shell_and_retrieval_activity(monkeypatch) -> None:
19
+ import scripts.main as main
20
+
21
+ async def fake_stream_chat(_request):
22
+ yield ChatEvent("thread_started", {"thread_id": "thread_kb"})
23
+ yield ChatEvent("message_started", {"message_id": "message_kb"})
24
+ yield ChatEvent(
25
+ "tool_call_started",
26
+ {
27
+ "message_id": "message_kb",
28
+ "call_id": "call_search",
29
+ "tool_name": "retrieve_tutor_context",
30
+ "args": {"query": "generate_next_queries_tool"},
31
+ "args_text": "generate_next_queries_tool",
32
+ },
33
+ )
34
+ yield ChatEvent(
35
+ "tool_call_completed",
36
+ {
37
+ "message_id": "message_kb",
38
+ "call_id": "call_search",
39
+ "tool_name": "retrieve_tutor_context",
40
+ "args": {"query": "generate_next_queries_tool"},
41
+ "args_text": "generate_next_queries_tool",
42
+ "output_text": (
43
+ '{"matches": [{"doc_id": "agentic_ai_engineering:lesson-18", '
44
+ '"title": "Lesson 18: Research Loop", '
45
+ '"url": "https://academy.towardsai.net/lesson-18", '
46
+ '"source": "agentic_ai_engineering", "score": 10.0}]}'
47
+ ),
48
+ },
49
+ )
50
+ yield ChatEvent(
51
+ "tool_call_started",
52
+ {
53
+ "message_id": "message_kb",
54
+ "call_id": "call_rg",
55
+ "tool_name": "run_kb_command",
56
+ "args": {"command": "rg generate_next_queries_tool raw"},
57
+ "args_text": "rg generate_next_queries_tool raw",
58
+ },
59
+ )
60
+ yield ChatEvent(
61
+ "tool_call_completed",
62
+ {
63
+ "message_id": "message_kb",
64
+ "call_id": "call_rg",
65
+ "tool_name": "run_kb_command",
66
+ "args": {"command": "rg generate_next_queries_tool raw"},
67
+ "args_text": "rg generate_next_queries_tool raw",
68
+ "output_text": "$ rg generate_next_queries_tool raw\ncwd: data/kb\nexit_code: 0\nstdout:\nraw/courses/agentic_ai_engineering/lesson-18.md:4:Call `generate_next_queries_tool`.",
69
+ },
70
+ )
71
+ yield ChatEvent(
72
+ "text_delta",
73
+ {
74
+ "message_id": "message_kb",
75
+ "text": "The symbol appears in Lesson 18: Research Loop.",
76
+ },
77
+ )
78
+ yield ChatEvent(
79
+ "source_match",
80
+ {
81
+ "message_id": "message_kb",
82
+ "doc_id": "agentic_ai_engineering:lesson-18",
83
+ "title": "Lesson 18: Research Loop",
84
+ "url": "https://academy.towardsai.net/lesson-18",
85
+ "source_key": "agentic_ai_engineering",
86
+ "source_label": "Agentic AI Engineering",
87
+ "score": 10.0,
88
+ "group": "courses",
89
+ },
90
+ )
91
+ yield ChatEvent(
92
+ "message_completed",
93
+ {
94
+ "message_id": "message_kb",
95
+ "thread_id": "thread_kb",
96
+ "answer": "The symbol appears in Lesson 18: Research Loop.",
97
+ },
98
+ )
99
+
100
+ monkeypatch.setattr(main, "stream_chat", fake_stream_chat)
101
+ port = free_port()
102
+ main.demo.launch(
103
+ server_name="127.0.0.1",
104
+ server_port=port,
105
+ prevent_thread_lock=True,
106
+ quiet=True,
107
+ )
108
+ try:
109
+ client = Client(f"http://127.0.0.1:{port}")
110
+ result = client.predict(
111
+ "Where is generate_next_queries_tool discussed?",
112
+ [],
113
+ ["Agentic AI Engineering"],
114
+ "google-genai:gemini-3.5-flash",
115
+ "",
116
+ False,
117
+ False,
118
+ api_name="/chat",
119
+ )
120
+ finally:
121
+ main.demo.close()
122
+
123
+ rendered = result[0] if isinstance(result, tuple) else str(result)
124
+ assert "Using `retrieve_tutor_context`" in rendered
125
+ assert "Using `run_kb_command`" in rendered
126
+ assert "rg generate_next_queries_tool raw" in rendered
127
+ assert "Lesson 18: Research Loop" in rendered
128
+ assert "Agentic AI Engineering" in rendered
129
+
130
+
131
+ @pytest.mark.skipif(
132
+ os.getenv("RUN_LIVE_GRADIO_CURL_E2E") != "1",
133
+ reason="Set RUN_LIVE_GRADIO_CURL_E2E=1 to run the live model-backed Gradio curl smoke test.",
134
+ )
135
+ def test_live_gradio_curl_can_use_retrieval_and_shell() -> None:
136
+ if os.getenv("RUN_LIVE_GRADIO_CURL_E2E") != "1":
137
+ pytest.skip("Set RUN_LIVE_GRADIO_CURL_E2E=1 to run this test.")
138
+ import scripts.main as main
139
+
140
+ import json
141
+ import shutil
142
+ import subprocess
143
+
144
+ if shutil.which("curl") is None:
145
+ pytest.skip("curl is required for the live Gradio curl smoke test.")
146
+ if not os.getenv("COHERE_API_KEY"):
147
+ pytest.skip("COHERE_API_KEY is required for live retrieval fallback.")
148
+ if not (os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")):
149
+ pytest.skip("GEMINI_API_KEY or GOOGLE_API_KEY is required for the live model.")
150
+ if not os.path.exists("data/kb/wiki/index.md"):
151
+ pytest.skip("data/kb artifacts must exist before running the live smoke test.")
152
+
153
+ port = free_port()
154
+ main.demo.launch(
155
+ server_name="127.0.0.1",
156
+ server_port=port,
157
+ prevent_thread_lock=True,
158
+ quiet=True,
159
+ )
160
+ try:
161
+ post = subprocess.run(
162
+ [
163
+ "curl",
164
+ "-s",
165
+ f"http://127.0.0.1:{port}/gradio_api/call/chat",
166
+ "-H",
167
+ "Content-Type: application/json",
168
+ "-d",
169
+ json.dumps(
170
+ {
171
+ "data": [
172
+ "Use both retrieve_tutor_context and run_kb_command to answer: how does PEFT configure LoRA with LoraConfig?",
173
+ [],
174
+ ["PEFT Docs", "Transformers Docs"],
175
+ os.getenv("LIVE_GRADIO_E2E_MODEL", "google-genai:gemini-3.5-flash"),
176
+ "",
177
+ False,
178
+ False,
179
+ ]
180
+ }
181
+ ),
182
+ ],
183
+ check=True,
184
+ capture_output=True,
185
+ text=True,
186
+ timeout=30,
187
+ )
188
+ event_id = json.loads(post.stdout)["event_id"]
189
+ stream = subprocess.run(
190
+ [
191
+ "curl",
192
+ "-N",
193
+ f"http://127.0.0.1:{port}/gradio_api/call/chat/{event_id}",
194
+ ],
195
+ check=True,
196
+ capture_output=True,
197
+ text=True,
198
+ timeout=180,
199
+ )
200
+ finally:
201
+ main.demo.close()
202
+
203
+ rendered = stream.stdout
204
+ assert "Using `retrieve_tutor_context`" in rendered
205
+ assert "Using `run_kb_command`" in rendered
206
+ assert "LoRA" in rendered
207
+ assert "PEFT" in rendered or "Transformers" in rendered
tests/test_gradio_presenter.py CHANGED
@@ -3,7 +3,11 @@ from __future__ import annotations
3
  import unittest
4
 
5
  from scripts.chat_types import ChatEvent
6
- from scripts.gradio_presenter import GradioPresenterState, SOURCES_HEADER
 
 
 
 
7
 
8
 
9
  class GradioPresenterTestCase(unittest.TestCase):
@@ -46,6 +50,43 @@ class GradioPresenterTestCase(unittest.TestCase):
46
  self.assertIn(SOURCES_HEADER, completed)
47
  self.assertIn("RAG overview", completed)
48
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
49
 
50
  if __name__ == "__main__":
51
  unittest.main()
 
3
  import unittest
4
 
5
  from scripts.chat_types import ChatEvent
6
+ from scripts.gradio_presenter import (
7
+ GradioPresenterState,
8
+ SOURCES_HEADER,
9
+ summarize_tool_result,
10
+ )
11
 
12
 
13
  class GradioPresenterTestCase(unittest.TestCase):
 
50
  self.assertIn(SOURCES_HEADER, completed)
51
  self.assertIn("RAG overview", completed)
52
 
53
+ def test_retrieval_summary_uses_tool_payload_matches(self) -> None:
54
+ event = ChatEvent(
55
+ "tool_call_completed",
56
+ {
57
+ "tool_name": "retrieve_tutor_context",
58
+ "output_text": (
59
+ '{"matches": [{"title": "LoRA", '
60
+ '"source_label": "PEFT Docs"}]}'
61
+ ),
62
+ },
63
+ )
64
+
65
+ self.assertEqual(
66
+ summarize_tool_result(event),
67
+ "_Found 1 match from PEFT Docs._",
68
+ )
69
+
70
+ def test_run_kb_command_summary_confirms_command_execution(self) -> None:
71
+ event = ChatEvent(
72
+ "tool_call_completed",
73
+ {
74
+ "tool_name": "run_kb_command",
75
+ "output_text": (
76
+ "$ rg LoraConfig raw\n"
77
+ "cwd: data/kb\n"
78
+ "exit_code: 0\n"
79
+ "stdout:\n"
80
+ "raw/docs/peft/lora.md:3:LoraConfig"
81
+ ),
82
+ },
83
+ )
84
+
85
+ self.assertEqual(
86
+ summarize_tool_result(event),
87
+ "_Ran `rg LoraConfig raw` with exit code 0._",
88
+ )
89
+
90
 
91
  if __name__ == "__main__":
92
  unittest.main()
tests/test_kb_shell.py ADDED
@@ -0,0 +1,102 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import json
4
+ import shutil
5
+ from pathlib import Path
6
+
7
+ import pytest
8
+
9
+ from data.scraping_scripts.build_kb_artifacts import build_kb_artifacts
10
+ from data.scraping_scripts.update_kb_wiki import update_kb_wiki
11
+ from scripts.kb_shell import KbCommandError, format_command_payload, run_kb_command
12
+
13
+
14
+ def write_jsonl(path: Path, rows: list[dict]) -> None:
15
+ with path.open("w", encoding="utf-8") as handle:
16
+ for row in rows:
17
+ handle.write(json.dumps(row) + "\n")
18
+
19
+
20
+ @pytest.fixture()
21
+ def kb_dir(tmp_path: Path) -> Path:
22
+ input_file = tmp_path / "all_sources_data.jsonl"
23
+ output_dir = tmp_path / "kb"
24
+ write_jsonl(
25
+ input_file,
26
+ [
27
+ {
28
+ "doc_id": "peft:lora",
29
+ "name": "LoRA",
30
+ "url": "https://example.com/peft/lora",
31
+ "source": "peft",
32
+ "source_path": "package_reference/lora.md",
33
+ "tokens": 300,
34
+ "retrieve_doc": True,
35
+ "content": "# LoRA\n\nUse `LoraConfig` with `get_peft_model`.",
36
+ },
37
+ {
38
+ "doc_id": "agentic_ai_engineering:lesson-18",
39
+ "name": "Lesson 18: Research Loop",
40
+ "url": "https://academy.towardsai.net/lesson-18",
41
+ "source": "agentic_ai_engineering",
42
+ "tokens": 500,
43
+ "retrieve_doc": True,
44
+ "content": "# Lesson 18\n\nCall `generate_next_queries_tool`.",
45
+ },
46
+ ],
47
+ )
48
+ build_kb_artifacts(input_file, output_dir)
49
+ update_kb_wiki(output_dir, seed_defaults=True, changed_only=False)
50
+ return output_dir
51
+
52
+
53
+ @pytest.mark.skipif(shutil.which("rg") is None, reason="ripgrep is required")
54
+ def test_run_kb_command_runs_rg_and_returns_plain_transcript(kb_dir: Path) -> None:
55
+ result = run_kb_command(
56
+ "rg LoraConfig raw/docs/peft",
57
+ root=kb_dir,
58
+ )
59
+
60
+ assert result.command == "rg LoraConfig raw/docs/peft"
61
+ assert result.argv[:4] == ["rg", "--color=never", "--line-number", "--no-heading"]
62
+ assert result.exit_code == 0
63
+ assert not result.timed_out
64
+ assert "LoraConfig" in result.stdout
65
+
66
+ payload = format_command_payload(result)
67
+ assert "$ rg LoraConfig raw/docs/peft" in payload
68
+ assert "exit_code: 0" in payload
69
+ assert "stdout:" in payload
70
+ assert "matches" not in payload
71
+
72
+
73
+ @pytest.mark.skipif(shutil.which("rg") is None, reason="ripgrep is required")
74
+ def test_run_kb_command_can_read_all_raw_sources(kb_dir: Path) -> None:
75
+ result = run_kb_command(
76
+ "rg generate_next_queries_tool raw/courses/agentic_ai_engineering",
77
+ root=kb_dir,
78
+ )
79
+
80
+ assert result.exit_code == 0
81
+ assert "generate_next_queries_tool" in result.stdout
82
+ assert "agentic_ai_engineering" in result.stdout
83
+
84
+
85
+ def test_run_kb_command_rejects_shell_chaining(kb_dir: Path) -> None:
86
+ with pytest.raises(KbCommandError):
87
+ run_kb_command("rg LoraConfig | head", root=kb_dir)
88
+
89
+
90
+ def test_run_kb_command_rejects_path_traversal(kb_dir: Path) -> None:
91
+ with pytest.raises(KbCommandError):
92
+ run_kb_command("cat ../outside.md", root=kb_dir)
93
+
94
+
95
+ def test_run_kb_command_rejects_unbounded_broad_raw_search(kb_dir: Path) -> None:
96
+ with pytest.raises(KbCommandError, match="requires -m"):
97
+ run_kb_command("rg LoraConfig raw", root=kb_dir)
98
+
99
+ result = run_kb_command("rg -m 20 LoraConfig raw", root=kb_dir)
100
+
101
+ assert result.exit_code == 0
102
+ assert "LoraConfig" in result.stdout
tests/test_process_md_files.py CHANGED
@@ -1,4 +1,8 @@
1
- from data.scraping_scripts.process_md_files import extract_title
 
 
 
 
2
 
3
 
4
  def test_extract_title_prefers_frontmatter_title() -> None:
@@ -56,3 +60,36 @@ First paragraph fallback
56
  """
57
 
58
  assert extract_title(content) == "First paragraph fallback"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from data.scraping_scripts.process_md_files import (
2
+ content_sha256,
3
+ extract_title,
4
+ stable_doc_id,
5
+ )
6
 
7
 
8
  def test_extract_title_prefers_frontmatter_title() -> None:
 
60
  """
61
 
62
  assert extract_title(content) == "First paragraph fallback"
63
+
64
+
65
+ def test_stable_doc_id_prefers_source_path() -> None:
66
+ doc_id = stable_doc_id(
67
+ source="transformers",
68
+ source_path="main_classes/model.md",
69
+ title="Model",
70
+ url="https://example.com/ignored",
71
+ content_hash=content_sha256("content"),
72
+ )
73
+
74
+ assert doc_id == "transformers:main-classes-model"
75
+
76
+
77
+ def test_stable_doc_id_uses_url_when_source_path_missing() -> None:
78
+ doc_id = stable_doc_id(
79
+ source="agentic_ai_engineering",
80
+ source_path="",
81
+ title="Lesson 18",
82
+ url="https://academy.towardsai.net/courses/take/agent-engineering/multimedia/70289117-lesson-18-the-research-loop",
83
+ content_hash=content_sha256("content"),
84
+ )
85
+
86
+ assert doc_id == (
87
+ "agentic_ai_engineering:courses-take-agent-engineering-multimedia-"
88
+ "70289117-lesson-18-the-research-loop"
89
+ )
90
+
91
+
92
+ def test_content_sha256_is_stable_and_prefixed() -> None:
93
+ assert content_sha256("same") == content_sha256("same")
94
+ assert content_sha256("same").startswith("sha256:")
95
+ assert content_sha256("same") != content_sha256("different")
tests/test_prompts.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ from pathlib import Path
4
+
5
+ from scripts.prompts import build_system_prompt, load_kb_agents_instructions
6
+
7
+
8
+ def test_build_system_prompt_includes_local_kb_agents_instructions(
9
+ tmp_path: Path,
10
+ monkeypatch,
11
+ ) -> None:
12
+ agents_path = tmp_path / "AGENTS.md"
13
+ agents_path.write_text(
14
+ "# Test KB Instructions\n\n"
15
+ "- Read `wiki/index.md` before broad questions.\n"
16
+ "- Treat `raw/` as source authority.\n",
17
+ encoding="utf-8",
18
+ )
19
+ monkeypatch.setenv("AI_TUTOR_KB_AGENTS_PATH", str(agents_path))
20
+
21
+ prompt = build_system_prompt("google-genai:gemini-3.5-flash", ())
22
+
23
+ assert "## Local KB Instructions" in prompt
24
+ assert "loaded from `data/kb/AGENTS.md`" in prompt
25
+ assert "# Test KB Instructions" in prompt
26
+ assert "Treat `raw/` as source authority." in prompt
27
+
28
+
29
+ def test_build_system_prompt_omits_local_kb_section_when_file_missing(
30
+ tmp_path: Path,
31
+ monkeypatch,
32
+ ) -> None:
33
+ monkeypatch.setenv("AI_TUTOR_KB_AGENTS_PATH", str(tmp_path / "missing.md"))
34
+
35
+ prompt = build_system_prompt("google-genai:gemini-3.5-flash", ())
36
+
37
+ assert "## Local KB Instructions" not in prompt
38
+
39
+
40
+ def test_load_kb_agents_instructions_returns_empty_for_missing_path(
41
+ tmp_path: Path,
42
+ ) -> None:
43
+ assert load_kb_agents_instructions(tmp_path / "missing.md") == ""
tests/test_update_kb_wiki.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import json
4
+ from pathlib import Path
5
+
6
+ from data.scraping_scripts.build_kb_artifacts import build_kb_artifacts
7
+ from data.scraping_scripts.update_kb_wiki import update_kb_wiki
8
+
9
+
10
+ def write_jsonl(path: Path, rows: list[dict]) -> None:
11
+ with path.open("w", encoding="utf-8") as handle:
12
+ for row in rows:
13
+ handle.write(json.dumps(row) + "\n")
14
+
15
+
16
+ def test_update_kb_wiki_seeds_navigation_pages(tmp_path: Path) -> None:
17
+ input_file = tmp_path / "all_sources_data.jsonl"
18
+ kb_dir = tmp_path / "kb"
19
+ write_jsonl(
20
+ input_file,
21
+ [
22
+ {
23
+ "doc_id": "peft:lora",
24
+ "name": "LoRA",
25
+ "url": "https://example.com/lora",
26
+ "source": "peft",
27
+ "source_path": "package_reference/lora.md",
28
+ "tokens": 300,
29
+ "retrieve_doc": True,
30
+ "content": "# LoRA\n\nUse `LoraConfig` for adapter fine-tuning.",
31
+ },
32
+ {
33
+ "doc_id": "agentic_ai_engineering:lesson-1",
34
+ "name": "Lesson 1",
35
+ "url": "https://example.com/lesson-1",
36
+ "source": "agentic_ai_engineering",
37
+ "tokens": 300,
38
+ "retrieve_doc": True,
39
+ "content": "# Lesson 1\n\nAgent lesson.",
40
+ }
41
+ ],
42
+ )
43
+ build_kb_artifacts(input_file, kb_dir)
44
+
45
+ update_kb_wiki(kb_dir, seed_defaults=True, changed_only=False)
46
+
47
+ assert (kb_dir / "AGENTS.md").exists()
48
+ assert (kb_dir / "wiki" / "index.md").exists()
49
+ assert (kb_dir / "wiki" / "frameworks" / "peft.md").exists()
50
+ assert (kb_dir / "wiki" / "courses" / "agentic_ai_engineering.md").exists()
51
+ agents = (kb_dir / "AGENTS.md").read_text(encoding="utf-8")
52
+ assert "## Exploration Budget" in agents
53
+ assert 'rg -n -m 20 "query terms"' in agents
54
+ assert "Stop browsing once you have one clearly relevant source page" in agents
55
+ index = (kb_dir / "wiki" / "index.md").read_text(encoding="utf-8")
56
+ assert "`wiki/courses/agentic_ai_engineering.md`" in index
57
+ assert "`wiki/frameworks/peft.md`" in index
58
+ topic = (kb_dir / "wiki" / "topics" / "lora.md").read_text(encoding="utf-8")
59
+ assert "LoRA" in topic
60
+ assert "`raw/docs/peft/package_reference/lora.md`" in topic
61
+ assert "lookup_tutor_symbol" not in topic
62
+
63
+
64
+ def test_update_kb_wiki_preserves_authored_topic_content_and_appends_log(tmp_path: Path) -> None:
65
+ input_file = tmp_path / "all_sources_data.jsonl"
66
+ kb_dir = tmp_path / "kb"
67
+ write_jsonl(
68
+ input_file,
69
+ [
70
+ {
71
+ "doc_id": "peft:lora",
72
+ "name": "LoRA",
73
+ "url": "https://example.com/lora",
74
+ "source": "peft",
75
+ "source_path": "package_reference/lora.md",
76
+ "tokens": 300,
77
+ "retrieve_doc": True,
78
+ "content": "# LoRA\n\nUse `LoraConfig` for adapter fine-tuning.",
79
+ }
80
+ ],
81
+ )
82
+ build_kb_artifacts(input_file, kb_dir)
83
+ topic_path = kb_dir / "wiki" / "topics" / "lora.md"
84
+ topic_path.parent.mkdir(parents=True)
85
+ topic_path.write_text("# Lora\n\nAuthored synthesis stays.\n", encoding="utf-8")
86
+
87
+ update_kb_wiki(kb_dir, seed_defaults=False, changed_only=True)
88
+ update_kb_wiki(kb_dir, seed_defaults=False, changed_only=True)
89
+
90
+ topic = topic_path.read_text(encoding="utf-8")
91
+ assert "Authored synthesis stays." in topic
92
+ assert "`raw/docs/peft/package_reference/lora.md`" in topic
93
+ log = (kb_dir / "wiki" / "log.md").read_text(encoding="utf-8")
94
+ assert log.count("Generated or refreshed") == 2