From 660d4e5294caa094520270c26ef48b8a80de2a20 Mon Sep 17 00:00:00 2001 From: Joe <1264204425@qq.com> Date: Thu, 20 Jun 2024 14:31:37 +0800 Subject: [PATCH] feat: add ops trace encrypt config decrypt_config obfuscate_config --- api/controllers/console/app/app.py | 13 ++- api/controllers/console/app/ops_trace.py | 1 - api/services/ops_trace/__init__.py | 0 api/services/ops_trace/langfuse_trace.py | 30 ++++++- api/services/ops_trace/langsmith_trace.py | 25 +++++- api/services/ops_trace/model.py | 26 ++++++ api/services/ops_trace/ops_trace_service.py | 89 +++------------------ 7 files changed, 90 insertions(+), 94 deletions(-) create mode 100644 api/services/ops_trace/__init__.py create mode 100644 api/services/ops_trace/model.py diff --git a/api/controllers/console/app/app.py b/api/controllers/console/app/app.py index a9e8b5fb4a2d33..2f49222450490f 100644 --- a/api/controllers/console/app/app.py +++ b/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"} diff --git a/api/controllers/console/app/ops_trace.py b/api/controllers/console/app/ops_trace.py index 1e96b6f98245e3..2a0681a9e60c26 100644 --- a/api/controllers/console/app/ops_trace.py +++ b/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"} diff --git a/api/services/ops_trace/__init__.py b/api/services/ops_trace/__init__.py new file mode 100644 index 00000000000000..e69de29bb2d1d6 diff --git a/api/services/ops_trace/langfuse_trace.py b/api/services/ops_trace/langfuse_trace.py index 485bb028bb27bd..9a565ec96406ff 100644 --- a/api/services/ops_trace/langfuse_trace.py +++ b/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) diff --git a/api/services/ops_trace/langsmith_trace.py b/api/services/ops_trace/langsmith_trace.py index a85d729854a11e..fa7d36aa85d498 100644 --- a/api/services/ops_trace/langsmith_trace.py +++ b/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) diff --git a/api/services/ops_trace/model.py b/api/services/ops_trace/model.py new file mode 100644 index 00000000000000..51dca08137e773 --- /dev/null +++ b/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 diff --git a/api/services/ops_trace/ops_trace_service.py b/api/services/ops_trace/ops_trace_service.py index ec2bc8d50d25b8..69dea3771fc6e4 100644 --- a/api/services/ops_trace/ops_trace_service.py +++ b/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,