Spaces:
Running
Running
| import streamlit as st | |
| from google import genai | |
| from google.genai import types | |
| from gemini_tools_defination import GeminiToolDefination | |
| from gemini_llm import GeminiLLM | |
| from gemini_history import ConversationBuilder | |
| from gemini_tools import GeminiCodeExecutionTool, GeminiThinking, GoogleSearchTool | |
| from utils import get_typing_indicator_html | |
| from gemini_tool_selector import GeminiToolSelector | |
| geminillm=GeminiLLM() | |
| geminitooldefination = GeminiToolDefination(geminillm) | |
| conversationbuilder = ConversationBuilder() | |
| geminitoolselector = GeminiToolSelector() | |
| client = geminillm.client | |
| if "chat" not in st.session_state: | |
| st.session_state.chat = client.chats.create(model=geminillm.model_id, config=geminitooldefination.config) | |
| if "messages" not in st.session_state: | |
| st.session_state.messages = [] | |
| # default message | |
| if not st.session_state.messages: | |
| with st.chat_message("assistant", avatar="src/images/ai.png"): | |
| st.markdown("👋 **Hi! I can help you think, code, or search real-time info. Just ask!**") | |
| # Tool name map | |
| tool_mapping = { | |
| "Auto": None, | |
| "Think": "GeminiThinking", | |
| "Search": "GoogleSearchTool", | |
| "Code": "GeminiCodeExecutionTool" | |
| } | |
| # display chat messages from history at every rerun | |
| for message in st.session_state.messages: | |
| avatar_path = ( | |
| "src/images/person_15454011.png" if message["role"] == "user" | |
| else "src/images/ai.png" | |
| ) | |
| with st.chat_message(message["role"], avatar=avatar_path): | |
| st.markdown(message["content"]) | |
| if "selected_tool" not in st.session_state: | |
| st.session_state.selected_tool = "Auto" | |
| st.sidebar.title("Tools") | |
| selected_tool = st.sidebar.selectbox( | |
| "Choose a tool:", | |
| options=list(tool_mapping.keys()), | |
| index=list(tool_mapping.keys()).index(st.session_state.selected_tool), | |
| ) | |
| # Update the session state with current selection | |
| st.session_state.selected_tool = selected_tool | |
| prompt = st.chat_input("What is up?") | |
| selected_tool_name = tool_mapping[st.session_state.selected_tool] | |
| if prompt: | |
| try: | |
| # show user message | |
| with st.chat_message("user", avatar="src/images/person_15454011.png"): | |
| st.markdown(prompt) | |
| # Save user message | |
| st.session_state.messages.append({"role": "user", "content": prompt}) | |
| # stream assistant response | |
| full_response = "" | |
| with st.chat_message("assistant", avatar="src/images/ai.png"): | |
| response_stream = st.session_state.chat.send_message_stream(prompt) | |
| response_placeholder = st.empty() | |
| if selected_tool_name: | |
| tool_args = {"query": prompt} | |
| full_response = geminitoolselector.handle_tool_response(selected_tool_name, tool_args, geminillm) | |
| # need to add response in chat history | |
| content_only_convo = conversationbuilder.build_with_content_only(full_response) | |
| st.session_state.chat.get_history().extend(content_only_convo) | |
| else: | |
| for chunk in response_stream: | |
| if chunk.text is None: # if response is None means a tool is hit | |
| tool_name = chunk.candidates[0].content.parts[0].function_call.name | |
| tool_args = chunk.candidates[0].content.parts[0].function_call.args | |
| response_placeholder = st.empty() | |
| # handle tool response | |
| streamed_output = geminitoolselector.handle_tool_response(tool_name, tool_args, geminillm) | |
| full_response = streamed_output | |
| # IMP..... add the response in the chathistory can see by st.session_state.chat.get_history() see complete history you will see the content is missing there and u can also more then one conttent | |
| content_only_convo = conversationbuilder.build_with_content_only(full_response) | |
| st.session_state.chat.get_history().extend(content_only_convo) | |
| # st.markdown(st.session_state.chat.get_history()) ## for reference how chat history is working | |
| else: | |
| full_response += chunk.text | |
| response_placeholder.markdown(full_response + "") # stream effect | |
| # save assistant message (load all the messages to show chat type view in UI ) | |
| st.session_state.messages.append({"role": "assistant", "content": full_response}) | |
| # print(st.session_state.chat.get_history()) ## for reference how chat history is working | |
| except Exception as e: | |
| st.error(f"An error occurred: {e}") | |