ops_trace_manager.py 29 KB

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