From 5aa6bdaed813e31c8f9ba5128454307ed02cc43e Mon Sep 17 00:00:00 2001 From: Vitali Khvatkov <1123272+Kvit@users.noreply.github.com> Date: Tue, 21 Jan 2025 11:47:43 -0600 Subject: [PATCH 1/3] Create ws_streaming_llamaindex.py Streaming chat with Lllamaindes RAG tool powered by OpenAI agent --- 02_chatbot/ws_streaming_llamaindex.py | 84 +++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) create mode 100644 02_chatbot/ws_streaming_llamaindex.py diff --git a/02_chatbot/ws_streaming_llamaindex.py b/02_chatbot/ws_streaming_llamaindex.py new file mode 100644 index 0000000..7f845f9 --- /dev/null +++ b/02_chatbot/ws_streaming_llamaindex.py @@ -0,0 +1,84 @@ +# Websocket steaming chat with Llamaindex Open AI Agent +## Requirements: +## llama-index-agent-openai + +from fasthtml.common import * +import asyncio +from llama_index.llms.openai import OpenAI +from llama_index.agent.openai import OpenAIAgent +import os + +# check if api key is in the environment variables +if 'OPENAI_API_KEY' not in os.environ: + raise ValueError("OPENAI_API_KEY not found in the environment variables") + +# Set up the app, including daisyui and tailwind for the chat component +tlink = Script(src="https://cdn.tailwindcss.com"), +dlink = Link(rel="stylesheet", href="https://cdn.jsdelivr.net/npm/daisyui@4.11.1/dist/full.min.css") +app = FastHTML(hdrs=(tlink, dlink, picolink), ws_hdr=True) + +# Initialize Llamaindex RAG OpenAI agent +llm = OpenAI(model="gpt-4o") +openai_agent = OpenAIAgent.from_tools(llm=llm) + +messages = [] + +# Chat message component (renders a chat bubble) +# Now with a unique ID for the content and the message +def ChatMessage(msg_idx, **kwargs): + msg = messages[msg_idx] + bubble_class = "chat-bubble-primary" if msg['role']=='user' else 'chat-bubble-secondary' + chat_class = "chat-end" if msg['role']=='user' else 'chat-start' + return Div(Div(msg['role'], cls="chat-header"), + Div(msg['content'], + id=f"chat-content-{msg_idx}", # Target if updating the content + cls=f"chat-bubble {bubble_class}"), + id=f"chat-message-{msg_idx}", # Target if replacing the whole message + cls=f"chat {chat_class}", **kwargs) + +# The input field for the user message. Also used to clear the +# input field after sending a message via an OOB swap +def ChatInput(): + return Input(type="text", name='msg', id='msg-input', + placeholder="Type a message", + cls="input input-bordered w-full", hx_swap_oob='true') + +# The main screen +@app.route("/") +def get(): + page = Body(H1('Chatbot Demo'), + Div(*[ChatMessage(msg) for msg in messages], + id="chatlist", cls="chat-box h-[73vh] overflow-y-auto"), + Form(Group(ChatInput(), Button("Send", cls="btn btn-primary")), + ws_send="", hx_ext="ws", ws_connect="/wscon", + cls="flex space-x-2 mt-2", + ), + cls="p-4 max-w-lg mx-auto", + ) # Open a websocket connection on page load + return Title('Chatbot Demo'), page + + +@app.ws('/wscon') +async def ws(msg:str, send): + messages.append({"role":"user", "content":msg}) + + # Send the user message to the user (updates the UI right away) + await send(Div(ChatMessage(len(messages)-1), hx_swap_oob='beforeend', id="chatlist")) + + # Send the clear input field command to the user + await send(ChatInput()) + + # Model response Async Streaming + response_stream = await openai_agent.astream_chat(msg) + + # Send an empty message with the assistant response + messages.append({"role":"assistant", "content":""}) + await send(Div(ChatMessage(len(messages)-1), hx_swap_oob='beforeend', id="chatlist")) + + # Fill in the message content + async for chunk in response_stream.async_response_gen(): + print(chunk_message, end='', flush=True) # test streaming with local print + messages[-1]["content"] += chunk_message + await send(Span(chunk_message, hx_swap_oob=swap, id=f"chat-content-{len(messages)-1}")) + +if __name__ == '__main__': uvicorn.run("ws_streaming:app", host='0.0.0.0', port=8000, reload=True) From 87a389a829a1220e7bfbb7752bf7d389d5bebdb2 Mon Sep 17 00:00:00 2001 From: Vitali Khvatkov <1123272+Kvit@users.noreply.github.com> Date: Tue, 21 Jan 2025 18:07:30 +0000 Subject: [PATCH 2/3] chunk message --- 02_chatbot/ws_streaming_llamaindex.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/02_chatbot/ws_streaming_llamaindex.py b/02_chatbot/ws_streaming_llamaindex.py index 7f845f9..f7faf64 100644 --- a/02_chatbot/ws_streaming_llamaindex.py +++ b/02_chatbot/ws_streaming_llamaindex.py @@ -76,7 +76,7 @@ async def ws(msg:str, send): await send(Div(ChatMessage(len(messages)-1), hx_swap_oob='beforeend', id="chatlist")) # Fill in the message content - async for chunk in response_stream.async_response_gen(): + async for chunk_message in response_stream.async_response_gen(): print(chunk_message, end='', flush=True) # test streaming with local print messages[-1]["content"] += chunk_message await send(Span(chunk_message, hx_swap_oob=swap, id=f"chat-content-{len(messages)-1}")) From d412e406edc9eebb9e8ca6014b954da16facc70d Mon Sep 17 00:00:00 2001 From: Vitali Khvatkov Date: Tue, 21 Jan 2025 19:02:55 +0000 Subject: [PATCH 3/3] updated llamaindex streaming example --- 02_chatbot/ws_streaming_llamaindex.py | 73 +++++++++++++++++---------- 1 file changed, 45 insertions(+), 28 deletions(-) diff --git a/02_chatbot/ws_streaming_llamaindex.py b/02_chatbot/ws_streaming_llamaindex.py index f7faf64..863eae4 100644 --- a/02_chatbot/ws_streaming_llamaindex.py +++ b/02_chatbot/ws_streaming_llamaindex.py @@ -3,25 +3,30 @@ ## llama-index-agent-openai from fasthtml.common import * -import asyncio from llama_index.llms.openai import OpenAI -from llama_index.agent.openai import OpenAIAgent +from llama_index.agent.openai import OpenAIAgent import os # check if api key is in the environment variables if 'OPENAI_API_KEY' not in os.environ: raise ValueError("OPENAI_API_KEY not found in the environment variables") + +# Set up the app, including daisyui and tailwind for the chat component +headers = ( + Script(src="https://cdn.tailwindcss.com"), + Link(rel="stylesheet", href="https://cdn.jsdelivr.net/npm/daisyui@4.11.1/dist/full.min.css"), +) + # Set up the app, including daisyui and tailwind for the chat component -tlink = Script(src="https://cdn.tailwindcss.com"), -dlink = Link(rel="stylesheet", href="https://cdn.jsdelivr.net/npm/daisyui@4.11.1/dist/full.min.css") -app = FastHTML(hdrs=(tlink, dlink, picolink), ws_hdr=True) +app = FastHTML(hdrs=headers, exts='ws', pico = False, debug=True) + # Initialize Llamaindex RAG OpenAI agent llm = OpenAI(model="gpt-4o") openai_agent = OpenAIAgent.from_tools(llm=llm) -messages = [] +messages=[] # Chat message component (renders a chat bubble) # Now with a unique ID for the content and the message @@ -30,9 +35,12 @@ def ChatMessage(msg_idx, **kwargs): bubble_class = "chat-bubble-primary" if msg['role']=='user' else 'chat-bubble-secondary' chat_class = "chat-end" if msg['role']=='user' else 'chat-start' return Div(Div(msg['role'], cls="chat-header"), - Div(msg['content'], - id=f"chat-content-{msg_idx}", # Target if updating the content - cls=f"chat-bubble {bubble_class}"), + + Div(msg['content'], # TODO: support markdown + id=f"chat-content-{msg_idx}", # Target if updating the content + cls=f"chat-bubble {bubble_class}" + ), + id=f"chat-message-{msg_idx}", # Target if replacing the whole message cls=f"chat {chat_class}", **kwargs) @@ -46,39 +54,48 @@ def ChatInput(): # The main screen @app.route("/") def get(): - page = Body(H1('Chatbot Demo'), - Div(*[ChatMessage(msg) for msg in messages], + page = Body( + # Chat messages + Div(*[ChatMessage(msg_idx) for msg_idx, msg in enumerate(messages)], id="chatlist", cls="chat-box h-[73vh] overflow-y-auto"), + + # Input field and send button Form(Group(ChatInput(), Button("Send", cls="btn btn-primary")), - ws_send="", hx_ext="ws", ws_connect="/wscon", - cls="flex space-x-2 mt-2", - ), - cls="p-4 max-w-lg mx-auto", - ) # Open a websocket connection on page load - return Title('Chatbot Demo'), page + ws_send=True, hx_ext="ws", ws_connect="/wscon", + cls="flex space-x-2 mt-2"), + + cls="p-4 max-w-lg mx-auto") + + return Titled("Chatbot Demo", page ) @app.ws('/wscon') async def ws(msg:str, send): - messages.append({"role":"user", "content":msg}) + + # add user message to messages list + messages.append({"role":"user", "content":msg.rstrip()}) + swap = 'beforeend' + # Send the user message to the user (updates the UI right away) - await send(Div(ChatMessage(len(messages)-1), hx_swap_oob='beforeend', id="chatlist")) + await send(Div(ChatMessage(len(messages)-1), hx_swap_oob=swap, id="chatlist")) # Send the clear input field command to the user - await send(ChatInput()) - - # Model response Async Streaming - response_stream = await openai_agent.astream_chat(msg) - + await send(ChatInput()) + # Send an empty message with the assistant response messages.append({"role":"assistant", "content":""}) - await send(Div(ChatMessage(len(messages)-1), hx_swap_oob='beforeend', id="chatlist")) + await send(Div(ChatMessage(len(messages)-1), hx_swap_oob=swap, id="chatlist")) + + # Start OpenAI agent async streaming chat + response_stream = await openai_agent.astream_chat(msg) - # Fill in the message content + # Get and process async streaming chat response chunks async for chunk_message in response_stream.async_response_gen(): - print(chunk_message, end='', flush=True) # test streaming with local print + print(chunk_message, end='', flush=True) # Check streaming response via print messages[-1]["content"] += chunk_message await send(Span(chunk_message, hx_swap_oob=swap, id=f"chat-content-{len(messages)-1}")) -if __name__ == '__main__': uvicorn.run("ws_streaming:app", host='0.0.0.0', port=8000, reload=True) + +if __name__ == '__main__': + serve() \ No newline at end of file