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

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

@@ -0,0 +1,125 @@
+# 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 langchain_core.chat_history import BaseChatMessageHistory
+
+from frontend.utils.title_summary import chain_title
+
+
+class LocalChatMessageHistory(BaseChatMessageHistory):
+    """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")
+
+        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"]:
+            messages = session["messages"] + [
+                {
+                    "type": "human",
+                    "content": "End of conversation - Create one single title",
+                }
+            ]
+            # Remove the tool calls from conversation
+            messages = [
+                msg
+                for msg in messages
+                if msg["type"] in ("ai", "human") and isinstance(msg["content"], str)
+            ]
+
+            response = chain_title.invoke(messages)
+            title = (
+                response.content.strip()
+                if isinstance(response.content, str)
+                else str(response.content)
+            )
+            session["title"] = title
+            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/frontend/utils/message_editing.py

@@ -0,0 +1,59 @@
+# 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
+
+
+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}"
+        if message_type == "human":
+            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 = st.session_state[button_id]
+        else:
+            st.session_state.user_chats[st.session_state["session_id"]]["messages"][
+                button_idx
+            ]["content"] = st.session_state[button_id]
+
+    @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]

src/frontends/streamlit/frontend/utils/multimodal_utils.py

@@ -0,0 +1,217 @@
+# 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 base64
+from typing import Any
+from urllib.parse import quote
+
+from google.cloud import storage
+
+HELP_MESSAGE_MULTIMODALITY = (
+    "For Gemini models to access the URIs you provide, store them in "
+    "Google Cloud Storage buckets within the same project used by Gemini."
+)
+
+HELP_GCS_CHECKBOX = (
+    "Enabling GCS upload will increase the app observability by avoiding"
+    " forwarding and logging large byte strings within the app."
+)
+
+
+def format_content(content: str | list[dict[str, Any]]) -> str:
+    """Formats content as a string, handling both text and multimedia inputs."""
+    if isinstance(content, str):
+        return content
+    if len(content) == 1 and content[0]["type"] == "text":
+        return content[0]["text"]
+    markdown = """Media:
+"""
+    text = ""
+    for part in content:
+        if part["type"] == "text":
+            text = part["text"]
+        # Local Images:
+        if part["type"] == "image_url":
+            image_url = part["image_url"]["url"]
+            image_markdown = f'<img src="{image_url}" width="100">'
+            markdown = (
+                markdown
+                + f"""
+- {image_markdown}
+"""
+            )
+        if part["type"] == "media":
+            # Local other media
+            if "data" in part:
+                markdown = markdown + f"- Local media: {part['file_name']}\n"
+            # From GCS:
+            if "file_uri" in part:
+                # GCS images
+                if "image" in part["mime_type"]:
+                    image_url = gs_uri_to_https_url(part["file_uri"])
+                    image_markdown = f'<img src="{image_url}" width="100">'
+                    markdown = (
+                        markdown
+                        + f"""
+- {image_markdown}
+"""
+                    )
+                # GCS other media
+                else:
+                    image_url = gs_uri_to_https_url(part["file_uri"])
+                    markdown = (
+                        markdown + f"- Remote media: "
+                        f"[{part['file_uri']}]({image_url})\n"
+                    )
+    markdown = (
+        markdown
+        + f"""
+
+{text}"""
+    )
+    return markdown
+
+
+def get_gcs_blob_mime_type(gcs_uri: str) -> str | None:
+    """Fetches the MIME type (content type) of a Google Cloud Storage blob.
+
+    Args:
+        gcs_uri (str): The GCS URI of the blob in the format "gs://bucket-name/object-name".
+
+    Returns:
+        str: The MIME type of the blob (e.g., "image/jpeg", "text/plain") if found,
+             or None if the blob does not exist or an error occurs.
+    """
+    storage_client = storage.Client()
+
+    try:
+        bucket_name, object_name = gcs_uri.replace("gs://", "").split("/", 1)
+
+        bucket = storage_client.bucket(bucket_name)
+        blob = bucket.blob(object_name)
+        blob.reload()
+        return blob.content_type
+    except Exception as e:
+        print(f"Error retrieving MIME type for {gcs_uri}: {e}")
+        return None  # Indicate failure
+
+
+def get_parts_from_files(
+    upload_gcs_checkbox: bool, uploaded_files: list[Any], gcs_uris: str
+) -> list[dict[str, Any]]:
+    """Processes uploaded files and GCS URIs to create a list of content parts."""
+    parts = []
+    # read from local directly
+    if not upload_gcs_checkbox:
+        for uploaded_file in uploaded_files:
+            im_bytes = uploaded_file.read()
+            if "image" in uploaded_file.type:
+                content = {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:{uploaded_file.type};base64,"
+                        f"{base64.b64encode(im_bytes).decode('utf-8')}"
+                    },
+                    "file_name": uploaded_file.name,
+                }
+            else:
+                content = {
+                    "type": "media",
+                    "data": base64.b64encode(im_bytes).decode("utf-8"),
+                    "file_name": uploaded_file.name,
+                    "mime_type": uploaded_file.type,
+                }
+
+            parts.append(content)
+    if gcs_uris != "":
+        for uri in gcs_uris.split(","):
+            content = {
+                "type": "media",
+                "file_uri": uri,
+                "mime_type": get_gcs_blob_mime_type(uri),
+            }
+            parts.append(content)
+    return parts
+
+
+def upload_bytes_to_gcs(
+    bucket_name: str,
+    blob_name: str,
+    file_bytes: bytes,
+    content_type: str | None = None,
+) -> str:
+    """Uploads a bytes object to Google Cloud Storage and returns the GCS URI.
+
+    Args:
+        bucket_name: The name of the GCS bucket.
+        blob_name: The desired name for the uploaded file in GCS.
+        file_bytes: The file's content as a bytes object.
+        content_type (optional): The MIME type of the file (e.g., "image/png").
+            If not provided, GCS will try to infer it.
+
+    Returns:
+        str: The GCS URI (gs://bucket_name/blob_name) of the uploaded file.
+
+    Raises:
+        GoogleCloudError: If there's an issue with the GCS operation.
+    """
+    storage_client = storage.Client()
+    bucket = storage_client.bucket(bucket_name)
+    blob = bucket.blob(blob_name)
+    blob.upload_from_string(data=file_bytes, content_type=content_type)
+    # Construct and return the GCS URI
+    gcs_uri = f"gs://{bucket_name}/{blob_name}"
+    return gcs_uri
+
+
+def gs_uri_to_https_url(gs_uri: str) -> str:
+    """Converts a GS URI to an HTTPS URL without authentication.
+
+    Args:
+        gs_uri: The GS URI in the format gs://<bucket>/<object>.
+
+    Returns:
+        The corresponding HTTPS URL, or None if the GS URI is invalid.
+    """
+
+    if not gs_uri.startswith("gs://"):
+        raise ValueError("Invalid GS URI format")
+
+    gs_uri = gs_uri[5:]
+
+    # Extract bucket and object names, then URL encode the object name
+    bucket_name, object_name = gs_uri.split("/", 1)
+    object_name = quote(object_name)
+
+    # Construct the HTTPS URL
+    https_url = f"https://storage.mtls.cloud.google.com/{bucket_name}/{object_name}"
+    return https_url
+
+
+def upload_files_to_gcs(st: Any, bucket_name: str, files_to_upload: list[Any]) -> None:
+    """Upload multiple files to Google Cloud Storage and store URIs in session state."""
+    bucket_name = bucket_name.replace("gs://", "")
+    uploaded_uris = []
+    for file in files_to_upload:
+        if file:
+            file_bytes = file.read()
+            gcs_uri = upload_bytes_to_gcs(
+                bucket_name=bucket_name,
+                blob_name=file.name,
+                file_bytes=file_bytes,
+                content_type=file.type,
+            )
+            uploaded_uris.append(gcs_uri)
+    st.session_state.uploader_key += 1
+    st.session_state["gcs_uris_to_be_sent"] = ",".join(uploaded_uris)

src/frontends/streamlit/frontend/utils/stream_handler.py

@@ -0,0 +1,282 @@
+# 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="unreachable"
+import importlib
+import json
+import uuid
+from collections.abc import Generator
+from typing import Any
+from urllib.parse import urljoin
+
+import google.auth
+import google.auth.transport.requests
+import google.oauth2.id_token
+import requests
+import streamlit as st
+import vertexai
+from google.auth.exceptions import DefaultCredentialsError
+from langchain_core.messages import AIMessage, ToolMessage
+from vertexai.preview import reasoning_engines
+
+from frontend.utils.multimodal_utils import format_content
+
+
+@st.cache_resource
+def get_remote_agent(remote_agent_engine_id: str) -> Any:
+    """Get cached remote agent instance."""
+    # Extract location and engine ID from the full resource ID.
+    parts = remote_agent_engine_id.split("/")
+    project_id = parts[1]
+    location = parts[3]
+    vertexai.init(project=project_id, location=location)
+    return reasoning_engines.ReasoningEngine(remote_agent_engine_id)
+
+
+@st.cache_resource
+def get_remote_url_config(url: str, authenticate_request: bool) -> dict[str, Any]:
+    """Get cached remote URL agent configuration."""
+    stream_url = urljoin(url, "stream_messages")
+    creds, _ = google.auth.default()
+    id_token = None
+    if authenticate_request:
+        auth_req = google.auth.transport.requests.Request()
+        try:
+            id_token = google.oauth2.id_token.fetch_id_token(auth_req, stream_url)
+        except DefaultCredentialsError:
+            creds.refresh(auth_req)
+            id_token = creds.id_token
+    return {
+        "url": stream_url,
+        "authenticate_request": authenticate_request,
+        "creds": creds,
+        "id_token": id_token,
+    }
+
+
+class Client:
+    """A client for streaming events from a server."""
+
+    def __init__(
+        self,
+        agent_callable_path: str | None = None,
+        remote_agent_engine_id: str | None = None,
+        url: str | None = None,
+        authenticate_request: bool = False,
+    ) -> None:
+        """Initialize the Client with appropriate configuration.
+
+        Args:
+            agent_callable_path: Path to local agent class
+            remote_agent_engine_id: ID of remote reasoning engine
+            url: URL for remote service
+            authenticate_request: Whether to authenticate requests to remote URL
+        """
+        if url:
+            remote_config = get_remote_url_config(url, authenticate_request)
+            self.url = remote_config["url"]
+            self.authenticate_request = remote_config["authenticate_request"]
+            self.creds = remote_config["creds"]
+            self.id_token = remote_config["id_token"]
+            self.agent = None
+        elif remote_agent_engine_id:
+            self.agent = get_remote_agent(remote_agent_engine_id)
+            self.url = None
+        else:
+            # Force reload all submodules to get latest changes
+            self.url = None
+            if agent_callable_path is None:
+                raise ValueError("agent_callable_path cannot be None")
+            module_path, class_name = agent_callable_path.rsplit(".", 1)
+            module = importlib.import_module(module_path)
+            self.agent = getattr(module, class_name)()
+            self.agent.set_up()
+
+    def log_feedback(self, feedback_dict: dict[str, Any], run_id: str) -> None:
+        """Log user feedback for a specific run."""
+        score = feedback_dict["score"]
+        if score == "😞":
+            score = 0.0
+        elif score == "🙁":
+            score = 0.25
+        elif score == "😐":
+            score = 0.5
+        elif score == "🙂":
+            score = 0.75
+        elif score == "😀":
+            score = 1.0
+        feedback_dict["score"] = score
+        feedback_dict["run_id"] = run_id
+        feedback_dict["log_type"] = "feedback"
+        feedback_dict.pop("type")
+        url = urljoin(self.url, "feedback")
+        headers = {
+            "Content-Type": "application/json",
+        }
+        if self.authenticate_request:
+            headers["Authorization"] = f"Bearer {self.id_token}"
+        requests.post(url, data=json.dumps(feedback_dict), headers=headers, timeout=10)
+
+    def stream_messages(
+        self, data: dict[str, Any]
+    ) -> Generator[dict[str, Any], None, None]:
+        """Stream events from the server, yielding parsed event data."""
+        if self.url:
+            headers = {
+                "Content-Type": "application/json",
+                "Accept": "text/event-stream",
+            }
+            if self.authenticate_request:
+                headers["Authorization"] = f"Bearer {self.id_token}"
+            with requests.post(
+                self.url, json={"input": data}, headers=headers, stream=True, timeout=10
+            ) as response:
+                for line in response.iter_lines():
+                    if line:
+                        try:
+                            event = json.loads(line.decode("utf-8"))
+                            yield event
+                        except json.JSONDecodeError:
+                            print(f"Failed to parse event: {line.decode('utf-8')}")
+        elif self.agent is not None:
+            yield from self.agent.stream_query(input=data)
+
+
+class StreamHandler:
+    """Handles streaming updates to a Streamlit interface."""
+
+    def __init__(self, st: Any, initial_text: str = "") -> None:
+        """Initialize the StreamHandler with Streamlit context and initial text."""
+        self.st = st
+        self.tool_expander = st.expander("Tool Calls:", expanded=False)
+        self.container = st.empty()
+        self.text = initial_text
+        self.tools_logs = initial_text
+
+    def new_token(self, token: str) -> None:
+        """Add a new token to the main text display."""
+        self.text += token
+        self.container.markdown(format_content(self.text), unsafe_allow_html=True)
+
+    def new_status(self, status_update: str) -> None:
+        """Add a new status update to the tool calls expander."""
+        self.tools_logs += status_update
+        self.tool_expander.markdown(status_update)
+
+
+class EventProcessor:
+    """Processes events from the stream and updates the UI accordingly."""
+
+    def __init__(self, st: Any, client: Client, stream_handler: StreamHandler) -> None:
+        """Initialize the EventProcessor with Streamlit context, client, and stream handler."""
+        self.st = st
+        self.client = client
+        self.stream_handler = stream_handler
+        self.final_content = ""
+        self.tool_calls: list[dict[str, Any]] = []
+        self.current_run_id: str | None = None
+        self.additional_kwargs: dict[str, Any] = {}
+
+    def process_events(self, run_id: str | None = None) -> None:
+        """Process events from the stream, handling each event type appropriately."""
+        messages = self.st.session_state.user_chats[
+            self.st.session_state["session_id"]
+        ]["messages"]
+        self.current_run_id = run_id or str(uuid.uuid4())
+        # Set run_id in session state at start of processing
+        self.st.session_state["run_id"] = self.current_run_id
+        stream = self.client.stream_messages(
+            data={
+                "messages": messages,
+                "config": {
+                    "run_id": self.current_run_id,
+                    "metadata": {
+                        "user_id": self.st.session_state["user_id"],
+                        "session_id": self.st.session_state["session_id"],
+                    },
+                },
+            }
+        )
+        # Each event is a tuple message, metadata. https://langchain-ai.github.io/langgraph/how-tos/streaming/#messages
+        for message, _ in stream:
+            if isinstance(message, dict):
+                if message.get("type") == "constructor":
+                    message = message["kwargs"]
+
+                    # Handle tool calls
+                    if message.get("tool_calls"):
+                        tool_calls = message["tool_calls"]
+                        ai_message = AIMessage(content="", tool_calls=tool_calls)
+                        self.tool_calls.append(ai_message.model_dump())
+                        for tool_call in tool_calls:
+                            msg = f"\n\nCalling tool: `{tool_call['name']}` with args: `{tool_call['args']}`"
+                            self.stream_handler.new_status(msg)
+
+                    # Handle tool responses
+                    elif message.get("tool_call_id"):
+                        content = message["content"]
+                        tool_call_id = message["tool_call_id"]
+                        tool_message = ToolMessage(
+                            content=content, type="tool", tool_call_id=tool_call_id
+                        ).model_dump()
+                        self.tool_calls.append(tool_message)
+                        msg = f"\n\nTool response: `{content}`"
+                        self.stream_handler.new_status(msg)
+
+                    # Handle AI responses
+                    elif content := message.get("content"):
+                        self.final_content += content
+                        self.stream_handler.new_token(content)
+
+        # Handle end of stream
+        if self.final_content:
+            final_message = AIMessage(
+                content=self.final_content,
+                id=self.current_run_id,
+                additional_kwargs=self.additional_kwargs,
+            ).model_dump()
+            session = self.st.session_state["session_id"]
+            self.st.session_state.user_chats[session]["messages"] = (
+                self.st.session_state.user_chats[session]["messages"] + self.tool_calls
+            )
+            self.st.session_state.user_chats[session]["messages"].append(final_message)
+            self.st.session_state.run_id = self.current_run_id
+
+
+def get_chain_response(st: Any, client: Client, stream_handler: StreamHandler) -> None:
+    """Process the chain response update the Streamlit UI.
+
+    This function initiates the event processing for a chain of operations,
+    involving an AI model's response generation and potential tool calls.
+    It creates an EventProcessor instance and starts the event processing loop.
+
+    Args:
+        st (Any): The Streamlit app instance, used for accessing session state
+                 and updating the UI.
+        client (Client): An instance of the Client class used to stream events
+                        from the server.
+        stream_handler (StreamHandler): An instance of the StreamHandler class
+                                      used to update the Streamlit UI with
+                                      streaming content.
+
+    Returns:
+        None
+
+    Side effects:
+        - Updates the Streamlit UI with streaming tokens and tool call information.
+        - Modifies the session state to include the final AI message and run ID.
+        - Handles various events like chain starts/ends, tool calls, and model outputs.
+    """
+    processor = EventProcessor(st, client, stream_handler)
+    processor.process_events()