Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

replace OpenAIRequest with more generalized Box<dyn CompletionRequest>

Browse source
This commit is contained in:
KCaverly 2023-10-22 14:33:19 +02:00
parent 05ae978cb7
commit d813ae8845
6 changed files with 58 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

20
crates/assistant/src/assistant_panel.rs
View file

@ -6,8 +6,11 @@ use crate::{
SavedMessage,
};
use ai::providers::open_ai::{
stream_completion, OpenAICompletionProvider, OpenAIRequest, RequestMessage, OPENAI_API_URL,
use ai::{
completion::CompletionRequest,
providers::open_ai::{
stream_completion, OpenAICompletionProvider, OpenAIRequest, RequestMessage, OPENAI_API_URL,
},
};
use ai::prompts::repository_context::PromptCodeSnippet;
@ -745,13 +748,14 @@ impl AssistantPanel {
content: prompt,
});
let request = OpenAIRequest {
let request = Box::new(OpenAIRequest {
model: model.full_name().into(),
messages,
stream: true,
stop: vec!["|END|>".to_string()],
temperature,
};
});
codegen.update(&mut cx, |codegen, cx| codegen.start(request, cx));
anyhow::Ok(())
})
@ -1735,7 +1739,7 @@ impl Conversation {
return Default::default();
};
let request = OpenAIRequest {
let request: Box<dyn CompletionRequest> = Box::new(OpenAIRequest {
model: self.model.full_name().to_string(),
messages: self
.messages(cx)
@ -1745,7 +1749,7 @@ impl Conversation {
stream: true,
stop: vec![],
temperature: 1.0,
};
});
let stream = stream_completion(api_key, cx.background().clone(), request);
let assistant_message = self
@ -2025,13 +2029,13 @@ impl Conversation {
"Summarize the conversation into a short title without punctuation"
.into(),
}));
let request = OpenAIRequest {
let request: Box<dyn CompletionRequest> = Box::new(OpenAIRequest {
model: self.model.full_name().to_string(),
messages: messages.collect(),
stream: true,
stop: vec![],
temperature: 1.0,
};
});
let stream = stream_completion(api_key, cx.background().clone(), request);
self.pending_summary = cx.spawn(|this, mut cx| {

25
crates/assistant/src/codegen.rs
View file

@ -1,6 +1,5 @@
use crate::streaming_diff::{Hunk, StreamingDiff};
use ai::completion::CompletionProvider;
use ai::providers::open_ai::OpenAIRequest;
use ai::completion::{CompletionProvider, CompletionRequest};
use anyhow::Result;
use editor::{multi_buffer, Anchor, MultiBuffer, MultiBufferSnapshot, ToOffset, ToPoint};
use futures::{channel::mpsc, SinkExt, Stream, StreamExt};
@ -96,7 +95,7 @@ impl Codegen {
self.error.as_ref()
}
pub fn start(&mut self, prompt: OpenAIRequest, cx: &mut ModelContext<Self>) {
pub fn start(&mut self, prompt: Box<dyn CompletionRequest>, cx: &mut ModelContext<Self>) {
let range = self.range();
let snapshot = self.snapshot.clone();
let selected_text = snapshot
@ -336,6 +335,7 @@ fn strip_markdown_codeblock(
#[cfg(test)]
mod tests {
use super::*;
use ai::providers::dummy::DummyCompletionRequest;
use futures::{
future::BoxFuture,
stream::{self, BoxStream},
@ -381,7 +381,10 @@ mod tests {
cx,
)
});
codegen.update(cx, |codegen, cx| codegen.start(Default::default(), cx));
let request = Box::new(DummyCompletionRequest {
name: "test".to_string(),
});
codegen.update(cx, |codegen, cx| codegen.start(request, cx));
let mut new_text = concat!(
" let mut x = 0;\n",
@ -443,7 +446,11 @@ mod tests {
cx,
)
});
codegen.update(cx, |codegen, cx| codegen.start(Default::default(), cx));
let request = Box::new(DummyCompletionRequest {
name: "test".to_string(),
});
codegen.update(cx, |codegen, cx| codegen.start(request, cx));
let mut new_text = concat!(
"t mut x = 0;\n",
@ -505,7 +512,11 @@ mod tests {
cx,
)
});
codegen.update(cx, |codegen, cx| codegen.start(Default::default(), cx));
let request = Box::new(DummyCompletionRequest {
name: "test".to_string(),
});
codegen.update(cx, |codegen, cx| codegen.start(request, cx));
let mut new_text = concat!(
"let mut x = 0;\n",
@ -617,7 +628,7 @@ mod tests {
impl CompletionProvider for TestCompletionProvider {
fn complete(
&self,
_prompt: OpenAIRequest,
_prompt: Box<dyn CompletionRequest>,
) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
let (tx, rx) = mpsc::channel(1);
*self.last_completion_tx.lock() = Some(tx);