updated vector_store db to leverage EMBEDDINGS_DIR path

This commit is contained in:
KCaverly 2023-06-29 15:18:32 -04:00
parent 0a7245a583
commit 39137fc19f
4 changed files with 26 additions and 17 deletions

View file

@ -6,6 +6,7 @@ lazy_static::lazy_static! {
pub static ref HOME: PathBuf = dirs::home_dir().expect("failed to determine home directory"); pub static ref HOME: PathBuf = dirs::home_dir().expect("failed to determine home directory");
pub static ref CONFIG_DIR: PathBuf = HOME.join(".config").join("zed"); pub static ref CONFIG_DIR: PathBuf = HOME.join(".config").join("zed");
pub static ref CONVERSATIONS_DIR: PathBuf = HOME.join(".config/zed/conversations"); pub static ref CONVERSATIONS_DIR: PathBuf = HOME.join(".config/zed/conversations");
pub static ref EMBEDDINGS_DIR: PathBuf = HOME.join(".config/zed/embeddings");
pub static ref LOGS_DIR: PathBuf = HOME.join("Library/Logs/Zed"); pub static ref LOGS_DIR: PathBuf = HOME.join("Library/Logs/Zed");
pub static ref SUPPORT_DIR: PathBuf = HOME.join("Library/Application Support/Zed"); pub static ref SUPPORT_DIR: PathBuf = HOME.join("Library/Application Support/Zed");
pub static ref LANGUAGES_DIR: PathBuf = HOME.join("Library/Application Support/Zed/languages"); pub static ref LANGUAGES_DIR: PathBuf = HOME.join("Library/Application Support/Zed/languages");

View file

@ -15,11 +15,6 @@ use sha1::{Digest, Sha1};
use crate::IndexedFile; use crate::IndexedFile;
// This is saving to a local database store within the users dev zed path
// Where do we want this to sit?
// Assuming near where the workspace DB sits.
pub const VECTOR_DB_URL: &str = "embeddings_db";
// Note this is not an appropriate document // Note this is not an appropriate document
#[derive(Debug)] #[derive(Debug)]
pub struct DocumentRecord { pub struct DocumentRecord {
@ -109,7 +104,7 @@ pub struct VectorDatabase {
} }
impl VectorDatabase { impl VectorDatabase {
pub fn new(path: &str) -> Result<Self> { pub fn new(path: String) -> Result<Self> {
let this = Self { let this = Self {
db: rusqlite::Connection::open(path)?, db: rusqlite::Connection::open(path)?,
}; };

View file

@ -6,16 +6,23 @@ mod modal;
mod vector_store_tests; mod vector_store_tests;
use anyhow::{anyhow, Result}; use anyhow::{anyhow, Result};
use db::{FileSha1, VectorDatabase, VECTOR_DB_URL}; use db::{FileSha1, VectorDatabase};
use embedding::{EmbeddingProvider, OpenAIEmbeddings}; use embedding::{EmbeddingProvider, OpenAIEmbeddings};
use gpui::{AppContext, Entity, ModelContext, ModelHandle, Task, ViewContext}; use gpui::{AppContext, Entity, ModelContext, ModelHandle, Task, ViewContext};
use language::{Language, LanguageRegistry}; use language::{Language, LanguageRegistry};
use modal::{SemanticSearch, SemanticSearchDelegate, Toggle}; use modal::{SemanticSearch, SemanticSearchDelegate, Toggle};
use project::{Fs, Project, WorktreeId}; use project::{Fs, Project, WorktreeId};
use smol::channel; use smol::channel;
use std::{cmp::Ordering, collections::HashMap, path::PathBuf, sync::Arc}; use std::{
cmp::Ordering,
collections::HashMap,
path::{Path, PathBuf},
sync::Arc,
};
use tree_sitter::{Parser, QueryCursor}; use tree_sitter::{Parser, QueryCursor};
use util::{http::HttpClient, ResultExt, TryFutureExt}; use util::{
channel::RELEASE_CHANNEL_NAME, http::HttpClient, paths::EMBEDDINGS_DIR, ResultExt, TryFutureExt,
};
use workspace::{Workspace, WorkspaceCreated}; use workspace::{Workspace, WorkspaceCreated};
#[derive(Debug)] #[derive(Debug)]
@ -31,11 +38,14 @@ pub fn init(
language_registry: Arc<LanguageRegistry>, language_registry: Arc<LanguageRegistry>,
cx: &mut AppContext, cx: &mut AppContext,
) { ) {
let db_file_path = EMBEDDINGS_DIR
.join(Path::new(RELEASE_CHANNEL_NAME.as_str()))
.join("embeddings_db");
let vector_store = cx.add_model(|_| { let vector_store = cx.add_model(|_| {
VectorStore::new( VectorStore::new(
fs, fs,
VECTOR_DB_URL.to_string(), db_file_path,
// Arc::new(DummyEmbeddings {}),
Arc::new(OpenAIEmbeddings { Arc::new(OpenAIEmbeddings {
client: http_client, client: http_client,
}), }),
@ -87,7 +97,7 @@ pub struct IndexedFile {
pub struct VectorStore { pub struct VectorStore {
fs: Arc<dyn Fs>, fs: Arc<dyn Fs>,
database_url: Arc<str>, database_url: Arc<PathBuf>,
embedding_provider: Arc<dyn EmbeddingProvider>, embedding_provider: Arc<dyn EmbeddingProvider>,
language_registry: Arc<LanguageRegistry>, language_registry: Arc<LanguageRegistry>,
worktree_db_ids: Vec<(WorktreeId, i64)>, worktree_db_ids: Vec<(WorktreeId, i64)>,
@ -104,13 +114,13 @@ pub struct SearchResult {
impl VectorStore { impl VectorStore {
fn new( fn new(
fs: Arc<dyn Fs>, fs: Arc<dyn Fs>,
database_url: String, database_url: PathBuf,
embedding_provider: Arc<dyn EmbeddingProvider>, embedding_provider: Arc<dyn EmbeddingProvider>,
language_registry: Arc<LanguageRegistry>, language_registry: Arc<LanguageRegistry>,
) -> Self { ) -> Self {
Self { Self {
fs, fs,
database_url: database_url.into(), database_url: Arc::new(database_url),
embedding_provider, embedding_provider,
language_registry, language_registry,
worktree_db_ids: Vec::new(), worktree_db_ids: Vec::new(),
@ -209,7 +219,10 @@ impl VectorStore {
.timer(std::time::Duration::from_secs(3)) .timer(std::time::Duration::from_secs(3))
.await; .await;
let db = VectorDatabase::new(&database_url)?; if let Some(db_directory) = database_url.parent() {
fs.create_dir(db_directory).await.log_err();
}
let db = VectorDatabase::new(database_url.to_string_lossy().into())?;
let worktrees = project.read_with(&cx, |project, cx| { let worktrees = project.read_with(&cx, |project, cx| {
project project
@ -372,7 +385,7 @@ impl VectorStore {
let documents = cx let documents = cx
.background() .background()
.spawn(async move { .spawn(async move {
let database = VectorDatabase::new(database_url.as_ref())?; let database = VectorDatabase::new(database_url.to_string_lossy().into())?;
let phrase_embedding = embedding_provider let phrase_embedding = embedding_provider
.embed_batch(vec![&phrase]) .embed_batch(vec![&phrase])

View file

@ -63,7 +63,7 @@ async fn test_vector_store(cx: &mut TestAppContext) {
let store = cx.add_model(|_| { let store = cx.add_model(|_| {
VectorStore::new( VectorStore::new(
fs.clone(), fs.clone(),
db_path.to_string_lossy().to_string(), db_path,
Arc::new(FakeEmbeddingProvider), Arc::new(FakeEmbeddingProvider),
languages, languages,
) )