import os
# from flask import Flask, request, jsonify
import openai
import pinecone
from tqdm.auto import tqdm
from time import sleep
# from pinecone.grpc import PineconeGRPC as Pinecone
from pinecone import Pinecone, ServerlessSpec
from groq import Groq
from dotenv import load_dotenv

load_dotenv("C:/Users/user/Documents/GitHub/Memory-Assist/test/test.env")
openai_key = os.getenv("openai_key")
openAI_client = openai.OpenAI(api_key=openai_key)
pckey = os.getenv("pinecone_key")
pc = Pinecone(api_key=pckey)

embed_model = "text-embedding-ada-002"

# UTILITY FUNCTIONS 👇
def get_index(device_id):
    index_name = f"id-{device_id}"
    if not pc.has_index(index_name):
        pc.create_index(
            name=index_name,
            dimension=1536,
            metric='cosine',
            spec=ServerlessSpec(cloud='aws', region='us-east-1')
        )
    
    return pc.Index(index_name)

def embed_text(text):
    res = openAI_client.embeddings.create(
        input=[text],
        model=embed_model
    )
    return res.data[0].embedding

def retreive_data(query, index):
    query_embeds = embed_text(query)
    results = index.query(
        vector=query_embeds,
        top_k=20,
        include_metadata=True
    )

    return [m['metadata']['text'] for m in results['matches']]

def reply(prompt):
    try:
        response = openAI_client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "You are a helpful assistant. "
                        "When answering questions, only use information from the context if it directly relates to the question. "
                        "Do not include irrelevant context information. "
                        "If the context contains relevant information, use it to answer the question, even if it may be incorrect. "
                        "If you notice the information from the context is incorrect, you may politely correct it after providing the context-based answer. "
                        "If the context does not contain information relevant to the question, answer based on your own knowledge without mentioning the context. "
                        "Do not say anything about the context lacking information."
                },
                {"role": "user", "content": f"{prompt}"},
            ]
        )

        return response.choices[0].message.content.strip()
    
    except openai.error.OpenAIError as e:
        return f"An error occurred: {e}"

# MAIN FUNCTIONS 👇
def query(query):
    device_id = query.get('device_id')
    if not device_id:
        return {"error": "Device ID is required"}, 400
    
    index = get_index(device_id)
    query_text = query.get('query')

    if not query_text:
        return {"error": "Query text is required"}, 400
    
    contexts = retreive_data(query_text, index)

    prompt = (
        "Context: \n" + "\n".join(contexts) +
        f"\nQuestion: {query_text}"
    )

    answer = reply(prompt)
    return {"answer": answer}, 200

query_data = {
    "device_id": "pc1-01",
    "query": "How much did Google pay to buy Youtube?"
}

output = query(query_data)
print(output)

I need to run the embed_model used in this code on Groq