llm request tracking, user context cleanup, minor fixes
This commit is contained in:
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user