Implement WebSocket service for handling conversations and integrate OpenAI client for text generation

src/services/_conversation_service.py

@@ -1,28 +1,18 @@
-import json
 from fastapi import WebSocket
 from bson import ObjectId
-from uuid import UUID
 from beanie import WriteRules
 
-from ._file_service import FileService
-from src.utils import OpenAIClient
-from src.repositories import ConversationRepository, MessageRepository
-from src.models import Conversation, Message, User
 from src.config import logger
+from src.utils import OpenAIClient
+from src.models import Conversation, User
+
+from ._websocket_service import WebSocketService
 
 
 class ConversationService:
     def __init__(self):
         self.openai_client = OpenAIClient
-        self.conversation_repository = ConversationRepository()
-        self.message_repository = MessageRepository()
-        self.file_service = FileService
-        self.functions_dictionary = {
-            "get_relevant_information": {
-                "service": self.file_service,
-                "function": "semantic_search",
-            },
-        }
+        self.websocket_service = WebSocketService
 
     async def __aenter__(self):
         return self
@@ -53,229 +43,5 @@ class ConversationService:
             )
 
     async def conversation(self, websocket: WebSocket):
-        try:
-            query_params = websocket.query_params
-            conversation_id = query_params.get("conversation_id")
-            modality = query_params.get("modality")
-
-            await websocket.accept()
-
-            if not conversation_id or not modality:
-                await websocket.close(
-                    code=1008, reason="Missing or invalid query parameters"
-                )
-                return
-
-            if modality == "text":
-                await self.handle_text_conversion(websocket, conversation_id)
-            elif modality == "voice":
-                await self.handle_voice_conversion(websocket, conversation_id)
-            else:
-                await websocket.close(code=1008, reason="Unsupported modality")
-                return
-        finally:
-            await self.handle_conversation_summary(websocket, conversation_id)
-            # await websocket.close()
-
-    async def handle_text_conversion(self, websocket, conversation_id):
-        while True:
-            data = await websocket.receive_text()
-            message = json.loads(data)
-
-            openai_server_events = {
-                "session_created": message["type"] == "session.created",
-                "function_call": (
-                    message["type"] == "response.function_call_arguments.done"
-                ),
-                "ai_response": (
-                    message["type"] == "response.done"
-                    and message["response"]["status"] == "completed"
-                    and message["response"]["output"][0]["type"] == "message"
-                    and message["response"]["output"][0]["role"] == "assistant"
-                ),
-            }
-
-            if openai_server_events["session_created"]:
-                user_query = input("Enter your query: ")
-                await self.handle_user_message(
-                    message_content=user_query, conversation_id=conversation_id
-                )
-
-                event_stage_1 = {
-                    "type": "conversation.item.create",
-                    "previous_item_id": None,
-                    "item": {
-                        "type": "message",
-                        "role": "user",
-                        "content": [
-                            {
-                                "type": "input_text",
-                                "text": user_query,
-                            }
-                        ],
-                    },
-                }
-                await websocket.send_text(json.dumps(event_stage_1))
-                event_response_2 = {"type": "response.create"}
-                await websocket.send_text(json.dumps(event_response_2))
-
-            if openai_server_events["function_call"]:
-                response = await self.handle_ai_function_call(message)
-                event_response = {
-                    "type": "conversation.item.create",
-                    "previous_item_id": None,
-                    "item": {
-                        "type": "message",
-                        "role": "user",
-                        "content": [
-                            {
-                                "type": "input_text",
-                                "text": response,
-                            }
-                        ],
-                    },
-                }
-                await websocket.send_text(json.dumps(event_response))
-                event_response_2 = {"type": "response.create"}
-                await websocket.send_text(json.dumps(event_response_2))
-
-            if openai_server_events["ai_response"]:
-                ai_response = message["response"]["output"][0]["content"][0]["text"]
-                await self.handle_ai_message(
-                    message_content=ai_response, conversation_id=conversation_id
-                )
-
-                user_query = input("Enter your query: ")
-                await self.handle_user_message(
-                    message_content=user_query, conversation_id=conversation_id
-                )
-
-                event_stage_1 = {
-                    "type": "conversation.item.create",
-                    "previous_item_id": None,
-                    "item": {
-                        "type": "message",
-                        "role": "user",
-                        "content": [
-                            {
-                                "type": "input_text",
-                                "text": user_query,
-                            }
-                        ],
-                    },
-                }
-                await websocket.send_text(json.dumps(event_stage_1))
-                event_response_2 = {"type": "response.create"}
-                await websocket.send_text(json.dumps(event_response_2))
-
-            elif message["type"] == "error":
-                logger.error(f"Error: {message['error']}")
-
-    async def handle_voice_conversion(self, websocket, conversation_id):
-        while True:
-            data = await websocket.receive_text()
-            message = json.loads(data)
-
-            openai_server_events = {
-                "function_call": (
-                    message["type"] == "response.function_call_arguments.done"
-                ),
-                "input_transcription": (
-                    message["type"]
-                    == "conversation.item.input_audio_transcription.completed"
-                ),
-                "output_transcription": (
-                    message["type"] == "response.done"
-                    and message["response"]["status"] == "completed"
-                    and message["response"]["output"][0]["type"] == "message"
-                ),
-            }
-
-            if openai_server_events["function_call"]:
-                response = await self.handle_ai_function_call(message)
-                event_response = {
-                    "type": "conversation.item.create",
-                    "previous_item_id": None,
-                    "item": {
-                        "type": "message",
-                        "role": "user",
-                        "content": [
-                            {
-                                "type": "input_text",
-                                "text": response,
-                            }
-                        ],
-                    },
-                }
-                await websocket.send_text(json.dumps(event_response))
-                event_response_2 = {"type": "response.create"}
-                await websocket.send_text(json.dumps(event_response_2))
-
-            elif openai_server_events["input_transcription"]:
-                user_message = message["transcript"]
-                await self.handle_user_message(
-                    message_content=user_message, conversation_id=conversation_id
-                )
-
-            elif openai_server_events["output_transcription"]:
-                ai_response = message["response"]["output"][0]["content"][0][
-                    "transcript"
-                ]
-                await self.handle_ai_message(
-                    message_content=ai_response, conversation_id=conversation_id
-                )
-
-            elif message["type"] == "error":
-                logger.error(f"Error: {message['error']}")
-
-    async def handle_user_message(self, message_content, conversation_id):
-        logger.info(f"User Query: {message_content}")
-        conversation_object = await Conversation.get(ObjectId(conversation_id))
-        message_object = Message(
-            conversation=conversation_object,
-            role="user",
-            content=message_content,
-        )
-        return await message_object.save(link_rule=WriteRules.WRITE)
-
-    async def handle_ai_message(self, message_content, conversation_id):
-        logger.info(f"AI Response: {message_content}")
-        conversation_object = await Conversation.get(ObjectId(conversation_id))
-        message_object = Message(
-            conversation=conversation_object,
-            role="assistant",
-            content=message_content,
-        )
-        return await message_object.save(link_rule=WriteRules.WRITE)
-
-    async def handle_conversation_summary(self, conversation_id):
-        conversation_object = await Conversation.get(ObjectId(conversation_id))
-        messages: list[Message] = await Message.find_many(
-            Message.conversation.id == conversation_object.id,
-            fetch_links=True,
-        ).to_list()
-        conversation_summary = "\n".join(
-            [f"{message.role}: {message.content}" for message in messages]
-        )
-        conversation_object.summary = conversation_summary
-        await conversation_object.save(link_rule=WriteRules.WRITE)
-
-    async def handle_ai_function_call(self, data):
-        tool_name = data["name"]
-        arguments = json.loads(data["arguments"])
-        logger.info(f"Function call: {tool_name} \nArguments: {arguments}")
-
-        function_info = self.functions_dictionary.get(tool_name)
-        if not function_info:
-            raise AttributeError(f"Function {tool_name} not found in dictionary")
-
-        service_class = function_info["service"]
-        async with service_class() as service:
-            func = getattr(service, function_info["function"], None)
-            if not func:
-                raise AttributeError(
-                    f"No such function {function_info['function']} in service"
-                )
-            response = await func(arguments["query"])
-            response = f"Context : {response}\n\n# Use this context to respond to the user query : {arguments['query']}"
-            return response
+        async with self.websocket_service() as ws_service:
+            await ws_service.handle_conversation(websocket)

