add: valuerror

wandb · Oct 28, 2024 · ee3cf42 · ee3cf42
1 parent c78c476
commit ee3cf42
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/weave/integrations/google_ai_studio/google_ai_studio_sdk.py b/weave/integrations/google_ai_studio/google_ai_studio_sdk.py
@@ -36,6 +36,8 @@ def gemini_on_finish(
         original_model_name = call.inputs["self"].model_name
     elif hasattr(call.inputs["self"], "model"):
         original_model_name = call.inputs["self"].model.model_name
+    else:
+        raise ValueError("Unknown model type")
     model_name = original_model_name.split("/")[-1]
     usage = {model_name: {"requests": 1}}
     summary_update = {"usage": usage}