From 8ff732b589e7e2499fdb9a65812d76b87682d093 Mon Sep 17 00:00:00 2001 From: Andrew Truong Date: Wed, 28 Aug 2024 15:12:25 -0400 Subject: [PATCH] chore(weave): Legacy Refactor pt22 (#2240) * split: further split trace and legacy errors * split: make separate trace and legacy utils * delete: extraneous import --- weave/tests/trace/op_versioning_solo.py | 1 - weave/tests/trace/test_op_versioning.py | 2 +- weave/trace/api.py | 3 ++- weave/trace/client_context/weave_client.py | 2 +- weave/trace/errors.py | 27 +++++++++++++++++----- weave/trace/feedback.py | 3 +-- weave/trace/rich_pydantic_util.py | 2 +- weave/trace/serve_fastapi.py | 5 ++-- weave/trace/util.py | 25 ++++++++++++++++++++ weave/trace/weave_client.py | 8 ++----- weave/trace/weave_init.py | 3 +-- 11 files changed, 57 insertions(+), 24 deletions(-) diff --git a/weave/tests/trace/op_versioning_solo.py b/weave/tests/trace/op_versioning_solo.py index bacc0e27155..ac3afc6dd39 100644 --- a/weave/tests/trace/op_versioning_solo.py +++ b/weave/tests/trace/op_versioning_solo.py @@ -1,7 +1,6 @@ import numpy as np import weave -from weave.legacy.weave import artifact_fs @weave.op() diff --git a/weave/tests/trace/test_op_versioning.py b/weave/tests/trace/test_op_versioning.py index 3fc023e235c..b77d35a49d1 100644 --- a/weave/tests/trace/test_op_versioning.py +++ b/weave/tests/trace/test_op_versioning.py @@ -5,7 +5,7 @@ import pytest import weave -from weave.legacy.weave import artifact_fs, derive_op, op_def +from weave.legacy.weave import artifact_fs, derive_op from weave.trace_server.trace_server_interface import FileContentReadReq, ObjReadReq diff --git a/weave/trace/api.py b/weave/trace/api.py index 1c7dbbb4bb0..cbb6fecacff 100644 --- a/weave/trace/api.py +++ b/weave/trace/api.py @@ -6,7 +6,8 @@ import time from typing import Any, Callable, Iterator, Optional, Union -from weave.legacy.weave import urls, util +from weave.legacy.weave import urls +from weave.trace import util from weave.trace.call_context import get_current_call from weave.trace.client_context import weave_client as weave_client_context diff --git a/weave/trace/client_context/weave_client.py b/weave/trace/client_context/weave_client.py index f74759f26ac..ee7e994d68d 100644 --- a/weave/trace/client_context/weave_client.py +++ b/weave/trace/client_context/weave_client.py @@ -2,7 +2,7 @@ from typing import TYPE_CHECKING, Optional from weave.legacy.weave import context_state -from weave.legacy.weave.errors import WeaveInitError +from weave.trace.errors import WeaveInitError if TYPE_CHECKING: from weave.trace.weave_client import WeaveClient diff --git a/weave/trace/errors.py b/weave/trace/errors.py index 6b5617614fb..752ac759e62 100644 --- a/weave/trace/errors.py +++ b/weave/trace/errors.py @@ -1,10 +1,25 @@ -class Error(Exception): - pass +class Error(Exception): ... -class InternalError(Error): - pass +class InternalError(Error): ... -class OpCallError(Error): - pass +class OpCallError(Error): ... + + +class WeaveTypeError(Error): ... + + +class WeaveSerializeError(Error): ... + + +class WeaveOpSerializeError(WeaveSerializeError): ... + + +class WeaveInitError(Error): ... + + +class WeaveDefinitionError(Error): ... + + +class WeaveWandbAuthenticationException(Error): ... diff --git a/weave/trace/feedback.py b/weave/trace/feedback.py index 99893df018d..312903f1921 100644 --- a/weave/trace/feedback.py +++ b/weave/trace/feedback.py @@ -5,8 +5,7 @@ from rich.table import Table -from weave.legacy.weave import util -from weave.trace import rich_pydantic_util +from weave.trace import rich_pydantic_util, 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 diff --git a/weave/trace/rich_pydantic_util.py b/weave/trace/rich_pydantic_util.py index ec88c0545db..03ce95efe9d 100644 --- a/weave/trace/rich_pydantic_util.py +++ b/weave/trace/rich_pydantic_util.py @@ -6,7 +6,7 @@ from rich.console import Console from rich.table import Table -from weave.legacy.weave import util +from weave.trace import util def dict_to_table(d: dict[str, Any]) -> Table: diff --git a/weave/trace/serve_fastapi.py b/weave/trace/serve_fastapi.py index 98adc7537b1..9351f968e7b 100644 --- a/weave/trace/serve_fastapi.py +++ b/weave/trace/serve_fastapi.py @@ -12,13 +12,12 @@ except ImportError: from typing_extensions import Annotated # type: ignore -from weave.legacy.weave import cache, op_args, pyfunc_type_util +from weave.legacy.weave import cache, op_args, pyfunc_type_util, weave_pydantic from weave.legacy.weave.wandb_api import WandbApiAsync +from weave.trace import errors from weave.trace.op import Op from weave.trace.refs import ObjectRef -from ..legacy.weave import errors, weave_pydantic - key_cache: cache.LruTimeWindowCache[str, typing.Optional[bool]] = ( cache.LruTimeWindowCache(datetime.timedelta(minutes=5)) ) diff --git a/weave/trace/util.py b/weave/trace/util.py index 9631569a6e4..b8fe3e25120 100644 --- a/weave/trace/util.py +++ b/weave/trace/util.py @@ -94,6 +94,31 @@ def run(self) -> None: self.context.run(super().run) +def is_colab(): # type: ignore + import importlib + + spec = importlib.util.find_spec("google.colab") + return bool(spec) + + +def is_notebook() -> bool: + if is_colab(): # type: ignore[no-untyped-call] + return True + try: + from IPython import get_ipython + except ImportError: + return False + else: + ip = get_ipython() + if ip is None: + return False + if "IPKernelApp" not in ip.config: + return False + # if "VSCODE_PID" in os.environ: + # return False + return True + + # rename for cleaner export ThreadPoolExecutor = ContextAwareThreadPoolExecutor Thread = ContextAwareThread diff --git a/weave/trace/weave_client.py b/weave/trace/weave_client.py index a16884b955d..6097a6d918a 100644 --- a/weave/trace/weave_client.py +++ b/weave/trace/weave_client.py @@ -51,10 +51,6 @@ TraceServerInterface, ) -if typing.TYPE_CHECKING: - from ..legacy.weave import ref_base - - # Controls if objects can have refs to projects not the WeaveClient project. # If False, object refs with with mismatching projects will be recreated. # If True, use existing ref to object in other project. @@ -904,10 +900,10 @@ def _set_call_display_name( def _remove_call_display_name(self, call: Call) -> None: self._set_call_display_name(call, None) - def _ref_input_to(self, ref: "ref_base.Ref") -> Sequence[Call]: + def _ref_input_to(self, ref: Ref) -> Sequence[Call]: raise NotImplementedError() - def _ref_value_input_to(self, ref: "ref_base.Ref") -> list[Call]: + def _ref_value_input_to(self, ref: Ref) -> list[Call]: raise NotImplementedError() def _ref_output_of(self, ref: ObjectRef) -> typing.Optional[Call]: diff --git a/weave/trace/weave_init.py b/weave/trace/weave_init.py index 55188cad443..603cf0b3592 100644 --- a/weave/trace/weave_init.py +++ b/weave/trace/weave_init.py @@ -1,7 +1,6 @@ import typing -from weave.legacy.weave import errors -from weave.trace import autopatch, init_message, trace_sentry, weave_client +from weave.trace import autopatch, errors, 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