From eee78ed817f297c273bb0e9ecc577654c53ca72a Mon Sep 17 00:00:00 2001
From: Andrew Truong <itsandrewtruong@gmail.com>
Date: Fri, 23 Aug 2024 18:20:23 -0400
Subject: [PATCH] delete: api; split out trace and query api from single api
 file

---
 weave/__init__.py                            | 16 ++++++----------
 weave/api.py                                 |  4 ----
 weave/deploy/modal/stub.py                   |  2 +-
 weave/tests/legacy/test_api.py               |  3 ++-
 weave/tests/legacy/test_arrow.py             |  2 +-
 weave/tests/legacy/test_arrow_vectorizer.py  |  3 +--
 weave/tests/legacy/test_async.py             |  4 +---
 weave/tests/legacy/test_basic_ops.py         |  3 +--
 weave/tests/legacy/test_compile.py           |  3 +--
 weave/tests/legacy/test_custom_types.py      |  2 +-
 weave/tests/legacy/test_decorators.py        |  3 +--
 weave/tests/legacy/test_derive_op.py         |  3 ++-
 weave/tests/legacy/test_examples.py          |  3 +--
 weave/tests/legacy/test_execute.py           |  3 +--
 weave/tests/legacy/test_file.py              |  4 +---
 weave/tests/legacy/test_list_arrow_compat.py |  2 +-
 weave/tests/legacy/test_media_user.py        |  3 +--
 weave/tests/legacy/test_mutations.py         |  3 +--
 weave/tests/legacy/test_node_ref.py          |  2 +-
 weave/tests/legacy/test_op_def.py            |  3 +--
 weave/tests/legacy/test_partial_object.py    |  2 +-
 weave/tests/legacy/test_run_segment.py       |  4 +---
 weave/tests/legacy/test_storage.py           |  2 +-
 weave/tests/legacy/test_table_ops.py         |  2 +-
 weave/tests/legacy/test_trace.py             |  2 +-
 weave/tests/legacy/test_wb_domain_types.py   |  3 +--
 weave/tests/legacy/test_weavejs_fixes.py     |  3 +--
 weave/tests/list_arrow_test_helpers.py       |  3 +--
 weave/tests/trace/test_server.py             |  2 +-
 weave/trace/cli.py                           |  3 +--
 30 files changed, 36 insertions(+), 61 deletions(-)
 delete mode 100644 weave/api.py

diff --git a/weave/__init__.py b/weave/__init__.py
index 0ed50283b77..9ea806d2e21 100644
--- a/weave/__init__.py
+++ b/weave/__init__.py
@@ -13,25 +13,21 @@
 _loading_builtins_token = _context_state.set_loading_built_ins()
 
 from weave.legacy import weave_types as types
-from .legacy import storage
-
-from .legacy.api import *
-from .trace.api import *
-
-from .legacy.errors import *
-
+from weave.legacy import storage
+from weave.legacy.api import *
+from weave.legacy.errors import *
 from weave.legacy import mappers_python_def
-
 from weave.legacy import wandb_api as _wandb_api
+from weave.legacy import context as _context
 
-from . import version
+from weave import version
 
 _wandb_api.init()
 
 # Ensure there is a client available for eager mode
-from weave.legacy import context as _context
 
 
+from weave.trace.api import *
 _context_state.clear_loading_built_ins(_loading_builtins_token)
 
 __version__ = version.VERSION
diff --git a/weave/api.py b/weave/api.py
deleted file mode 100644
index 892dd0ec968..00000000000
--- a/weave/api.py
+++ /dev/null
@@ -1,4 +0,0 @@
-"""These are the top-level functions in the `import weave` namespace."""
-
-from .legacy.api import *
-from .trace.api import *
diff --git a/weave/deploy/modal/stub.py b/weave/deploy/modal/stub.py
index 2dee6e838e6..89eeff34bcc 100644
--- a/weave/deploy/modal/stub.py
+++ b/weave/deploy/modal/stub.py
@@ -26,7 +26,7 @@
 @stub.function(image=image, secret=Secret.from_dotenv(__file__))
 @asgi_app(label=safe_name(uri.name))
 def fastapi_app() -> FastAPI:
