| 1 | //! Text chat workflows for `DeepSeek` and DeepSeek-compatible APIs. |
| 2 | |
| 3 | use std::collections::HashMap; |
| 4 | use std::io::{self, Write}; |
| 5 | use std::path::Path; |
| 6 | use std::time::Instant; |
| 7 | |
| 8 | use anyhow::{Context, Result}; |
| 9 | use colored::{ColoredString, Colorize}; |
| 10 | use rustyline::completion::{Completer, Pair}; |
| 11 | use rustyline::error::ReadlineError; |
| 12 | use rustyline::highlight::Highlighter; |
| 13 | use rustyline::hint::Hinter; |
| 14 | use rustyline::history::DefaultHistory; |
| 15 | use rustyline::validate::Validator; |
| 16 | use rustyline::{Context as RlContext, Editor, Helper}; |
| 17 | use serde_json::{Value, json}; |
| 18 | |
| 19 | use crate::client::DeepSeekClient; |
| 20 | use crate::models::{ |
| 21 | CacheControl, ContentBlock, ContentBlockStart, Delta, Message, MessageRequest, StreamEvent, |
| 22 | SystemBlock, SystemPrompt, Tool, Usage, |
| 23 | }; |
| 24 | use crate::palette; |
| 25 | use crate::utils::pretty_json; |
| 26 | |
| 27 | // === Types === |
| 28 | |
| 29 | /// Options for running text chat sessions. |
| 30 | #[allow(clippy::struct_excessive_bools)] |
| 31 | pub struct TextChatOptions { |
| 32 | pub model: String, |
| 33 | pub prompt: Option<String>, |
| 34 | pub system: Option<String>, |
| 35 | pub stream: bool, |
| 36 | pub temperature: Option<f32>, |
| 37 | pub top_p: Option<f32>, |
| 38 | pub max_tokens: u32, |
| 39 | pub cache_prompt: bool, |
| 40 | pub cache_system: bool, |
| 41 | pub cache_tools: bool, |
| 42 | pub tools: Option<Vec<Tool>>, |
| 43 | pub tool_choice: Option<Value>, |
| 44 | } |
| 45 | |
| 46 | // === Public API === |
| 47 | |
| 48 | pub async fn run_deepseek_chat(client: &DeepSeekClient, options: TextChatOptions) -> Result<()> { |
| 49 | let mut messages: Vec<Message> = Vec::new(); |
| 50 | let mut stats = SessionStats::new(); |
| 51 | |
| 52 | print_banner("DeepSeek Compatible API"); |
| 53 | print_session_info( |
| 54 | &options, |
| 55 | messages.len(), |
| 56 | options.tools.as_ref().map_or(0, std::vec::Vec::len), |
| 57 | ); |
| 58 | |
| 59 | if let Some(prompt) = options.prompt.as_deref() { |
| 60 | process_deepseek_turn(client, &options, &mut messages, prompt, &mut stats).await?; |
| 61 | } else { |
| 62 | let mut rl = create_editor()?; |
| 63 | while let Some(line) = read_prompt(&mut rl)? { |
| 64 | if handle_line_deepseek(line, client, &options, &mut messages, &mut stats).await? { |
| 65 | break; |
| 66 | } |
| 67 | } |
| 68 | } |
| 69 | |
| 70 | Ok(()) |
| 71 | } |
| 72 | |
| 73 | pub async fn run_official_chat(client: &DeepSeekClient, options: TextChatOptions) -> Result<()> { |
| 74 | let mut messages: Vec<Value> = Vec::new(); |
| 75 | let mut stats = SessionStats::new(); |
| 76 | |
| 77 | if let Some(system) = options.system.clone() { |
| 78 | messages.push(json!({ "role": "system", "content": system })); |
| 79 | } |
| 80 | |
| 81 | print_banner("Official API"); |
| 82 | print_session_info( |
| 83 | &options, |
| 84 | messages.len(), |
| 85 | options.tools.as_ref().map_or(0, std::vec::Vec::len), |
| 86 | ); |
| 87 | |
| 88 | if let Some(prompt) = options.prompt.as_deref() { |
| 89 | process_official_turn(client, &options, &mut messages, prompt, &mut stats).await?; |
| 90 | } else { |
| 91 | let mut rl = create_editor()?; |
| 92 | while let Some(line) = read_prompt(&mut rl)? { |
| 93 | if handle_line_official( |
| 94 | line, |
| 95 | client, |
| 96 | &options, |
| 97 | &mut messages, |
| 98 | &mut stats, |
| 99 | options.system.as_deref(), |
| 100 | ) |
| 101 | .await? |
| 102 | { |
| 103 | break; |
| 104 | } |
| 105 | } |
| 106 | } |
| 107 | |
| 108 | Ok(()) |
| 109 | } |
| 110 | |
| 111 | pub fn load_tools( |
| 112 | tools_file: Option<&Path>, |
| 113 | tools_json: Option<&str>, |
| 114 | ) -> Result<Option<Vec<Tool>>> { |
| 115 | let tools = if let Some(raw_json) = tools_json { |
| 116 | let parsed: Vec<Tool> = serde_json::from_str(raw_json) |
| 117 | .context("Failed to parse tools_json: expected an array of tool definitions.")?; |
| 118 | Some(parsed) |
| 119 | } else if let Some(path) = tools_file { |
| 120 | let contents = std::fs::read_to_string(path) |
| 121 | .with_context(|| format!("Failed to read tools file: {}", path.display()))?; |
| 122 | let parsed: Vec<Tool> = serde_json::from_str(&contents) |
| 123 | .with_context(|| format!("Failed to parse tools file: {}", path.display()))?; |
| 124 | Some(parsed) |
| 125 | } else { |
| 126 | None |
| 127 | }; |
| 128 | |
| 129 | Ok(tools) |
| 130 | } |
| 131 | |
| 132 | pub fn parse_tool_choice(choice: Option<&str>) -> Result<Option<Value>> { |
| 133 | let Some(choice) = choice else { |
| 134 | return Ok(None); |
| 135 | }; |
| 136 | let trimmed = choice.trim(); |
| 137 | if trimmed.starts_with('{') || trimmed.starts_with('[') { |
| 138 | let value: Value = |
| 139 | serde_json::from_str(trimmed).context("Failed to parse tool_choice: expected JSON.")?; |
| 140 | return Ok(Some(value)); |
| 141 | } |
| 142 | |
| 143 | let value = match trimmed { |
| 144 | "auto" | "none" | "any" => json!({ "type": trimmed }), |
| 145 | _ => json!({ "type": "tool", "name": trimmed }), |
| 146 | }; |
| 147 | Ok(Some(value)) |
| 148 | } |
| 149 | |
| 150 | #[allow(clippy::too_many_lines)] |
| 151 | async fn process_deepseek_turn( |
| 152 | client: &DeepSeekClient, |
| 153 | options: &TextChatOptions, |
| 154 | messages: &mut Vec<Message>, |
| 155 | user_input: &str, |
| 156 | stats: &mut SessionStats, |
| 157 | ) -> Result<()> { |
| 158 | let cache_control = if options.cache_prompt { |
| 159 | Some(CacheControl { |
| 160 | cache_type: "ephemeral".to_string(), |
| 161 | }) |
| 162 | } else { |
| 163 | None |
| 164 | }; |
| 165 | |
| 166 | messages.push(Message { |
| 167 | role: "user".to_string(), |
| 168 | content: vec![ContentBlock::Text { |
| 169 | text: user_input.to_string(), |
| 170 | cache_control, |
| 171 | }], |
| 172 | }); |
| 173 | |
| 174 | let request = MessageRequest { |
| 175 | model: options.model.clone(), |
| 176 | messages: messages.clone(), |
| 177 | max_tokens: options.max_tokens, |
| 178 | system: build_system_prompt(options.system.as_deref(), options.cache_system), |
| 179 | tools: cache_tools(options.tools.clone(), options.cache_tools), |
| 180 | tool_choice: options.tool_choice.clone(), |
| 181 | metadata: None, |
| 182 | thinking: None, |
| 183 | reasoning_effort: None, |
| 184 | stream: Some(options.stream), |
| 185 | temperature: options.temperature, |
| 186 | top_p: options.top_p, |
| 187 | }; |
| 188 | |
| 189 | if options.stream { |
| 190 | let stream = client.create_message_stream(request).await?; |
| 191 | tokio::pin!(stream); |
| 192 | |
| 193 | let mut current_thinking = String::new(); |
| 194 | let mut current_text = String::new(); |
| 195 | let mut block_types: HashMap<u32, String> = HashMap::new(); |
| 196 | let mut tool_blocks: HashMap<u32, (String, String, String)> = HashMap::new(); |
| 197 | let mut is_thinking = false; |
| 198 | |
| 199 | while let Some(event) = futures_util::StreamExt::next(&mut stream).await { |
| 200 | let event = event?; |
| 201 | match event { |
| 202 | StreamEvent::ContentBlockStart { |
| 203 | index, |
| 204 | content_block, |
| 205 | } => match content_block { |
| 206 | ContentBlockStart::Thinking { .. } => { |
| 207 | is_thinking = true; |
| 208 | block_types.insert(index, "thinking".to_string()); |
| 209 | println!("{}", ds_sky("Thinking 💭").dimmed()); |
| 210 | } |
| 211 | ContentBlockStart::Text { .. } => { |
| 212 | if is_thinking { |
| 213 | println!(); |
| 214 | is_thinking = false; |
| 215 | } |
| 216 | block_types.insert(index, "text".to_string()); |
| 217 | } |
| 218 | ContentBlockStart::ToolUse { id, name, .. } => { |
| 219 | block_types.insert(index, "tool_use".to_string()); |
| 220 | tool_blocks.insert(index, (id, name.clone(), String::new())); |
| 221 | println!( |
| 222 | "{} {}", |
| 223 | ds_blue("Tool Call:").bold(), |
| 224 | ds_blue(&name).bold() |
| 225 | ); |
| 226 | } |
| 227 | }, |
| 228 | StreamEvent::ContentBlockDelta { index, delta } => match delta { |
| 229 | Delta::ThinkingDelta { thinking } => { |
| 230 | print!("{}", ds_sky(&thinking).dimmed()); |
| 231 | io::stdout().flush()?; |
| 232 | current_thinking.push_str(&thinking); |
| 233 | } |
| 234 | Delta::TextDelta { text } => { |
| 235 | print!("{text}"); |
| 236 | io::stdout().flush()?; |
| 237 | current_text.push_str(&text); |
| 238 | } |
| 239 | Delta::InputJsonDelta { partial_json } => { |
| 240 | if let Some((_id, _name, json)) = tool_blocks.get_mut(&index) { |
| 241 | json.push_str(&partial_json); |
| 242 | } |
| 243 | } |
| 244 | }, |
| 245 | StreamEvent::ContentBlockStop { index } => { |
| 246 | if let Some(block_type) = block_types.get(&index) |
| 247 | && block_type == "tool_use" |
| 248 | && let Some((_id, name, json_str)) = tool_blocks.get(&index) |
| 249 | { |
| 250 | if let Ok(parsed) = serde_json::from_str::<Value>(json_str) { |
| 251 | println!("{} {}", ds_blue("Tool Input:"), pretty_json(&parsed)); |
| 252 | } else if !json_str.is_empty() { |
| 253 | println!("{} {}", ds_blue("Tool Input:"), json_str); |
| 254 | } |
| 255 | println!("{}", ds_blue(&format!("Tool End: {name}")).dimmed()); |
| 256 | } |
| 257 | } |
| 258 | StreamEvent::MessageDelta { |
| 259 | usage: Some(usage), .. |
| 260 | } => { |
| 261 | stats.update(&usage); |
| 262 | } |
| 263 | _ => {} |
| 264 | } |
| 265 | } |
| 266 | println!(); |
| 267 | |
| 268 | let mut blocks = Vec::new(); |
| 269 | if !current_thinking.is_empty() { |
| 270 | blocks.push(ContentBlock::Thinking { |
| 271 | thinking: current_thinking, |
| 272 | }); |
| 273 | } |
| 274 | if !current_text.is_empty() { |
| 275 | blocks.push(ContentBlock::Text { |
| 276 | text: current_text, |
| 277 | cache_control: None, |
| 278 | }); |
| 279 | } |
| 280 | for (_index, (id, name, input)) in tool_blocks { |
| 281 | let parsed = serde_json::from_str::<Value>(&input).unwrap_or(Value::String(input)); |
| 282 | blocks.push(ContentBlock::ToolUse { |
| 283 | id, |
| 284 | name, |
| 285 | input: parsed, |
| 286 | caller: None, |
| 287 | }); |
| 288 | } |
| 289 | |
| 290 | messages.push(Message { |
| 291 | role: "assistant".to_string(), |
| 292 | content: blocks, |
| 293 | }); |
| 294 | } else { |
| 295 | let response = client.create_message(request).await?; |
| 296 | for block in &response.content { |
| 297 | match block { |
| 298 | ContentBlock::Thinking { thinking } => { |
| 299 | println!("{}", ds_sky("\nThinking 💭").dimmed()); |
| 300 | println!("{}", ds_sky(thinking).dimmed()); |
| 301 | } |
| 302 | ContentBlock::Text { text, .. } => { |
| 303 | println!("{text}"); |
| 304 | } |
| 305 | ContentBlock::ToolUse { name, input, .. } => { |
| 306 | println!( |
| 307 | "{} {}", |
| 308 | ds_blue("Tool Call:").bold(), |
| 309 | ds_blue(name).bold() |
| 310 | ); |
| 311 | println!("{}", pretty_json(input)); |
| 312 | } |
| 313 | ContentBlock::ToolResult { content, .. } => { |
| 314 | if let Ok(value) = serde_json::from_str::<Value>(content) { |
| 315 | println!("{}", pretty_json(&value)); |
| 316 | } else { |
| 317 | println!("{content}"); |
| 318 | } |
| 319 | } |
| 320 | } |
| 321 | } |
| 322 | |
| 323 | messages.push(Message { |
| 324 | role: "assistant".to_string(), |
| 325 | content: response.content, |
| 326 | }); |
| 327 | stats.update(&response.usage); |
| 328 | } |
| 329 | |
| 330 | Ok(()) |
| 331 | } |
| 332 | |
| 333 | async fn process_official_turn( |
| 334 | client: &DeepSeekClient, |
| 335 | options: &TextChatOptions, |
| 336 | messages: &mut Vec<Value>, |
| 337 | user_input: &str, |
| 338 | stats: &mut SessionStats, |
| 339 | ) -> Result<()> { |
| 340 | messages.push(json!({ "role": "user", "content": user_input })); |
| 341 | |
| 342 | let request = json!({ |
| 343 | "model": options.model, |
| 344 | "messages": messages, |
| 345 | "stream": false, |
| 346 | "max_tokens": options.max_tokens, |
| 347 | "temperature": options.temperature, |
| 348 | "top_p": options.top_p, |
| 349 | "tools": options.tools, |
| 350 | "tool_choice": options.tool_choice, |
| 351 | }); |
| 352 | |
| 353 | let response: Value = client |
| 354 | .post_json("/v1/text/chatcompletion_v2", &request) |
| 355 | .await?; |
| 356 | if let Some(text) = extract_text_from_response(&response) { |
| 357 | println!("{text}"); |
| 358 | messages.push(json!({ "role": "assistant", "content": text })); |
| 359 | } else { |
| 360 | println!("{}", pretty_json(&response)); |
| 361 | } |
| 362 | update_stats_from_official_response(&response, stats); |
| 363 | |
| 364 | Ok(()) |
| 365 | } |
| 366 | |
| 367 | fn extract_text_from_response(response: &Value) -> Option<String> { |
| 368 | let choices = response.get("choices")?.as_array()?; |
| 369 | let choice = choices.first()?; |
| 370 | if let Some(message) = choice.get("message") |
| 371 | && let Some(content) = message.get("content") |
| 372 | && let Some(text) = content.as_str() |
| 373 | { |
| 374 | return Some(text.to_string()); |
| 375 | } |
| 376 | if let Some(text) = choice.get("text").and_then(|v| v.as_str()) { |
| 377 | return Some(text.to_string()); |
| 378 | } |
| 379 | None |
| 380 | } |
| 381 | |
| 382 | fn build_system_prompt(system: Option<&str>, cache_system: bool) -> Option<SystemPrompt> { |
| 383 | let text = system?; |
| 384 | if !cache_system { |
| 385 | return Some(SystemPrompt::Text(text.to_string())); |
| 386 | } |
| 387 | let blocks = vec![SystemBlock { |
| 388 | block_type: "text".to_string(), |
| 389 | text: text.to_string(), |
| 390 | cache_control: Some(CacheControl { |
| 391 | cache_type: "ephemeral".to_string(), |
| 392 | }), |
| 393 | }]; |
| 394 | Some(SystemPrompt::Blocks(blocks)) |
| 395 | } |
| 396 | |
| 397 | fn cache_tools(tools: Option<Vec<Tool>>, cache_tools: bool) -> Option<Vec<Tool>> { |
| 398 | if !cache_tools { |
| 399 | return tools; |
| 400 | } |
| 401 | let mut tools = tools?; |
| 402 | if let Some(last) = tools.last_mut() { |
| 403 | last.cache_control = Some(CacheControl { |
| 404 | cache_type: "ephemeral".to_string(), |
| 405 | }); |
| 406 | } |
| 407 | Some(tools) |
| 408 | } |
| 409 | |
| 410 | fn update_stats_from_official_response(response: &Value, stats: &mut SessionStats) { |
| 411 | let usage = response.get("usage").and_then(|value| value.as_object()); |
| 412 | if let Some(usage) = usage { |
| 413 | let input = usage |
| 414 | .get("input_tokens") |
| 415 | .or_else(|| usage.get("prompt_tokens")) |
| 416 | .and_then(serde_json::Value::as_u64) |
| 417 | .and_then(|v| u32::try_from(v).ok()) |
| 418 | .unwrap_or(0); |
| 419 | let output = usage |
| 420 | .get("output_tokens") |
| 421 | .or_else(|| usage.get("completion_tokens")) |
| 422 | .and_then(serde_json::Value::as_u64) |
| 423 | .and_then(|v| u32::try_from(v).ok()) |
| 424 | .unwrap_or(0); |
| 425 | let total = usage |
| 426 | .get("total_tokens") |
| 427 | .and_then(serde_json::Value::as_u64) |
| 428 | .and_then(|v| u32::try_from(v).ok()) |
| 429 | .unwrap_or_else(|| input.saturating_add(output)); |
| 430 | stats.add_counts(input, output, Some(total)); |
| 431 | } |
| 432 | } |
| 433 | |
| 434 | fn matches_exit(input: &str) -> bool { |
| 435 | let normalized = input.trim().to_lowercase(); |
| 436 | matches!(normalized.as_str(), "exit" | "quit" | "q" | "/exit") |
| 437 | } |
| 438 | |
| 439 | fn handle_command_deepseek( |
| 440 | input: &str, |
| 441 | messages: &mut Vec<Message>, |
| 442 | options: Option<&TextChatOptions>, |
| 443 | stats: &mut SessionStats, |
| 444 | ) -> bool { |
| 445 | let trimmed = input.trim(); |
| 446 | if !trimmed.starts_with('/') { |
| 447 | return false; |
| 448 | } |
| 449 | |
| 450 | match trimmed { |
| 451 | "/help" => { |
| 452 | print_help(); |
| 453 | } |
| 454 | "/history" => { |
| 455 | println!("Messages: {}", messages.len()); |
| 456 | } |
| 457 | "/stats" => { |
| 458 | print_stats(stats); |
| 459 | } |
| 460 | "/clear" => { |
| 461 | messages.clear(); |
| 462 | stats.reset(); |
| 463 | if let Some(options) = options { |
| 464 | print_session_info( |
| 465 | options, |
| 466 | messages.len(), |
| 467 | options.tools.as_ref().map_or(0, std::vec::Vec::len), |
| 468 | ); |
| 469 | } |
| 470 | } |
| 471 | _ => { |
| 472 | println!("Unknown command. Type /help for available commands."); |
| 473 | } |
| 474 | } |
| 475 | true |
| 476 | } |
| 477 | |
| 478 | fn handle_command_official( |
| 479 | input: &str, |
| 480 | messages: &mut Vec<Value>, |
| 481 | options: Option<&TextChatOptions>, |
| 482 | stats: &mut SessionStats, |
| 483 | system_prompt: Option<&str>, |
| 484 | ) -> bool { |
| 485 | let trimmed = input.trim(); |
| 486 | if !trimmed.starts_with('/') { |
| 487 | return false; |
| 488 | } |
| 489 | |
| 490 | match trimmed { |
| 491 | "/help" => { |
| 492 | print_help(); |
| 493 | } |
| 494 | "/history" => { |
| 495 | println!("Messages: {}", messages.len()); |
| 496 | } |
| 497 | "/stats" => { |
| 498 | print_stats(stats); |
| 499 | } |
| 500 | "/clear" => { |
| 501 | messages.clear(); |
| 502 | if let Some(system) = system_prompt { |
| 503 | messages.push(json!({ "role": "system", "content": system })); |
| 504 | } |
| 505 | stats.reset(); |
| 506 | if let Some(options) = options { |
| 507 | print_session_info( |
| 508 | options, |
| 509 | messages.len(), |
| 510 | options.tools.as_ref().map_or(0, std::vec::Vec::len), |
| 511 | ); |
| 512 | } |
| 513 | } |
| 514 | _ => { |
| 515 | println!("Unknown command. Type /help for available commands."); |
| 516 | } |
| 517 | } |
| 518 | true |
| 519 | } |
| 520 | |
| 521 | fn print_banner(mode: &str) { |
| 522 | println!("{}", ds_blue("DeepSeek TUI").bold()); |
| 523 | println!("Mode: {mode}"); |
| 524 | println!("Type /help for commands. Use /exit to quit.\n"); |
| 525 | } |
| 526 | |
| 527 | fn print_help() { |
| 528 | println!("{}", ds_sky("Commands:").bold()); |
| 529 | println!(" /help Show this help"); |
| 530 | println!(" /clear Clear history (keeps system prompt)"); |
| 531 | println!(" /history Show message count"); |
| 532 | println!(" /stats Show token stats"); |
| 533 | println!(" /exit Exit session"); |
| 534 | } |
| 535 | |
| 536 | fn print_session_info(options: &TextChatOptions, messages: usize, tools: usize) { |
| 537 | let width = 56usize; |
| 538 | let header = "Session Info"; |
| 539 | println!("┌{}┐", "─".repeat(width)); |
| 540 | println!("│{:^width$}│", ds_blue(header).bold(), width = width); |
| 541 | println!("├{}┤", "─".repeat(width)); |
| 542 | println!( |
| 543 | "│ {:<width$}│", |
| 544 | format!("Model: {}", options.model), |
| 545 | width = width - 1 |
| 546 | ); |
| 547 | println!( |
| 548 | "│ {:<width$}│", |
| 549 | format!("Messages: {}", messages), |
| 550 | width = width - 1 |
| 551 | ); |
| 552 | println!( |
| 553 | "│ {:<width$}│", |
| 554 | format!("Tools: {}", tools), |
| 555 | width = width - 1 |
| 556 | ); |
| 557 | println!("└{}┘", "─".repeat(width)); |
| 558 | println!(); |
| 559 | } |
| 560 | |
| 561 | fn print_stats(stats: &SessionStats) { |
| 562 | let elapsed = stats.started.elapsed(); |
| 563 | let seconds = elapsed.as_secs(); |
| 564 | let hours = seconds / 3600; |
| 565 | let minutes = (seconds % 3600) / 60; |
| 566 | let secs = seconds % 60; |
| 567 | |
| 568 | println!("{}", ds_sky("Session Stats").bold()); |
| 569 | println!(" Duration: {hours:02}:{minutes:02}:{secs:02}"); |
| 570 | println!(" Input tokens: {}", stats.input_tokens); |
| 571 | println!(" Output tokens: {}", stats.output_tokens); |
| 572 | if stats.total_tokens > 0 { |
| 573 | println!(" Total tokens: {}", stats.total_tokens); |
| 574 | } |
| 575 | } |
| 576 | |
| 577 | fn ds_blue(text: &str) -> ColoredString { |
| 578 | let (r, g, b) = palette::DEEPSEEK_BLUE_RGB; |
| 579 | text.truecolor(r, g, b) |
| 580 | } |
| 581 | |
| 582 | fn ds_sky(text: &str) -> ColoredString { |
| 583 | let (r, g, b) = palette::DEEPSEEK_SKY_RGB; |
| 584 | text.truecolor(r, g, b) |
| 585 | } |
| 586 | |
| 587 | fn ds_red(text: &str) -> ColoredString { |
| 588 | let (r, g, b) = palette::DEEPSEEK_RED_RGB; |
| 589 | text.truecolor(r, g, b) |
| 590 | } |
| 591 | |
| 592 | struct SessionStats { |
| 593 | started: Instant, |
| 594 | input_tokens: u32, |
| 595 | output_tokens: u32, |
| 596 | total_tokens: u32, |
| 597 | } |
| 598 | |
| 599 | impl SessionStats { |
| 600 | fn new() -> Self { |
| 601 | Self { |
| 602 | started: Instant::now(), |
| 603 | input_tokens: 0, |
| 604 | output_tokens: 0, |
| 605 | total_tokens: 0, |
| 606 | } |
| 607 | } |
| 608 | |
| 609 | fn update(&mut self, usage: &Usage) { |
| 610 | self.add_counts(usage.input_tokens, usage.output_tokens, None); |
| 611 | } |
| 612 | |
| 613 | fn add_counts(&mut self, input: u32, output: u32, total: Option<u32>) { |
| 614 | self.input_tokens = self.input_tokens.saturating_add(input); |
| 615 | self.output_tokens = self.output_tokens.saturating_add(output); |
| 616 | let total = total.unwrap_or_else(|| input.saturating_add(output)); |
| 617 | self.total_tokens = self.total_tokens.saturating_add(total); |
| 618 | } |
| 619 | |
| 620 | fn reset(&mut self) { |
| 621 | self.started = Instant::now(); |
| 622 | self.input_tokens = 0; |
| 623 | self.output_tokens = 0; |
| 624 | self.total_tokens = 0; |
| 625 | } |
| 626 | } |
| 627 | |
| 628 | #[derive(Clone)] |
| 629 | struct CommandCompleter { |
| 630 | commands: Vec<String>, |
| 631 | } |
| 632 | |
| 633 | impl Helper for CommandCompleter {} |
| 634 | impl Hinter for CommandCompleter { |
| 635 | type Hint = String; |
| 636 | } |
| 637 | impl Highlighter for CommandCompleter {} |
| 638 | impl Validator for CommandCompleter {} |
| 639 | |
| 640 | impl Completer for CommandCompleter { |
| 641 | type Candidate = Pair; |
| 642 | |
| 643 | fn complete( |
| 644 | &self, |
| 645 | line: &str, |
| 646 | pos: usize, |
| 647 | _ctx: &RlContext<'_>, |
| 648 | ) -> Result<(usize, Vec<Pair>), ReadlineError> { |
| 649 | if !line.trim_start().starts_with('/') { |
| 650 | return Ok((pos, Vec::new())); |
| 651 | } |
| 652 | let start = line.rfind('/').unwrap_or(0); |
| 653 | let prefix = &line[start..pos]; |
| 654 | let matches = self |
| 655 | .commands |
| 656 | .iter() |
| 657 | .filter(|cmd| cmd.starts_with(prefix)) |
| 658 | .map(|cmd| Pair { |
| 659 | display: cmd.clone(), |
| 660 | replacement: cmd.clone(), |
| 661 | }) |
| 662 | .collect(); |
| 663 | Ok((start, matches)) |
| 664 | } |
| 665 | } |
| 666 | |
| 667 | fn create_editor() -> Result<Editor<CommandCompleter, DefaultHistory>> { |
| 668 | let helper = CommandCompleter { |
| 669 | commands: vec![ |
| 670 | "/help".to_string(), |
| 671 | "/clear".to_string(), |
| 672 | "/history".to_string(), |
| 673 | "/stats".to_string(), |
| 674 | "/exit".to_string(), |
| 675 | ], |
| 676 | }; |
| 677 | let mut editor = Editor::new()?; |
| 678 | editor.set_helper(Some(helper)); |
| 679 | if let Some(path) = history_path() { |
| 680 | let _ = editor.load_history(&path); |
| 681 | } |
| 682 | Ok(editor) |
| 683 | } |
| 684 | |
| 685 | fn read_prompt(editor: &mut Editor<CommandCompleter, DefaultHistory>) -> Result<Option<String>> { |
| 686 | match editor.readline("You> ") { |
| 687 | Ok(line) => { |
| 688 | let trimmed = line.trim().to_string(); |
| 689 | if !trimmed.is_empty() { |
| 690 | editor.add_history_entry(trimmed.as_str())?; |
| 691 | if let Some(path) = history_path() { |
| 692 | let _ = editor.append_history(&path); |
| 693 | } |
| 694 | } |
| 695 | Ok(Some(trimmed)) |
| 696 | } |
| 697 | Err(ReadlineError::Interrupted) => Ok(Some(String::new())), |
| 698 | Err(ReadlineError::Eof) => Ok(None), |
| 699 | Err(err) => Err(err.into()), |
| 700 | } |
| 701 | } |
| 702 | |
| 703 | fn history_path() -> Option<std::path::PathBuf> { |
| 704 | dirs::home_dir().map(|home| { |
| 705 | let dir = home.join(".deepseek"); |
| 706 | let _ = std::fs::create_dir_all(&dir); |
| 707 | dir.join("history") |
| 708 | }) |
| 709 | } |
| 710 | |
| 711 | async fn handle_line_deepseek( |
| 712 | line: String, |
| 713 | client: &DeepSeekClient, |
| 714 | options: &TextChatOptions, |
| 715 | messages: &mut Vec<Message>, |
| 716 | stats: &mut SessionStats, |
| 717 | ) -> Result<bool> { |
| 718 | let input = line.trim(); |
| 719 | if input.is_empty() { |
| 720 | return Ok(false); |
| 721 | } |
| 722 | if matches_exit(input) { |
| 723 | return Ok(true); |
| 724 | } |
| 725 | if handle_command_deepseek(input, messages, Some(options), stats) { |
| 726 | return Ok(false); |
| 727 | } |
| 728 | if let Err(error) = process_deepseek_turn(client, options, messages, input, stats).await { |
| 729 | eprintln!("{} {}", ds_red("Error:").bold(), error); |
| 730 | } |
| 731 | Ok(false) |
| 732 | } |
| 733 | |
| 734 | async fn handle_line_official( |
| 735 | line: String, |
| 736 | client: &DeepSeekClient, |
| 737 | options: &TextChatOptions, |
| 738 | messages: &mut Vec<Value>, |
| 739 | stats: &mut SessionStats, |
| 740 | system_prompt: Option<&str>, |
| 741 | ) -> Result<bool> { |
| 742 | let input = line.trim(); |
| 743 | if input.is_empty() { |
| 744 | return Ok(false); |
| 745 | } |
| 746 | if matches_exit(input) { |
| 747 | return Ok(true); |
| 748 | } |
| 749 | if handle_command_official(input, messages, Some(options), stats, system_prompt) { |
| 750 | return Ok(false); |
| 751 | } |
| 752 | if let Err(error) = process_official_turn(client, options, messages, input, stats).await { |
| 753 | eprintln!("{} {}", ds_red("Error:").bold(), error); |
| 754 | } |
| 755 | Ok(false) |
| 756 | } |
| 757 |