361 lines
12 KiB
Python
361 lines
12 KiB
Python
"""
|
||
使用装饰器实现全栈函数调用追踪
|
||
解决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)
|