services_config.py
5.96 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
"""
Services configuration - single source for translation, embedding, rerank providers.
All provider selection and endpoint resolution is centralized here.
Priority: env vars > config.yaml > defaults.
"""
from __future__ import annotations
import os
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path
from typing import Any, Dict, Optional
import yaml
@dataclass
class ServiceConfig:
"""Config for one capability (translation/embedding/rerank)."""
provider: str
providers: Dict[str, Any] = field(default_factory=dict)
def get_provider_cfg(self) -> Dict[str, Any]:
"""Get config for current provider."""
p = (self.provider or "").strip().lower()
return self.providers.get(p, {}) if isinstance(self.providers, dict) else {}
def _load_services_raw(config_path: Optional[Path] = None) -> Dict[str, Any]:
"""Load services block from config.yaml."""
if config_path is None:
config_path = Path(__file__).parent / "config.yaml"
path = Path(config_path)
if not path.exists():
return {}
try:
with open(path, "r", encoding="utf-8") as f:
data = yaml.safe_load(f)
except Exception:
return {}
services = data.get("services") if isinstance(data, dict) else {}
return services if isinstance(services, dict) else {}
def _resolve_translation() -> ServiceConfig:
raw = _load_services_raw()
cfg = raw.get("translation", {}) if isinstance(raw.get("translation"), dict) else {}
providers = cfg.get("providers", {}) if isinstance(cfg.get("providers"), dict) else {}
provider = (
os.getenv("TRANSLATION_PROVIDER")
or cfg.get("provider")
or "direct"
)
provider = str(provider).strip().lower()
# Env override for http base_url
env_url = os.getenv("TRANSLATION_SERVICE_URL")
if env_url and provider == "http":
providers = dict(providers)
providers["http"] = dict(providers.get("http", {}))
providers["http"]["base_url"] = env_url.rstrip("/")
return ServiceConfig(provider=provider, providers=providers)
def _resolve_embedding() -> ServiceConfig:
raw = _load_services_raw()
cfg = raw.get("embedding", {}) if isinstance(raw.get("embedding"), dict) else {}
providers = cfg.get("providers", {}) if isinstance(cfg.get("providers"), dict) else {}
provider = (
os.getenv("EMBEDDING_PROVIDER")
or cfg.get("provider")
or "http"
)
provider = str(provider).strip().lower()
env_url = os.getenv("EMBEDDING_SERVICE_URL")
if env_url and provider == "http":
providers = dict(providers)
providers["http"] = dict(providers.get("http", {}))
providers["http"]["base_url"] = env_url.rstrip("/")
return ServiceConfig(provider=provider, providers=providers)
def _resolve_rerank() -> ServiceConfig:
raw = _load_services_raw()
cfg = raw.get("rerank", {}) if isinstance(raw.get("rerank"), dict) else {}
providers = cfg.get("providers", {}) if isinstance(cfg.get("providers"), dict) else {}
provider = (
os.getenv("RERANK_PROVIDER")
or cfg.get("provider")
or "http"
)
provider = str(provider).strip().lower()
env_url = os.getenv("RERANKER_SERVICE_URL")
if env_url:
url = env_url.rstrip("/")
if not url.endswith("/rerank"):
url = f"{url}/rerank" if "/rerank" not in url else url
providers = dict(providers)
providers["http"] = dict(providers.get("http", {}))
providers["http"]["base_url"] = url.replace("/rerank", "")
providers["http"]["service_url"] = url
return ServiceConfig(provider=provider, providers=providers)
def get_rerank_backend_config() -> tuple[str, dict]:
"""
Resolve reranker backend name and config for the reranker service process.
Returns (backend_name, backend_cfg).
Env RERANK_BACKEND overrides config.
"""
raw = _load_services_raw()
cfg = raw.get("rerank", {}) if isinstance(raw.get("rerank"), dict) else {}
backends = cfg.get("backends", {}) if isinstance(cfg.get("backends"), dict) else {}
name = (
os.getenv("RERANK_BACKEND")
or cfg.get("backend")
or "bge"
)
name = str(name).strip().lower()
backend_cfg = backends.get(name, {}) if isinstance(backends.get(name), dict) else {}
return name, backend_cfg
@lru_cache(maxsize=1)
def get_translation_config() -> ServiceConfig:
"""Get translation service config."""
return _resolve_translation()
@lru_cache(maxsize=1)
def get_embedding_config() -> ServiceConfig:
"""Get embedding service config."""
return _resolve_embedding()
@lru_cache(maxsize=1)
def get_rerank_config() -> ServiceConfig:
"""Get rerank service config."""
return _resolve_rerank()
def get_translation_base_url() -> str:
"""Resolve translation HTTP base URL (for http provider)."""
base = (
os.getenv("TRANSLATION_SERVICE_URL")
or get_translation_config().providers.get("http", {}).get("base_url")
or "http://127.0.0.1:6006"
)
return str(base).rstrip("/")
def get_embedding_base_url() -> str:
"""Resolve embedding HTTP base URL."""
base = (
os.getenv("EMBEDDING_SERVICE_URL")
or get_embedding_config().providers.get("http", {}).get("base_url")
or "http://127.0.0.1:6005"
)
return str(base).rstrip("/")
def get_rerank_service_url() -> str:
"""Resolve rerank service URL (full path including /rerank)."""
base = (
os.getenv("RERANKER_SERVICE_URL")
or get_rerank_config().providers.get("http", {}).get("service_url")
or get_rerank_config().providers.get("http", {}).get("base_url")
or "http://127.0.0.1:6007"
)
base = str(base).rstrip("/")
return base if base.endswith("/rerank") else f"{base}/rerank"
def clear_services_cache() -> None:
"""Clear cached config (for tests)."""
get_translation_config.cache_clear()
get_embedding_config.cache_clear()
get_rerank_config.cache_clear()