added gemini and openai rag

rag.py

@@ -0,0 +1,101 @@
+import requests
+import os
+import google.generativeai as genai
+
+from typing import List
+from utils import encode_image
+from PIL import Image
+
+class Rag:
+
+    def get_answer_from_gemini(self, query, imagePaths):
+        try:
+            genai.configure(api_key=os.environ['GEMINI_API_KEY'])
+            model = genai.GenerativeModel('gemini-1.5-flash')
+            
+            images = [Image.open(path) for path in imagePaths]
+            
+
+            chat = model.start_chat()
+            response = chat.send_message([*images, query])
+
+            answer = response.text
+
+            print(answer)
+            
+            return answer
+        
+        except Exception as e:
+            print(f"An error occurred while querying Gemini: {e}")
+            return f"Error: {str(e)}"
+        
+
+    def get_answer_from_openai(self, query, imagesPaths):
+        try:
+            print(f"Querying LLM for query={query}, imagesPaths={imagesPaths}")
+    
+            payload = self.__get_openai_api_payload(query, imagesPaths)
+
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {os.environ['OPENAI_API_KEY']}"
+            }
+    
+            response = requests.post(
+                url="https://api.openai.com/v1/chat/completions",
+                headers=headers,
+                json=payload
+            )
+            response.raise_for_status()  # Raise an HTTPError for bad responses
+    
+            answer = response.json()["choices"][0]["message"]["content"]
+    
+            print(answer)
+    
+            return answer
+    
+        except Exception as e:
+            print(f"An error occurred while querying OpenAI: {e}")
+            return None
+
+
+    def __get_openai_api_payload(self, query:str, imagesPaths:List[str]):
+        image_payload = []
+
+        for imagePath in imagesPaths:
+            base64_image = encode_image(imagePath)
+            image_payload.append({
+                "type": "image_url",
+                "image_url": {
+                    "url": f"data:image/jpeg;base64,{base64_image}"
+                }
+            })
+
+        payload = {
+            "model": "gpt-4o",
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": query
+                        },
+                        *image_payload
+                    ]
+                }
+            ],
+            "max_tokens": 1024
+        }
+
+        return payload
+    
+
+
+# if __name__ == "__main__":
+#     rag = Rag()
+    
+#     query = "Based on attached images, how many new cases were reported during second wave peak"
+#     imagesPaths = ["covid_slides_page_8.png", "covid_slides_page_8.png"]
+    
+#     rag.get_answer_from_gemini(query, imagesPaths)

requirements.txt

@@ -6,3 +6,4 @@ colpali_engine==0.3.4
 tqdm==4.66.5
 pillow==10.4.0
 spaces==0.30.4
+google-generativeai==0.8.3

utils.py

@@ -0,0 +1,5 @@
+import base64
+
+def encode_image(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode('utf-8')