chore(weave): Legacy Refactor pt3

weave/conftest.py

@@ -26,10 +26,11 @@
     trace_server_interface as tsi,
 )
 
-from . import autopatch, environment, logs
+from . import environment, logs
 from .tests import fixture_fakewandb
 from .tests.trace.trace_server_clickhouse_conftest import *
 from .tests.wandb_system_tests_conftest import *
+from .trace import autopatch
 
 logs.configure_logger()
 

weave/integrations/langchain/langchain.py

@@ -35,12 +35,12 @@
 from contextvars import ContextVar
 from uuid import UUID
 
-from weave import call_context
 from weave.client_context import weave_client as weave_client_context
 from weave.integrations.integration_utilities import (
     make_pythonic_function_name,
     truncate_op_name,
 )
+from weave.trace import call_context
 from weave.trace.patcher import Patcher
 from weave.weave_client import Call
 

weave/legacy/monitoring/monitor.py

@@ -11,14 +11,14 @@
 import typing
 import uuid
 
-from weave import call_context, errors
+from weave import errors
 from weave.legacy import (
     graph,
     run_streamtable_span,
     stream_data_interfaces
 )
 from weave.legacy.wandb_interface.wandb_stream_table import StreamTable
-from weave.trace import context as trace_context
+from weave.trace import call_context, context as trace_context
 from weave.client_context import weave_client as weave_client_context
 
 logger = logging.getLogger(__name__)

weave/legacy/monitoring/openai/openai.py

@@ -10,7 +10,7 @@
 from openai import AsyncStream, Stream
 from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
 from packaging import version
-from weave import call_context
+from weave.trace import call_context
 from weave.client_context import weave_client as weave_client_context
 from weave.legacy.monitoring.monitor import _get_global_monitor
 from weave.legacy.monitoring.openai.models import *

weave/parallelism.py

@@ -3,7 +3,6 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import Callable, Generator, Iterator, Optional, TypeVar
 
-from weave import call_context
 from weave.legacy import (
     cache,
     context,
@@ -13,6 +12,7 @@
     memo,
     wandb_api,
 )
+from weave.trace import call_context
 
 # Must be power of 2
 MAX_PARALLELISM = 16

weave/trace/autopatch.py

@@ -0,0 +1,53 @@
+"""Basic autopatching of trackable libraries.
+
+This module should not require any dependencies beyond the standard library. It should
+check if libraries are installed and imported and patch in the case that they are.
+"""
+
+
+def autopatch() -> None:
+    from ..integrations.anthropic.anthropic_sdk import anthropic_patcher
+    from ..integrations.cerebras.cerebras_sdk import cerebras_patcher
+    from ..integrations.cohere.cohere_sdk import cohere_patcher
+    from ..integrations.dspy.dspy_sdk import dspy_patcher
+    from ..integrations.groq.groq_sdk import groq_patcher
+    from ..integrations.langchain.langchain import langchain_patcher
+    from ..integrations.litellm.litellm import litellm_patcher
+    from ..integrations.llamaindex.llamaindex import llamaindex_patcher
+    from ..integrations.mistral.mistral import mistral_patcher
+    from ..integrations.openai.openai_sdk import openai_patcher
+
+    openai_patcher.attempt_patch()
+    mistral_patcher.attempt_patch()
+    litellm_patcher.attempt_patch()
+    llamaindex_patcher.attempt_patch()
+    langchain_patcher.attempt_patch()
+    anthropic_patcher.attempt_patch()
+    groq_patcher.attempt_patch()
+    dspy_patcher.attempt_patch()
+    cerebras_patcher.attempt_patch()
+    cohere_patcher.attempt_patch()
+
+
+def reset_autopatch() -> None:
+    from ..integrations.anthropic.anthropic_sdk import anthropic_patcher
+    from ..integrations.cerebras.cerebras_sdk import cerebras_patcher
+    from ..integrations.cohere.cohere_sdk import cohere_patcher
+    from ..integrations.dspy.dspy_sdk import dspy_patcher
+    from ..integrations.groq.groq_sdk import groq_patcher
+    from ..integrations.langchain.langchain import langchain_patcher
+    from ..integrations.litellm.litellm import litellm_patcher
+    from ..integrations.llamaindex.llamaindex import llamaindex_patcher
+    from ..integrations.mistral.mistral import mistral_patcher
+    from ..integrations.openai.openai_sdk import openai_patcher
+
+    openai_patcher.undo_patch()
+    mistral_patcher.undo_patch()
+    litellm_patcher.undo_patch()
+    llamaindex_patcher.undo_patch()
+    langchain_patcher.undo_patch()
+    anthropic_patcher.undo_patch()
+    groq_patcher.undo_patch()
+    dspy_patcher.undo_patch()
+    cerebras_patcher.undo_patch()
+    cohere_patcher.undo_patch()

