from openkosmos_ai.nodes.config import ModelGatewayConfig
from openkosmos_ai.common.model import RerankResult, EmbeddingResult, RelevanceResult
from openkosmos_ai.nodes.api_client import ApiClientNode


class ModelGatewayNode(ApiClientNode[ModelGatewayConfig]):

    def __init__(self, config: ModelGatewayConfig):
        super().__init__(config)

    def rerank(self, query: str, documents: list[str], top_n=3) -> RerankResult:
        request_body = {
            "model": self.config().rerank_model,
            "query": query,
            "documents": documents,
            "top_n": top_n
        }
        return (self.post("/rerank")
                .json_data(request_body)
                .retrieve()
                .map(lambda result:
                     RerankResult(results=[RelevanceResult(
                         index=r["index"],
                         relevance_score=r["relevance_score"])
                         for r in result["results"]], metadata=result["meta"]))
                .body()
                )

    def embeddings(self, input: str) -> EmbeddingResult:
        request_body = {
            "model": self.config().embedding_model,
            "input": input,
            "encoding_format": "float"
        }
        return EmbeddingResult.model_validate_json(
            self.post("/embeddings").json_data(request_body).retrieve().body())
