move to trace: client_context

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
 

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
 

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):

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__)
 

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 *

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
 

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,

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)
     )

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

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.

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}...")

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

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

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

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()

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

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

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

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 (

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