/ examples / middleware / injected_state.py
injected_state.py
 1  """
 2  Injected State Example - PraisonAI Agents
 3  
 4  Demonstrates injecting agent state into tools without exposing it in the schema.
 5  """
 6  
 7  from praisonaiagents import Agent, tool
 8  from praisonaiagents.tools import Injected
 9  from praisonaiagents.tools.injected import AgentState, with_injection_context
10  
11  # Tool with injected state - state param is NOT in the public schema
12  @tool
13  def show_context(query: str, state: Injected[dict]) -> str:
14      """Show the current agent context."""
15      session_id = state.get('session_id', 'unknown')
16      agent_id = state.get('agent_id', 'unknown')
17      return f"Query: {query}, Session: {session_id}, Agent: {agent_id}"
18  
19  # Create agent with the tool
20  agent = Agent(
21      name="ContextBot",
22      instructions="You help show context information.",
23      tools=[show_context],
24      session_id="my-session-123"
25  )
26  
27  if __name__ == "__main__":
28      # Verify injected param is not in schema
29      schema = show_context.get_schema()
30      params = schema['function']['parameters']['properties']
31      print(f"Schema params: {list(params.keys())}")
32      assert 'state' not in params, "state should NOT be in schema"
33      print("āœ“ 'state' correctly excluded from schema")
34      
35      # Test with manual injection context
36      mock_state = AgentState(
37          agent_id="test-agent",
38          run_id="run-1",
39          session_id="session-abc"
40      )
41      
42      with with_injection_context(mock_state):
43          result = show_context(query="hello")
44          print(f"Result: {result}")
45      
46      # Test via agent.execute_tool
47      result = agent.execute_tool("show_context", {"query": "test"})
48      print(f"Agent result: {result}")
49      
50      print("\nāœ“ Injected state example complete")