agent-starter-pack 0.2.3__py3-none-any.whl → 0.3.1__py3-none-any.whl

src/frontends/streamlit_adk/frontend/streamlit_app.py

@@ -0,0 +1,314 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# mypy: disable-error-code="arg-type"
+import json
+import uuid
+from collections.abc import Sequence
+from functools import partial
+from typing import Any
+
+import streamlit as st
+from google.adk.events.event import Event
+from google.genai import types
+from streamlit_feedback import streamlit_feedback
+
+from frontend.side_bar import SideBar
+from frontend.style.app_markdown import MARKDOWN_STR
+from frontend.utils.local_chat_history import LocalChatMessageHistory
+from frontend.utils.message_editing import MessageEditing
+from frontend.utils.multimodal_utils import format_content, get_parts_from_files
+from frontend.utils.stream_handler import Client, StreamHandler, get_chain_response
+
+USER = "my_user"
+EMPTY_CHAT_NAME = "Empty chat"
+
+
+def setup_page() -> None:
+    """Configure the Streamlit page settings."""
+    st.set_page_config(
+        page_title="Playground",
+        layout="wide",
+        initial_sidebar_state="auto",
+        menu_items=None,
+    )
+    st.title("Playground")
+    st.markdown(MARKDOWN_STR, unsafe_allow_html=True)
+
+
+def initialize_session_state() -> None:
+    """Initialize the session state with default values."""
+    if "user_chats" not in st.session_state:
+        st.session_state["session_id"] = str(uuid.uuid4())
+        st.session_state.uploader_key = 0
+        st.session_state.invocation_id = None
+        st.session_state.user_id = USER
+        st.session_state["gcs_uris_to_be_sent"] = ""
+        st.session_state.modified_prompt = None
+        st.session_state.session_db = LocalChatMessageHistory(
+            session_id=st.session_state["session_id"],
+            user_id=st.session_state["user_id"],
+        )
+        st.session_state.user_chats = (
+            st.session_state.session_db.get_all_conversations()
+        )
+        st.session_state.user_chats[st.session_state["session_id"]] = {
+            "title": EMPTY_CHAT_NAME,
+            "messages": [],
+        }
+
+
+def display_messages() -> None:
+    """Display all messages in the current chat session."""
+    messages = st.session_state.user_chats[st.session_state["session_id"]]["messages"]
+    tool_calls_map = {}  # Map tool_call_id to tool call input
+    for i, message in enumerate(messages):
+        # Convert message to Event if it's not already
+        event = message if isinstance(message, Event) else Event.model_validate(message)
+
+        # Check if this is a model message with function calls
+        if hasattr(event.content, "parts") and event.content.parts:
+            for part in event.content.parts:
+                if hasattr(part, "function_call") and part.function_call:
+                    # Store function call info for later matching with responses
+                    tool_calls_map[part.function_call.id] = {
+                        "id": part.function_call.id,
+                        "name": part.function_call.name,
+                        "args": part.function_call.args,
+                    }
+
+        # Check if this is a message with function responses
+        function_responses = []
+        if hasattr(event.content, "parts") and event.content.parts:
+            for part in event.content.parts:
+                if hasattr(part, "function_response") and part.function_response:
+                    function_responses.append(part.function_response)
+
+        # Display function responses if any
+        for function_response in function_responses:
+            tool_call_id = function_response.id
+            if tool_call_id in tool_calls_map:
+                # Display the tool output and remove from map
+                tool_call = tool_calls_map.pop(tool_call_id, None)
+                if tool_call:
+                    display_tool_output(
+                        tool_call,
+                        {
+                            "type": "tool",
+                            "content": function_response.response,
+                            "tool_call_id": tool_call_id,
+                        },
+                    )
+
+        # Display regular chat messages (model or user)
+        if hasattr(event.content, "role") and event.content.role in ["model", "user"]:
+            # Only display if there's text content (skip pure function call messages)
+            has_text_content = False
+            if hasattr(event.content, "parts"):
+                for part in event.content.parts:
+                    if hasattr(part, "text") and part.text:
+                        has_text_content = True
+                        break
+
+            if has_text_content:
+                display_chat_message(message, i)
+
+
+def display_chat_message(message: dict[str, Any], index: int) -> None:
+    """Display a single chat message with edit, refresh, and delete options."""
+    role = "assistant" if message["content"]["role"] == "model" else "user"
+    chat_message = st.chat_message(role)
+    with chat_message:
+        st.markdown(format_content(message["content"]["parts"]), unsafe_allow_html=True)
+        col1, col2, col3 = st.columns([2, 2, 94])
+        display_message_buttons(message, index, col1, col2, col3)
+
+
+def display_message_buttons(
+    message: dict[str, Any], index: int, col1: Any, col2: Any, col3: Any
+) -> None:
+    """Display edit, refresh, and delete buttons for a chat message."""
+    edit_button = f"{index}_edit"
+    refresh_button = f"{index}_refresh"
+    delete_button = f"{index}_delete"
+
+    # Extract content from the message, handling the new event structure
+    content = ""
+    if isinstance(message, dict):
+        if "content" in message:
+            if isinstance(message["content"], dict) and "parts" in message["content"]:
+                parts = message["content"]["parts"]
+                if parts:
+                    if isinstance(parts, list):
+                        for part in parts:
+                            if isinstance(part, dict) and "text" in part:
+                                content += (
+                                    part["text"] if part["text"] is not None else ""
+                                )
+                    elif isinstance(parts, str):
+                        content = parts
+            elif isinstance(message["content"], str):
+                content = message["content"]
+
+    with col1:
+        st.button(label="✎", key=edit_button, type="primary")
+    if message["content"]["role"] == "user":
+        with col2:
+            st.button(
+                label="⟳",
+                key=refresh_button,
+                type="primary",
+                on_click=partial(MessageEditing.refresh_message, st, index, content),
+            )
+        with col3:
+            st.button(
+                label="X",
+                key=delete_button,
+                type="primary",
+                on_click=partial(MessageEditing.delete_message, st, index),
+            )
+
+    if st.session_state[edit_button]:
+        st.text_area(
+            "Edit your message:",
+            value=content,
+            key=f"edit_box_{index}",
+            on_change=partial(
+                MessageEditing.edit_message, st, index, message["content"]["role"]
+            ),
+        )
+
+
+def display_tool_output(
+    tool_call_input: dict[str, Any], tool_call_output: dict[str, Any]
+) -> None:
+    """Display the input and output of a tool call in an expander."""
+    tool_expander = st.expander(label="Tool Calls:", expanded=False)
+    with tool_expander:
+        msg = (
+            f"\n\nEnding tool: `{tool_call_input}` with\n **args:**\n"
+            f"```\n{json.dumps(tool_call_input, indent=2)}\n```\n"
+            f"\n\n**output:**\n "
+            f"```\n{json.dumps(tool_call_output, indent=2)}\n```"
+        )
+        st.markdown(msg, unsafe_allow_html=True)
+
+
+def handle_user_input(side_bar: SideBar) -> None:
+    """Process user input, generate AI response, and update chat history."""
+    prompt = st.chat_input() or st.session_state.modified_prompt
+    if prompt:
+        st.session_state.modified_prompt = None
+        parts = get_parts_from_files(
+            upload_gcs_checkbox=st.session_state.checkbox_state,
+            uploaded_files=side_bar.uploaded_files,
+            gcs_uris=side_bar.gcs_uris,
+        )
+        st.session_state["gcs_uris_to_be_sent"] = ""
+        parts.append(types.Part(text=prompt))
+        st.session_state.user_chats[st.session_state["session_id"]]["messages"].append(
+            Event(
+                content=types.Content(parts=parts, role="user"), author="user"
+            ).model_dump()
+        )
+        display_user_input(parts)
+        generate_ai_response(
+            remote_agent_engine_id=side_bar.remote_agent_engine_id,
+            agent_callable_path=side_bar.agent_callable_path,
+            url=side_bar.url_input_field,
+            authenticate_request=side_bar.should_authenticate_request,
+        )
+        update_chat_title()
+        if len(parts) > 1:
+            st.session_state.uploader_key += 1
+        st.rerun()
+
+
+def display_user_input(parts: Sequence[dict[str, Any]]) -> None:
+    """Display the user's input in the chat interface."""
+    human_message = st.chat_message("human")
+    with human_message:
+        existing_user_input = format_content(parts)
+        st.markdown(existing_user_input, unsafe_allow_html=True)
+
+
+def generate_ai_response(
+    remote_agent_engine_id: str | None = None,
+    agent_callable_path: str | None = None,
+    url: str | None = None,
+    authenticate_request: bool = False,
+) -> None:
+    """Generate and display the AI's response to the user's input."""
+    ai_message = st.chat_message("ai")
+    with ai_message:
+        status = st.status("Generating answer🤖")
+        stream_handler = StreamHandler(st=st)
+        client = Client(
+            remote_agent_engine_id=remote_agent_engine_id,
+            agent_callable_path=agent_callable_path,
+            url=url,
+            authenticate_request=authenticate_request,
+        )
+        get_chain_response(st=st, client=client, stream_handler=stream_handler)
+        status.update(label="Finished!", state="complete", expanded=False)
+
+
+def update_chat_title() -> None:
+    """Update the chat title if it's currently empty."""
+    if (
+        st.session_state.user_chats[st.session_state["session_id"]]["title"]
+        == EMPTY_CHAT_NAME
+    ):
+        st.session_state.session_db.set_title(
+            st.session_state.user_chats[st.session_state["session_id"]]
+        )
+    st.session_state.session_db.upsert_session(
+        st.session_state.user_chats[st.session_state["session_id"]]
+    )
+
+
+def display_feedback(side_bar: SideBar) -> None:
+    """Display a feedback component and log the feedback if provided."""
+    if st.session_state.invocation_id is not None:
+        feedback = streamlit_feedback(
+            feedback_type="faces",
+            optional_text_label="[Optional] Please provide an explanation",
+            key=f"feedback-{st.session_state.invocation_id}",
+        )
+        if feedback is not None:
+            client = Client(
+                remote_agent_engine_id=side_bar.remote_agent_engine_id,
+                agent_callable_path=side_bar.agent_callable_path,
+                url=side_bar.url_input_field,
+                authenticate_request=side_bar.should_authenticate_request,
+            )
+            client.log_feedback(
+                feedback_dict=feedback,
+                invocation_id=st.session_state.invocation_id,
+            )
+
+
+def main() -> None:
+    """Main function to set up and run the Streamlit app."""
+    setup_page()
+    initialize_session_state()
+    side_bar = SideBar(st=st)
+    side_bar.init_side_bar()
+    display_messages()
+    handle_user_input(side_bar=side_bar)
+    display_feedback(side_bar=side_bar)
+
+
+if __name__ == "__main__":
+    main()

src/frontends/streamlit_adk/frontend/style/app_markdown.py

@@ -0,0 +1,37 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+MARKDOWN_STR = """
+<style>
+button[kind="primary"] {
+    background: none!important;
+    border: 0;
+    padding: 20!important;
+    color: grey !important;
+    text-decoration: none;
+    cursor: pointer;
+    border: none !important;
+    # float: right;
+}
+button[kind="primary"]:hover {
+    text-decoration: none;
+    color: white !important;
+}
+button[kind="primary"]:focus {
+    outline: none !important;
+    box-shadow: none !important;
+    color:  !important;
+}
+</style>
+"""

src/frontends/streamlit_adk/frontend/utils/chat_utils.py

@@ -0,0 +1,84 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from pathlib import Path
+from typing import Any
+
+import yaml
+from google.adk.events.event import Event
+
+SAVED_CHAT_PATH = str(os.getcwd()) + "/.saved_chats"
+
+
+def clean_text(text: str) -> str:
+    """Preprocess the input text by removing leading and trailing newlines."""
+    if not text:
+        return text
+
+    if text.startswith("\n"):
+        text = text[1:]
+    if text.endswith("\n"):
+        text = text[:-1]
+    return text
+
+
+def sanitize_messages(
+    messages: list[Event | dict],
+) -> list[Event | dict]:
+    """Preprocess and fix the content of messages."""
+    for message in messages:
+        # Handle Event objects
+        if isinstance(message, Event):
+            if hasattr(message.content, "parts"):
+                for part in message.content.parts:
+                    if hasattr(part, "text") and part.text:
+                        part.text = clean_text(part.text)
+        # Handle dictionary format
+        elif isinstance(message, dict):
+            if "content" in message:
+                if (
+                    isinstance(message["content"], dict)
+                    and "parts" in message["content"]
+                ):
+                    for part in message["content"]["parts"]:
+                        if isinstance(part, dict) and "text" in part:
+                            part["text"] = clean_text(part["text"])
+                elif isinstance(message["content"], str):
+                    message["content"] = clean_text(message["content"])
+                elif isinstance(message["content"], list):
+                    for part in message["content"]:
+                        if part.get("type") == "text":
+                            part["text"] = clean_text(part["text"])
+    return messages
+
+
+def save_chat(st: Any) -> None:
+    """Save the current chat session to a YAML file."""
+    Path(SAVED_CHAT_PATH).mkdir(parents=True, exist_ok=True)
+    session_id = st.session_state["session_id"]
+    session = st.session_state.user_chats[session_id]
+    messages = session.get("messages", [])
+    if len(messages) > 0:
+        session["messages"] = sanitize_messages(session["messages"])
+        filename = f"{session_id}.yaml"
+        with open(Path(SAVED_CHAT_PATH) / filename, "w") as file:
+            yaml.dump(
+                [session],
+                file,
+                allow_unicode=True,
+                default_flow_style=False,
+                encoding="utf-8",
+            )
+        st.toast(f"Chat saved to path: ↓ {Path(SAVED_CHAT_PATH) / filename}")

