simple-chatbot-openai/app.py

57 lines
1.9 KiB
Python
Raw Normal View History

2024-10-30 17:42:30 +07:00
import gradio as gr
from openai import OpenAI
# Initialize the OpenAI client
client = OpenAI(
2024-12-11 07:46:06 +00:00
base_url='https://hub.societyai.com/models/llama-3-2-3b/openai/v1',
2024-10-30 17:42:30 +07:00
)
2024-12-30 12:21:42 +00:00
with gr.Blocks() as demo:
2024-10-30 17:42:30 +07:00
chatbot = gr.Chatbot(type="messages")
msg = gr.Textbox()
clear = gr.Button("Clear")
2024-12-30 12:37:29 +00:00
def chat(user_message, history: list):
"""
1) Appends the user message to the conversation history
2) Sends the conversation history to the vLLM API
3) Streams back the assistant's response
"""
# 1) Append the user message
history = history + [{"role": "user", "content": user_message}]
# 2) Prepare an empty assistant entry for streaming
2024-10-30 17:42:30 +07:00
history.append({"role": "assistant", "content": ""})
2024-12-30 12:37:29 +00:00
# 3) Stream the assistant's response
2024-10-30 17:42:30 +07:00
try:
completion = client.chat.completions.create(
2024-12-30 12:37:29 +00:00
model="llama-3.2-3B-instruct",
2024-10-30 17:42:30 +07:00
messages=history,
stream=True
)
for chunk in completion:
delta = chunk.choices[0].delta
content = getattr(delta, 'content', '')
if content:
history[-1]['content'] += content
2024-12-30 12:37:29 +00:00
# Yield both the cleared textbox ("") and updated history
yield "", history
2024-10-30 17:42:30 +07:00
except Exception as e:
# Handle exceptions and display an error message
history[-1]['content'] += f"\n[Error]: {str(e)}"
2024-12-30 12:37:29 +00:00
yield "", history
2024-10-30 17:42:30 +07:00
2024-12-30 12:37:29 +00:00
# Wire up the single chat function to the UI
msg.submit(
fn=chat, # single consolidated function
inputs=[msg, chatbot], # pass user message and current chatbot history
outputs=[msg, chatbot] # clear the message box and update the chatbot
2024-10-30 17:42:30 +07:00
)
2024-12-30 12:37:29 +00:00
2024-10-30 17:42:30 +07:00
clear.click(lambda: None, None, chatbot, queue=False)
if __name__ == "__main__":
demo.launch()