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

from dotenv import load_dotenv
load_dotenv(dotenv_path="D:/projects/TwT/Functions/memory-assist-twt/.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)

index_name = os.getenv('index_name') or 'id-pc1-01'

embed_model = "text-embedding-ada-002"
CHAT_WINDOW_SIZE = int(os.getenv('CHAT_WINDOW_SIZE'))
MAX_TOKENS = int(os.getenv('MAX_TOKENS'))

# UTILITY FUNCTIONS 👇
def get_index():
    try:
        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)
    
    except Exception as e:
        logging.error("Error while getting index.")

def embed_text(text):   # convert text to embeddings
    try:
        res = openAI_client.embeddings.create(
            input=[text],
            model=embed_model
        )
        return res.data[0].embedding
    except Exception as e:
        logging.error("Error while embedding text.")

def retrieve_data(query, index, ns):
    """
    Performs similarity search on pinecone and retrieves relevant data.

    Args:
        query (string): Users query.
        index (index): Pinecone index.
        ns (string): Users device id (used as namespace).
    
    Returns:
        list: A list of query-relevant data from pinecone.
    """
    try:
        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']]
    except Exception as e:
        logging.error("Error while retrieving data from pinecone.")

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. Use context to answer questions if relevant, even if incorrect."
                        "Correct mistakes politely after using context. If irrelevant, rely on your own knowledge without mentioning the lack of context."},
                {"role": "user", "content": f"{prompt}"},
            ]
        )

        return response.choices[0].message.content.strip()
    
    except openai.APIError as e:
        logging.error(f"Error while getting reply from LLM: {str(e)}")
        return f"An error occurred: {e}"

def update_summary(summary, query, text):
    try:
        prompt = f"""Update the summary below with the latest conversation. Prioritize new information without losing existing summary.\n
                    Summary: {summary}\n
                    New Conversation:\nUser: {query}\nBot: {text}\n
                    Updated Summary:"""

        response = openAI_client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {'role': "system", "content": "You are a professional summarizer. Your task is to concisely summarize conversations into key points, prioritizing new information while retaining essential details from previous summaries."},
                {"role": "user", "content": f"{prompt}"},
            ],
            max_tokens=MAX_TOKENS
        )

        return response.choices[0].message.content.strip()
    
    except openai.APIError as e:
        logging.error(f"Error while getting reply from LLM while updating summary: {str(e)}")
        return f"An error occurred: {e}"

# CRUD 👇
def store_chat(cs_payload, device_id):
    """
    Asychronously store chat.
    """
    try:
        chat_storage_url = os.getenv('chat_storage_url')
        response = requests.post(chat_storage_url, json=cs_payload)
        response.raise_for_status()
        logging.info(f"Chat successfully stored for {device_id}.")
    except Exception as e:
        logging.error(f"Failed to store chat asychronously for {device_id}: {str(e)}")

# MAIN FUNCTION 👇
def main(req: func.HttpRequest) -> func.HttpResponse:   # query LLM based on context similar to the query received from pinecone
    """
    Generates an output based on user query

    Args(req: func.HttpRequest):
        device_id(string): Users device id (used as namespace)
        chat_session_id(Int): Chat session ID on users account
        query(string): Users query that needs to be resolved
    
    Returns(func.HttpResponse):
        string: Solution to users query
    """
    try:
        query = req.get_json()
        device_id = query.get('device_id')
        chat_session_id = query.get('chat_session_id')
        if not device_id:
            logging.error("No user ID found.")
            return func.HttpResponse(
                json.dumps({"error": "Device ID is required"}),
                status_code=400,
                mimetype="application/json"
            )
        
        if not chat_session_id:
            logging.error("No chat session ID found.")
            return func.HttpResponse(
                json.dumps({"error": "Chat session ID is required"}),
                status_code=400,
                mimetype="application/json"
            )
        
        try:
            index = get_index()
        except Exception as e:
            logging.error(f"{str(e)}")
        query_text = query.get('query')

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

        contexts = retrieve_data(query_text, index, ns=device_id)

        chat_retrieval_url = os.getenv('chat_retrieval_url')
        cr_payload = {
            'device_id': device_id,
            'chat_session_id': chat_session_id,
            'chat_window_size': CHAT_WINDOW_SIZE
        }
        chat_retrieval_response = requests.post(chat_retrieval_url, json=cr_payload)
        chat_retrieval_response.raise_for_status()

        chat_data = chat_retrieval_response.json()
        previous_chat = chat_data['chats']
        summary = chat_data['summary']

        prompt = (
            "Context: \n" + "\n".join(contexts) +
            "Previous chat for context: \n" + previous_chat +
            "Summary of the whole chat: \n" + summary +
            f"\nQuestion: {query_text}"
        )

        answer = reply(prompt)
        new_summary = update_summary(summary, query, answer)
        cs_payload = {
            'device_id': device_id,
            'chat_session_id': chat_session_id,
            'user_query': query_text,
            'bot_reply': answer,
            'summary': new_summary
        }

        storage_thread = threading.Thread(target=store_chat, args=(cs_payload, device_id))
        storage_thread.daemon = True
        storage_thread.start()

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

This is an API I use to query a chatbot, it takes 5-6 seconds to reply at least. Does this code look good or is it unprofessional? Can I make this any faster?