update

api/core/ops/ops_trace_manager.py

@@ -16,11 +16,7 @@ from sqlalchemy.orm import Session
 from core.helper.encrypter import decrypt_token, encrypt_token, obfuscated_token
 from core.ops.entities.config_entity import (
     OPS_FILE_PATH,
-    LangfuseConfig,
-    LangSmithConfig,
-    OpikConfig,
     TracingProviderEnum,
-    WeaveConfig,
 )
 from core.ops.entities.trace_entity import (
     DatasetRetrievalTraceInfo,
@@ -41,57 +37,59 @@ from models.workflow import WorkflowAppLog, WorkflowRun
 from tasks.ops_trace_task import process_trace_tasks
 
 
-def build_langfuse_trace_instance(config: LangfuseConfig):
+class OpsTraceProviderConfigMap(dict[str, dict[str, Any]]):
+    def __getitem__(self, provider: str) -> dict[str, Any]:
+        match provider:
+            case TracingProviderEnum.LANGFUSE:
+                from core.ops.entities.config_entity import LangfuseConfig
                 from core.ops.langfuse_trace.langfuse_trace import LangFuseDataTrace
 
-    return LangFuseDataTrace(config)
+                return {
-
-
-def build_langsmith_trace_instance(config: LangfuseConfig):
-    from core.ops.langfuse_trace.langfuse_trace import LangFuseDataTrace
-
-    return LangFuseDataTrace(config)
-
-
-def build_opik_trace_instance(config: LangSmithConfig):
-    from core.ops.langsmith_trace.langsmith_trace import LangSmithDataTrace
-
-    return LangSmithDataTrace(config)
-
-
-def build_weave_trace_instance(config: WeaveConfig):
-    from core.ops.weave_trace.weave_trace import WeaveDataTrace
-
-    return WeaveDataTrace(config)
-
-
-provider_config_map: dict[str, dict[str, Any]] = {
-    TracingProviderEnum.LANGFUSE: {
                     "config_class": LangfuseConfig,
                     "secret_keys": ["public_key", "secret_key"],
                     "other_keys": ["host", "project_key"],
-        "trace_instance": lambda config: build_langfuse_trace_instance(config),
+                    "trace_instance": LangFuseDataTrace,
-    },
+                }
-    TracingProviderEnum.LANGSMITH: {
+
+            case TracingProviderEnum.LANGSMITH:
+                from core.ops.entities.config_entity import LangSmithConfig
+                from core.ops.langsmith_trace.langsmith_trace import LangSmithDataTrace
+
+                return {
                     "config_class": LangSmithConfig,
                     "secret_keys": ["api_key"],
                     "other_keys": ["project", "endpoint"],
-        "trace_instance": lambda config: build_langsmith_trace_instance(config),
+                    "trace_instance": LangSmithDataTrace,
-    },
+                }
-    TracingProviderEnum.OPIK: {
+
+            case TracingProviderEnum.OPIK:
+                from core.ops.entities.config_entity import OpikConfig
+                from core.ops.opik_trace.opik_trace import OpikDataTrace
+
+                return {
                     "config_class": OpikConfig,
                     "secret_keys": ["api_key"],
                     "other_keys": ["project", "url", "workspace"],
-        "trace_instance": lambda config: build_opik_trace_instance(config),
+                    "trace_instance": OpikDataTrace,
-    },
+                }
-    TracingProviderEnum.WEAVE: {
+
+            case TracingProviderEnum.WEAVE:
+                from core.ops.entities.config_entity import WeaveConfig
+                from core.ops.weave_trace.weave_trace import WeaveDataTrace
+
+                return {
                     "config_class": WeaveConfig,
                     "secret_keys": ["api_key"],
                     "other_keys": ["project", "entity", "endpoint"],
-        "trace_instance": lambda config: build_weave_trace_instance(config),
+                    "trace_instance": WeaveDataTrace,
-    },
                 }
 
+            case _:
+                raise KeyError(f"Unsupported tracing provider: {provider}")
+
+
+provider_config_map: dict[str, dict[str, Any]] = OpsTraceProviderConfigMap()
+
 
 class OpsTraceManager:
     ops_trace_instances_cache: LRUCache = LRUCache(maxsize=128)