| from typing import Any |
|
|
| import httpx |
|
|
| from app.config import Settings |
|
|
|
|
| class EmbeddingServiceError(RuntimeError): |
| pass |
|
|
|
|
| class JinaEmbeddingService: |
| def __init__(self, settings: Settings, client: httpx.AsyncClient | None = None) -> None: |
| self.settings = settings |
| self._client = client |
|
|
| @property |
| def client(self) -> httpx.AsyncClient: |
| if self._client is None: |
| self._client = httpx.AsyncClient(timeout=self.settings.request_timeout_seconds) |
| return self._client |
|
|
| async def embed_query(self, text: str) -> list[float]: |
| embeddings = await self.embed_texts([text], task=self.settings.jina_query_task) |
| return embeddings[0] |
|
|
| async def embed_passages(self, texts: list[str]) -> list[list[float]]: |
| return await self.embed_texts(texts, task=self.settings.jina_passage_task) |
|
|
| async def embed_texts(self, texts: list[str], *, task: str) -> list[list[float]]: |
| cleaned = [text.strip() for text in texts if text and text.strip()] |
| if not cleaned: |
| return [] |
|
|
| payload: dict[str, Any] = { |
| "model": self.settings.jina_embedding_model, |
| "input": [ |
| _with_retrieval_prefix(text, task, self.settings.jina_embedding_model) |
| for text in cleaned |
| ], |
| "dimensions": self.settings.jina_embedding_dimensions, |
| "task": task, |
| "truncate": True, |
| } |
| headers = { |
| "Authorization": f"Bearer {self.settings.jina_api_key}", |
| "Content-Type": "application/json", |
| } |
|
|
| try: |
| response = await self.client.post( |
| self.settings.jina_embedding_endpoint, |
| headers=headers, |
| json=payload, |
| ) |
| response.raise_for_status() |
| except httpx.HTTPError as exc: |
| raise EmbeddingServiceError(f"Jina embedding request failed: {exc}") from exc |
|
|
| body = response.json() |
| data = body.get("data") |
| if not isinstance(data, list): |
| raise EmbeddingServiceError("Jina embedding response did not include a data array") |
|
|
| embeddings = [_embedding_from_item(item) for item in sorted(data, key=_embedding_index)] |
| if len(embeddings) != len(cleaned): |
| raise EmbeddingServiceError("Jina embedding response count did not match input count") |
| return embeddings |
|
|
|
|
| def _with_retrieval_prefix(text: str, task: str, model: str) -> str: |
| if "v5" not in model: |
| return text |
| if task == "retrieval.query" and not text.startswith("Query:"): |
| return f"Query: {text}" |
| if task == "retrieval.passage" and not text.startswith("Document:"): |
| return f"Document: {text}" |
| return text |
|
|
|
|
| def _embedding_index(item: Any) -> int: |
| if isinstance(item, dict): |
| return int(item.get("index") or 0) |
| return 0 |
|
|
|
|
| def _embedding_from_item(item: Any) -> list[float]: |
| if not isinstance(item, dict) or not isinstance(item.get("embedding"), list): |
| raise EmbeddingServiceError("Jina embedding item did not include an embedding array") |
| return [float(value) for value in item["embedding"]] |
|
|
