streaming_chat/
streaming_chat.rs

1#![allow(dead_code)]
2#![allow(unused_variables)]
3use helios_engine::config::LLMConfig;
4/// Example: Using streaming responses with Helios Engine
5///
6/// This example demonstrates how to use the streaming API to get
7/// real-time responses from the LLM, including detection of thinking tags.
8use 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    // Setup LLM configuration
17    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}