opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-08 11:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(api): type Chroma and AnalyticDB config params dicts with TypedDicts (#34678)

Browse Source
This commit is contained in:
Statxc 2026-04-07 10:27:12 -03:00 committed by GitHub
parent e645cbd8f8
commit c44ddd9831
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 24 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
api/core/rag/datasource/vdb/analyticdb/analyticdb_vector_openapi.py
View File

@ -1,5 +1,5 @@
import json import json
from typing import Any from typing import Any, TypedDict
from pydantic import BaseModel, model_validator from pydantic import BaseModel, model_validator
@ -13,6 +13,13 @@ from core.rag.models.document import Document
from extensions.ext_redis import redis_client from extensions.ext_redis import redis_client
class AnalyticdbClientParamsDict(TypedDict):
access_key_id: str
access_key_secret: str
region_id: str
read_timeout: int
class AnalyticdbVectorOpenAPIConfig(BaseModel): class AnalyticdbVectorOpenAPIConfig(BaseModel):
access_key_id: str access_key_id: str
access_key_secret: str access_key_secret: str
@ -44,13 +51,14 @@ class AnalyticdbVectorOpenAPIConfig(BaseModel):
raise ValueError("config ANALYTICDB_NAMESPACE_PASSWORD is required") raise ValueError("config ANALYTICDB_NAMESPACE_PASSWORD is required")
return values return values
def to_analyticdb_client_params(self): def to_analyticdb_client_params(self) -> AnalyticdbClientParamsDict:
return { result: AnalyticdbClientParamsDict = {
"access_key_id": self.access_key_id, "access_key_id": self.access_key_id,
"access_key_secret": self.access_key_secret, "access_key_secret": self.access_key_secret,
"region_id": self.region_id, "region_id": self.region_id,
"read_timeout": self.read_timeout, "read_timeout": self.read_timeout,
} }
return result
class AnalyticdbVectorOpenAPI: class AnalyticdbVectorOpenAPI:

17
api/core/rag/datasource/vdb/chroma/chroma_vector.py
View File

@ -1,5 +1,5 @@
import json import json
from typing import Any from typing import Any, TypedDict
import chromadb import chromadb
from chromadb import QueryResult, Settings from chromadb import QueryResult, Settings
@ -15,6 +15,15 @@ from extensions.ext_redis import redis_client
from models.dataset import Dataset from models.dataset import Dataset
class ChromaParamsDict(TypedDict):
host: str
port: int
ssl: bool
tenant: str
database: str
settings: Settings
class ChromaConfig(BaseModel): class ChromaConfig(BaseModel):
host: str host: str
port: int port: int
@ -23,14 +32,13 @@ class ChromaConfig(BaseModel):
auth_provider: str | None = None auth_provider: str | None = None
auth_credentials: str | None = None auth_credentials: str | None = None
def to_chroma_params(self): def to_chroma_params(self) -> ChromaParamsDict:
settings = Settings( settings = Settings(
# auth # auth
chroma_client_auth_provider=self.auth_provider, chroma_client_auth_provider=self.auth_provider,
chroma_client_auth_credentials=self.auth_credentials, chroma_client_auth_credentials=self.auth_credentials,
) )
result: ChromaParamsDict = {
return {
"host": self.host, "host": self.host,
"port": self.port, "port": self.port,
"ssl": False, "ssl": False,
@ -38,6 +46,7 @@ class ChromaConfig(BaseModel):
"database": self.database, "database": self.database,
"settings": settings, "settings": settings,
} }
return result
class ChromaVector(BaseVector): class ChromaVector(BaseVector):