""" 使用装饰器实现全栈函数调用追踪 解决LangGraph无法捕获子函数内部状态的问题 """ import functools import inspect import asyncio import contextvars from typing import Callable, Any, Dict, Optional, List from datetime import datetime from collections import deque # 使用 contextvars 实现请求级别的隔离，避免多个并发请求互相干扰 _event_callbacks: contextvars.ContextVar[List[Callable]] = contextvars.ContextVar('event_callbacks', default=None) _current_stream_handler: contextvars.ContextVar[Optional[Any]] = contextvars.ContextVar('current_stream_handler', default=None) def register_event_callback(callback: Callable[[Dict[str, Any]], None]): """ 注册事件回调函数（使用 contextvars 实现请求隔离） """ callbacks = _event_callbacks.get() if callbacks is None: callbacks = [] _event_callbacks.set(callbacks) callbacks.append(callback) def unregister_event_callback(callback: Callable[[Dict[str, Any]], None]): """ 移除事件回调函数（使用 contextvars 实现请求隔离） """ callbacks = _event_callbacks.get() if callbacks and callback in callbacks: callbacks.remove(callback) def clear_event_callbacks(): """ 清空所有事件回调（使用 contextvars 实现请求隔离） """ callbacks = _event_callbacks.get() if callbacks: callbacks.clear() def get_all_callbacks(): """ 获取所有注册的回调函数（使用 contextvars 实现请求隔离） """ callbacks = _event_callbacks.get() return callbacks.copy() if callbacks else [] def extract_function_description(func: Callable) -> str: """ 提取函数的描述（docstring的第一行或前100个字符） """ if func.__doc__: doc = func.__doc__.strip() # 取第一行或前100个字符 first_line = doc.split('\n')[0].strip() return first_line[:100] if len(first_line) > 100 else first_line return "" def _extract_details_from_result(result: Any) -> str: """ 从函数返回值中提取 details，顺序固定为：先 details 再中文。 1) 先看返回的 dict 里是否有「details 类」字段（__details__、_details、details 或键名含 "details"）， 且其值为非空，则用该值； 2) 若没有或为空，再走原本的中文兜底：从返回结构中找一段包含中文的内容作为 details。 """ default = "完成" if not result: return default # 第一步：优先用 details 类字段（仅当有值且非空时采用） if isinstance(result, dict): for key in ("__details__", "_details", "details"): if key in result and result[key] is not None: s = str(result[key]).strip() if s: return str(result[key]) for k, v in result.items(): if "details" in k.lower() and v is not None: s = str(v).strip() if s: return str(v) # 第二步：保留原本的中文兜底 def _contains_chinese(text: str) -> bool: if not text: return False return any("\u4e00" <= c <= "\u9fff" for c in str(text)) def _find_chinese(obj: Any, max_length: int = 200) -> Optional[str]: if isinstance(obj, str): return (obj[:max_length] + "...") if _contains_chinese(obj) and len(obj) > max_length else ( obj if _contains_chinese(obj) else None) if isinstance(obj, dict): for key, value in obj.items(): if "details" in key.lower(): continue found = _find_chinese(value, max_length) if found: return found if isinstance(obj, list): for item in obj: found = _find_chinese(item, max_length) if found: return found return None chinese = _find_chinese(result) return chinese if chinese else default def track_function_calls(func: Callable) -> Callable: """ 装饰器：追踪函数调用过程，支持同步和异步函数 Args: func: 被装饰的函数 """ @functools.wraps(func) def sync_wrapper(*args, **kwargs): # 获取函数签名信息 sig = inspect.signature(func) bound_args = sig.bind(*args, **kwargs) bound_args.apply_defaults() func_name = func.__name__ module_name = func.__module__ # 尝试从函数参数中获取 state state = None if args and isinstance(args[0], dict): state = args[0] elif 'state' in bound_args.arguments: state = bound_args.arguments.get('state') func_description = extract_function_description(func) # 准备函数参数信息 params = {k: v for k, v in bound_args.arguments.items()} try: # 执行原函数 result = func(*args, **kwargs) # 构建 details：优先用返回值中的 details 类字段（__details__、_details、details 或键名含 details） details = _extract_details_from_result(result) # 合并函数执行事件（只在执行完成后发送一次） execution_event = { "type": "function_execution", "title": func_description, "details": details, "result": result, # 添加函数返回结果 "timestamp": datetime.now().isoformat() } # 调用所有注册的回调（从上下文变量获取） callbacks = _event_callbacks.get() or [] for callback in callbacks: try: callback(execution_event) except Exception as e: print(f"回调执行出错: {e}") return result except Exception as e: # 发送函数执行失败事件 error_event = { "type": "function_error", "title": func_description, "details": f"错误: {str(e)}", "timestamp": datetime.now().isoformat() } callbacks = _event_callbacks.get() or [] for callback in callbacks: try: callback(error_event) except Exception as ex: print(f"回调执行出错: {ex}") raise @functools.wraps(func) async def async_wrapper(*args, **kwargs): # 获取函数签名信息 sig = inspect.signature(func) bound_args = sig.bind(*args, **kwargs) bound_args.apply_defaults() func_name = func.__name__ module_name = func.__module__ # 尝试从函数参数中获取 state state = None if args and isinstance(args[0], dict): state = args[0] elif 'state' in bound_args.arguments: state = bound_args.arguments.get('state') func_description = extract_function_description(func) # 准备函数参数信息 params = {k: v for k, v in bound_args.arguments.items()} try: # 执行原函数 result = await func(*args, **kwargs) # 构建 details：优先用返回值中的 details 类字段（__details__、_details、details 或键名含 details） details = _extract_details_from_result(result) # 合并函数执行事件（只在执行完成后发送一次） execution_event = { "type": "function_execution", "title": func_description, "details": details, "result": result, # 添加函数返回结果 "timestamp": datetime.now().isoformat() } # 调用所有注册的回调（从上下文变量获取） callbacks = _event_callbacks.get() or [] for callback in callbacks: try: callback(execution_event) except Exception as e: print(f"回调执行出错: {e}") return result except Exception as e: # 发送函数执行失败事件 error_event = { "type": "function_error", "title": func_description, "details": f"错误: {str(e)}", "timestamp": datetime.now().isoformat() } callbacks = _event_callbacks.get() or [] for callback in callbacks: try: callback(error_event) except Exception as ex: print(f"回调执行出错: {ex}") raise # 根据函数是否为协程函数返回对应的装饰器 if asyncio.iscoroutinefunction(func): return async_wrapper else: return sync_wrapper # 事件处理器 def console_event_handler(event: Dict[str, Any]): """ 控制台事件处理器 """ event_type = event["type"] title = event.get("title", "") details = event.get("details", "") if event_type == "function_execution": print(f"✅ 函数执行: {title}") if details: print(f" 详情: {details}") elif event_type == "function_error": print(f"❌ 函数错误: {title}") if details: print(f" 详情: {details}") class StreamEventHandler: """ 流式事件处理器，将装饰器事件放入队列供流式输出使用 """ def __init__(self, queue: asyncio.Queue): self.queue = queue def __call__(self, event: Dict[str, Any]): """ 处理事件，将事件放入队列 注意：这是同步方法，但队列是异步的，需要特殊处理 """ try: # 使用线程安全的方式将事件放入队列 # 如果队列已满，使用 put_nowait 并忽略错误 try: self.queue.put_nowait(event) except asyncio.QueueFull: # 队列已满，跳过此事件（不应该发生，因为队列大小足够） pass except Exception as e: # 忽略错误，避免影响主流程 pass def send_stream_content(self, content: str, title: str = ""): """ 发送流式内容事件 Args: content: 流式内容片段 title: 函数描述/标题（可选） """ try: event = { "type": "stream_content", "content": content, "title": title, "timestamp": datetime.now().isoformat() } self.queue.put_nowait(event) except Exception as e: # 忽略错误，避免影响主流程 pass def create_stream_event_handler() -> tuple[StreamEventHandler, asyncio.Queue]: """ 创建流式事件处理器和对应的队列（使用 contextvars 实现请求隔离） Returns: (handler, queue): 事件处理器和事件队列 """ queue = asyncio.Queue(maxsize=1000) # 设置足够大的队列大小 handler = StreamEventHandler(queue) _current_stream_handler.set(handler) # 保存到上下文变量 return handler, queue def get_current_stream_handler() -> Optional[StreamEventHandler]: """ 获取当前活动的流式事件处理器（使用 contextvars 实现请求隔离） Returns: 当前流式事件处理器，如果没有则返回 None """ return _current_stream_handler.get() def clear_current_stream_handler(): """ 清除当前活动的流式事件处理器（使用 contextvars 实现请求隔离） """ _current_stream_handler.set(None)