src/services/_websocket_service.py

@@ -0,0 +1,264 @@
+import json
+from fastapi import WebSocket
+from bson import ObjectId
+from beanie import WriteRules
+
+from ._file_service import FileService
+from src.utils import OpenAIClient
+from src.repositories import ConversationRepository, MessageRepository
+from src.models import Conversation, Message
+from src.config import logger
+
+
+class WebSocketService:
+    def __init__(self):
+        self.openai_client = OpenAIClient
+        self.conversation_repository = ConversationRepository()
+        self.message_repository = MessageRepository()
+        self.file_service = FileService
+        self.functions_dictionary = {
+            "get_relevant_information": {
+                "service": self.file_service,
+                "function": "semantic_search",
+            },
+        }
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        pass
+
+    async def handle_conversation(self, websocket: WebSocket):
+        try:
+            query_params = websocket.query_params
+            conversation_id = query_params.get("conversation_id")
+            modality = query_params.get("modality")
+
+            await websocket.accept()
+
+            if not conversation_id or not modality:
+                await websocket.close(
+                    code=1008, reason="Missing or invalid query parameters"
+                )
+                return
+
+            if modality == "text":
+                await self.handle_text_conversion(websocket, conversation_id)
+            elif modality == "voice":
+                await self.handle_voice_conversion(websocket, conversation_id)
+            else:
+                await websocket.close(code=1008, reason="Unsupported modality")
+                return
+        finally:
+            await self.handle_conversation_summary(conversation_id)
+
+    async def handle_text_conversion(self, websocket, conversation_id):
+        while True:
+            data = await websocket.receive_text()
+            message = json.loads(data)
+
+            openai_server_events = {
+                "session_created": message["type"] == "session.created",
+                "function_call": (
+                    message["type"] == "response.function_call_arguments.done"
+                ),
+                "ai_response": (
+                    message["type"] == "response.done"
+                    and message["response"]["status"] == "completed"
+                    and message["response"]["output"][0]["type"] == "message"
+                    and message["response"]["output"][0]["role"] == "assistant"
+                ),
+            }
+
+            if openai_server_events["session_created"]:
+                user_query = input("Enter your query: ")
+                await self.handle_user_message(
+                    message_content=user_query, conversation_id=conversation_id
+                )
+
+                event_stage_1 = {
+                    "type": "conversation.item.create",
+                    "previous_item_id": None,
+                    "item": {
+                        "type": "message",
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "input_text",
+                                "text": user_query,
+                            }
+                        ],
+                    },
+                }
+                await websocket.send_text(json.dumps(event_stage_1))
+                event_response_2 = {"type": "response.create"}
+                await websocket.send_text(json.dumps(event_response_2))
+
+            if openai_server_events["function_call"]:
+                response = await self.handle_ai_function_call(message)
+                event_response = {
+                    "type": "conversation.item.create",
+                    "previous_item_id": None,
+                    "item": {
+                        "type": "message",
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "input_text",
+                                "text": response,
+                            }
+                        ],
+                    },
+                }
+                await websocket.send_text(json.dumps(event_response))
+                event_response_2 = {"type": "response.create"}
+                await websocket.send_text(json.dumps(event_response_2))
+
+            if openai_server_events["ai_response"]:
+                ai_response = message["response"]["output"][0]["content"][0]["text"]
+                await self.handle_ai_message(
+                    message_content=ai_response, conversation_id=conversation_id
+                )
+
+                user_query = input("Enter your query: ")
+                await self.handle_user_message(
+                    message_content=user_query, conversation_id=conversation_id
+                )
+
+                event_stage_1 = {
+                    "type": "conversation.item.create",
+                    "previous_item_id": None,
+                    "item": {
+                        "type": "message",
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "input_text",
+                                "text": user_query,
+                            }
+                        ],
+                    },
+                }
+                await websocket.send_text(json.dumps(event_stage_1))
+                event_response_2 = {"type": "response.create"}
+                await websocket.send_text(json.dumps(event_response_2))
+
+            elif message["type"] == "error":
+                logger.error(f"Error: {message['error']}")
+
+    async def handle_voice_conversion(self, websocket, conversation_id):
+        while True:
+            data = await websocket.receive_text()
+            message = json.loads(data)
+
+            openai_server_events = {
+                "function_call": (
+                    message["type"] == "response.function_call_arguments.done"
+                ),
+                "input_transcription": (
+                    message["type"]
+                    == "conversation.item.input_audio_transcription.completed"
+                ),
+                "output_transcription": (
+                    message["type"] == "response.done"
+                    and message["response"]["status"] == "completed"
+                    and message["response"]["output"][0]["type"] == "message"
+                ),
+            }
+
+            if openai_server_events["function_call"]:
+                response = await self.handle_ai_function_call(message)
+                event_response = {
+                    "type": "conversation.item.create",
+                    "previous_item_id": None,
+                    "item": {
+                        "type": "message",
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "input_text",
+                                "text": response,
+                            }
+                        ],
+                    },
+                }
+                await websocket.send_text(json.dumps(event_response))
+                event_response_2 = {"type": "response.create"}
+                await websocket.send_text(json.dumps(event_response_2))
+
+            elif openai_server_events["input_transcription"]:
+                user_message = message["transcript"]
+                await self.handle_user_message(
+                    message_content=user_message, conversation_id=conversation_id
+                )
+
+            elif openai_server_events["output_transcription"]:
+                ai_response = message["response"]["output"][0]["content"][0][
+                    "transcript"
+                ]
+                await self.handle_ai_message(
+                    message_content=ai_response, conversation_id=conversation_id
+                )
+
+            elif message["type"] == "error":
+                logger.error(f"Error: {message['error']}")
+
+    async def handle_user_message(self, message_content, conversation_id):
+        logger.info(f"User Query: {message_content}")
+        conversation_object = await Conversation.get(ObjectId(conversation_id))
+        message_object = Message(
+            conversation=conversation_object,
+            role="user",
+            content=message_content,
+        )
+        return await message_object.save(link_rule=WriteRules.WRITE)
+
+    async def handle_ai_message(self, message_content, conversation_id):
+        logger.info(f"AI Response: {message_content}")
+        conversation_object = await Conversation.get(ObjectId(conversation_id))
+        message_object = Message(
+            conversation=conversation_object,
+            role="assistant",
+            content=message_content,
+        )
+        return await message_object.save(link_rule=WriteRules.WRITE)
+
+    async def handle_conversation_summary(self, conversation_id):
+        conversation_object = await Conversation.get(ObjectId(conversation_id))
+        messages: list[Message] = await Message.find_many(
+            Message.conversation.id == conversation_object.id,
+            fetch_links=True,
+        ).to_list()
+
+        conversation_history = "\n".join(
+            [f"{message.role}: {message.content}" for message in messages]
+        )
+        query = "Generate Conversation Summary\n\n" + conversation_history
+        async with self.openai_client() as client:
+            conversation_summary = await client.text_generation(query=query)
+
+        logger.info(f"Conversation Summary: {conversation_summary}")
+
+        conversation_object.summary = conversation_summary
+        await conversation_object.save(link_rule=WriteRules.WRITE)
+
+    async def handle_ai_function_call(self, data):
+        tool_name = data["name"]
+        arguments = json.loads(data["arguments"])
+        logger.info(f"Function call: {tool_name} \nArguments: {arguments}")
+
+        function_info = self.functions_dictionary.get(tool_name)
+        if not function_info:
+            raise AttributeError(f"Function {tool_name} not found in dictionary")
+
+        service_class = function_info["service"]
+        async with service_class() as service:
+            func = getattr(service, function_info["function"], None)
+            if not func:
+                raise AttributeError(
+                    f"No such function {function_info['function']} in service"
+                )
+            response = await func(arguments["query"])
+            response = f"Context : {response}\n\n# Use this context to respond to the user query : {arguments['query']}"
+            return response

