diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 20e1dd2..cedceaa 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -49,7 +49,7 @@ jobs: OPENAI_BASE_URL: ${{ secrets.OPENAI_BASE_URL }} OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} DEFAULT_EMBED_MODEL: text-embedding-ada-002 - DEFAULT_CHAT_MODEL: gpt-3.5-turbo + DEFAULT_CHAT_MODEL: gpt-4o-mini HF_TOKENIZER_PATH: 01-ai/Yi-6B-Chat DEFAULT_CHUNK_SIZE: 300 DEFAULT_CHUNK_OVERLAP: 100 diff --git a/src/cardinal/vectorstore/auto.py b/src/cardinal/vectorstore/auto.py index f907bb1..7131553 100644 --- a/src/cardinal/vectorstore/auto.py +++ b/src/cardinal/vectorstore/auto.py @@ -29,12 +29,10 @@ def create(cls, name: str, texts: Sequence[str], data: Sequence[T], drop_old: Op return _get_vectorstore().create(name, texts, data, drop_old) def insert(self, texts: Sequence[str], data: Sequence[T]) -> None: - self._vectorstore.insert(texts, data) - return self._flush() + return self._vectorstore.insert(texts, data) def delete(self, condition: "Condition") -> None: - self._vectorstore.delete(condition) - return self._flush() + return self._vectorstore.delete(condition) def search( self, query: str, top_k: Optional[int] = 4, condition: Optional["Condition"] = None @@ -47,11 +45,8 @@ def exists(self) -> bool: def destroy(self) -> None: return self._vectorstore.destroy() - def _flush(self) -> None: - if isinstance(self._vectorstore, Milvus): - return self._vectorstore.store.flush() - else: - return None + def flush(self) -> None: + return self._vectorstore.flush() _vectorstores: Dict[str, Type["VectorStore"]] = {} diff --git a/src/cardinal/vectorstore/chroma.py b/src/cardinal/vectorstore/chroma.py index ee6a914..47690da 100644 --- a/src/cardinal/vectorstore/chroma.py +++ b/src/cardinal/vectorstore/chroma.py @@ -145,3 +145,6 @@ def destroy(self) -> None: client = _get_chroma_client() client.delete_collection(self.name) self.store = None + + def flush(self) -> None: + return None diff --git a/src/cardinal/vectorstore/milvus.py b/src/cardinal/vectorstore/milvus.py index 1f5f8c6..f1a22cb 100644 --- a/src/cardinal/vectorstore/milvus.py +++ b/src/cardinal/vectorstore/milvus.py @@ -177,3 +177,7 @@ def destroy(self): self._try_init_and_check_exists() self.store.drop() self.store = None + + def flush(self) -> None: + return self.store.flush() + \ No newline at end of file diff --git a/tests/vectorstore/test_vector_store.py b/tests/vectorstore/test_vector_store.py index 1960379..d63af0e 100644 --- a/tests/vectorstore/test_vector_store.py +++ b/tests/vectorstore/test_vector_store.py @@ -11,12 +11,13 @@ class Animal(BaseModel): def test_vector_store(): - vectorStore = AutoVectorStore[Animal](name="test") - assert(not vectorStore.exists()) # False - vectorStore.insert(texts=texts, data=data) - vectorStore.delete(AutoCondition(key="name", value="dog", op=Operator.Eq)) - assert(vectorStore.search(query="dog", top_k=2)[0][0] == data[2]) - assert(vectorStore.search(query="dog", top_k=2)[1][0] == data[1]) + vectorstore = AutoVectorStore[Animal](name="test") + assert(not vectorstore.exists()) # False + vectorstore.insert(texts=texts, data=data) + vectorstore.delete(AutoCondition(key="name", value="dog", op=Operator.Eq)) + vectorstore.flush() + assert(vectorstore.search(query="dog", top_k=2)[0][0] == data[2]) + assert(vectorstore.search(query="dog", top_k=2)[1][0] == data[1]) # [(Animal(name='puppy'), 0.8510237336158752), (Animal(name='llama'), 1.1970627307891846)] - assert(vectorStore.exists()) # True - vectorStore.destroy() + assert(vectorstore.exists()) # True + vectorstore.destroy()