105 lines
2.8 KiB
Python
105 lines
2.8 KiB
Python
import time
|
|
import uuid
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from queue import Queue
|
|
|
|
import gradio as gr
|
|
from gradio import ChatMessage
|
|
|
|
from flow import create_flow
|
|
|
|
# create global thread pool
|
|
chatflow_thread_pool = ThreadPoolExecutor(
|
|
max_workers=5,
|
|
thread_name_prefix="chatflow_worker",
|
|
)
|
|
|
|
|
|
def chat_fn(message, history, uuid):
|
|
"""
|
|
Main chat function that handles the conversation flow and message processing.
|
|
|
|
Args:
|
|
message (str): The current user message
|
|
history (list): Previous conversation history
|
|
uuid (UUID): Unique identifier for the conversation
|
|
|
|
Yields:
|
|
ChatMessage: Streams of thought process and chat responses
|
|
"""
|
|
# Log conversation details
|
|
print(f"Conversation ID: {str(uuid)}\nHistory: {history}\nQuery: {message}\n---")
|
|
|
|
# Initialize queues for chat messages and flow thoughts
|
|
chat_queue = Queue()
|
|
flow_queue = Queue()
|
|
|
|
# Create shared context for the flow
|
|
shared = {
|
|
"conversation_id": str(uuid),
|
|
"query": message,
|
|
"history": history,
|
|
"queue": chat_queue,
|
|
"flow_queue": flow_queue,
|
|
}
|
|
|
|
# Create and run the chat flow in a separate thread
|
|
chat_flow = create_flow()
|
|
chatflow_thread_pool.submit(chat_flow.run, shared)
|
|
|
|
# Initialize thought response tracking
|
|
start_time = time.time()
|
|
thought_response = ChatMessage(
|
|
content="", metadata={"title": "Flow Log", "id": 0, "status": "pending"}
|
|
)
|
|
yield thought_response
|
|
|
|
# Process and accumulate thoughts from the flow queue
|
|
accumulated_thoughts = ""
|
|
while True:
|
|
thought = flow_queue.get()
|
|
if thought is None:
|
|
break
|
|
accumulated_thoughts += f"- {thought}\n\n"
|
|
thought_response.content = accumulated_thoughts.strip()
|
|
yield thought_response
|
|
flow_queue.task_done()
|
|
|
|
# Mark thought processing as complete and record duration
|
|
thought_response.metadata["status"] = "done"
|
|
thought_response.metadata["duration"] = time.time() - start_time
|
|
yield thought_response
|
|
|
|
# Process and yield chat messages from the chat queue
|
|
while True:
|
|
msg = chat_queue.get()
|
|
if msg is None:
|
|
break
|
|
chat_response = [thought_response, ChatMessage(content=msg)]
|
|
yield chat_response
|
|
chat_queue.task_done()
|
|
|
|
|
|
def clear_fn():
|
|
print("Clearing conversation")
|
|
return uuid.uuid4()
|
|
|
|
|
|
with gr.Blocks(fill_height=True, theme="ocean") as demo:
|
|
uuid_state = gr.State(uuid.uuid4())
|
|
demo.load(clear_fn, outputs=[uuid_state])
|
|
|
|
chatbot = gr.Chatbot(type="messages", scale=1)
|
|
chatbot.clear(clear_fn, outputs=[uuid_state])
|
|
|
|
gr.ChatInterface(
|
|
fn=chat_fn,
|
|
type="messages",
|
|
additional_inputs=[uuid_state],
|
|
chatbot=chatbot,
|
|
title="PocketFlow Gradio Demo",
|
|
)
|
|
|
|
|
|
demo.launch()
|