agent-starter-pack 0.0.1b0__py3-none-any.whl

src/frontends/streamlit/frontend/side_bar.py

@@ -0,0 +1,213 @@
+# 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.
+
+# ruff: noqa: RUF015
+import json
+import os
+import uuid
+from typing import Any
+
+from frontend.utils.chat_utils import save_chat
+from frontend.utils.multimodal_utils import (
+    HELP_GCS_CHECKBOX,
+    HELP_MESSAGE_MULTIMODALITY,
+    upload_files_to_gcs,
+)
+
+EMPTY_CHAT_NAME = "Empty chat"
+NUM_CHAT_IN_RECENT = 3
+DEFAULT_BASE_URL = "http://localhost:8000/"
+
+DEFAULT_REMOTE_AGENT_ENGINE_ID = "N/A"
+if os.path.exists("deployment_metadata.json"):
+    with open("deployment_metadata.json") as f:
+        DEFAULT_REMOTE_AGENT_ENGINE_ID = json.load(f)["remote_agent_engine_id"]
+DEFAULT_AGENT_CALLABLE_PATH = "app.agent_engine_app.AgentEngineApp"
+
+
+class SideBar:
+    """Manages the sidebar components of the Streamlit application."""
+
+    def __init__(self, st: Any) -> None:
+        """
+        Initialize the SideBar.
+
+        Args:
+            st (Any): The Streamlit object for rendering UI components.
+        """
+        self.st = st
+
+    def init_side_bar(self) -> None:
+        """Initialize and render the sidebar components."""
+        with self.st.sidebar:
+            default_agent_type = (
+                "Remote URL" if os.path.exists("Dockerfile") else "Local Agent"
+            )
+            use_agent_path = self.st.selectbox(
+                "Select Agent Type",
+                ["Local Agent", "Remote Engine ID", "Remote URL"],
+                index=["Local Agent", "Remote Engine ID", "Remote URL"].index(
+                    default_agent_type
+                ),
+            )
+
+            if use_agent_path == "Local Agent":
+                self.agent_callable_path = self.st.text_input(
+                    label="Agent Callable Path",
+                    value=os.environ.get(
+                        "AGENT_CALLABLE_PATH", DEFAULT_AGENT_CALLABLE_PATH
+                    ),
+                )
+                self.remote_agent_engine_id = None
+                self.url_input_field = None
+                self.should_authenticate_request = False
+            elif use_agent_path == "Remote Engine ID":
+                self.remote_agent_engine_id = self.st.text_input(
+                    label="Remote Agent Engine ID",
+                    value=os.environ.get(
+                        "REMOTE_AGENT_ENGINE_ID", DEFAULT_REMOTE_AGENT_ENGINE_ID
+                    ),
+                )
+                self.agent_callable_path = None
+                self.url_input_field = None
+                self.should_authenticate_request = False
+            else:
+                self.url_input_field = self.st.text_input(
+                    label="Service URL",
+                    value=os.environ.get("SERVICE_URL", DEFAULT_BASE_URL),
+                )
+                self.should_authenticate_request = self.st.checkbox(
+                    label="Authenticate request",
+                    value=False,
+                    help="If checked, any request to the server will contain an"
+                    "Identity token to allow authentication. "
+                    "See the Cloud Run documentation to know more about authentication:"
+                    "https://cloud.google.com/run/docs/authenticating/service-to-service",
+                )
+                self.agent_callable_path = None
+                self.remote_agent_engine_id = None
+
+            col1, col2, col3 = self.st.columns(3)
+            with col1:
+                if self.st.button("+ New chat"):
+                    if (
+                        len(
+                            self.st.session_state.user_chats[
+                                self.st.session_state["session_id"]
+                            ]["messages"]
+                        )
+                        > 0
+                    ):
+                        self.st.session_state.run_id = None
+
+                        self.st.session_state["session_id"] = str(uuid.uuid4())
+                        self.st.session_state.session_db.get_session(
+                            session_id=self.st.session_state["session_id"],
+                        )
+                        self.st.session_state.user_chats[
+                            self.st.session_state["session_id"]
+                        ] = {
+                            "title": EMPTY_CHAT_NAME,
+                            "messages": [],
+                        }
+
+            with col2:
+                if self.st.button("Delete chat"):
+                    self.st.session_state.run_id = None
+                    self.st.session_state.session_db.clear()
+                    self.st.session_state.user_chats.pop(
+                        self.st.session_state["session_id"]
+                    )
+                    if len(self.st.session_state.user_chats) > 0:
+                        chat_id = list(self.st.session_state.user_chats.keys())[0]
+                        self.st.session_state["session_id"] = chat_id
+                        self.st.session_state.session_db.get_session(
+                            session_id=self.st.session_state["session_id"],
+                        )
+                    else:
+                        self.st.session_state["session_id"] = str(uuid.uuid4())
+                        self.st.session_state.user_chats[
+                            self.st.session_state["session_id"]
+                        ] = {
+                            "title": EMPTY_CHAT_NAME,
+                            "messages": [],
+                        }
+            with col3:
+                if self.st.button("Save chat"):
+                    save_chat(self.st)
+
+            self.st.subheader("Recent")  # Style the heading
+
+            all_chats = list(reversed(self.st.session_state.user_chats.items()))
+            for chat_id, chat in all_chats[:NUM_CHAT_IN_RECENT]:
+                if self.st.button(chat["title"], key=chat_id):
+                    self.st.session_state.run_id = None
+                    self.st.session_state["session_id"] = chat_id
+                    self.st.session_state.session_db.get_session(
+                        session_id=self.st.session_state["session_id"],
+                    )
+
+            with self.st.expander("Other chats"):
+                for chat_id, chat in all_chats[NUM_CHAT_IN_RECENT:]:
+                    if self.st.button(chat["title"], key=chat_id):
+                        self.st.session_state.run_id = None
+                        self.st.session_state["session_id"] = chat_id
+                        self.st.session_state.session_db.get_session(
+                            session_id=self.st.session_state["session_id"],
+                        )
+
+            self.st.divider()
+            self.st.header("Upload files from local")
+            bucket_name = self.st.text_input(
+                label="GCS Bucket for upload",
+                value=os.environ.get("BUCKET_NAME", "gs://your-bucket-name"),
+            )
+            if "checkbox_state" not in self.st.session_state:
+                self.st.session_state.checkbox_state = True
+
+            self.st.session_state.checkbox_state = self.st.checkbox(
+                "Upload to GCS first (suggested)", value=False, help=HELP_GCS_CHECKBOX
+            )
+
+            self.uploaded_files = self.st.file_uploader(
+                label="Send files from local",
+                accept_multiple_files=True,
+                key=f"uploader_images_{self.st.session_state.uploader_key}",
+                type=[
+                    "png",
+                    "jpg",
+                    "jpeg",
+                    "txt",
+                    "docx",
+                    "pdf",
+                    "rtf",
+                    "csv",
+                    "tsv",
+                    "xlsx",
+                ],
+            )
+            if self.uploaded_files and self.st.session_state.checkbox_state:
+                upload_files_to_gcs(self.st, bucket_name, self.uploaded_files)
+
+            self.st.divider()
+
+            self.st.header("Upload files from GCS")
+            self.gcs_uris = self.st.text_area(
+                "GCS uris (comma-separated)",
+                value=self.st.session_state["gcs_uris_to_be_sent"],
+                key=f"upload_text_area_{self.st.session_state.uploader_key}",
+                help=HELP_MESSAGE_MULTIMODALITY,
+            )
+
+            self.st.caption(f"Note: {HELP_MESSAGE_MULTIMODALITY}")

