From 34958baf4a8861b309cc49ded8eab177cce25323 Mon Sep 17 00:00:00 2001
From: Andrew Truong <itsandrewtruong@gmail.com>
Date: Mon, 26 Aug 2024 09:34:39 -0400
Subject: [PATCH] split: further split trace and legacy env

---
 weave/trace/env.py     | 4 ++++
 weave/trace/op_type.py | 2 +-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/weave/trace/env.py b/weave/trace/env.py
index 0d75c3c93f9..90396ed245f 100644
--- a/weave/trace/env.py
+++ b/weave/trace/env.py
@@ -5,3 +5,7 @@
 
 def get_weave_parallelism() -> int:
     return int(os.getenv(WEAVE_PARALLELISM, "20"))
+
+
+def wandb_production() -> bool:
+    return os.getenv("WEAVE_ENV") == "wandb_production"
diff --git a/weave/trace/op_type.py b/weave/trace/op_type.py
index c570e60436c..6705ac98c35 100644
--- a/weave/trace/op_type.py
+++ b/weave/trace/op_type.py
@@ -14,6 +14,7 @@
 from typing import Any, Callable, Optional, Union, get_args, get_origin
 
 from weave.legacy.weave import artifact_fs, context_state, errors, storage
+from weave.trace import env as environment
 from weave.trace.ipython import (
     ClassNotFoundError,
     get_class_source,
@@ -22,7 +23,6 @@
 from weave.trace.op import Op
 from weave.trace.refs import ObjectRef
 
-from ..legacy.weave import environment
 from . import serializer
 
 WEAVE_OP_PATTERN = re.compile(r"@weave\.op(\(\))?")