From 79964f0b3278149315d220ff0d045361765ddebe Mon Sep 17 00:00:00 2001
From: Bharat Ramanathan <bharat.ramanathan@wandb.com>
Date: Mon, 24 Feb 2025 15:35:02 +0530
Subject: [PATCH] feat: add weave to tracing providers configs and manager

---
 api/core/ops/entities/config_entity.py | 21 +++++++++++++++++++++
 api/core/ops/ops_trace_manager.py      |  8 ++++++++
 2 files changed, 29 insertions(+)

diff --git a/api/core/ops/entities/config_entity.py b/api/core/ops/entities/config_entity.py
index b484242b61..142cd85f31 100644
--- a/api/core/ops/entities/config_entity.py
+++ b/api/core/ops/entities/config_entity.py
@@ -7,6 +7,7 @@ class TracingProviderEnum(Enum):
     LANGFUSE = "langfuse"
     LANGSMITH = "langsmith"
     OPIK = "opik"
+    WEAVE = "weave"
 
 
 class BaseTracingConfig(BaseModel):
@@ -87,6 +88,26 @@ class OpikConfig(BaseTracingConfig):
 
         return v
 
+class WeaveConfig(BaseTracingConfig):
+    """
+    Model class for Weave tracing config.
+    """
+
+    api_key: str
+    entity: str | None = None
+    project: str
+    endpoint: str = "https://trace.wandb.ai"
+
+    @field_validator("endpoint")
+    @classmethod
+    def set_value(cls, v, info: ValidationInfo):
+        if v is None or v == "":
+            v = "https://trace.wandb.ai"
+        if not v.startswith("https://"):
+            raise ValueError("endpoint must start with https://")
+
+        return v
+
 
 OPS_FILE_PATH = "ops_trace/"
 OPS_TRACE_FAILED_KEY = "FAILED_OPS_TRACE"
diff --git a/api/core/ops/ops_trace_manager.py b/api/core/ops/ops_trace_manager.py
index c153e3f9dd..8f9741393a 100644
--- a/api/core/ops/ops_trace_manager.py
+++ b/api/core/ops/ops_trace_manager.py
@@ -18,6 +18,7 @@ from core.ops.entities.config_entity import (
     LangfuseConfig,
     LangSmithConfig,
     OpikConfig,
+    WeaveConfig,
     TracingProviderEnum,
 )
 from core.ops.entities.trace_entity import (
@@ -33,6 +34,7 @@ from core.ops.entities.trace_entity import (
 )
 from core.ops.langfuse_trace.langfuse_trace import LangFuseDataTrace
 from core.ops.langsmith_trace.langsmith_trace import LangSmithDataTrace
+from core.ops.weave_trace.weave_trace import WeaveDataTrace
 from core.ops.opik_trace.opik_trace import OpikDataTrace
 from core.ops.utils import get_message_data
 from extensions.ext_database import db
@@ -60,6 +62,12 @@ provider_config_map: dict[str, dict[str, Any]] = {
         "other_keys": ["project", "url", "workspace"],
         "trace_instance": OpikDataTrace,
     },
+    TracingProviderEnum.WEAVE.value: {
+        "config_class": WeaveConfig,
+        "secret_keys": ["api_key"],
+        "other_keys": ["project", "entity", "endpoint"],
+        "trace_instance": WeaveDataTrace,
+    },
 }