Skip to main content

netget/llm/
response_handler.rs

1//! Common LLM response handling logic
2
3use super::ollama_client::LlmResponse;
4use crate::state::app_state::AppState;
5use std::sync::Arc;
6use tracing::info;
7
8/// Result of processing an LLM response - contains flags and data for protocol-specific handling
9pub struct ProcessedResponse {
10    /// Data to send over the connection
11    pub output: Option<String>,
12    /// Whether to close the connection
13    pub close_connection: bool,
14    /// Whether to wait for more data
15    pub wait_for_more: bool,
16    /// Whether to shutdown the server
17    pub shutdown_server: bool,
18}
19
20/// Handle common LLM response actions (memory updates, logging)
21/// Returns the processed response with flags for protocol-specific handling
22pub async fn handle_llm_response(
23    response: LlmResponse,
24    app_state: &Arc<AppState>,
25) -> ProcessedResponse {
26    // Get first server ID for memory updates (legacy support)
27    let server_id = app_state.get_first_server_id().await;
28
29    // Handle memory updates
30    if let Some(set_mem) = response.set_memory {
31        if let Some(sid) = server_id {
32            app_state.set_memory(sid, set_mem).await;
33        }
34    }
35    if let Some(append_mem) = response.append_memory {
36        if let Some(sid) = server_id {
37            let current = app_state.get_memory(sid).await.unwrap_or_default();
38            let new_memory = if current.is_empty() {
39                append_mem
40            } else {
41                format!("{}\n{}", current, append_mem)
42            };
43            app_state.set_memory(sid, new_memory).await;
44        }
45    }
46
47    // Handle log messages
48    if let Some(log_msg) = response.log_message {
49        info!("{}", log_msg);
50    }
51
52    ProcessedResponse {
53        output: response.output,
54        close_connection: response.close_connection,
55        wait_for_more: response.wait_for_more,
56        shutdown_server: response.shutdown_server,
57    }
58}