feat: add ops trace encrypt config decrypt_config obfuscate_config

api/controllers/console/app/app.py

@@ -306,14 +306,11 @@ def post(self, app_id):
         parser.add_argument('tracing_provider', type=str, required=True, location='json')
         args = parser.parse_args()
 
-        try:
-            OpsTraceService.update_app_tracing_config(
-                app_id=app_id,
-                enabled=args['enabled'],
-                tracing_provider=args['tracing_provider'],
-            )
-        except Exception as e:
-            raise e
+        OpsTraceService.update_app_tracing_config(
+            app_id=app_id,
+            enabled=args['enabled'],
+            tracing_provider=args['tracing_provider'],
+        )
 
         return {"result": "success"}
 

api/controllers/console/app/ops_trace.py

@@ -71,7 +71,6 @@ def patch(self, app_id):
                 tracing_provider=args['tracing_provider'],
                 tracing_config=args['tracing_config']
             )
-            print("==============", result)
             if not result:
                 raise TracingConfigNotExist()
             return {"result": "success"}

api/services/ops_trace/langfuse_trace.py

@@ -1,4 +1,5 @@
 import json
+import logging
 import os
 from datetime import datetime, timedelta
 from enum import Enum
@@ -8,14 +9,17 @@
 from pydantic import BaseModel, Field, field_validator
 from pydantic_core.core_schema import ValidationInfo
 
+from core.helper.encrypter import decrypt_token, encrypt_token, obfuscated_token
 from core.moderation.base import ModerationInputsResult
 from extensions.ext_database import db
 from models.dataset import Document
 from models.model import Message, MessageAgentThought, MessageFile
 from models.workflow import WorkflowNodeExecution, WorkflowRun
 from services.ops_trace.base_trace_instance import BaseTraceInstance
+from services.ops_trace.model import LangfuseConfig
 from services.ops_trace.utils import filter_none_values, replace_text_with_content
 
