From 7d95a002c720cd235f29a0fdee541b896fc530ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Wed, 11 Feb 2026 11:24:26 +0100 Subject: [PATCH] fix: improve model configuration for litellm gateway (#345) * fix: improve model configuration for litellm gateway * fix: add missing config imports for Cohere and LiteLLM providers Add missing DEFAULT_* and ENV_* constants to cross_encoder.py and embeddings.py imports: - DEFAULT_RERANKER_COHERE_MODEL - DEFAULT_LITELLM_API_BASE - DEFAULT_RERANKER_LITELLM_MODEL - DEFAULT_EMBEDDINGS_COHERE_MODEL - DEFAULT_EMBEDDINGS_LITELLM_MODEL - ENV_RERANKER_COHERE_MODEL This fixes NameError failures in test-api, test-hindsight-all, and test-upgrade CI jobs. --- hindsight-api/hindsight_api/config.py | 45 +++++++++++++++++-- .../hindsight_api/engine/cross_encoder.py | 27 ++++++----- .../hindsight_api/engine/embeddings.py | 28 ++++++------ hindsight-api/hindsight_api/main.py | 10 +++++ .../docs/developer/configuration.md | 25 ++++++----- 5 files changed, 91 insertions(+), 44 deletions(-) diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index 92002731..f69ab94d 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -71,17 +71,28 @@ ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY" ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL" ENV_EMBEDDINGS_OPENAI_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL" -ENV_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY" +# Cohere configuration (separate for embeddings and reranker) +ENV_EMBEDDINGS_COHERE_API_KEY = "HINDSIGHT_API_EMBEDDINGS_COHERE_API_KEY" ENV_EMBEDDINGS_COHERE_MODEL = "HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL" ENV_EMBEDDINGS_COHERE_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_COHERE_BASE_URL" +ENV_RERANKER_COHERE_API_KEY = "HINDSIGHT_API_RERANKER_COHERE_API_KEY" ENV_RERANKER_COHERE_MODEL = "HINDSIGHT_API_RERANKER_COHERE_MODEL" ENV_RERANKER_COHERE_BASE_URL = "HINDSIGHT_API_RERANKER_COHERE_BASE_URL" -# LiteLLM gateway configuration (for embeddings and reranker via LiteLLM proxy) +# Deprecated: Legacy shared Cohere API key (for backward compatibility) +ENV_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY" + +# LiteLLM configuration (separate for embeddings and reranker) +ENV_EMBEDDINGS_LITELLM_API_BASE = "HINDSIGHT_API_EMBEDDINGS_LITELLM_API_BASE" +ENV_EMBEDDINGS_LITELLM_API_KEY = "HINDSIGHT_API_EMBEDDINGS_LITELLM_API_KEY" +ENV_EMBEDDINGS_LITELLM_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL" +ENV_RERANKER_LITELLM_API_BASE = "HINDSIGHT_API_RERANKER_LITELLM_API_BASE" +ENV_RERANKER_LITELLM_API_KEY = "HINDSIGHT_API_RERANKER_LITELLM_API_KEY" +ENV_RERANKER_LITELLM_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_MODEL" + +# Deprecated: Legacy shared LiteLLM config (for backward compatibility) ENV_LITELLM_API_BASE = "HINDSIGHT_API_LITELLM_API_BASE" ENV_LITELLM_API_KEY = "HINDSIGHT_API_LITELLM_API_KEY" -ENV_EMBEDDINGS_LITELLM_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL" -ENV_RERANKER_LITELLM_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_MODEL" ENV_RERANKER_PROVIDER = "HINDSIGHT_API_RERANKER_PROVIDER" ENV_RERANKER_LOCAL_MODEL = "HINDSIGHT_API_RERANKER_LOCAL_MODEL" @@ -395,7 +406,12 @@ class HindsightConfig: embeddings_local_force_cpu: bool embeddings_tei_url: str | None embeddings_openai_base_url: str | None + embeddings_cohere_api_key: str | None + embeddings_cohere_model: str embeddings_cohere_base_url: str | None + embeddings_litellm_api_base: str + embeddings_litellm_api_key: str | None + embeddings_litellm_model: str # Reranker reranker_provider: str @@ -406,7 +422,12 @@ class HindsightConfig: reranker_tei_batch_size: int reranker_tei_max_concurrent: int reranker_max_candidates: int + reranker_cohere_api_key: str | None + reranker_cohere_model: str reranker_cohere_base_url: str | None + reranker_litellm_api_base: str + reranker_litellm_api_key: str | None + reranker_litellm_model: str # Server host: str @@ -588,7 +609,15 @@ class HindsightConfig: in ("true", "1"), embeddings_tei_url=os.getenv(ENV_EMBEDDINGS_TEI_URL), embeddings_openai_base_url=os.getenv(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None, + # Cohere embeddings (with backward-compatible fallback to shared API key) + embeddings_cohere_api_key=os.getenv(ENV_EMBEDDINGS_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY), + embeddings_cohere_model=os.getenv(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL), embeddings_cohere_base_url=os.getenv(ENV_EMBEDDINGS_COHERE_BASE_URL) or None, + # LiteLLM embeddings (with backward-compatible fallback to shared config) + embeddings_litellm_api_base=os.getenv(ENV_EMBEDDINGS_LITELLM_API_BASE) + or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE), + embeddings_litellm_api_key=os.getenv(ENV_EMBEDDINGS_LITELLM_API_KEY) or os.getenv(ENV_LITELLM_API_KEY), + embeddings_litellm_model=os.getenv(ENV_EMBEDDINGS_LITELLM_MODEL, DEFAULT_EMBEDDINGS_LITELLM_MODEL), # Reranker reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER), reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL), @@ -605,7 +634,15 @@ class HindsightConfig: os.getenv(ENV_RERANKER_TEI_MAX_CONCURRENT, str(DEFAULT_RERANKER_TEI_MAX_CONCURRENT)) ), reranker_max_candidates=int(os.getenv(ENV_RERANKER_MAX_CANDIDATES, str(DEFAULT_RERANKER_MAX_CANDIDATES))), + # Cohere reranker (with backward-compatible fallback to shared API key) + reranker_cohere_api_key=os.getenv(ENV_RERANKER_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY), + reranker_cohere_model=os.getenv(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL), reranker_cohere_base_url=os.getenv(ENV_RERANKER_COHERE_BASE_URL) or None, + # LiteLLM reranker (with backward-compatible fallback to shared config) + reranker_litellm_api_base=os.getenv(ENV_RERANKER_LITELLM_API_BASE) + or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE), + reranker_litellm_api_key=os.getenv(ENV_RERANKER_LITELLM_API_KEY) or os.getenv(ENV_LITELLM_API_KEY), + reranker_litellm_model=os.getenv(ENV_RERANKER_LITELLM_MODEL, DEFAULT_RERANKER_LITELLM_MODEL), # Server host=os.getenv(ENV_HOST, DEFAULT_HOST), port=int(os.getenv(ENV_PORT, DEFAULT_PORT)), diff --git a/hindsight-api/hindsight_api/engine/cross_encoder.py b/hindsight-api/hindsight_api/engine/cross_encoder.py index c7691045..31ded872 100644 --- a/hindsight-api/hindsight_api/engine/cross_encoder.py +++ b/hindsight-api/hindsight_api/engine/cross_encoder.py @@ -27,14 +27,10 @@ from ..config import ( DEFAULT_RERANKER_PROVIDER, DEFAULT_RERANKER_TEI_BATCH_SIZE, DEFAULT_RERANKER_TEI_MAX_CONCURRENT, - ENV_COHERE_API_KEY, - ENV_LITELLM_API_BASE, - ENV_LITELLM_API_KEY, - ENV_RERANKER_COHERE_BASE_URL, + ENV_RERANKER_COHERE_API_KEY, ENV_RERANKER_COHERE_MODEL, ENV_RERANKER_FLASHRANK_CACHE_DIR, ENV_RERANKER_FLASHRANK_MODEL, - ENV_RERANKER_LITELLM_MODEL, ENV_RERANKER_LOCAL_FORCE_CPU, ENV_RERANKER_LOCAL_MAX_CONCURRENT, ENV_RERANKER_LOCAL_MODEL, @@ -849,21 +845,24 @@ def create_cross_encoder_from_env() -> CrossEncoderModel: force_cpu=config.reranker_local_force_cpu, ) elif provider == "cohere": - api_key = os.environ.get(ENV_COHERE_API_KEY) + api_key = config.reranker_cohere_api_key if not api_key: - raise ValueError(f"{ENV_COHERE_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'cohere'") - model = os.environ.get(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL) - base_url = os.environ.get(ENV_RERANKER_COHERE_BASE_URL) or None - return CohereCrossEncoder(api_key=api_key, model=model, base_url=base_url) + raise ValueError(f"{ENV_RERANKER_COHERE_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'cohere'") + return CohereCrossEncoder( + api_key=api_key, + model=config.reranker_cohere_model, + base_url=config.reranker_cohere_base_url, + ) elif provider == "flashrank": model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL) cache_dir = os.environ.get(ENV_RERANKER_FLASHRANK_CACHE_DIR, DEFAULT_RERANKER_FLASHRANK_CACHE_DIR) return FlashRankCrossEncoder(model_name=model, cache_dir=cache_dir) elif provider == "litellm": - api_base = os.environ.get(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE) - api_key = os.environ.get(ENV_LITELLM_API_KEY) - model = os.environ.get(ENV_RERANKER_LITELLM_MODEL, DEFAULT_RERANKER_LITELLM_MODEL) - return LiteLLMCrossEncoder(api_base=api_base, api_key=api_key, model=model) + return LiteLLMCrossEncoder( + api_base=config.reranker_litellm_api_base, + api_key=config.reranker_litellm_api_key, + model=config.reranker_litellm_model, + ) elif provider == "rrf": return RRFPassthroughCrossEncoder() else: diff --git a/hindsight-api/hindsight_api/engine/embeddings.py b/hindsight-api/hindsight_api/engine/embeddings.py index 5c2eec53..125d8953 100644 --- a/hindsight-api/hindsight_api/engine/embeddings.py +++ b/hindsight-api/hindsight_api/engine/embeddings.py @@ -24,10 +24,7 @@ from ..config import ( DEFAULT_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_PROVIDER, DEFAULT_LITELLM_API_BASE, - ENV_COHERE_API_KEY, - ENV_EMBEDDINGS_COHERE_BASE_URL, - ENV_EMBEDDINGS_COHERE_MODEL, - ENV_EMBEDDINGS_LITELLM_MODEL, + ENV_EMBEDDINGS_COHERE_API_KEY, ENV_EMBEDDINGS_LOCAL_FORCE_CPU, ENV_EMBEDDINGS_LOCAL_MODEL, ENV_EMBEDDINGS_OPENAI_API_KEY, @@ -35,8 +32,6 @@ from ..config import ( ENV_EMBEDDINGS_OPENAI_MODEL, ENV_EMBEDDINGS_PROVIDER, ENV_EMBEDDINGS_TEI_URL, - ENV_LITELLM_API_BASE, - ENV_LITELLM_API_KEY, ENV_LLM_API_KEY, ) @@ -754,17 +749,20 @@ def create_embeddings_from_env() -> Embeddings: base_url = os.environ.get(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None return OpenAIEmbeddings(api_key=api_key, model=model, base_url=base_url) elif provider == "cohere": - api_key = os.environ.get(ENV_COHERE_API_KEY) + api_key = config.embeddings_cohere_api_key if not api_key: - raise ValueError(f"{ENV_COHERE_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'cohere'") - model = os.environ.get(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL) - base_url = os.environ.get(ENV_EMBEDDINGS_COHERE_BASE_URL) or None - return CohereEmbeddings(api_key=api_key, model=model, base_url=base_url) + raise ValueError(f"{ENV_EMBEDDINGS_COHERE_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'cohere'") + return CohereEmbeddings( + api_key=api_key, + model=config.embeddings_cohere_model, + base_url=config.embeddings_cohere_base_url, + ) elif provider == "litellm": - api_base = os.environ.get(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE) - api_key = os.environ.get(ENV_LITELLM_API_KEY) - model = os.environ.get(ENV_EMBEDDINGS_LITELLM_MODEL, DEFAULT_EMBEDDINGS_LITELLM_MODEL) - return LiteLLMEmbeddings(api_base=api_base, api_key=api_key, model=model) + return LiteLLMEmbeddings( + api_base=config.embeddings_litellm_api_base, + api_key=config.embeddings_litellm_api_key, + model=config.embeddings_litellm_model, + ) else: raise ValueError( f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei', 'openai', 'cohere', 'litellm'" diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index 7bc5df6d..65df255a 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/hindsight-api/hindsight_api/main.py @@ -199,7 +199,12 @@ def main(): embeddings_local_force_cpu=config.embeddings_local_force_cpu, embeddings_tei_url=config.embeddings_tei_url, embeddings_openai_base_url=config.embeddings_openai_base_url, + embeddings_cohere_api_key=config.embeddings_cohere_api_key, + embeddings_cohere_model=config.embeddings_cohere_model, embeddings_cohere_base_url=config.embeddings_cohere_base_url, + embeddings_litellm_api_base=config.embeddings_litellm_api_base, + embeddings_litellm_api_key=config.embeddings_litellm_api_key, + embeddings_litellm_model=config.embeddings_litellm_model, reranker_provider=config.reranker_provider, reranker_local_model=config.reranker_local_model, reranker_local_force_cpu=config.reranker_local_force_cpu, @@ -208,7 +213,12 @@ def main(): reranker_tei_batch_size=config.reranker_tei_batch_size, reranker_tei_max_concurrent=config.reranker_tei_max_concurrent, reranker_max_candidates=config.reranker_max_candidates, + reranker_cohere_api_key=config.reranker_cohere_api_key, + reranker_cohere_model=config.reranker_cohere_model, reranker_cohere_base_url=config.reranker_cohere_base_url, + reranker_litellm_api_base=config.reranker_litellm_api_base, + reranker_litellm_api_key=config.reranker_litellm_api_key, + reranker_litellm_model=config.reranker_litellm_model, host=args.host, port=args.port, log_level=args.log_level, diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index 9eaf5ad2..c2711ecf 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -273,11 +273,11 @@ export HINDSIGHT_API_RETAIN_LLM_MAX_BACKOFF=120.0 # Cap at 2min instead of 1m | `HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY` | OpenAI API key (falls back to `HINDSIGHT_API_LLM_API_KEY`) | - | | `HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL` | OpenAI embedding model | `text-embedding-3-small` | | `HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL` | Custom base URL for OpenAI-compatible API (e.g., Azure OpenAI) | - | -| `HINDSIGHT_API_COHERE_API_KEY` | Cohere API key (shared for embeddings and reranker) | - | +| `HINDSIGHT_API_EMBEDDINGS_COHERE_API_KEY` | Cohere API key for embeddings | - | | `HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL` | Cohere embedding model | `embed-english-v3.0` | | `HINDSIGHT_API_EMBEDDINGS_COHERE_BASE_URL` | Custom base URL for Cohere-compatible API (e.g., Azure-hosted) | - | -| `HINDSIGHT_API_LITELLM_API_BASE` | LiteLLM proxy base URL (shared for embeddings and reranker) | `http://localhost:4000` | -| `HINDSIGHT_API_LITELLM_API_KEY` | LiteLLM proxy API key (optional, depends on proxy config) | - | +| `HINDSIGHT_API_EMBEDDINGS_LITELLM_API_BASE` | LiteLLM proxy base URL for embeddings | `http://localhost:4000` | +| `HINDSIGHT_API_EMBEDDINGS_LITELLM_API_KEY` | LiteLLM proxy API key for embeddings (optional, depends on proxy config) | - | | `HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL` | LiteLLM embedding model (use provider prefix, e.g., `cohere/embed-english-v3.0`) | `text-embedding-3-small` | ```bash @@ -302,19 +302,19 @@ export HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080 # Cohere - cloud-based embeddings export HINDSIGHT_API_EMBEDDINGS_PROVIDER=cohere -export HINDSIGHT_API_COHERE_API_KEY=your-api-key +export HINDSIGHT_API_EMBEDDINGS_COHERE_API_KEY=your-api-key export HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL=embed-english-v3.0 # 1024 dimensions # Azure-hosted Cohere - embeddings via custom endpoint export HINDSIGHT_API_EMBEDDINGS_PROVIDER=cohere -export HINDSIGHT_API_COHERE_API_KEY=your-azure-api-key +export HINDSIGHT_API_EMBEDDINGS_COHERE_API_KEY=your-azure-api-key export HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL=embed-english-v3.0 export HINDSIGHT_API_EMBEDDINGS_COHERE_BASE_URL=https://your-azure-cohere-endpoint.com # LiteLLM proxy - unified gateway for multiple providers export HINDSIGHT_API_EMBEDDINGS_PROVIDER=litellm -export HINDSIGHT_API_LITELLM_API_BASE=http://localhost:4000 -export HINDSIGHT_API_LITELLM_API_KEY=your-litellm-key # optional +export HINDSIGHT_API_EMBEDDINGS_LITELLM_API_BASE=http://localhost:4000 +export HINDSIGHT_API_EMBEDDINGS_LITELLM_API_KEY=your-litellm-key # optional export HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL=text-embedding-3-small # or cohere/embed-english-v3.0 ``` @@ -344,8 +344,11 @@ Supported OpenAI embedding dimensions: | `HINDSIGHT_API_RERANKER_TEI_URL` | TEI server URL | - | | `HINDSIGHT_API_RERANKER_TEI_BATCH_SIZE` | Batch size for TEI reranking | `128` | | `HINDSIGHT_API_RERANKER_TEI_MAX_CONCURRENT` | Max concurrent TEI reranking requests | `8` | +| `HINDSIGHT_API_RERANKER_COHERE_API_KEY` | Cohere API key for reranking | - | | `HINDSIGHT_API_RERANKER_COHERE_MODEL` | Cohere rerank model | `rerank-english-v3.0` | | `HINDSIGHT_API_RERANKER_COHERE_BASE_URL` | Custom base URL for Cohere-compatible API (e.g., Azure-hosted) | - | +| `HINDSIGHT_API_RERANKER_LITELLM_API_BASE` | LiteLLM proxy base URL for reranking | `http://localhost:4000` | +| `HINDSIGHT_API_RERANKER_LITELLM_API_KEY` | LiteLLM proxy API key for reranking (optional, depends on proxy config) | - | | `HINDSIGHT_API_RERANKER_LITELLM_MODEL` | LiteLLM rerank model (use provider prefix, e.g., `cohere/rerank-english-v3.0`) | `cohere/rerank-english-v3.0` | | `HINDSIGHT_API_RERANKER_FLASHRANK_MODEL` | FlashRank model for fast CPU-based reranking | `ms-marco-MiniLM-L-12-v2` | | `HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR` | Cache directory for FlashRank models | System default | @@ -361,19 +364,19 @@ export HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081 # Cohere - cloud-based reranking export HINDSIGHT_API_RERANKER_PROVIDER=cohere -export HINDSIGHT_API_COHERE_API_KEY=your-api-key # shared with embeddings +export HINDSIGHT_API_RERANKER_COHERE_API_KEY=your-api-key export HINDSIGHT_API_RERANKER_COHERE_MODEL=rerank-english-v3.0 # Azure-hosted Cohere - reranking via custom endpoint export HINDSIGHT_API_RERANKER_PROVIDER=cohere -export HINDSIGHT_API_COHERE_API_KEY=your-azure-api-key +export HINDSIGHT_API_RERANKER_COHERE_API_KEY=your-azure-api-key export HINDSIGHT_API_RERANKER_COHERE_MODEL=rerank-english-v3.0 export HINDSIGHT_API_RERANKER_COHERE_BASE_URL=https://your-azure-cohere-endpoint.com # LiteLLM proxy - unified gateway for multiple reranking providers export HINDSIGHT_API_RERANKER_PROVIDER=litellm -export HINDSIGHT_API_LITELLM_API_BASE=http://localhost:4000 -export HINDSIGHT_API_LITELLM_API_KEY=your-litellm-key # optional +export HINDSIGHT_API_RERANKER_LITELLM_API_BASE=http://localhost:4000 +export HINDSIGHT_API_RERANKER_LITELLM_API_KEY=your-litellm-key # optional export HINDSIGHT_API_RERANKER_LITELLM_MODEL=cohere/rerank-english-v3.0 # or voyage/rerank-2, together_ai/... ```