llm request tracking, user context cleanup, minor fixes

This commit is contained in:
2026-07-11 00:23:14 +01:00
parent 5936cf0b2e
commit 587958ffe0
24 changed files with 402 additions and 347 deletions
+29 -53
View File
@@ -1,11 +1,12 @@
//! Transparent logging wrapper for any [`ChatbotClient`].
//!
//! [`LoggingChatbotClient`] intercepts every `chat_with_tools` call, captures
//! the raw HTTP request/response from the inner provider via `chat_with_tools_raw`,
//! then persists a row to `llm_requests` asynchronously (fire-and-forget).
//! [`LoggingChatbotClient`] intercepts every `chat_with_tools_raw` call, captures
//! the raw HTTP request/response from the inner provider, persists a **metadata-only**
//! row to `llm_requests` in `system.db` (fire-and-forget), then returns the raw data
//! to the caller so it can write the **payload** to the user's own database.
//!
//! The LLM loop is completely unaware of this: it only holds an
//! `Arc<dyn ChatbotClient>` and calls `chat_with_tools` as usual.
//! The split keeps conversation content (payloads) behind the user key while
//! metadata (cost, tokens, timing) stays in the admin-readable registry.
use std::sync::Arc;
use std::time::Instant;
@@ -21,26 +22,10 @@ use super::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Messa
// ─────────────────────────────────────────────────────────────────────────────
/// Controls which parts of the HTTP exchange are persisted per row.
#[derive(Debug, Clone, Copy)]
pub struct LogSaveFlags {
pub request_payload: bool,
pub response_payload: bool,
pub request_headers: bool,
pub response_headers: bool,
}
impl Default for LogSaveFlags {
fn default() -> Self {
Self { request_payload: true, response_payload: true, request_headers: true, response_headers: true }
}
}
pub struct LoggingChatbotClient {
inner: Arc<dyn ChatbotClient>,
pool: Arc<SqlitePool>,
model_name: String,
flags: LogSaveFlags,
}
impl LoggingChatbotClient {
@@ -48,9 +33,8 @@ impl LoggingChatbotClient {
inner: Arc<dyn ChatbotClient>,
pool: Arc<SqlitePool>,
model_name: impl Into<String>,
flags: LogSaveFlags,
) -> Self {
Self { inner, pool, model_name: model_name.into(), flags }
Self { inner, pool, model_name: model_name.into() }
}
}
@@ -65,20 +49,35 @@ impl ChatbotClient for LoggingChatbotClient {
self.inner.chat(messages, options).await
}
/// Intercepts the call, delegates to `inner.chat_with_tools_raw` to capture
/// HTTP wire data, then spawns a fire-and-forget DB write before returning.
/// Passthrough that drops the raw meta. Used by callers that do not need
/// payload capture (e.g. the compactor).
async fn chat_with_tools(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<LlmTurn> {
let (turn, _) = self.chat_with_tools_raw(messages, tools, options).await?;
Ok(turn)
}
/// Intercepts the call, delegates to `inner.chat_with_tools_raw` to capture
/// HTTP wire data, writes a **metadata-only** row to `system.db`, then returns
/// the raw data so the caller can persist payloads to the user's own database.
async fn chat_with_tools_raw(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
let start = Instant::now();
let result = self.inner.chat_with_tools_raw(messages, tools, options).await;
let duration_ms = start.elapsed().as_millis() as i64;
let session_id = options.session_id;
let stack_id = options.stack_id;
let user_id = options.user_id.clone();
let request_id = options.request_id.clone();
let model_name = self.model_name.clone();
let pool = Arc::clone(&self.pool);
@@ -90,24 +89,13 @@ impl ChatbotClient for LoggingChatbotClient {
(*input_tokens, *output_tokens, *cache_read_tokens, *cache_creation_tokens),
};
let meta = meta.unwrap_or_default();
let flags = self.flags;
let request_json = if flags.request_payload {
meta.request_body.map(|v| v.to_string()).unwrap_or_default()
} else { String::new() };
let request_headers = if flags.request_headers { meta.request_headers.map(|v| v.to_string()) } else { None };
let response_json = if flags.response_payload { meta.response_body.map(|v| v.to_string()) } else { None };
let response_headers = if flags.response_headers { meta.response_headers.map(|v| v.to_string()) } else { None };
tokio::spawn(async move {
if let Err(e) = llm_requests::insert(&pool, llm_requests::LlmRequestRow {
request_id,
user_id,
session_id,
stack_id,
model_name,
request_json,
request_headers,
response_json,
response_headers,
error_text: None,
input_tokens: input_tokens.map(|n| n as i64),
output_tokens: output_tokens.map(|n| n as i64),
@@ -119,7 +107,7 @@ impl ChatbotClient for LoggingChatbotClient {
}
});
Ok(turn)
Ok((turn, meta))
}
Err(e) => {
@@ -127,13 +115,11 @@ impl ChatbotClient for LoggingChatbotClient {
tokio::spawn(async move {
if let Err(log_err) = llm_requests::insert(&pool, llm_requests::LlmRequestRow {
request_id,
user_id,
session_id,
stack_id,
model_name,
request_json: String::new(),
request_headers: None,
response_json: None,
response_headers: None,
error_text: Some(error_text),
input_tokens: None,
output_tokens: None,
@@ -149,14 +135,4 @@ impl ChatbotClient for LoggingChatbotClient {
}
}
}
/// Expose raw metadata so this wrapper can itself be wrapped if needed.
async fn chat_with_tools_raw(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
self.inner.chat_with_tools_raw(messages, tools, options).await
}
}