streaming_chat/
streaming_chat.rs1#![allow(dead_code)]
2#![allow(unused_variables)]
3use helios_engine::config::LLMConfig;
4use helios_engine::{ChatMessage, ChatSession, LLMClient};
9use std::io::{self, Write};
10
11#[tokio::main]
12async fn main() -> helios_engine::Result<()> {
13 println!("š Helios Engine - Streaming Example");
14 println!("=====================================\n");
15
16 let llm_config = LLMConfig {
18 model_name: "gpt-3.5-turbo".to_string(),
19 base_url: "https://api.openai.com/v1".to_string(),
20 api_key: std::env::var("OPENAI_API_KEY")
21 .unwrap_or_else(|_| "your-api-key-here".to_string()),
22 temperature: 0.7,
23 max_tokens: 2048,
24 };
25
26 let client = LLMClient::new(helios_engine::llm::LLMProviderType::Remote(llm_config)).await?;
27
28 println!("Example 1: Simple Streaming Response");
29 println!("======================================\n");
30
31 let messages = vec![
32 ChatMessage::system("You are a helpful assistant."),
33 ChatMessage::user("Write a short poem about coding."),
34 ];
35
36 print!("Assistant: ");
37 io::stdout().flush()?;
38
39 let response = client
40 .chat_stream(messages, None, |chunk| {
41 print!("{}", chunk);
42 io::stdout().flush().unwrap();
43 })
44 .await?;
45
46 println!("\n\n");
47
48 println!("Example 2: Interactive Streaming Chat");
49 println!("======================================\n");
50
51 let mut session = ChatSession::new().with_system_prompt("You are a helpful coding assistant.");
52
53 let questions = vec![
54 "What is Rust?",
55 "What are its main benefits?",
56 "Show me a simple example.",
57 ];
58
59 for question in questions {
60 println!("User: {}", question);
61 session.add_user_message(question);
62
63 print!("Assistant: ");
64 io::stdout().flush()?;
65
66 let response = client
67 .chat_stream(session.get_messages(), None, |chunk| {
68 print!("{}", chunk);
69 io::stdout().flush().unwrap();
70 })
71 .await?;
72
73 session.add_assistant_message(&response.content);
74 println!("\n");
75 }
76
77 println!("\nExample 3: Streaming with Thinking Tags");
78 println!("=========================================\n");
79 println!("When using models that support thinking tags (like o1),");
80 println!("you can detect and display them during streaming.\n");
81
82 struct ThinkingTracker {
83 in_thinking: bool,
84 thinking_buffer: String,
85 }
86
87 impl ThinkingTracker {
88 fn new() -> Self {
89 Self {
90 in_thinking: false,
91 thinking_buffer: String::new(),
92 }
93 }
94
95 fn process_chunk(&mut self, chunk: &str) -> String {
96 let mut output = String::new();
97 let mut chars = chunk.chars().peekable();
98
99 while let Some(c) = chars.next() {
100 if c == '<' {
101 let remaining: String = chars.clone().collect();
102 if remaining.starts_with("thinking>") {
103 self.in_thinking = true;
104 self.thinking_buffer.clear();
105 output.push_str("\nš [Thinking");
106 for _ in 0..9 {
107 chars.next();
108 }
109 continue;
110 } else if remaining.starts_with("/thinking>") {
111 self.in_thinking = false;
112 output.push_str("]\n");
113 for _ in 0..10 {
114 chars.next();
115 }
116 continue;
117 }
118 }
119
120 if self.in_thinking {
121 self.thinking_buffer.push(c);
122 if self.thinking_buffer.len() % 3 == 0 {
123 output.push('.');
124 }
125 } else {
126 output.push(c);
127 }
128 }
129
130 output
131 }
132 }
133
134 let messages = vec![ChatMessage::user(
135 "Solve this problem: What is 15 * 234 + 89?",
136 )];
137
138 let mut tracker = ThinkingTracker::new();
139 print!("Assistant: ");
140 io::stdout().flush()?;
141
142 let _response = client
143 .chat_stream(messages, None, |chunk| {
144 let output = tracker.process_chunk(chunk);
145 print!("{}", output);
146 io::stdout().flush().unwrap();
147 })
148 .await?;
149
150 println!("\n\nā
Streaming examples completed!");
151 println!("\nKey benefits of streaming:");
152 println!(" ⢠Real-time response display");
153 println!(" ⢠Better user experience for long responses");
154 println!(" ⢠Ability to show thinking/reasoning process");
155 println!(" ⢠Early cancellation possible (future feature)");
156
157 Ok(())
158}