from flask import current_app
from langchain_mistralai import MistralAIEmbeddings
from typing import List, Any
import time

from common.eveai_model.eveai_embedding_base import EveAIEmbeddings
from common.utils.business_event_context import current_event
from mistralai import Mistral


class TrackedMistralAIEmbeddings(EveAIEmbeddings):
    def __init__(self, model: str = "mistral_embed"):
        api_key = current_app.config['MISTRAL_API_KEY']
        self.client = Mistral(
            api_key=api_key
        )
        self.model = model
        super().__init__()

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        start_time = time.time()
        result = self.client.embeddings.create(
            model=self.model,
            inputs=texts
        )
        end_time = time.time()

        metrics = {
            'total_tokens': result.usage.total_tokens,
            'prompt_tokens': result.usage.prompt_tokens,  # For embeddings, all tokens are prompt tokens
            'completion_tokens': result.usage.completion_tokens,
            'time_elapsed': end_time - start_time,
            'interaction_type': 'Embedding',
        }
        current_event.log_llm_metrics(metrics)

        embeddings = [embedding.embedding for embedding in result.data]

        return embeddings