update cohere client v2 chat

wandb · Aug 20, 2024 · bbbfd92 · bbbfd92
1 parent 4e78195
commit bbbfd92
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/weave/integrations/cohere/cohere_sdk.py b/weave/integrations/cohere/cohere_sdk.py
@@ -72,5 +72,17 @@ def wrapper(fn: typing.Callable) -> typing.Callable:
             "AsyncClient.chat_stream",
             cohere_stream_wrapper("cohere.AsyncClient.chat_stream"),
         ),
+        # Add patch for cohere v2
+        SymbolPatcher(
+            lambda: importlib.import_module("cohere"),
+            "ClientV2.chat",
+            cohere_wrapper("cohere.ClientV2.chat"),
+        ),
+        # Add patch for cohre v2 async chat method
+        SymbolPatcher(
+            lambda: importlib.import_module("cohere"),
+            "AsyncClientV2.chat",
+            cohere_wrapper("cohere.AsyncClientV2.chat"),
+        ),
     ]
 )