Merge pull request #66 from SilasMarvin/release/v0.6.2

Release/v0.6.2
This commit is contained in:
Silas Marvin
2024-08-27 09:13:59 -07:00
committed by GitHub
20 changed files with 115 additions and 120 deletions

View File

@@ -1,6 +1,6 @@
[package] [package]
name = "lsp-ai" name = "lsp-ai"
version = "0.6.1" version = "0.6.2"
description.workspace = true description.workspace = true
repository.workspace = true repository.workspace = true

View File

@@ -14,12 +14,12 @@ const fn true_default() -> bool {
} }
#[derive(Clone, Debug, Deserialize, Serialize)] #[derive(Clone, Debug, Deserialize, Serialize)]
pub struct PostProcess { pub(crate) struct PostProcess {
pub extractor: Option<String>, pub(crate) extractor: Option<String>,
#[serde(default = "true_default")] #[serde(default = "true_default")]
pub remove_duplicate_start: bool, pub(crate) remove_duplicate_start: bool,
#[serde(default = "true_default")] #[serde(default = "true_default")]
pub remove_duplicate_end: bool, pub(crate) remove_duplicate_end: bool,
} }
impl Default for PostProcess { impl Default for PostProcess {
@@ -34,7 +34,7 @@ impl Default for PostProcess {
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize)]
#[serde(tag = "type")] #[serde(tag = "type")]
pub enum ValidSplitter { pub(crate) enum ValidSplitter {
#[serde(rename = "tree_sitter")] #[serde(rename = "tree_sitter")]
TreeSitter(TreeSitter), TreeSitter(TreeSitter),
#[serde(rename = "text_splitter")] #[serde(rename = "text_splitter")]
@@ -56,11 +56,11 @@ const fn chunk_overlap_default() -> usize {
} }
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize)]
pub struct TreeSitter { pub(crate) struct TreeSitter {
#[serde(default = "chunk_size_default")] #[serde(default = "chunk_size_default")]
pub chunk_size: usize, pub(crate) chunk_size: usize,
#[serde(default = "chunk_overlap_default")] #[serde(default = "chunk_overlap_default")]
pub chunk_overlap: usize, pub(crate) chunk_overlap: usize,
} }
impl Default for TreeSitter { impl Default for TreeSitter {
@@ -73,39 +73,39 @@ impl Default for TreeSitter {
} }
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize)]
pub struct TextSplitter { pub(crate) struct TextSplitter {
#[serde(default = "chunk_size_default")] #[serde(default = "chunk_size_default")]
pub chunk_size: usize, pub(crate) chunk_size: usize,
} }
#[derive(Debug, Clone, Deserialize, Default)] #[derive(Debug, Clone, Deserialize, Default)]
pub struct EmbeddingPrefix { pub(crate) struct EmbeddingPrefix {
#[serde(default)] #[serde(default)]
pub storage: String, pub(crate) storage: String,
#[serde(default)] #[serde(default)]
pub retrieval: String, pub(crate) retrieval: String,
} }
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize)]
pub struct OllamaEmbeddingModel { pub(crate) struct OllamaEmbeddingModel {
// The generate endpoint, default: 'http://localhost:11434/api/embeddings' // The generate endpoint, default: 'http://localhost:11434/api/embeddings'
pub endpoint: Option<String>, pub(crate) endpoint: Option<String>,
// The model name // The model name
pub model: String, pub(crate) model: String,
// The prefix to apply to the embeddings // The prefix to apply to the embeddings
#[serde(default)] #[serde(default)]
pub prefix: EmbeddingPrefix, pub(crate) prefix: EmbeddingPrefix,
} }
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize)]
#[serde(tag = "type")] #[serde(tag = "type")]
pub enum ValidEmbeddingModel { pub(crate) enum ValidEmbeddingModel {
#[serde(rename = "ollama")] #[serde(rename = "ollama")]
Ollama(OllamaEmbeddingModel), Ollama(OllamaEmbeddingModel),
} }
#[derive(Debug, Clone, Copy, Deserialize)] #[derive(Debug, Clone, Copy, Deserialize)]
pub enum VectorDataType { pub(crate) enum VectorDataType {
#[serde(rename = "f32")] #[serde(rename = "f32")]
F32, F32,
#[serde(rename = "binary")] #[serde(rename = "binary")]
@@ -114,11 +114,11 @@ pub enum VectorDataType {
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize)]
pub(crate) struct VectorStore { pub(crate) struct VectorStore {
pub crawl: Option<Crawl>, pub(crate) crawl: Option<Crawl>,
#[serde(default)] #[serde(default)]
pub splitter: ValidSplitter, pub(crate) splitter: ValidSplitter,
pub embedding_model: ValidEmbeddingModel, pub(crate) embedding_model: ValidEmbeddingModel,
pub data_type: VectorDataType, pub(crate) data_type: VectorDataType,
} }
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize)]
@@ -265,20 +265,20 @@ const fn n_ctx_default() -> u32 {
#[cfg(feature = "llama_cpp")] #[cfg(feature = "llama_cpp")]
#[derive(Clone, Debug, Deserialize)] #[derive(Clone, Debug, Deserialize)]
#[serde(deny_unknown_fields)] #[serde(deny_unknown_fields)]
pub struct LLaMACPP { pub(crate) struct LLaMACPP {
// Which model to use // Which model to use
pub repository: Option<String>, pub(crate) repository: Option<String>,
pub name: Option<String>, pub(crate) name: Option<String>,
pub file_path: Option<String>, pub(crate) file_path: Option<String>,
// The layers to put on the GPU // The layers to put on the GPU
#[serde(default = "n_gpu_layers_default")] #[serde(default = "n_gpu_layers_default")]
pub n_gpu_layers: u32, pub(crate) n_gpu_layers: u32,
// The context size // The context size
#[serde(default = "n_ctx_default")] #[serde(default = "n_ctx_default")]
pub n_ctx: u32, pub(crate) n_ctx: u32,
// The maximum requests per second // The maximum requests per second
#[serde(default = "max_requests_per_second_default")] #[serde(default = "max_requests_per_second_default")]
pub max_requests_per_second: f32, pub(crate) max_requests_per_second: f32,
} }
#[derive(Clone, Debug, Deserialize)] #[derive(Clone, Debug, Deserialize)]
@@ -349,7 +349,7 @@ pub(crate) struct Completion {
} }
#[derive(Clone, Debug, Deserialize)] #[derive(Clone, Debug, Deserialize)]
pub struct Chat { pub(crate) struct Chat {
// The trigger text // The trigger text
pub(crate) trigger: String, pub(crate) trigger: String,
// The name to display in the editor // The name to display in the editor
@@ -362,7 +362,7 @@ pub struct Chat {
} }
#[derive(Clone, Debug, Deserialize)] #[derive(Clone, Debug, Deserialize)]
pub struct Action { pub(crate) struct Action {
// The name to display in the editor // The name to display in the editor
pub(crate) action_display_name: String, pub(crate) action_display_name: String,
// The model key to use // The model key to use
@@ -395,13 +395,13 @@ pub(crate) struct ValidClientParams {
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct Config { pub(crate) struct Config {
pub(crate) config: ValidConfig, pub(crate) config: ValidConfig,
pub(crate) client_params: ValidClientParams, pub(crate) client_params: ValidClientParams,
} }
impl Config { impl Config {
pub fn new(mut args: Value) -> Result<Self> { pub(crate) fn new(mut args: Value) -> Result<Self> {
// Validate that the models specified are there so we can unwrap // Validate that the models specified are there so we can unwrap
let configuration_args = args let configuration_args = args
.as_object_mut() .as_object_mut()
@@ -422,23 +422,19 @@ impl Config {
// Helpers for the backends /////////// // Helpers for the backends ///////////
/////////////////////////////////////// ///////////////////////////////////////
pub fn get_chats(&self) -> &Vec<Chat> { pub(crate) fn get_chats(&self) -> &Vec<Chat> {
&self.config.chats &self.config.chats
} }
pub fn get_actions(&self) -> &Vec<Action> { pub(crate) fn get_actions(&self) -> &Vec<Action> {
&self.config.actions &self.config.actions
} }
pub fn is_completions_enabled(&self) -> bool { pub(crate) fn get_completions_post_process(&self) -> Option<&PostProcess> {
self.config.completion.is_some()
}
pub fn get_completions_post_process(&self) -> Option<&PostProcess> {
self.config.completion.as_ref().map(|x| &x.post_process) self.config.completion.as_ref().map(|x| &x.post_process)
} }
pub fn get_completion_transformer_max_requests_per_second(&self) -> anyhow::Result<f32> { pub(crate) fn get_completion_transformer_max_requests_per_second(&self) -> anyhow::Result<f32> {
match &self match &self
.config .config
.models .models
@@ -470,7 +466,7 @@ impl Config {
// For teesting use only // For teesting use only
#[cfg(test)] #[cfg(test)]
impl Config { impl Config {
pub fn default_with_file_store_without_models() -> Self { pub(crate) fn default_with_file_store_without_models() -> Self {
Self { Self {
config: ValidConfig { config: ValidConfig {
memory: ValidMemoryBackend::FileStore(FileStore { crawl: None }), memory: ValidMemoryBackend::FileStore(FileStore { crawl: None }),

View File

@@ -22,7 +22,7 @@ impl Crawl {
} }
#[instrument(skip(self, f))] #[instrument(skip(self, f))]
pub fn maybe_do_crawl( pub(crate) fn maybe_do_crawl(
&mut self, &mut self,
triggered_file: Option<String>, triggered_file: Option<String>,
mut f: impl FnMut(&config::Crawl, &str) -> anyhow::Result<bool>, mut f: impl FnMut(&config::Crawl, &str) -> anyhow::Result<bool>,

View File

@@ -5,12 +5,12 @@ pub(crate) enum GenerationStream {}
#[derive(Debug, PartialEq, Clone, Deserialize, Serialize)] #[derive(Debug, PartialEq, Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")] #[serde(rename_all = "camelCase")]
pub struct GenerationStreamParams { pub(crate) struct GenerationStreamParams {
pub partial_result_token: ProgressToken, pub(crate) partial_result_token: ProgressToken,
// This field was "mixed-in" from TextDocumentPositionParams // This field was "mixed-in" from TextDocumentPositionParams
#[serde(flatten)] #[serde(flatten)]
pub text_document_position: TextDocumentPositionParams, pub(crate) text_document_position: TextDocumentPositionParams,
} }
#[derive(Debug, PartialEq, Clone, Deserialize, Serialize)] #[derive(Debug, PartialEq, Clone, Deserialize, Serialize)]

View File

@@ -2,7 +2,7 @@ use crate::config::ValidEmbeddingModel;
mod ollama; mod ollama;
pub fn normalize(mut vector: Vec<f32>) -> Vec<f32> { fn normalize(mut vector: Vec<f32>) -> Vec<f32> {
let magnitude = (vector.iter().map(|&x| x * x).sum::<f32>()).sqrt(); let magnitude = (vector.iter().map(|&x| x * x).sum::<f32>()).sqrt();
if magnitude != 0.0 { if magnitude != 0.0 {
@@ -15,13 +15,13 @@ pub fn normalize(mut vector: Vec<f32>) -> Vec<f32> {
} }
#[derive(Clone, Copy)] #[derive(Clone, Copy)]
pub enum EmbeddingPurpose { pub(crate) enum EmbeddingPurpose {
Storage, Storage,
Retrieval, Retrieval,
} }
#[async_trait::async_trait] #[async_trait::async_trait]
pub trait EmbeddingModel { pub(crate) trait EmbeddingModel {
async fn embed( async fn embed(
&self, &self,
batch: Vec<&str>, batch: Vec<&str>,

View File

@@ -8,29 +8,29 @@ use crate::config;
use super::{normalize, EmbeddingModel, EmbeddingPurpose}; use super::{normalize, EmbeddingModel, EmbeddingPurpose};
#[derive(Deserialize)] #[derive(Deserialize)]
pub struct Embed { pub(crate) struct Embed {
embedding: Vec<f32>, embedding: Vec<f32>,
} }
#[derive(Deserialize)] #[derive(Deserialize)]
pub struct EmbedError { pub(crate) struct EmbedError {
error: Value, error: Value,
} }
#[derive(Deserialize)] #[derive(Deserialize)]
#[serde(untagged)] #[serde(untagged)]
pub enum EmbedResponse { pub(crate) enum EmbedResponse {
Success(Embed), Success(Embed),
Error(EmbedError), Error(EmbedError),
Other(HashMap<String, Value>), Other(HashMap<String, Value>),
} }
pub struct Ollama { pub(crate) struct Ollama {
config: config::OllamaEmbeddingModel, config: config::OllamaEmbeddingModel,
} }
impl Ollama { impl Ollama {
pub fn new(config: config::OllamaEmbeddingModel) -> Self { pub(crate) fn new(config: config::OllamaEmbeddingModel) -> Self {
Self { config } Self { config }
} }
} }

View File

@@ -28,7 +28,7 @@ impl AdditionalFileStoreParams {
} }
#[derive(Clone)] #[derive(Clone)]
pub struct File { pub(crate) struct File {
rope: Rope, rope: Rope,
tree: Option<Tree>, tree: Option<Tree>,
} }
@@ -38,11 +38,11 @@ impl File {
Self { rope, tree } Self { rope, tree }
} }
pub fn rope(&self) -> &Rope { pub(crate) fn rope(&self) -> &Rope {
&self.rope &self.rope
} }
pub fn tree(&self) -> Option<&Tree> { pub(crate) fn tree(&self) -> Option<&Tree> {
self.tree.as_ref() self.tree.as_ref()
} }
} }
@@ -278,15 +278,15 @@ impl FileStore {
}) })
} }
pub fn file_map(&self) -> &RwLock<HashMap<String, File>> { pub(crate) fn file_map(&self) -> &RwLock<HashMap<String, File>> {
&self.file_map &self.file_map
} }
pub fn contains_file(&self, uri: &str) -> bool { pub(crate) fn contains_file(&self, uri: &str) -> bool {
self.file_map.read().contains_key(uri) self.file_map.read().contains_key(uri)
} }
pub fn position_to_byte(&self, position: &TextDocumentPositionParams) -> anyhow::Result<usize> { pub(crate) fn position_to_byte(&self, position: &TextDocumentPositionParams) -> anyhow::Result<usize> {
let file_map = self.file_map.read(); let file_map = self.file_map.read();
let uri = position.text_document.uri.to_string(); let uri = position.text_document.uri.to_string();
let file = file_map let file = file_map
@@ -494,7 +494,7 @@ impl MemoryBackend for FileStore {
// For testing use only // For testing use only
#[cfg(test)] #[cfg(test)]
impl FileStore { impl FileStore {
pub fn default_with_filler_file() -> anyhow::Result<Self> { pub(crate) fn default_with_filler_file() -> anyhow::Result<Self> {
let config = Config::default_with_file_store_without_models(); let config = Config::default_with_file_store_without_models();
let file_store_config = if let config::ValidMemoryBackend::FileStore(file_store_config) = let file_store_config = if let config::ValidMemoryBackend::FileStore(file_store_config) =
config.config.memory.clone() config.config.memory.clone()

View File

@@ -11,7 +11,7 @@ mod postgresml;
mod vector_store; mod vector_store;
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub enum PromptType { pub(crate) enum PromptType {
ContextAndCode, ContextAndCode,
FIM, FIM,
} }
@@ -33,20 +33,20 @@ impl From<&Value> for MemoryRunParams {
} }
#[derive(Debug)] #[derive(Debug)]
pub struct ContextAndCodePrompt { pub(crate) struct ContextAndCodePrompt {
pub context: String, pub(crate) context: String,
pub code: String, pub(crate) code: String,
pub selected_text: Option<String>, pub(crate) selected_text: Option<String>,
} }
#[derive(Debug)] #[derive(Debug)]
pub struct FIMPrompt { pub(crate) struct FIMPrompt {
pub prompt: String, pub(crate) prompt: String,
pub suffix: String, pub(crate) suffix: String,
} }
#[derive(Debug)] #[derive(Debug)]
pub enum Prompt { pub(crate) enum Prompt {
FIM(FIMPrompt), FIM(FIMPrompt),
ContextAndCode(ContextAndCodePrompt), ContextAndCode(ContextAndCodePrompt),
} }
@@ -96,10 +96,7 @@ impl<'a> TryFrom<&'a Prompt> for &'a FIMPrompt {
} }
#[async_trait::async_trait] #[async_trait::async_trait]
pub trait MemoryBackend { pub(crate) trait MemoryBackend {
async fn init(&self) -> anyhow::Result<()> {
Ok(())
}
fn opened_text_document(&self, params: DidOpenTextDocumentParams) -> anyhow::Result<()>; fn opened_text_document(&self, params: DidOpenTextDocumentParams) -> anyhow::Result<()>;
fn code_action_request( fn code_action_request(
&self, &self,
@@ -144,7 +141,7 @@ impl TryFrom<Config> for Box<dyn MemoryBackend + Send + Sync> {
// easier to just pass in a default prompt. // easier to just pass in a default prompt.
#[cfg(test)] #[cfg(test)]
impl Prompt { impl Prompt {
pub fn default_with_cursor() -> Self { pub(crate) fn default_with_cursor() -> Self {
Self::ContextAndCode(ContextAndCodePrompt { Self::ContextAndCode(ContextAndCodePrompt {
context: r#"def test_context():\n pass"#.to_string(), context: r#"def test_context():\n pass"#.to_string(),
code: r#"def test_code():\n <CURSOR>"#.to_string(), code: r#"def test_code():\n <CURSOR>"#.to_string(),
@@ -152,14 +149,14 @@ impl Prompt {
}) })
} }
pub fn default_fim() -> Self { pub(crate) fn default_fim() -> Self {
Self::FIM(FIMPrompt { Self::FIM(FIMPrompt {
prompt: r#"def test_context():\n pass"#.to_string(), prompt: r#"def test_context():\n pass"#.to_string(),
suffix: r#"def test_code():\n "#.to_string(), suffix: r#"def test_code():\n "#.to_string(),
}) })
} }
pub fn default_without_cursor() -> Self { pub(crate) fn default_without_cursor() -> Self {
Self::ContextAndCode(ContextAndCodePrompt { Self::ContextAndCode(ContextAndCodePrompt {
context: r#"def test_context():\n pass"#.to_string(), context: r#"def test_context():\n pass"#.to_string(),
code: r#"def test_code():\n "#.to_string(), code: r#"def test_code():\n "#.to_string(),

View File

@@ -79,7 +79,7 @@ pub(crate) struct PostgresML {
impl PostgresML { impl PostgresML {
#[instrument] #[instrument]
pub fn new( pub(crate) fn new(
mut postgresml_config: config::PostgresML, mut postgresml_config: config::PostgresML,
configuration: Config, configuration: Config,
) -> anyhow::Result<Self> { ) -> anyhow::Result<Self> {

View File

@@ -329,7 +329,7 @@ impl VectorStoreInner {
} }
} }
pub struct VectorStore { pub(crate) struct VectorStore {
file_store: Arc<FileStore>, file_store: Arc<FileStore>,
crawl: Option<Arc<Mutex<Crawl>>>, crawl: Option<Arc<Mutex<Crawl>>>,
splitter: Arc<Box<dyn Splitter + Send + Sync>>, splitter: Arc<Box<dyn Splitter + Send + Sync>>,
@@ -340,7 +340,7 @@ pub struct VectorStore {
} }
impl VectorStore { impl VectorStore {
pub fn new( pub(crate) fn new(
mut vector_store_config: config::VectorStore, mut vector_store_config: config::VectorStore,
config: Config, config: Config,
) -> anyhow::Result<Self> { ) -> anyhow::Result<Self> {
@@ -703,7 +703,7 @@ impl MemoryBackend for VectorStore {
// Get the embedding // Get the embedding
let embedding = self let embedding = self
.embedding_model .embedding_model
.embed(vec![&query], EmbeddingPurpose::Storage) .embed(vec![&query], EmbeddingPurpose::Retrieval)
.await? .await?
.into_iter() .into_iter()
.nth(0) .nth(0)

View File

@@ -6,13 +6,13 @@ mod text_splitter;
mod tree_sitter; mod tree_sitter;
#[derive(Debug, Serialize)] #[derive(Debug, Serialize)]
pub struct ByteRange { pub(crate) struct ByteRange {
pub start_byte: usize, pub(crate) start_byte: usize,
pub end_byte: usize, pub(crate) end_byte: usize,
} }
impl ByteRange { impl ByteRange {
pub fn new(start_byte: usize, end_byte: usize) -> Self { pub(crate) fn new(start_byte: usize, end_byte: usize) -> Self {
Self { Self {
start_byte, start_byte,
end_byte, end_byte,
@@ -21,9 +21,9 @@ impl ByteRange {
} }
#[derive(Serialize)] #[derive(Serialize)]
pub struct Chunk { pub(crate) struct Chunk {
pub text: String, pub(crate) text: String,
pub range: ByteRange, pub(crate) range: ByteRange,
} }
impl Chunk { impl Chunk {
@@ -32,7 +32,7 @@ impl Chunk {
} }
} }
pub trait Splitter { pub(crate) trait Splitter {
fn split(&self, file: &File) -> Vec<Chunk>; fn split(&self, file: &File) -> Vec<Chunk>;
fn split_file_contents(&self, uri: &str, contents: &str) -> Vec<Chunk>; fn split_file_contents(&self, uri: &str, contents: &str) -> Vec<Chunk>;

View File

@@ -11,7 +11,7 @@ fn template_name_from_template_string(template: &str) -> String {
xxhash_rust::xxh3::xxh3_64(template.as_bytes()).to_string() xxhash_rust::xxh3::xxh3_64(template.as_bytes()).to_string()
} }
pub fn apply_chat_template( pub(crate) fn apply_chat_template(
template: &str, template: &str,
chat_messages: Vec<ChatMessage>, chat_messages: Vec<ChatMessage>,
bos_token: &str, bos_token: &str,

View File

@@ -57,7 +57,7 @@ struct AnthropicChatMessage {
} }
#[derive(Deserialize, Serialize)] #[derive(Deserialize, Serialize)]
pub struct ChatError { pub(crate) struct ChatError {
error: Value, error: Value,
} }

View File

@@ -23,23 +23,23 @@ const fn max_new_tokens_default() -> usize {
// NOTE: We cannot deny unknown fields as the provided parameters may contain other fields relevant to other processes // NOTE: We cannot deny unknown fields as the provided parameters may contain other fields relevant to other processes
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
pub struct LLaMACPPRunParams { pub(crate) struct LLaMACPPRunParams {
pub fim: Option<FIM>, pub(crate) fim: Option<FIM>,
messages: Option<Vec<ChatMessage>>, messages: Option<Vec<ChatMessage>>,
chat_template: Option<String>, // A Jinja template chat_template: Option<String>, // A Jinja template
chat_format: Option<String>, // The name of a template in llamacpp chat_format: Option<String>, // The name of a template in llamacpp
#[serde(default = "max_new_tokens_default")] #[serde(default = "max_new_tokens_default")]
pub max_tokens: usize, pub(crate) max_tokens: usize,
// TODO: Explore other arguments // TODO: Explore other arguments
} }
pub struct LLaMACPP { pub(crate) struct LLaMACPP {
model: Model, model: Model,
} }
impl LLaMACPP { impl LLaMACPP {
#[instrument] #[instrument]
pub fn new(configuration: config::LLaMACPP) -> anyhow::Result<Self> { pub(crate) fn new(configuration: config::LLaMACPP) -> anyhow::Result<Self> {
let model_path = match ( let model_path = match (
&configuration.file_path, &configuration.file_path,
&configuration.repository, &configuration.repository,

View File

@@ -17,14 +17,14 @@ use super::LLaMACPPRunParams;
static BACKEND: Lazy<LlamaBackend> = Lazy::new(|| LlamaBackend::init().unwrap()); static BACKEND: Lazy<LlamaBackend> = Lazy::new(|| LlamaBackend::init().unwrap());
pub struct Model { pub(crate) struct Model {
model: LlamaModel, model: LlamaModel,
n_ctx: NonZeroU32, n_ctx: NonZeroU32,
} }
impl Model { impl Model {
#[instrument] #[instrument]
pub fn new(model_path: PathBuf, config: &config::LLaMACPP) -> anyhow::Result<Self> { pub(crate) fn new(model_path: PathBuf, config: &config::LLaMACPP) -> anyhow::Result<Self> {
// Initialize the model_params // Initialize the model_params
let model_params = LlamaModelParams::default().with_n_gpu_layers(config.n_gpu_layers); let model_params = LlamaModelParams::default().with_n_gpu_layers(config.n_gpu_layers);
@@ -42,7 +42,7 @@ impl Model {
} }
#[instrument(skip(self))] #[instrument(skip(self))]
pub fn complete(&self, prompt: &str, params: LLaMACPPRunParams) -> anyhow::Result<String> { pub(crate) fn complete(&self, prompt: &str, params: LLaMACPPRunParams) -> anyhow::Result<String> {
info!("Completing with llama.cpp with prompt:\n{prompt}"); info!("Completing with llama.cpp with prompt:\n{prompt}");
// initialize the context // initialize the context
@@ -129,7 +129,7 @@ impl Model {
} }
#[instrument(skip(self))] #[instrument(skip(self))]
pub fn apply_chat_template( pub(crate) fn apply_chat_template(
&self, &self,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
template: Option<String>, template: Option<String>,
@@ -144,13 +144,13 @@ impl Model {
} }
#[instrument(skip(self))] #[instrument(skip(self))]
pub fn get_eos_token(&self) -> anyhow::Result<String> { pub(crate) fn get_eos_token(&self) -> anyhow::Result<String> {
let token = self.model.token_eos(); let token = self.model.token_eos();
Ok(self.model.token_to_str(token, Special::Tokenize)?) Ok(self.model.token_to_str(token, Special::Tokenize)?)
} }
#[instrument(skip(self))] #[instrument(skip(self))]
pub fn get_bos_token(&self) -> anyhow::Result<String> { pub(crate) fn get_bos_token(&self) -> anyhow::Result<String> {
let token = self.model.token_bos(); let token = self.model.token_bos();
Ok(self.model.token_to_str(token, Special::Tokenize)?) Ok(self.model.token_to_str(token, Special::Tokenize)?)
} }

View File

@@ -19,7 +19,7 @@ mod ollama;
mod open_ai; mod open_ai;
#[async_trait::async_trait] #[async_trait::async_trait]
pub trait TransformerBackend { pub(crate) trait TransformerBackend {
async fn do_completion( async fn do_completion(
&self, &self,
prompt: &Prompt, prompt: &Prompt,
@@ -38,6 +38,7 @@ pub trait TransformerBackend {
params: Value, params: Value,
) -> anyhow::Result<DoGenerationResponse>; ) -> anyhow::Result<DoGenerationResponse>;
#[allow(dead_code)]
async fn do_generate_stream( async fn do_generate_stream(
&self, &self,
request: &GenerationStreamRequest, request: &GenerationStreamRequest,

View File

@@ -69,7 +69,7 @@ enum OllamaChatResponse {
impl Ollama { impl Ollama {
#[instrument] #[instrument]
pub fn new(configuration: config::Ollama) -> Self { pub(crate) fn new(configuration: config::Ollama) -> Self {
Self { configuration } Self { configuration }
} }

View File

@@ -106,7 +106,7 @@ pub(crate) enum OpenAIChatResponse {
impl OpenAI { impl OpenAI {
#[instrument] #[instrument]
pub fn new(configuration: config::OpenAI) -> Self { pub(crate) fn new(configuration: config::OpenAI) -> Self {
Self { configuration } Self { configuration }
} }

View File

@@ -54,37 +54,37 @@ impl GenerationRequest {
// The generate stream is not yet ready but we don't want to remove it // The generate stream is not yet ready but we don't want to remove it
#[allow(dead_code)] #[allow(dead_code)]
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct GenerationStreamRequest { pub(crate) struct GenerationStreamRequest {
id: RequestId, id: RequestId,
params: GenerationStreamParams, params: GenerationStreamParams,
} }
impl GenerationStreamRequest { impl GenerationStreamRequest {
pub fn new(id: RequestId, params: GenerationStreamParams) -> Self { pub(crate) fn new(id: RequestId, params: GenerationStreamParams) -> Self {
Self { id, params } Self { id, params }
} }
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct CodeActionRequest { pub(crate) struct CodeActionRequest {
id: RequestId, id: RequestId,
params: CodeActionParams, params: CodeActionParams,
} }
impl CodeActionRequest { impl CodeActionRequest {
pub fn new(id: RequestId, params: CodeActionParams) -> Self { pub(crate) fn new(id: RequestId, params: CodeActionParams) -> Self {
Self { id, params } Self { id, params }
} }
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct CodeActionResolveRequest { pub(crate) struct CodeActionResolveRequest {
id: RequestId, id: RequestId,
params: CodeAction, params: CodeAction,
} }
impl CodeActionResolveRequest { impl CodeActionResolveRequest {
pub fn new(id: RequestId, params: CodeAction) -> Self { pub(crate) fn new(id: RequestId, params: CodeAction) -> Self {
Self { id, params } Self { id, params }
} }
} }
@@ -112,16 +112,17 @@ impl WorkerRequest {
} }
} }
pub struct DoCompletionResponse { pub(crate) struct DoCompletionResponse {
pub insert_text: String, pub(crate) insert_text: String,
} }
pub struct DoGenerationResponse { pub(crate) struct DoGenerationResponse {
pub generated_text: String, pub(crate) generated_text: String,
} }
pub struct DoGenerationStreamResponse { #[allow(dead_code)]
pub generated_text: String, pub(crate) struct DoGenerationStreamResponse {
pub(crate) generated_text: String,
} }
fn post_process_start(response: String, front: &str) -> String { fn post_process_start(response: String, front: &str) -> String {

View File

@@ -77,7 +77,7 @@ pub(crate) fn parse_tree(
.with_context(|| format!("parsing tree failed for {uri}")) .with_context(|| format!("parsing tree failed for {uri}"))
} }
pub fn format_file_chunk(uri: &str, excerpt: &str, root_uri: Option<&str>) -> String { pub(crate) fn format_file_chunk(uri: &str, excerpt: &str, root_uri: Option<&str>) -> String {
let path = match root_uri { let path = match root_uri {
Some(root_uri) => { Some(root_uri) => {
if uri.starts_with(root_uri) { if uri.starts_with(root_uri) {