#[cfg(test)] #[path = "backend_test.rs"] mod tests; use anyhow::Result; use async_trait::async_trait; use strum::EnumIter; use strum::EnumVariantNames; use strum::IntoEnumIterator; use tokio::sync::mpsc; use super::Author; use super::EditorContext; use super::Event; #[derive(Clone, Debug, PartialEq, Eq, EnumIter, EnumVariantNames, strum::Display)] #[strum(serialize_all = "lowercase")] pub enum BackendName { LangChain, Ollama, OpenAI, Claude, Gemini, } impl BackendName { pub fn parse(text: String) -> Option { return BackendName::iter().find(|e| return e.to_string() == text); } } pub struct BackendPrompt { pub text: String, pub backend_context: String, } impl BackendPrompt { pub fn new(text: String, backend_context: String) -> BackendPrompt { return BackendPrompt { text, backend_context, }; } pub fn append_chat_context(&mut self, editor_context: &Option) { if let Some(context) = editor_context { let lang = &context.language; let code = &context.code; let system_prompt = format!(". The coding language is {lang}. Add language to any code blocks."); self.text += &system_prompt; if !code.is_empty() { let code_prompt = format!(" The code is the following:\n{code}"); self.text += &code_prompt; } } else { self.text += ". Add language to any code blocks." } } } #[derive(Clone)] pub struct BackendResponse { pub author: Author, pub text: String, pub done: bool, pub context: Option, } #[async_trait] pub trait Backend { /// Returns the name of the backend fn name(&self) -> BackendName; /// Used at startup to verify all configurations are available to work with /// the backend. async fn health_check(&self) -> Result<()>; /// Called when using the `/modellist` slash commands to provide all /// available models for the backend. async fn list_models<'a>(&'a self) -> Result>; /// Requests completions from the backend. Completion results may be /// streamed back to the UI by passing each response through a channel. /// /// Upon receiving all results, a final `done` boolean /// is provided as the last message to the channel. /// /// In order for a backend to maintain history, a context array may be /// provided by the backend. This can be passed alongside the `done` /// boolean, and will be provided on the next prompt to the backend. async fn get_completion<'a>( &self, prompt: BackendPrompt, tx: &'a mpsc::UnboundedSender, ) -> Result<()>; } pub type BackendBox = Box;