diff --git a/weave/legacy/codify.py b/weave/legacy/codify.py index 04fac73cb3c..0395775e034 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 cbe4293e893..30cbdf4d98f 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 7e04250e4f9..2a0f686d676 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 2e42b15b1be..99fadbcb16e 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 c7f0dc8b6cb..e193e4b1e38 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 13c0dcd35de..337bb66cccd 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 b0af4369b9f..4eaf2ca215f 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 0757c9b03d3..38cd09d3e7e 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 c57d8a296aa..449d82ed786 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 e05958a2f9c..ccd0d67f101 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 29488de46f2..dbf15feeadd 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 7364ebcd409..a445e748527 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 964aeb4cab9..bb5fccf4525 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 f4123b18a54..8f3a11a2452 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 1987645a7c9..72d945d39f0 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 5417cb96de7..33ace6416b4 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 80bcccfb04f..b7fc045cb22 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 50f7d284acc..08aa87a1aff 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 29061705502..62f4a42134f 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 1fd73868f8a..3c2de9a04bc 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/tests/legacy/test_derive_op.py b/weave/tests/legacy/test_derive_op.py index c3bc56848be..85b86659543 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 9beb8445a63..18822a65478 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_op_behaviors.py b/weave/tests/legacy/test_op_behaviors.py index 263f8541647..e49ff51c9c1 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 2eadbe3c55d..cede44a1109 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 66ad75c99d5..7c08a59b982 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 85b34ced7f9..3ce4594856f 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 ec2d5bb74bd..5674b65204e 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 68a7edbe53c..2b6e07b294b 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 db96977c0a4..461fe55989d 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