chore(weave): add weave.finish

weave/api.py

@@ -137,6 +137,8 @@ def from_pandas(df):
 
 #### Newer API below
 
+_current_inited_client = None
+
 
 def init(project_name: str) -> _weave_client.WeaveClient:
     """Initialize weave tracking, logging to a wandb project.
@@ -153,11 +155,13 @@ def init(project_name: str) -> _weave_client.WeaveClient:
     Returns:
         A Weave client.
     """
+    global _current_inited_client
     # This is the stream-table backend. Disabling it in favor of the new
     # trace-server backend.
     # return _weave_init.init_wandb(project_name).client
     # return _weave_init.init_trace_remote(project_name).client
-    return _weave_init.init_weave(project_name).client
+    _current_inited_client = _weave_init.init_weave(project_name)
+    return _current_inited_client.client
 
 
 @contextlib.contextmanager
@@ -350,4 +354,16 @@ def run():
     raise ValueError("Should not reach here")
 
 
+def finish() -> None:
+    """Stops logging to weave.
+
+    Following finish, calls of weave.op() decorated functions will no longer be logged. You will need to run weave.init() again to resume logging.
+
+    """
+    global _current_inited_client
+    if _current_inited_client is not None:
+        _weave_init.finish(_current_inited_client)
+        _current_inited_client = None
+
+
 __docspec__ = [init, publish, ref]

weave/autopatch.py

@@ -21,6 +21,19 @@ def autopatch_openai() -> None:
         patch()
 
 
+def unpatch_openai() -> None:
+    try:
+        import openai  # type: ignore
+    except ImportError:
+        pass
+    else:
+        if openai.__version__ < "1":
+            return
+        from weave.monitoring.openai import unpatch
+
+        unpatch()
+
+
 def autopatch() -> None:
     autopatch_openai()
 
@@ -31,3 +44,15 @@ def autopatch() -> None:
     mistral_patcher.attempt_patch()
     litellm_patcher.attempt_patch()
     llamaindex_patcher.attempt_patch()
+
+
+def reset_autopatch() -> None:
+    unpatch_openai()
+
+    from .integrations.mistral.mistral import mistral_patcher
+    from .integrations.litellm.litellm import litellm_patcher
+    from .integrations.llamaindex.llamaindex import llamaindex_patcher
+
+    mistral_patcher.undo_patch()
+    litellm_patcher.undo_patch()
+    llamaindex_patcher.undo_patch()

weave/monitoring/openai/openai.py

@@ -204,7 +204,7 @@ def _patch() -> None:
 
 def unpatch() -> None:
     # if _get_global_monitor() is not None:
-    info("Unpatching OpenAI completions")
+    # info("Unpatching OpenAI completions")
     openai.resources.chat.completions.Completions.create = old_create
     openai.resources.chat.completions.AsyncCompletions.create = old_async_create
 

weave/weave_init.py

@@ -159,3 +159,9 @@ def init_local() -> InitializedClient:
     server.setup_tables()
     client = weave_client.WeaveClient("none", "none", server)
     return InitializedClient(client)
+
+
+def finish(init_client: InitializedClient) -> None:
+    init_client.reset()
+    autopatch.reset_autopatch()
+    trace_sentry.global_trace_sentry.end_session()