chore: use in-memory db by default

williamfzc · Nov 23, 2023 · c91688f · c91688f
1 parent 3a093cf
commit c91688f
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/examples/read.py b/examples/read.py
@@ -2,6 +2,7 @@
 from srctag.tagger import Tagger
 
 storage = Storage()
+storage.config.db_path = "./chroma"
 tagger = Tagger()
 tagger.config.tags = ["storage"]
 tag_dict = tagger.tag(storage)

diff --git a/examples/write.py b/examples/write.py
@@ -9,4 +9,5 @@
 
 ctx = collector.collect_metadata()
 storage = Storage()
+storage.config.db_path = "./chroma"
 storage.embed_ctx(ctx)
diff --git a/srctag/storage.py b/srctag/storage.py
@@ -13,7 +13,7 @@
 
 
 class StorageConfig(BaseSettings):
-    db_path: str = "./chroma"
+    db_path: str = ""
     collection_name: str = "default_collection"
 
     # English: paraphrase-MiniLM-L6-v2
@@ -33,7 +33,12 @@ def __init__(self, config: StorageConfig = None):
     def init_chroma(self):
         if self.chromadb and self.chromadb_collection:
             return
-        self.chromadb = chromadb.PersistentClient(path=self.config.db_path)
+
+        if self.config.db_path:
+            self.chromadb = chromadb.PersistentClient(path=self.config.db_path)
+        else:
+            self.chromadb = chromadb.Client()
+
         self.chromadb_collection = self.chromadb.get_or_create_collection(
             self.config.collection_name,
             embedding_function=SentenceTransformerEmbeddingFunction(