from typing import Optional

import requests

from openkosmos_ai.common.model import AIBaseConfig
from openkosmos_ai.nodes.base_node import BaseStateNode


class LitellmClientConfig(AIBaseConfig):
    base_url: str
    api_key: str
    rerank_model: Optional[str] = "bge-reranker-v2-m3"
    embedding_model: Optional[str] = "bge-m3"


class LitellmClientNode(BaseStateNode[LitellmClientConfig]):

    def __init__(self, config: LitellmClientConfig):
        super().__init__(config)
        self.header = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {config.api_key}"
        }

    def rerank(self, documents: list[str], query: str, top_n=3):
        return requests.post(self.config().base_url + "/v1/rerank",
                             headers=self.header,
                             json={
                                 "model": self.config().rerank_model,
                                 "query": query,
                                 "documents": documents,
                                 "top_n": top_n
                             }).json()["results"]

    def embeddings(self, input: str):
        return requests.post(self.config().base_url + "/v1/embeddings",
                             headers=self.header,
                             json={
                                 "model": self.config().embeddings_model,
                                 "input": input,
                                 "encoding_format": "float"
                             }).json()
