From dfa4cf0bb8994a2df30a6529eb82452753dee95a Mon Sep 17 00:00:00 2001
From: Andrew Truong <itsandrewtruong@gmail.com>
Date: Tue, 27 Aug 2024 18:55:44 -0400
Subject: [PATCH] chore(weave): Legacy Refactor pt7 (#2225)

* move to legacy: types_numpy

* move to legacy: registry_mem

* move to legacy: weave_types
---
 weave/legacy/artifact_fs.py                    |  2 +-
 weave/legacy/codify.py                         |  4 ++--
 weave/legacy/compile.py                        |  2 +-
 weave/legacy/compile_domain.py                 |  4 ++--
 weave/legacy/decorator_class.py                |  4 ++--
 weave/legacy/decorator_op.py                   |  3 +--
 weave/legacy/derive_op.py                      |  2 +-
 weave/legacy/dispatch.py                       |  4 ++--
 weave/legacy/ecosystem/root.py                 |  4 ++--
 weave/legacy/ecosystem/wandb/wandb_objs.py     |  2 +-
 weave/legacy/execute.py                        |  2 +-
 weave/legacy/execute_fast.py                   |  2 +-
 .../tagging/process_opdef_output_type.py       |  2 +-
 weave/legacy/op_def_type.py                    |  4 ++--
 weave/legacy/ops_arrow/vectorize.py            |  3 +--
 .../run_history/history_op_common.py           |  2 +-
 weave/legacy/ops_primitives/weave_api.py       |  2 +-
 weave/legacy/panels_py/generator_templates.py  |  4 ++--
 weave/legacy/propagate_gql_keys.py             |  4 ++--
 weave/{ => legacy}/registry_mem.py             | 18 +++++++++---------
 weave/legacy/stitch.py                         |  4 ++--
 weave/{ => legacy}/types_numpy.py              | 18 +++++++++---------
 weave/{ => legacy}/weave_pydantic.py           |  4 ++--
 weave/query_api.py                             |  2 +-
 weave/serve_fastapi.py                         |  3 ++-
 weave/tests/legacy/test_derive_op.py           |  2 +-
 weave/tests/legacy/test_mappability.py         |  2 +-
 weave/tests/legacy/test_numpy.py               |  2 +-
 weave/tests/legacy/test_op_behaviors.py        |  3 ++-
 weave/tests/legacy/test_op_coverage.py         |  2 +-
 weave/tests/legacy/test_op_serialization.py    |  3 ++-
 weave/tests/legacy/test_serialize.py           |  4 ++--
 weave/tests/legacy/test_wb_domain_ops.py       |  2 +-
 weave/weave_internal.py                        |  2 +-
 weave/weave_server.py                          | 10 ++++++++--
 weave/weave_types.py                           |  2 +-
 36 files changed, 73 insertions(+), 66 deletions(-)
 rename weave/{ => legacy}/registry_mem.py (92%)
 rename weave/{ => legacy}/types_numpy.py (68%)
 rename weave/{ => legacy}/weave_pydantic.py (93%)

diff --git a/weave/legacy/artifact_fs.py b/weave/legacy/artifact_fs.py
index 53e249e09de6..b753b5019748 100644
--- a/weave/legacy/artifact_fs.py
+++ b/weave/legacy/artifact_fs.py
@@ -266,7 +266,7 @@ def type(self) -> types.Type:
 
         ot = self._outer_type
         if self.extra is not None:
-            from weave import types_numpy
+            from weave.legacy import types_numpy
 
             if not types.is_list_like(ot) and isinstance(
                 ot, types_numpy.NumpyArrayType
diff --git a/weave/legacy/codify.py b/weave/legacy/codify.py
index 04fac73cb3c0..0395775e0347 100644
--- a/weave/legacy/codify.py
+++ b/weave/legacy/codify.py
@@ -6,8 +6,8 @@
 
 import black
 
-from weave import registry_mem, storage, weave_types
-from weave.legacy import graph
+from weave import storage, weave_types
+from weave.legacy import graph, registry_mem
 
 from . import codifiable_value_mixin
 
diff --git a/weave/legacy/compile.py b/weave/legacy/compile.py
index cbe4293e893b..30cbdf4d98f4 100644
--- a/weave/legacy/compile.py
+++ b/weave/legacy/compile.py
@@ -8,7 +8,6 @@
 from weave import (
     engine_trace,
     errors,
-    registry_mem,
     weave_internal,
 )
 from weave import weave_types as types
@@ -26,6 +25,7 @@
     op_args,
     partial_object,
     propagate_gql_keys,
+    registry_mem,
     serialize,
     stitch,
     value_or_error,
diff --git a/weave/legacy/compile_domain.py b/weave/legacy/compile_domain.py
index 7e04250e4f95..2a0f686d676c 100644
--- a/weave/legacy/compile_domain.py
+++ b/weave/legacy/compile_domain.py
@@ -2,9 +2,9 @@
 
 import graphql
 
-from weave import errors, registry_mem
+from weave import errors
 from weave import weave_types as types
-from weave.legacy import gql_op_plugin, gql_to_weave, graph, op_args, stitch
+from weave.legacy import gql_op_plugin, gql_to_weave, graph, op_args, stitch, registry_mem
 from weave.legacy.input_provider import InputAndStitchProvider
 
 if typing.TYPE_CHECKING:
diff --git a/weave/legacy/decorator_class.py b/weave/legacy/decorator_class.py
index 2e42b15b1bea..99fadbcb16ed 100644
--- a/weave/legacy/decorator_class.py
+++ b/weave/legacy/decorator_class.py
@@ -1,9 +1,9 @@
 import inspect
 import typing
 
-from weave import errors, registry_mem
+from weave import errors
 from weave import weave_types as types
-from weave.legacy import context_state, derive_op, op_def
+from weave.legacy import context_state, derive_op, op_def, registry_mem
 
 # Contrary to the way it is read, the weave.class() decorator runs AFTER the
 # inner methods are defined. Therefore, this function runs after the ops are
diff --git a/weave/legacy/decorator_op.py b/weave/legacy/decorator_op.py
index c7f0dc8b6cb2..e193e4b1e385 100644
--- a/weave/legacy/decorator_op.py
+++ b/weave/legacy/decorator_op.py
@@ -4,9 +4,8 @@
 
 from typing_extensions import ParamSpec
 
-from weave import registry_mem
 from weave import weave_types as types
-from weave.legacy import context_state, derive_op, op_args, op_def, pyfunc_type_util
+from weave.legacy import context_state, derive_op, op_args, op_def, pyfunc_type_util, registry_mem
 
 if typing.TYPE_CHECKING:
     from weave.legacy.gql_op_plugin import GqlOpPlugin
diff --git a/weave/legacy/derive_op.py b/weave/legacy/derive_op.py
index 13c0dcd35dea..337bb66cccdd 100644
--- a/weave/legacy/derive_op.py
+++ b/weave/legacy/derive_op.py
@@ -9,12 +9,12 @@
 from weave import (
     errors,
     parallelism,
-    registry_mem,
     storage,
     weave_internal,
 )
 from weave import weave_types as types
 from weave.legacy import (
+    registry_mem,
     box,
     context_state,
     execute_fast,
diff --git a/weave/legacy/dispatch.py b/weave/legacy/dispatch.py
index b0af4369b9fd..4eaf2ca215f9 100644
--- a/weave/legacy/dispatch.py
+++ b/weave/legacy/dispatch.py
@@ -5,9 +5,9 @@
 import typing
 from dataclasses import dataclass
 
-from weave import errors, registry_mem, util
+from weave import errors, util
 from weave import weave_types as types
-from weave.legacy import graph, memo, op_args, op_def, pyfunc_type_util
+from weave.legacy import graph, memo, op_args, op_def, pyfunc_type_util, registry_mem
 from weave.legacy.language_features.tagging.is_tag_getter import is_tag_getter
 from weave.legacy.language_features.tagging.tagged_value_type import TaggedValueType
 
diff --git a/weave/legacy/ecosystem/root.py b/weave/legacy/ecosystem/root.py
index 0757c9b03d3f..38cd09d3e7eb 100644
--- a/weave/legacy/ecosystem/root.py
+++ b/weave/legacy/ecosystem/root.py
@@ -22,7 +22,7 @@
 """
 )
 
-from weave import registry_mem
+from weave.legacy import registry_mem
 
 op_org_name = registry_mem.memory_registry.get_op("user-name")
 
@@ -63,7 +63,7 @@ def ops(self) -> list[op_def.OpDef]:
 # objects.
 @weave.op(name="op-ecosystem", render_info={"type": "function"})
 def ecosystem() -> Ecosystem:
-    from weave import registry_mem
+    from weave.legacy import registry_mem
 
     return Ecosystem(
         _orgs=[],
diff --git a/weave/legacy/ecosystem/wandb/wandb_objs.py b/weave/legacy/ecosystem/wandb/wandb_objs.py
index c57d8a296aae..449d82ed786c 100644
--- a/weave/legacy/ecosystem/wandb/wandb_objs.py
+++ b/weave/legacy/ecosystem/wandb/wandb_objs.py
@@ -1,7 +1,7 @@
 import typing
 
 import weave
-from weave import registry_mem
+from weave.legacy import registry_mem
 from weave.legacy.ops_domain import run_ops, wb_domain_types
 
 # We can't chain ops called .name() because of a weird bug :( [its a field on VarNode].
diff --git a/weave/legacy/execute.py b/weave/legacy/execute.py
index e05958a2f9c9..ccd0d67f101d 100644
--- a/weave/legacy/execute.py
+++ b/weave/legacy/execute.py
@@ -15,7 +15,6 @@
     errors,
     parallelism,
     ref_base,
-    registry_mem,
     trace_local,
 )
 from weave import weave_types as types
@@ -27,6 +26,7 @@
 # Trace / cache
 # Language Features
 from weave.legacy import (
+    registry_mem,
     box,
     compile,
     context,
diff --git a/weave/legacy/execute_fast.py b/weave/legacy/execute_fast.py
index 29488de46f28..dbf15feeadd8 100644
--- a/weave/legacy/execute_fast.py
+++ b/weave/legacy/execute_fast.py
@@ -4,11 +4,11 @@
     engine_trace,
     errors,
     ref_base,
-    registry_mem,
     weave_internal,
 )
 from weave import weave_types as types
 from weave.legacy import (
+    registry_mem,
     box,
     compile,
     forward_graph,
diff --git a/weave/legacy/language_features/tagging/process_opdef_output_type.py b/weave/legacy/language_features/tagging/process_opdef_output_type.py
index 7364ebcd409f..a445e7485279 100644
--- a/weave/legacy/language_features/tagging/process_opdef_output_type.py
+++ b/weave/legacy/language_features/tagging/process_opdef_output_type.py
@@ -5,7 +5,7 @@
 
 import typing
 
-from weave import registry_mem
+from weave.legacy import registry_mem
 from weave import weave_types as types
 from weave.legacy import graph
 from weave.legacy.language_features.tagging.opdef_util import (
diff --git a/weave/legacy/op_def_type.py b/weave/legacy/op_def_type.py
index 964aeb4cab98..bb5fccf45258 100644
--- a/weave/legacy/op_def_type.py
+++ b/weave/legacy/op_def_type.py
@@ -11,9 +11,9 @@
 from _ast import AsyncFunctionDef, ExceptHandler
 from typing import Any
 
-from weave import environment, errors, registry_mem, storage
+from weave import environment, errors, storage
 from weave import weave_types as types
-from weave.legacy import artifact_fs, artifact_local, context_state, infer_types
+from weave.legacy import artifact_fs, artifact_local, context_state, infer_types, registry_mem
 
 if typing.TYPE_CHECKING:
     from .op_def import OpDef
diff --git a/weave/legacy/ops_arrow/vectorize.py b/weave/legacy/ops_arrow/vectorize.py
index f4123b18a54a..8f3a11a24527 100644
--- a/weave/legacy/ops_arrow/vectorize.py
+++ b/weave/legacy/ops_arrow/vectorize.py
@@ -8,10 +8,9 @@
 
 from weave import (
     errors,
-    registry_mem,
     weave_internal,
 )
-from weave.legacy import weavify
+from weave.legacy import weavify, registry_mem
 from weave import weave_types as types
 from weave.query_api import op, use
 from weave.legacy import dispatch, graph, graph_debug, op_args, op_def
diff --git a/weave/legacy/ops_domain/run_history/history_op_common.py b/weave/legacy/ops_domain/run_history/history_op_common.py
index 1987645a7c90..72d945d39f0d 100644
--- a/weave/legacy/ops_domain/run_history/history_op_common.py
+++ b/weave/legacy/ops_domain/run_history/history_op_common.py
@@ -7,12 +7,12 @@
 from weave import (
     engine_trace,
     errors,
-    registry_mem,
     util,
 )
 from weave import weave_types as types
 from weave.query_api import use
 from weave.legacy import (
+    registry_mem,
     _dict_utils,
     artifact_base,
     artifact_fs,
diff --git a/weave/legacy/ops_primitives/weave_api.py b/weave/legacy/ops_primitives/weave_api.py
index 5417cb96de7b..33ace6416b4d 100644
--- a/weave/legacy/ops_primitives/weave_api.py
+++ b/weave/legacy/ops_primitives/weave_api.py
@@ -5,13 +5,13 @@
 from weave import (
     errors,
     ref_base,
-    registry_mem,
     storage,
     weave_internal,
 )
 from weave import weave_types as types
 from weave.query_api import mutation, op, weave_class
 from weave.legacy import (
+    registry_mem,
     artifact_fs,
     artifact_local,
     artifact_wandb,
diff --git a/weave/legacy/panels_py/generator_templates.py b/weave/legacy/panels_py/generator_templates.py
index 80bcccfb04f1..b7fc045cb225 100644
--- a/weave/legacy/panels_py/generator_templates.py
+++ b/weave/legacy/panels_py/generator_templates.py
@@ -18,8 +18,8 @@
 import dataclasses
 import typing
 
-from weave import registry_mem, weave_types
-from weave.legacy import decorator_op, graph
+from weave import weave_types
+from weave.legacy import decorator_op, graph, registry_mem
 
 
 @dataclasses.dataclass
diff --git a/weave/legacy/propagate_gql_keys.py b/weave/legacy/propagate_gql_keys.py
index 50f7d284accb..08aa87a1affd 100644
--- a/weave/legacy/propagate_gql_keys.py
+++ b/weave/legacy/propagate_gql_keys.py
@@ -1,8 +1,8 @@
 import typing
 
-from weave import registry_mem
+
 from weave import weave_types as types
-from weave.legacy import gql_op_plugin, graph, input_provider, op_def, partial_object
+from weave.legacy import gql_op_plugin, graph, input_provider, op_def, partial_object, registry_mem
 
 
 def _propagate_gql_keys_for_node(
diff --git a/weave/registry_mem.py b/weave/legacy/registry_mem.py
similarity index 92%
rename from weave/registry_mem.py
rename to weave/legacy/registry_mem.py
index 290617055025..62f4a42134f4 100644
--- a/weave/registry_mem.py
+++ b/weave/legacy/registry_mem.py
@@ -28,7 +28,7 @@ class Registry:
     # the registry over HTTP.
     _updated_at: float
 
-    def __init__(self):
+    def __init__(self):  # type: ignore
         self._types = {}
         self._ops = {}
         self._ops_by_common_name = {}
@@ -41,8 +41,8 @@ def mark_updated(self) -> None:
     def updated_at(self) -> float:
         return self._updated_at
 
-    def register_op(self, op: "OpDef", location=None):
-        if context_state.get_no_op_register():
+    def register_op(self, op: "OpDef", location=None):  # type: ignore
+        if context_state.get_no_op_register():  # type: ignore[no-untyped-call]
             return op
         self.mark_updated()
         # do not save built-in ops today
@@ -97,7 +97,7 @@ def get_op(self, uri: str) -> "OpDef":
             raise errors.WeaveMissingOpDefError("Op not registered: %s" % uri)
         return res
 
-    def find_op_by_fn(self, lazy_local_fn):
+    def find_op_by_fn(self, lazy_local_fn):  # type: ignore
         for op_def in self._op_versions.values():
             if op_def.call_fn == lazy_local_fn:
                 return op_def
@@ -111,7 +111,7 @@ def find_ops_by_common_name(self, common_name: str) -> typing.List["OpDef"]:
         return ops
 
     def find_chainable_ops(self, arg0_type: weave_types.Type) -> typing.List["OpDef"]:
-        def is_chainable(op):
+        def is_chainable(op):  # type: ignore
             if not isinstance(op.input_type, op_args.OpNamedArgs):
                 return False
             args = list(op.input_type.arg_types.values())
@@ -119,9 +119,9 @@ def is_chainable(op):
                 return False
             return args[0].assign_type(arg0_type)
 
-        return [op for op in self._ops.values() if is_chainable(op)]
+        return [op for op in self._ops.values() if is_chainable(op)]  # type: ignore[no-untyped-call]
 
-    def load_saved_ops(self):
+    def load_saved_ops(self):  # type: ignore
         from weave.legacy import op_def_type
 
         for op_ref in storage.objects(op_def_type.OpDefType()):
@@ -149,7 +149,7 @@ def list_packages(self) -> typing.List["OpDef"]:
         ]
         return packages
 
-    def rename_op(self, name, new_name):
+    def rename_op(self, name, new_name):  # type: ignore
         """Internal use only, used during op bootstrapping at decorator time"""
         self.mark_updated()
         op = self._ops.pop(name)
@@ -186,4 +186,4 @@ def rename_op(self, name, new_name):
 
 
 # Processes have a singleton MemoryRegistry
-memory_registry = Registry()
+memory_registry = Registry()  # type: ignore[no-untyped-call]
diff --git a/weave/legacy/stitch.py b/weave/legacy/stitch.py
index 1fd73868f8af..3c2de9a04bcc 100644
--- a/weave/legacy/stitch.py
+++ b/weave/legacy/stitch.py
@@ -22,10 +22,10 @@
 import dataclasses
 import typing
 
-from weave.legacy import graph, op_def
+from weave.legacy import graph, op_def, registry_mem
 from weave.legacy.language_features.tagging import opdef_util
 
-from .. import errors, registry_mem
+from .. import errors
 from .. import weave_types as types
 from . import _dict_utils
 
diff --git a/weave/types_numpy.py b/weave/legacy/types_numpy.py
similarity index 68%
rename from weave/types_numpy.py
rename to weave/legacy/types_numpy.py
index 1d998500b9e5..2577e1119a76 100644
--- a/weave/types_numpy.py
+++ b/weave/legacy/types_numpy.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from . import weave_types as types
+from .. import weave_types as types
 
 
 # TODO: this doesn't match how extra works for list types...
@@ -8,22 +8,22 @@ class NumpyArrayType(types.Type):
     instance_classes = np.ndarray
     name = "WeaveNDArray"
 
-    def __init__(self, dtype="x", shape=(0,)):
+    def __init__(self, dtype="x", shape=(0,)):  # type: ignore
         self.dtype = dtype
         self.shape = shape
 
     @classmethod
-    def type_of_instance(cls, obj):
+    def type_of_instance(cls, obj):  # type: ignore
         return cls(obj.dtype, obj.shape)
 
-    def _to_dict(self):
+    def _to_dict(self):  # type: ignore
         return {"dtype": str(self.dtype), "shape": self.shape}
 
     @classmethod
-    def from_dict(cls, d):
+    def from_dict(cls, d):  # type: ignore
         return cls(np.dtype(d.get("dtype", "object")), d["shape"])
 
-    def _assign_type_inner(self, next_type):
+    def _assign_type_inner(self, next_type):  # type: ignore
         if not isinstance(next_type, NumpyArrayType):
             return False
         if (
@@ -33,13 +33,13 @@ def _assign_type_inner(self, next_type):
             return False
         return True
 
-    def save_instance(self, obj, artifact, name):
+    def save_instance(self, obj, artifact, name):  # type: ignore
         with artifact.new_file(f"{name}.npz", binary=True) as f:
             np.savez_compressed(f, arr=obj)
 
-    def load_instance(self, artifact, name, extra=None):
+    def load_instance(self, artifact, name, extra=None):  # type: ignore
         with artifact.open(f"{name}.npz", binary=True) as f:
             return np.load(f)["arr"]
 
-    def __str__(self):
+    def __str__(self):  # type: ignore
         return "<NumpyArrayType %s %s>" % (self.dtype, self.shape)
diff --git a/weave/weave_pydantic.py b/weave/legacy/weave_pydantic.py
similarity index 93%
rename from weave/weave_pydantic.py
rename to weave/legacy/weave_pydantic.py
index 19d200ad2fd1..086feb8d2fa4 100644
--- a/weave/weave_pydantic.py
+++ b/weave/legacy/weave_pydantic.py
@@ -1,7 +1,7 @@
 from pydantic import BaseModel, create_model
 
-from . import weave_types as types
-from .legacy import infer_types
+from .. import weave_types as types
+from . import infer_types
 
 
 def weave_type_to_pydantic(
diff --git a/weave/query_api.py b/weave/query_api.py
index 03ff7562fdfd..908c7e526fb9 100644
--- a/weave/query_api.py
+++ b/weave/query_api.py
@@ -24,7 +24,7 @@
 from . import weave_types as types
 
 # needed to enable automatic numpy serialization
-from . import types_numpy as _types_numpy
+from .legacy import types_numpy as _types_numpy
 
 from . import errors
 from weave.legacy.decorators import weave_class, mutation, type
diff --git a/weave/serve_fastapi.py b/weave/serve_fastapi.py
index d6b23431349d..ca63901d03b4 100644
--- a/weave/serve_fastapi.py
+++ b/weave/serve_fastapi.py
@@ -17,7 +17,8 @@
 from weave.trace.op import Op
 from weave.trace.refs import ObjectRef
 
-from . import errors, weave_pydantic
+from . import errors
+from .legacy import weave_pydantic
 
 key_cache: cache.LruTimeWindowCache[str, typing.Optional[bool]] = (
     cache.LruTimeWindowCache(datetime.timedelta(minutes=5))
diff --git a/weave/tests/legacy/test_derive_op.py b/weave/tests/legacy/test_derive_op.py
index c3bc56848be0..85b866595433 100644
--- a/weave/tests/legacy/test_derive_op.py
+++ b/weave/tests/legacy/test_derive_op.py
@@ -1,5 +1,5 @@
 from ... import api as weave
-from ... import registry_mem
+from ...legacy import registry_mem
 
 
 def test_mapped_add():
diff --git a/weave/tests/legacy/test_mappability.py b/weave/tests/legacy/test_mappability.py
index 9beb8445a634..18822a65478f 100644
--- a/weave/tests/legacy/test_mappability.py
+++ b/weave/tests/legacy/test_mappability.py
@@ -2,8 +2,8 @@
 from weave.legacy import context_state as _context
 from weave.legacy import graph
 
-from ... import registry_mem
 from ... import weave_types as types
+from ...legacy import registry_mem
 from ...weave_internal import make_const_node
 
 _loading_builtins_token = _context.set_loading_built_ins()
diff --git a/weave/tests/legacy/test_numpy.py b/weave/tests/legacy/test_numpy.py
index cdb89250c9e0..3ac8858e42c4 100644
--- a/weave/tests/legacy/test_numpy.py
+++ b/weave/tests/legacy/test_numpy.py
@@ -3,8 +3,8 @@
 from weave.legacy import artifact_fs, artifact_wandb
 from weave.legacy.ops_domain import table
 
-from ... import types_numpy as numpy_types
 from ... import weave_types as types
+from ...legacy import types_numpy as numpy_types
 
 
 def test_construct_numpy_type():
diff --git a/weave/tests/legacy/test_op_behaviors.py b/weave/tests/legacy/test_op_behaviors.py
index 263f8541647d..e49ff51c9c1c 100644
--- a/weave/tests/legacy/test_op_behaviors.py
+++ b/weave/tests/legacy/test_op_behaviors.py
@@ -25,7 +25,8 @@
     TaggedValueType,
 )
 
-from ... import registry_mem, storage, weave_internal
+from ... import storage, weave_internal
+from ...legacy import registry_mem
 from ..concrete_tagged_value import (
     TaggedValue,
     concrete_from_tagstore,
diff --git a/weave/tests/legacy/test_op_coverage.py b/weave/tests/legacy/test_op_coverage.py
index 2eadbe3c55d4..cede44a11091 100644
--- a/weave/tests/legacy/test_op_coverage.py
+++ b/weave/tests/legacy/test_op_coverage.py
@@ -1,4 +1,4 @@
-from ... import registry_mem
+from ...legacy import registry_mem
 
 
 def make_error_message(missing_ops, section_name):
diff --git a/weave/tests/legacy/test_op_serialization.py b/weave/tests/legacy/test_op_serialization.py
index 66ad75c99d59..7c08a59b9828 100644
--- a/weave/tests/legacy/test_op_serialization.py
+++ b/weave/tests/legacy/test_op_serialization.py
@@ -1,8 +1,9 @@
 import pytest
 
 import weave
+import weave.legacy
 
-ops = weave.registry_mem.memory_registry.list_ops()
+ops = weave.legacy.registry_mem.memory_registry.list_ops()
 
 
 def output_type_dict_is_const_function_node(output_type_dict):
diff --git a/weave/tests/legacy/test_serialize.py b/weave/tests/legacy/test_serialize.py
index 85b34ced7f93..3ce4594856f9 100644
--- a/weave/tests/legacy/test_serialize.py
+++ b/weave/tests/legacy/test_serialize.py
@@ -2,9 +2,9 @@
 
 import weave
 from weave import query_api as api
-from weave import registry_mem, weave_internal
+from weave import weave_internal
 from weave import weave_types as types
-from weave.legacy import graph, op_args, ops, serialize
+from weave.legacy import graph, op_args, ops, registry_mem, serialize
 from weave.legacy.ops_primitives import list_
 from weave.weave_internal import make_const_node
 
diff --git a/weave/tests/legacy/test_wb_domain_ops.py b/weave/tests/legacy/test_wb_domain_ops.py
index ec2d5bb74bd8..5674b65204e3 100644
--- a/weave/tests/legacy/test_wb_domain_ops.py
+++ b/weave/tests/legacy/test_wb_domain_ops.py
@@ -9,7 +9,7 @@
 from weave.legacy.ops_domain import wb_domain_types
 from weave.legacy.ops_primitives import _dict_utils
 
-from ... import registry_mem
+from ...legacy import registry_mem
 from .. import fixture_fakewandb as fwb
 
 """
diff --git a/weave/weave_internal.py b/weave/weave_internal.py
index 68a7edbe53c1..2b6e07b294b2 100644
--- a/weave/weave_internal.py
+++ b/weave/weave_internal.py
@@ -150,7 +150,7 @@ def define_fn(
 # and the function doesn't explicitly operate on tagged values. this ensures that the input tags
 # are propagated appropriately to the output type of the function.
 def refine_graph(node: graph.Node) -> graph.Node:
-    from .registry_mem import memory_registry
+    from .legacy.registry_mem import memory_registry
 
     if isinstance(node, (graph.ConstNode, graph.VoidNode, graph.VarNode)):
         return node
diff --git a/weave/weave_server.py b/weave/weave_server.py
index db96977c0a4b..461fe55989d0 100644
--- a/weave/weave_server.py
+++ b/weave/weave_server.py
@@ -31,12 +31,18 @@
     errors,
     filesystem,
     logs,
-    registry_mem,
     server,
     storage,
     util,
 )
-from weave.legacy import context_state, graph, value_or_error, wandb_api, weavejs_fixes
+from weave.legacy import (
+    context_state,
+    graph,
+    registry_mem,
+    value_or_error,
+    wandb_api,
+    weavejs_fixes,
+)
 from weave.legacy.language_features.tagging import tag_store
 from weave.server_error_handling import client_safe_http_exceptions_as_werkzeug
 
diff --git a/weave/weave_types.py b/weave/weave_types.py
index 990ec2586a86..bee5df4bd115 100644
--- a/weave/weave_types.py
+++ b/weave/weave_types.py
@@ -1057,7 +1057,7 @@ def property_types(self) -> dict[str, Type]:
 
     @classmethod
     def typeclass_of_class(cls, check_class):
-        from . import weave_pydantic
+        from .legacy import weave_pydantic
 
         if not issubclass(check_class, pydantic.BaseModel):
             return cls