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

src/deployment_targets/agent_engine/app/agent_engine_app.py

@@ -26,16 +26,12 @@ import vertexai
 from google.cloud import logging as google_cloud_logging
 from langchain_core.runnables import RunnableConfig
 from traceloop.sdk import Instruments, Traceloop
-from vertexai.preview import reasoning_engines
+from vertexai import agent_engines
 
 from app.utils.gcs import create_bucket_if_not_exists
 from app.utils.tracing import CloudTraceLoggingSpanExporter
 from app.utils.typing import Feedback, InputChat, dumpd, ensure_valid_config
 
-logging.basicConfig(
-    level=logging.INFO,
-)
-
 
 class AgentEngineApp:
     """Class for managing agent engine functionality."""
@@ -124,6 +120,8 @@ class AgentEngineApp:
         **kwargs: Any,
     ) -> Any:
         """Process a single input and return the agent's response."""
+        config = ensure_valid_config(config)
+        self.set_tracing_properties(config=config)
         return dumpd(self.runnable.invoke(input=input, config=config, **kwargs))
 
     def register_operations(self) -> Mapping[str, Sequence]:
@@ -151,7 +149,7 @@ def deploy_agent_engine_app(
     requirements_file: str = ".requirements.txt",
     extra_packages: list[str] = ["./app"],
     env_vars: dict[str, str] | None = None,
-) -> reasoning_engines.ReasoningEngine:
+) -> agent_engines.AgentEngine:
     """Deploy the agent engine app to Vertex AI."""
 
     staging_bucket = f"gs://{project}-agent-engine"
@@ -169,7 +167,7 @@ def deploy_agent_engine_app(
 
     # Common configuration for both create and update operations
     agent_config = {
-        "reasoning_engine": agent,
+        "agent_engine": agent,
         "display_name": agent_name,
         "description": "This is a sample custom application in Agent Engine that uses LangGraph",
         "extra_packages": extra_packages,
@@ -178,10 +176,7 @@ def deploy_agent_engine_app(
     agent_config["requirements"] = requirements
 
     # Check if an agent with this name already exists
-    existing_agents = reasoning_engines.ReasoningEngine.list(
-        filter=f"display_name={agent_name}"
-    )
-
+    existing_agents = list(agent_engines.list(filter=f"display_name={agent_name}"))
     if existing_agents:
         # Update the existing agent with new configuration
         logging.info(f"Updating existing agent: {agent_name}")
@@ -189,7 +184,7 @@ def deploy_agent_engine_app(
     else:
         # Create a new agent if none exists
         logging.info(f"Creating new agent: {agent_name}")
-        remote_agent = reasoning_engines.ReasoningEngine.create(**agent_config)
+        remote_agent = agent_engines.create(**agent_config)
 
     config = {
         "remote_agent_engine_id": remote_agent.resource_name,

src/deployment_targets/cloud_run/Dockerfile

@@ -14,7 +14,7 @@
 
 FROM python:3.11-slim
 
-RUN pip install --no-cache-dir uv
+RUN pip install --no-cache-dir uv==0.6.12
 
 WORKDIR /code
 

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

@@ -28,7 +28,7 @@ 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 vertexai import agent_engines
 
 from frontend.utils.multimodal_utils import format_content
 
@@ -43,7 +43,7 @@ def get_remote_agent(remote_agent_engine_id: str) -> Any:
     project_id = parts[1]
     location = parts[3]
     vertexai.init(project=project_id, location=location)
-    return reasoning_engines.ReasoningEngine(remote_agent_engine_id)
+    return agent_engines.AgentEngine(remote_agent_engine_id)
 
 
 @st.cache_resource
@@ -159,7 +159,7 @@ class Client:
             if self.authenticate_request:
                 headers["Authorization"] = f"Bearer {self.id_token}"
             with requests.post(
-                self.url, json=data, headers=headers, stream=True, timeout=10
+                self.url, json=data, headers=headers, stream=True, timeout=60
             ) as response:
                 for line in response.iter_lines():
                     if line: