diff --git a/cookbook/providers/vertexai/README.md b/cookbook/providers/vertexai/README.md index 06674b896..ed7cb9328 100644 --- a/cookbook/providers/vertexai/README.md +++ b/cookbook/providers/vertexai/README.md @@ -24,13 +24,13 @@ pip install -U google-cloud-aiplatform duckduckgo-search yfinance phidata - Streaming on ```shell -python cookbook/providers/google/basic_stream.py +python cookbook/providers/vertexai/basic_stream.py ``` - Streaming off ```shell -python cookbook/providers/google/basic.py +python cookbook/providers/vertexai/basic.py ``` ### 5. Run Agent with Tools @@ -38,47 +38,47 @@ python cookbook/providers/google/basic.py - Yahoo Finance with streaming on ```shell -python cookbook/providers/google/agent_stream.py +python cookbook/providers/vertexai/agent_stream.py ``` - Yahoo Finance without streaming ```shell -python cookbook/providers/google/agent.py +python cookbook/providers/vertexai/agent.py ``` - Finance Agent ```shell -python cookbook/providers/google/finance_agent.py +python cookbook/providers/vertexai/finance_agent.py ``` - Web Search Agent ```shell -python cookbook/providers/google/web_search.py +python cookbook/providers/vertexai/web_search.py ``` - Data Analysis Agent ```shell -python cookbook/providers/google/data_analyst.py +python cookbook/providers/vertexai/data_analyst.py ``` ### 6. Run Agent that returns structured output ```shell -python cookbook/providers/google/structured_output.py +python cookbook/providers/vertexai/structured_output.py ``` ### 7. Run Agent that uses storage ```shell -python cookbook/providers/google/storage.py +python cookbook/providers/vertexai/storage.py ``` ### 8. Run Agent that uses knowledge ```shell -python cookbook/providers/google/knowledge.py +python cookbook/providers/vertexai/knowledge.py ``` diff --git a/phi/model/vertexai/gemini.py b/phi/model/vertexai/gemini.py index 5755d6d72..3e010e3fb 100644 --- a/phi/model/vertexai/gemini.py +++ b/phi/model/vertexai/gemini.py @@ -21,9 +21,6 @@ Content, Part, ) - from google.cloud.aiplatform_v1beta1.types.prediction_service import GenerateContentResponse - - UsageMetadata = GenerateContentResponse.UsageMetadata except ImportError: logger.error("`google-cloud-aiplatform` not installed") @@ -38,7 +35,7 @@ class MessageData: response_role: Optional[str] = None response_parts: Optional[List] = None response_tool_calls: List[Dict[str, Any]] = field(default_factory=list) - response_usage: Optional[UsageMetadata] = None + response_usage: Optional[Dict[str, Any]] = None response_tool_call_block = None @@ -271,7 +268,7 @@ def _update_usage_metrics( self, assistant_message: Message, metrics: Metrics, - usage: Optional[UsageMetadata] = None, + usage: Optional[Dict[str, Any]] = None, ) -> None: """ Update usage metrics for the assistant message. @@ -279,7 +276,7 @@ def _update_usage_metrics( Args: assistant_message: Message object containing the response content metrics: Metrics object containing the usage metrics - usage: UsageMetadata object containing the usage metrics + usage: Dict[str, Any object containing the usage metrics """ assistant_message.metrics["time"] = metrics.response_timer.elapsed self.metrics.setdefault("response_times", []).append(metrics.response_timer.elapsed)