test

wandb · Dec 5, 2024 · cc468aa · cc468aa
1 parent ba05dd8
commit cc468aa
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 29 deletions.
diff --git a/weave/integrations/openai/openai_sdk.py b/weave/integrations/openai/openai_sdk.py
@@ -375,28 +375,29 @@ def _openai_stream_options_is_set(inputs: dict) -> bool:
     return wrapper
 
 
-symbol_patchers = [
-    # Patch the Completions.create method
-    SymbolPatcher(
-        lambda: importlib.import_module("openai.resources.chat.completions"),
-        "Completions.create",
-        create_wrapper_sync(name="openai.chat.completions.create"),
-    ),
-    SymbolPatcher(
-        lambda: importlib.import_module("openai.resources.chat.completions"),
-        "AsyncCompletions.create",
-        create_wrapper_async(name="openai.chat.completions.create"),
-    ),
-    SymbolPatcher(
-        lambda: importlib.import_module("openai.resources.beta.chat.completions"),
-        "Completions.parse",
-        create_wrapper_sync(name="openai.beta.chat.completions.parse"),
-    ),
-    SymbolPatcher(
-        lambda: importlib.import_module("openai.resources.beta.chat.completions"),
-        "AsyncCompletions.parse",
-        create_wrapper_async(name="openai.beta.chat.completions.parse"),
-    ),
-]
-
-openai_patcher = MultiPatcher(symbol_patchers)  # type: ignore
+def get_openai_patcher() -> MultiPatcher:
+    symbol_patchers = [
+        # Patch the Completions.create method
+        SymbolPatcher(
+            lambda: importlib.import_module("openai.resources.chat.completions"),
+            "Completions.create",
+            create_wrapper_sync(name="openai.chat.completions.create"),
+        ),
+        SymbolPatcher(
+            lambda: importlib.import_module("openai.resources.chat.completions"),
+            "AsyncCompletions.create",
+            create_wrapper_async(name="openai.chat.completions.create"),
+        ),
+        SymbolPatcher(
+            lambda: importlib.import_module("openai.resources.beta.chat.completions"),
+            "Completions.parse",
+            create_wrapper_sync(name="openai.beta.chat.completions.parse"),
+        ),
+        SymbolPatcher(
+            lambda: importlib.import_module("openai.resources.beta.chat.completions"),
+            "AsyncCompletions.parse",
+            create_wrapper_async(name="openai.beta.chat.completions.parse"),
+        ),
+    ]
+
+    return MultiPatcher(symbol_patchers)  # type: ignore
diff --git a/weave/trace/autopatch.py b/weave/trace/autopatch.py
@@ -20,10 +20,10 @@ def autopatch() -> None:
     from weave.integrations.llamaindex.llamaindex import llamaindex_patcher
     from weave.integrations.mistral import mistral_patcher
     from weave.integrations.notdiamond.tracing import notdiamond_patcher
-    from weave.integrations.openai.openai_sdk import openai_patcher
+    from weave.integrations.openai.openai_sdk import get_openai_patcher
     from weave.integrations.vertexai.vertexai_sdk import vertexai_patcher
 
-    openai_patcher.attempt_patch()
+    get_openai_patcher().attempt_patch()
     mistral_patcher.attempt_patch()
     litellm_patcher.attempt_patch()
     llamaindex_patcher.attempt_patch()
@@ -54,10 +54,10 @@ def reset_autopatch() -> None:
     from weave.integrations.llamaindex.llamaindex import llamaindex_patcher
     from weave.integrations.mistral import mistral_patcher
     from weave.integrations.notdiamond.tracing import notdiamond_patcher
-    from weave.integrations.openai.openai_sdk import openai_patcher
+    from weave.integrations.openai.openai_sdk import get_openai_patcher
     from weave.integrations.vertexai.vertexai_sdk import vertexai_patcher
 
-    openai_patcher.undo_patch()
+    get_openai_patcher().undo_patch()
     mistral_patcher.undo_patch()
     litellm_patcher.undo_patch()
     llamaindex_patcher.undo_patch()