ruff format

data/scraping_scripts/add_context_to_nodes.py

@@ -19,19 +19,19 @@ from pydantic import BaseModel, Field
 from tenacity import retry, retry_if_exception, stop_after_attempt
 from tqdm.asyncio import tqdm
 
-from scripts.chroma_rag import ChunkRecord, build_chunk_records, format_chunk_for_retrieval
+from scripts.chroma_rag import (
+    ChunkRecord,
+    build_chunk_records,
+    format_chunk_for_retrieval,
+)
 
 load_dotenv(".env")
 
 CONTEXT_MODEL = os.getenv("GEMINI_CONTEXT_MODEL", "gemini-3.1-flash-lite")
 CONTEXT_MAX_OUTPUT_TOKENS = 1000
 CONTEXT_TPM_LIMIT = int(os.getenv("GEMINI_CONTEXT_TPM_LIMIT", "30000000"))
-CONTEXT_TPM_SAFETY_MARGIN = float(
-    os.getenv("GEMINI_CONTEXT_TPM_SAFETY_MARGIN", "0.8")
-)
-CONTEXT_TPM_WINDOW_SECONDS = float(
-    os.getenv("GEMINI_CONTEXT_TPM_WINDOW_SECONDS", "60")
-)
+CONTEXT_TPM_SAFETY_MARGIN = float(os.getenv("GEMINI_CONTEXT_TPM_SAFETY_MARGIN", "0.8"))
+CONTEXT_TPM_WINDOW_SECONDS = float(os.getenv("GEMINI_CONTEXT_TPM_WINDOW_SECONDS", "60"))
 CONTEXT_RETRY_ATTEMPTS = int(os.getenv("GEMINI_CONTEXT_RETRY_ATTEMPTS", "8"))
 DEFAULT_SEMAPHORE_LIMIT = int(os.getenv("GEMINI_CONTEXT_CONCURRENCY", "50"))
 MAX_DOCUMENT_TOKENS = 120_000
@@ -192,7 +192,7 @@ def wait_for_genai_retry(retry_state) -> float:
     if retry_delay is not None:
         return retry_delay + random.uniform(1.0, 4.0)
 
-    exponential_delay = min(60, max(4, 2 ** retry_state.attempt_number))
+    exponential_delay = min(60, max(4, 2**retry_state.attempt_number))
     return exponential_delay + random.uniform(0.5, 2.0)
 
 

data/scraping_scripts/build_kb_artifacts.py

@@ -217,9 +217,7 @@ def unique_path(candidate: Path, used_paths: set[Path], content_hash: str) -> Pa
     candidate = candidate.with_name(f"{stem}-{suffix}{candidate.suffix}")
     counter = 2
     while candidate in used_paths:
-        candidate = candidate.with_name(
-            f"{stem}-{suffix}-{counter}{candidate.suffix}"
-        )
+        candidate = candidate.with_name(f"{stem}-{suffix}-{counter}{candidate.suffix}")
         counter += 1
     used_paths.add(candidate)
     return candidate
