import os
import azure.functions as func
import openai
from pinecone import Pinecone, ServerlessSpec
import logging
import json

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):   # get pinecone index for an input 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):   # convert text to embeddings
    res = openAI_client.embeddings.create(
        input=[text],
        model=embed_model
    )
    return res.data[0].embedding

def retreive_data(query, index, ns):    # retreive data from pinecone database
    query_embeds = embed_text(query)
    results = index.query(
        namespace=ns,
        vector=query_embeds,
        top_k=20,
        include_metadata=True
    )

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

def reply(prompt):  # get reply from LLM
    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 FUNCTION 👇
def main(req: func.HttpRequest) -> func.HttpResponse:   # query LLM based on context similar to the query received from pinecone
    try:
        query = req.get_json()
        device_id = query.get('device_id')
        if not device_id:
            return func.HttpResponse(json.dumps({"error": "Device ID is required"}), status_code=400, mimetype="application/json")
        
        index = get_index(device_id)
        query_text = query.get('query')

        if not query_text:
            return func.HttpResponse(json.dumps({"error": "Query text is required"}), status_code=400, mimetype="application/json")

        ns = query.get('namespace')    
        contexts = retreive_data(query_text, index, ns)

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

        answer = reply(prompt)
        return func.HttpResponse(json.dumps({"answer": answer}), status_code=200, mimetype="application/json")
    
    except Exception as e:
        logging.error(f"Error while querying")
        return func.HttpResponse(json.dumps({"error": str(e)}), status_code=500, mimetype="application/json")

ive developed a code to query data from pinecone database and ive already integrated with the frontend of a chatbot app, one problem seems to be that the LLM isn't holding onto any context of the chat, how can I implement that? also, tell me if my code is nice.