chore(weave): Legacy Refactor pt9 (#2227)

* move to trace: table

* move to trace: trace_sentry

* move to trace: trace_api

* move to trace: serve_fastapi

* move to trace: ref_util

* move to trace: exception

weave/__init__.py

@@ -16,7 +16,7 @@
 from . import storage
 
 from .query_api import *
-from .trace_api import *
+from .trace.api import *
 
 from .errors import *
 
@@ -64,7 +64,7 @@
 # Special object informing doc generation tooling which symbols
 # to document & to associate with this module.
 __docspec__ = [
-    # Re-exported from trace_api
+    # Re-exported from trace.api
     init,
     publish,
     ref,

weave/api.py

@@ -1,4 +1,4 @@
 """These are the top-level functions in the `import weave` namespace."""
 
 from .query_api import *
-from .trace_api import *
+from .trace.api import *

weave/deploy/modal/stub.py

@@ -27,7 +27,7 @@
 @asgi_app(label=safe_name(uri.name))
 def fastapi_app() -> FastAPI:
     from weave import api
-    from weave.serve_fastapi import object_method_app
+    from weave.trace.serve_fastapi import object_method_app
 
     uri_ref = parse_uri(os.environ["MODEL_REF"])
     if not isinstance(uri_ref, ObjectRef):

weave/legacy/arrow/list_.py

@@ -12,7 +12,6 @@
 
 from weave import (
     errors,
-    ref_util,
     weave_internal,
 )
 from weave import weave_types as types
@@ -40,6 +39,7 @@
     tag_store,
     tagged_value_type,
 )
+from weave.trace import ref_util
 
 
 def reverse_dict(d: dict) -> dict:

weave/legacy/artifact_fs.py

@@ -6,10 +6,11 @@
 import os
 import typing
 
-from weave import errors, ref_util
+from weave import errors
 from weave import weave_types as types
 from weave.legacy import artifact_base, file_base, object_context, ref_base, uris
 from weave.legacy.language_features.tagging import tag_store
+from weave.trace import ref_util
 
 if typing.TYPE_CHECKING:
     from weave.legacy import graph

weave/legacy/box.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from weave import ref_util
+from weave.trace import ref_util
 from weave.legacy import context_state
 
 

weave/legacy/object_type_ref_util.py

@@ -1,6 +1,6 @@
 import typing
 
-from weave import ref_util
+from weave.trace import ref_util
 from weave.legacy import context_state
 
 

weave/tests/legacy/test_refs.py

@@ -4,13 +4,14 @@
 from weave.flow.obj import Object
 from weave.legacy import artifact_local
 from weave.legacy import ops_arrow as arrow
+from weave.trace import ref_util
 from weave.trace_server.refs_internal import (
     DICT_KEY_EDGE_NAME,
     LIST_INDEX_EDGE_NAME,
     OBJECT_ATTR_EDGE_NAME,
 )
 
-from ... import ref_util, storage
+from ... import storage
 
 
 def test_laref_artifact_version_1():

weave/trace_api.py → weave/trace/api.py

@@ -11,18 +11,18 @@
 
 # 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.
-from . import (
+from .. import (
     type_serializers,  # noqa: F401
     urls,
     util,
     weave_init,
 )
+from . import context, weave_client
+from .constants import TRACE_OBJECT_EMOJI
+from .op import Op, op
+from .refs import ObjectRef, parse_uri
+from .settings import UserSettings, parse_and_apply_settings
 from .table import Table
-from .trace import context, weave_client
-from .trace.constants import TRACE_OBJECT_EMOJI
-from .trace.op import Op, op
-from .trace.refs import ObjectRef, parse_uri
-from .trace.settings import UserSettings, parse_and_apply_settings
 
 
 def init(

weave/ref_util.py → weave/trace/ref_util.py

@@ -3,8 +3,7 @@
 from urllib import parse
 
 from weave.legacy import box
-
-from .trace_server import refs_internal
+from weave.trace_server import refs_internal
 
 DICT_KEY_EDGE_NAME = refs_internal.DICT_KEY_EDGE_NAME
 LIST_INDEX_EDGE_NAME = refs_internal.LIST_INDEX_EDGE_NAME
@@ -23,9 +22,7 @@ def parse_local_ref_str(s: str) -> typing.Tuple[str, typing.Optional[list[str]]]
 def val_with_relative_ref(
     parent_object: typing.Any, child_object: typing.Any, ref_extra_parts: list[str]
 ) -> typing.Any:
-    from weave.legacy import context_state
-
-    from .legacy import ref_base
+    from weave.legacy import context_state, ref_base
 
     # If we already have a ref, resolve it
     if isinstance(child_object, ref_base.Ref):
@@ -34,7 +31,7 @@ def val_with_relative_ref(
     # Only do this if ref_tracking_enabled right now. I just want to
     # avoid introducing new behavior into W&B prod for the moment.
     if context_state.ref_tracking_enabled():
-        from . import storage
+        from .. import storage
 
         child_ref = storage.get_ref(child_object)
         parent_ref = ref_base.get_ref(parent_object)

weave/serve_fastapi.py → weave/trace/serve_fastapi.py

@@ -17,8 +17,8 @@
 from weave.trace.op import Op
 from weave.trace.refs import ObjectRef
 
-from . import errors
-from .legacy import weave_pydantic
+from .. import errors
+from ..legacy import weave_pydantic
 
 key_cache: cache.LruTimeWindowCache[str, typing.Optional[bool]] = (
     cache.LruTimeWindowCache(datetime.timedelta(minutes=5))

weave/trace_sentry.py → weave/trace/trace_sentry.py

@@ -90,7 +90,7 @@ def setup(self) -> None:
         to avoid the possibility of interfering with the user's
         own Sentry SDK setup.
         """
-        from . import version
+        from .. import version
 
         client = sentry_sdk.Client(
             dsn=self.dsn,

weave/trace/vals.py

@@ -9,7 +9,6 @@
 from pydantic import v1 as pydantic_v1
 
 from weave.client_context.weave_client import get_weave_client
-from weave.table import Table
 from weave.trace import box
 from weave.trace.errors import InternalError
 from weave.trace.object_record import ObjectRecord
@@ -24,6 +23,7 @@
     TableRef,
 )
 from weave.trace.serialize import from_json
+from weave.trace.table import Table
 from weave.trace_server.trace_server_interface import (
     ObjReadReq,
     TableQueryReq,

weave/trace/weave_client.py

@@ -10,11 +10,10 @@
 import pydantic
 from requests import HTTPError
 
-from weave import trace_sentry, urls, version
+from weave import urls, version
 from weave.client_context import weave_client as weave_client_context
-from weave.exception import exception_to_json_str
-from weave.table import Table
-from weave.trace import call_context
+from weave.trace import call_context, trace_sentry
+from weave.trace.exception import exception_to_json_str
 from weave.trace.feedback import FeedbackQuery, RefFeedbackQuery
 from weave.trace.object_record import (
     ObjectRecord,
@@ -26,6 +25,7 @@
 from weave.trace.refs import CallRef, ObjectRef, OpRef, Ref, TableRef
 from weave.trace.serialize import from_json, isinstance_namedtuple, to_json
 from weave.trace.serializer import get_serializer_for_obj
+from weave.trace.table import Table
 from weave.trace.vals import WeaveObject, WeaveTable, make_trace_obj
 from weave.trace_server.ids import generate_id
 from weave.trace_server.trace_server_interface import (

weave/weave_init.py

@@ -2,8 +2,8 @@
 
 from weave.client_context import weave_client as weave_client_context
 
-from . import errors, init_message, trace_sentry
-from .trace import autopatch, weave_client
+from . import errors, init_message
+from .trace import autopatch, trace_sentry, weave_client
 from .trace_server import remote_http_trace_server, sqlite_trace_server
 
 _current_inited_client = None