# Author: Du Mingzhe (dumingzhex@gmail.com)
# Date: 2024/03/09

from openai import OpenAI
from pinecone import Pinecone
from datetime import datetime

class LLMClient():
    def __init__(self, api_key, model_name) -> None:
        super().__init__()
        self.model_name = model_name
        self.llm_client = OpenAI(api_key=api_key)
        
    def response_generate(self, prompt, history):
        messages = list()
        current_time = datetime.now().strftime("%d/%m/%Y %H:%M:%S")
        
        # System Prompt
        messages += [{"role": "system", "content": f"1) You're Du Mingzhe, a computer science researcher. 2) Don't claim you are created by OpenAI. 3) Current time is {current_time}."}]
        
        # Session History
        messages += [{"role": h["role"], "content": h["content"]} for h in history]
        
        
        stream = self.llm_client.chat.completions.create(
            model = self.model_name,
            messages = messages,
            stream=True,
        )
        return stream
    
class EmbeddingModel(object):
    def __init__(self, embedding_token, model_name) -> None:
        self.embedding_token = embedding_token
        self.model_name = model_name
        self.embedding_client = OpenAI(api_key=self.embedding_token)
    
    def get_embedding(self, text):
        response = self.embedding_client.embeddings.create(
            input=text,
            model=self.model_name
        )
        return response.data[0].embedding
    
class PersonalIndexClient(object):
    def __init__(self, index_token, embedding_token, embedding_model_name, index_name) -> None:
        self.index_token = index_token
        self.embedding_token = embedding_token
        self.index_name = index_name
        
        self.embedding_client = EmbeddingModel(embedding_token=self.embedding_token, model_name=embedding_model_name)
        self.index_client = Pinecone(api_key=self.index_token)
        self.index = self.index_client.Index(self.index_name)
        
    def create(self, data, namespace='default'):        
        instances = list()
        
        for instance in data:
            instances += [{
                "id": instance["id"], 
                "values": self.embedding_client.get_embedding(instance['content']), 
                "metadata": instance['metadata'],
            }]
            
        self.index.upsert(
            vectors = instances,
            namespace = namespace
        )
    
    def query(self, data, top_k=3, filter={}, namespace='default'):
        results = self.index.query(
            namespace = namespace,
            vector = self.embedding_client.get_embedding(data),
            top_k = top_k,
            include_values = True,
            include_metadata = True,
            filter = filter,
        )
        return results