| import streamlit as st |
| from agent import chatbot, classification_llm |
| from langchain_core.messages import HumanMessage, AIMessage, BaseMessage, SystemMessage |
| import uuid |
| import asyncio |
|
|
| def generate_thread_id(): |
| thread_id= uuid.uuid4() |
| return thread_id |
|
|
| def reset_chat(): |
| thread_id=uuid.uuid4() |
| st.session_state['thread_id']=thread_id |
| add_thread(st.session_state['thread_id']) |
| st.session_state['message_history']=[] |
|
|
| def add_thread(thread_id): |
| if thread_id not in st.session_state['chat_threads']: |
| st.session_state['chat_threads'].append(thread_id) |
| st.session_state['thread_titles'][thread_id]=f"New Chat {len(st.session_state['chat_threads'])}" |
|
|
| def load_conversation(thread_id): |
| try: |
| state= chatbot.get_state(config={'configurable' : {'thread_id': thread_id}}) |
| raw_messages = state.values.get('messages', []) if state else [] |
| return [msg for msg in raw_messages if isinstance(msg, BaseMessage)] |
| except Exception as e: |
| print(f"Error loading conversation for thread {thread_id}: {e}") |
| return [] |
|
|
| def generate_title(query): |
| print("--- Generating Title ---") |
| try: |
| prompt = f"Summarize this query into a very short title (max 5 words): {query}" |
| response = classification_llm.invoke(prompt) |
| title = response.content.strip().strip('"') |
| return title if title else "Chat" |
| except Exception as e: |
| print(f"Error generating title: {e}") |
| return "Chat" |
|
|
| if 'message_history' not in st.session_state: st.session_state['message_history']=[] |
| if 'thread_id' not in st.session_state: st.session_state['thread_id']=generate_thread_id() |
| if 'chat_threads' not in st.session_state: st.session_state['chat_threads']=[] |
| if 'thread_titles' not in st.session_state: st.session_state['thread_titles']={} |
| add_thread(st.session_state['thread_id']) |
|
|
| st.sidebar.title("IIITDMJ Chatbot") |
| if st.sidebar.button("➕ New Chat"): |
| reset_chat() |
| st.rerun() |
| st.sidebar.header("My Conversations") |
| for thread_id in st.session_state['chat_threads'][::-1]: |
| title=st.session_state['thread_titles'].get(thread_id,"Untitled Chat") |
| if st.sidebar.button(title, key=f"thread_{thread_id}", use_container_width=True): |
| st.session_state['thread_id']=thread_id |
| messages= load_conversation(thread_id) |
| temp_messages = [] |
| for msg in messages: |
| if isinstance(msg, SystemMessage): continue |
| role = 'user' if isinstance(msg, HumanMessage) else 'assistant' |
| temp_messages.append({'role': role, 'content': msg.content}) |
| st.session_state['message_history'] = temp_messages |
| st.rerun() |
|
|
| st.title("IIITDMJ College Assistant") |
| st.caption("This bot uses a local vector store and LangGraph to answer your questions.") |
|
|
| for message in st.session_state['message_history']: |
| with st.chat_message(message['role']): |
| if message['role'] == 'assistant': |
| st.markdown(f"<div style='font-size: 15px;'>{message['content']}</div>", unsafe_allow_html=True) |
| else: |
| st.markdown(message['content']) |
|
|
| user_input=st.chat_input("Ask about IIITDMJ...") |
|
|
| if user_input: |
|
|
| CONFIG={'configurable' : {'thread_id': st.session_state['thread_id']}} |
|
|
| st.session_state['message_history'].append({'role':'user','content':user_input}) |
| with st.chat_message('user'): |
| st.markdown(user_input) |
|
|
| with st.chat_message('assistant'): |
| placeholder = st.empty() |
| ai_message_content = "" |
|
|
| try: |
| print(f"\n--- Streaming response for Thread ID: {st.session_state['thread_id']} ---") |
|
|
| async def stream_agent_events(stream_placeholder): |
| local_ai_message_content_streamed = "" |
| local_final_node_output = None |
| local_final_node_name = "" |
|
|
| async for event in chatbot.astream_events( |
| {'messages': [HumanMessage(content=user_input)]}, |
| config=CONFIG, |
| version="v1" |
| ): |
| kind = event["event"] |
| name = event["name"] |
|
|
| if kind == "on_chat_model_stream": |
| if name in ("generate_answer", "generate_synthesized_answer", "handle_chat"): |
| chunk_content = event["data"]["chunk"].content |
| if chunk_content: |
| local_ai_message_content_streamed += chunk_content |
| stream_placeholder.markdown(f"<div style='font-size: 15px;'>{local_ai_message_content_streamed}▌</div>", unsafe_allow_html=True) |
|
|
| if kind == "on_chain_end": |
| if name in ("generate_answer", "generate_synthesized_answer", "handle_chat"): |
| if "output" in event.get("data", {}) and isinstance(event["data"]["output"], dict): |
| local_final_node_output = event["data"]["output"] |
| local_final_node_name = name |
| print(f"--- Captured final output from node: {name} ---") |
| |
| return local_ai_message_content_streamed, local_final_node_output, local_final_node_name |
|
|
|
|
| streamed_content, final_output, final_name = asyncio.run(stream_agent_events(placeholder)) |
|
|
| if not streamed_content and final_output: |
| print(f"--- Using fallback: No stream content captured. Using final output from {final_name}. ---") |
| if "messages" in final_output and final_output["messages"]: |
| ai_message_content = final_output["messages"][-1].content |
| placeholder.markdown(f"<div style='font-size: 15px;'>{ai_message_content}</div>", unsafe_allow_html=True) |
| else: |
| print(f"--- Fallback failed: Final output from {final_name} had unexpected format: {final_output} ---") |
| ai_message_content = "Sorry, I couldn't generate a response (fallback error)." |
| placeholder.markdown(ai_message_content) |
|
|
| elif streamed_content: |
| ai_message_content = streamed_content |
| placeholder.markdown(f"<div style='font-size: 15px;'>{ai_message_content}</div>", unsafe_allow_html=True) |
| else: |
| print("--- Fallback failed: No stream content and no final output captured. ---") |
| ai_message_content = "Sorry, I couldn't generate a response (capture error)." |
| placeholder.markdown(ai_message_content) |
|
|
| except Exception as e: |
| st.error(f"An error occurred: {e}") |
| print(f"ERROR DURING STREAM/FALLBACK: {e}") |
| ai_message_content = "Sorry, I encountered an error during execution." |
| placeholder.markdown(ai_message_content) |
|
|
| if not ai_message_content: |
| ai_message_content = "Sorry, I couldn't generate a response." |
|
|
| st.session_state['message_history'].append({'role':'assistant','content':ai_message_content}) |
|
|
| current_id=st.session_state['thread_id'] |
| current_title=st.session_state['thread_titles'].get(current_id,"New Chat") |
| if current_title.startswith("New Chat") and len(st.session_state['message_history']) <= 2: |
| summarized_title = generate_title(user_input) |
| st.session_state['thread_titles'][current_id] = summarized_title |
| st.rerun() |
