diff --git a/src/praisonai-agents/praisonaiagents/agent/chat_mixin.py b/src/praisonai-agents/praisonaiagents/agent/chat_mixin.py index 7e8385e70..fabe0336e 100644 --- a/src/praisonai-agents/praisonaiagents/agent/chat_mixin.py +++ b/src/praisonai-agents/praisonaiagents/agent/chat_mixin.py @@ -2039,14 +2039,115 @@ Your Goal: {self.goal}""" return fn(*args, **kwargs) return {**llm_kwargs, 'execute_tool_fn': tracked}, ran + def _custom_llm_hook_messages(self, llm_kwargs): + """The custom-LLM request as the message list BEFORE_LLM / AFTER_LLM hooks see.""" + messages = [] + if llm_kwargs.get('system_prompt'): + messages.append({"role": "system", "content": llm_kwargs['system_prompt']}) + messages.extend(llm_kwargs.get('chat_history') or []) + messages.append({"role": "user", "content": llm_kwargs.get('prompt')}) + return messages + + def _custom_llm_before_input(self, llm_kwargs): + """BEFORE_LLM input for the custom-LLM path, or ``None`` when no hook is registered.""" + from ..hooks import HookEvent, BeforeLLMInput + if not self._hook_runner.registry.has_hooks(HookEvent.BEFORE_LLM): + return None + return BeforeLLMInput( + session_id=getattr(self, '_session_id', 'default'), + cwd=os.getcwd(), + event_name=HookEvent.BEFORE_LLM, + timestamp=str(time.time()), + agent_name=self.name, + messages=self._custom_llm_hook_messages(llm_kwargs), + model=self.llm if isinstance(self.llm, str) else str(self.llm), + temperature=llm_kwargs.get('temperature'), + ) + + def _adopt_custom_llm_before(self, llm_kwargs, before, results): + """Apply BEFORE_LLM ``results`` like ``_chat_completion``: return ``(llm_kwargs, block_message)``.""" + if self._hook_runner.is_blocked(results): + reason = next( + (getattr(r.output, "reason", None) for r in results + if r.output and getattr(r.output, "is_denied", lambda: False)()), + None, + ) or "Blocked by hook" + logging.warning(f"Agent {self.name} LLM request blocked by BEFORE_LLM hook: {reason}") + return llm_kwargs, f"[LLM request blocked by hook: {reason}]" + messages = list(before.messages or []) + if messages and messages != self._custom_llm_hook_messages(llm_kwargs): + # A hook rewrote the messages (e.g. PII redactor): map them back. + system = None + if isinstance(messages[0], dict) and messages[0].get("role") == "system": + system = messages.pop(0).get("content") + last = messages.pop() if messages else {} + if not isinstance(last, dict) or last.get("role") != "user": + # ``LLM.get_response`` always sends the prompt as a user turn; fail closed + # rather than silently re-role a rewritten assistant/tool message. + logging.warning(f"Agent {self.name} BEFORE_LLM rewrite must end with a user message") + return llm_kwargs, "[LLM request blocked by hook: rewritten messages must end with a user message]" + llm_kwargs = {**llm_kwargs, 'system_prompt': system, 'chat_history': messages, + 'prompt': last.get("content")} + return llm_kwargs, None + + def _custom_llm_after_input(self, llm_kwargs, response, tally, start_time): + """AFTER_LLM input for the custom-LLM path, or ``None`` when no hook is registered.""" + from ..hooks import HookEvent, AfterLLMInput + if not self._hook_runner.registry.has_hooks(HookEvent.AFTER_LLM): + return None + return AfterLLMInput( + session_id=getattr(self, '_session_id', 'default'), + cwd=os.getcwd(), + event_name=HookEvent.AFTER_LLM, + timestamp=str(time.time()), + agent_name=self.name, + messages=self._custom_llm_hook_messages(llm_kwargs), + response=str(response), + model=self.llm if isinstance(self.llm, str) else str(self.llm), + tokens_used=sum(tally), + latency_ms=(time.time() - start_time) * 1000, + ) + def _custom_llm_get_response(self, llm_kwargs): + """Call the custom LLM, firing BEFORE_LLM / AFTER_LLM like ``_chat_completion``.""" + from ..hooks import HookEvent + before = self._custom_llm_before_input(llm_kwargs) + if before is not None: + results = self._hook_runner.execute_sync(HookEvent.BEFORE_LLM, before) + llm_kwargs, blocked = self._adopt_custom_llm_before(llm_kwargs, before, results) + if blocked: + return blocked + start_time, tally = time.time(), [] + response = self._custom_llm_dispatch(llm_kwargs, tally) + after = self._custom_llm_after_input(llm_kwargs, response, tally, start_time) + if after is not None: + self._hook_runner.execute_sync(HookEvent.AFTER_LLM, after) + return response + + async def _custom_llm_get_response_async(self, llm_kwargs): + """Async twin of :meth:`_custom_llm_get_response`.""" + from ..hooks import HookEvent + before = self._custom_llm_before_input(llm_kwargs) + if before is not None: + results = await self._hook_runner.execute(HookEvent.BEFORE_LLM, before) + llm_kwargs, blocked = self._adopt_custom_llm_before(llm_kwargs, before, results) + if blocked: + return blocked + start_time, tally = time.time(), [] + response = await self._custom_llm_dispatch_async(llm_kwargs, tally) + after = self._custom_llm_after_input(llm_kwargs, response, tally, start_time) + if after is not None: + await self._hook_runner.execute(HookEvent.AFTER_LLM, after) + return response + + def _custom_llm_dispatch(self, llm_kwargs, tally=None): """Call ``llm_instance.get_response`` walking ``fallback_models`` on outage.""" llm, index, current = self.llm_instance, 0, self.llm stream_cb = self.stream_emitter.emit if hasattr(self, 'stream_emitter') else None llm_kwargs, tools_ran = self._track_tool_runs(llm_kwargs) while True: try: - with self._llm_usage_sink(llm): + with self._llm_usage_sink(llm, tally): return llm.get_response(**llm_kwargs) except (ToolExecutionError, BudgetExceededError, InterruptedError): raise @@ -2059,14 +2160,14 @@ Your Goal: {self.goal}""" self._emit_model_fallback(current, next_model, category, index, stream_cb) llm, index, current = self._fallback_llm_for(next_model), index + 1, next_model - async def _custom_llm_get_response_async(self, llm_kwargs): - """Async twin of :meth:`_custom_llm_get_response`.""" + async def _custom_llm_dispatch_async(self, llm_kwargs, tally=None): + """Async twin of :meth:`_custom_llm_dispatch`.""" llm, index, current = self.llm_instance, 0, self.llm stream_cb = self.stream_emitter.emit if hasattr(self, 'stream_emitter') else None llm_kwargs, tools_ran = self._track_tool_runs(llm_kwargs) while True: try: - with self._llm_usage_sink(llm): + with self._llm_usage_sink(llm, tally): return await llm.get_response_async(**llm_kwargs) except (ToolExecutionError, BudgetExceededError, InterruptedError): raise @@ -2288,9 +2389,10 @@ Your Goal: {self.goal}""" self._on_budget_exceeded(current_cost, self._max_budget) @contextlib.contextmanager - def _llm_usage_sink(self, llm=None): + def _llm_usage_sink(self, llm=None, tally=None): """Route every call ``llm`` (default: the custom LLM) makes in this block into ``_record_llm_usage``. + Each call's total tokens are also appended to ``tally`` when given. Refuses to start when ``max_budget`` is already spent (``"stop"`` mode), so retrying after ``BudgetExceededError`` sends no further paid request. """ @@ -2309,10 +2411,12 @@ Your Goal: {self.goal}""" llm = llm if llm is not None else self.llm_instance # A fallback model is priced as itself, not as the primary ``self.llm``. model = None if llm is self.llm_instance else getattr(llm, 'model', None) - token = _usage_sink.set(( - llm, - lambda response: self._record_llm_usage(*self._extract_llm_usage(response, model)), - )) + def record(response): + usage = self._extract_llm_usage(response, model) + if tally is not None: + tally.append(usage[0] + usage[1]) + self._record_llm_usage(*usage) + token = _usage_sink.set((llm, record)) try: yield finally: diff --git a/src/praisonai-agents/tests/unit/agent/test_llm_hooks_custom_llm_path.py b/src/praisonai-agents/tests/unit/agent/test_llm_hooks_custom_llm_path.py new file mode 100644 index 000000000..8369f3558 --- /dev/null +++ b/src/praisonai-agents/tests/unit/agent/test_llm_hooks_custom_llm_path.py @@ -0,0 +1,137 @@ +"""BEFORE_LLM / AFTER_LLM on the custom-LLM (``provider/model``) path (#5821). + +That path hands the turn to ``LLM.get_response*`` and skips +``_chat_completion``, so it must fire the same hooks itself, once per turn, +with ``tokens_used`` taken from the provider's ``usage``. +``litellm.completion`` is stubbed; nothing leaves the machine. +""" +import asyncio + +import pytest + +pytest.importorskip("litellm") + +from praisonaiagents import Agent # noqa: E402 +from praisonaiagents.hooks import HookEvent, HookRegistry, HookResult # noqa: E402 + +MODEL = "anthropic/claude-3-5-sonnet-20240620" + + +def _litellm(): + import litellm + return litellm + + +def _response(model, n, tool_call=False): + message = {"role": "assistant", "content": "ok"} + if tool_call: + message = {"role": "assistant", "content": None, "tool_calls": [{ + "id": f"call_{n}", "type": "function", + "function": {"name": "step", "arguments": "{}"}}]} + usage = {"prompt_tokens": 100 + n, "completion_tokens": 10 + n, + "total_tokens": 110 + 2 * n} + return _litellm().ModelResponse(model=model, usage=usage, choices=[{ + "index": 0, "finish_reason": "tool_calls" if tool_call else "stop", + "message": message}]) + + +@pytest.fixture +def calls(monkeypatch): + """Stub litellm; answer with a tool call while ``tool_rounds`` remain.""" + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test") + monkeypatch.setattr(Agent, "_default_model", None) + monkeypatch.setattr(Agent, "_default_model_checked", False) + state = {"n": 0, "tool_rounds": 0, "messages": []} + + def fake(*a, **k): + state["n"] += 1 + state["messages"].append(k.get("messages")) + if k.get("tools") and state["tool_rounds"] > 0: + state["tool_rounds"] -= 1 + return _response(k.get("model"), state["n"], tool_call=True) + return _response(k.get("model"), state["n"]) + + async def afake(*a, **k): + return fake(*a, **k) + + litellm = _litellm() + monkeypatch.setattr(litellm, "completion", fake) + monkeypatch.setattr(litellm, "acompletion", afake) + return state + + +def _agent(registry, **kw): + return Agent(name="p", instructions="x", llm=MODEL, output="silent", + hooks=registry, **kw) + + +def _recorder(): + registry, seen = HookRegistry(), [] + for event in (HookEvent.BEFORE_LLM, HookEvent.AFTER_LLM): + registry.register_function(event, lambda data: seen.append(data) or HookResult.allow()) + return registry, seen + + +def step() -> str: + """One step of work.""" + return "done" + + +@pytest.mark.parametrize("run", [ + lambda a: a.chat("hi"), + lambda a: asyncio.run(a.achat("hi")), +], ids=["sync", "async"]) +def test_hooks_fire_with_provider_tokens(calls, run): + registry, seen = _recorder() + assert run(_agent(registry)) == "ok" + assert [e.event_name for e in seen] == [HookEvent.BEFORE_LLM, HookEvent.AFTER_LLM] + before, after = seen + assert before.model == MODEL and before.messages[-1]["content"] == "hi" + assert after.tokens_used == 101 + 11 + assert after.response == "ok" + + +def test_tokens_cover_every_call_in_a_tool_round_trip(calls): + calls["tool_rounds"] = 1 + registry, seen = _recorder() + assert _agent(registry, tools=[step]).chat("go") == "ok" + after = [e for e in seen if e.event_name == HookEvent.AFTER_LLM] + assert len(after) == 1 # Same once-per-turn granularity as the OpenAI-client path. + assert after[0].tokens_used == (101 + 11) + (102 + 12) + + +def test_blocking_before_llm_skips_the_provider(calls): + registry = HookRegistry() + registry.register_function(HookEvent.BEFORE_LLM, lambda data: HookResult.deny("nope")) + result = _agent(registry).chat("hi") + assert calls["n"] == 0 + assert "nope" in result + + +def test_before_llm_message_rewrite_reaches_the_provider(calls): + registry = HookRegistry() + + def redact(data): + messages = [{**m, "content": str(m["content"]).replace("my secret", "[redacted]")} + for m in data.messages] + return HookResult(decision="allow", modified_input={"messages": messages}) + + registry.register_function(HookEvent.BEFORE_LLM, redact, sequential=True) + assert _agent(registry).chat("my secret") == "ok" + sent = calls["messages"][0] + assert sent[0]["role"] == "system" + assert all("my secret" not in str(m.get("content")) for m in sent) + assert any(m.get("content") == "[redacted]" for m in sent) + + +def test_rewrite_ending_in_non_user_message_fails_closed(calls): + registry = HookRegistry() + + def append_assistant(data): + messages = list(data.messages) + [{"role": "assistant", "content": "injected"}] + return HookResult(decision="allow", modified_input={"messages": messages}) + + registry.register_function(HookEvent.BEFORE_LLM, append_assistant, sequential=True) + result = _agent(registry).chat("hi") + assert calls["n"] == 0 + assert "must end with a user message" in result