+logger = logging.getLogger(__name__)
 
 def validate_input_output(v, field_name):
     """
@@ -696,7 +700,7 @@ def add_trace(self, langfuse_trace_data: Optional[LangfuseTrace] = None):
         )
         try:
             self.langfuse_client.trace(**format_trace_data)
-            print("LangFuse Trace created successfully")
+            logger.debug("LangFuse Trace created successfully")
         except Exception as e:
             raise f"LangFuse Failed to create trace: {str(e)}"
 
@@ -706,7 +710,7 @@ def add_span(self, langfuse_span_data: Optional[LangfuseSpan] = None):
         )
         try:
             self.langfuse_client.span(**format_span_data)
-            print("LangFuse Span created successfully")
+            logger.debug("LangFuse Span created successfully")
         except Exception as e:
             raise f"LangFuse Failed to create span: {str(e)}"
 
@@ -727,7 +731,7 @@ def add_generation(
         )
         try:
             self.langfuse_client.generation(**format_generation_data)
-            print("LangFuse Generation created successfully")
+            logger.debug("LangFuse Generation created successfully")
         except Exception as e:
             raise f"LangFuse Failed to create generation: {str(e)}"
 
@@ -746,5 +750,23 @@ def api_check(self):
         try:
             return self.langfuse_client.auth_check()
         except Exception as e:
-            print(f"LangFuse API check failed: {str(e)}")
+            logger.debug(f"LangFuse API check failed: {str(e)}")
             return False
+
+    @classmethod
+    def obfuscate_config(cls, config: LangfuseConfig):
+        public_key = obfuscated_token(config.public_key)
+        secret_key = obfuscated_token(config.secret_key)
+        return LangfuseConfig(public_key=public_key, secret_key=secret_key, host=config.host)
+
+    @classmethod
+    def encrypt_config(cls, tenant_id, config: LangfuseConfig):
+        decrypt_public_key = encrypt_token(tenant_id, config.public_key)
+        decrypt_secret_key = encrypt_token(tenant_id, config.secret_key)
+        return LangfuseConfig(public_key=decrypt_public_key, secret_key=decrypt_secret_key, host=config.host)
+
+    @classmethod
+    def decrypt_config(cls, tenant_id, config: LangfuseConfig):
+        decrypt_public_key = decrypt_token(tenant_id, config.public_key)
+        decrypt_secret_key = decrypt_token(tenant_id, config.secret_key)
+        return LangfuseConfig(public_key=decrypt_public_key, secret_key=decrypt_secret_key, host=config.host)

api/services/ops_trace/langsmith_trace.py

@@ -1,4 +1,5 @@
 import json
+import logging
 import os
 from datetime import datetime, timedelta
 from enum import Enum
@@ -8,14 +9,17 @@
 from pydantic import BaseModel, Field, field_validator
 from pydantic_core.core_schema import ValidationInfo
 
+from core.helper.encrypter import decrypt_token, encrypt_token, obfuscated_token
 from core.moderation.base import ModerationInputsResult
 from extensions.ext_database import db
 from models.dataset import Document
 from models.model import Message, MessageAgentThought, MessageFile
 from models.workflow import WorkflowNodeExecution, WorkflowRun
 from services.ops_trace.base_trace_instance import BaseTraceInstance
+from services.ops_trace.model import LangSmithConfig
 from services.ops_trace.utils import filter_none_values, replace_text_with_content
 
+logger = logging.getLogger(__name__)
 
 class LangSmithRunType(str, Enum):
     tool = "tool"
@@ -588,7 +592,7 @@ def add_run(self, run_data: LangSmithRunModel):
         data = filter_none_values(data)
         try:
             self.langsmith_client.create_run(**data)
-            print("LangSmith Run created successfully.")
+            logger.debug("LangSmith Run created successfully.")
         except Exception as e:
             raise f"LangSmith Failed to create run: {str(e)}"
 
@@ -597,7 +601,7 @@ def update_run(self, update_run_data: LangSmithRunUpdateModel):
         data = filter_none_values(data)
         try:
             self.langsmith_client.update_run(**data)
-            print("LangSmith Run updated successfully.")
+            logger.debug("LangSmith Run updated successfully.")
         except Exception as e:
             raise f"LangSmith Failed to update run: {str(e)}"
 
@@ -608,5 +612,20 @@ def api_check(self):
             self.langsmith_client.delete_project(project_name=random_project_name)
             return True
         except Exception as e:
-            print(f"LangSmith API check failed: {str(e)}")
+            logger.debug(f"LangSmith API check failed: {str(e)}")
             return False
+
+    @classmethod
+    def obfuscate_config(cls, config: LangSmithConfig):
+        api_key = obfuscated_token(config.api_key)
+        return LangSmithConfig(api_key=api_key, project=config.project, endpoint=config.endpoint)
+
+    @classmethod
+    def encrypt_config(cls, tenant_id, config: LangSmithConfig):
+        api_key = encrypt_token(tenant_id, config.api_key)
+        return LangSmithConfig(api_key=api_key, project=config.project, endpoint=config.endpoint)
+
+    @classmethod
+    def decrypt_config(cls, tenant_id, config: LangSmithConfig):
+        api_key = decrypt_token(tenant_id, config.api_key)
+        return LangSmithConfig(api_key=api_key, project=config.project, endpoint=config.endpoint)

api/services/ops_trace/model.py

@@ -0,0 +1,26 @@
+from enum import Enum
+
+from pydantic import BaseModel
+
+
+class TracingProviderEnum(Enum):
+    LANGFUSE = 'langfuse'
+    LANGSMITH = 'langsmith'
+
+
+class LangfuseConfig(BaseModel):
+    """
+    Model class for Langfuse tracing config.
+    """
+    public_key: str
+    secret_key: str
+    host: str
+
+
+class LangSmithConfig(BaseModel):
+    """
+    Model class for Langsmith tracing config.
+    """
+    api_key: str
+    project: str
+    endpoint: str

api/services/ops_trace/ops_trace_service.py

@@ -1,38 +1,12 @@
 import json
-from enum import Enum
 from typing import Union
 from uuid import UUID
 
-from pydantic import BaseModel
-
-from core.helper.encrypter import decrypt_token, encrypt_token, obfuscated_token
 from extensions.ext_database import db
 from models.model import App, AppModelConfig, Conversation, Message, TraceAppConfig
 from services.ops_trace.langfuse_trace import LangFuseDataTrace
 from services.ops_trace.langsmith_trace import LangSmithDataTrace
-
-
-class TracingProviderEnum(Enum):
-    LANGFUSE = 'langfuse'
-    LANGSMITH = 'langsmith'
-
-
-class LangfuseConfig(BaseModel):
-    """
-    Model class for Langfuse tracing config.
-    """
-    public_key: str
-    secret_key: str
-    host: str
-
-
-class LangSmithConfig(BaseModel):
-    """
-    Model class for Langsmith tracing config.
-    """
-    api_key: str
-    project: str
-    endpoint: str
+from services.ops_trace.model import LangfuseConfig, LangSmithConfig, TracingProviderEnum
 
 
 class OpsTraceService:
@@ -162,25 +136,12 @@ def encrypt_tracing_config(cls, tenant_id: str, tracing_provider: str, tracing_c
         """
         if tracing_provider == TracingProviderEnum.LANGFUSE.value:
             tracing_config = LangfuseConfig(**tracing_config)
