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)
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.
     

      #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']
        latest_message={"content":self.generate_template(latest_message,task_id)}
        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}'


        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={'Content-Type': 'application/json'}) as response:
                try:
                    temp=  await response.json()
                    return temp["candidates"][0]
                except Exception as e:
                    return f"Error ⚠️ {e} {temp}"