import os
#print(os.environ['ELASTIC_CLOUD_ID'])
#time.sleep(10)
from elasticsearch import Elasticsearch

from elasticsearch_llm_cache.elasticsearch_llm_cache import (
    ElasticsearchLLMCache,  # Import the class from the file
)

from pprint import pprint
import time

# Initialize Elasticsearch client
es_client = Elasticsearch(
    cloud_id=os.environ['ELASTIC_CLOUD_ID'],
    basic_auth=(os.environ['ELASTIC_USER'], os.environ['ELASTIC_PASSWORD']),
    request_timeout=30)

# Initialize your caching class
cache = ElasticsearchLLMCache(es_client=es_client, index_name="llm_cache_test")
#print(cache)

# Example prompt, response, and prompt vector
example_prompt = "What is the capital of France?"
example_response = "The capital of France is Paris."
print(f'example prompt: {example_prompt}')
print(f'example LLM Response: {example_response}')

#
print ('first pass attempt')
q1= 'What is the capital of France?'
print(f'query 1: {q1}')
query_result_1 = cache.query(prompt_text=q1)
print("Query result 1:") 
pprint(query_result_1)
print()


# Add document to cache
add_result = cache.add(prompt=example_prompt, response=example_response)
print("Add result to cache: ")
pprint(add_result)
print()
time.sleep(5)

# Query the cache (simulating later, separate operations)
print()
print ('Second pass attempt')
print("Testing cached similar results\n")
print()

q1= 'What is the capital of France?'
print(f'query 1: {q1}')
query_result_1 = cache.query(prompt_text=q1)
print("Query result 1:") 
pprint(query_result_1)
print()

q2= 'What is the currency of the UK?'
print(f'query 1: {q2}')
query_result_2 = cache.query(prompt_text=q2)
print("Query result 2:") 
pprint(query_result_2)
