Update services/chat_service.py

services/chat_service.py

@@ -41,6 +41,7 @@ class ConversationManager:
             del self.conversations[session_id]
 
 class ChatService:
+    """Main chat service that coordinates responses"""
     """Main chat service that coordinates responses"""
     def __init__(
         self,
@@ -55,7 +56,91 @@ class ChatService:
         self.pdf_service = pdf_service
         self.faq_service = faq_service
         self.conversation_manager = ConversationManager()
-        
+
+    def construct_system_prompt(self, context: str) -> str:
+        """Constructs the system message."""
+        return (
+            "You are a friendly bot specializing in Bofrost products. "
+            "Return comprehensive German answers. Always add product IDs. "
+            "Use the following product descriptions:\n\n"
+            f"{context}\n\n"
+        )
+
+    def construct_prompt(
+        self, 
+        user_input: str, 
+        context: str, 
+        chat_history: List[Tuple[str, str]], 
+        max_history_turns: int = 1
+    ) -> str:
+        """Constructs the full prompt."""
+        # System message
+        system_message = self.construct_system_prompt(context)
+
+        # Start with system message
+        prompt = f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>"
+
+        # Add chat history (limit to last `max_history_turns` interactions)
+        for user_msg, assistant_msg in chat_history[-max_history_turns:]:
+            prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{user_msg}<|eot_id|>"
+            prompt += f"<|start_header_id|>assistant<|end_header_id|>\n\n{assistant_msg}<|eot_id|>"
+
+        # Add the current user input
+        prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{user_input}<|eot_id|>"
+        prompt += "<|start_header_id|>assistant<|end_header_id|>\n\n"
+
+        return prompt
+
+    async def chat(
+        self,
+        user_input: str,
+        session_id: Any,
+        max_length: int = 1000
+    ) -> Tuple[str, List[Tuple[str, str]], Dict[str, List[Dict[str, Any]]]]:
+        """Main chat method that coordinates the entire conversation flow."""
+        try:
+            # Ensure session_id is a string
+            if not isinstance(session_id, str):
+                session_id = str(session_id)
+    
+            # Get chat history
+            chat_history_raw = self.conversation_manager.get_history(session_id)
+            chat_history = [
+                (entry['user_input'], entry['response']) for entry in chat_history_raw
+            ]
+    
+            # Search all sources
+            search_results = await self.search_all_sources(user_input)
+    
+            # Build context
+            context = self.build_context(search_results, chat_history_raw)
+    
+            # Construct the prompt
+            prompt = self.construct_prompt(user_input, context, chat_history)
+    
+            # Generate response
+            response = await self.generate_response(prompt, max_length)
+    
+            # Store interaction
+            self.conversation_manager.add_interaction(
+                session_id,
+                user_input,
+                response,
+                {'search_results': search_results}
+            )
+    
+            # Prepare the chat history for Gradio
+            formatted_history = [
+                (entry['user_input'], entry['response']) for entry in self.conversation_manager.get_history(session_id)
+            ]
+    
+            return response, formatted_history, search_results
+    
+        except Exception as e:
+            logger.error(f"Error in chat: {e}")
+            raise
+
+            
     async def search_all_sources(
         self,
         query: str,
@@ -174,48 +259,3 @@ class ChatService:
             logger.error(f"Error generating response: {e}")
             raise
 
-    async def chat(
-        self,
-        user_input: str,
-        session_id: Any,  # Allow any type to ensure flexibility, validate later
-        max_length: int = 1000
-    ) -> Tuple[str, List[Dict[str, Any]]]:
-        """Main chat method that coordinates the entire conversation flow"""
-        try:
-            # Ensure session_id is a string
-            if not isinstance(session_id, str):
-                session_id = str(session_id)
-    
-            # Get chat history
-            chat_history = self.conversation_manager.get_history(session_id)
-    
-            # Search all sources
-            search_results = await self.search_all_sources(user_input)
-    
-            # Build context
-            context = self.build_context(search_results, chat_history)
-    
-            # Create prompt
-            prompt = (
-                f"Context:\n{context}\n\n"
-                f"User: {user_input}\n"
-                "Assistant:"
-            )
-    
-            # Generate response
-            response = await self.generate_response(prompt, max_length)
-    
-            # Store interaction
-            self.conversation_manager.add_interaction(
-                session_id,
-                user_input,
-                response,
-                {'search_results': search_results}
-            )
-    
-            return response, self.conversation_manager.get_history(session_id), search_results
-
-    
-        except Exception as e:
-            logger.error(f"Error in chat: {e}")
-            raise