From 34b584f140131fc94a7e27e75732a20fd74c3a4c Mon Sep 17 00:00:00 2001
From: ysolanky <yash@phidata.com>
Date: Tue, 26 Nov 2024 23:50:27 -0500
Subject: [PATCH] update

---
 cookbook/providers/vertexai/README.md | 20 ++++++++++----------
 phi/model/vertexai/gemini.py          |  9 +++------
 2 files changed, 13 insertions(+), 16 deletions(-)

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)