Spaces:
Running
Running
File size: 3,134 Bytes
2aa4095 |
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 |
import streamlit as st
from google import genai
from google.genai import types
from gemini_tools import GeminiCodeExecutionTool, GeminiThinking, GoogleSearchTool
from utils import get_typing_indicator_html
class GeminiToolSelector():
def __init__(self) -> None:
pass
def handle_tool_response(self, tool_name, tool_args, geminillm):
if tool_name == "GeminiThinking":
gen = GeminiThinking(geminillm).activate_stream(tool_args['query'])
typing_placeholder = st.empty()
thought_placeholder = st.empty()
response_placeholder = st.empty()
thought_buffer = ""
streamed_output = ""
typing_placeholder.markdown(get_typing_indicator_html("Thinking"), unsafe_allow_html=True)
for chunk in gen:
if chunk["thought"]:
thought_buffer += chunk["text"] + "\n"
with thought_placeholder.expander("Thoughts (click to expand)", expanded=False):
st.markdown(thought_buffer)
else:
break
typing_placeholder.empty()
if chunk and not chunk["thought"]:
streamed_output += chunk["text"]
response_placeholder.markdown(streamed_output)
for chunk in gen:
if not chunk["thought"]:
streamed_output += chunk["text"]
response_placeholder.markdown(streamed_output)
return streamed_output
elif tool_name == "GoogleSearchTool":
gen = GoogleSearchTool(geminillm).activate_stream(tool_args['query'])
typing_placeholder = st.empty()
typing_placeholder.markdown(get_typing_indicator_html("Searching Web"), unsafe_allow_html=True)
response_placeholder = st.empty()
try:
first_chunk = next(gen)
except StopIteration:
first_chunk = ""
typing_placeholder.empty()
streamed_output = first_chunk
response_placeholder.markdown(streamed_output)
for chunk in gen:
streamed_output += chunk
response_placeholder.markdown(streamed_output)
return streamed_output
elif tool_name == "GeminiCodeExecutionTool":
gen = GeminiCodeExecutionTool(geminillm).activate_stream(tool_args['query'])
typing_placeholder = st.empty()
typing_placeholder.markdown(get_typing_indicator_html("Writing Code"), unsafe_allow_html=True)
response_placeholder = st.empty()
try:
first_chunk = next(gen)
except StopIteration:
first_chunk = ""
typing_placeholder.empty()
streamed_output = first_chunk
response_placeholder.markdown(streamed_output)
for chunk in gen:
streamed_output += chunk
response_placeholder.markdown(streamed_output)
return streamed_output
else:
return "Unknown tool triggered."
|