mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 00:28:31 +00:00
Lorenze/ensure hooks work with lite agents flows (#3981)
* liteagent support hooks * wip llm.call hooks work - needs tests for this * fix tests * fixed more * more tool hooks test cassettes
This commit is contained in:
@@ -309,3 +309,188 @@ class TestLLMHooksIntegration:
|
||||
clear_all_llm_call_hooks()
|
||||
hooks = get_before_llm_call_hooks()
|
||||
assert len(hooks) == 0
|
||||
|
||||
@pytest.mark.vcr()
|
||||
def test_lite_agent_hooks_integration_with_real_llm(self):
|
||||
"""Test that LiteAgent executes before/after LLM call hooks and prints messages correctly."""
|
||||
import os
|
||||
from crewai.lite_agent import LiteAgent
|
||||
|
||||
# Skip if no API key available
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
pytest.skip("OPENAI_API_KEY not set - skipping real LLM test")
|
||||
|
||||
# Track hook invocations
|
||||
hook_calls = {"before": [], "after": []}
|
||||
|
||||
def before_llm_call_hook(context: LLMCallHookContext) -> bool:
|
||||
"""Log and verify before hook execution."""
|
||||
print(f"\n[BEFORE HOOK] Agent: {context.agent.role if context.agent else 'None'}")
|
||||
print(f"[BEFORE HOOK] Iterations: {context.iterations}")
|
||||
print(f"[BEFORE HOOK] Message count: {len(context.messages)}")
|
||||
print(f"[BEFORE HOOK] Messages: {context.messages}")
|
||||
|
||||
# Track the call
|
||||
hook_calls["before"].append({
|
||||
"iterations": context.iterations,
|
||||
"message_count": len(context.messages),
|
||||
"has_task": context.task is not None,
|
||||
"has_crew": context.crew is not None,
|
||||
})
|
||||
|
||||
return True # Allow execution
|
||||
|
||||
def after_llm_call_hook(context: LLMCallHookContext) -> str | None:
|
||||
"""Log and verify after hook execution."""
|
||||
print(f"\n[AFTER HOOK] Agent: {context.agent.role if context.agent else 'None'}")
|
||||
print(f"[AFTER HOOK] Iterations: {context.iterations}")
|
||||
print(f"[AFTER HOOK] Response: {context.response[:100] if context.response else 'None'}...")
|
||||
print(f"[AFTER HOOK] Final message count: {len(context.messages)}")
|
||||
|
||||
# Track the call
|
||||
hook_calls["after"].append({
|
||||
"iterations": context.iterations,
|
||||
"has_response": context.response is not None,
|
||||
"response_length": len(context.response) if context.response else 0,
|
||||
})
|
||||
|
||||
# Optionally modify response
|
||||
if context.response:
|
||||
return f"[HOOKED] {context.response}"
|
||||
return None
|
||||
|
||||
# Register hooks
|
||||
register_before_llm_call_hook(before_llm_call_hook)
|
||||
register_after_llm_call_hook(after_llm_call_hook)
|
||||
|
||||
try:
|
||||
# Create LiteAgent
|
||||
lite_agent = LiteAgent(
|
||||
role="Test Assistant",
|
||||
goal="Answer questions briefly",
|
||||
backstory="You are a helpful test assistant",
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Verify hooks are loaded
|
||||
assert len(lite_agent.before_llm_call_hooks) > 0, "Before hooks not loaded"
|
||||
assert len(lite_agent.after_llm_call_hooks) > 0, "After hooks not loaded"
|
||||
|
||||
# Execute with a simple prompt
|
||||
result = lite_agent.kickoff("Say 'Hello World' and nothing else")
|
||||
|
||||
|
||||
# Verify hooks were called
|
||||
assert len(hook_calls["before"]) > 0, "Before hook was never called"
|
||||
assert len(hook_calls["after"]) > 0, "After hook was never called"
|
||||
|
||||
# Verify context had correct attributes for LiteAgent (used in flows)
|
||||
# LiteAgent doesn't have task/crew context, unlike agents in CrewBase
|
||||
before_call = hook_calls["before"][0]
|
||||
assert before_call["has_task"] is False, "Task should be None for LiteAgent in flows"
|
||||
assert before_call["has_crew"] is False, "Crew should be None for LiteAgent in flows"
|
||||
assert before_call["message_count"] > 0, "Should have messages"
|
||||
|
||||
# Verify after hook received response
|
||||
after_call = hook_calls["after"][0]
|
||||
assert after_call["has_response"] is True, "After hook should have response"
|
||||
assert after_call["response_length"] > 0, "Response should not be empty"
|
||||
|
||||
# Verify response was modified by after hook
|
||||
# Note: The hook modifies the raw LLM response, but LiteAgent then parses it
|
||||
# to extract the "Final Answer" portion. We check the messages to see the modification.
|
||||
assert len(result.messages) > 2, "Should have assistant message in messages"
|
||||
last_message = result.messages[-1]
|
||||
assert last_message["role"] == "assistant", "Last message should be from assistant"
|
||||
assert "[HOOKED]" in last_message["content"], "Hook should have modified the assistant message"
|
||||
|
||||
|
||||
finally:
|
||||
# Clean up hooks
|
||||
unregister_before_llm_call_hook(before_llm_call_hook)
|
||||
unregister_after_llm_call_hook(after_llm_call_hook)
|
||||
|
||||
@pytest.mark.vcr()
|
||||
def test_direct_llm_call_hooks_integration(self):
|
||||
"""Test that hooks work for direct llm.call() without agents."""
|
||||
import os
|
||||
from crewai.llm import LLM
|
||||
|
||||
# Skip if no API key available
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
pytest.skip("OPENAI_API_KEY not set - skipping real LLM test")
|
||||
|
||||
# Track hook invocations
|
||||
hook_calls = {"before": [], "after": []}
|
||||
|
||||
def before_hook(context: LLMCallHookContext) -> bool:
|
||||
"""Log and verify before hook execution."""
|
||||
print(f"\n[BEFORE HOOK] Agent: {context.agent}")
|
||||
print(f"[BEFORE HOOK] Task: {context.task}")
|
||||
print(f"[BEFORE HOOK] Crew: {context.crew}")
|
||||
print(f"[BEFORE HOOK] LLM: {context.llm}")
|
||||
print(f"[BEFORE HOOK] Iterations: {context.iterations}")
|
||||
print(f"[BEFORE HOOK] Message count: {len(context.messages)}")
|
||||
|
||||
# Track the call
|
||||
hook_calls["before"].append({
|
||||
"agent": context.agent,
|
||||
"task": context.task,
|
||||
"crew": context.crew,
|
||||
"llm": context.llm is not None,
|
||||
"message_count": len(context.messages),
|
||||
})
|
||||
|
||||
return True # Allow execution
|
||||
|
||||
def after_hook(context: LLMCallHookContext) -> str | None:
|
||||
"""Log and verify after hook execution."""
|
||||
print(f"\n[AFTER HOOK] Agent: {context.agent}")
|
||||
print(f"[AFTER HOOK] Response: {context.response[:100] if context.response else 'None'}...")
|
||||
|
||||
# Track the call
|
||||
hook_calls["after"].append({
|
||||
"has_response": context.response is not None,
|
||||
"response_length": len(context.response) if context.response else 0,
|
||||
})
|
||||
|
||||
# Modify response
|
||||
if context.response:
|
||||
return f"[HOOKED] {context.response}"
|
||||
return None
|
||||
|
||||
# Register hooks
|
||||
register_before_llm_call_hook(before_hook)
|
||||
register_after_llm_call_hook(after_hook)
|
||||
|
||||
try:
|
||||
# Create LLM and make direct call
|
||||
llm = LLM(model="gpt-4o-mini")
|
||||
result = llm.call([{"role": "user", "content": "Say hello"}])
|
||||
|
||||
print(f"\n[TEST] Final result: {result}")
|
||||
|
||||
# Verify hooks were called
|
||||
assert len(hook_calls["before"]) > 0, "Before hook was never called"
|
||||
assert len(hook_calls["after"]) > 0, "After hook was never called"
|
||||
|
||||
# Verify context had correct attributes for direct LLM calls
|
||||
before_call = hook_calls["before"][0]
|
||||
assert before_call["agent"] is None, "Agent should be None for direct LLM calls"
|
||||
assert before_call["task"] is None, "Task should be None for direct LLM calls"
|
||||
assert before_call["crew"] is None, "Crew should be None for direct LLM calls"
|
||||
assert before_call["llm"] is True, "LLM should be present"
|
||||
assert before_call["message_count"] > 0, "Should have messages"
|
||||
|
||||
# Verify after hook received response
|
||||
after_call = hook_calls["after"][0]
|
||||
assert after_call["has_response"] is True, "After hook should have response"
|
||||
assert after_call["response_length"] > 0, "Response should not be empty"
|
||||
|
||||
# Verify response was modified by after hook
|
||||
assert "[HOOKED]" in result, "Response should be modified by after hook"
|
||||
|
||||
finally:
|
||||
# Clean up hooks
|
||||
unregister_before_llm_call_hook(before_hook)
|
||||
unregister_after_llm_call_hook(after_hook)
|
||||
|
||||
Reference in New Issue
Block a user