import dashscope
from dashscope import TextEmbedding
from typing import List


class EmbeddingService:
    def __init__(self):
        from app.core.config import settings
        self.model = settings.embedding_model
        self.dimension = settings.embedding_dim
        dashscope.api_key = settings.dashscope_api_key

    def embed_texts(self, texts: List[str]) -> List[List[float]]:
        """批量文本嵌入"""
        response = TextEmbedding.call(
            model=self.model,
            input=texts,
        )

        if response.status_code == 200:
            embeddings = []
            for item in response.output.embeddings:
                embeddings.append(item.embedding)
            return embeddings
        raise Exception(f"Embedding failed: {response.code}")

    def embed_single(self, text: str) -> List[float]:
        """单个文本嵌入"""
        embeddings = self.embed_texts([text])
        return embeddings[0]


embedding_service = EmbeddingService()