fix:multi client error (#1532)

This commit is contained in:
Aries-ckt 2024-05-16 23:23:24 +08:00 committed by GitHub
parent af700d1dc0
commit c500fb9060
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -3,6 +3,7 @@
import copy import copy
import logging import logging
import os import os
from collections import defaultdict
from typing import Any, Dict, List, Optional, Tuple, Type, cast from typing import Any, Dict, List, Optional, Tuple, Type, cast
from dbgpt.core import Chunk, Embeddings from dbgpt.core import Chunk, Embeddings
@ -21,6 +22,7 @@ from dbgpt.util.i18n_utils import _
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
connector: Dict[str, Tuple[Type, Type]] = {} connector: Dict[str, Tuple[Type, Type]] = {}
pools = defaultdict(dict)
def _load_vector_options() -> List[OptionValue]: def _load_vector_options() -> List[OptionValue]:
@ -115,8 +117,15 @@ class VectorStoreConnector:
config_dict = vector_store_config.dict() config_dict = vector_store_config.dict()
config.llm_client = config_dict.get("llm_client", None) config.llm_client = config_dict.get("llm_client", None)
config.model_name = config_dict.get("model_name", None) config.model_name = config_dict.get("model_name", None)
if (
self.client = self.connector_class(config) vector_store_type in pools
and config.name in pools[vector_store_type]
):
self.client = pools[vector_store_type][config.name]
else:
client = self.connector_class(config)
pools[vector_store_type][config.name] = self.client = client
self.client = client
except Exception as e: except Exception as e:
logger.error("connect vector store failed: %s", e) logger.error("connect vector store failed: %s", e)
raise e raise e