Fix API contract bugs and refactor Model to persist knowledge base state across requests

[happy4pizza]

This PR fixes several bugs introduced during the MLflow 2.x to 3.x migration and completes an architectural refactor to make the pipeline functional end to end.

Bug Fixes

• Fixed update_kb command not being handled in predict()- the model was treating it as a search query, running the full RAG pipeline against it, and returning a response with no status key, so kb_ready never flipped to True in the UI
• Fixed query payload shape - chat queries were sending force_regenerate as a top-level column which violated the MLflow schema (query and payload required). All query data now lives inside the payload JSON field
• Fixed KeyError: 'config' on every chat query - the query path was calling _build_transient_kb() which expected payload["config"] and payload["secrets"], but chat payloads only contain the question
• Fixed TypeError when displaying metrics - added a None guard before formatting generation time, faithfulness, and relevance so the UI doesn't crash if any metric is missing
• Fixed missing logger in main.py -logger was referenced before being defined

Architectural Refactor

• Added _initialize_kb() which builds the KB (knowledge base) once on update_kb and stores all state (bm25_index, doc_map, rag_pipeline) on the Model instance
• Removed _build_transient_kb() from the query path entirely - regular queries now call self.rag_pipeline.generate() directly against the already-built KB
• Added a warmup query at the end of _initialize_kb() to pre-compile Triton CUDA kernels so the first real user query returns in seconds rather than minutes
• Added a guard that returns an error if the user queries before syncing
• Used a persistent temp directory for the wiki clone so image paths remain valid across requests
• Added automatic model installation - the notebook now checks for the model weights at the expected datafabric path on startup. If not found, it downloads automatically from the Hugging Face Hub before proceeding.
• Added image_files to the API response - the model now returns the original image filenames (e.g. image-68e01678-....png) alongside the existing base64 blobs in a new image_files field. This enables external evaluation tooling to validate image retrieval against expected sources without requiring direct access to the internal vector store, and unblocks image hit rate, precision, and recall metrics in the evaluation harness

Documentation

• Updated README.md to reflect the current AI Studio deployment workflow, including corrected setup steps, and hardware requirements aligned with the actual CUDA version in use

Result

Query latency drops from 3+ minutes (due to wiki cloning, re-embedding, and Triton JIT compilation on every request) to seconds. The KB rebuilds only when the user explicitly clicks "Sync", and Triton kernels are pre-compiled during sync so the first real query is fast approx. 10 seconds on testing.

[github-actions]

Welcome to our 'AI-Blueprints' repo 🚀😁🚀