-            encrypt_public_key = encrypt_token(tenant_id, tracing_config.public_key)
-            encrypt_secret_key = encrypt_token(tenant_id, tracing_config.secret_key)
-            tracing_config = LangfuseConfig(
-                public_key=encrypt_public_key,
-                secret_key=encrypt_secret_key,
-                host=tracing_config.host
-            )
+            tracing_config = LangFuseDataTrace.encrypt_config(tenant_id, tracing_config)
         elif tracing_provider == TracingProviderEnum.LANGSMITH.value:
             tracing_config = LangSmithConfig(**tracing_config)
-            encrypt_api_key = encrypt_token(tenant_id, tracing_config.api_key)
-            tracing_config = LangSmithConfig(
-                api_key=encrypt_api_key,
-                project=tracing_config.project,
-                endpoint=tracing_config.endpoint
-            )
+            tracing_config = LangSmithDataTrace.encrypt_config(tenant_id, tracing_config)
 
-        if isinstance(tracing_config, BaseModel):
-            return tracing_config.dict()
-        return tracing_config
+        return tracing_config.model_dump()
 
     @classmethod
     def decrypt_tracing_config(cls, tenant_id: str, tracing_provider: str, tracing_config: dict):
@@ -193,25 +154,12 @@ def decrypt_tracing_config(cls, tenant_id: str, tracing_provider: str, tracing_c
         """
         if tracing_provider == TracingProviderEnum.LANGFUSE.value:
             tracing_config = LangfuseConfig(**tracing_config)
-            decrypt_public_key = decrypt_token(tenant_id, tracing_config.public_key)
-            decrypt_secret_key = decrypt_token(tenant_id, tracing_config.secret_key)
-            tracing_config = LangfuseConfig(
-                public_key=decrypt_public_key,
-                secret_key=decrypt_secret_key,
-                host=tracing_config.host
-            )
+            tracing_config = LangFuseDataTrace.decrypt_config(tenant_id, tracing_config)
         elif tracing_provider == TracingProviderEnum.LANGSMITH.value:
             tracing_config = LangSmithConfig(**tracing_config)
-            decrypt_api_key = decrypt_token(tenant_id, tracing_config.api_key)
-            tracing_config = LangSmithConfig(
-                api_key=decrypt_api_key,
-                project=tracing_config.project,
-                endpoint=tracing_config.endpoint
-            )
+            tracing_config = LangSmithDataTrace.decrypt_config(tenant_id, tracing_config)
 
-        if isinstance(tracing_config, BaseModel):
-            return tracing_config.dict()
-        return tracing_config
+        return tracing_config.model_dump()
 
     @classmethod
     def obfuscated_decrypt_token(cls, tracing_provider: str, decrypt_tracing_config:dict):
@@ -221,28 +169,14 @@ def obfuscated_decrypt_token(cls, tracing_provider: str, decrypt_tracing_config:
         :param decrypt_tracing_config: tracing config
         :return:
         """
+        obfuscate_config = None
         if tracing_provider == TracingProviderEnum.LANGFUSE.value:
             decrypt_tracing_config = LangfuseConfig(**decrypt_tracing_config)
-            decrypt_public_key = decrypt_tracing_config.public_key
-            decrypt_secret_key = decrypt_tracing_config.secret_key
-            obfuscated_public_key = obfuscated_token(decrypt_public_key)
-            obfuscated_secret_key = obfuscated_token(decrypt_secret_key)
-            decrypt_tracing_config = LangfuseConfig(
-                public_key=obfuscated_public_key,
-                secret_key=obfuscated_secret_key,
-                host=decrypt_tracing_config.host
-            )
+            obfuscate_config = LangFuseDataTrace.obfuscate_config(decrypt_tracing_config)
         elif tracing_provider == TracingProviderEnum.LANGSMITH.value:
             decrypt_tracing_config = LangSmithConfig(**decrypt_tracing_config)
-            decrypt_api_key = decrypt_tracing_config.api_key
-            obfuscated_api_key = obfuscated_token(decrypt_api_key)
-            decrypt_tracing_config = LangSmithConfig(
-                api_key=obfuscated_api_key,
-                project=decrypt_tracing_config.project,
-                endpoint=decrypt_tracing_config.endpoint
-            )
-
-        return decrypt_tracing_config.dict()
+            obfuscate_config = LangSmithDataTrace.obfuscate_config(decrypt_tracing_config)
+        return obfuscate_config.model_dump()
 
     @classmethod
     def get_decrypted_tracing_config(cls, app_id: str, tracing_provider: str):
@@ -321,7 +255,6 @@ def get_ops_trace_instance(
                 langsmith_api_key = decrypt_trace_config.get('api_key')
                 langsmith_project = decrypt_trace_config.get('project')
                 langsmith_endpoint = decrypt_trace_config.get('endpoint')
-                print(langsmith_api_key, langsmith_project, langsmith_endpoint)
                 tracing_instance = LangSmithDataTrace(
                     langsmith_api_key,
                     langsmith_project,