Google検索のGroundingとThinking modelを使ったStreamlitチャット

https://youtu.be/ZTENmcmgwy4

GOOGLE_API_KEY = ""

import streamlit as st
import streamlit.components.v1 as components
from google import genai
from google.genai import types

client = genai.Client(api_key=GOOGLE_API_KEY, http_options={"api_version": "v1alpha"})

st.title("Chat with Gemini")

if "chat" not in st.session_state:
    st.session_state.chat = client.chats.create(
        model="gemini-2.0-flash-thinking-exp",
    )
chat = st.session_state.chat

if "messages" not in st.session_state:
    st.session_state.messages = []
messages = st.session_state.messages

for message in messages:
    if "grounding_metadata" == message["role"]:
        components.html(message["content"], scrolling=True)
    else:
        with st.chat_message(message["role"]):
            st.write(message["content"])

if prompt := st.chat_input():

    messages.append({"role": "user", "content": prompt})

    with st.chat_message("user"):
        st.write(prompt)

    with st.chat_message("assistant"):

        response = chat.send_message_stream(
            prompt,
            config=types.GenerateContentConfig(
                tools=[types.Tool(google_search=types.GoogleSearchRetrieval)]
            ),
        )

        message_write = st.empty()
        text = ""
        rendered_content = None

        for chunk in response:
            text = text + chunk.text
            message_write.write(text)

            if chunk.candidates[0].grounding_metadata:
                grounding_metadata = chunk.candidates[0].grounding_metadata

                if grounding_metadata.grounding_chunks:
                    for chunk in grounding_metadata.grounding_chunks:
                        if chunk.web:
                            st.link_button(chunk.web.title, chunk.web.uri)

                if grounding_metadata.search_entry_point:
                    if grounding_metadata.search_entry_point.rendered_content:
                        rendered_content = (
                            grounding_metadata.search_entry_point.rendered_content
                        )

                        components.html(rendered_content, scrolling=True)

        messages.append({"role": "assistant", "content": text})
        if rendered_content:
            messages.append({"role": "grounding_metadata", "content": rendered_content})