src/frontends/streamlit/frontend/streamlit_app.py

@@ -0,0 +1,263 @@
+# 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 langchain_core.messages import HumanMessage
+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.run_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):
+        if message["type"] in ["ai", "human"] and message["content"]:
+            display_chat_message(message, i)
+        elif message.get("tool_calls"):
+            # Store each tool call input mapped by its ID
+            for tool_call in message["tool_calls"]:
+                tool_calls_map[tool_call["id"]] = tool_call
+        elif message["type"] == "tool":
+            # Look up the corresponding tool call input by ID
+            tool_call_id = message["tool_call_id"]
+            if tool_call_id in tool_calls_map:
+                display_tool_output(tool_calls_map[tool_call_id], message)
+            else:
+                st.error(f"Could not find tool call input for ID: {tool_call_id}")
+        else:
+            st.error(f"Unexpected message type: {message['type']}")
+            st.write("Full messages list:", messages)
+            raise ValueError(f"Unexpected message type: {message['type']}")
+
+
+def display_chat_message(message: dict[str, Any], index: int) -> None:
+    """Display a single chat message with edit, refresh, and delete options."""
+    chat_message = st.chat_message(message["type"])
+    with chat_message:
+        st.markdown(format_content(message["content"]), 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"
+    content = (
+        message["content"]
+        if isinstance(message["content"], str)
+        else message["content"][-1]["text"]
+    )
+
+    with col1:
+        st.button(label="✎", key=edit_button, type="primary")
+    if message["type"] == "human":
+        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["type"]),
+        )
+
+
+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({"type": "text", "text": prompt})
+        st.session_state.user_chats[st.session_state["session_id"]]["messages"].append(
+            HumanMessage(content=parts).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.run_id is not None:
+        feedback = streamlit_feedback(
+            feedback_type="faces",
+            optional_text_label="[Optional] Please provide an explanation",
+            key=f"feedback-{st.session_state.run_id}",
+        )
+        if feedback is not None:
+            client = Client(
+                agent_callable_path=side_bar.agent_callable_path,
+                remote_agent_engine_id=side_bar.remote_agent_engine_id,
+            )
+            client.log_feedback(
+                feedback_dict=feedback,
+                run_id=st.session_state.run_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/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/frontend/utils/chat_utils.py

@@ -0,0 +1,67 @@
+# 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
+
+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[dict[str, str | list[dict[str, str]]]],
+) -> list[dict[str, str | list[dict[str, str]]]]:
+    """Preprocess and fix the content of messages."""
+    for message in messages:
+        if isinstance(message["content"], list):
+            for part in message["content"]:
+                if part["type"] == "text":
+                    part["text"] = clean_text(part["text"])
+        else:
+            message["content"] = clean_text(message["content"])
+    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}")