src/frontends/streamlit_adk/frontend/utils/local_chat_history.py

@@ -0,0 +1,110 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from datetime import datetime
+
+import yaml
+
+from frontend.utils.title_summary import DummySummarizer, TitleGenerator
+
+
+class LocalChatMessageHistory:
+    """Manages local storage and retrieval of chat message history."""
+
+    def __init__(
+        self,
+        user_id: str,
+        session_id: str = "default",
+        base_dir: str = ".streamlit_chats",
+    ) -> None:
+        self.user_id = user_id
+        self.session_id = session_id
+        self.base_dir = base_dir
+        self.user_dir = os.path.join(self.base_dir, self.user_id)
+        self.session_file = os.path.join(self.user_dir, f"{session_id}.yaml")
+        try:
+            self.title_generator = TitleGenerator()
+        except:  # noqa: E722
+            self.title_generator: TitleGenerator = DummySummarizer()  # type: ignore
+        os.makedirs(self.user_dir, exist_ok=True)
+
+    def get_session(self, session_id: str) -> None:
+        """Updates the session ID and file path for the current session."""
+        self.session_id = session_id
+        self.session_file = os.path.join(self.user_dir, f"{session_id}.yaml")
+
+    def get_all_conversations(self) -> dict[str, dict]:
+        """Retrieves all conversations for the current user."""
+        conversations = {}
+        for filename in os.listdir(self.user_dir):
+            if filename.endswith(".yaml"):
+                file_path = os.path.join(self.user_dir, filename)
+                with open(file_path) as f:
+                    conversation = yaml.safe_load(f)
+                    if not isinstance(conversation, list) or len(conversation) > 1:
+                        raise ValueError(
+                            f"""Invalid format in {file_path}.
+                        YAML file can only contain one conversation with the following
+                        structure.
+                          - messages:
+                              - content: [message text]
+                              - type: (human or ai)"""
+                        )
+                    conversation = conversation[0]
+                    if "title" not in conversation:
+                        conversation["title"] = filename
+                conversations[filename[:-5]] = conversation
+        return dict(
+            sorted(conversations.items(), key=lambda x: x[1].get("update_time", ""))
+        )
+
+    def upsert_session(self, session: dict) -> None:
+        """Updates or inserts a session into the local storage."""
+        session["update_time"] = datetime.now().isoformat()
+        with open(self.session_file, "w") as f:
+            yaml.dump(
+                [session],
+                f,
+                allow_unicode=True,
+                default_flow_style=False,
+                encoding="utf-8",
+            )
+
+    def set_title(self, session: dict) -> None:
+        """
+        Set the title for the given session.
+
+        This method generates a title for the session based on its messages.
+        If the session has messages, it appends a special message to prompt
+        for title creation, generates the title using a title chain, and
+        updates the session with the new title.
+
+        Args:
+            session (dict): A dictionary containing session information,
+                            including messages.
+
+        Returns:
+            None
+        """
+        if session["messages"]:
+            session["title"] = self.title_generator.summarize(
+                events=session["messages"]
+            )
+            self.upsert_session(session)
+
+    def clear(self) -> None:
+        """Removes the current session file if it exists."""
+        if os.path.exists(self.session_file):
+            os.remove(self.session_file)

