From 3b04a9ec4f36dc203fde9995e744326f624b7781 Mon Sep 17 00:00:00 2001 From: Andrew Truong Date: Tue, 27 Aug 2024 15:54:35 -0400 Subject: [PATCH 1/3] move to trace: autopatch,call_context --- weave/autopatch.py | 53 ----------------------- weave/conftest.py | 3 +- weave/integrations/langchain/langchain.py | 2 +- weave/legacy/monitoring/monitor.py | 4 +- weave/legacy/monitoring/openai/openai.py | 2 +- weave/parallelism.py | 2 +- weave/trace/autopatch.py | 53 +++++++++++++++++++++++ weave/{ => trace}/call_context.py | 2 +- weave/trace/op.py | 3 +- weave/trace_api.py | 2 +- weave/weave_client.py | 3 +- weave/weave_init.py | 3 +- 12 files changed, 67 insertions(+), 65 deletions(-) delete mode 100644 weave/autopatch.py create mode 100644 weave/trace/autopatch.py rename weave/{ => trace}/call_context.py (98%) diff --git a/weave/autopatch.py b/weave/autopatch.py deleted file mode 100644 index f9445881321c..000000000000 --- a/weave/autopatch.py +++ /dev/null @@ -1,53 +0,0 @@ -"""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() diff --git a/weave/conftest.py b/weave/conftest.py index b7b94fd2e702..0b401ce0ff42 100644 --- a/weave/conftest.py +++ b/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() diff --git a/weave/integrations/langchain/langchain.py b/weave/integrations/langchain/langchain.py index 2c5fc3c2e03a..52e7be8ac398 100644 --- a/weave/integrations/langchain/langchain.py +++ b/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 diff --git a/weave/legacy/monitoring/monitor.py b/weave/legacy/monitoring/monitor.py index 69fe290a023a..ff0e67ba3592 100644 --- a/weave/legacy/monitoring/monitor.py +++ b/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__) diff --git a/weave/legacy/monitoring/openai/openai.py b/weave/legacy/monitoring/openai/openai.py index 6e91314253d4..c002ec85c4ca 100644 --- a/weave/legacy/monitoring/openai/openai.py +++ b/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 * diff --git a/weave/parallelism.py b/weave/parallelism.py index dab68b09e09d..a11dcfeac76a 100644 --- a/weave/parallelism.py +++ b/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 diff --git a/weave/trace/autopatch.py b/weave/trace/autopatch.py new file mode 100644 index 000000000000..b80015c394e3 --- /dev/null +++ b/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() diff --git a/weave/call_context.py b/weave/trace/call_context.py similarity index 98% rename from weave/call_context.py rename to weave/trace/call_context.py index 2c3860e80c3d..74d9d9ed1e3c 100644 --- a/weave/call_context.py +++ b/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=[] diff --git a/weave/trace/op.py b/weave/trace/op.py index 37494d80ebe2..d217ce7c9a77 100644 --- a/weave/trace/op.py +++ b/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 diff --git a/weave/trace_api.py b/weave/trace_api.py index 22a8b87af603..a586cfe46395 100644 --- a/weave/trace_api.py +++ b/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. diff --git a/weave/weave_client.py b/weave/weave_client.py index d65275180d94..f2790e36f12b 100644 --- a/weave/weave_client.py +++ b/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.object_record import ( ObjectRecord, dataclass_object_record, diff --git a/weave/weave_init.py b/weave/weave_init.py index d2af6bf5eb77..6febcbc3c88b 100644 --- a/weave/weave_init.py +++ b/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 From 196df4d4a7f6f9d4fcde8536833f5277aa5717b6 Mon Sep 17 00:00:00 2001 From: Andrew Truong Date: Wed, 21 Aug 2024 19:02:21 -0400 Subject: [PATCH 2/3] move to trace: feedback --- weave/{ => trace}/feedback.py | 2 +- weave/weave_client.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename weave/{ => trace}/feedback.py (99%) diff --git a/weave/feedback.py b/weave/trace/feedback.py similarity index 99% rename from weave/feedback.py rename to weave/trace/feedback.py index feffab692342..5f911ed4e392 100644 --- a/weave/feedback.py +++ b/weave/trace/feedback.py @@ -13,7 +13,7 @@ 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]): diff --git a/weave/weave_client.py b/weave/weave_client.py index f2790e36f12b..1d136f7a9947 100644 --- a/weave/weave_client.py +++ b/weave/weave_client.py @@ -12,9 +12,9 @@ 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, From cc5cc4447b9cfd25d8ee563c2b8c9665bd1a9d4f Mon Sep 17 00:00:00 2001 From: Andrew Truong Date: Wed, 21 Aug 2024 20:08:17 -0400 Subject: [PATCH 3/3] move to trace: rich_* --- weave/trace/feedback.py | 6 +++--- weave/{ => trace}/rich_container.py | 2 +- weave/{ => trace}/rich_pydantic_util.py | 2 +- weave/{refs.py => trace/rich_refs.py} | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) rename weave/{ => trace}/rich_container.py (97%) rename weave/{ => trace}/rich_pydantic_util.py (98%) rename weave/{refs.py => trace/rich_refs.py} (95%) diff --git a/weave/trace/feedback.py b/weave/trace/feedback.py index 5f911ed4e392..01103f82d024 100644 --- a/weave/trace/feedback.py +++ b/weave/trace/feedback.py @@ -5,11 +5,11 @@ 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 diff --git a/weave/rich_container.py b/weave/trace/rich_container.py similarity index 97% rename from weave/rich_container.py rename to weave/trace/rich_container.py index 626e290a449e..a8a9565faf16 100644 --- a/weave/rich_container.py +++ b/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") diff --git a/weave/rich_pydantic_util.py b/weave/trace/rich_pydantic_util.py similarity index 98% rename from weave/rich_pydantic_util.py rename to weave/trace/rich_pydantic_util.py index 33114d9b92a5..1c2c1474ef71 100644 --- a/weave/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 . import util +from .. import util def dict_to_table(d: dict[str, Any]) -> Table: diff --git a/weave/refs.py b/weave/trace/rich_refs.py similarity index 95% rename from weave/refs.py rename to weave/trace/rich_refs.py index af6e858ffe23..c72e07f4227e 100644 --- a/weave/refs.py +++ b/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