From 99ed17802df6b2036c5c06e00fa8b4c5c70322e6 Mon Sep 17 00:00:00 2001
From: Andrew Truong <itsandrewtruong@gmail.com>
Date: Sat, 24 Aug 2024 11:46:46 -0400
Subject: [PATCH] move to trace: client_context

---
 weave/integrations/langchain/langchain.py        | 2 +-
 weave/integrations/llamaindex/llamaindex.py      | 2 +-
 weave/legacy/mappers_python_def.py               | 2 +-
 weave/legacy/monitoring/monitor.py               | 2 +-
 weave/legacy/monitoring/openai/openai.py         | 2 +-
 weave/legacy/ref_base.py                         | 2 +-
 weave/legacy/storage.py                          | 2 +-
 weave/tests/legacy/test_weave_api.py             | 2 +-
 weave/tests/trace/test_client_trace.py           | 6 +++---
 weave/trace/api.py                               | 2 +-
 weave/trace/cli.py                               | 2 +-
 weave/{ => trace}/client_context/weave_client.py | 0
 weave/trace/custom_objs.py                       | 2 +-
 weave/trace/feedback.py                          | 2 +-
 weave/trace/op.py                                | 2 +-
 weave/trace/refs.py                              | 2 +-
 weave/trace/rich_refs.py                         | 2 +-
 weave/trace/tests/test_vals.py                   | 2 +-
 weave/trace/vals.py                              | 2 +-
 weave/trace/weave_client.py                      | 2 +-
 weave/trace/weave_init.py                        | 2 +-
 21 files changed, 22 insertions(+), 22 deletions(-)
 rename weave/{ => trace}/client_context/weave_client.py (100%)

diff --git a/weave/integrations/langchain/langchain.py b/weave/integrations/langchain/langchain.py
index 9d2ee657d9ef..0b9743123909 100644
--- a/weave/integrations/langchain/langchain.py
+++ b/weave/integrations/langchain/langchain.py
@@ -35,12 +35,12 @@
 from contextvars import ContextVar
 from uuid import UUID
 
-from weave.client_context import weave_client as weave_client_context
 from weave.integrations.integration_utilities import (
     make_pythonic_function_name,
     truncate_op_name,
 )
 from weave.trace import call_context
+from weave.trace.client_context import weave_client as weave_client_context
 from weave.trace.patcher import Patcher
 from weave.trace.weave_client import Call
 
diff --git a/weave/integrations/llamaindex/llamaindex.py b/weave/integrations/llamaindex/llamaindex.py
index acb22efb3c31..df30a6b6abb3 100644
--- a/weave/integrations/llamaindex/llamaindex.py
+++ b/weave/integrations/llamaindex/llamaindex.py
@@ -1,4 +1,4 @@
-from weave.client_context import weave_client as weave_client_context
+from weave.trace.client_context import weave_client as weave_client_context
 from weave.trace.patcher import Patcher
 from weave.trace.weave_client import Call
 
diff --git a/weave/legacy/mappers_python_def.py b/weave/legacy/mappers_python_def.py
index 5a1ee099220e..c50bdc648ac3 100644
--- a/weave/legacy/mappers_python_def.py
+++ b/weave/legacy/mappers_python_def.py
@@ -19,7 +19,7 @@
 from weave.legacy import timestamp as weave_timestamp
 from weave.legacy.language_features.tagging import tagged_value_type
 from weave.legacy.partial_object import PartialObject, PartialObjectType
-from weave.client_context import weave_client as weave_client_context
+from weave.trace.client_context import weave_client as weave_client_context
 
 
 class TypedDictToPyDict(mappers_weave.TypedDictMapper):
diff --git a/weave/legacy/monitoring/monitor.py b/weave/legacy/monitoring/monitor.py
index 5a707af576ca..518383ab59f7 100644
--- a/weave/legacy/monitoring/monitor.py
+++ b/weave/legacy/monitoring/monitor.py
@@ -19,7 +19,7 @@
 )
 from weave.legacy.wandb_interface.wandb_stream_table import StreamTable
 from weave.trace import call_context, context as trace_context
-from weave.client_context import weave_client as weave_client_context
+from weave.trace.client_context import weave_client as weave_client_context
 
 logger = logging.getLogger(__name__)
 
diff --git a/weave/legacy/monitoring/openai/openai.py b/weave/legacy/monitoring/openai/openai.py
index c002ec85c4ca..f41f551d9949 100644
--- a/weave/legacy/monitoring/openai/openai.py
+++ b/weave/legacy/monitoring/openai/openai.py
@@ -11,7 +11,7 @@
 from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
 from packaging import version
 from weave.trace import call_context
-from weave.client_context import weave_client as weave_client_context
+from weave.trace.client_context import weave_client as weave_client_context
 from weave.legacy.monitoring.monitor import _get_global_monitor
 from weave.legacy.monitoring.openai.models import *
 from weave.legacy.monitoring.openai.util import *
