import os import getpass import gradio as gr from typing import Sequence from typing_extensions import Annotated, TypedDict from langchain_google_genai import ChatGoogleGenerativeAI from langchain_core.messages import ( HumanMessage, AIMessage, trim_messages, BaseMessage, ) from langgraph.checkpoint.memory import MemorySaver from langgraph.graph import START, StateGraph from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory from langgraph.graph.message import add_messages # Define the State Schema """ A state schema is like a blueprint that defines how we organize and store information in our chatbot. It helps the chatbot keep track of important details, such as the messages exchanged and other constraints. Messages: This is a list holding all messages that have been sent and received during the conversation. One can think of each message as a piece of the chat history. Language: This is a simple text field that stores the user's preferred language for communication. This is to ensure that the chatbot answers back in the correct language. """ class State(TypedDict): messages: Annotated[Sequence[BaseMessage], add_messages] language: str # Configure the AI Model """ initialize chatbot's core functionalities in App """ class Model: def __init__(self): self.model = ChatGoogleGenerativeAI( model="gemini-1.5-flash", temperature=0, max_tokens=None, timeout=None, max_retries=2, google_api_key = os.getenv("GOOGLE_API_KEY") ) # Define Prompt Template """ The chatbot uses a prompt template to define its behavior """ self.prompt_template = ChatPromptTemplate.from_messages( [ ( "system", "You are a Medical Assistant. Answer all the only medical questions {language}.", ), MessagesPlaceholder(variable_name="messages"), ] ) # Implement the Workflow """ 1. self.workflow = StateGraph(state_schema=State) This line creates a new workflow called workflow using a `StateGraph`. A StateGraph is a structure that helps in managing the flow of data and actions in a program. The `state_schema=State` part means that the workflow will use a specific format or structure for the data it will handle, defined by `State`. 2. self.workflow.add_edge(START, "model") Here, an "edge" which is similar to a connection or a path is added to the workflow. 3. self.workflow.add_node("model", self.call_model) This line introduces a new "node" (which is a specific task or function) into the workflow. 4. self.memory = MemorySaver() This line establishes a memory storage system called memory with the MemorySaver. 5. self.app = self.workflow.compile(checkpointer=self.memory) this line compiles the workflow into a runnable application called app. checkpointer=self.memory part means that the workflow will use the MemorySaver to keep track of its state while it runs. """ self.workflow = StateGraph(state_schema=State) self.workflow.add_edge(START, "model") self.workflow.add_node("model", self.call_model) self.memory = MemorySaver() self.app = self.workflow.compile(checkpointer=self.memory) # Setup Trimmer """ `self.trimmer` to handle message histories """ self.trimmer = trim_messages( max_tokens=65, strategy="last", token_counter=self.model, include_system=True, allow_partial=False, start_on="human", ) #Call the AI Model """ `call_model` function processes inputs, generates responses, and updates memory """ def call_model(self, state: State): if not state["messages"] or len(state["messages"]) == 1: state["messages"] = self.chat_message_history.messages + state["messages"] print(state["messages"]) trimmed_messages = self.trimmer.invoke(state["messages"]) prompt = self.prompt_template.invoke( {"messages": trimmed_messages, "language": state["language"]} ) 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": [response]} def invoke(self, query, config, chat_message_history): self.chat_message_history = chat_message_history input_messages = [HumanMessage(query)] response = self.app.invoke( {"messages": input_messages, "language": "en"}, config ) return response["messages"][-1].content # Create Chat Interface Class class ChatInterface: def __init__(self): self.model = Model() self.active_users = {} def get_chat_history(self, user_id): return MongoDBChatMessageHistory( session_id=user_id, connection_string=os.getenv("Mongo_URI"), database_name="Chatbot", collection_name="chat_histories", ) def chat(self, message, history, user_id, language="English"): try: if not user_id: return "", history, "Please enter a user ID first." config = {"configurable": {"thread_id": user_id}} chat_history = self.get_chat_history(user_id) response = self.model.invoke(message, config, chat_history) return "", history + [(message, response)], "" except Exception as e: return "", history, f"Error: {str(e)}" def reset_chat(self, user_id): try: if not user_id: return None, "Please enter a user ID first" chat_history = self.get_chat_history(user_id) chat_history.clear() return None, f"Chat history cleared for user {user_id}" except Exception as e: return None, f"Error clearing chat: {str(e)}" # Create the Gradio Interface def create_gradio_interface(): chat_interface = ChatInterface() custom_css = """ .gradio-container { max-width: 800px !important; margin: auto !important; padding: 20px !important; background-color: #f7f7f7 !important; } .chat-message { padding: 15px !important; border-radius: 10px !important; margin: 5px !important; } .message-wrap { display: flex !important; flex-direction: column !important; gap: 10px !important; } """ with gr.Blocks(css="footer {visibility: hidden}") as demo: with gr.Column(elem_id="chat-container"): gr.Markdown( """ # 🤖 CareBot - Your Personal Medical AI Assistant Welcome to your personal Medical AI assistant! Enter your ID and start chatting. """ ) with gr.Row(): with gr.Column(scale=3): user_id = gr.Textbox( label="🆔 User ID", placeholder="Enter Your User ID", elem_id="user-id" ) with gr.Column(scale=1): language = gr.Dropdown( choices=["English", "Spanish", "French", "German", "Chinese"], value="English", label="🌐 Language", elem_id="language" ) chatbot = gr.Chatbot( label="💬 Chat History", height=500, bubble_full_width=False, show_label=True, elem_id="chatbot", type = "messages" ) with gr.Row(): msg = gr.Textbox( label="✍️ Message", placeholder="Type your message here...", show_label=True, elem_id="message" ) submit = gr.Button("🚀 Send", variant="primary") clear = gr.Button("🗑️ Clear Chat", variant="secondary") status_msg = gr.Markdown("") # Event handlers submit.click( fn=chat_interface.chat, inputs=[msg, chatbot, user_id, language], outputs=[msg, chatbot, status_msg] ) msg.submit( fn=chat_interface.chat, inputs=[msg, chatbot, user_id, language], outputs=[msg, chatbot, status_msg] ) clear.click( fn=chat_interface.reset_chat, inputs=[user_id], outputs=[chatbot, status_msg] ) return demo if __name__ == "__main__": demo = create_gradio_interface() demo.launch(share=True)