Chat bot with tools

level 2 conversation bot.py

@@ -0,0 +1,132 @@
+import gradio as gr
+from langchain_ollama import ChatOllama
+from langchain_core.tools import tool
+from langgraph.prebuilt import ToolNode, tools_condition
+from langgraph.graph import StateGraph, START, END
+from langgraph.graph.message import MessagesState
+from langgraph.checkpoint.memory import MemorySaver
+from langchain_core.messages import (
+    convert_to_openai_messages,
+    SystemMessage,
+    HumanMessage,
+)
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add a and b.
+
+    Args:
+        a: first int
+        b: second int
+    """
+    return a + b
+
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply a and b.
+
+    Args:
+        a: first int
+        b: second int
+    """
+    return a * b
+
+@tool
+def divide(a: int, b: int) -> int:
+    """Divide a by b.
+
+    Args:
+        a: first int
+        b: second int
+    """
+    return a / b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract b from a.
+
+    Args:
+        a: first int
+        b: second int
+    """
+    return a - b
+
+@tool
+def square(a: int) -> int:
+    """Square a.
+
+    Args:
+        a: first int
+    """
+    return a * a
+
+def create_conversation_graph():
+    """
+    Create a conversational graph with a memory saver.
+    """
+    memory = MemorySaver()
+    tools = [add, multiply, divide, subtract, square]
+
+    llm = ChatOllama(model="qwen2.5:3b", temperature=0.5)
+    llm_with_tools = llm.bind_tools(tools)
+
+    sys_msg = SystemMessage(content="You are a helpful assistant tasked with performing arithmetic on a set of inputs.")
+
+    def assistant(state: MessagesState) -> MessagesState:
+        return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}
+
+    builder = StateGraph(MessagesState)
+
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode(tools))
+
+    builder.add_edge(START, "assistant")
+    builder.add_conditional_edges("assistant", tools_condition)
+    builder.add_edge("tools", "assistant")
+
+    graph = builder.compile(checkpointer=memory)
+    return graph
+
+
+def create_chat_interface():
+    """
+    Create and configure the chat interface with the conversation graph.
+    """
+    graph = create_conversation_graph()
+
+    # Specify a thread id
+    thread_id = "123"
+    config = {"configurable": {"thread_id": thread_id}}
+
+    def chat_with_assistant(message, history):
+        """
+        Chat with the assistant using the conversational graph.
+        """
+        # Create a MessagesState with a HumanMessage
+        messages_state = MessagesState(messages=[HumanMessage(content=message)])
+
+        # Invoke the graph with the properly formatted input
+        response = graph.invoke(messages_state, config)
+
+        for msg in response["messages"]:
+            msg.pretty_print()
+
+        # Extract the last message from the response's messages list
+        ai_message = response["messages"][-1]
+
+        # Return just the content of the AI message
+        return convert_to_openai_messages(ai_message)
+
+
+    demo = gr.ChatInterface(
+        fn=chat_with_assistant,
+        type="messages",
+        title="Conversational Bot",
+        description="Ask anything you want",
+        examples=["Hello", "What is your name?", "What is the weather in Tokyo?"],
+    )
+    return demo
+
+if __name__ == "__main__":
+    demo = create_chat_interface()
+    demo.launch()