@@ -373,8 +371,11 @@ def write_raw_markdown(
     for record in records:
         original = find_original_markdown(record, original_indexes)
         if original:
-            output_path = raw_dir / "docs" / record.source / safe_relative_path(
-                original.source_path
+            output_path = (
+                raw_dir
+                / "docs"
+                / record.source
+                / safe_relative_path(original.source_path)
             )
             output_path = unique_path(output_path, used_paths, record.content_hash)
             source_path = original.source_path
@@ -386,7 +387,9 @@ def write_raw_markdown(
             original_path = ""
             content = record.content
             if record.source_group == "docs":
-                fallback_counts[record.source] = fallback_counts.get(record.source, 0) + 1
+                fallback_counts[record.source] = (
+                    fallback_counts.get(record.source, 0) + 1
+                )
 
         output_path.parent.mkdir(parents=True, exist_ok=True)
         markdown = markdown_with_frontmatter(
@@ -417,7 +420,9 @@ def write_raw_markdown(
 def iter_markdown_headings(path: Path) -> Iterable[dict[str, Any]]:
     in_fence = False
     fence_char = ""
-    for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
+    for line_number, line in enumerate(
+        path.read_text(encoding="utf-8").splitlines(), 1
+    ):
         fence_match = FENCE_RE.match(line)
         if fence_match:
             current = fence_match.group(1)[0]
@@ -464,7 +469,9 @@ def extract_symbols(text: str) -> set[str]:
     }
 
 
-def write_generated_indexes(manifest: list[dict[str, Any]], generated_dir: Path) -> None:
+def write_generated_indexes(
+    manifest: list[dict[str, Any]], generated_dir: Path
+) -> None:
     if generated_dir.exists():
         shutil.rmtree(generated_dir)
     generated_dir.mkdir(parents=True, exist_ok=True)
@@ -511,7 +518,9 @@ def write_generated_indexes(manifest: list[dict[str, Any]], generated_dir: Path)
         for row in heading_rows:
             handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")
 
-    with (generated_dir / "symbols.tsv").open("w", encoding="utf-8", newline="") as handle:
+    with (generated_dir / "symbols.tsv").open(
+        "w", encoding="utf-8", newline=""
+    ) as handle:
         writer = csv.DictWriter(
             handle,
             fieldnames=["symbol", "source", "title", "path", "heading", "doc_id"],
@@ -523,7 +532,9 @@ def write_generated_indexes(manifest: list[dict[str, Any]], generated_dir: Path)
 
 def build_kb_artifacts(input_file: Path, output_dir: Path) -> dict[str, int]:
     rows = load_jsonl(input_file)
-    records = [normalize_record(row) for row in rows if str(row.get("content") or "").strip()]
+    records = [
+        normalize_record(row) for row in rows if str(row.get("content") or "").strip()
+    ]
     manifest = write_raw_markdown(records, input_file=input_file, output_dir=output_dir)
     write_generated_indexes(manifest, output_dir / GENERATED_DIR_NAME)
     return {"documents": len(records), "manifest_rows": len(manifest)}

data/scraping_scripts/create_vector_stores.py

@@ -77,7 +77,9 @@ SOURCE_CONFIGS = vector_store_source_configs()
 
 def load_or_create_chunk_records(source: str) -> list[Any]:
     config = SOURCE_CONFIGS[source]
-    if source == "all_sources" and os.path.exists("data/all_sources_contextual_nodes.pkl"):
+    if source == "all_sources" and os.path.exists(
+        "data/all_sources_contextual_nodes.pkl"
+    ):
         with open("data/all_sources_contextual_nodes.pkl", "rb") as handle:
             records = pickle.load(handle)
         active_records = []

data/scraping_scripts/github_to_markdown_ai_docs.py

@@ -22,7 +22,7 @@ Example:
 
 This will download and process the documentation files for both TRL and PEFT libraries.
 
-Note: 
+Note:
 - Ensure you have set the GITHUB_TOKEN variable with your GitHub Personal Access Token.
 - The script creates a 'data' directory in the current working directory to store the downloaded files.
 - Each source's files are stored in a subdirectory named '<repo>_md_files'.
@@ -363,9 +363,7 @@ def process_source(source: str):
             os.makedirs(target_dir, exist_ok=True)
             fetch_files(api_url, target_dir, source_extensions, local_dir)
     else:
-        api_url = (
-            f"https://api.github.com/repos/{config['owner']}/{config['repo']}/contents/{config['path']}"
-        )
+        api_url = f"https://api.github.com/repos/{config['owner']}/{config['repo']}/contents/{config['path']}"
         fetch_files(api_url, local_dir, source_extensions, local_dir)
 
     save_source_extension_manifest(local_dir, source_extensions)

data/scraping_scripts/retire_source_workflow.py

@@ -85,7 +85,9 @@ def ensure_required_files_exist(data_repo_id: str) -> None:
         if local_path.exists():
             continue
 
-        print(f"{remote_filename} not found locally. Downloading from {data_repo_id}...")
+        print(
+            f"{remote_filename} not found locally. Downloading from {data_repo_id}..."
+        )
         hf_hub_download(
             token=os.getenv("HF_TOKEN"),
             repo_id=data_repo_id,
@@ -95,7 +97,9 @@ def ensure_required_files_exist(data_repo_id: str) -> None:
         )
 
 
-def filter_all_sources_jsonl(sources_to_retire: set[str], dry_run: bool) -> Counter[str]:
+def filter_all_sources_jsonl(
+    sources_to_retire: set[str], dry_run: bool
+) -> Counter[str]:
     if not ALL_SOURCES_JSONL.exists():
         raise SystemExit(f"Missing required file: {ALL_SOURCES_JSONL}")
 
@@ -272,7 +276,9 @@ def update_source_registry(sources: list[str], dry_run: bool) -> None:
 
 def rebuild_vector_store() -> None:
     if not os.getenv("COHERE_API_KEY"):
-        raise SystemExit("COHERE_API_KEY is required to rebuild the Chroma vector store.")
+        raise SystemExit(
+            "COHERE_API_KEY is required to rebuild the Chroma vector store."
+        )
 
     print("Rebuilding Chroma vector store for all_sources...")
     result = run_module("data.scraping_scripts.create_vector_stores", "all_sources")
@@ -318,7 +324,9 @@ def delete_remote_source_files(
 
 def upload_vector_store(vector_repo_id: str) -> None:
     print(f"Uploading rebuilt vector store to {vector_repo_id}...")
-    result = run_module("data.scraping_scripts.upload_dbs_to_hf", "--repo", vector_repo_id)
+    result = run_module(
+        "data.scraping_scripts.upload_dbs_to_hf", "--repo", vector_repo_id
+    )
     if result.returncode != 0:
         raise SystemExit("Error uploading vector store. Check output above.")
 
@@ -406,9 +414,8 @@ def main() -> None:
     missing_required_files = [
         path for path in (ALL_SOURCES_JSONL, CONTEXTUAL_NODES) if not path.exists()
     ]
-    needs_data_repo = (
-        (not args.skip_download and bool(missing_required_files))
-        or (not args.dry_run and not args.skip_upload and not args.skip_data_upload)
+    needs_data_repo = (not args.skip_download and bool(missing_required_files)) or (
+        not args.dry_run and not args.skip_upload and not args.skip_data_upload
     )
     needs_vector_repo = (
         not args.dry_run and not args.skip_upload and not args.skip_vector_rebuild

data/scraping_scripts/update_kb_wiki.py

@@ -64,7 +64,9 @@ def generated_block(content: str) -> str:
     return f"{AUTO_START}\n{content.rstrip()}\n{AUTO_END}"
 
 
-def write_generated_section(path: Path, header: str, generated: str, *, overwrite: bool) -> None:
+def write_generated_section(
+    path: Path, header: str, generated: str, *, overwrite: bool
+) -> None:
     path.parent.mkdir(parents=True, exist_ok=True)
     block = generated_block(generated)
     if not path.exists() or overwrite:
@@ -86,7 +88,9 @@ def top_sources(manifest: list[dict[str, Any]]) -> list[tuple[str, int]]:
     return sorted(counts.items(), key=lambda item: (-item[1], item[0]))
 
 
-def seed_index(kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite: bool) -> None:
+def seed_index(
+    kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite: bool
+) -> None:
     group_by_source = {
         str(row.get("source") or "unknown"): str(row.get("source_group") or "docs")
         for row in manifest
@@ -94,7 +98,9 @@ def seed_index(kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite: bool)
     source_lines = []
     for source, count in top_sources(manifest):
         folder = "courses" if group_by_source.get(source) == "courses" else "frameworks"
-        source_lines.append(f"- {source}: `wiki/{folder}/{source}.md` - {count} corpus pages")
+        source_lines.append(
+            f"- {source}: `wiki/{folder}/{source}.md` - {count} corpus pages"
+        )
     topic_lines = [f"- {topic}: `wiki/topics/{topic}.md`" for topic in TOPIC_KEYWORDS]
     header = """# AI Tutor KB Index
 
@@ -136,7 +142,9 @@ def seed_log(kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite: bool) -
     path.write_text(f"{existing}\n\n{entry.rstrip()}\n", encoding="utf-8")
 
 
-def seed_source_pages(kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite: bool) -> None:
+def seed_source_pages(
+    kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite: bool
+) -> None:
     by_source: dict[str, list[dict[str, Any]]] = defaultdict(list)
     for row in manifest:
         by_source[str(row.get("source") or "unknown")].append(row)
@@ -145,8 +153,7 @@ def seed_source_pages(kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite
         rows = sorted(rows, key=lambda item: str(item.get("title") or ""))
         sample = rows[:20]
         page_links = [
-            f"- {row.get('title')}: `{shell_path(row, kb_dir)}`"
-            for row in sample
+            f"- {row.get('title')}: `{shell_path(row, kb_dir)}`" for row in sample
         ]
         group = str(rows[0].get("source_group") or "docs")
         folder = "courses" if group == "courses" else "frameworks"
@@ -190,13 +197,12 @@ def matching_topic_rows(
     return [row for _score, row in scored[:limit]]
 
 
-def seed_topic_pages(kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite: bool) -> None:
+def seed_topic_pages(
+    kb_dir: Path, manifest: list[dict[str, Any]], *, overwrite: bool
+) -> None:
     for topic, keywords in TOPIC_KEYWORDS.items():
         rows = matching_topic_rows(manifest, keywords)
-        links = [
-            f"- {row.get('title')}: `{shell_path(row, kb_dir)}`"
-            for row in rows
-        ]
+        links = [f"- {row.get('title')}: `{shell_path(row, kb_dir)}`" for row in rows]
         header = f"""# {topic.replace("-", " ").title()}
 
 Use this page as a starting map for questions involving: {", ".join(keywords)}.

scripts/api.py

@@ -95,7 +95,9 @@ app.add_middleware(
 
 
 def sse_frame(payload: dict[str, Any] | str) -> str:
-    data = payload if isinstance(payload, str) else json.dumps(payload, ensure_ascii=False)
+    data = (
+        payload if isinstance(payload, str) else json.dumps(payload, ensure_ascii=False)
+    )
     return f"data: {data}\n\n"
 
 
@@ -138,11 +140,14 @@ def build_chat_request(payload: ApiChatRequest) -> ChatRequest:
 
     allowed_source_keys = set(AVAILABLE_SOURCES)
     requested_source_keys = payload.sourceKeys or list(DEFAULT_SELECTED_SOURCE_KEYS)
-    source_keys = tuple(
-        dict.fromkeys(
-            key for key in requested_source_keys if key in allowed_source_keys
+    source_keys = (
+        tuple(
+            dict.fromkeys(
+                key for key in requested_source_keys if key in allowed_source_keys
+            )
         )
-    ) or DEFAULT_SELECTED_SOURCE_KEYS
+        or DEFAULT_SELECTED_SOURCE_KEYS
+    )
     model_name = (payload.model or DEFAULT_MODEL_NAME).strip()
     if payload.enabledTools is None:
         enabled_tools = tuple(
@@ -226,7 +231,9 @@ class UIMessageStreamEncoder:
         if event.type == "reasoning_delta":
             if not self.active_reasoning_id:
                 self.active_reasoning_id = f"reasoning_{uuid4().hex[:8]}"
-                parts.append({"type": "reasoning-start", "id": self.active_reasoning_id})
+                parts.append(
+                    {"type": "reasoning-start", "id": self.active_reasoning_id}
+                )
             parts.append(
                 {
                     "type": "reasoning-delta",
@@ -281,7 +288,9 @@ class UIMessageStreamEncoder:
                 "group": str(event.data.get("group", "")),
             }
             if call_id:
-                self.source_matches_by_call_id.setdefault(call_id, []).append(source_data)
+                self.source_matches_by_call_id.setdefault(call_id, []).append(
+                    source_data
+                )
             parts.append(
                 {
                     "type": "source-url",
@@ -328,7 +337,7 @@ class UIMessageStreamEncoder:
                         "id": self.text_block_id,
                         "delta": answer,
                     }
-            )
+                )
             if self.text_block_id:
                 parts.append({"type": "text-end", "id": self.text_block_id})
             parts.append({"type": "finish-step"})

scripts/chat_service.py

@@ -64,7 +64,9 @@ THOUGHTS_BLOCK_START = "<!-- GEMINI_THOUGHTS_START -->"
 THOUGHTS_BLOCK_END = "<!-- GEMINI_THOUGHTS_END -->"
 ANSWER_HEADER = "**Answer**"
 LEGACY_THOUGHTS_DETAILS_OPEN = "<details><summary>Gemini thoughts</summary>"
-LEGACY_THOUGHTS_DETAILS_OPEN_EXPANDED = "<details open><summary>Gemini thoughts</summary>"
+LEGACY_THOUGHTS_DETAILS_OPEN_EXPANDED = (
+    "<details open><summary>Gemini thoughts</summary>"
+)
 CHECKPOINTER = InMemorySaver()
 _RETRIEVER_INIT_LOCK = Lock()
 KB_TOOL_NAMES = ("run_kb_command",)
@@ -403,7 +405,9 @@ def collect_retrieval_source_matches(payload: str) -> list[SourceMatch]:
     return matches
 
 
-def _record_evidence(target: dict[str, SourceMatch], matches: list[SourceMatch]) -> None:
+def _record_evidence(
+    target: dict[str, SourceMatch], matches: list[SourceMatch]
+) -> None:
     for match in matches:
         key = source_match_key(match)
         existing = target.get(key)
@@ -704,9 +708,7 @@ class SourcePreferenceMiddleware(AgentMiddleware):
             label = SOURCE_KEY_TO_LABEL.get(key, key)
             group = "courses" if key in COURSE_SOURCE_KEYS else "docs"
             wiki_dir = "courses" if key in COURSE_SOURCE_KEYS else "frameworks"
-            lines.append(
-                f"- {label}: `raw/{group}/{key}/`, `wiki/{wiki_dir}/{key}.md`"
-            )
+            lines.append(f"- {label}: `raw/{group}/{key}/`, `wiki/{wiki_dir}/{key}.md`")
         lines.append("")
         lines.append(
             "Only branch out to other KB sources if these don't have the answer."
@@ -1098,9 +1100,7 @@ async def stream_chat(request: ChatRequest) -> AsyncIterator[ChatEvent]:
 
                 step = str(metadata.get("langgraph_step", ""))
                 if include_reasoning:
-                    thought_text = "\n\n".join(
-                        extract_thought_summaries(token.content)
-                    )
+                    thought_text = "\n\n".join(extract_thought_summaries(token.content))
                     if thought_text:
                         yield ChatEvent(
                             "reasoning_delta",
@@ -1157,9 +1157,7 @@ async def stream_chat(request: ChatRequest) -> AsyncIterator[ChatEvent]:
                             "call_id": google_search_call_id,
                             "tool_name": GOOGLE_SEARCH_TOOL_NAME,
                             "args": {
-                                "query": "; ".join(new_queries)
-                                if new_queries
-                                else ""
+                                "query": "; ".join(new_queries) if new_queries else ""
                             },
                             "args_text": "; ".join(new_queries),
                         },
@@ -1245,9 +1243,7 @@ async def stream_chat(request: ChatRequest) -> AsyncIterator[ChatEvent]:
                             "call_id": google_search_call_id,
                             "tool_name": GOOGLE_SEARCH_TOOL_NAME,
                             "args": {
-                                "query": "; ".join(new_queries)
-                                if new_queries
-                                else ""
+                                "query": "; ".join(new_queries) if new_queries else ""
                             },
                             "args_text": "; ".join(new_queries),
                         },
@@ -1311,9 +1307,7 @@ async def stream_chat(request: ChatRequest) -> AsyncIterator[ChatEvent]:
             output_text = "Google search ran but returned no grounding results."
         else:
             plural = "" if google_search_match_count == 1 else "s"
-            output_text = (
-                f"Google search returned {google_search_match_count} web result{plural}."
-            )
+            output_text = f"Google search returned {google_search_match_count} web result{plural}."
         yield ChatEvent(
             "tool_call_completed",
             {

scripts/chroma_rag.py

@@ -515,8 +515,7 @@ def heading_aware_markdown_chunks(
         for split_unit in split_units:
             size = unit_size(split_unit.text)
             heading_changed = (
-                current_heading_path
-                and split_unit.heading_path != current_heading_path
+                current_heading_path and split_unit.heading_path != current_heading_path
             )
             would_exceed = current_parts and current_size + size > budget_size
             if heading_changed or would_exceed:
@@ -591,7 +590,7 @@ def build_chunk_records(
             }
             chunk_records.append(
                 ChunkRecord(
-                    chunk_id=f'{document["doc_id"]}:{index}',
+                    chunk_id=f"{document['doc_id']}:{index}",
                     doc_id=document["doc_id"],
                     text=chunk.text,
                     metadata=metadata,
@@ -836,7 +835,7 @@ def _wait_for_cohere_retry(
     if retry_after is not None:
         delay = retry_after
     else:
-        delay = min(window_seconds, max(15.0, 2.0 ** attempt))
+        delay = min(window_seconds, max(15.0, 2.0**attempt))
 
     time.sleep(delay + random.uniform(0.5, 2.0))
 
@@ -948,7 +947,10 @@ def embed_texts(
                     )
                     break
                 except Exception as exc:
-                    if not _is_cohere_rate_limit_error(exc) or attempt == retry_attempts:
+                    if (
+                        not _is_cohere_rate_limit_error(exc)
+                        or attempt == retry_attempts
+                    ):
                         raise
                     _wait_for_cohere_retry(exc, attempt, window_seconds)
 
@@ -1096,7 +1098,9 @@ def reciprocal_rank_fusion(
             current = representatives.get(key)
             if current is None:
                 representatives[key] = result
-            elif result.retrieval_method == "bm25" and current.retrieval_method != "bm25":
+            elif (
+                result.retrieval_method == "bm25" and current.retrieval_method != "bm25"
+            ):
                 representatives[key] = result
             elif (
                 result.retrieval_method == current.retrieval_method
@@ -1357,7 +1361,9 @@ def ensure_parent_dir(path: str) -> None:
     Path(path).parent.mkdir(parents=True, exist_ok=True)
 
 
-def save_document_dict(document_dict: dict[str, dict[str, Any]], output_file: str) -> None:
+def save_document_dict(
+    document_dict: dict[str, dict[str, Any]], output_file: str
+) -> None:
     ensure_parent_dir(output_file)
     with open(output_file, "wb") as handle:
         pickle.dump(document_dict, handle)

scripts/gradio_presenter.py

@@ -124,7 +124,9 @@ def summarize_tool_result(event: ChatEvent) -> str:
         if not matches:
             try:
                 payload = json.loads(str(event.data.get("output_text", "")))
-                matches = payload.get("matches", []) if isinstance(payload, dict) else []
+                matches = (
+                    payload.get("matches", []) if isinstance(payload, dict) else []
+                )
                 match_count = len(matches)
             except json.JSONDecodeError:
                 matches = []
@@ -221,11 +223,7 @@ def render_activity_block(events: list[ActivityEvent]) -> str:
         return ""
 
     rendered_sections = "\n\n".join(section for section in sections if section)
-    return (
-        f"{ACTIVITY_BLOCK_START}\n\n"
-        f"{rendered_sections}\n\n"
-        f"{ACTIVITY_BLOCK_END}"
-    )
+    return f"{ACTIVITY_BLOCK_START}\n\n{rendered_sections}\n\n{ACTIVITY_BLOCK_END}"
 
 
 def format_sources(matches_by_doc_id: dict[str, dict[str, Any]]) -> str:
@@ -273,7 +271,9 @@ class GradioPresenterState:
     activity_events: list[ActivityEvent] = field(default_factory=list)
     answer_chunks: list[str] = field(default_factory=list)
     completed_answer: str = ""
-    tool_matches_by_call_id: dict[str, list[dict[str, Any]]] = field(default_factory=dict)
+    tool_matches_by_call_id: dict[str, list[dict[str, Any]]] = field(
+        default_factory=dict
+    )
 
     def apply(self, event: ChatEvent) -> None:
         if event.type == "thread_started":

scripts/kb_manifest.py

@@ -62,7 +62,12 @@ def load_manifest_entries(kb_dir: str = KB_DIR) -> tuple[KbManifestEntry, ...]:
 
 def manifest_indexes(
     kb_dir: str = KB_DIR,
-) -> tuple[dict[str, KbManifestEntry], dict[str, KbManifestEntry], dict[str, KbManifestEntry], dict[str, KbManifestEntry]]:
+) -> tuple[
+    dict[str, KbManifestEntry],
+    dict[str, KbManifestEntry],
+    dict[str, KbManifestEntry],
+    dict[str, KbManifestEntry],
+]:
     by_doc_id: dict[str, KbManifestEntry] = {}
     by_url: dict[str, KbManifestEntry] = {}
     by_path: dict[str, KbManifestEntry] = {}
@@ -81,7 +86,9 @@ def manifest_indexes(
     return by_doc_id, by_url, by_path, by_title
 
 
-def source_match_from_manifest(entry: KbManifestEntry, *, score: float = 1.0) -> SourceMatch:
+def source_match_from_manifest(
+    entry: KbManifestEntry, *, score: float = 1.0
+) -> SourceMatch:
     return SourceMatch(
         doc_id=entry.doc_id,
         title=entry.title,
@@ -89,7 +96,9 @@ def source_match_from_manifest(entry: KbManifestEntry, *, score: float = 1.0) ->
         source_key=entry.source,
         source_label=SOURCE_KEY_TO_LABEL.get(entry.source, entry.source),
         score=score,
-        group="courses" if entry.source in COURSE_SOURCE_KEYS else entry.source_group or "docs",
+        group="courses"
+        if entry.source in COURSE_SOURCE_KEYS
+        else entry.source_group or "docs",
     )
 
 
@@ -167,7 +176,9 @@ def citation_dedupe_key(match: SourceMatch) -> str:
     return normalize_url(match.url) or match.doc_id or match.title
 
 
-def source_match_payload(match: SourceMatch, *, message_id: str, call_id: str = "") -> dict[str, Any]:
+def source_match_payload(
+    match: SourceMatch, *, message_id: str, call_id: str = ""
+) -> dict[str, Any]:
     payload: dict[str, Any] = {
         "message_id": message_id,
         "doc_id": match.doc_id,

scripts/kb_shell.py

@@ -41,7 +41,9 @@ RG_FLAG_OPTIONS = frozenset(
         "--no-ignore",
     }
 )
-GREP_FLAG_OPTIONS = frozenset({"-i", "--ignore-case", "-w", "--word-regexp", "-F", "-E"})
+GREP_FLAG_OPTIONS = frozenset(
+    {"-i", "--ignore-case", "-w", "--word-regexp", "-F", "-E"}
+)
 LS_FLAG_OPTIONS = frozenset({"-1", "-a", "-l", "-la", "-al"})
 WC_FLAG_OPTIONS = frozenset({"-l", "-w", "-c", "-m"})
 
@@ -117,7 +119,9 @@ def _is_broad_raw_path(path: str) -> bool:
     return normalized in {".", "raw", "raw/courses", "raw/docs"}
 
 
-def _reject_unbounded_raw_search(command: str, paths: list[str], has_max_count: bool) -> None:
+def _reject_unbounded_raw_search(
+    command: str, paths: list[str], has_max_count: bool
+) -> None:
     if has_max_count:
         return
     if any(_is_broad_raw_path(path) for path in paths):
@@ -230,7 +234,9 @@ def _build_find(tokens: list[str], root: Path) -> list[str]:
             argv.extend([token, _safe_pattern_value(tokens[idx + 1], token)])
             idx += 2
             continue
-        raise KbCommandError("`find` supports only path, -maxdepth, -mindepth, -type, -name, and -iname")
+        raise KbCommandError(
+            "`find` supports only path, -maxdepth, -mindepth, -type, -name, and -iname"
+        )
     return argv
 
 
@@ -291,7 +297,9 @@ def _build_wc(tokens: list[str], root: Path) -> list[str]:
     return [*argv, *paths]
 
 
-def build_kb_command_argv(command: str, *, root: Path | None = None) -> tuple[list[str], Path]:
+def build_kb_command_argv(
+    command: str, *, root: Path | None = None
+) -> tuple[list[str], Path]:
     resolved_root = _resolve_root(root)
     try:
         tokens = shlex.split(command)
@@ -304,7 +312,9 @@ def build_kb_command_argv(command: str, *, root: Path | None = None) -> tuple[li
     executable = tokens[0]
     if executable not in SUPPORTED_COMMANDS:
         supported = ", ".join(sorted(SUPPORTED_COMMANDS))
-        raise KbCommandError(f"Unsupported KB command `{executable}`. Supported: {supported}")
+        raise KbCommandError(
+            f"Unsupported KB command `{executable}`. Supported: {supported}"
+        )
 
     builders = {
         "rg": _build_rg,
@@ -361,7 +371,9 @@ def run_kb_command(
         exit_code = 124
         stdout = exc.stdout if isinstance(exc.stdout, str) else ""
         stderr = exc.stderr if isinstance(exc.stderr, str) else ""
-        stderr = (stderr + "\n" if stderr else "") + f"Command timed out after {timeout}s."
+        stderr = (
+            stderr + "\n" if stderr else ""
+        ) + f"Command timed out after {timeout}s."
 
     stdout, stdout_truncated = _truncate(stdout, output_limit)
     stderr, stderr_truncated = _truncate(stderr, min(output_limit, 8_000))

scripts/probe_gemini_cache.py

@@ -69,12 +69,16 @@ def main() -> None:
     client = genai.Client(api_key=api_key)
 
     print(f"Model: {MODEL}")
-    print(f"Stable prefix (system + LONG_CONTEXT) repeats. Only the final user_text differs.")
+    print(
+        "Stable prefix (system + LONG_CONTEXT) repeats. Only the final user_text differs."
+    )
     print()
 
     for i in range(3):
-        result = call(client, user_text=f"Question {i+1}: what is RAG in one sentence?")
-        print(f"Call {i+1}: {result}")
+        result = call(
+            client, user_text=f"Question {i + 1}: what is RAG in one sentence?"
+        )
+        print(f"Call {i + 1}: {result}")
         time.sleep(0.5)
 
 

scripts/prompts.py

@@ -205,7 +205,9 @@ def build_system_prompt(model_name: str, enabled_tools: tuple[str, ...]) -> str:
         KB_USAGE_SECTION,
     ]
     if usage_sections:
-        parts.append("## When to use web search / URL reading\n\n" + "\n\n".join(usage_sections))
+        parts.append(
+            "## When to use web search / URL reading\n\n" + "\n\n".join(usage_sections)
+        )
         parts.append(
             "Prefer `retrieve_tutor_context` first when the question is clearly about\n"
             "course material. Combine tools when it helps (e.g. retrieve corpus\n"

tests/test_api.py

@@ -56,7 +56,10 @@ class ApiTestCase(unittest.TestCase):
             any(source["selectedByDefault"] for source in retrieval["sources"])
         )
         self.assertTrue(
-            all(source["group"] in {"courses", "docs"} for source in retrieval["sources"])
+            all(
+                source["group"] in {"courses", "docs"}
+                for source in retrieval["sources"]
+            )
         )
         # Gemini is the default model, so web search + url reading are present.
         tool_keys = {tool["key"] for tool in tools}
@@ -124,8 +127,13 @@ class ApiTestCase(unittest.TestCase):
                     "output_text": "payload",
                 },
             )
-            yield ChatEvent("text_delta", {"message_id": "message_1", "text": "RAG combines "})
-            yield ChatEvent("text_delta", {"message_id": "message_1", "text": "retrieval with generation."})
+            yield ChatEvent(
+                "text_delta", {"message_id": "message_1", "text": "RAG combines "}
+            )
+            yield ChatEvent(
+                "text_delta",
+                {"message_id": "message_1", "text": "retrieval with generation."},
+            )
             yield ChatEvent(
                 "message_completed",
                 {
@@ -189,7 +197,9 @@ class ApiTestCase(unittest.TestCase):
 
         with patch("scripts.api.stream_chat", broken_stream_chat):
             with TestClient(app) as client:
-                with client.stream("POST", "/api/chat", json={"query": "Hello"}) as response:
+                with client.stream(
+                    "POST", "/api/chat", json={"query": "Hello"}
+                ) as response:
                     body = "".join(response.iter_text())
 
         self.assertEqual(response.status_code, 200)
@@ -208,7 +218,9 @@ class ApiTestCase(unittest.TestCase):
     def test_chat_stream_restarts_reasoning_after_tool_activity(self) -> None:
         async def fake_stream_chat(_request):
             yield ChatEvent("message_started", {"message_id": "message_1"})
-            yield ChatEvent("reasoning_delta", {"message_id": "message_1", "text": "First thought"})
+            yield ChatEvent(
+                "reasoning_delta", {"message_id": "message_1", "text": "First thought"}
+            )
             yield ChatEvent(
                 "tool_call_started",
                 {
@@ -227,8 +239,12 @@ class ApiTestCase(unittest.TestCase):
                     "output_text": "payload",
                 },
             )
-            yield ChatEvent("reasoning_delta", {"message_id": "message_1", "text": "Second thought"})
-            yield ChatEvent("text_delta", {"message_id": "message_1", "text": "Final answer"})
+            yield ChatEvent(
+                "reasoning_delta", {"message_id": "message_1", "text": "Second thought"}
+            )
+            yield ChatEvent(
+                "text_delta", {"message_id": "message_1", "text": "Final answer"}
+            )
             yield ChatEvent(
                 "message_completed",
                 {
@@ -239,7 +255,9 @@ class ApiTestCase(unittest.TestCase):
 
         with patch("scripts.api.stream_chat", fake_stream_chat):
             with TestClient(app) as client:
-                with client.stream("POST", "/api/chat", json={"query": "Hello"}) as response:
+                with client.stream(
+                    "POST", "/api/chat", json={"query": "Hello"}
+                ) as response:
                     body = "".join(response.iter_text())
 
         self.assertEqual(response.status_code, 200)
@@ -249,7 +267,9 @@ class ApiTestCase(unittest.TestCase):
 
         self.assertEqual(part_types.count("reasoning-start"), 2)
         self.assertEqual(part_types.count("reasoning-end"), 2)
-        self.assertLess(part_types.index("tool-input-start"), part_types.index("text-start"))
+        self.assertLess(
+            part_types.index("tool-input-start"), part_types.index("text-start")
+        )
 
 
 LIVE_API_E2E = pytest.mark.skipif(

tests/test_build_kb_artifacts.py

@@ -89,7 +89,10 @@ Use `LoraConfig` with `get_peft_model`.
     assert manifest_rows[0]["original_path"] == docs_page.as_posix()
 
     markdown_path = Path(manifest_rows[0]["path"])
-    assert markdown_path == output_dir / "raw" / "docs" / "temp_docs" / "package_reference" / "lora.mdx"
+    assert (
+        markdown_path
+        == output_dir / "raw" / "docs" / "temp_docs" / "package_reference" / "lora.mdx"
+    )
     markdown = markdown_path.read_text(encoding="utf-8")
     assert 'doc_id: "temp_docs:docs-package-reference-lora"' in markdown
     assert 'source_path: "package_reference/lora.mdx"' in markdown

tests/test_chat_service.py

@@ -209,8 +209,12 @@ class ChatServiceTestCase(unittest.TestCase):
         )
         # Enabling Gemini web tools adds exactly one toggle-specific middleware
         # (GeminiServerSideToolsMiddleware) on top of the shared base middlewares.
-        web_middleware = [type(m).__name__ for m in created_agents[0].kwargs["middleware"]]
-        plain_middleware = [type(m).__name__ for m in created_agents[1].kwargs["middleware"]]
+        web_middleware = [
+            type(m).__name__ for m in created_agents[0].kwargs["middleware"]
+        ]
+        plain_middleware = [
+            type(m).__name__ for m in created_agents[1].kwargs["middleware"]
+        ]
         self.assertIn("GeminiServerSideToolsMiddleware", web_middleware)
         self.assertNotIn("GeminiServerSideToolsMiddleware", plain_middleware)
         self.assertEqual(len(web_middleware), len(plain_middleware) + 1)
@@ -382,11 +386,7 @@ class ChatServiceTestCase(unittest.TestCase):
             if event.type == "tool_call_completed"
             and event.data.get("tool_name") == "run_kb_command"
         ]
-        source_matches = [
-            event
-            for event in events
-            if event.type == "source_match"
-        ]
+        source_matches = [event for event in events if event.type == "source_match"]
 
         self.assertEqual(started[0].data["args_text"], "rg LoraConfig raw")
         self.assertIn("rg LoraConfig raw", completed[0].data["output_text"])

tests/test_chroma_rag.py

@@ -86,7 +86,9 @@ After the example.
         ]
         index = BM25Index.build(records)
 
-        hits = index.search("AutoModel.from_pretrained", allowed_sources=["transformers"])
+        hits = index.search(
+            "AutoModel.from_pretrained", allowed_sources=["transformers"]
+        )
 
         self.assertEqual([record.chunk_id for record, _score in hits], ["a"])
         self.assertEqual(index.search("AutoModel", allowed_sources=["langchain"]), [])
@@ -159,7 +161,9 @@ After the example.
 
         self.assertGreaterEqual(len(records), 1)
         setup_record = next(
-            record for record in records if record.metadata["heading_path"] == "Guide > Setup"
+            record
+            for record in records
+            if record.metadata["heading_path"] == "Guide > Setup"
         )
         self.assertIn("raw_text", setup_record.metadata)
         self.assertIn("Title: Guide", setup_record.text)

tests/test_gradio_kb_e2e.py

@@ -172,7 +172,9 @@ def test_live_gradio_curl_can_use_retrieval_and_shell() -> None:
                             "Use both retrieve_tutor_context and run_kb_command to answer: how does PEFT configure LoRA with LoraConfig?",
                             [],
                             ["PEFT Docs", "Transformers Docs"],
-                            os.getenv("LIVE_GRADIO_E2E_MODEL", "google-genai:gemini-3.5-flash"),
+                            os.getenv(
+                                "LIVE_GRADIO_E2E_MODEL", "google-genai:gemini-3.5-flash"
+                            ),
                             "",
                             False,
                             False,

tests/test_gradio_presenter.py

@@ -56,8 +56,7 @@ class GradioPresenterTestCase(unittest.TestCase):
             {
                 "tool_name": "retrieve_tutor_context",
                 "output_text": (
-                    '{"matches": [{"title": "LoRA", '
-                    '"source_label": "PEFT Docs"}]}'
+                    '{"matches": [{"title": "LoRA", "source_label": "PEFT Docs"}]}'
                 ),
             },
         )

tests/test_retire_source_workflow.py

@@ -6,7 +6,7 @@ from data.scraping_scripts.retire_source_workflow import (
 
 
 def test_remove_source_from_registry_text_removes_active_source_entries():
-    registry = '''SOURCE_CONFIGS = {
+    registry = """SOURCE_CONFIGS = {
     "openai_cookbooks": {
         "output_file": "data/openai_cookbooks_data.jsonl",
         "nested": {"keep": "balanced"},
@@ -32,7 +32,7 @@ UI_SOURCE_KEYS = (
     "openai_cookbooks",
     "langchain",
 )
-'''
+"""
 
     updated, changed = remove_source_from_registry_text(
         registry,

tests/test_update_kb_wiki.py

@@ -37,7 +37,7 @@ def test_update_kb_wiki_seeds_navigation_pages(tmp_path: Path) -> None:
                 "tokens": 300,
                 "retrieve_doc": True,
                 "content": "# Lesson 1\n\nAgent lesson.",
-            }
+            },
         ],
     )
     build_kb_artifacts(input_file, kb_dir)
@@ -61,7 +61,9 @@ def test_update_kb_wiki_seeds_navigation_pages(tmp_path: Path) -> None:
     assert "lookup_tutor_symbol" not in topic
 
 
-def test_update_kb_wiki_preserves_authored_topic_content_and_appends_log(tmp_path: Path) -> None:
+def test_update_kb_wiki_preserves_authored_topic_content_and_appends_log(
+    tmp_path: Path,
+) -> None:
     input_file = tmp_path / "all_sources_data.jsonl"
     kb_dir = tmp_path / "kb"
     write_jsonl(