feat:Add vector type and vector factory class for vector database integration

2025-05-14 21:30:50 +08:00
parent 4ce1e88750
commit 6f869ded43
7 changed files with 63 additions and 34 deletions
@@ -0,0 +1,48 @@
+from open_webui.retrieval.vector.main import VectorDBBase
+from open_webui.retrieval.vector.type import VectorType
+from open_webui.config import VECTOR_DB
+
+
+class Vector:
+
+    @staticmethod
+    def get_vector(vector_type: str) -> VectorDBBase:
+        """
+        get vector db instance by vector type
+        """
+        match vector_type:
+            case VectorType.MILVUS:
+                from open_webui.retrieval.vector.dbs.milvus import MilvusClient
+
+                return MilvusClient()
+            case VectorType.QDRANT:
+                from open_webui.retrieval.vector.dbs.qdrant import QdrantClient
+
+                return QdrantClient()
+            case VectorType.PINECONE:
+                from open_webui.retrieval.vector.dbs.pinecone import PineconeClient
+
+                return PineconeClient()
+            case VectorType.OPENSEARCH:
+                from open_webui.retrieval.vector.dbs.opensearch import OpenSearchClient
+
+                return OpenSearchClient()
+            case VectorType.PGVECTOR:
+                from open_webui.retrieval.vector.dbs.pgvector import PgvectorClient
+
+                return PgvectorClient()
+            case VectorType.ELASTICSEARCH:
+                from open_webui.retrieval.vector.dbs.elasticsearch import (
+                    ElasticsearchClient,
+                )
+
+                return ElasticsearchClient()
+            case VectorType.CHROMA:
+                from open_webui.retrieval.vector.dbs.chroma import ChromaClient
+
+                return ChromaClient()
+            case _:
+                raise ValueError(f"Unsupported vector type: {vector_type}")
+
+
+VECTOR_DB_CLIENT = Vector.get_vector(VECTOR_DB)