github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-08-14 14:34:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: auto register for vectorstores

Browse Source
This commit is contained in:
csunny 2023-10-08 18:47:49 +08:00
parent c28f18e792
commit 232661a4ac
1 changed files with 5 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
pilot/vector_store/connector.py
View File

@ -1,8 +1,7 @@
from pilot.vector_store.chroma_store import ChromaStore
from pilot import vector_store from pilot import vector_store
from pilot.vector_store.base import VectorStoreBase from pilot.vector_store.base import VectorStoreBase
connector = {"Chroma": ChromaStore} connector = {}
class VectorStoreConnector: class VectorStoreConnector:
"""VectorStoreConnector, can connect different vector db provided load document api_v1 and similar search api_v1. """VectorStoreConnector, can connect different vector db provided load document api_v1 and similar search api_v1.
@ -22,7 +21,8 @@ class VectorStoreConnector:
self.connector_class = connector.get(vector_store_type) self.connector_class = connector.get(vector_store_type)
else: else:
raise Exception(f"Vector Type Not support. {0}", vector_store_type) raise Exception(f"Vector Type Not support. {0}", vector_store_type)
print(self.connector_class)
self.client = self.connector_class(ctx) self.client = self.connector_class(ctx)
@ -55,4 +55,5 @@ class VectorStoreConnector:
def _register(self): def _register(self):
for cls in vector_store.__all__: for cls in vector_store.__all__:
if issubclass(getattr(vector_store, cls), VectorStoreBase): if issubclass(getattr(vector_store, cls), VectorStoreBase):
connector.update({cls, getattr(vector_store, cls)}) _k, _v = cls, getattr(vector_store, cls)
connector.update({_k: _v})