hz/ebiz-dify-ai
1
0
Fork 0
mirror of http://112.124.100.131/huang.ze/ebiz-dify-ai.git synced 2025-12-09 10:56:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add WEAVIATE_BATCH_SIZE (#349)

Browse Source
This commit is contained in:
John Wang
2023-06-13 14:49:40 +08:00
committed by GitHub
parent 6a3ab36101
commit cd136fb293
4 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
api/core/vector_store/vector_store.py
View File

@@ -27,7 +27,8 @@ class VectorStore:
self._client = WeaviateVectorStoreClient(
endpoint=app.config['WEAVIATE_ENDPOINT'],
api_key=app.config['WEAVIATE_API_KEY'],
grpc_enabled=app.config['WEAVIATE_GRPC_ENABLED']
grpc_enabled=app.config['WEAVIATE_GRPC_ENABLED'],
batch_size=app.config['WEAVIATE_BATCH_SIZE']
)
elif self._vector_store == 'qdrant':
self._client = QdrantVectorStoreClient(

8
api/core/vector_store/weaviate_vector_store_client.py
View File

@@ -18,10 +18,10 @@ from llama_index.readers.weaviate.utils import (
class WeaviateVectorStoreClient(BaseVectorStoreClient):
def __init__(self, endpoint: str, api_key: str, grpc_enabled: bool):
self._client = self.init_from_config(endpoint, api_key, grpc_enabled)
def __init__(self, endpoint: str, api_key: str, grpc_enabled: bool, batch_size: int):
self._client = self.init_from_config(endpoint, api_key, grpc_enabled, batch_size)
def init_from_config(self, endpoint: str, api_key: str, grpc_enabled: bool):
def init_from_config(self, endpoint: str, api_key: str, grpc_enabled: bool, batch_size: int):
auth_config = weaviate.auth.AuthApiKey(api_key=api_key)
weaviate.connect.connection.has_grpc = grpc_enabled
@@ -36,7 +36,7 @@ class WeaviateVectorStoreClient(BaseVectorStoreClient):
client.batch.configure(
# `batch_size` takes an `int` value to enable auto-batching
# (`None` is used for manual batching)
batch_size=100,
batch_size=batch_size,
# dynamically update the `batch_size` based on import speed
dynamic=True,
# `timeout_retries` takes an `int` value to retry on time outs