src/frontends/streamlit_adk/frontend/utils/message_editing.py

@@ -0,0 +1,61 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# fmt: off
+
+from typing import Any
+
+from google.adk.events.event import Event
+
+
+class MessageEditing:
+    """Provides methods for editing, refreshing, and deleting chat messages."""
+
+    @staticmethod
+    def edit_message(st: Any, button_idx: int, message_type: str) -> None:
+        """Edit a message in the chat history."""
+        button_id = f"edit_box_{button_idx}"
+        # Handle Event type messages
+        message = st.session_state.user_chats[st.session_state["session_id"]]["messages"][button_idx]
+        # Convert to Event if it's not already
+        event = message if isinstance(message, Event) else Event.model_validate(message)
+        # Update the text content in the event
+        if hasattr(event.content, 'parts'):
+            for part in event.content.parts:
+                if hasattr(part, 'text'):
+                    part.text = st.session_state[button_id]
+                    break
+        # Update the message in the session state
+        st.session_state.user_chats[st.session_state["session_id"]]["messages"][button_idx] = event.model_dump()
+
+    @staticmethod
+    def refresh_message(st: Any, button_idx: int, content: str) -> None:
+        """Refresh a message in the chat history."""
+        messages = st.session_state.user_chats[st.session_state["session_id"]][
+            "messages"
+        ]
+        st.session_state.user_chats[st.session_state["session_id"]][
+            "messages"
+        ] = messages[:button_idx]
+        st.session_state.modified_prompt = content
+
+    @staticmethod
+    def delete_message(st: Any, button_idx: int) -> None:
+        """Delete a message from the chat history."""
+        messages = st.session_state.user_chats[st.session_state["session_id"]][
+            "messages"
+        ]
+        st.session_state.user_chats[st.session_state["session_id"]][
+            "messages"
+        ] = messages[:button_idx]