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 CONFIG_DIR: PathBuf = HOME.join(".config").join("zed");
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 SUPPORT_DIR: PathBuf = HOME.join("Library/Application Support/Zed");
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;
// 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
#[derive(Debug)]
pub struct DocumentRecord {
@ -109,7 +104,7 @@ pub struct VectorDatabase {
}
impl VectorDatabase {
pub fn new(path: &str) -> Result<Self> {
pub fn new(path: String) -> Result<Self> {
let this = Self {
db: rusqlite::Connection::open(path)?,
};

View File

@ -6,16 +6,23 @@ mod modal;
mod vector_store_tests;
use anyhow::{anyhow, Result};
use db::{FileSha1, VectorDatabase, VECTOR_DB_URL};
use db::{FileSha1, VectorDatabase};
use embedding::{EmbeddingProvider, OpenAIEmbeddings};
use gpui::{AppContext, Entity, ModelContext, ModelHandle, Task, ViewContext};
use language::{Language, LanguageRegistry};
use modal::{SemanticSearch, SemanticSearchDelegate, Toggle};
use project::{Fs, Project, WorktreeId};
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 util::{http::HttpClient, ResultExt, TryFutureExt};
use util::{
channel::RELEASE_CHANNEL_NAME, http::HttpClient, paths::EMBEDDINGS_DIR, ResultExt, TryFutureExt,
};
use workspace::{Workspace, WorkspaceCreated};
#[derive(Debug)]
@ -31,11 +38,14 @@ pub fn init(
language_registry: Arc<LanguageRegistry>,
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(|_| {
VectorStore::new(
fs,
VECTOR_DB_URL.to_string(),
// Arc::new(DummyEmbeddings {}),
db_file_path,
Arc::new(OpenAIEmbeddings {
client: http_client,
}),
@ -87,7 +97,7 @@ pub struct IndexedFile {
pub struct VectorStore {
fs: Arc<dyn Fs>,
database_url: Arc<str>,
database_url: Arc<PathBuf>,
embedding_provider: Arc<dyn EmbeddingProvider>,
language_registry: Arc<LanguageRegistry>,
worktree_db_ids: Vec<(WorktreeId, i64)>,
@ -104,13 +114,13 @@ pub struct SearchResult {
impl VectorStore {
fn new(
fs: Arc<dyn Fs>,
database_url: String,
database_url: PathBuf,
embedding_provider: Arc<dyn EmbeddingProvider>,
language_registry: Arc<LanguageRegistry>,
) -> Self {
Self {
fs,
database_url: database_url.into(),
database_url: Arc::new(database_url),
embedding_provider,
language_registry,
worktree_db_ids: Vec::new(),
@ -209,7 +219,10 @@ impl VectorStore {
.timer(std::time::Duration::from_secs(3))
.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| {
project
@ -372,7 +385,7 @@ impl VectorStore {
let documents = cx
.background()
.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
.embed_batch(vec![&phrase])

View File

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