From e0e74874090b9dc8c6754055141b3e588bbc4e8b Mon Sep 17 00:00:00 2001 From: cv70 <2548604505@qq.com> Date: Thu, 21 May 2026 14:27:20 +0800 Subject: [PATCH] feat: add agent lifecycle hook trigger mechanism --- agent_loop.py | 14 +++- agentmain.py | 5 ++ llmcore.py | 2 +- plugins/hooks.py | 67 ++++++++++++++++ plugins/langfuse_tracing.py | 148 +++++++++++++++++++++--------------- 5 files changed, 169 insertions(+), 67 deletions(-) create mode 100644 plugins/hooks.py diff --git a/agent_loop.py b/agent_loop.py index 3011735f..7afb5850 100644 --- a/agent_loop.py +++ b/agent_loop.py @@ -1,6 +1,8 @@ import json, re, os from dataclasses import dataclass from typing import Any, Optional +try: from plugins.hooks import trigger as _hook +except ImportError: _hook = lambda *a, **k: None @dataclass class StepOutcome: data: Any @@ -12,16 +14,14 @@ def try_call_generator(func, *args, **kwargs): return ret class BaseHandler: - def tool_before_callback(self, tool_name, args, response): pass - def tool_after_callback(self, tool_name, args, response, ret): pass def turn_end_callback(self, response, tool_calls, tool_results, turn, next_prompt, exit_reason): return next_prompt def dispatch(self, tool_name, args, response, index=0, tool_num=1): method_name = f"do_{tool_name}" if hasattr(self, method_name): args['_index'] = index; args['_tool_num'] = tool_num - prer = yield from try_call_generator(self.tool_before_callback, tool_name, args, response) + _hook('tool_before', locals()) ret = yield from try_call_generator(getattr(self, method_name), args, response) - _ = yield from try_call_generator(self.tool_after_callback, tool_name, args, response, ret) + _hook('tool_after', locals()) return ret elif tool_name == 'bad_json': return StepOutcome(None, next_prompt=args.get('msg', 'bad_json'), should_exit=False) else: @@ -46,6 +46,7 @@ def agent_runner_loop(client, system_prompt, user_input, handler, tools_schema, {"role": "user", "content": initial_user_content if initial_user_content is not None else user_input} ] turn = 0; handler.max_turns = max_turns + _hook('agent_before', locals()) while turn < handler.max_turns: turn += 1; turnstr = f'LLM Running (Turn {turn}) ...' if handler.parent.task_dir: turnstr = f'Turn {turn} ...' @@ -53,6 +54,8 @@ def agent_runner_loop(client, system_prompt, user_input, handler, tools_schema, if yield_info: yield {'turn': turn} yield f"\n\n{turnstr}\n\n" if turn%10 == 0: client.last_tools = '' # 每10轮重置一次工具描述 + _hook('turn_before', locals()) + _hook('llm_before', locals()) response_gen = client.chat(messages=messages, tools=tools_schema) if verbose: response = yield from response_gen @@ -61,6 +64,7 @@ def agent_runner_loop(client, system_prompt, user_input, handler, tools_schema, response = exhaust(response_gen) cleaned = _clean_content(response.content) if cleaned: yield cleaned + '\n' + _hook('llm_after', locals()) if not response.tool_calls: tool_calls = [{'tool_name': 'no_tool', 'args': {}}] else: tool_calls = [{'tool_name': tc.function.name, 'args': json.loads(tc.function.arguments), 'id': tc.id} @@ -96,8 +100,10 @@ def agent_runner_loop(client, system_prompt, user_input, handler, tools_schema, if len(handler._done_hooks) == 0 or exit_reason.get('result', '') == 'EXITED': break next_prompts.add(handler._done_hooks.pop(0)) next_prompt = handler.turn_end_callback(response, tool_calls, tool_results, turn, '\n'.join(next_prompts), exit_reason) + _hook('turn_after', locals()) messages = [{"role": "user", "content": next_prompt, "tool_results": tool_results}] # just new message, history is kept in *Session if exit_reason: handler.turn_end_callback(response, tool_calls, tool_results, turn, '', exit_reason) + _hook('agent_after', locals()) return exit_reason or {'result': 'MAX_TURNS_EXCEEDED'} def _clean_content(text): diff --git a/agentmain.py b/agentmain.py index 1403f4b8..a4f433b7 100644 --- a/agentmain.py +++ b/agentmain.py @@ -8,6 +8,11 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) from llmcore import reload_mykeys, LLMSession, ToolClient, ClaudeSession, MixinSession, NativeToolClient, NativeClaudeSession, NativeOAISession, resolve_client from agent_loop import agent_runner_loop +try: + from plugins.hooks import discover_and_load as _discover_plugins + _discover_plugins() +except Exception: + pass from ga import GenericAgentHandler, smart_format, get_global_memory, format_error, consume_file script_dir = os.path.dirname(os.path.abspath(__file__)) diff --git a/llmcore.py b/llmcore.py index 2e239d32..c81fb0bb 100644 --- a/llmcore.py +++ b/llmcore.py @@ -22,7 +22,7 @@ def reload_mykeys(): print(f'[Info] Load mykeys from {_mykey_path}') globals().update(mykeys=mk) if mk.get('langfuse_config'): - try: from plugins import langfuse_tracing + try: from plugins.hooks import load as _hl; _hl('langfuse_tracing') except Exception: pass return mk, True diff --git a/plugins/hooks.py b/plugins/hooks.py new file mode 100644 index 00000000..655fa377 --- /dev/null +++ b/plugins/hooks.py @@ -0,0 +1,67 @@ +import os +import sys +import importlib + +# 模块级注册表: event_name -> [callback, ...] +_registry = {} +_PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def register(event): + def decorator(fn): + _registry.setdefault(event, []).append(fn) + return fn + return decorator + + +def trigger(event, ctx: dict): + for fn in _registry.get(event, []): + try: + r = fn(ctx) + if isinstance(r, dict): + ctx = r + except Exception as e: + sys.stderr.write(f"[hooks] {event} callback error: {e}\n") + return ctx + + +def unregister(event, fn): + try: + _registry[event] = [f for f in _registry[event] if f is not fn] + except KeyError: + pass + + +def clear(event=None): + if event: + _registry.pop(event, None) + else: + _registry.clear() + + +def has(event): + return bool(_registry.get(event)) + + +def discover_and_load(plugin_dir=None): + if plugin_dir is None: + plugin_dir = os.path.join(_PROJECT_ROOT, 'plugins') + if not os.path.isdir(plugin_dir): + return + parent = os.path.dirname(plugin_dir) + if parent not in sys.path: + sys.path.insert(0, parent) + for fn in sorted(os.listdir(plugin_dir)): + if fn.startswith('_') or not fn.endswith('.py'): + continue + name = fn[:-3] + load(name) + + +def load(name): + try: + importlib.import_module(f'plugins.{name}') + return True + except Exception as e: + sys.stderr.write(f"[hooks] plugin '{name}' load failed: {e}\n") + return False \ No newline at end of file diff --git a/plugins/langfuse_tracing.py b/plugins/langfuse_tracing.py index 08e0e397..3d503183 100644 --- a/plugins/langfuse_tracing.py +++ b/plugins/langfuse_tracing.py @@ -1,9 +1,11 @@ -"""Opt-in Langfuse tracing. Self-activates on import if langfuse_config exists in mykey. +"""Langfuse tracing via hook system. Self-activates on import if langfuse_config exists in mykey. -Hooks only via monkey-patch so core files stay untouched: -- agent_loop.agent_runner_loop -> outer agent trace (parent of all below) -- llmcore._write_llm_log -> generation span (Prompt=start, Response=end) -- BaseHandler.tool_before/after -> tool span +Replaces old monkey-patch approach with hooks on: + - agent_before / agent_after -> agent trace + - llm_before / llm_after -> generation span + - tool_before / tool_after -> tool span + +Usage tracking (SSE parser wrapping) stays as internal llmcore patch. """ import threading, sys @@ -16,21 +18,82 @@ except Exception: _lf = None if _lf: - import llmcore, agent_loop + import plugins.hooks as hooks, llmcore _tls = threading.local() - _orig_log = llmcore._write_llm_log - def _patched_log(label, content, log_path=None): + # ── Agent trace ────────────────────────────────────────────── + + @hooks.register('agent_before') + def _on_agent_before(ctx): try: - if label == 'Prompt': - _tls.gen = _lf.start_observation(name='llm.chat', as_type='generation', input=content[:20000]) - _tls.usage = None - elif label == 'Response' and getattr(_tls, 'gen', None) is not None: - _tls.gen.update(output=content[:20000], usage_details=getattr(_tls, 'usage', None)) - _tls.gen.end(); _tls.gen = None - except Exception: pass - return _orig_log(label, content, log_path) - llmcore._write_llm_log = _patched_log + _tls.trace_obs = _lf.start_observation( + name='agent.task', as_type='agent', + input={'user_input': ctx.get('user_input', '')}) + except Exception: + _tls.trace_obs = None + + @hooks.register('agent_after') + def _on_agent_after(ctx): + try: + obs = getattr(_tls, 'trace_obs', None) + if obs: + obs.update(output=ctx.get('exit_reason')) + obs.end() + _tls.trace_obs = None + _lf.flush() + except Exception: + pass + + # ── LLM generation span (replaces _write_llm_log patch) ───── + + @hooks.register('llm_before') + def _on_llm_before(ctx): + try: + _tls.gen = _lf.start_observation( + name='llm.chat', as_type='generation', + input=str(ctx.get('messages', ''))[:20000]) + _tls._usage = None + except Exception: + _tls.gen = None + + @hooks.register('llm_after') + def _on_llm_after(ctx): + try: + gen = getattr(_tls, 'gen', None) + if gen: + gen.update(output=str(ctx.get('response', ''))[:20000], + usage_details=getattr(_tls, '_usage', None)) + gen.end() + _tls.gen = None + except Exception: + pass + + # ── Tool spans (replaces tool_before/after_callback patches) ─ + + @hooks.register('tool_before') + def _on_tool_before(ctx): + try: + name = ctx.get('tool_name', '?') + args = {k: v for k, v in (ctx.get('args') or {}).items() if not k.startswith('_')} + if not hasattr(_tls, 'tstack'): _tls.tstack = [] + _tls.tstack.append(_lf.start_observation(name=name, as_type='tool', input=args)) + except Exception: + pass + + @hooks.register('tool_after') + def _on_tool_after(ctx): + try: + stack = getattr(_tls, 'tstack', []) + if stack: + sp = stack.pop() + ret = ctx.get('ret') + out = {'data': ret.data, 'next_prompt': ret.next_prompt, + 'should_exit': ret.should_exit} if ret else None + sp.update(output=out); sp.end() + except Exception: + pass + + # ── Usage tracking: tee SSE data for token counts ─────────── def _extract_usage(buf): u = {} @@ -72,51 +135,12 @@ if _lf: for ln in resp_lines: buf.append(ln); yield ln ret = yield from orig(tee(), *a, **kw) - try: _tls.usage = _extract_usage(buf) - except Exception: pass + try: + _tls._usage = _extract_usage(buf) + except Exception: + pass return ret return wrapped + llmcore._parse_claude_sse = _wrap_parser(llmcore._parse_claude_sse) - llmcore._parse_openai_sse = _wrap_parser(llmcore._parse_openai_sse) - - _orig_before = agent_loop.BaseHandler.tool_before_callback - _orig_after = agent_loop.BaseHandler.tool_after_callback - - def _patched_before(self, tool_name, args, response): - try: - if not hasattr(_tls, 'tstack'): _tls.tstack = [] - a = {k: v for k, v in args.items() if k != '_index'} - _tls.tstack.append(_lf.start_observation(name=tool_name, as_type='tool', input=a)) - except Exception: pass - return _orig_before(self, tool_name, args, response) - - def _patched_after(self, tool_name, args, response, ret): - try: - if getattr(_tls, 'tstack', None): - sp = _tls.tstack.pop() - out = {'data': ret.data, 'next_prompt': ret.next_prompt, 'should_exit': ret.should_exit} if ret else None - sp.update(output=out); sp.end() - except Exception: pass - return _orig_after(self, tool_name, args, response, ret) - - agent_loop.BaseHandler.tool_before_callback = _patched_before - agent_loop.BaseHandler.tool_after_callback = _patched_after - - _orig_loop = agent_loop.agent_runner_loop - def _patched_loop(client, system_prompt, user_input, handler, tools_schema, *a, **kw): - try: cm = _lf.start_as_current_observation(name='agent.task', as_type='agent', input={'user_input': user_input}) - except Exception: cm = None - if cm is None: - ret = yield from _orig_loop(client, system_prompt, user_input, handler, tools_schema, *a, **kw); return ret - with cm as sp: - ret = yield from _orig_loop(client, system_prompt, user_input, handler, tools_schema, *a, **kw) - try: sp.update(output=ret) - except Exception: pass - try: _lf.flush() - except Exception: pass - return ret - agent_loop.agent_runner_loop = _patched_loop - for _m in list(sys.modules.values()): - if _m and getattr(_m, 'agent_runner_loop', None) is _orig_loop: - try: setattr(_m, 'agent_runner_loop', _patched_loop) - except Exception: pass + llmcore._parse_openai_sse = _wrap_parser(llmcore._parse_openai_sse) \ No newline at end of file