from elasticsearch import Elasticsearch
from FLUJOS_DATOS.TENSOR_FLOW.generate_embeddings import get_embeddings

# Configura la conexión a Elasticsearch
es = Elasticsearch([{'host': 'localhost', 'port': 9200, 'scheme': 'http'}])


def save_embedding(text, embeddings):
    body = {
        "text": text,
        "embedding": embeddings.tolist()  # Convertir el embedding (numpy array) a lista
    }
    response = es.index(index="text_embeddings", body=body)
    return response



def retrieve_embedding(text):
    query = {
        "query": {
            "match": {
                "text": text
            }
        }
    }
    response = es.search(index="text_embeddings", body=query)
    if response['hits']['total']['value'] > 0:
        return response['hits']['hits'][0]['_source']['embedding']
    else:
        return None