src/utils/_openai_client.py

@@ -2,6 +2,7 @@ import os
 import httpx
 from fastapi import HTTPException
 from aiortc import RTCPeerConnection, RTCSessionDescription
+from openai import AsyncOpenAI
 
 from src.config import logger
 
@@ -27,6 +28,7 @@ _OPENAI_TOOLS = [
 
 class OpenAIClient:
     def __init__(self):
+        self.client = None
         self.session_url = "https://api.openai.com/v1/realtime/sessions"
         self.model = "gpt-4o-mini-realtime-preview-2024-12-17"
         self.webrtc_url = f"https://api.openai.com/v1/realtime?model={self.model}"
@@ -40,6 +42,7 @@ class OpenAIClient:
 
     async def __aenter__(self):
         self.system_prompt = await self.prompt_loader("system_prompt.md")
+        self.client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))
         return self
 
     async def __aexit__(self, *args):
@@ -51,6 +54,19 @@ class OpenAIClient:
             prompt = file.read()
         return prompt
 
+    async def text_generation(self, query: str):
+
+        completion = await self.client.chat.completions.create(
+            model="gpt-4o-mini-2024-07-18",
+            messages=[
+                {
+                    "role": "user",
+                    "content": query,
+                },
+            ],
+        )
+        return completion.choices[0].message.content
+
     async def create_openai_session(self, text_mode_only=False):
         headers = {
             "Authorization": f'Bearer {os.getenv("OPENAI_API_KEY")}',