feat: add editing and MCP command handling, enhance SSE streaming with usage tracking
- Implemented `Edit` and `McpAdd` commands in the command parser and handler. - Added a new `stream` module to the runtime for handling streaming events. - Enhanced `SseParser` to parse usage information from SSE events. - Introduced `ToolCallAccumulator` for tracking tool calls independently. - Updated `AppStateRest` to include `app_config` and `MiscState` to track `effort_level` and `selected_index`. - Modified `LlmClient` to support streaming responses with usage tracking. - Improved error handling and retry logic in the streaming API calls. - Added tests for new features and improved markdown rendering in the chat view.
This commit is contained in:
asepharyana
+166
-16
@@ -1,8 +1,10 @@
|
||||
use std::time::Duration;
|
||||
use anyhow::Result;
|
||||
|
||||
use crate::app::runtime::stream::{SseParser, StreamEvent};
|
||||
use crate::app::runtime::stream::turn::StreamedTurn;
|
||||
use crate::dto::chat::message::ChatMessage;
|
||||
use crate::dto::provider::request::ToolDef;
|
||||
use crate::dto::provider::request::{ChatRequest, StreamOptions, ToolDef};
|
||||
|
||||
const DEFAULT_BASE_URL: &str = "https://opencode.ai/zen/v1";
|
||||
const DEFAULT_MODEL: &str = "deepseek-v4-flash-free";
|
||||
@@ -17,7 +19,7 @@ pub struct LlmClient {
|
||||
}
|
||||
|
||||
impl LlmClient {
|
||||
pub fn new(api_key: String, model: String) -> Self {
|
||||
pub fn new(api_key: String, model: String, base_url: Option<String>) -> Self {
|
||||
let model = if model.is_empty() {
|
||||
DEFAULT_MODEL.to_string()
|
||||
} else {
|
||||
@@ -31,7 +33,7 @@ impl LlmClient {
|
||||
LlmClient {
|
||||
client,
|
||||
api_key,
|
||||
base_url: DEFAULT_BASE_URL.to_string(),
|
||||
base_url: base_url.filter(|s| !s.is_empty()).unwrap_or_else(|| DEFAULT_BASE_URL.to_string()),
|
||||
model,
|
||||
}
|
||||
}
|
||||
@@ -46,7 +48,7 @@ impl LlmClient {
|
||||
messages: &[ChatMessage],
|
||||
tools: Option<Vec<ToolDef>>,
|
||||
) -> Result<ChatMessage> {
|
||||
let req = crate::dto::provider::request::ChatRequest {
|
||||
let req = ChatRequest {
|
||||
model: self.model.clone(),
|
||||
messages: messages.to_vec(),
|
||||
max_tokens: Some(4096),
|
||||
@@ -55,18 +57,130 @@ impl LlmClient {
|
||||
stream: Some(false),
|
||||
top_p: None,
|
||||
stop: None,
|
||||
stream_options: None,
|
||||
};
|
||||
|
||||
let url = format!("{}/chat/completions", self.base_url);
|
||||
let mut http_req = self.client
|
||||
.post(&url)
|
||||
.header("Content-Type", "application/json");
|
||||
let max_retries = 10;
|
||||
let mut attempt = 0;
|
||||
|
||||
loop {
|
||||
attempt += 1;
|
||||
|
||||
let mut http_req = self.client
|
||||
.post(&url)
|
||||
.header("Content-Type", "application/json");
|
||||
|
||||
if !self.api_key.is_empty() {
|
||||
http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key));
|
||||
}
|
||||
|
||||
let result = (|| -> Result<ChatMessage> {
|
||||
let resp = http_req.json(&req).send().map_err(|e| {
|
||||
if e.is_timeout() {
|
||||
anyhow::anyhow!("API request timed out after {:?}. Check your network or try again.", REQUEST_TIMEOUT)
|
||||
} else if e.is_connect() {
|
||||
anyhow::anyhow!("Could not connect to {}. Is the URL correct and is the service reachable?", self.base_url)
|
||||
} else {
|
||||
anyhow::anyhow!("API request failed: {}", e)
|
||||
}
|
||||
})?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().unwrap_or_default();
|
||||
anyhow::bail!("API error {} from {}: {}", status, self.base_url, body);
|
||||
}
|
||||
|
||||
let data: crate::dto::provider::response::ChatResponse = resp.json()?;
|
||||
let message = data
|
||||
.choices
|
||||
.into_iter()
|
||||
.next()
|
||||
.map(|c| c.message)
|
||||
.ok_or_else(|| anyhow::anyhow!("API response had no choices"))?;
|
||||
Ok(message)
|
||||
})();
|
||||
|
||||
match result {
|
||||
Ok(msg) => return Ok(msg),
|
||||
Err(e) => {
|
||||
if attempt >= max_retries {
|
||||
return Err(e);
|
||||
}
|
||||
eprintln!("Warning: {}. Retrying {}/{}...", e, attempt, max_retries);
|
||||
std::thread::sleep(Duration::from_secs(2));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Streaming variant of `chat_with_tools`. Feeds SSE chunks into an `SseParser` /
|
||||
/// `StreamedTurn` and invokes `on_event` for every parsed `StreamEvent` as it arrives,
|
||||
/// so the caller can push incremental UI updates in real time. Returns the fully
|
||||
/// assembled assistant message plus token usage (prompt, completion) if the server
|
||||
/// reported it. Retries the whole request only if no event has been observed yet
|
||||
/// (once tokens start arriving, a partial turn cannot be safely replayed).
|
||||
pub fn chat_with_tools_streaming(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
tools: Option<Vec<ToolDef>>,
|
||||
temperature: Option<f32>,
|
||||
max_tokens: Option<u32>,
|
||||
mut on_event: impl FnMut(&StreamEvent),
|
||||
) -> Result<(ChatMessage, Option<(u64, u64)>)> {
|
||||
let req = ChatRequest {
|
||||
model: self.model.clone(),
|
||||
messages: messages.to_vec(),
|
||||
max_tokens: Some(max_tokens.unwrap_or(4096)),
|
||||
temperature: Some(temperature.unwrap_or(0.7)),
|
||||
tools,
|
||||
stream: Some(true),
|
||||
top_p: None,
|
||||
stop: None,
|
||||
stream_options: Some(StreamOptions { include_usage: true }),
|
||||
};
|
||||
|
||||
let url = format!("{}/chat/completions", self.base_url);
|
||||
let max_retries = 10;
|
||||
let mut attempt = 0;
|
||||
let mut started = false;
|
||||
|
||||
loop {
|
||||
attempt += 1;
|
||||
let mut wrapped = |event: &StreamEvent| {
|
||||
started = true;
|
||||
on_event(event);
|
||||
};
|
||||
match self.try_stream_once(&req, &url, &mut wrapped) {
|
||||
Ok(result) => return Ok(result),
|
||||
Err(e) => {
|
||||
if started || attempt >= max_retries {
|
||||
return Err(e);
|
||||
}
|
||||
eprintln!("Warning: {}. Retrying {}/{}...", e, attempt, max_retries);
|
||||
std::thread::sleep(Duration::from_secs(2));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn try_stream_once(
|
||||
&self,
|
||||
req: &ChatRequest,
|
||||
url: &str,
|
||||
on_event: &mut dyn FnMut(&StreamEvent),
|
||||
) -> Result<(ChatMessage, Option<(u64, u64)>)> {
|
||||
use std::io::Read;
|
||||
|
||||
let mut http_req = self.client
|
||||
.post(url)
|
||||
.header("Content-Type", "application/json");
|
||||
if !self.api_key.is_empty() {
|
||||
http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key));
|
||||
}
|
||||
|
||||
let resp = http_req.json(&req).send().map_err(|e| {
|
||||
let resp = http_req.json(req).send().map_err(|e| {
|
||||
if e.is_timeout() {
|
||||
anyhow::anyhow!("API request timed out after {:?}. Check your network or try again.", REQUEST_TIMEOUT)
|
||||
} else if e.is_connect() {
|
||||
@@ -82,13 +196,49 @@ impl LlmClient {
|
||||
anyhow::bail!("API error {} from {}: {}", status, self.base_url, body);
|
||||
}
|
||||
|
||||
let data: crate::dto::provider::response::ChatResponse = resp.json()?;
|
||||
let message = data
|
||||
.choices
|
||||
.into_iter()
|
||||
.next()
|
||||
.map(|c| c.message)
|
||||
.ok_or_else(|| anyhow::anyhow!("API response had no choices"))?;
|
||||
Ok(message)
|
||||
let mut turn = StreamedTurn::new();
|
||||
let mut usage: Option<(u64, u64)> = None;
|
||||
let mut parser = SseParser::new();
|
||||
let mut reader = resp;
|
||||
let mut byte_buf: Vec<u8> = Vec::new();
|
||||
let mut chunk_buf = [0u8; 4096];
|
||||
|
||||
loop {
|
||||
let n = reader.read(&mut chunk_buf)
|
||||
.map_err(|e| anyhow::anyhow!("stream read error: {}", e))?;
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
byte_buf.extend_from_slice(&chunk_buf[..n]);
|
||||
let valid_len = match std::str::from_utf8(&byte_buf) {
|
||||
Ok(s) => s.len(),
|
||||
Err(e) => e.valid_up_to(),
|
||||
};
|
||||
if valid_len == 0 {
|
||||
continue;
|
||||
}
|
||||
let text = String::from_utf8_lossy(&byte_buf[..valid_len]).into_owned();
|
||||
byte_buf.drain(..valid_len);
|
||||
|
||||
for event in parser.feed(&text) {
|
||||
on_event(&event);
|
||||
match &event {
|
||||
StreamEvent::Usage { prompt_tokens, completion_tokens, .. } => {
|
||||
usage = Some((*prompt_tokens, *completion_tokens));
|
||||
}
|
||||
StreamEvent::Error(msg) => {
|
||||
anyhow::bail!("stream error: {}", msg);
|
||||
}
|
||||
StreamEvent::Done => {
|
||||
turn.apply_event(&event);
|
||||
return Ok((turn.build_assistant_message(), usage));
|
||||
}
|
||||
_ => turn.apply_event(&event),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
turn.is_complete = true;
|
||||
Ok((turn.build_assistant_message(), usage))
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user