Use chroma for temp vectordb

Author: ajhai

llmstack/common/blocks/data/store/vectorstore/chroma.py

@@ -2,6 +2,7 @@
 from typing import Any, Tuple
 from typing import List
 from uuid import uuid4
+import uuid
 
 import chromadb
 from pydantic import BaseModel
@@ -13,6 +14,7 @@
 class ChromaConfiguration(BaseModel):
     _type = 'Chroma'
     anonymized_telemetry = False
+    is_persistent = True
 
 
 class Chroma(VectorStoreInterface):
@@ -23,8 +25,12 @@ class Chroma(VectorStoreInterface):
     def __init__(self, *args, **kwargs) -> None:
         configuration = ChromaConfiguration(**kwargs)
         db_settings = chromadb.config.Settings(**configuration.dict())
-        self._client = chromadb.PersistentClient(
-            path=settings.DEFAULT_VECTOR_DATABASE_PATH, settings=db_settings) if settings.DEFAULT_VECTOR_DATABASE_PATH else chromadb.Client(settings=db_settings)
+
+        if db_settings.is_persistent:
+            self._client = chromadb.PersistentClient(
+                path=settings.DEFAULT_VECTOR_DATABASE_PATH, settings=db_settings) if settings.DEFAULT_VECTOR_DATABASE_PATH else chromadb.Client(settings=db_settings)
+        else:
+            self._client = chromadb.EphemeralClient(settings=db_settings)
 
     def add_text(self, index_name: str, document: Document, **kwargs: Any):
         content_key = document.page_content_key
@@ -101,3 +107,9 @@ def similarity_search(self, index_name: str, document_query: DocumentQuery, **kw
             )
 
         return result
+
+    def create_temp_index(self):
+        index_name = 'Temp_{}'.format(str(uuid.uuid4())).replace('-', '_')
+        self.create_index(schema='', index_name=index_name)
+
+        return index_name

llmstack/processors/providers/promptly/data_uri_text_extract.py

@@ -8,8 +8,8 @@
 from pydantic import conint
 from pydantic import Field
 
-from llmstack.common.blocks.data.store.vectorstore import Document
-from llmstack.common.blocks.data.store.vectorstore.temp_weaviate import TempWeaviate
+from llmstack.common.blocks.data.store.vectorstore import Document, DocumentQuery
+from llmstack.common.blocks.data.store.vectorstore.chroma import Chroma
 from llmstack.common.utils.text_extract import extract_text_from_b64_json, ExtraParams
 from llmstack.common.utils.splitter import SpacyTextSplitter
 from llmstack.common.utils.utils import validate_parse_data_uri
@@ -93,22 +93,8 @@ def session_data_to_persist(self) -> dict:
 
     def process(self) -> str:
         openai_api_key = self._env.get('openai_api_key', None)
-        weaviate_url = self._env['weaviate_url']
-        weaviate_api_key = self._env.get('weaviate_api_key', None)
-        azure_openai_api_key = self._env.get('azure_openai_api_key', None)
-        weaviate_embedding_endpoint = self._env['weaviate_embedding_endpoint']
-        weaviate_text2vec_config = self._env['weaviate_text2vec_config']
-
         query = self._input.query
-
-        self.temp_store = TempWeaviate(
-            url=weaviate_url,
-            openai_key=openai_api_key,
-            azure_openai_key=azure_openai_api_key,
-            weaviate_rw_api_key=weaviate_api_key,
-            weaviate_embedding_endpoint=weaviate_embedding_endpoint,
-            weaviate_text2vec_config=weaviate_text2vec_config,
-        )
+        self.temp_store = Chroma(is_persistent=False)
 
         file = self._input.file or None
         if (file is None or file == '') and self._input.file_data:
@@ -157,13 +143,15 @@ def process(self) -> str:
                 ).split_text(text)
                 futures = [
                     executor.submit(
-                        self.temp_store.add_content,
-                        index_name, text_chunk, source=file_name,
+                        self.temp_store.add_text,
+                        index_name, Document(page_content_key="content", page_content=text_chunk, metadata={
+                            'source': file_name}),
                     ) for text_chunk in text_chunks
                 ]
                 concurrent.futures.wait(futures)