diff --git a/weave/legacy/ref_base.py b/weave/legacy/ref_base.py
index 234b9f4aedec..107ca9ca6c11 100644
--- a/weave/legacy/ref_base.py
+++ b/weave/legacy/ref_base.py
@@ -5,7 +5,7 @@
 import weakref
 from typing import Sequence
 
-from weave.client_context import weave_client as weave_client_context
+from weave.trace.client_context import weave_client as weave_client_context
 from weave.legacy import box, context_state, object_context, uris
 from weave.legacy.language_features.tagging import tag_store
 
diff --git a/weave/legacy/storage.py b/weave/legacy/storage.py
index e13f102bb1b5..cd14615af6c7 100644
--- a/weave/legacy/storage.py
+++ b/weave/legacy/storage.py
@@ -7,7 +7,7 @@
 import re
 import typing
 
-from weave.client_context import weave_client as weave_client_context
+from weave.trace.client_context import weave_client as weave_client_context
 from weave.legacy import (
     artifact_base,
     artifact_fs,
diff --git a/weave/tests/legacy/test_weave_api.py b/weave/tests/legacy/test_weave_api.py
index ad1996ccc970..90fd26ebb248 100644
--- a/weave/tests/legacy/test_weave_api.py
+++ b/weave/tests/legacy/test_weave_api.py
@@ -39,7 +39,7 @@ def test_weave_finish_unsets_client(client):
     def foo():
         return 1
 
-    weave.client_context.weave_client.set_weave_client_global(None)
+    weave.trace.client_context.weave_client.set_weave_client_global(None)
     weave.trace.weave_init._current_inited_client = (
         weave.trace.weave_init.InitializedClient(client)
     )
diff --git a/weave/tests/trace/test_client_trace.py b/weave/tests/trace/test_client_trace.py
index 57b7e697fe5b..dd638e7fbe41 100644
--- a/weave/tests/trace/test_client_trace.py
+++ b/weave/tests/trace/test_client_trace.py
@@ -1558,12 +1558,12 @@ def test_ref_get_no_client(trace_init_client):
 
 @contextmanager
 def _no_graph_client():
-    client = weave.client_context.weave_client.get_weave_client()
-    weave.client_context.weave_client.set_weave_client_global(None)
+    client = weave.trace.client_context.weave_client.get_weave_client()
+    weave.trace.client_context.weave_client.set_weave_client_global(None)
     try:
         yield
     finally:
-        weave.client_context.weave_client.set_weave_client_global(client)
+        weave.trace.client_context.weave_client.set_weave_client_global(client)
 
 
 @contextmanager
diff --git a/weave/trace/api.py b/weave/trace/api.py
index 4b1929770ab9..916d27c28410 100644
--- a/weave/trace/api.py
+++ b/weave/trace/api.py
@@ -6,9 +6,9 @@
 import time
 from typing import Any, Callable, Iterator, Optional, Union
 
-from weave.client_context import weave_client as weave_client_context
 from weave.legacy import urls, util
 from weave.trace.call_context import get_current_call
+from weave.trace.client_context import weave_client as weave_client_context
 
 # TODO: type_serializers is imported here to trigger registration of the image serializer.
 # There is probably a better place for this, but including here for now to get the fix in.
diff --git a/weave/trace/cli.py b/weave/trace/cli.py
index 97aa7aad6355..59a6118d12c1 100644
--- a/weave/trace/cli.py
+++ b/weave/trace/cli.py
@@ -134,7 +134,7 @@ def gcp(
 )
 @click.option("--dev", is_flag=True, help="Run the function locally.")
 def modal(model_ref: str, project: str, auth_entity: str, dev: bool = False) -> None:
-    from .deploy import modal as mdp
+    from weave.deploy import modal as mdp
 
     if dev:
         print(f"Developing model {model_ref}...")
diff --git a/weave/client_context/weave_client.py b/weave/trace/client_context/weave_client.py
similarity index 100%
rename from weave/client_context/weave_client.py
rename to weave/trace/client_context/weave_client.py
diff --git a/weave/trace/custom_objs.py b/weave/trace/custom_objs.py
index 6a8c0022d3ea..99b7d41eb601 100644
--- a/weave/trace/custom_objs.py
+++ b/weave/trace/custom_objs.py
@@ -4,9 +4,9 @@
 import tempfile
 from typing import Any, Dict, Generator, Iterator, Mapping, Optional, Union
 
-from weave.client_context.weave_client import require_weave_client
 from weave.legacy import artifact_fs
 from weave.trace import op_type  # noqa: F401, Must import this to register op save/load
+from weave.trace.client_context.weave_client import require_weave_client
 from weave.trace.op import Op, op
 from weave.trace.refs import ObjectRef, parse_uri
 from weave.trace.serializer import get_serializer_by_id, get_serializer_for_obj
diff --git a/weave/trace/feedback.py b/weave/trace/feedback.py
index 8f2eafd378f3..a9a386aa002d 100644
--- a/weave/trace/feedback.py
+++ b/weave/trace/feedback.py
@@ -5,9 +5,9 @@
 
 from rich.table import Table
 
-from weave.client_context import weave_client as weave_client_context
 from weave.legacy import util
 from weave.trace import rich_pydantic_util
+from weave.trace.client_context import weave_client as weave_client_context
 from weave.trace.refs import parse_uri
 from weave.trace.rich_container import AbstractRichContainer
 from weave.trace.rich_refs import Refs
diff --git a/weave/trace/op.py b/weave/trace/op.py
index 29babb6b2d9d..3dca692b28d2 100644
--- a/weave/trace/op.py
+++ b/weave/trace/op.py
@@ -17,9 +17,9 @@
     runtime_checkable,
 )
 