-    from weave import api
+    from weave.trace import api
     from weave.trace.serve_fastapi import object_method_app
 
     uri_ref = parse_uri(os.environ["MODEL_REF"])
diff --git a/weave/tests/legacy/test_api.py b/weave/tests/legacy/test_api.py
index dce30746540..75111a6fc8c 100644
--- a/weave/tests/legacy/test_api.py
+++ b/weave/tests/legacy/test_api.py
@@ -1,6 +1,7 @@
 import shutil
 
-from ... import api as weave
+from weave.legacy import api as weave
+
 from ...legacy.show import _show_params
 
 
diff --git a/weave/tests/legacy/test_arrow.py b/weave/tests/legacy/test_arrow.py
index 7e8a7703e64..e5554ab165c 100644
--- a/weave/tests/legacy/test_arrow.py
+++ b/weave/tests/legacy/test_arrow.py
@@ -7,6 +7,7 @@
 import pytest
 from PIL import Image
 
+from weave.legacy import api as weave
 from weave.legacy import (
     box,
     context_state,
@@ -36,7 +37,6 @@
 from weave.legacy.ops_primitives import list_, make_list
 from weave.tests import list_arrow_test_helpers as lath
 
-from ... import api as weave
 from ...tests import tag_test_util as ttu
 from ...tests import weavejs_ops
 from ...tests.legacy import test_wb
diff --git a/weave/tests/legacy/test_arrow_vectorizer.py b/weave/tests/legacy/test_arrow_vectorizer.py
index 8f41805c840..1bb224fdb4f 100644
--- a/weave/tests/legacy/test_arrow_vectorizer.py
+++ b/weave/tests/legacy/test_arrow_vectorizer.py
@@ -5,6 +5,7 @@
 import pytest
 from pyarrow import compute as pc
 
+from weave.legacy import api as weave
 from weave.legacy import box, dispatch, errors, ops, weave_internal
 from weave.legacy import ops_arrow as arrow
 from weave.legacy import weave_types as types
@@ -18,8 +19,6 @@
 from weave.legacy.ops_domain import wb_domain_types as wdt
 from weave.legacy.ops_primitives import Boolean, Number, date, dict_, list_
 
-from ... import api as weave
-
 string_ops_test_cases = [
     ("eq-scalar", lambda x: x == "bc", [True, False, False]),
     ("ne-scalar", lambda x: x != "bc", [False, True, True]),
diff --git a/weave/tests/legacy/test_async.py b/weave/tests/legacy/test_async.py
index 976e7e9398d..e0a9632154e 100644
--- a/weave/tests/legacy/test_async.py
+++ b/weave/tests/legacy/test_async.py
@@ -1,8 +1,6 @@
 import pytest
 
-from weave.legacy import async_demo, ops, runs
-
-from ... import api, storage
+from weave.legacy import api, async_demo, ops, runs, storage
 
 
 def test_run_basic():
diff --git a/weave/tests/legacy/test_basic_ops.py b/weave/tests/legacy/test_basic_ops.py
index 19c2d3020b8..e8868da0e03 100644
--- a/weave/tests/legacy/test_basic_ops.py
+++ b/weave/tests/legacy/test_basic_ops.py
@@ -1,10 +1,9 @@
+from weave.legacy import api as weave
 from weave.legacy import box, ops
 from weave.legacy.ops_primitives import number
 from weave.legacy.ops_primitives.string import *
 from weave.legacy.weave_internal import make_const_node
 
-from ... import api as weave
-
 
 def test_number_ops():
     nine = make_const_node(weave.types.Number(), 9)
diff --git a/weave/tests/legacy/test_compile.py b/weave/tests/legacy/test_compile.py
index 2bdde6920af..dc5679d769d 100644
--- a/weave/tests/legacy/test_compile.py
+++ b/weave/tests/legacy/test_compile.py
@@ -4,14 +4,13 @@
 import weave
 from weave.legacy import async_demo, compile, graph
 from weave.legacy import weave_types as types
+from weave.legacy.api import use
 from weave.legacy.dispatch import RuntimeOutputNode
 from weave.legacy.ops_arrow import to_arrow
 from weave.legacy.ops_arrow.vectorize import raise_on_python_bailout
 from weave.legacy.wandb_interface.wandb_stream_table import StreamTable
 from weave.legacy.weave_internal import const, define_fn, make_const_node
 
-from ...api import use
-
 
 def test_automatic_await_compile():
     twelve = async_demo.slowmult(3, 4, 0.01)
diff --git a/weave/tests/legacy/test_custom_types.py b/weave/tests/legacy/test_custom_types.py
index 5f8940785fa..f5b15ff28ba 100644
--- a/weave/tests/legacy/test_custom_types.py
+++ b/weave/tests/legacy/test_custom_types.py
@@ -1,10 +1,10 @@
 import pytest
 from PIL import Image
 
+from weave.legacy import api as weave
 from weave.legacy import context_state as _context
 from weave.legacy import ops_arrow
 
-from ... import api as weave
 from ... import errors
 from .. import geom
 
diff --git a/weave/tests/legacy/test_decorators.py b/weave/tests/legacy/test_decorators.py
index cc1a2e1846b..7682136833f 100644
--- a/weave/tests/legacy/test_decorators.py
+++ b/weave/tests/legacy/test_decorators.py
@@ -1,9 +1,8 @@
+from weave.legacy import api as weave
 from weave.legacy import storage
 from weave.legacy import weave_types as types
 from weave.legacy.decorator_op import op
 
-from ... import api as weave
-
 
 def test_function_op_name():
     @op()
diff --git a/weave/tests/legacy/test_derive_op.py b/weave/tests/legacy/test_derive_op.py
index 85b86659543..637255a8ca7 100644
--- a/weave/tests/legacy/test_derive_op.py
+++ b/weave/tests/legacy/test_derive_op.py
@@ -1,4 +1,5 @@
-from ... import api as weave
+from weave.legacy import api as weave
+
 from ...legacy import registry_mem
 
 
diff --git a/weave/tests/legacy/test_examples.py b/weave/tests/legacy/test_examples.py
index b9324c781f8..1fc1a60d773 100644
--- a/weave/tests/legacy/test_examples.py
+++ b/weave/tests/legacy/test_examples.py
@@ -1,10 +1,9 @@
 import math
 import typing
 
+from weave.legacy import api as weave
 from weave.legacy import context, context_state
 
-from ... import api as weave
-
 
 class XOnly(typing.TypedDict):
     x: float
diff --git a/weave/tests/legacy/test_execute.py b/weave/tests/legacy/test_execute.py
index a36eec66b6b..8c77e72a21e 100644
--- a/weave/tests/legacy/test_execute.py
+++ b/weave/tests/legacy/test_execute.py
@@ -4,10 +4,9 @@
 import pytest
 
 import weave
-from weave.legacy import environment, execute, ops, weave_internal
+from weave.legacy import api, environment, execute, ops, weave_internal
 from weave.legacy import weave_types as types
 
-from ... import api
 from . import test_wb
 
 execute_test_count_op_run_count = 0
diff --git a/weave/tests/legacy/test_file.py b/weave/tests/legacy/test_file.py
index 0fed90136c3..d8561dae6a0 100644
--- a/weave/tests/legacy/test_file.py
+++ b/weave/tests/legacy/test_file.py
@@ -1,9 +1,7 @@
 import pytest
 
 import weave
-from weave.legacy import context_state, environment, ops
-
-from ... import api, errors
+from weave.legacy import api, context_state, environment, errors, ops
 
 
 def test_dir():
diff --git a/weave/tests/legacy/test_list_arrow_compat.py b/weave/tests/legacy/test_list_arrow_compat.py
index 08287ed6da9..37680040db3 100644
--- a/weave/tests/legacy/test_list_arrow_compat.py
+++ b/weave/tests/legacy/test_list_arrow_compat.py
@@ -3,6 +3,7 @@
 import numpy as np
 import pytest
 
+from weave.legacy import api as weave
 from weave.legacy import box, ops, weave_internal
 from weave.legacy import ops_arrow as arrow
 from weave.legacy import weave_types as types
@@ -13,7 +14,6 @@
 )
 from weave.legacy.ops_primitives import dict_, list_
 
-from ... import api as weave
 from ...tests import tag_test_util as ttu
 from .. import list_arrow_test_helpers as lath
 
diff --git a/weave/tests/legacy/test_media_user.py b/weave/tests/legacy/test_media_user.py
index edbceb07217..8bb37deb4f4 100644
--- a/weave/tests/legacy/test_media_user.py
+++ b/weave/tests/legacy/test_media_user.py
@@ -1,10 +1,9 @@
 from PIL import Image
 
+from weave.legacy import api as weave
 from weave.legacy import context_state
 from weave.legacy.ops_primitives import geom as media_user
 
-from ... import api as weave
-
 
 def test_im_with_metadata():
     base_im = Image.linear_gradient("L").resize((32, 32))
diff --git a/weave/tests/legacy/test_mutations.py b/weave/tests/legacy/test_mutations.py
index a83ff5a0272..453d8af7ec5 100644
--- a/weave/tests/legacy/test_mutations.py
+++ b/weave/tests/legacy/test_mutations.py
@@ -1,7 +1,6 @@
+from weave.legacy import api as weave
 from weave.legacy import ops, storage, weave_internal
 
-from ... import api as weave
-
 
 def test_autocommit(cereal_csv):
     csv = ops.local_path(cereal_csv).readcsv()
diff --git a/weave/tests/legacy/test_node_ref.py b/weave/tests/legacy/test_node_ref.py
index eaacc2c3190..0a8d56f1e7e 100644
--- a/weave/tests/legacy/test_node_ref.py
+++ b/weave/tests/legacy/test_node_ref.py
@@ -1,6 +1,6 @@
+from weave.legacy import api as weave
 from weave.legacy import graph
 
-from ... import api as weave
 from ...legacy import node_ref
 
 
diff --git a/weave/tests/legacy/test_op_def.py b/weave/tests/legacy/test_op_def.py
index aa24f185282..cb05fe0135c 100644
--- a/weave/tests/legacy/test_op_def.py
+++ b/weave/tests/legacy/test_op_def.py
@@ -5,10 +5,9 @@
 
 import pytest
 
+from weave.legacy import api as weave
 from weave.legacy import context_state
 
-from ... import api as weave
-
 _loading_builtins_token = context_state.set_loading_built_ins()
 
 
diff --git a/weave/tests/legacy/test_partial_object.py b/weave/tests/legacy/test_partial_object.py
index 8096e6061b5..22eac741e3b 100644
--- a/weave/tests/legacy/test_partial_object.py
+++ b/weave/tests/legacy/test_partial_object.py
@@ -1,3 +1,4 @@
+from weave.legacy import api as weave
 from weave.legacy import compile, ops
 from weave.legacy import weave_types as types
 from weave.legacy.language_features.tagging.tagged_value_type import TaggedValueType
@@ -5,7 +6,6 @@
 from weave.legacy.ops_domain.project_ops import root_all_projects
 from weave.legacy.ops_domain.report_ops import root_all_reports
 
-from ... import api as weave
 from .test_wb import table_mock1_no_display_name
 
 
diff --git a/weave/tests/legacy/test_run_segment.py b/weave/tests/legacy/test_run_segment.py
index 6ab3815c200..fd92b66c552 100644
--- a/weave/tests/legacy/test_run_segment.py
+++ b/weave/tests/legacy/test_run_segment.py
@@ -6,13 +6,11 @@
 import pytest
 
 import weave
-from weave.legacy import ops, storage, weave_internal
+from weave.legacy import api, ops, storage, weave_internal
 from weave.legacy import weave_types as types
 from weave.legacy.ops_arrow import ArrowWeaveList, arrow_as_array
 from weave.legacy.ops_domain.run_segment import RunSegment
 
-from ... import api
-
 N_NUMERIC_METRICS = 99  # number of numerical columns in the metrics table
 
 
diff --git a/weave/tests/legacy/test_storage.py b/weave/tests/legacy/test_storage.py
index ab91835e742..29dda1e8963 100644
--- a/weave/tests/legacy/test_storage.py
+++ b/weave/tests/legacy/test_storage.py
@@ -7,6 +7,7 @@
 import pytest
 import wandb
 
+from weave.legacy import api as weave
 from weave.legacy import (
     artifact_mem,
     artifact_wandb,
@@ -18,7 +19,6 @@
 from weave.legacy.arrow import list_ as arrow
 from weave.legacy.weave_internal import make_const_node
 
-from ... import api as weave
 from ...legacy.weavejs_fixes import recursively_unwrap_unions
 from . import test_helpers
 
diff --git a/weave/tests/legacy/test_table_ops.py b/weave/tests/legacy/test_table_ops.py
index c5ece54c3ab..46f4af2e436 100644
--- a/weave/tests/legacy/test_table_ops.py
+++ b/weave/tests/legacy/test_table_ops.py
@@ -3,6 +3,7 @@
 
 import pytest
 
+from weave.legacy import api as weave
 from weave.legacy import (
     box,
     context,
@@ -15,7 +16,6 @@
 from weave.legacy import weave_types as types
 from weave.legacy.ops_domain import table as table_ops
 
-from ... import api as weave
 from .. import weavejs_ops
 
 TABLE_TYPES = ["list", "pandas", "sql"]
diff --git a/weave/tests/legacy/test_trace.py b/weave/tests/legacy/test_trace.py
index 629ca15a4c3..a0c308aab12 100644
--- a/weave/tests/legacy/test_trace.py
+++ b/weave/tests/legacy/test_trace.py
@@ -1,9 +1,9 @@
 import re
 
+from weave.legacy import api as weave
 from weave.legacy import graph, storage
 from weave.legacy.weave_internal import make_const_node
 
-from ... import api as weave
 from ...legacy import trace_legacy
 
 
diff --git a/weave/tests/legacy/test_wb_domain_types.py b/weave/tests/legacy/test_wb_domain_types.py
index b035ebd3b7a..ed550677062 100644
--- a/weave/tests/legacy/test_wb_domain_types.py
+++ b/weave/tests/legacy/test_wb_domain_types.py
@@ -1,9 +1,8 @@
+from weave.legacy import api as weave
 from weave.legacy import storage
 from weave.legacy import weave_types as types
 from weave.legacy.ops_domain import wb_domain_types as wdt
 
-from ... import api as weave
-
 
 def test_with_keys_assignability():
     org_type = wdt.OrgType
diff --git a/weave/tests/legacy/test_weavejs_fixes.py b/weave/tests/legacy/test_weavejs_fixes.py
index c40a9d3145f..94bc36039a8 100644
--- a/weave/tests/legacy/test_weavejs_fixes.py
+++ b/weave/tests/legacy/test_weavejs_fixes.py
@@ -3,6 +3,7 @@
 import pytest
 
 from weave.legacy import (
+    api,
     context_state,
     mappers_python,
     ops,
@@ -11,8 +12,6 @@
 )
 from weave.legacy import weave_types as types
 
-from ... import api
-
 
 @pytest.mark.skip(
     "calling custom ops with graph engine is broken right now. Not needed for weaveflow"
diff --git a/weave/tests/list_arrow_test_helpers.py b/weave/tests/list_arrow_test_helpers.py
index 8c2fedcd4a9..ab7e19d8eb1 100644
--- a/weave/tests/list_arrow_test_helpers.py
+++ b/weave/tests/list_arrow_test_helpers.py
@@ -1,8 +1,7 @@
+from weave.legacy import api as weave
 from weave.legacy import ops_arrow as arrow
 from weave.legacy.ops_primitives import list_
 
-from .. import api as weave
-
 
 class ListLikeNodeInterface:
     @staticmethod
diff --git a/weave/tests/trace/test_server.py b/weave/tests/trace/test_server.py
index aaaf1f08ff7..936eb7aa360 100644
--- a/weave/tests/trace/test_server.py
+++ b/weave/tests/trace/test_server.py
@@ -6,7 +6,7 @@
 import pytest
 import requests
 
-from weave import api as weave
+from weave.legacy import api as weave
 from weave.legacy import client as _client
 from weave.legacy import context_state, ops
 from weave.legacy import server as _server
diff --git a/weave/trace/cli.py b/weave/trace/cli.py
index 59a6118d12c..2204e2c4c76 100644
--- a/weave/trace/cli.py
+++ b/weave/trace/cli.py
@@ -5,10 +5,9 @@
 
 from weave import __version__
 from weave.deploy import gcp as google
+from weave.trace import api
 from weave.trace.refs import ObjectRef, parse_uri
 
-from . import api
-
 # from .model_server import app
 
 # TODO: does this work?