agentic office

network-poc/hub/Cargo.toml

@@ -5,7 +5,7 @@ edition = "2024"
 
 [dependencies]
 axum = { version = "0.7.4", features = ["ws", "macros"] }
-tokio = { version = "1.36.0", features = ["full"] }
+tokio = { version = "1.36.0", features = ["full", "sync"] }
 tower-http = { version = "0.5.2", features = ["fs", "cors", "trace"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"

network-poc/hub/src/main.rs

@@ -339,6 +339,7 @@ async fn main() {
         .route("/api/sessions", get(api_sessions))
         .route("/api/pairs", get(api_pairs))
         .route("/api/stats", get(api_stats))
+        .route("/api/v1/chat/completions", axum::routing::post(api_chat_completions))
         .route("/admin", get(admin_page))
         .nest_service("/", ServeDir::new(std::env::var("STATIC_DIR").unwrap_or_else(|_| "../static".to_string())))
         .with_state(state);
@@ -820,3 +821,50 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, ip: IpAddr) {
     broadcast_stats(&state).await;
     sender_task.abort();
 }
+#[derive(serde::Deserialize)]
+struct ChatCompletionRequest {
+    model: String,
+    prompt: String,
+    task_id: String,
+}
+
+#[derive(serde::Serialize)]
+struct ChatCompletionResponse {
+    response: String,
+    model: String,
+    tokens_generated: u64,
+}
+
+async fn api_chat_completions(
+    axum::extract::State(state): axum::extract::State<Arc<AppState>>,
+    axum::Json(payload): axum::Json<ChatCompletionRequest>,
+) -> axum::response::Response {
+    let msg = serde_json::json!({
+        "type": "llm_prompt",
+        "prompt": payload.prompt,
+        "model": payload.model,
+        "task_id": payload.task_id,
+    });
+    
+    let mut rx = state.stats_tx.subscribe();
+    let _ = state.stats_tx.send(msg.to_string());
+    
+    while let Ok(msg_str) = rx.recv().await {
+        if let Ok(v) = serde_json::from_str::<serde_json::Value>(&msg_str) {
+            if v["type"].as_str() == Some("llm_done") {
+                if let Some(tid) = v["task_id"].as_str() {
+                    if tid == payload.task_id {
+                        let res = ChatCompletionResponse {
+                            response: v["response"].as_str().unwrap_or("").to_string(),
+                            model: v["model"].as_str().unwrap_or("").to_string(),
+                            tokens_generated: v["tokens_generated"].as_u64().unwrap_or(0),
+                        };
+                        return axum::Json(res).into_response();
+                    }
+                }
+            }
+        }
+    }
+    
+    (axum::http::StatusCode::INTERNAL_SERVER_ERROR, "Network Error").into_response()
+}