wren
41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from fastapi_modules.fastapi_leaudit.rag_engine.config import RAG_CONFIG
|
|
|
|
_instance: Any | None = None
|
|
|
|
|
|
def init_chroma() -> Any:
|
|
global _instance
|
|
if _instance is not None:
|
|
return _instance
|
|
|
|
import chromadb # lazy import to avoid hard failure before feature is enabled
|
|
import chromadb.config
|
|
|
|
host = RAG_CONFIG["CHROMA_HOST"]
|
|
if host:
|
|
token = RAG_CONFIG.get("CHROMA_TOKEN", "")
|
|
header = RAG_CONFIG.get("CHROMA_AUTH_HEADER", "X-Chroma-Token")
|
|
settings = (
|
|
chromadb.config.Settings(
|
|
chroma_client_auth_provider="chromadb.auth.token_authn.TokenAuthClientProvider",
|
|
chroma_client_auth_credentials=token,
|
|
chroma_auth_token_transport_header=header,
|
|
)
|
|
if token
|
|
else chromadb.config.Settings()
|
|
)
|
|
_instance = chromadb.HttpClient(host=host, port=RAG_CONFIG["CHROMA_PORT"], settings=settings)
|
|
else:
|
|
_instance = chromadb.PersistentClient(path=RAG_CONFIG["CHROMA_PERSIST_DIR"])
|
|
return _instance
|
|
|
|
|
|
def get_chroma() -> Any:
|
|
if _instance is None:
|
|
return init_chroma()
|
|
return _instance
|