initial commit - signoff

Signed-off-by: michaelfeil <me@michaelfeil.eu>

Author: michaelfeil

lib/llm/src/kv_router.rs

@@ -376,19 +376,45 @@ impl KvRouter {
     }
 }
 
-// NOTE: this would not be usable for now, should deprecate
+// NOTE: KVRouter works like a PushRouter, 
+// but without the reverse proxy functionality, but based on contract of 3 request types
 #[async_trait]
 impl AsyncEngine<SingleIn<RouterRequest>, ManyOut<Annotated<RouterResponse>>, Error> for KvRouter {
     async fn generate(
         &self,
         request: SingleIn<RouterRequest>,
     ) -> Result<ManyOut<Annotated<RouterResponse>>> {
         let (request, ctx) = request.into_parts();
-        let (worker_id, _) = self
-            .find_best_match(ctx.id(), &request.tokens, None, true)
-            .await?;
+        let context_id = ctx.context().id().to_string();
+        // Handle different request types
+        let response = match request {
+            RouterRequest::New { tokens } => {
+                let (worker_id, overlap_blocks) = self
+                    .find_best_match(&context_id, &tokens, None, true)
+                    .await?;
+
+                RouterResponse {
+                    worker_id,
+                    overlap_blocks,
+                }
+            }
+            RouterRequest::MarkPrefill {} => {
+                self.mark_prefill_completed(&context_id).await;
+
+                RouterResponse {
+                    worker_id: -1,
+                    overlap_blocks: 0,
+                }
+            }
+            RouterRequest::MarkFree {} => {
+                self.free(&context_id).await;
+                RouterResponse {
+                    worker_id: -1,
+                    overlap_blocks: 0,
+                }
+            }
+        };
 
-        let response = RouterResponse { worker_id };
         let response = Annotated::from_data(response);
         let stream = stream::iter(vec![response]);
         Ok(ResponseStream::new(Box::pin(stream), ctx.context()))

lib/llm/src/kv_router/protocols.rs

@@ -5,14 +5,28 @@ use crate::tokens::{SequenceHash, Token};
 use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
-#[derive(Debug, Clone, Serialize, Deserialize, Default)]
-pub struct RouterRequest {
-    pub tokens: Vec<Token>,
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "method", rename_all = "snake_case")]
+pub enum RouterRequest {
+    // ini
+    #[serde(rename = "new")]
+    New {
+        tokens: Vec<Token>,
+    },
+    MarkPrefill {},
+    MarkFree {},
+}
+
+impl Default for RouterRequest {
+    fn default() -> Self {
+        RouterRequest::New { tokens: vec![] }
+    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct RouterResponse {
     pub worker_id: i64,
+    pub overlap_blocks: u32,
 }
 
 #[derive(Debug)]