Merge pull request #6103 from MervinPraison/claude/issue-5821-20261028

fix: fire BEFORE_LLM/AFTER_LLM hooks on the custom-LLM path

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:

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