kubeflow / kubeflow/docs-agent

feat: Implement Token-Aware Sliding Window Memory Manager and Global Embeddings for HTTP API

Open
#136 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
42
Forks
111
Avg merge
6d 23m
Merged PRs (30d)
2

Description

### Background & Context
As the Kubeflow Docs Agent scales to handle more complex, multi-turn diagnostic conversations, the payload sent to the LLM backend (via KServe) grows exponentially. The current architecture in `server-https/app.py` blindly appends all user inputs, assistant responses, and massive `milvus_search` tool-call results into a single, unbounded `messages` list.

### The Core Problems

**1. Unbounded Context Window & API Crashes**
Every LLM has a strict maximum context window (e.g., 8192 tokens for Llama 3 models). Because there is currently no eviction strategy or context sanitation layer, long conversations or tool calls that return large `top_k` text blocks from Milvus will inevitably breach this limit. When this happens, the KServe API throws a fatal `HTTP 400/500 Bad Request` error, immediately crashing the application and permanently breaking the user's current session.

**2. Synchronous ML Model Instantiation (Performance Bottleneck)**
Similar to the issue previously identified in the WebSocket server, the HTTP API currently re-initializes `SentenceTransformer(EMBEDDING_MODEL)` inside the synchronous `milvus_search` function. This forces the server to load a multi-hundred-megabyte ML model into memory from scratch on *every single tool execution*, resulting in severe latency spikes and a high risk of Out-Of-Memory (OOM) crashes under concurrent user load.

### Proposed Architecture & Implementation Strategy

To resolve both the stability and performance bottlenecks, I propose introducing a resilient memory management layer and porting the global singleton pattern to the REST API.

**Part 1: `ContextWindowManager` (Sliding Window Algorithm)**
I will implement an Object-Oriented context management class that intercepts the payload before it is streamed to the LLM.
* **Token Heuristics:** The manager will utilize a lightweight character-to-token heuristic to calculate the exact footprint of the `messages` array prior to transmission.
* **FIFO Eviction:** If the payload approaches a predefined safe threshold (e.g., ~7000 tokens), the manager will dynamically pop the oldest conversational turns from the front of the queue.
* **Absolute System Preservation:** The `system` prompt will be permanently preserved at index `[0]` to ensure the agent never loses its foundational instructions.
* **Orphaned Tool Protection (Edge Case Handling):** LLM APIs enforce strict conversational grammar. If the algorithm evicts an `assistant` message that contains `tool_calls`, it must simultaneously detect and evict the subsequent `tool` response. Failure to do so results in "orphaned" tool payloads that KServe will reject. This manager will include explicit logic to prevent orphan corruption.

**Part 2: Global Encoder Singleton**
I will refactor the embedding model initialization out of the `milvus_search` execution path and into the global scope. The `GLOBAL_ENCODER` will load exactly once at server startup, allowing the tool to reference the existing model in memory, reducing tool execution latency from seconds to milliseconds.

### Expected Impact
* **100% Conversational Uptime:** Users can maintain infinitely long conversations without ever triggering a context overflow crash.
* **Massive Latency Reduction:** Search operations will execute exponentially faster.
* **Enterprise Stability:** The FastAPI event loop will be protected from heavy, repetitive I/O blocking.

I already have this implementation fully tested locally and will append these changes to my existing open Pull Request shortly to patch both the WebSocket and HTTP endpoints simultaneously.

Contributor guide

Open the contributing guide

Research direction

Start in server-https/app.py, focusing on the messages handling and the synchronous milvus_search embedding path. Review the proposed ContextWindowManager behavior and global encoder initialization; done means bounded context preserves the system prompt and valid tool-call pairs, while the embedding model loads once rather than per search.

Written by the indexing model from the issue text.

Assessment

Tech stack
fastapi, python
Domain
api, backend, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.