Rag is ready?

App/Chat/utils/RAG.py

@@ -1,32 +1,65 @@
 import aiohttp
-import asyncio,pprint
+import asyncio
+import json,os
+import yaml
 import google.generativeai as palm
-from langchain.chains.question_answering import load_qa_chain
-from langchain.llms import GooglePalm
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain import PromptTemplate
-import os
-PALM_API = ''
-API_KEY=os.environ.get("PALM_API",PALM_API)
+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='chat-bison-001', 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(question,task_id):
+        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.
+       Make it short and to the point, not more than 1 PARAGRAPH. 
 
-def count_tokens(text):
-    return palm.count_message_tokens(prompt=text)['token_count']
-llm = GooglePalm(
-    google_api_key=API_KEY,       **{ "safety_settings": [
-            {"category": "HARM_CATEGORY_DEROGATORY", "threshold": 4},
-            {"category": "HARM_CATEGORY_TOXICITY", "threshold": 4},
-            {"category": "HARM_CATEGORY_VIOLENCE", "threshold": 4},
-            {"category": "HARM_CATEGORY_SEXUAL", "threshold": 4},
-            {"category": "HARM_CATEGORY_MEDICAL", "threshold": 4},
-            {"category": "HARM_CATEGORY_DANGEROUS", "threshold": 4},
-        ]})
-text_splitter = RecursiveCharacterTextSplitter(separators=["\n\n", "\n","."], chunk_size=40_000, chunk_overlap=500)
-with open('./sample.txt', 'r') as file:
-    essay = file.read()
+      #Context
+      {context_yaml}
 
-docs = text_splitter.create_documents([essay])
-for doc in docs:
-    print(count_tokens(doc.page_content))
+      #Question
+      {question}
+        '''
+        return Template
+
+    async def generate_message(self,  messages,task_id):
+        latest_message = messages[-1]['content']
+        latest_message={"content":self.generate_template(latest_message,task_id)}
+        messages[-1]=latest_message
+        url = f'https://generativelanguage.googleapis.com/v1beta3/models/{self.model}:generateMessage?key={self.api_key}'
+        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=data, headers={'Content-Type': 'application/json'}) as response:
+                try:
+                    return await response.json()['candidates']
+                except Exception as e:
+                    return f"Error ⚠️ {e}"
+                
+