refactor: use string instead of vector of string

This commit is contained in:
iff 2025-04-10 01:31:41 +02:00
parent 6a4b51d340
commit 32e2cb71e8
2 changed files with 21 additions and 23 deletions

View file

@ -38,14 +38,14 @@ enum State {
} }
pub struct Buffer { pub struct Buffer {
pub buf: Vec<String>, pub buf: String,
state: State, state: State,
} }
impl Buffer { impl Buffer {
pub fn new() -> Self { pub fn new() -> Self {
Buffer { Buffer {
buf: vec![], buf: String::new(),
state: State::Write, state: State::Write,
} }
} }
@ -53,12 +53,12 @@ impl Buffer {
match self.state { match self.state {
State::Write => self.proc_write(data), State::Write => self.proc_write(data),
State::Think => self.proc_think(data), State::Think => self.proc_think(data),
State::Buf => self.buf.push(data.to_string()), State::Buf => self.buf.push_str(data),
} }
} }
pub fn print_return_remain(&mut self) -> String { pub fn print_return_remain(&mut self) -> String {
let buffered = self.buf.join("").trim().to_string(); let buffered = self.buf.trim().to_string();
self.buf.clear(); self.buf.clear();
if self.state == State::Buf { if self.state == State::Buf {
return buffered; return buffered;
@ -75,15 +75,14 @@ impl Buffer {
fn proc_write(&mut self, data: &str) { fn proc_write(&mut self, data: &str) {
if !data.contains("\n") { if !data.contains("\n") {
self.buf.push(data.to_string()); self.buf.push_str(data);
let buffered = self.buf.join("").trim().to_string(); let buffered = self.buf.trim().to_string();
let filled = fill(&buffered); let filled = fill(&buffered);
if let Some(filled) = filled { if let Some(filled) = filled {
self.buf.clear(); self.buf.clear();
let formatted = clear_format(&filled); let formatted = clear_format(&filled);
eprint!("{}", formatted); eprint!("{}", formatted);
self.buf self.buf.push_str(formatted.split_once("\n").unwrap().1);
.push(formatted.split_once("\n").unwrap().1.to_string());
std::io::stdout().flush().unwrap(); std::io::stdout().flush().unwrap();
return; return;
} }
@ -97,8 +96,8 @@ impl Buffer {
let lines = data.split_once("\n").unwrap(); let lines = data.split_once("\n").unwrap();
let first = lines.0; let first = lines.0;
let last = lines.1; let last = lines.1;
self.buf.push(first.to_string()); self.buf.push_str(first);
let buffered = self.buf.join("").trim().to_string(); let buffered = self.buf.trim().to_string();
self.buf.clear(); self.buf.clear();
if buffered.ends_with("<note>") { if buffered.ends_with("<note>") {
let warn = format!("{}:", t!("ai-suggestion")) let warn = format!("{}:", t!("ai-suggestion"))
@ -148,15 +147,14 @@ impl Buffer {
fn proc_think(&mut self, data: &str) { fn proc_think(&mut self, data: &str) {
if !data.contains("\n") { if !data.contains("\n") {
self.buf.push(data.to_string()); self.buf.push_str(data);
let buffered = self.buf.join("").trim().to_string(); let buffered = self.buf.trim().to_string();
let filled = fill(&buffered); let filled = fill(&buffered);
if let Some(filled) = filled { if let Some(filled) = filled {
self.buf.clear(); self.buf.clear();
let formatted = clear_format(&filled); let formatted = clear_format(&filled);
eprint!("{}", formatted); eprint!("{}", formatted);
self.buf self.buf.push_str(formatted.split_once("\n").unwrap().1);
.push(formatted.split_once("\n").unwrap().1.to_string());
std::io::stdout().flush().unwrap(); std::io::stdout().flush().unwrap();
return; return;
} }
@ -170,8 +168,8 @@ impl Buffer {
let lines = data.split_once("\n").unwrap(); let lines = data.split_once("\n").unwrap();
let first = lines.0; let first = lines.0;
let last = lines.1; let last = lines.1;
self.buf.push(first.to_string()); self.buf.push_str(first);
let buffered = self.buf.join("").trim().to_string(); let buffered = self.buf.trim().to_string();
self.buf.clear(); self.buf.clear();
if buffered.ends_with("</think>") { if buffered.ends_with("</think>") {
let tag = "</think>"; let tag = "</think>";

View file

@ -130,27 +130,27 @@ pub async fn ai_suggestion(last_command: &str, error_msg: &str) {
.await; .await;
let mut stream = res.unwrap().bytes_stream(); let mut stream = res.unwrap().bytes_stream();
let mut json_buffer = vec![]; let mut json_buffer = String::new();
let mut buffer = buffer::Buffer::new(); let mut buffer = buffer::Buffer::new();
while let Some(item) = stream.next().await { while let Some(item) = stream.next().await {
let item = item.unwrap(); let item = item.unwrap();
let str = std::str::from_utf8(&item).unwrap(); let str = std::str::from_utf8(&item).unwrap();
if json_buffer.is_empty() { if json_buffer.is_empty() {
json_buffer.push(str.to_string()); json_buffer.push_str(str);
continue; continue;
} }
if !str.contains("\n\ndata: {") { if !str.contains("\n\ndata: {") {
json_buffer.push(str.to_string()); json_buffer.push_str(str);
continue; continue;
} }
let data_loc = str.find("\n\ndata: {").unwrap(); let data_loc = str.find("\n\ndata: {").unwrap();
let split = str.split_at(data_loc); let split = str.split_at(data_loc);
json_buffer.push(split.0.to_string()); json_buffer.push_str(split.0);
let working_str = json_buffer.join(""); let working_str = json_buffer.clone();
json_buffer.clear(); json_buffer.clear();
json_buffer.push(split.1.to_string()); json_buffer.push_str(split.1);
for part in working_str.split("\n\n") { for part in working_str.split("\n\n") {
if let Some(data) = part.strip_prefix("data: ") { if let Some(data) = part.strip_prefix("data: ") {
@ -168,7 +168,7 @@ pub async fn ai_suggestion(last_command: &str, error_msg: &str) {
} }
} }
if !json_buffer.is_empty() { if !json_buffer.is_empty() {
let working_str = json_buffer.join(""); let working_str = json_buffer.clone();
for part in working_str.split("\n\n") { for part in working_str.split("\n\n") {
if let Some(data) = part.strip_prefix("data: ") { if let Some(data) = part.strip_prefix("data: ") {
if data == "[DONE]" { if data == "[DONE]" {