add knowledge base wiki exploration tool
Browse files- .gitignore +6 -0
- README.md +3 -1
- data/.gitattributes +65 -0
- data/scraping_scripts/README.md +79 -0
- data/scraping_scripts/add_course_workflow.py +23 -0
- data/scraping_scripts/build_kb_artifacts.py +546 -0
- data/scraping_scripts/kb_agents_template.md +50 -0
- data/scraping_scripts/process_md_files.py +48 -10
- data/scraping_scripts/update_docs_workflow.py +38 -0
- data/scraping_scripts/update_kb_wiki.py +302 -0
- data/scraping_scripts/upload_dbs_to_hf.py +2 -1
- scripts/chat_service.py +442 -242
- scripts/gradio_presenter.py +37 -1
- scripts/kb_manifest.py +172 -0
- scripts/kb_shell.py +394 -0
- scripts/probe_gemini_cache.py +82 -0
- scripts/prompts.py +69 -1
- scripts/setup.py +60 -10
- tests/manual_e2e_langsmith.md +345 -0
- tests/test_api.py +186 -0
- tests/test_build_kb_artifacts.py +106 -0
- tests/test_chat_service.py +188 -6
- tests/test_gradio_kb_e2e.py +207 -0
- tests/test_gradio_presenter.py +42 -1
- tests/test_kb_shell.py +102 -0
- tests/test_process_md_files.py +38 -1
- tests/test_prompts.py +43 -0
- tests/test_update_kb_wiki.py +94 -0
.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
|
| 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":
|
| 294 |
-
|
| 295 |
-
|
| 296 |
-
|
| 297 |
-
|
| 298 |
-
|
| 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=["*.
|
| 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
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
| 264 |
-
|
| 265 |
-
payload: str,
|
| 266 |
-
) -> list[SourceMatch]:
|
| 267 |
-
updated: list[SourceMatch] = []
|
| 268 |
for match in parse_tool_payload(payload):
|
| 269 |
-
|
| 270 |
-
|
| 271 |
-
|
| 272 |
-
|
| 273 |
-
|
| 274 |
-
|
| 275 |
-
|
| 276 |
-
|
|
|
|
|
|
|
| 277 |
)
|
| 278 |
-
|
| 279 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 280 |
continue
|
| 281 |
-
|
| 282 |
-
|
| 283 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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] = [
|
| 417 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
-
|
|
|
|
|
|
|
| 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 |
-
|
| 722 |
-
{"thread_id": active_thread_id},
|
| 723 |
-
)
|
| 724 |
-
yield ChatEvent(
|
| 725 |
-
"message_started",
|
| 726 |
-
{"message_id": message_id},
|
| 727 |
-
)
|
| 728 |
|
| 729 |
-
|
| 730 |
-
|
| 731 |
-
|
| 732 |
-
|
| 733 |
-
|
| 734 |
-
|
| 735 |
-
|
| 736 |
-
|
| 737 |
-
|
| 738 |
-
|
| 739 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 740 |
|
| 741 |
-
|
| 742 |
-
|
| 743 |
-
|
| 744 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 745 |
yield ChatEvent(
|
| 746 |
-
"
|
| 747 |
{
|
| 748 |
"message_id": message_id,
|
| 749 |
-
"
|
| 750 |
-
"
|
|
|
|
|
|
|
| 751 |
},
|
| 752 |
)
|
| 753 |
|
| 754 |
-
|
| 755 |
-
|
| 756 |
-
|
| 757 |
-
|
| 758 |
-
|
| 759 |
-
|
| 760 |
-
"
|
| 761 |
-
|
| 762 |
-
|
| 763 |
-
|
| 764 |
-
|
| 765 |
-
|
| 766 |
-
)
|
| 767 |
|
| 768 |
-
|
| 769 |
-
|
| 770 |
-
|
| 771 |
-
|
| 772 |
-
|
| 773 |
-
|
| 774 |
-
|
| 775 |
-
|
| 776 |
-
|
| 777 |
-
"text": text_delta,
|
| 778 |
-
},
|
| 779 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 780 |
|
| 781 |
-
|
| 782 |
-
|
| 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 |
-
|
| 824 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 825 |
|
| 826 |
-
|
| 827 |
-
|
| 828 |
-
|
| 829 |
-
|
| 830 |
-
|
| 831 |
-
|
| 832 |
-
|
| 833 |
-
|
| 834 |
-
|
| 835 |
-
|
| 836 |
-
|
| 837 |
-
|
|
|
|
| 838 |
yield ChatEvent(
|
| 839 |
-
"
|
| 840 |
{
|
| 841 |
"message_id": message_id,
|
| 842 |
-
"
|
| 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 |
-
|
| 854 |
-
|
| 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 |
-
|
| 916 |
-
|
| 917 |
-
|
| 918 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 919 |
)
|
| 920 |
-
|
| 921 |
-
|
| 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":
|
| 929 |
-
"tool_name":
|
| 930 |
-
"args":
|
| 931 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 932 |
},
|
| 933 |
)
|
| 934 |
-
|
| 935 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 936 |
yield ChatEvent(
|
| 937 |
-
"
|
| 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 |
-
|
| 972 |
-
|
| 973 |
-
|
|
|
|
|
|
|
| 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(
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
| 48 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
-
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
-
(
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 38 |
)
|
| 39 |
self.assertEqual(
|
| 40 |
effective_tool_names(
|
| 41 |
"anthropic:claude-haiku-4-5",
|
| 42 |
("web_search", "url_context", "web_fetch"),
|
| 43 |
),
|
| 44 |
-
(
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
-
[
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|