Browse Source
feat(jina): add multi-key failover support for embedding, reranking and wiki sync
master
feat(jina): add multi-key failover support for embedding, reranking and wiki sync
master
nwhco
3 weeks ago
6 changed files with 259 additions and 160 deletions
-
3config/production.env
-
36src/knowledge/embedding_factory.py
-
163src/knowledge/multi_jina_embedder.py
-
11src/knowledge/sync_wiki.py
-
21src/utils/jina_keys.py
-
125src/utils/reranker.py
@ -1,67 +1,55 @@ |
|||
import yaml |
|||
import os |
|||
from typing import Optional |
|||
from agno.knowledge.embedder.openai import OpenAIEmbedder |
|||
from agno.knowledge.embedder.jina import JinaEmbedder |
|||
from pathlib import Path |
|||
# If Agno supports generic OpenAI-like embedders, we use OpenAIEmbedder with base_url |
|||
from agno.knowledge.embedder.openai import OpenAIEmbedder |
|||
from src.knowledge.multi_jina_embedder import MultiKeyJinaEmbedder |
|||
from src.utils.jina_keys import get_jina_keys |
|||
|
|||
class EmbeddingFactory: |
|||
def __init__(self): |
|||
# Get the directory where this file (factory.py) is located |
|||
current_file_path = Path(__file__).resolve() |
|||
|
|||
# Navigate up to the project root |
|||
# If structure is: /app/src/models/factory.py |
|||
# .parent = models, .parent = src, .parent = app (root) |
|||
project_root = current_file_path.parent.parent.parent |
|||
|
|||
# Construct the absolute path |
|||
config_path = project_root / 'config' / 'embeddings.yaml' |
|||
|
|||
print(f"Loading config from: {config_path}") # Debug log |
|||
with open(config_path) as f: |
|||
# Simple variable expansion for ${VAR} |
|||
print(f"Loading config from: {config_path}") |
|||
with open(config_path, "r", encoding="utf-8") as f: |
|||
content = f.read() |
|||
for key, val in os.environ.items(): |
|||
content = content.replace(f"${{{key}}}", val) |
|||
self.config = yaml.safe_load(content) |
|||
|
|||
def get_embedder(self, model_name: Optional[str] = None): |
|||
# 1. Default Logic |
|||
if model_name is None: |
|||
model_name = self.config['embeddings']['default'] |
|||
|
|||
|
|||
|
|||
models_config = self.config['embeddings']['models'] |
|||
if model_name not in models_config: |
|||
raise ValueError(f"Embedding model '{model_name}' not found in config.") |
|||
|
|||
config = models_config[model_name] |
|||
provider = config['provider'] |
|||
# # 2. Provider Logic |
|||
api_key_env = config.get('api_key') |
|||
if api_key_env and api_key_env.startswith("${"): |
|||
if api_key_env and str(api_key_env).startswith("${"): |
|||
api_key = os.getenv(api_key_env[2:-1]) |
|||
else: |
|||
api_key = api_key_env |
|||
|
|||
# CASE B: OpenAI (Official) |
|||
if provider == "openai": |
|||
return OpenAIEmbedder( |
|||
id=config['id'], |
|||
dimensions=config['dimensions'], |
|||
api_key=api_key |
|||
) |
|||
|
|||
# CASE C: OpenAI Compatible (Jina API, etc.) |
|||
elif provider == "jinaai": |
|||
return JinaEmbedder( |
|||
all_keys = get_jina_keys() |
|||
if api_key and api_key not in all_keys: |
|||
all_keys.insert(0, api_key) |
|||
return MultiKeyJinaEmbedder( |
|||
id=config['id'], |
|||
dimensions=config['dimensions'], |
|||
api_key=api_key |
|||
api_keys=all_keys, |
|||
api_key=all_keys[0] if all_keys else None |
|||
) |
|||
|
|||
print(f"Unknown provider type: {provider}") |
|||
raise ValueError(f"Unknown provider type: {provider}") |
|||
@ -0,0 +1,163 @@ |
|||
import os |
|||
import requests |
|||
import aiohttp |
|||
import logging |
|||
from dataclasses import dataclass, field |
|||
from typing import List, Optional, Dict, Any, Tuple |
|||
from agno.knowledge.embedder.jina import JinaEmbedder |
|||
from src.utils.jina_keys import get_jina_keys |
|||
|
|||
logger = logging.getLogger(__name__) |
|||
|
|||
@dataclass |
|||
class MultiKeyJinaEmbedder(JinaEmbedder): |
|||
""" |
|||
JinaEmbedder subclass that supports multiple API keys with automatic failover/rotation. |
|||
If a key hits balance limit (402), forbidden (403), unauthorized (401), or rate limit (429), |
|||
it automatically switches to the next available key and retries. |
|||
""" |
|||
api_keys: List[str] = field(default_factory=get_jina_keys) |
|||
_current_key_idx: int = field(default=0, init=False, repr=False) |
|||
|
|||
def __post_init__(self): |
|||
if not self.api_keys: |
|||
self.api_keys = get_jina_keys() |
|||
if self.api_key and self.api_key not in self.api_keys: |
|||
self.api_keys.insert(0, self.api_key) |
|||
if self.api_keys: |
|||
self.api_key = self.api_keys[0] |
|||
|
|||
def _get_active_key(self) -> str: |
|||
if not self.api_keys: |
|||
if self.api_key: |
|||
return self.api_key |
|||
raise ValueError("No Jina API keys provided in JINA_API_KEYS or JINA_API_KEY") |
|||
return self.api_keys[self._current_key_idx % len(self.api_keys)] |
|||
|
|||
def _rotate_key(self) -> str: |
|||
if len(self.api_keys) > 1: |
|||
prev = self._get_active_key() |
|||
self._current_key_idx = (self._current_key_idx + 1) % len(self.api_keys) |
|||
new_key = self._get_active_key() |
|||
self.api_key = new_key |
|||
logger.warning(f"🔄 Rotating Jina API Key from {prev[:12]}... to {new_key[:12]}...") |
|||
return new_key |
|||
return self._get_active_key() |
|||
|
|||
def _get_headers_for_key(self, key: str) -> Dict[str, str]: |
|||
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {key}"} |
|||
if self.headers: |
|||
headers.update(self.headers) |
|||
return headers |
|||
|
|||
def _response(self, text: str) -> Dict[str, Any]: |
|||
data = { |
|||
"model": self.id, |
|||
"late_chunking": self.late_chunking, |
|||
"dimensions": self.dimensions, |
|||
"embedding_type": self.embedding_type, |
|||
"input": [text], |
|||
} |
|||
if self.user is not None: |
|||
data["user"] = self.user |
|||
if self.request_params: |
|||
data.update(self.request_params) |
|||
|
|||
keys_to_try = max(len(self.api_keys), 1) |
|||
last_exception = None |
|||
|
|||
for _ in range(keys_to_try): |
|||
key = self._get_active_key() |
|||
try: |
|||
headers = self._get_headers_for_key(key) |
|||
response = requests.post(self.base_url, headers=headers, json=data, timeout=self.timeout or 30.0) |
|||
if response.status_code in (401, 402, 403, 429): |
|||
logger.warning(f"⚠️ Jina key {key[:12]}... failed with HTTP {response.status_code}: {response.text[:120]}") |
|||
self._rotate_key() |
|||
continue |
|||
response.raise_for_status() |
|||
return response.json() |
|||
except Exception as e: |
|||
last_exception = e |
|||
logger.warning(f"⚠️ Exception with Jina key {key[:12]}...: {e}. Trying next key...") |
|||
self._rotate_key() |
|||
|
|||
if last_exception: |
|||
raise last_exception |
|||
raise RuntimeError("All Jina API keys failed") |
|||
|
|||
async def _async_response(self, text: str) -> Dict[str, Any]: |
|||
data = { |
|||
"model": self.id, |
|||
"late_chunking": self.late_chunking, |
|||
"dimensions": self.dimensions, |
|||
"embedding_type": self.embedding_type, |
|||
"input": [text], |
|||
} |
|||
if self.user is not None: |
|||
data["user"] = self.user |
|||
if self.request_params: |
|||
data.update(self.request_params) |
|||
|
|||
timeout = aiohttp.ClientTimeout(total=self.timeout or 30.0) |
|||
keys_to_try = max(len(self.api_keys), 1) |
|||
last_exception = None |
|||
|
|||
for _ in range(keys_to_try): |
|||
key = self._get_active_key() |
|||
try: |
|||
headers = self._get_headers_for_key(key) |
|||
async with aiohttp.ClientSession(timeout=timeout) as session: |
|||
async with session.post(self.base_url, headers=headers, json=data) as response: |
|||
if response.status in (401, 402, 403, 429): |
|||
logger.warning(f"⚠️ Async Jina key {key[:12]}... failed with HTTP {response.status}") |
|||
self._rotate_key() |
|||
continue |
|||
response.raise_for_status() |
|||
return await response.json() |
|||
except Exception as e: |
|||
last_exception = e |
|||
logger.warning(f"⚠️ Async Jina exception with key {key[:12]}...: {e}. Trying next key...") |
|||
self._rotate_key() |
|||
|
|||
if last_exception: |
|||
raise last_exception |
|||
raise RuntimeError("All Jina API keys failed in async request") |
|||
|
|||
async def _async_batch_response(self, texts: List[str]) -> Dict[str, Any]: |
|||
data = { |
|||
"model": self.id, |
|||
"late_chunking": self.late_chunking, |
|||
"dimensions": self.dimensions, |
|||
"embedding_type": self.embedding_type, |
|||
"input": texts, |
|||
} |
|||
if self.user is not None: |
|||
data["user"] = self.user |
|||
if self.request_params: |
|||
data.update(self.request_params) |
|||
|
|||
timeout = aiohttp.ClientTimeout(total=self.timeout or 60.0) |
|||
keys_to_try = max(len(self.api_keys), 1) |
|||
last_exception = None |
|||
|
|||
for _ in range(keys_to_try): |
|||
key = self._get_active_key() |
|||
try: |
|||
headers = self._get_headers_for_key(key) |
|||
async with aiohttp.ClientSession(timeout=timeout) as session: |
|||
async with session.post(self.base_url, headers=headers, json=data) as response: |
|||
if response.status in (401, 402, 403, 429): |
|||
logger.warning(f"⚠️ Async batch Jina key {key[:12]}... failed with HTTP {response.status}") |
|||
self._rotate_key() |
|||
continue |
|||
response.raise_for_status() |
|||
return await response.json() |
|||
except Exception as e: |
|||
last_exception = e |
|||
logger.warning(f"⚠️ Async batch Jina exception with key {key[:12]}...: {e}. Trying next key...") |
|||
self._rotate_key() |
|||
|
|||
if last_exception: |
|||
raise last_exception |
|||
raise RuntimeError("All Jina API keys failed in async batch request") |
|||
@ -0,0 +1,21 @@ |
|||
import os |
|||
from typing import List |
|||
|
|||
def get_jina_keys() -> List[str]: |
|||
""" |
|||
Returns a list of Jina API keys from environment variables. |
|||
Supports both JINA_API_KEYS (comma-separated) and JINA_API_KEY. |
|||
""" |
|||
keys = [] |
|||
keys_env = os.getenv("JINA_API_KEYS", "") |
|||
if keys_env: |
|||
for k in keys_env.split(","): |
|||
cleaned = k.strip() |
|||
if cleaned and cleaned not in keys: |
|||
keys.append(cleaned) |
|||
|
|||
single_key = os.getenv("JINA_API_KEY", "").strip() |
|||
if single_key and single_key not in keys: |
|||
keys.append(single_key) |
|||
|
|||
return keys |
|||
Write
Preview
Loading…
Cancel
Save
Reference in new issue