56 lines
2 KiB
Python
56 lines
2 KiB
Python
import logging
|
|
import os
|
|
from typing import Any
|
|
|
|
EMBEDDINGS_DENSE_MODEL = "Qwen/Qwen3-Embedding-0.6B"
|
|
SPARSE_MODEL_NAME = "Qdrant/bm25"
|
|
RERANKER_MODEL = "nvidia/llama-nemotron-rerank-1b-v2"
|
|
|
|
HOST = os.getenv("HOST", "0.0.0.0")
|
|
PORT = int(os.getenv("PORT", "8003"))
|
|
|
|
API_KEY = os.getenv("API_KEY")
|
|
EMBEDDINGS_DENSE_URL = os.getenv("EMBEDDINGS_DENSE_URL")
|
|
RERANKER_URL = os.getenv("RERANKER_URL")
|
|
QDRANT_URL = os.getenv("QDRANT_URL")
|
|
QDRANT_COLLECTION_NAME = os.getenv("QDRANT_COLLECTION_NAME", "evaluation")
|
|
QDRANT_DENSE_VECTOR_NAME = os.getenv("QDRANT_DENSE_VECTOR_NAME", "dense")
|
|
QDRANT_SPARSE_VECTOR_NAME = os.getenv("QDRANT_SPARSE_VECTOR_NAME", "sparse")
|
|
OPEN_API_LOGIN = os.getenv("OPEN_API_LOGIN")
|
|
OPEN_API_PASSWORD = os.getenv("OPEN_API_PASSWORD")
|
|
|
|
DENSE_PREFETCH_K = 50
|
|
SPARSE_PREFETCH_K = 100
|
|
RETRIEVE_K = 80
|
|
RERANK_LIMIT = 60
|
|
TOP_K = 50
|
|
|
|
HTTP_TIMEOUT = 30.0
|
|
HTTP_MAX_RETRIES = 2
|
|
|
|
REQUIRED_ENV_VARS = ["EMBEDDINGS_DENSE_URL", "RERANKER_URL", "QDRANT_URL"]
|
|
|
|
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
|
|
logger = logging.getLogger("search-service")
|
|
|
|
|
|
def validate_required_env() -> None:
|
|
if bool(OPEN_API_LOGIN) != bool(OPEN_API_PASSWORD):
|
|
raise RuntimeError("OPEN_API_LOGIN and OPEN_API_PASSWORD must be set together")
|
|
if not API_KEY and not (OPEN_API_LOGIN and OPEN_API_PASSWORD):
|
|
raise RuntimeError("Either API_KEY or OPEN_API_LOGIN and OPEN_API_PASSWORD must be set")
|
|
missing = [name for name in REQUIRED_ENV_VARS if not os.getenv(name)]
|
|
if missing:
|
|
logger.error("Empty required env vars: %s", ", ".join(missing))
|
|
raise RuntimeError(f"Empty required env vars: {', '.join(missing)}")
|
|
|
|
|
|
def get_upstream_kwargs() -> dict[str, Any]:
|
|
headers = {"Content-Type": "application/json"}
|
|
kwargs: dict[str, Any] = {"headers": headers}
|
|
if OPEN_API_LOGIN and OPEN_API_PASSWORD:
|
|
kwargs["auth"] = (OPEN_API_LOGIN, OPEN_API_PASSWORD)
|
|
return kwargs
|
|
if API_KEY:
|
|
headers["Authorization"] = f"Bearer {API_KEY}"
|
|
return kwargs
|