DataOps-platform/app/services/embedding/qwen.py

from __future__ import annotations

import hashlib

import requests


class QwenEmbeddingError(RuntimeError):
    pass


class QwenEmbeddingClient:
    def __init__(self, *, api_key: str, base_url: str, model: str, dimension: int = 1024, timeout: int = 30):
        if not api_key or not base_url or not model:
            raise QwenEmbeddingError("Qwen embedding configuration is incomplete")
        self.api_key, self.base_url, self.model = api_key, base_url.rstrip("/"), model
        self.dimension, self.timeout = dimension, timeout

    @property
    def profile_key(self) -> str:
        return f"qwen:{self.model}:{self.dimension}"

    def cache_key(self, text: str) -> str:
        return hashlib.sha256(f"{self.model}:{self.dimension}:{text}".encode()).hexdigest()

    def embed(self, texts: list[str]) -> list[list[float]]:
        try:
            response = requests.post(f"{self.base_url}/embeddings", headers={"Authorization": "Bearer " + self.api_key}, json={"model": self.model, "input": texts, "dimensions": self.dimension}, timeout=self.timeout)
            response.raise_for_status()
            vectors = [item["embedding"] for item in sorted(response.json()["data"], key=lambda item: item["index"])]
        except requests.RequestException as exc:
            raise QwenEmbeddingError("Qwen embedding request failed") from exc
        if len(vectors) != len(texts) or any(len(vector) != self.dimension for vector in vectors):
            raise QwenEmbeddingError("Qwen embedding dimension mismatch")
        return vectors