File size: 2,231 Bytes
6369fd4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
import gradio as gr
from langchain_ollama import ChatOllama
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,
)
def create_conversation_graph():
"""
Create a conversational graph with a memory saver.
"""
memory = MemorySaver()
llm = ChatOllama(model="gemma3:4b", temperature=0)
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.invoke([sys_msg] + state["messages"])]}
builder = StateGraph(MessagesState)
builder.add_node("assistant", assistant)
builder.add_edge(START, "assistant")
builder.add_edge("assistant", END)
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)
# 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()
|