App/Chat/utils/RAG.py

import aiohttp
import asyncio
import json, os
import yaml
import google.generativeai as palm
from App.Embedding.utils.Initialize import search

PALM_API = ""
API_KEY = os.environ.get("PALM_API", PALM_API)
GroqAPIKey = os.environ.get("Groq_API", "")
# palm.configure(api_key=API_KEY)


class GenerativeAIAssistant:
    def __init__(
        self,
        api_key=API_KEY,
        model="gemini-pro",
        temperature=0.85,
        candidate_count=1,
        top_k=40,
        top_p=0.95,
    ):
        self.api_key = api_key
        self.model = model
        self.temperature = temperature
        self.candidate_count = candidate_count
        self.top_k = top_k
        self.top_p = top_p
        self.examples = [
            {
                "input": {"content": "hello"},
                "output": {"content": "Hello to you too! How can I help you today?"},
            }
        ]
        self.context = "You are a helpful assistant"

    def generate_template(self, question, task_id, summary=None):
        if summary == None:
            self.context = "You are a helpful assistant"
        else:
            self.context = summary
        contexts = search(question, task_id=task_id)
        context_yaml = ""
        for context in contexts:
            context_yaml += "\n" + yaml.dump(context)
        Template = f""" 
      #Instructions
      You are given the following context in yaml of a transcript of a youtube video, the start and end times are indicated and the text that was said is also given. You are also given a question, use the context to answer the question in a consise manner, make it short and to the point, don't provide additional details. If its just a conversational question,ignore the context, respond with the appropriate response. eg Hello hi, how are you? etc.
     

      #Context
      {context_yaml}

      #Your response be as short as possible and to the point.
      
      #Question
      {question}
        """
        return Template

    async def generate_message(self, messages, task_id="ok"):
        user_message = messages[-1]
        # latest_message = messages[-1]["parts"][0]["text"]
      
        user_message["content"]=self.generate_template(user_message["content"], task_id)
        headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {GroqAPIKey}",
        }
        # user_message["parts"][0]["text"] = latest_message
        messages[-1]=user_message
        # url = f'https://generativelanguage.googleapis.com/v1beta/models/{self.model}:generateContent?key={self.api_key}'
        url = "https://api.groq.com/openai/v1/chat/completions"
        payload = {
            "messages": messages,
            "model": "mixtral-8x7b-32768",
        }
        # payload = {
        #     "contents": messages,
        #     "generationConfig": {
        #         "temperature": 0.9,
        #         "topK": 1,
        #         "topP": 1,
        #         "maxOutputTokens": 2048,
        #         "stopSequences": [],
        #     },
        #     "safetySettings": [
        #         {
        #             "category": "HARM_CATEGORY_HARASSMENT",
        #             "threshold": "BLOCK_ONLY_HIGH",
        #         },
        #         {
        #             "category": "HARM_CATEGORY_HATE_SPEECH",
        #             "threshold": "BLOCK_ONLY_HIGH",
        #         },
        #         {
        #             "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
        #             "threshold": "BLOCK_ONLY_HIGH",
        #         },
        #         {
        #             "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
        #             "threshold": "BLOCK_ONLY_HIGH",
        #         },
        #     ],
        # }

        # data = {
        #     "prompt": {
        #         "context": self.context,
        #         "examples": self.examples,
        #         "messages": messages
        #     },
        #     "temperature": self.temperature,
        #     "top_k": self.top_k,
        #     "top_p": self.top_p,
        #     "candidate_count": self.candidate_count
        # }

        async with aiohttp.ClientSession() as session:
            async with session.post(
                url, json=payload, headers=headers
            ) as response:
                try:
                    temp = await response.json()
                    return temp["choices"][0]['message']["content"]
                except Exception as e:
                    return f"Error ⚠️ {e} {temp}"