Spaces:
Running
Running
File size: 5,801 Bytes
b026171 |
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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
# chat_column.py
import streamlit as st
# Assuming BASE_PROMPT is imported or defined elsewhere if not passed explicitly
# from prompt import BASE_PROMPT # Or pass it as an argument
def render_chat_column(st, llm_client, model_option, max_tokens, BASE_PROMPT):
"""Renders the chat history, input, and LLM prompt generation column."""
st.header("💬 Chat & Prompt Generation")
# --- Display Chat History ---
# (This part remains the same)
for message in st.session_state.messages:
avatar = '🤖' if message["role"] == "assistant" else '🦔'
with st.chat_message(message["role"], avatar=avatar):
st.markdown(message["content"])
# --- Chat Input and LLM Call ---
if prompt := st.chat_input("Enter topic to generate image prompt..."):
if len(prompt.strip()) == 0:
st.warning("Please enter a topic.", icon="⚠️")
elif len(prompt) > 4000: # Example length limit
st.error("Input is too long (max 4000 chars).", icon="🚨")
else:
# Add user message to history and display FIRST
# It's important to add the user message *before* sending it to the API
st.session_state.messages.append(
{"role": "user", "content": prompt})
with st.chat_message("user", avatar='🦔'):
st.markdown(prompt)
# Generate and display assistant response
try:
with st.chat_message("assistant", avatar="🤖"):
response_placeholder = st.empty()
response_placeholder.markdown("Generating prompt... ▌")
full_response = ""
# --- MODIFICATION START ---
# Construct messages for API including the conversation history
# 1. Start with the system prompt
messages_for_api = [
{"role": "system", "content": BASE_PROMPT}]
# 2. Add all messages from the session state (history)
# This now includes the user message we just added above.
messages_for_api.extend(st.session_state.messages)
# 3. Filter out any potential empty messages (just in case)
# This step might be less critical now but is good practice.
messages_for_api = [
m for m in messages_for_api if m.get("content")]
# --- MODIFICATION END ---
stream_kwargs = {
"model": model_option,
"messages": messages_for_api, # <--- Now contains history!
"max_tokens": max_tokens,
"stream": True,
}
# Assuming llm_client is correctly initialized (OpenAI or Cerebras)
response_stream = llm_client.chat.completions.create(
**stream_kwargs)
# --- (Rest of the streaming and response handling code remains the same) ---
for chunk in response_stream:
chunk_content = ""
try:
if chunk.choices and chunk.choices[0].delta:
chunk_content = chunk.choices[0].delta.content or ""
except (AttributeError, IndexError):
chunk_content = "" # Handle potential errors gracefully
if chunk_content:
full_response += chunk_content
response_placeholder.markdown(full_response + "▌")
# Final response display
response_placeholder.markdown(full_response)
# Add assistant response to history
# Check if the last message isn't already the assistant's response to avoid duplicates if rerun happens unexpectedly
if not st.session_state.messages or st.session_state.messages[-1]['role'] != 'assistant':
st.session_state.messages.append(
{"role": "assistant", "content": full_response})
elif st.session_state.messages[-1]['role'] == 'assistant':
# If last message is assistant, update it (useful if streaming was interrupted/retried)
st.session_state.messages[-1]['content'] = full_response
# No longer updating image prompt text area here (based on previous request)
# Rerun might still cause subtle issues with message duplication if not handled carefully,
# The check above helps mitigate this. Consider removing rerun if it causes problems.
# st.rerun() # Keeping rerun commented out for now based on potential issues
except Exception as e:
st.error(
f"Error during LLM response generation: {str(e)}", icon="🚨")
# Clean up potentially failed message
# Ensure we only pop if the *last* message is the user's (meaning the assistant failed)
if st.session_state.messages and st.session_state.messages[-1]["role"] == "user":
# Maybe add a placeholder error message for the assistant instead of popping user?
# For now, let's not pop the user's message. The error message itself indicates failure.
pass
# Or if the assistant message was partially added:
elif st.session_state.messages and st.session_state.messages[-1]["role"] == "assistant" and not full_response:
st.session_state.messages.pop()
|