update

phidatahq · Nov 27, 2024 · 34b584f · 34b584f
1 parent 17c84c0
commit 34b584f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 16 deletions.
diff --git a/cookbook/providers/vertexai/README.md b/cookbook/providers/vertexai/README.md
@@ -24,61 +24,61 @@ 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
 
 - 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
@@ -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,15 +268,15 @@ 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.
 
         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)