weave/call_context.py → weave/trace/call_context.py

@@ -6,7 +6,7 @@
 
 if typing.TYPE_CHECKING:
     # from .run import Run
-    from .weave_client import Call
+    from ..weave_client import Call
 
 _call_stack: contextvars.ContextVar[list["Call"]] = contextvars.ContextVar(
     "call", default=[]

weave/feedback.py → weave/trace/feedback.py

@@ -5,15 +5,15 @@
 
 from rich.table import Table
 
-from weave import rich_pydantic_util
 from weave.client_context import weave_client as weave_client_context
-from weave.refs import Refs
-from weave.rich_container import AbstractRichContainer
+from weave.trace import rich_pydantic_util
 from weave.trace.refs import parse_uri
+from weave.trace.rich_container import AbstractRichContainer
+from weave.trace.rich_refs import Refs
 from weave.trace_server import trace_server_interface as tsi
 from weave.trace_server.interface.query import Query
 
-from . import util
+from .. import util
 
 
 class Feedbacks(AbstractRichContainer[tsi.Feedback]):

weave/trace/op.py

@@ -17,10 +17,9 @@
     runtime_checkable,
 )
 
-from weave import call_context
 from weave.client_context import weave_client as weave_client_context
 from weave.legacy import context_state
-from weave.trace import box, settings
+from weave.trace import box, call_context, settings
 from weave.trace.context import call_attributes
 from weave.trace.errors import OpCallError
 from weave.trace.refs import ObjectRef

weave/rich_container.py → weave/trace/rich_container.py

@@ -3,7 +3,7 @@
 
 from rich.table import Table
 
-from weave import rich_pydantic_util
+from weave.trace import rich_pydantic_util
 
 T = TypeVar("T")
 

weave/rich_pydantic_util.py → weave/trace/rich_pydantic_util.py

@@ -6,7 +6,7 @@
 from rich.console import Console
 from rich.table import Table
 
-from . import util
+from .. import util
 
 
 def dict_to_table(d: dict[str, Any]) -> Table:

weave/refs.py → weave/trace/rich_refs.py

@@ -5,8 +5,8 @@
 from rich.table import Table
 
 from weave.client_context import weave_client as weave_client_context
-from weave.rich_container import AbstractRichContainer
 from weave.trace.refs import AnyRef, CallRef, parse_uri
+from weave.trace.rich_container import AbstractRichContainer
 from weave.trace.vals import WeaveObject
 
 

weave/trace_api.py

@@ -6,8 +6,8 @@
 import time
 from typing import Any, Callable, Iterator, Optional, Union
 
-from weave.call_context import get_current_call
 from weave.client_context import weave_client as weave_client_context
+from weave.trace.call_context import get_current_call
 
 # 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.

weave/weave_client.py

@@ -9,11 +9,12 @@
 import pydantic
 from requests import HTTPError
 
-from weave import call_context, trace_sentry, urls, version
+from weave import trace_sentry, urls, version
 from weave.client_context import weave_client as weave_client_context
 from weave.exception import exception_to_json_str
-from weave.feedback import FeedbackQuery, RefFeedbackQuery
 from weave.table import Table
+from weave.trace import call_context
+from weave.trace.feedback import FeedbackQuery, RefFeedbackQuery
 from weave.trace.object_record import (
     ObjectRecord,
     dataclass_object_record,

weave/weave_init.py

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