-            documents: List[Document] = self.temp_store.search_temp_index(
-                self.storage_index_name, query, self._config.document_limit,
+            documents: List[Document] = self.temp_store.hybrid_search(
+                self.storage_index_name, document_query=DocumentQuery(
+                    query=query, limit=self._config.document_limit),
             )
 
             async_to_sync(self._output_stream.write)(

llmstack/processors/providers/promptly/http_uri_text_extract.py

@@ -5,10 +5,9 @@
 from asgiref.sync import async_to_sync
 from pydantic import conint
 from pydantic import Field
-from pydantic import HttpUrl
 
-from llmstack.common.blocks.data.store.vectorstore import Document
-from llmstack.common.blocks.data.store.vectorstore.temp_weaviate import TempWeaviate
+from llmstack.common.blocks.data.store.vectorstore import Document, DocumentQuery
+from llmstack.common.blocks.data.store.vectorstore.chroma import Chroma
 from llmstack.common.utils.text_extract import extract_text_from_url, ExtraParams
 from llmstack.common.utils.splitter import SpacyTextSplitter
 from llmstack.processors.providers.api_processor_interface import ApiProcessorInterface, ApiProcessorSchema
@@ -74,23 +73,10 @@ def session_data_to_persist(self) -> dict:
 
     def process(self) -> HttpUriTextExtractorOutput:
         openai_api_key = self._env.get('openai_api_key', None)
-        weaviate_url = self._env['weaviate_url']
-        weaviate_api_key = self._env.get('weaviate_api_key', None)
-        azure_openai_api_key = self._env.get('azure_openai_api_key', None)
-        weaviate_embedding_endpoint = self._env['weaviate_embedding_endpoint']
-        weaviate_text2vec_config = self._env['weaviate_text2vec_config']
 
         query = self._input.query
         url = self._input.url.strip().rstrip()
-
-        self.temp_store = TempWeaviate(
-            url=weaviate_url,
-            openai_key=openai_api_key,
-            weaviate_rw_api_key=weaviate_api_key,
-            azure_openai_key=azure_openai_api_key,
-            weaviate_embedding_endpoint=weaviate_embedding_endpoint,
-            weaviate_text2vec_config=weaviate_text2vec_config,
-        )
+        self.temp_store = Chroma(is_persistent=False)
 
         if (query is None or query == '') and url == self.url and self.extracted_text is not None:
             async_to_sync(self._output_stream.write)(
@@ -100,8 +86,9 @@ def process(self) -> HttpUriTextExtractorOutput:
             return output
 
         if query and self.storage_index_name and url == self.url:
-            documents: List[Document] = self.temp_store.search_temp_index(
-                self.storage_index_name, query, self._config.document_limit,
+            documents: List[Document] = self.temp_store.hybrid_search(
+                self.storage_index_name, document_query=DocumentQuery(
+                    query=query, limit=self._config.document_limit),
             )
             for document in documents:
                 async_to_sync(self._output_stream.write)(
@@ -124,11 +111,13 @@ def process(self) -> HttpUriTextExtractorOutput:
             index_name = self.temp_store.create_temp_index()
             self.storage_index_name = index_name
             for text_chunk in SpacyTextSplitter(separator='\n', chunk_size=self._config.text_chunk_size).split_text(text):
-                self.temp_store.add_content(
-                    index_name, text_chunk, source=self.url,
+                self.temp_store.add_text(
+                    index_name, Document(page_content_key="content", page_content=text_chunk, metadata={
+                                         'source': self.url}),
                 )
-            documents: List[Document] = self.temp_store.search_temp_index(
-                self.storage_index_name, query, self._config.document_limit,
+            documents: List[Document] = self.temp_store.hybrid_search(
+                self.storage_index_name, document_query=DocumentQuery(
+                    query=query, limit=self._config.document_limit),
             )
 
             for document in documents: