refac: conditional import of chromadb

This commit is contained in:
Rômulo Mendes Figueiredo 2025-02-25 11:35:36 -03:00
parent 6fedd72e39
commit b52a7e3543

View File

@ -9,7 +9,6 @@ from pathlib import Path
from typing import Generic, Optional, TypeVar from typing import Generic, Optional, TypeVar
from urllib.parse import urlparse from urllib.parse import urlparse
import chromadb
import requests import requests
from pydantic import BaseModel from pydantic import BaseModel
from sqlalchemy import JSON, Column, DateTime, Integer, func from sqlalchemy import JSON, Column, DateTime, Integer, func
@ -1497,22 +1496,24 @@ Ensure that the tools are effectively utilized to achieve the highest-quality an
VECTOR_DB = os.environ.get("VECTOR_DB", "chroma") VECTOR_DB = os.environ.get("VECTOR_DB", "chroma")
# Chroma # Chroma
CHROMA_DATA_PATH = f"{DATA_DIR}/vector_db" if VECTOR_DB == "chroma":
CHROMA_TENANT = os.environ.get("CHROMA_TENANT", chromadb.DEFAULT_TENANT) import chromadb
CHROMA_DATABASE = os.environ.get("CHROMA_DATABASE", chromadb.DEFAULT_DATABASE) CHROMA_DATA_PATH = f"{DATA_DIR}/vector_db"
CHROMA_HTTP_HOST = os.environ.get("CHROMA_HTTP_HOST", "") CHROMA_TENANT = os.environ.get("CHROMA_TENANT", chromadb.DEFAULT_TENANT)
CHROMA_HTTP_PORT = int(os.environ.get("CHROMA_HTTP_PORT", "8000")) CHROMA_DATABASE = os.environ.get("CHROMA_DATABASE", chromadb.DEFAULT_DATABASE)
CHROMA_CLIENT_AUTH_PROVIDER = os.environ.get("CHROMA_CLIENT_AUTH_PROVIDER", "") CHROMA_HTTP_HOST = os.environ.get("CHROMA_HTTP_HOST", "")
CHROMA_CLIENT_AUTH_CREDENTIALS = os.environ.get("CHROMA_CLIENT_AUTH_CREDENTIALS", "") CHROMA_HTTP_PORT = int(os.environ.get("CHROMA_HTTP_PORT", "8000"))
# Comma-separated list of header=value pairs CHROMA_CLIENT_AUTH_PROVIDER = os.environ.get("CHROMA_CLIENT_AUTH_PROVIDER", "")
CHROMA_HTTP_HEADERS = os.environ.get("CHROMA_HTTP_HEADERS", "") CHROMA_CLIENT_AUTH_CREDENTIALS = os.environ.get("CHROMA_CLIENT_AUTH_CREDENTIALS", "")
if CHROMA_HTTP_HEADERS: # Comma-separated list of header=value pairs
CHROMA_HTTP_HEADERS = os.environ.get("CHROMA_HTTP_HEADERS", "")
if CHROMA_HTTP_HEADERS:
CHROMA_HTTP_HEADERS = dict( CHROMA_HTTP_HEADERS = dict(
[pair.split("=") for pair in CHROMA_HTTP_HEADERS.split(",")] [pair.split("=") for pair in CHROMA_HTTP_HEADERS.split(",")]
) )
else: else:
CHROMA_HTTP_HEADERS = None CHROMA_HTTP_HEADERS = None
CHROMA_HTTP_SSL = os.environ.get("CHROMA_HTTP_SSL", "false").lower() == "true" CHROMA_HTTP_SSL = os.environ.get("CHROMA_HTTP_SSL", "false").lower() == "true"
# this uses the model defined in the Dockerfile ENV variable. If you dont use docker or docker based deployments such as k8s, the default embedding model will be used (sentence-transformers/all-MiniLM-L6-v2) # this uses the model defined in the Dockerfile ENV variable. If you dont use docker or docker based deployments such as k8s, the default embedding model will be used (sentence-transformers/all-MiniLM-L6-v2)
# Milvus # Milvus