Refactor IPC and DTO structures; remove unused code and streamline message handling
- Removed unused structs and methods from `response.rs`, `usage.rs`, and `client.rs`. - Simplified `Connection` handling in `conn.rs` to only support Unix sockets. - Updated `IpcServer` to exclusively use Unix sockets and removed TCP handling. - Cleaned up `editlog.rs` by removing loading and recent entry methods. - Refactored `memory.rs` to eliminate unused functions related to lesson promotion and retrospective creation. - Enhanced `search.rs` to support multiple search providers and improved error handling. - Updated chat view logic to simplify message display and improve user experience. - Removed deprecated modules and constants from various files to streamline the codebase.
This commit is contained in:
+139
-115
@@ -15,7 +15,6 @@ const MAX_AGENT_STEPS: usize = 40;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum Action {
|
||||
Quit,
|
||||
ForceQuit,
|
||||
SwitchMode(ModeKind),
|
||||
SubmitInput(String),
|
||||
@@ -32,19 +31,10 @@ pub enum Action {
|
||||
CloseOverlay,
|
||||
ToggleYoloArm,
|
||||
CycleAgentMode,
|
||||
ToolResult {
|
||||
tool_call_id: String,
|
||||
output: String,
|
||||
is_error: bool,
|
||||
},
|
||||
StreamToken(String),
|
||||
StreamDone,
|
||||
StreamError(String),
|
||||
SystemNote {
|
||||
kind: String,
|
||||
message: String,
|
||||
},
|
||||
RunCommand(String),
|
||||
QuitConfirm,
|
||||
Resize(u16, u16),
|
||||
Tick,
|
||||
@@ -60,15 +50,13 @@ pub enum Action {
|
||||
LessonReject {
|
||||
name: String,
|
||||
},
|
||||
StartOAuth {
|
||||
provider: String,
|
||||
},
|
||||
}
|
||||
|
||||
pub fn apply_action(state: &mut AppStateRest, action: Action) {
|
||||
match action {
|
||||
Action::Quit => {
|
||||
save_current_session(state);
|
||||
auto_create_retrospective(state);
|
||||
state.quit = true;
|
||||
}
|
||||
Action::ForceQuit => {
|
||||
save_current_session(state);
|
||||
state.quit = true;
|
||||
@@ -106,8 +94,8 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) {
|
||||
if let Some(ref mut rt) = state.session_runtime {
|
||||
rt.push_message(ChatMessage::user(text));
|
||||
}
|
||||
state.misc.thinking = true;
|
||||
spawn_turn(state);
|
||||
state.push_transcript(ChatMessageDisplay::new(Role::Assistant, "Thinking...".to_string()));
|
||||
state.dirty = true;
|
||||
}
|
||||
Action::InsertChar(c) => {
|
||||
@@ -137,9 +125,7 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) {
|
||||
state.dirty = true;
|
||||
}
|
||||
Action::ScrollUp => {
|
||||
let total = state.transcript_cache.messages.len();
|
||||
state.scroll.scroll_up();
|
||||
state.scroll.scroll_down(total);
|
||||
state.dirty = true;
|
||||
}
|
||||
Action::ScrollDown => {
|
||||
@@ -166,51 +152,6 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) {
|
||||
state.push_toast(toast);
|
||||
state.dirty = true;
|
||||
}
|
||||
Action::ToolResult {
|
||||
tool_call_id,
|
||||
output,
|
||||
is_error,
|
||||
} => {
|
||||
if let Some(ref mut rt) = state.session_runtime {
|
||||
rt.push_message(ChatMessage::tool_result(tool_call_id.clone(), output.clone()));
|
||||
rt.tool_call_results.push(
|
||||
crate::app::state::runtime::ToolCallResult {
|
||||
tool_call_id,
|
||||
tool_name: String::new(),
|
||||
output,
|
||||
is_error,
|
||||
duration_ms: 0,
|
||||
},
|
||||
);
|
||||
}
|
||||
state.dirty = true;
|
||||
}
|
||||
Action::StreamToken(token) => {
|
||||
if let Some(ref mut rt) = state.session_runtime {
|
||||
let found = rt.messages.iter_mut().rev().find(|m| {
|
||||
matches!(m.role, crate::dto::chat::message::Role::Assistant)
|
||||
});
|
||||
if let Some(last) = found {
|
||||
let current = last.content.take().unwrap_or_default();
|
||||
last.content = Some(current + &token);
|
||||
} else {
|
||||
let mut msg = ChatMessage::assistant(None);
|
||||
msg.content = Some(token);
|
||||
rt.push_message(msg);
|
||||
}
|
||||
}
|
||||
state.dirty = true;
|
||||
}
|
||||
Action::StreamDone => {
|
||||
state.dirty = true;
|
||||
}
|
||||
Action::StreamError(msg) => {
|
||||
let toast = crate::app::state::types::Toast::new(
|
||||
crate::app::state::types::ToastKind::Error,
|
||||
msg,
|
||||
);
|
||||
state.push_toast(toast);
|
||||
}
|
||||
Action::SystemNote { kind: _kind, message } => {
|
||||
let toast = crate::app::state::types::Toast::new(
|
||||
crate::app::state::types::ToastKind::Info,
|
||||
@@ -218,12 +159,6 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) {
|
||||
);
|
||||
state.push_toast(toast);
|
||||
}
|
||||
Action::RunCommand(text) => {
|
||||
if let Some(ref mut rt) = state.session_runtime {
|
||||
rt.push_message(ChatMessage::user(text));
|
||||
}
|
||||
state.dirty = true;
|
||||
}
|
||||
Action::QuitConfirm => {
|
||||
state.misc.overlay = Overlay::QuitConfirm;
|
||||
state.dirty = true;
|
||||
@@ -271,6 +206,26 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) {
|
||||
}
|
||||
state.dirty = true;
|
||||
}
|
||||
Action::StartOAuth { provider } => {
|
||||
let turn_events = state.turn_events.clone();
|
||||
let provider_clone = provider.clone();
|
||||
std::thread::spawn(move || {
|
||||
let result = run_oauth_flow(&provider_clone);
|
||||
let message = match result {
|
||||
Ok(msg) => msg,
|
||||
Err(e) => format!("OAuth login failed: {}", e),
|
||||
};
|
||||
if let Ok(mut q) = turn_events.lock() {
|
||||
q.push_back(TurnEvent::SystemNote {
|
||||
kind: "oauth".to_string(),
|
||||
message,
|
||||
});
|
||||
}
|
||||
});
|
||||
let toast = Toast::new(ToastKind::Info, format!("Opening browser for {} login...", provider));
|
||||
state.push_toast(toast);
|
||||
state.dirty = true;
|
||||
}
|
||||
Action::Tick => {
|
||||
let now_ms = chrono::Utc::now().timestamp_millis();
|
||||
state.misc.drain_expired_toasts(now_ms);
|
||||
@@ -290,27 +245,17 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) {
|
||||
for event in events {
|
||||
match event {
|
||||
TurnEvent::AssistantMessage(msg) => {
|
||||
state.misc.thinking = false;
|
||||
let display_content = msg.content.clone().unwrap_or_default();
|
||||
if !display_content.is_empty() {
|
||||
let replaced = if let Some(last) = state.transcript_cache.messages.last_mut() {
|
||||
if last.role == Role::Assistant && last.content == "Thinking..." {
|
||||
last.content = display_content.clone();
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
} else {
|
||||
false
|
||||
};
|
||||
if !replaced {
|
||||
state.push_transcript(ChatMessageDisplay::new(Role::Assistant, display_content));
|
||||
}
|
||||
state.push_transcript(ChatMessageDisplay::new(Role::Assistant, display_content));
|
||||
}
|
||||
if let Some(ref mut rt) = state.session_runtime {
|
||||
rt.push_message(msg);
|
||||
}
|
||||
}
|
||||
TurnEvent::ToolResult { tool_call_id, tool_name, output, is_error, path } => {
|
||||
state.misc.thinking = false;
|
||||
let display_path = path.unwrap_or_default();
|
||||
let display = if tool_name == "read" {
|
||||
let line_count = output.lines().count();
|
||||
@@ -385,6 +330,7 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) {
|
||||
turn_finished = true;
|
||||
}
|
||||
TurnEvent::Done => {
|
||||
state.misc.thinking = false;
|
||||
turn_finished = true;
|
||||
}
|
||||
}
|
||||
@@ -454,6 +400,9 @@ fn spawn_turn(state: &AppStateRest) {
|
||||
let events_q = turn_events.clone();
|
||||
|
||||
std::thread::spawn(move || {
|
||||
let db = crate::model::msglog::open_or_create(&edit_session_dir)
|
||||
.ok()
|
||||
.map(|c| std::sync::Arc::new(std::sync::Mutex::new(c)));
|
||||
let tc = TurnCtx {
|
||||
client: crate::service::provider::LlmClient::new(api_key, model),
|
||||
tdefs: tool_defs,
|
||||
@@ -463,6 +412,7 @@ fn spawn_turn(state: &AppStateRest) {
|
||||
workspace_roots,
|
||||
edit_log_session_dir: edit_session_dir,
|
||||
session_id,
|
||||
db,
|
||||
};
|
||||
let result = run_agent_turn(tc, &messages, &events_q);
|
||||
if let Err(e) = result {
|
||||
@@ -485,6 +435,15 @@ struct TurnCtx {
|
||||
workspace_roots: Vec<std::path::PathBuf>,
|
||||
edit_log_session_dir: std::path::PathBuf,
|
||||
session_id: String,
|
||||
db: Option<std::sync::Arc<std::sync::Mutex<rusqlite::Connection>>>,
|
||||
}
|
||||
|
||||
fn archive_message(db: &Option<std::sync::Arc<std::sync::Mutex<rusqlite::Connection>>>, session_id: &str, msg: &ChatMessage) {
|
||||
if let Some(ref arc) = db {
|
||||
if let Ok(conn) = arc.lock() {
|
||||
let _ = crate::model::msglog::insert_message(&conn, session_id, msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn run_agent_turn(
|
||||
@@ -495,6 +454,7 @@ fn run_agent_turn(
|
||||
let mut msgs = messages.to_vec();
|
||||
let mut edits_this_turn = 0u32;
|
||||
let mut tool_only_rounds = 0usize;
|
||||
let mut prev_shaped = false;
|
||||
|
||||
let system_text = format!(
|
||||
"{}\n\n{}",
|
||||
@@ -502,20 +462,37 @@ fn run_agent_turn(
|
||||
crate::resources::SYSTEM_TOOLS,
|
||||
);
|
||||
if !msgs.iter().any(|m| matches!(m.role, crate::dto::chat::message::Role::System)) {
|
||||
msgs.insert(0, ChatMessage::system(system_text));
|
||||
let sys = ChatMessage::system(system_text);
|
||||
archive_message(&tc.db, &tc.session_id, &sys);
|
||||
msgs.insert(0, sys);
|
||||
}
|
||||
|
||||
for _step in 0..MAX_AGENT_STEPS {
|
||||
if tool_only_rounds >= MAX_TOOL_ONLY_TURNS {
|
||||
msgs.push(ChatMessage::user(
|
||||
let stop_msg = ChatMessage::user(
|
||||
"Stop calling tools. Respond naturally now.".to_string(),
|
||||
));
|
||||
);
|
||||
archive_message(&tc.db, &tc.session_id, &stop_msg);
|
||||
msgs.push(stop_msg);
|
||||
tool_only_rounds = 0;
|
||||
}
|
||||
|
||||
let wire_msgs = if crate::app::runtime::shortsend::should_shape(msgs.len(), prev_shaped) {
|
||||
let total_chars: usize = msgs.iter()
|
||||
.filter_map(|m| m.content.as_deref())
|
||||
.map(|c| c.len())
|
||||
.sum();
|
||||
let token_estimate = total_chars / 4;
|
||||
prev_shaped = true;
|
||||
crate::app::runtime::shortsend::shape_messages(&msgs, token_estimate)
|
||||
} else {
|
||||
prev_shaped = false;
|
||||
msgs.clone()
|
||||
};
|
||||
|
||||
let response = tc
|
||||
.client
|
||||
.chat_with_tools(&msgs, Some(tc.tdefs.clone()))?;
|
||||
.chat_with_tools(&wire_msgs, Some(tc.tdefs.clone()))?;
|
||||
|
||||
let has_tool_calls = response.tool_calls.is_some()
|
||||
&& response.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty());
|
||||
@@ -524,6 +501,7 @@ fn run_agent_turn(
|
||||
if has_tool_calls {
|
||||
tool_only_rounds += 1;
|
||||
let tool_calls = response.tool_calls.clone().unwrap_or_default();
|
||||
archive_message(&tc.db, &tc.session_id, &response);
|
||||
msgs.push(response);
|
||||
for tool_call in tool_calls {
|
||||
let tool_name = tool_call.function.name.clone();
|
||||
@@ -581,10 +559,12 @@ fn run_agent_turn(
|
||||
}
|
||||
|
||||
let tool_msg = ChatMessage::tool_result(tool_call.id.clone(), output);
|
||||
archive_message(&tc.db, &tc.session_id, &tool_msg);
|
||||
msgs.push(tool_msg);
|
||||
}
|
||||
} else {
|
||||
if !content.is_empty() {
|
||||
archive_message(&tc.db, &tc.session_id, &response);
|
||||
if let Ok(mut q) = events_q.lock() {
|
||||
q.push_back(TurnEvent::AssistantMessage(response));
|
||||
}
|
||||
@@ -698,36 +678,80 @@ fn save_current_session(state: &AppStateRest) {
|
||||
}
|
||||
}
|
||||
|
||||
fn auto_create_retrospective(state: &mut AppStateRest) {
|
||||
if state.session_runtime.is_none() {
|
||||
return;
|
||||
}
|
||||
let session = crate::model::session::Session::new(
|
||||
state.session_id.clone(),
|
||||
"session".to_string(),
|
||||
);
|
||||
match crate::model::memory::auto_create_retrospective(&state.session_dir, &session) {
|
||||
Ok(Some(retro)) => {
|
||||
state.push_toast(crate::app::state::types::Toast::new(
|
||||
crate::app::state::types::ToastKind::Info,
|
||||
format!("Retrospective created: {}", retro.name),
|
||||
));
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(e) => {
|
||||
let _ = e;
|
||||
}
|
||||
}
|
||||
let lessons: Vec<crate::model::memory::Memory> = crate::model::memory::Memory::list(&state.memory_dir)
|
||||
.iter()
|
||||
.filter_map(|n| crate::model::memory::Memory::read(&state.memory_dir, n).ok())
|
||||
.filter(|m| m.kind == "lesson")
|
||||
.collect();
|
||||
if let Some(global_dir) = dirs::data_dir().map(|d| d.join("zesdex")) {
|
||||
for lesson in &lessons {
|
||||
if lesson.scope.as_deref() != Some("global") {
|
||||
let _ = crate::model::memory::promote_with_consensus(&global_dir, lesson);
|
||||
fn run_oauth_flow(provider: &str) -> anyhow::Result<String> {
|
||||
use crate::service::oauth::manager::{OAuthConfig, OAuthManager};
|
||||
use crate::service::oauth::loopback::LoopbackServer;
|
||||
use crate::service::oauth::pkce::CodeVerifier;
|
||||
|
||||
let config = match provider {
|
||||
"zen" | "opencode" => OAuthConfig {
|
||||
auth_url: "https://opencode.ai/zen/oauth/authorize".to_string(),
|
||||
token_url: "https://opencode.ai/zen/oauth/token".to_string(),
|
||||
client_id: std::env::var("ZEN_CLIENT_ID")
|
||||
.unwrap_or_else(|_| "zesdex".to_string()),
|
||||
client_secret: std::env::var("ZEN_CLIENT_SECRET").ok(),
|
||||
scopes: vec!["openid".to_string(), "profile".to_string(), "email".to_string()],
|
||||
},
|
||||
"openai" => OAuthConfig {
|
||||
auth_url: "https://auth0.openai.com/authorize".to_string(),
|
||||
token_url: "https://auth0.openai.com/oauth/token".to_string(),
|
||||
client_id: std::env::var("OPENAI_CLIENT_ID")
|
||||
.unwrap_or_else(|_| "zesdex".to_string()),
|
||||
client_secret: std::env::var("OPENAI_CLIENT_SECRET").ok(),
|
||||
scopes: vec!["openid".to_string(), "profile".to_string(), "email".to_string()],
|
||||
},
|
||||
other => {
|
||||
let auth_url = std::env::var(format!("{}_AUTH_URL", other.to_uppercase()))
|
||||
.map_err(|_| anyhow::anyhow!("unknown provider '{}'. Set {}_AUTH_URL env var.", other, other.to_uppercase()))?;
|
||||
let token_url = std::env::var(format!("{}_TOKEN_URL", other.to_uppercase()))
|
||||
.map_err(|_| anyhow::anyhow!("{}_TOKEN_URL not set", other.to_uppercase()))?;
|
||||
let client_id = std::env::var(format!("{}_CLIENT_ID", other.to_uppercase()))
|
||||
.unwrap_or_else(|_| "zesdex".to_string());
|
||||
OAuthConfig {
|
||||
auth_url,
|
||||
token_url,
|
||||
client_id,
|
||||
client_secret: std::env::var(format!("{}_CLIENT_SECRET", other.to_uppercase())).ok(),
|
||||
scopes: vec!["openid".to_string(), "profile".to_string(), "email".to_string()],
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let server = LoopbackServer::bind()?;
|
||||
let redirect_uri = server.redirect_uri();
|
||||
|
||||
let verifier = CodeVerifier::new();
|
||||
let challenge = verifier.challenge();
|
||||
let state_token = format!("{:x}", sha2::Sha256::digest(rand_bytes(16)));
|
||||
|
||||
let mut manager = OAuthManager::new(config.clone());
|
||||
let auth_url = manager.build_auth_url(&redirect_uri, &state_token, challenge.as_str());
|
||||
|
||||
let _ = webbrowser::open(&auth_url);
|
||||
|
||||
let code = server.wait_for_code(120_000)?;
|
||||
|
||||
manager.exchange_code(&code, &redirect_uri, verifier.as_str())
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
if let Some(ref token) = manager.token {
|
||||
let token_path = dirs::config_dir()
|
||||
.unwrap_or_else(|| std::path::PathBuf::from("."))
|
||||
.join("zesdex")
|
||||
.join(format!("oauth_{}.json", provider));
|
||||
if let Some(parent) = token_path.parent() {
|
||||
let _ = std::fs::create_dir_all(parent);
|
||||
}
|
||||
let _ = std::fs::write(&token_path, serde_json::to_string_pretty(token).unwrap_or_default());
|
||||
}
|
||||
|
||||
Ok(format!("Successfully authenticated with {}.", provider))
|
||||
}
|
||||
|
||||
fn rand_bytes(n: usize) -> Vec<u8> {
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
let seed = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().subsec_nanos();
|
||||
(0..n).map(|i| ((seed >> (i % 4 * 8)) ^ (i as u32 * 2654435761)) as u8).collect()
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -41,10 +41,7 @@ pub fn apply_command(command: Command) -> Vec<Action> {
|
||||
}]
|
||||
}
|
||||
Command::Login { provider } => {
|
||||
vec![Action::SystemNote {
|
||||
kind: "oauth".to_string(),
|
||||
message: format!("OAuth login flow started for {}", provider),
|
||||
}]
|
||||
vec![Action::StartOAuth { provider }]
|
||||
}
|
||||
Command::Unknown(cmd) => {
|
||||
vec![Action::SystemNote {
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use crate::app::state::runtime::TurnEvent;
|
||||
|
||||
const FAST_POLL_MS: u64 = 8;
|
||||
const SLOW_POLL_MS: u64 = 100;
|
||||
const IDLE_THRESHOLD_MS: u64 = 500;
|
||||
|
||||
pub struct EventLoop {
|
||||
last_activity: Instant,
|
||||
fast_poll_until: Option<Instant>,
|
||||
}
|
||||
|
||||
impl EventLoop {
|
||||
pub fn new() -> Self {
|
||||
EventLoop {
|
||||
last_activity: Instant::now(),
|
||||
fast_poll_until: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn poll_interval(&self) -> Duration {
|
||||
if let Some(fast_until) = self.fast_poll_until {
|
||||
if Instant::now() < fast_until {
|
||||
return Duration::from_millis(FAST_POLL_MS);
|
||||
}
|
||||
}
|
||||
Duration::from_millis(SLOW_POLL_MS)
|
||||
}
|
||||
|
||||
pub fn mark_active(&mut self) {
|
||||
self.last_activity = Instant::now();
|
||||
self.fast_poll_until = Some(Instant::now() + Duration::from_millis(IDLE_THRESHOLD_MS));
|
||||
}
|
||||
|
||||
pub fn is_idle(&self) -> bool {
|
||||
self.last_activity.elapsed().as_millis() as u64 > IDLE_THRESHOLD_MS
|
||||
}
|
||||
|
||||
pub fn drain_events(
|
||||
events: &std::sync::Mutex<VecDeque<TurnEvent>>,
|
||||
) -> Vec<TurnEvent> {
|
||||
events.lock().map(|mut q| q.drain(..).collect()).unwrap_or_default()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for EventLoop {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
pub mod actions;
|
||||
pub mod commands;
|
||||
pub mod event_loop;
|
||||
pub mod stream;
|
||||
pub mod shortsend;
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
pub mod turn;
|
||||
pub mod tools;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum StreamEvent {
|
||||
Token(String),
|
||||
Reasoning(String),
|
||||
ToolCallDelta {
|
||||
index: usize,
|
||||
id: Option<String>,
|
||||
name: Option<String>,
|
||||
arguments_delta: String,
|
||||
},
|
||||
Usage {
|
||||
prompt_tokens: u64,
|
||||
completion_tokens: u64,
|
||||
total_tokens: u64,
|
||||
},
|
||||
Done,
|
||||
Error(String),
|
||||
}
|
||||
|
||||
pub struct SseParser {
|
||||
buffer: String,
|
||||
event_type: Option<String>,
|
||||
data_lines: Vec<String>,
|
||||
}
|
||||
|
||||
impl SseParser {
|
||||
pub fn new() -> Self {
|
||||
SseParser {
|
||||
buffer: String::new(),
|
||||
event_type: None,
|
||||
data_lines: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn feed(&mut self, chunk: &str) -> Vec<StreamEvent> {
|
||||
self.buffer.push_str(chunk);
|
||||
let mut events = Vec::new();
|
||||
while let Some(line_end) = self.buffer.find('\n') {
|
||||
let line = self.buffer[..line_end].trim_end_matches('\r').to_string();
|
||||
self.buffer = self.buffer[line_end + 1..].to_string();
|
||||
if line.is_empty() {
|
||||
if let Some(event) = self.flush_event() {
|
||||
events.push(event);
|
||||
}
|
||||
} else if let Some(ty) = line.strip_prefix("event: ") {
|
||||
self.event_type = Some(ty.trim().to_string());
|
||||
} else if let Some(data) = line.strip_prefix("data: ") {
|
||||
self.data_lines.push(data.to_string());
|
||||
} else if line.starts_with("data:") {
|
||||
self.data_lines.push(String::new());
|
||||
}
|
||||
}
|
||||
events
|
||||
}
|
||||
|
||||
fn flush_event(&mut self) -> Option<StreamEvent> {
|
||||
let data = self.data_lines.join("\n");
|
||||
self.data_lines.clear();
|
||||
let event_type = self.event_type.take().unwrap_or_default();
|
||||
if data.is_empty() || data == "[DONE]" {
|
||||
if data == "[DONE]" {
|
||||
return Some(StreamEvent::Done);
|
||||
}
|
||||
return None;
|
||||
}
|
||||
let value: Value = serde_json::from_str(&data).ok()?;
|
||||
match event_type.as_str() {
|
||||
"message.stop" => Some(StreamEvent::Done),
|
||||
"message.start" => None,
|
||||
"message.delta" | "" => {
|
||||
let delta = value.get("delta").or_else(|| value.get("choices"))?;
|
||||
if let Some(choices) = delta.as_array() {
|
||||
let choice = choices.first()?;
|
||||
let delta = choice.get("delta")?;
|
||||
if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
|
||||
return Some(StreamEvent::Token(content.to_string()));
|
||||
}
|
||||
if let Some(reasoning) = delta.get("reasoning_content").and_then(|r| r.as_str()) {
|
||||
return Some(StreamEvent::Reasoning(reasoning.to_string()));
|
||||
}
|
||||
if let Some(tool_calls) = delta.get("tool_calls").and_then(|tc| tc.as_array()) {
|
||||
for tc in tool_calls {
|
||||
let index = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
|
||||
let id = tc.get("id").and_then(|i| i.as_str()).map(|s| s.to_string());
|
||||
let name = tc.get("function")
|
||||
.and_then(|f| f.get("name"))
|
||||
.and_then(|n| n.as_str())
|
||||
.map(|s| s.to_string());
|
||||
let args_delta = tc.get("function")
|
||||
.and_then(|f| f.get("arguments"))
|
||||
.and_then(|a| a.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
return Some(StreamEvent::ToolCallDelta {
|
||||
index,
|
||||
id,
|
||||
name,
|
||||
arguments_delta: args_delta,
|
||||
});
|
||||
}
|
||||
}
|
||||
let finish = choice.get("finish_reason");
|
||||
if let Some(reason) = finish.and_then(|r| r.as_str()) {
|
||||
if reason == "stop" || reason == "tool_calls" {
|
||||
return Some(StreamEvent::Done);
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
|
||||
return Some(StreamEvent::Token(content.to_string()));
|
||||
}
|
||||
None
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn reset(&mut self) {
|
||||
self.buffer.clear();
|
||||
self.event_type = None;
|
||||
self.data_lines.clear();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse_stream_chunk(data: &str) -> Option<StreamEvent> {
|
||||
let value: Value = serde_json::from_str(data).ok()?;
|
||||
if value == Value::Null {
|
||||
return None;
|
||||
}
|
||||
let choices = value.get("choices")?.as_array()?;
|
||||
let choice = choices.first()?;
|
||||
let delta = choice.get("delta")?;
|
||||
if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
|
||||
return Some(StreamEvent::Token(content.to_string()));
|
||||
}
|
||||
if let Some(reasoning) = delta.get("reasoning_content").and_then(|r| r.as_str()) {
|
||||
return Some(StreamEvent::Reasoning(reasoning.to_string()));
|
||||
}
|
||||
if let Some(finish) = choice.get("finish_reason").and_then(|r| r.as_str()) {
|
||||
if finish == "stop" || finish == "tool_calls" {
|
||||
return Some(StreamEvent::Done);
|
||||
}
|
||||
}
|
||||
if let Some(tool_calls) = delta.get("tool_calls").and_then(|tc| tc.as_array()) {
|
||||
if let Some(tc) = tool_calls.first() {
|
||||
let index = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
|
||||
let id = tc.get("id").and_then(|i| i.as_str()).map(|s| s.to_string());
|
||||
let name = tc.get("function")
|
||||
.and_then(|f| f.get("name"))
|
||||
.and_then(|n| n.as_str())
|
||||
.map(|s| s.to_string());
|
||||
let args = tc.get("function")
|
||||
.and_then(|f| f.get("arguments"))
|
||||
.and_then(|a| a.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
return Some(StreamEvent::ToolCallDelta {
|
||||
index,
|
||||
id,
|
||||
name,
|
||||
arguments_delta: args,
|
||||
});
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
use super::turn::ParsedToolCall;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
pub struct ToolCallAccumulator {
|
||||
calls: Vec<ParsedToolCall>,
|
||||
}
|
||||
|
||||
impl ToolCallAccumulator {
|
||||
pub fn new() -> Self {
|
||||
ToolCallAccumulator { calls: Vec::new() }
|
||||
}
|
||||
|
||||
pub fn add_delta(
|
||||
&mut self,
|
||||
index: usize,
|
||||
id: Option<&str>,
|
||||
name: Option<&str>,
|
||||
arguments_delta: &str,
|
||||
) {
|
||||
while self.calls.len() <= index {
|
||||
self.calls.push(ParsedToolCall {
|
||||
id: String::new(),
|
||||
name: String::new(),
|
||||
arguments: String::new(),
|
||||
is_complete: false,
|
||||
});
|
||||
}
|
||||
let tc = &mut self.calls[index];
|
||||
if let Some(new_id) = id {
|
||||
if !new_id.is_empty() {
|
||||
tc.id = new_id.to_string();
|
||||
}
|
||||
}
|
||||
if let Some(new_name) = name {
|
||||
if !new_name.is_empty() {
|
||||
tc.name = new_name.to_string();
|
||||
}
|
||||
}
|
||||
tc.arguments.push_str(arguments_delta);
|
||||
}
|
||||
|
||||
pub fn calls(&self) -> &[ParsedToolCall] {
|
||||
&self.calls
|
||||
}
|
||||
|
||||
pub fn is_complete(&self) -> bool {
|
||||
!self.calls.is_empty() && self.calls.iter().all(|tc| !tc.name.is_empty() && !tc.arguments.is_empty())
|
||||
}
|
||||
|
||||
pub fn reset(&mut self) {
|
||||
self.calls.clear();
|
||||
}
|
||||
|
||||
pub fn pending_args(&self) -> Vec<Value> {
|
||||
self.calls
|
||||
.iter()
|
||||
.filter(|tc| !tc.name.is_empty())
|
||||
.map(|tc| {
|
||||
json!({
|
||||
"tool_call_id": tc.id,
|
||||
"name": tc.name,
|
||||
"arguments": tc.arguments,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ToolCallAccumulator {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
use super::StreamEvent;
|
||||
use crate::dto::chat::message::ChatMessage;
|
||||
use crate::dto::chat::tool::{ToolCall, ToolFunction};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct StreamedTurn {
|
||||
pub messages: Vec<ChatMessage>,
|
||||
pub tool_calls: Vec<ParsedToolCall>,
|
||||
pub is_complete: bool,
|
||||
pub accumulated_content: String,
|
||||
pub accumulated_reasoning: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ParsedToolCall {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub arguments: String,
|
||||
pub is_complete: bool,
|
||||
}
|
||||
|
||||
impl ParsedToolCall {
|
||||
pub fn try_parse(&self) -> Option<Value> {
|
||||
serde_json::from_str(&self.arguments).ok()
|
||||
}
|
||||
}
|
||||
|
||||
impl StreamedTurn {
|
||||
pub fn new() -> Self {
|
||||
StreamedTurn {
|
||||
messages: Vec::new(),
|
||||
tool_calls: Vec::new(),
|
||||
is_complete: false,
|
||||
accumulated_content: String::new(),
|
||||
accumulated_reasoning: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn apply_event(&mut self, event: &StreamEvent) {
|
||||
match event {
|
||||
StreamEvent::Token(token) => {
|
||||
self.accumulated_content.push_str(token);
|
||||
}
|
||||
StreamEvent::Reasoning(reasoning) => {
|
||||
self.accumulated_reasoning.push_str(reasoning);
|
||||
}
|
||||
StreamEvent::ToolCallDelta {
|
||||
index,
|
||||
id,
|
||||
name,
|
||||
arguments_delta,
|
||||
} => {
|
||||
while self.tool_calls.len() <= *index {
|
||||
self.tool_calls.push(ParsedToolCall {
|
||||
id: String::new(),
|
||||
name: String::new(),
|
||||
arguments: String::new(),
|
||||
is_complete: false,
|
||||
});
|
||||
}
|
||||
let tc = &mut self.tool_calls[*index];
|
||||
if let Some(new_id) = id {
|
||||
if !new_id.is_empty() {
|
||||
tc.id = new_id.clone();
|
||||
}
|
||||
}
|
||||
if let Some(new_name) = name {
|
||||
if !new_name.is_empty() {
|
||||
tc.name = new_name.clone();
|
||||
}
|
||||
}
|
||||
tc.arguments.push_str(arguments_delta);
|
||||
}
|
||||
StreamEvent::Done => {
|
||||
self.is_complete = true;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_assistant_message(&self) -> ChatMessage {
|
||||
let mut msg = if self.tool_calls.is_empty() {
|
||||
ChatMessage::assistant(None)
|
||||
} else {
|
||||
let tool_dtos: Vec<ToolCall> = self.tool_calls
|
||||
.iter()
|
||||
.filter(|tc| !tc.name.is_empty())
|
||||
.map(|tc| {
|
||||
let args_value: serde_json::Value = serde_json::from_str(&tc.arguments)
|
||||
.unwrap_or(serde_json::Value::String(tc.arguments.clone()));
|
||||
ToolCall {
|
||||
id: tc.id.clone(),
|
||||
type_: "function".to_string(),
|
||||
function: ToolFunction {
|
||||
name: tc.name.clone(),
|
||||
arguments: args_value,
|
||||
},
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let mut msg = ChatMessage::assistant(None);
|
||||
if !tool_dtos.is_empty() {
|
||||
msg.tool_calls = Some(tool_dtos);
|
||||
}
|
||||
msg
|
||||
};
|
||||
let content = if self.accumulated_content.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(self.accumulated_content.clone())
|
||||
};
|
||||
msg.content = content;
|
||||
msg
|
||||
}
|
||||
|
||||
pub fn has_tool_calls(&self) -> bool {
|
||||
self.tool_calls.iter().any(|tc| !tc.name.is_empty())
|
||||
}
|
||||
|
||||
pub fn content(&self) -> &str {
|
||||
&self.accumulated_content
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for StreamedTurn {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user