streaming: add chat_with_tools_raw_streaming to LM Studio + LoggingChatbotClient
LM Studio: delegate streaming to the inner OpenAI client. LoggingChatbotClient: extract shared log_and_return helper, add streaming override so deltas are forwarded and metadata still logged.
This commit is contained in:
@@ -1,7 +1,8 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
|
use tokio::sync::mpsc;
|
||||||
|
|
||||||
use crate::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, openai::OpenAiClient};
|
use crate::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, StreamDelta, openai::OpenAiClient};
|
||||||
|
|
||||||
/// LM Studio client.
|
/// LM Studio client.
|
||||||
///
|
///
|
||||||
@@ -48,4 +49,17 @@ impl ChatbotClient for LmStudioClient {
|
|||||||
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
|
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
|
||||||
self.inner.chat_with_tools_raw(messages, tools, options).await
|
self.inner.chat_with_tools_raw(messages, tools, options).await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// LM Studio is OpenAI-compatible: streaming forwards to the inner client.
|
||||||
|
/// If a local build rejects `stream_options`, the inner pre-delta buffered
|
||||||
|
/// retry covers it transparently.
|
||||||
|
async fn chat_with_tools_raw_streaming(
|
||||||
|
&self,
|
||||||
|
messages: &[Value],
|
||||||
|
tools: &[Value],
|
||||||
|
options: &ChatOptions,
|
||||||
|
delta_tx: mpsc::Sender<StreamDelta>,
|
||||||
|
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
|
||||||
|
self.inner.chat_with_tools_raw_streaming(messages, tools, options, delta_tx).await
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,16 +9,17 @@
|
|||||||
//! metadata (cost, tokens, timing) stays in the admin-readable registry.
|
//! metadata (cost, tokens, timing) stays in the admin-readable registry.
|
||||||
|
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::Instant;
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use sqlx::SqlitePool;
|
use sqlx::SqlitePool;
|
||||||
|
use tokio::sync::mpsc;
|
||||||
use tracing::warn;
|
use tracing::warn;
|
||||||
|
|
||||||
use crate::db::llm_requests;
|
use crate::db::llm_requests;
|
||||||
|
|
||||||
use super::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message};
|
use super::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, StreamDelta};
|
||||||
|
|
||||||
// ─────────────────────────────────────────────────────────────────────────────
|
// ─────────────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
@@ -36,43 +37,16 @@ impl LoggingChatbotClient {
|
|||||||
) -> Self {
|
) -> Self {
|
||||||
Self { inner, pool, model_name: model_name.into() }
|
Self { inner, pool, model_name: model_name.into() }
|
||||||
}
|
}
|
||||||
}
|
|
||||||
|
|
||||||
#[async_trait]
|
/// Shared logging tail of both raw entry points: writes the metadata-only
|
||||||
impl ChatbotClient for LoggingChatbotClient {
|
/// row to `system.db` (fire-and-forget), then passes the result through.
|
||||||
/// Passthrough — logging only applies to the tool-calling path.
|
async fn log_and_return(
|
||||||
async fn chat(
|
|
||||||
&self,
|
&self,
|
||||||
messages: &[Message],
|
|
||||||
options: &ChatOptions,
|
|
||||||
) -> anyhow::Result<ChatResponse> {
|
|
||||||
self.inner.chat(messages, options).await
|
|
||||||
}
|
|
||||||
|
|
||||||
/// 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,
|
options: &ChatOptions,
|
||||||
|
duration: Duration,
|
||||||
|
result: anyhow::Result<(LlmTurn, Option<LlmRawMeta>)>,
|
||||||
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
|
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
|
||||||
let start = Instant::now();
|
let duration_ms = duration.as_millis() as i64;
|
||||||
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 session_id = options.session_id;
|
||||||
let stack_id = options.stack_id;
|
let stack_id = options.stack_id;
|
||||||
@@ -136,3 +110,57 @@ impl ChatbotClient for LoggingChatbotClient {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl ChatbotClient for LoggingChatbotClient {
|
||||||
|
/// Passthrough — logging only applies to the tool-calling path.
|
||||||
|
async fn chat(
|
||||||
|
&self,
|
||||||
|
messages: &[Message],
|
||||||
|
options: &ChatOptions,
|
||||||
|
) -> anyhow::Result<ChatResponse> {
|
||||||
|
self.inner.chat(messages, options).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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;
|
||||||
|
self.log_and_return(options, start.elapsed(), result).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Streaming twin of `chat_with_tools_raw`: forwards `delta_tx` untouched to
|
||||||
|
/// the inner client (deltas are not logged — only the final turn is), then
|
||||||
|
/// applies the same metadata logging. Without this override the trait
|
||||||
|
/// default would silently fall back to the buffered call.
|
||||||
|
async fn chat_with_tools_raw_streaming(
|
||||||
|
&self,
|
||||||
|
messages: &[Value],
|
||||||
|
tools: &[Value],
|
||||||
|
options: &ChatOptions,
|
||||||
|
delta_tx: mpsc::Sender<StreamDelta>,
|
||||||
|
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
|
||||||
|
let start = Instant::now();
|
||||||
|
let result = self.inner.chat_with_tools_raw_streaming(messages, tools, options, delta_tx).await;
|
||||||
|
self.log_and_return(options, start.elapsed(), result).await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user