Update Chatbot.py

Chatbot.py

@@ -5,7 +5,6 @@ from typing import Sequence
 from typing_extensions import Annotated, TypedDict
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_core.messages import (
-    SystemMessage,
     HumanMessage,
     AIMessage,
     trim_messages,
@@ -56,7 +55,7 @@ class Model:
             [
                 (
                     "system",
-                    "You are a Medical assistant. Answer Only Medical related questions to the best of your ability in {language}.",
+                    "You are a helpful assistant. Answer all questions to the best of your ability in {language}.",
                 ),
                 MessagesPlaceholder(variable_name="messages"),
             ]
@@ -115,29 +114,12 @@ class Model:
         prompt = self.prompt_template.invoke(
             {"messages": trimmed_messages, "language": state["language"]}
         )
-        formatted_messages = []
-        for message in prompt.messages:
-            if isinstance(message, HumanMessage):
-                formatted_messages.append({
-                    "role": "user",
-                    "content": message.content
-                })
-            elif isinstance(message, AIMessage):
-                formatted_messages.append({
-                    "role": "assistant",
-                    "content": message.content
-                })
-            else:  # System message
-                formatted_messages.append({
-                    "role": "system",
-                    "content": message.content
-                })
         response = self.model.invoke(prompt)
 
         self.chat_message_history.add_user_message(prompt.messages[-1].content)
         self.chat_message_history.add_ai_message(response.content)
 
-        return {"messages": [AIMessage(content=response.content)]}
+        return {"messages": [response]}
 
     def invoke(self, query, config, chat_message_history):
         self.chat_message_history = chat_message_history