refactor(rag): switch to dify_config. (#6410)

Co-authored-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
Poorandy 2024-07-18 18:40:36 +08:00 • committed by GitHub
commit c8f5dfcf17
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 121 additions and 131 deletions

View file

@ -3,9 +3,9 @@ from typing import Any, Optional
import chromadb
from chromadb import QueryResult, Settings
from flask import current_app
from pydantic import BaseModel
from configs import dify_config
from core.rag.datasource.entity.embedding import Embeddings
from core.rag.datasource.vdb.vector_base import BaseVector
from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
@ -133,15 +133,14 @@ class ChromaVectorFactory(AbstractVectorFactory):
}
dataset.index_struct = json.dumps(index_struct_dict)
config = current_app.config
return ChromaVector(
collection_name=collection_name,
config=ChromaConfig(
host=config.get('CHROMA_HOST'),
port=int(config.get('CHROMA_PORT')),
tenant=config.get('CHROMA_TENANT', chromadb.DEFAULT_TENANT),
database=config.get('CHROMA_DATABASE', chromadb.DEFAULT_DATABASE),
auth_provider=config.get('CHROMA_AUTH_PROVIDER'),
auth_credentials=config.get('CHROMA_AUTH_CREDENTIALS'),
host=dify_config.CHROMA_HOST,
port=dify_config.CHROMA_PORT,
tenant=dify_config.CHROMA_TENANT or chromadb.DEFAULT_TENANT,
database=dify_config.CHROMA_DATABASE or chromadb.DEFAULT_DATABASE,
auth_provider=dify_config.CHROMA_AUTH_PROVIDER,
auth_credentials=dify_config.CHROMA_AUTH_CREDENTIALS,
),
)