Skip to content

Commit c007069

Browse files
pythopspythops
authored
Merge pull request #15 from pythops/support_stream_response
Support stream response
2 parents f90264b + bbb8b90 commit c007069

9 files changed

Lines changed: 482 additions & 331 deletions

File tree

‎Cargo.lock‎

Lines changed: 380 additions & 293 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "tenere"
3-
version = "0.6.0"
3+
version = "0.7.0"
44
authors = ["pythops <contact@pythops.com>"]
55
license = "AGPLv3"
66
edition = "2021"
@@ -21,4 +21,6 @@ ansi-to-tui = "3.1.0"
2121
clap = { version = "4", features = ["derive", "cargo"] }
2222
toml = { version = "0.7" }
2323
serde = { version = "1.0", features = ["derive"] }
24-
dirs = "5.0.0"
24+
dirs = "5.0.1"
25+
regex = "1.9.3"
26+
colored = "2.0.4"

‎src/app.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ pub struct App {
3535
pub focused_block: FocusedBlock,
3636
pub show_help_popup: bool,
3737
pub llm_messages: Vec<HashMap<String, String>>,
38+
pub answer: String,
3839
pub history: Vec<Vec<String>>,
3940
pub show_history_popup: bool,
4041
pub history_thread_index: usize,
@@ -55,6 +56,7 @@ impl App {
5556
focused_block: FocusedBlock::Prompt,
5657
show_help_popup: false,
5758
llm_messages: Vec::new(),
59+
answer: String::new(),
5860
history: Vec::new(),
5961
show_history_popup: false,
6062
history_thread_index: 0,

‎src/chatgpt.rs‎

Lines changed: 43 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,15 @@
1+
use crate::event::Event;
2+
use regex::Regex;
3+
use std::{thread, time};
4+
15
use crate::config::ChatGPTConfig;
2-
use crate::llm::LLM;
6+
use crate::llm::{LLMAnswer, LLM};
37
use reqwest::header::HeaderMap;
48
use serde_json::{json, Value};
59
use std;
610
use std::collections::HashMap;
11+
use std::io::Read;
12+
use std::sync::mpsc::Sender;
713

814
#[derive(Clone, Debug)]
915
pub struct ChatGPT {
@@ -41,12 +47,13 @@ impl LLM for ChatGPT {
4147
fn ask(
4248
&self,
4349
chat_messages: Vec<HashMap<String, String>>,
44-
) -> Result<String, Box<dyn std::error::Error>> {
50+
sender: &Sender<Event>,
51+
) -> Result<(), Box<dyn std::error::Error>> {
4552
let mut headers = HeaderMap::new();
46-
headers.insert("Content-Type", "application/json".parse().unwrap());
53+
headers.insert("Content-Type", "application/json".parse()?);
4754
headers.insert(
4855
"Authorization",
49-
format!("Bearer {}", self.openai_api_key).parse().unwrap(),
56+
format!("Bearer {}", self.openai_api_key).parse()?,
5057
);
5158

5259
let mut messages: Vec<HashMap<String, String>> = vec![
@@ -63,9 +70,12 @@ impl LLM for ChatGPT {
6370

6471
let body: Value = json!({
6572
"model": "gpt-3.5-turbo",
66-
"messages": messages
73+
"messages": messages,
74+
"stream": true,
6775
});
6876

77+
let mut buffer = String::new();
78+
6979
let response = self
7080
.client
7181
.post(&self.url)
@@ -74,17 +84,35 @@ impl LLM for ChatGPT {
7484
.send()?;
7585

7686
match response.error_for_status() {
77-
Ok(res) => {
78-
let response_body: Value = res.json()?;
79-
let answer = response_body["choices"][0]["message"]["content"]
80-
.as_str()
81-
.unwrap()
82-
.trim_matches('"')
83-
.to_string();
84-
85-
Ok(answer)
87+
Ok(mut res) => {
88+
let _answser = res.read_to_string(&mut buffer)?;
89+
90+
let re = Regex::new(r"data:\s(.*)").unwrap();
91+
92+
sender.send(Event::LLMEvent(LLMAnswer::StartAnswer))?;
93+
94+
for captures in re.captures_iter(&buffer) {
95+
if let Some(data_json) = captures.get(1) {
96+
if data_json.as_str() == "[DONE]" {
97+
sender.send(Event::LLMEvent(LLMAnswer::EndAnswer)).unwrap();
98+
break;
99+
}
100+
let x: Value = serde_json::from_str(data_json.as_str()).unwrap();
101+
102+
let msg = x["choices"][0]["delta"]["content"].as_str().unwrap_or("\n");
103+
104+
if msg != "null" {
105+
sender
106+
.send(Event::LLMEvent(LLMAnswer::Answer(msg.to_string())))
107+
.unwrap();
108+
}
109+
thread::sleep(time::Duration::from_millis(100));
110+
}
111+
}
86112
}
87-
Err(e) => Err(Box::new(e)),
113+
Err(e) => return Err(Box::new(e)),
88114
}
115+
116+
Ok(())
89117
}
90118
}

‎src/event.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use crate::app::AppResult;
2+
use crate::llm::LLMAnswer;
23
use crate::notification::Notification;
34
use crossterm::event::{self, Event as CrosstermEvent, KeyEvent, MouseEvent};
45
use std::sync::mpsc;
@@ -11,7 +12,7 @@ pub enum Event {
1112
Key(KeyEvent),
1213
Mouse(MouseEvent),
1314
Resize(u16, u16),
14-
LLMAnswer(String),
15+
LLMEvent(LLMAnswer),
1516
Notification(Notification),
1617
}
1718

‎src/handler.rs‎

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
1+
use crate::llm::LLMAnswer;
12
use crate::{
23
app::{App, AppResult, FocusedBlock, Mode},
34
event::Event,
45
};
6+
use colored::*;
57

68
use crate::llm::LLM;
79
use crossterm::event::{KeyCode, KeyEvent, KeyModifiers};
@@ -46,17 +48,22 @@ pub fn handle_key_events(
4648

4749
let llm_messages = app.llm_messages.clone();
4850

51+
app.spinner.active = true;
52+
app.chat.push("🤖: ".to_string());
53+
4954
thread::spawn(move || {
50-
let response = llm.ask(llm_messages.to_vec());
51-
sender
52-
.send(Event::LLMAnswer(match response {
53-
Ok(answer) => answer,
54-
Err(e) => e.to_string(),
55-
}))
56-
.unwrap();
55+
let res = llm.ask(llm_messages.to_vec(), &sender);
56+
if let Err(e) = res {
57+
sender
58+
.send(Event::LLMEvent(LLMAnswer::StartAnswer))
59+
.unwrap();
60+
sender
61+
.send(Event::LLMEvent(LLMAnswer::Answer(
62+
e.to_string().red().to_string(),
63+
)))
64+
.unwrap();
65+
}
5766
});
58-
app.spinner.active = true;
59-
app.chat.push("🤖: Waiting ..".to_string());
6067
}
6168

6269
// scroll down

‎src/llm.rs‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,24 @@
11
use crate::chatgpt::ChatGPT;
22
use crate::config::Config;
3+
use crate::event::Event;
34
use serde::Deserialize;
45
use std::collections::HashMap;
6+
use std::sync::mpsc::Sender;
57

68
use std::sync::Arc;
79
pub trait LLM: Send + Sync {
810
fn ask(
911
&self,
1012
chat_messages: Vec<HashMap<String, String>>,
11-
) -> Result<String, Box<dyn std::error::Error>>;
13+
sender: &Sender<Event>,
14+
) -> Result<(), Box<dyn std::error::Error>>;
15+
}
16+
17+
#[derive(Clone, Debug)]
18+
pub enum LLMAnswer {
19+
StartAnswer,
20+
Answer(String),
21+
EndAnswer,
1222
}
1323

1424
#[derive(Deserialize, Debug)]

‎src/main.rs‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use tenere::cli;
55
use tenere::config::Config;
66
use tenere::event::{Event, EventHandler};
77
use tenere::handler::handle_key_events;
8+
use tenere::llm::LLMAnswer;
89
use tenere::tui::Tui;
910
use tui::backend::CrosstermBackend;
1011
use tui::Terminal;
@@ -37,15 +38,22 @@ fn main() -> AppResult<()> {
3738
}
3839
Event::Mouse(_) => {}
3940
Event::Resize(_, _) => {}
40-
Event::LLMAnswer(answer) => {
41-
app.chat.pop();
42-
app.spinner.active = false;
43-
app.chat.push(format!("🤖: {}\n", answer));
44-
app.chat.push("\n".to_string());
41+
Event::LLMEvent(LLMAnswer::Answer(answer)) => {
42+
app.answer.push_str(answer.as_str());
43+
}
44+
Event::LLMEvent(LLMAnswer::EndAnswer) => {
4545
let mut conv: HashMap<String, String> = HashMap::new();
4646
conv.insert("role".to_string(), "user".to_string());
47-
conv.insert("content".to_string(), answer);
47+
conv.insert("content".to_string(), app.answer.clone());
4848
app.llm_messages.push(conv);
49+
app.chat.push(app.answer.clone());
50+
app.chat.push("\n".to_string());
51+
app.answer.clear();
52+
}
53+
Event::LLMEvent(LLMAnswer::StartAnswer) => {
54+
app.spinner.active = false;
55+
app.chat.pop();
56+
app.chat.push("🤖: ".to_string());
4957
}
5058
Event::Notification(notification) => {
5159
app.notifications.push(notification);

‎src/ui.rs‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,9 @@ pub fn render<B: Backend>(app: &mut App, frame: &mut Frame<'_, B>) {
192192

193193
// Chat block
194194
let chat = {
195-
let messages: String = app.chat.iter().map(|m| m.to_string()).collect();
195+
let mut messages: String = app.chat.iter().map(|m| m.to_string()).collect();
196+
197+
messages.push_str(app.answer.as_str());
196198

197199
let messages_height = {
198200
let mut height: u16 = 0;
@@ -203,6 +205,11 @@ pub fn render<B: Backend>(app: &mut App, frame: &mut Frame<'_, B>) {
203205
height += line.width() as u16 / app_area.width;
204206
}
205207
}
208+
209+
for line in app.answer.lines() {
210+
height += 1;
211+
height += line.width() as u16 / app_area.width;
212+
}
206213
height
207214
};
208215

@@ -219,12 +226,11 @@ pub fn render<B: Backend>(app: &mut App, frame: &mut Frame<'_, B>) {
219226
scroll = height_diff + app.scroll;
220227
}
221228

222-
// // scroll up case
229+
// scroll up case
223230
if height_diff > 0 && -app.scroll > height_diff {
224231
app.scroll = -height_diff;
225232
}
226-
//
227-
// // Scroll down case
233+
// Scroll down case
228234
if height_diff > 0 && app.scroll > 0 {
229235
app.scroll = 0;
230236
}

0 commit comments

Comments
 (0)