您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符

ops_trace_manager.py 33KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853
  1. import json
  2. import logging
  3. import os
  4. import queue
  5. import threading
  6. import time
  7. from datetime import timedelta
  8. from typing import Any, Optional, Union
  9. from uuid import UUID, uuid4
  10. from cachetools import LRUCache
  11. from flask import current_app
  12. from sqlalchemy import select
  13. from sqlalchemy.orm import Session
  14. from core.helper.encrypter import decrypt_token, encrypt_token, obfuscated_token
  15. from core.ops.entities.config_entity import (
  16. OPS_FILE_PATH,
  17. TracingProviderEnum,
  18. )
  19. from core.ops.entities.trace_entity import (
  20. DatasetRetrievalTraceInfo,
  21. GenerateNameTraceInfo,
  22. MessageTraceInfo,
  23. ModerationTraceInfo,
  24. SuggestedQuestionTraceInfo,
  25. TaskData,
  26. ToolTraceInfo,
  27. TraceTaskName,
  28. WorkflowTraceInfo,
  29. )
  30. from core.ops.utils import get_message_data
  31. from extensions.ext_database import db
  32. from extensions.ext_storage import storage
  33. from models.model import App, AppModelConfig, Conversation, Message, MessageFile, TraceAppConfig
  34. from models.workflow import WorkflowAppLog, WorkflowRun
  35. from tasks.ops_trace_task import process_trace_tasks
  36. class OpsTraceProviderConfigMap(dict[str, dict[str, Any]]):
  37. def __getitem__(self, provider: str) -> dict[str, Any]:
  38. match provider:
  39. case TracingProviderEnum.LANGFUSE:
  40. from core.ops.entities.config_entity import LangfuseConfig
  41. from core.ops.langfuse_trace.langfuse_trace import LangFuseDataTrace
  42. return {
  43. "config_class": LangfuseConfig,
  44. "secret_keys": ["public_key", "secret_key"],
  45. "other_keys": ["host", "project_key"],
  46. "trace_instance": LangFuseDataTrace,
  47. }
  48. case TracingProviderEnum.LANGSMITH:
  49. from core.ops.entities.config_entity import LangSmithConfig
  50. from core.ops.langsmith_trace.langsmith_trace import LangSmithDataTrace
  51. return {
  52. "config_class": LangSmithConfig,
  53. "secret_keys": ["api_key"],
  54. "other_keys": ["project", "endpoint"],
  55. "trace_instance": LangSmithDataTrace,
  56. }
  57. case TracingProviderEnum.OPIK:
  58. from core.ops.entities.config_entity import OpikConfig
  59. from core.ops.opik_trace.opik_trace import OpikDataTrace
  60. return {
  61. "config_class": OpikConfig,
  62. "secret_keys": ["api_key"],
  63. "other_keys": ["project", "url", "workspace"],
  64. "trace_instance": OpikDataTrace,
  65. }
  66. case TracingProviderEnum.WEAVE:
  67. from core.ops.entities.config_entity import WeaveConfig
  68. from core.ops.weave_trace.weave_trace import WeaveDataTrace
  69. return {
  70. "config_class": WeaveConfig,
  71. "secret_keys": ["api_key"],
  72. "other_keys": ["project", "entity", "endpoint"],
  73. "trace_instance": WeaveDataTrace,
  74. }
  75. case _:
  76. raise KeyError(f"Unsupported tracing provider: {provider}")
  77. provider_config_map: dict[str, dict[str, Any]] = OpsTraceProviderConfigMap()
  78. class OpsTraceManager:
  79. ops_trace_instances_cache: LRUCache = LRUCache(maxsize=128)
  80. @classmethod
  81. def encrypt_tracing_config(
  82. cls, tenant_id: str, tracing_provider: str, tracing_config: dict, current_trace_config=None
  83. ):
  84. """
  85. Encrypt tracing config.
  86. :param tenant_id: tenant id
  87. :param tracing_provider: tracing provider
  88. :param tracing_config: tracing config dictionary to be encrypted
  89. :param current_trace_config: current tracing configuration for keeping existing values
  90. :return: encrypted tracing configuration
  91. """
  92. # Get the configuration class and the keys that require encryption
  93. config_class, secret_keys, other_keys = (
  94. provider_config_map[tracing_provider]["config_class"],
  95. provider_config_map[tracing_provider]["secret_keys"],
  96. provider_config_map[tracing_provider]["other_keys"],
  97. )
  98. new_config = {}
  99. # Encrypt necessary keys
  100. for key in secret_keys:
  101. if key in tracing_config:
  102. if "*" in tracing_config[key]:
  103. # If the key contains '*', retain the original value from the current config
  104. new_config[key] = current_trace_config.get(key, tracing_config[key])
  105. else:
  106. # Otherwise, encrypt the key
  107. new_config[key] = encrypt_token(tenant_id, tracing_config[key])
  108. for key in other_keys:
  109. new_config[key] = tracing_config.get(key, "")
  110. # Create a new instance of the config class with the new configuration
  111. encrypted_config = config_class(**new_config)
  112. return encrypted_config.model_dump()
  113. @classmethod
  114. def decrypt_tracing_config(cls, tenant_id: str, tracing_provider: str, tracing_config: dict):
  115. """
  116. Decrypt tracing config
  117. :param tenant_id: tenant id
  118. :param tracing_provider: tracing provider
  119. :param tracing_config: tracing config
  120. :return:
  121. """
  122. config_class, secret_keys, other_keys = (
  123. provider_config_map[tracing_provider]["config_class"],
  124. provider_config_map[tracing_provider]["secret_keys"],
  125. provider_config_map[tracing_provider]["other_keys"],
  126. )
  127. new_config = {}
  128. for key in secret_keys:
  129. if key in tracing_config:
  130. new_config[key] = decrypt_token(tenant_id, tracing_config[key])
  131. for key in other_keys:
  132. new_config[key] = tracing_config.get(key, "")
  133. return config_class(**new_config).model_dump()
  134. @classmethod
  135. def obfuscated_decrypt_token(cls, tracing_provider: str, decrypt_tracing_config: dict):
  136. """
  137. Decrypt tracing config
  138. :param tracing_provider: tracing provider
  139. :param decrypt_tracing_config: tracing config
  140. :return:
  141. """
  142. config_class, secret_keys, other_keys = (
  143. provider_config_map[tracing_provider]["config_class"],
  144. provider_config_map[tracing_provider]["secret_keys"],
  145. provider_config_map[tracing_provider]["other_keys"],
  146. )
  147. new_config = {}
  148. for key in secret_keys:
  149. if key in decrypt_tracing_config:
  150. new_config[key] = obfuscated_token(decrypt_tracing_config[key])
  151. for key in other_keys:
  152. new_config[key] = decrypt_tracing_config.get(key, "")
  153. return config_class(**new_config).model_dump()
  154. @classmethod
  155. def get_decrypted_tracing_config(cls, app_id: str, tracing_provider: str):
  156. """
  157. Get decrypted tracing config
  158. :param app_id: app id
  159. :param tracing_provider: tracing provider
  160. :return:
  161. """
  162. trace_config_data: Optional[TraceAppConfig] = (
  163. db.session.query(TraceAppConfig)
  164. .filter(TraceAppConfig.app_id == app_id, TraceAppConfig.tracing_provider == tracing_provider)
  165. .first()
  166. )
  167. if not trace_config_data:
  168. return None
  169. # decrypt_token
  170. app = db.session.query(App).filter(App.id == app_id).first()
  171. if not app:
  172. raise ValueError("App not found")
  173. tenant_id = app.tenant_id
  174. decrypt_tracing_config = cls.decrypt_tracing_config(
  175. tenant_id, tracing_provider, trace_config_data.tracing_config
  176. )
  177. return decrypt_tracing_config
  178. @classmethod
  179. def get_ops_trace_instance(
  180. cls,
  181. app_id: Optional[Union[UUID, str]] = None,
  182. ):
  183. """
  184. Get ops trace through model config
  185. :param app_id: app_id
  186. :return:
  187. """
  188. if isinstance(app_id, UUID):
  189. app_id = str(app_id)
  190. if app_id is None:
  191. return None
  192. app: Optional[App] = db.session.query(App).filter(App.id == app_id).first()
  193. if app is None:
  194. return None
  195. app_ops_trace_config = json.loads(app.tracing) if app.tracing else None
  196. if app_ops_trace_config is None:
  197. return None
  198. if not app_ops_trace_config.get("enabled"):
  199. return None
  200. tracing_provider = app_ops_trace_config.get("tracing_provider")
  201. if tracing_provider is None:
  202. return None
  203. try:
  204. provider_config_map[tracing_provider]
  205. except KeyError:
  206. return None
  207. # decrypt_token
  208. decrypt_trace_config = cls.get_decrypted_tracing_config(app_id, tracing_provider)
  209. if not decrypt_trace_config:
  210. return None
  211. trace_instance, config_class = (
  212. provider_config_map[tracing_provider]["trace_instance"],
  213. provider_config_map[tracing_provider]["config_class"],
  214. )
  215. decrypt_trace_config_key = str(decrypt_trace_config)
  216. tracing_instance = cls.ops_trace_instances_cache.get(decrypt_trace_config_key)
  217. if tracing_instance is None:
  218. # create new tracing_instance and update the cache if it absent
  219. tracing_instance = trace_instance(config_class(**decrypt_trace_config))
  220. cls.ops_trace_instances_cache[decrypt_trace_config_key] = tracing_instance
  221. logging.info(f"new tracing_instance for app_id: {app_id}")
  222. return tracing_instance
  223. @classmethod
  224. def get_app_config_through_message_id(cls, message_id: str):
  225. app_model_config = None
  226. message_data = db.session.query(Message).filter(Message.id == message_id).first()
  227. if not message_data:
  228. return None
  229. conversation_id = message_data.conversation_id
  230. conversation_data = db.session.query(Conversation).filter(Conversation.id == conversation_id).first()
  231. if not conversation_data:
  232. return None
  233. if conversation_data.app_model_config_id:
  234. app_model_config = (
  235. db.session.query(AppModelConfig)
  236. .filter(AppModelConfig.id == conversation_data.app_model_config_id)
  237. .first()
  238. )
  239. elif conversation_data.app_model_config_id is None and conversation_data.override_model_configs:
  240. app_model_config = conversation_data.override_model_configs
  241. return app_model_config
  242. @classmethod
  243. def update_app_tracing_config(cls, app_id: str, enabled: bool, tracing_provider: str):
  244. """
  245. Update app tracing config
  246. :param app_id: app id
  247. :param enabled: enabled
  248. :param tracing_provider: tracing provider
  249. :return:
  250. """
  251. # auth check
  252. try:
  253. provider_config_map[tracing_provider]
  254. except KeyError:
  255. raise ValueError(f"Invalid tracing provider: {tracing_provider}")
  256. app_config: Optional[App] = db.session.query(App).filter(App.id == app_id).first()
  257. if not app_config:
  258. raise ValueError("App not found")
  259. app_config.tracing = json.dumps(
  260. {
  261. "enabled": enabled,
  262. "tracing_provider": tracing_provider,
  263. }
  264. )
  265. db.session.commit()
  266. @classmethod
  267. def get_app_tracing_config(cls, app_id: str):
  268. """
  269. Get app tracing config
  270. :param app_id: app id
  271. :return:
  272. """
  273. app: Optional[App] = db.session.query(App).filter(App.id == app_id).first()
  274. if not app:
  275. raise ValueError("App not found")
  276. if not app.tracing:
  277. return {"enabled": False, "tracing_provider": None}
  278. app_trace_config = json.loads(app.tracing)
  279. return app_trace_config
  280. @staticmethod
  281. def check_trace_config_is_effective(tracing_config: dict, tracing_provider: str):
  282. """
  283. Check trace config is effective
  284. :param tracing_config: tracing config
  285. :param tracing_provider: tracing provider
  286. :return:
  287. """
  288. config_type, trace_instance = (
  289. provider_config_map[tracing_provider]["config_class"],
  290. provider_config_map[tracing_provider]["trace_instance"],
  291. )
  292. tracing_config = config_type(**tracing_config)
  293. return trace_instance(tracing_config).api_check()
  294. @staticmethod
  295. def get_trace_config_project_key(tracing_config: dict, tracing_provider: str):
  296. """
  297. get trace config is project key
  298. :param tracing_config: tracing config
  299. :param tracing_provider: tracing provider
  300. :return:
  301. """
  302. config_type, trace_instance = (
  303. provider_config_map[tracing_provider]["config_class"],
  304. provider_config_map[tracing_provider]["trace_instance"],
  305. )
  306. tracing_config = config_type(**tracing_config)
  307. return trace_instance(tracing_config).get_project_key()
  308. @staticmethod
  309. def get_trace_config_project_url(tracing_config: dict, tracing_provider: str):
  310. """
  311. get trace config is project key
  312. :param tracing_config: tracing config
  313. :param tracing_provider: tracing provider
  314. :return:
  315. """
  316. config_type, trace_instance = (
  317. provider_config_map[tracing_provider]["config_class"],
  318. provider_config_map[tracing_provider]["trace_instance"],
  319. )
  320. tracing_config = config_type(**tracing_config)
  321. return trace_instance(tracing_config).get_project_url()
  322. class TraceTask:
  323. def __init__(
  324. self,
  325. trace_type: Any,
  326. message_id: Optional[str] = None,
  327. workflow_run: Optional[WorkflowRun] = None,
  328. conversation_id: Optional[str] = None,
  329. user_id: Optional[str] = None,
  330. timer: Optional[Any] = None,
  331. **kwargs,
  332. ):
  333. self.trace_type = trace_type
  334. self.message_id = message_id
  335. self.workflow_run_id = workflow_run.id if workflow_run else None
  336. self.conversation_id = conversation_id
  337. self.user_id = user_id
  338. self.timer = timer
  339. self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001")
  340. self.app_id = None
  341. self.kwargs = kwargs
  342. def execute(self):
  343. return self.preprocess()
  344. def preprocess(self):
  345. preprocess_map = {
  346. TraceTaskName.CONVERSATION_TRACE: lambda: self.conversation_trace(**self.kwargs),
  347. TraceTaskName.WORKFLOW_TRACE: lambda: self.workflow_trace(
  348. workflow_run_id=self.workflow_run_id, conversation_id=self.conversation_id, user_id=self.user_id
  349. ),
  350. TraceTaskName.MESSAGE_TRACE: lambda: self.message_trace(message_id=self.message_id),
  351. TraceTaskName.MODERATION_TRACE: lambda: self.moderation_trace(
  352. message_id=self.message_id, timer=self.timer, **self.kwargs
  353. ),
  354. TraceTaskName.SUGGESTED_QUESTION_TRACE: lambda: self.suggested_question_trace(
  355. message_id=self.message_id, timer=self.timer, **self.kwargs
  356. ),
  357. TraceTaskName.DATASET_RETRIEVAL_TRACE: lambda: self.dataset_retrieval_trace(
  358. message_id=self.message_id, timer=self.timer, **self.kwargs
  359. ),
  360. TraceTaskName.TOOL_TRACE: lambda: self.tool_trace(
  361. message_id=self.message_id, timer=self.timer, **self.kwargs
  362. ),
  363. TraceTaskName.GENERATE_NAME_TRACE: lambda: self.generate_name_trace(
  364. conversation_id=self.conversation_id, timer=self.timer, **self.kwargs
  365. ),
  366. }
  367. return preprocess_map.get(self.trace_type, lambda: None)()
  368. # process methods for different trace types
  369. def conversation_trace(self, **kwargs):
  370. return kwargs
  371. def workflow_trace(
  372. self,
  373. *,
  374. workflow_run_id: str | None,
  375. conversation_id: str | None,
  376. user_id: str | None,
  377. ):
  378. if not workflow_run_id:
  379. return {}
  380. with Session(db.engine) as session:
  381. workflow_run_stmt = select(WorkflowRun).where(WorkflowRun.id == workflow_run_id)
  382. workflow_run = session.scalars(workflow_run_stmt).first()
  383. if not workflow_run:
  384. raise ValueError("Workflow run not found")
  385. workflow_id = workflow_run.workflow_id
  386. tenant_id = workflow_run.tenant_id
  387. workflow_run_id = workflow_run.id
  388. workflow_run_elapsed_time = workflow_run.elapsed_time
  389. workflow_run_status = workflow_run.status
  390. workflow_run_inputs = workflow_run.inputs_dict
  391. workflow_run_outputs = workflow_run.outputs_dict
  392. workflow_run_version = workflow_run.version
  393. error = workflow_run.error or ""
  394. total_tokens = workflow_run.total_tokens
  395. file_list = workflow_run_inputs.get("sys.file") or []
  396. query = workflow_run_inputs.get("query") or workflow_run_inputs.get("sys.query") or ""
  397. # get workflow_app_log_id
  398. workflow_app_log_data_stmt = select(WorkflowAppLog.id).where(
  399. WorkflowAppLog.tenant_id == tenant_id,
  400. WorkflowAppLog.app_id == workflow_run.app_id,
  401. WorkflowAppLog.workflow_run_id == workflow_run.id,
  402. )
  403. workflow_app_log_id = session.scalar(workflow_app_log_data_stmt)
  404. # get message_id
  405. message_id = None
  406. if conversation_id:
  407. message_data_stmt = select(Message.id).where(
  408. Message.conversation_id == conversation_id,
  409. Message.workflow_run_id == workflow_run_id,
  410. )
  411. message_id = session.scalar(message_data_stmt)
  412. metadata = {
  413. "workflow_id": workflow_id,
  414. "conversation_id": conversation_id,
  415. "workflow_run_id": workflow_run_id,
  416. "tenant_id": tenant_id,
  417. "elapsed_time": workflow_run_elapsed_time,
  418. "status": workflow_run_status,
  419. "version": workflow_run_version,
  420. "total_tokens": total_tokens,
  421. "file_list": file_list,
  422. "triggered_from": workflow_run.triggered_from,
  423. "user_id": user_id,
  424. }
  425. workflow_trace_info = WorkflowTraceInfo(
  426. workflow_data=workflow_run.to_dict(),
  427. conversation_id=conversation_id,
  428. workflow_id=workflow_id,
  429. tenant_id=tenant_id,
  430. workflow_run_id=workflow_run_id,
  431. workflow_run_elapsed_time=workflow_run_elapsed_time,
  432. workflow_run_status=workflow_run_status,
  433. workflow_run_inputs=workflow_run_inputs,
  434. workflow_run_outputs=workflow_run_outputs,
  435. workflow_run_version=workflow_run_version,
  436. error=error,
  437. total_tokens=total_tokens,
  438. file_list=file_list,
  439. query=query,
  440. metadata=metadata,
  441. workflow_app_log_id=workflow_app_log_id,
  442. message_id=message_id,
  443. start_time=workflow_run.created_at,
  444. end_time=workflow_run.finished_at,
  445. )
  446. return workflow_trace_info
  447. def message_trace(self, message_id: str | None):
  448. if not message_id:
  449. return {}
  450. message_data = get_message_data(message_id)
  451. if not message_data:
  452. return {}
  453. conversation_mode_stmt = select(Conversation.mode).where(Conversation.id == message_data.conversation_id)
  454. conversation_mode = db.session.scalars(conversation_mode_stmt).all()
  455. if not conversation_mode or len(conversation_mode) == 0:
  456. return {}
  457. conversation_mode = conversation_mode[0]
  458. created_at = message_data.created_at
  459. inputs = message_data.message
  460. # get message file data
  461. message_file_data = db.session.query(MessageFile).filter_by(message_id=message_id).first()
  462. file_list = []
  463. if message_file_data and message_file_data.url is not None:
  464. file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else ""
  465. file_list.append(file_url)
  466. metadata = {
  467. "conversation_id": message_data.conversation_id,
  468. "ls_provider": message_data.model_provider,
  469. "ls_model_name": message_data.model_id,
  470. "status": message_data.status,
  471. "from_end_user_id": message_data.from_end_user_id,
  472. "from_account_id": message_data.from_account_id,
  473. "agent_based": message_data.agent_based,
  474. "workflow_run_id": message_data.workflow_run_id,
  475. "from_source": message_data.from_source,
  476. "message_id": message_id,
  477. }
  478. message_tokens = message_data.message_tokens
  479. message_trace_info = MessageTraceInfo(
  480. message_id=message_id,
  481. message_data=message_data.to_dict(),
  482. conversation_model=conversation_mode,
  483. message_tokens=message_tokens,
  484. answer_tokens=message_data.answer_tokens,
  485. total_tokens=message_tokens + message_data.answer_tokens,
  486. error=message_data.error or "",
  487. inputs=inputs,
  488. outputs=message_data.answer,
  489. file_list=file_list,
  490. start_time=created_at,
  491. end_time=created_at + timedelta(seconds=message_data.provider_response_latency),
  492. metadata=metadata,
  493. message_file_data=message_file_data,
  494. conversation_mode=conversation_mode,
  495. )
  496. return message_trace_info
  497. def moderation_trace(self, message_id, timer, **kwargs):
  498. moderation_result = kwargs.get("moderation_result")
  499. if not moderation_result:
  500. return {}
  501. inputs = kwargs.get("inputs")
  502. message_data = get_message_data(message_id)
  503. if not message_data:
  504. return {}
  505. metadata = {
  506. "message_id": message_id,
  507. "action": moderation_result.action,
  508. "preset_response": moderation_result.preset_response,
  509. "query": moderation_result.query,
  510. }
  511. # get workflow_app_log_id
  512. workflow_app_log_id = None
  513. if message_data.workflow_run_id:
  514. workflow_app_log_data = (
  515. db.session.query(WorkflowAppLog).filter_by(workflow_run_id=message_data.workflow_run_id).first()
  516. )
  517. workflow_app_log_id = str(workflow_app_log_data.id) if workflow_app_log_data else None
  518. moderation_trace_info = ModerationTraceInfo(
  519. message_id=workflow_app_log_id or message_id,
  520. inputs=inputs,
  521. message_data=message_data.to_dict(),
  522. flagged=moderation_result.flagged,
  523. action=moderation_result.action,
  524. preset_response=moderation_result.preset_response,
  525. query=moderation_result.query,
  526. start_time=timer.get("start"),
  527. end_time=timer.get("end"),
  528. metadata=metadata,
  529. )
  530. return moderation_trace_info
  531. def suggested_question_trace(self, message_id, timer, **kwargs):
  532. suggested_question = kwargs.get("suggested_question", [])
  533. message_data = get_message_data(message_id)
  534. if not message_data:
  535. return {}
  536. metadata = {
  537. "message_id": message_id,
  538. "ls_provider": message_data.model_provider,
  539. "ls_model_name": message_data.model_id,
  540. "status": message_data.status,
  541. "from_end_user_id": message_data.from_end_user_id,
  542. "from_account_id": message_data.from_account_id,
  543. "agent_based": message_data.agent_based,
  544. "workflow_run_id": message_data.workflow_run_id,
  545. "from_source": message_data.from_source,
  546. }
  547. # get workflow_app_log_id
  548. workflow_app_log_id = None
  549. if message_data.workflow_run_id:
  550. workflow_app_log_data = (
  551. db.session.query(WorkflowAppLog).filter_by(workflow_run_id=message_data.workflow_run_id).first()
  552. )
  553. workflow_app_log_id = str(workflow_app_log_data.id) if workflow_app_log_data else None
  554. suggested_question_trace_info = SuggestedQuestionTraceInfo(
  555. message_id=workflow_app_log_id or message_id,
  556. message_data=message_data.to_dict(),
  557. inputs=message_data.message,
  558. outputs=message_data.answer,
  559. start_time=timer.get("start"),
  560. end_time=timer.get("end"),
  561. metadata=metadata,
  562. total_tokens=message_data.message_tokens + message_data.answer_tokens,
  563. status=message_data.status,
  564. error=message_data.error,
  565. from_account_id=message_data.from_account_id,
  566. agent_based=message_data.agent_based,
  567. from_source=message_data.from_source,
  568. model_provider=message_data.model_provider,
  569. model_id=message_data.model_id,
  570. suggested_question=suggested_question,
  571. level=message_data.status,
  572. status_message=message_data.error,
  573. )
  574. return suggested_question_trace_info
  575. def dataset_retrieval_trace(self, message_id, timer, **kwargs):
  576. documents = kwargs.get("documents")
  577. message_data = get_message_data(message_id)
  578. if not message_data:
  579. return {}
  580. metadata = {
  581. "message_id": message_id,
  582. "ls_provider": message_data.model_provider,
  583. "ls_model_name": message_data.model_id,
  584. "status": message_data.status,
  585. "from_end_user_id": message_data.from_end_user_id,
  586. "from_account_id": message_data.from_account_id,
  587. "agent_based": message_data.agent_based,
  588. "workflow_run_id": message_data.workflow_run_id,
  589. "from_source": message_data.from_source,
  590. }
  591. dataset_retrieval_trace_info = DatasetRetrievalTraceInfo(
  592. message_id=message_id,
  593. inputs=message_data.query or message_data.inputs,
  594. documents=[doc.model_dump() for doc in documents] if documents else [],
  595. start_time=timer.get("start"),
  596. end_time=timer.get("end"),
  597. metadata=metadata,
  598. message_data=message_data.to_dict(),
  599. )
  600. return dataset_retrieval_trace_info
  601. def tool_trace(self, message_id, timer, **kwargs):
  602. tool_name = kwargs.get("tool_name", "")
  603. tool_inputs = kwargs.get("tool_inputs", {})
  604. tool_outputs = kwargs.get("tool_outputs", {})
  605. message_data = get_message_data(message_id)
  606. if not message_data:
  607. return {}
  608. tool_config = {}
  609. time_cost = 0
  610. error = None
  611. tool_parameters = {}
  612. created_time = message_data.created_at
  613. end_time = message_data.updated_at
  614. agent_thoughts = message_data.agent_thoughts
  615. for agent_thought in agent_thoughts:
  616. if tool_name in agent_thought.tools:
  617. created_time = agent_thought.created_at
  618. tool_meta_data = agent_thought.tool_meta.get(tool_name, {})
  619. tool_config = tool_meta_data.get("tool_config", {})
  620. time_cost = tool_meta_data.get("time_cost", 0)
  621. end_time = created_time + timedelta(seconds=time_cost)
  622. error = tool_meta_data.get("error", "")
  623. tool_parameters = tool_meta_data.get("tool_parameters", {})
  624. metadata = {
  625. "message_id": message_id,
  626. "tool_name": tool_name,
  627. "tool_inputs": tool_inputs,
  628. "tool_outputs": tool_outputs,
  629. "tool_config": tool_config,
  630. "time_cost": time_cost,
  631. "error": error,
  632. "tool_parameters": tool_parameters,
  633. }
  634. file_url = ""
  635. message_file_data = db.session.query(MessageFile).filter_by(message_id=message_id).first()
  636. if message_file_data:
  637. message_file_id = message_file_data.id if message_file_data else None
  638. type = message_file_data.type
  639. created_by_role = message_file_data.created_by_role
  640. created_user_id = message_file_data.created_by
  641. file_url = f"{self.file_base_url}/{message_file_data.url}"
  642. metadata.update(
  643. {
  644. "message_file_id": message_file_id,
  645. "created_by_role": created_by_role,
  646. "created_user_id": created_user_id,
  647. "type": type,
  648. }
  649. )
  650. tool_trace_info = ToolTraceInfo(
  651. message_id=message_id,
  652. message_data=message_data.to_dict(),
  653. tool_name=tool_name,
  654. start_time=timer.get("start") if timer else created_time,
  655. end_time=timer.get("end") if timer else end_time,
  656. tool_inputs=tool_inputs,
  657. tool_outputs=tool_outputs,
  658. metadata=metadata,
  659. message_file_data=message_file_data,
  660. error=error,
  661. inputs=message_data.message,
  662. outputs=message_data.answer,
  663. tool_config=tool_config,
  664. time_cost=time_cost,
  665. tool_parameters=tool_parameters,
  666. file_url=file_url,
  667. )
  668. return tool_trace_info
  669. def generate_name_trace(self, conversation_id, timer, **kwargs):
  670. generate_conversation_name = kwargs.get("generate_conversation_name")
  671. inputs = kwargs.get("inputs")
  672. tenant_id = kwargs.get("tenant_id")
  673. if not tenant_id:
  674. return {}
  675. start_time = timer.get("start")
  676. end_time = timer.get("end")
  677. metadata = {
  678. "conversation_id": conversation_id,
  679. "tenant_id": tenant_id,
  680. }
  681. generate_name_trace_info = GenerateNameTraceInfo(
  682. conversation_id=conversation_id,
  683. inputs=inputs,
  684. outputs=generate_conversation_name,
  685. start_time=start_time,
  686. end_time=end_time,
  687. metadata=metadata,
  688. tenant_id=tenant_id,
  689. )
  690. return generate_name_trace_info
  691. trace_manager_timer: Optional[threading.Timer] = None
  692. trace_manager_queue: queue.Queue = queue.Queue()
  693. trace_manager_interval = int(os.getenv("TRACE_QUEUE_MANAGER_INTERVAL", 5))
  694. trace_manager_batch_size = int(os.getenv("TRACE_QUEUE_MANAGER_BATCH_SIZE", 100))
  695. class TraceQueueManager:
  696. def __init__(self, app_id=None, user_id=None):
  697. global trace_manager_timer
  698. self.app_id = app_id
  699. self.user_id = user_id
  700. self.trace_instance = OpsTraceManager.get_ops_trace_instance(app_id)
  701. self.flask_app = current_app._get_current_object() # type: ignore
  702. if trace_manager_timer is None:
  703. self.start_timer()
  704. def add_trace_task(self, trace_task: TraceTask):
  705. global trace_manager_timer, trace_manager_queue
  706. try:
  707. if self.trace_instance:
  708. trace_task.app_id = self.app_id
  709. trace_manager_queue.put(trace_task)
  710. except Exception as e:
  711. logging.exception(f"Error adding trace task, trace_type {trace_task.trace_type}")
  712. finally:
  713. self.start_timer()
  714. def collect_tasks(self):
  715. global trace_manager_queue
  716. tasks: list[TraceTask] = []
  717. while len(tasks) < trace_manager_batch_size and not trace_manager_queue.empty():
  718. task = trace_manager_queue.get_nowait()
  719. tasks.append(task)
  720. trace_manager_queue.task_done()
  721. return tasks
  722. def run(self):
  723. try:
  724. tasks = self.collect_tasks()
  725. if tasks:
  726. self.send_to_celery(tasks)
  727. except Exception as e:
  728. logging.exception("Error processing trace tasks")
  729. def start_timer(self):
  730. global trace_manager_timer
  731. if trace_manager_timer is None or not trace_manager_timer.is_alive():
  732. trace_manager_timer = threading.Timer(trace_manager_interval, self.run)
  733. trace_manager_timer.name = f"trace_manager_timer_{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}"
  734. trace_manager_timer.daemon = False
  735. trace_manager_timer.start()
  736. def send_to_celery(self, tasks: list[TraceTask]):
  737. with self.flask_app.app_context():
  738. for task in tasks:
  739. if task.app_id is None:
  740. continue
  741. file_id = uuid4().hex
  742. trace_info = task.execute()
  743. task_data = TaskData(
  744. app_id=task.app_id,
  745. trace_info_type=type(trace_info).__name__,
  746. trace_info=trace_info.model_dump() if trace_info else None,
  747. )
  748. file_path = f"{OPS_FILE_PATH}{task.app_id}/{file_id}.json"
  749. storage.save(file_path, task_data.model_dump_json().encode("utf-8"))
  750. file_info = {
  751. "file_id": file_id,
  752. "app_id": task.app_id,
  753. }
  754. process_trace_tasks.delay(file_info)