-from weave.client_context import weave_client as weave_client_context
 from weave.legacy import context_state
 from weave.trace import box, call_context, settings
+from weave.trace.client_context import weave_client as weave_client_context
 from weave.trace.context import call_attributes
 from weave.trace.errors import OpCallError
 from weave.trace.refs import ObjectRef
diff --git a/weave/trace/refs.py b/weave/trace/refs.py
index 0df24902462f..0634de90a809 100644
--- a/weave/trace/refs.py
+++ b/weave/trace/refs.py
@@ -63,7 +63,7 @@ def get(self) -> Any:
         # Move import here so that it only happens when the function is called.
         # This import is invalid in the trace server and represents a dependency
         # that should be removed.
-        from weave.client_context.weave_client import get_weave_client
+        from weave.trace.client_context.weave_client import get_weave_client
         from weave.trace.weave_init import init_weave
 
         gc = get_weave_client()
diff --git a/weave/trace/rich_refs.py b/weave/trace/rich_refs.py
index c72e07f4227e..aab9f11d7f84 100644
--- a/weave/trace/rich_refs.py
+++ b/weave/trace/rich_refs.py
@@ -4,7 +4,7 @@
 
 from rich.table import Table
 
-from weave.client_context import weave_client as weave_client_context
+from weave.trace.client_context import weave_client as weave_client_context
 from weave.trace.refs import AnyRef, CallRef, parse_uri
 from weave.trace.rich_container import AbstractRichContainer
 from weave.trace.vals import WeaveObject
diff --git a/weave/trace/tests/test_vals.py b/weave/trace/tests/test_vals.py
index 37a3f9c75d49..32e1b815fed2 100644
--- a/weave/trace/tests/test_vals.py
+++ b/weave/trace/tests/test_vals.py
@@ -31,6 +31,6 @@ def my_op(obj: MyObj) -> None:
     # simulate a situation where the client is closed
     # but a reference to a trace object still exists.
 
-    weave.client_context.weave_client.set_weave_client_global(None)
+    weave.trace.client_context.weave_client.set_weave_client_global(None)
 
     assert my_obj.val == 1
diff --git a/weave/trace/vals.py b/weave/trace/vals.py
index c92f9a024ec6..a54d01b06266 100644
--- a/weave/trace/vals.py
+++ b/weave/trace/vals.py
@@ -8,8 +8,8 @@
 from pydantic import BaseModel
 from pydantic import v1 as pydantic_v1
 
-from weave.client_context.weave_client import get_weave_client
 from weave.trace import box
+from weave.trace.client_context.weave_client import get_weave_client
 from weave.trace.errors import InternalError
 from weave.trace.object_record import ObjectRecord
 from weave.trace.op import Op, call, maybe_bind_method
diff --git a/weave/trace/weave_client.py b/weave/trace/weave_client.py
index b856c059df46..11029cc64fde 100644
--- a/weave/trace/weave_client.py
+++ b/weave/trace/weave_client.py
@@ -10,9 +10,9 @@
 from requests import HTTPError
 
 from weave import version
-from weave.client_context import weave_client as weave_client_context
 from weave.legacy import urls
 from weave.trace import call_context, trace_sentry
+from weave.trace.client_context import weave_client as weave_client_context
 from weave.trace.exception import exception_to_json_str
 from weave.trace.feedback import FeedbackQuery, RefFeedbackQuery
 from weave.trace.object_record import (
diff --git a/weave/trace/weave_init.py b/weave/trace/weave_init.py
index 2bf83c97080e..07de4b40045f 100644
--- a/weave/trace/weave_init.py
+++ b/weave/trace/weave_init.py
@@ -1,8 +1,8 @@
 import typing
 
-from weave.client_context import weave_client as weave_client_context
 from weave.legacy import errors
 from weave.trace import autopatch, init_message, trace_sentry, weave_client
+from weave.trace.client_context import weave_client as weave_client_context
 from weave.trace_server import remote_http_trace_server, sqlite_trace_server
 
 _current_inited_client = None