Spaces:
Sleeping
Sleeping
File size: 2,980 Bytes
cbdc2c4 d1b9872 d6e8585 6fc6581 234eac0 d1b9872 d6e8585 719dfe3 28f3481 234eac0 d6e8585 cbdc2c4 d6e8585 cbdc2c4 d6e8585 d1b9872 d6e8585 d1b9872 234eac0 6fc6581 d6e8585 cbdc2c4 d6e8585 d1b9872 d6e8585 d1b9872 d6e8585 d1b9872 d6e8585 d1b9872 d6e8585 cbdc2c4 28f3481 d6e8585 |
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 77 78 79 80 81 82 |
import uuid
import os
import json
from langchain_core.messages import HumanMessage, AIMessage, AIMessageChunk
from langchain.schema.runnable.config import RunnableConfig
import chainlit as cl
import sys
from pathlib import Path
# Add the project root to the Python path
project_root = str(Path(__file__).parent.parent.parent)
if project_root not in sys.path:
sys.path.append(project_root)
from agent import create_agent_graph, get_checkpointer
from agent.utils.state import AgentState
SHORT_TERM_MEMORY_DB_PATH = "data/short_term.db"
os.makedirs(os.path.dirname(SHORT_TERM_MEMORY_DB_PATH), exist_ok=True)
@cl.on_chat_start
async def on_chat_start():
"""Initialize the chat session"""
try:
# Generate and store a session ID
session_id = str(uuid.uuid4())
cl.user_session.set("session_id", session_id)
welcome_message = cl.Message(
content="Hello! I'm your chief joy officer, here to help you with finding fun things to do in London!",
author="Assistant"
)
await welcome_message.send()
except Exception as e:
print(f"Error in chat initialization: {str(e)}")
await cl.Message(
content="I apologize, but I encountered an error during initialization. Please try refreshing the page.",
author="System"
).send()
@cl.on_message
async def on_message(message: cl.Message):
"""Handle incoming messages and stream responses"""
# Initialize response message
msg = cl.Message(content="")
try:
async with cl.Step(type="run"):
async with get_checkpointer(SHORT_TERM_MEMORY_DB_PATH) as saver:
# Create graph with memory
graph = await create_agent_graph(saver)
# Get session ID
session_id = cl.user_session.get("session_id")
# Process through graph with current message
async for chunk in graph.astream(
{"messages": [HumanMessage(content=message.content)]},
{"configurable": {"thread_id": session_id}},
stream_mode="messages"
):
if chunk[1]["langgraph_node"] == "agent" and isinstance(
chunk[0], (AIMessageChunk, AIMessage)
):
await msg.stream_token(chunk[0].content)
# Get final state
final_state = await graph.aget_state(
config={"configurable": {"thread_id": session_id}}
)
# Send the final message
await msg.send()
except Exception as e:
print(f"Error in message handler: {str(e)}")
await cl.Message(
content="I apologize, but I encountered an error processing your message. Please try again.",
author="System"
).send() |