Update handler.py

handler.py

@@ -1,8 +1,5 @@
-from typing import Dict, Any, List
-from environs import Env
 import json
-import requests
-
+from typing import Dict, List, Any
 
 def download_env_file(url: str, local_path: str):
     response = requests.get(url)
@@ -15,32 +12,51 @@ env_file_url = "https://www.dropbox.com/scl/fi/21ldek2cdsak2v3mhyy5x/openai.env?
 local_env_path = "openai.env"
 download_env_file(env_file_url, local_env_path)
 
-from coresugg import create_conversation_starter_prompt, generate_conversation_starters, NUMBER_OF_MESSAGES_FOR_CONTEXT as NUMBER_OF_MESSAGES_FOR_CONTEXT_SUGG
-from corechat import get_conversation_suggestions, NUMBER_OF_MESSAGES_FOR_CONTEXT as NUMBER_OF_MESSAGES_FOR_CONTEXT_CHAT
-
-env = Env()
+# Importing modules from both scripts
+from coresugg import ConversationPayload as ConversationPayloadSugg, create_conversation_starter_prompt, generate_conversation_starters, NUMBER_OF_MESSAGES_FOR_CONTEXT as NUMBER_OF_MESSAGES_FOR_CONTEXT_SUGG
+from corechat import ConversationPayload as ConversationPayloadChat, get_conversation_suggestions, NUMBER_OF_MESSAGES_FOR_CONTEXT as NUMBER_OF_MESSAGES_FOR_CONTEXT_CHAT
 
 class EndpointHandler:
-    async def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        if "FromUserKavasQuestions" in data and "ToUserKavasQuestions" in data and "Chatmood" in data:
-            from_user_questions = data["FromUserKavasQuestions"][-NUMBER_OF_MESSAGES_FOR_CONTEXT_SUGG:]
-            to_user_questions = data["ToUserKavasQuestions"][-NUMBER_OF_MESSAGES_FOR_CONTEXT_SUGG:]
-            ai_prompt = create_conversation_starter_prompt(from_user_questions + to_user_questions, data["Chatmood"])
-            conversation_starters = generate_conversation_starters(ai_prompt)
-            return [{"conversation_starters": conversation_starters}]
-
-        elif "LastChatMessages" in data:
-            last_chat_messages = data["LastChatMessages"][-NUMBER_OF_MESSAGES_FOR_CONTEXT_CHAT:]
-            suggestions = get_conversation_suggestions(last_chat_messages)
-            return [{"version": "1.0.0-alpha", "suggested_responses": suggestions}]
-
-        elif "FromUserKavasQuestions" in data and "Chatmood" in data:
-            prompt = create_conversation_starter_prompt(data["FromUserKavasQuestions"], data["Chatmood"])
+    def __init__(self):
+        pass
+
+    def integration(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        payload = ConversationPayloadSugg(**data)
+        from_user_questions = payload.FromUserKavasQuestions[-NUMBER_OF_MESSAGES_FOR_CONTEXT_SUGG:]
+        to_user_questions = payload.ToUserKavasQuestions[-NUMBER_OF_MESSAGES_FOR_CONTEXT_SUGG:]
+        ai_prompt = create_conversation_starter_prompt(from_user_questions + to_user_questions, payload.Chatmood)
+        conversation_starters = generate_conversation_starters(ai_prompt)
+        return {"conversation_starters": conversation_starters}
+
+    def chat_integration(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        payload = ConversationPayloadChat(**data)
+        last_chat_messages = payload.LastChatMessages[-NUMBER_OF_MESSAGES_FOR_CONTEXT_CHAT:]
+        suggestions = get_conversation_suggestions(last_chat_messages)
+        return {"version": "1.0.0-alpha", "suggested_responses": suggestions}
+
+    def upload(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        if "file" not in data:
+            raise Exception("No file provided")
+
+        file_data = data["file"]
+        try:
+            json_data = json.loads(file_data)
+        except json.JSONDecodeError:
+            raise Exception("Invalid JSON format.")
+
+        if "FromUserKavasQuestions" in json_data and "Chatmood" in json_data:
+            prompt = create_conversation_starter_prompt(
+                json_data["FromUserKavasQuestions"],
+                json_data["Chatmood"]
+            )
             starter_suggestion = generate_conversation_starters(prompt)
-            return [{"conversation_starter": starter_suggestion}]
-
+            return {"conversation_starter": starter_suggestion}
+        elif "LastChatMessages" in json_data:
+            last_chat_messages = json_data["LastChatMessages"][-NUMBER_OF_MESSAGES_FOR_CONTEXT_CHAT:]
+            response = {
+                "version": "1.0.0-alpha",
+                "suggested_responses": get_conversation_suggestions(last_chat_messages)
+            }
+            return response
         else:
-            raise ValueError("Invalid JSON structure")
-
-if __name__ == "__main__":
-    handler = EndpointHandler()
+            raise Exception("Invalid JSON structure.")