/ examples / middleware / basic_middleware.py
basic_middleware.py
 1  """
 2  Basic Middleware Example - PraisonAI Agents
 3  
 4  Demonstrates before/after hooks and wrap decorators for model and tool calls.
 5  """
 6  
 7  from praisonaiagents import Agent, tool
 8  from praisonaiagents.hooks import (
 9      before_model, after_model, wrap_model_call,
10      before_tool, after_tool, wrap_tool_call,
11      InvocationContext, ModelRequest, ModelResponse
12  )
13  
14  # Simple tool for demonstration
15  @tool
16  def get_weather(city: str) -> str:
17      """Get weather for a city."""
18      return f"Sunny, 22°C in {city}"
19  
20  # Before model hook - adds context
21  @before_model
22  def add_context(request: ModelRequest) -> ModelRequest:
23      print(f"[before_model] Adding context to request")
24      return request
25  
26  # After model hook - logs response
27  @after_model  
28  def log_response(response: ModelResponse) -> ModelResponse:
29      print(f"[after_model] Response received")
30      return response
31  
32  # Wrap tool call - retry on error
33  @wrap_tool_call
34  def retry_on_error(tool_call, call_next):
35      print(f"[wrap_tool_call] Executing tool")
36      try:
37          return call_next(tool_call)
38      except Exception as e:
39          print(f"[wrap_tool_call] Retrying after error: {e}")
40          return call_next(tool_call)
41  
42  # Create agent with hooks
43  agent = Agent(
44      name="WeatherBot",
45      instructions="You help with weather queries.",
46      tools=[get_weather],
47      hooks=[add_context, log_response, retry_on_error]
48  )
49  
50  if __name__ == "__main__":
51      # Test the tool directly
52      result = get_weather("London")
53      print(f"Weather: {result}")
54      
55      print("\n✓ Middleware example complete")