Merge pull request strands-agents#31 from mkmeral/fix-scripts

Fix tool calling

Author: mehtarac

src/strands/experimental/bidirectional_streaming/agent/agent.py

@@ -18,6 +18,7 @@
 from typing import Any, AsyncIterable, Callable
 
 from .... import _identifier
+from ....hooks.registry import HookRegistry
 from ....telemetry.metrics import EventLoopMetrics
 from ....tools.caller import _ToolCaller
 from ....tools.executors import ConcurrentToolExecutor
@@ -122,6 +123,7 @@ def __init__(
         # Initialize other components
         self.event_loop_metrics = EventLoopMetrics()
         self._tool_caller = _ToolCaller(self)
+        self.hooks = HookRegistry()
 
         # connection management
         self._agent_loop: "BidirectionalConnection" | None = None

src/strands/experimental/bidirectional_streaming/scripts/test_bidi.py

@@ -6,7 +6,7 @@
 
 sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent.parent))
 
-from strands.experimental.bidirectional_streaming.agent.agent import BidirectionalAgent
+from strands.experimental.bidirectional_streaming.agent.agent import BidiAgent
 from strands.experimental.bidirectional_streaming.models.novasonic import BidiNovaSonicModel
 from strands.experimental.bidirectional_streaming.io.audio import AudioIO
 from strands_tools import calculator
@@ -20,8 +20,8 @@ async def main():
     adapter = AudioIO()
     model = BidiNovaSonicModel(region="us-east-1")
 
-    async with BidirectionalAgent(model=model, tools=[calculator]) as agent:
-        print("New BidirectionalAgent Experience")
+    async with BidiAgent(model=model, tools=[calculator]) as agent:
+        print("New BidiAgent Experience")
         print("Try asking: 'What is 25 times 8?' or 'Calculate the square root of 144'")
         await agent.run(io_channels=[adapter])
 

src/strands/experimental/bidirectional_streaming/scripts/test_bidi_novasonic.py

@@ -17,7 +17,7 @@
 import pyaudio
 from strands_tools import calculator
 
-from strands.experimental.bidirectional_streaming.agent.agent import BidirectionalAgent
+from strands.experimental.bidirectional_streaming.agent.agent import BidiAgent
 from strands.experimental.bidirectional_streaming.models.novasonic import BidiNovaSonicModel
 
 
@@ -130,23 +130,22 @@ async def receive(agent, context):
     """Receive and process events from agent."""
     try:
         async for event in agent.receive():
-            # Get event type
             event_type = event.get("type", "unknown")
 
-            # Handle audio stream events (bidirectional_audio_stream)
-            if event_type == "bidirectional_audio_stream":
+            # Handle audio stream events (bidi_audio_stream)
+            if event_type == "bidi_audio_stream":
                 if not context.get("interrupted", False):
                     # Decode base64 audio string to bytes for playback
                     audio_b64 = event["audio"]
                     audio_data = base64.b64decode(audio_b64)
                     context["audio_out"].put_nowait(audio_data)
 
-            # Handle interruption events (bidirectional_interruption)
-            elif event_type == "bidirectional_interruption":
+            # Handle interruption events (bidi_interruption)
+            elif event_type == "bidi_interruption":
                 context["interrupted"] = True
 
-            # Handle transcript events (bidirectional_transcript_stream)
-            elif event_type == "bidirectional_transcript_stream":
+            # Handle transcript events (bidi_transcript_stream)
+            elif event_type == "bidi_transcript_stream":
                 text_content = event.get("text", "")
                 role = event.get("role", "unknown")
 
@@ -156,10 +155,29 @@ async def receive(agent, context):
                 elif role == "assistant":
                     print(f"Assistant: {text_content}")
 
-            # Handle turn complete events (bidirectional_turn_complete)
-            elif event_type == "bidirectional_turn_complete":
+            # Handle response complete events (bidi_response_complete)
+            elif event_type == "bidi_response_complete":
                 # Reset interrupted state since the turn is complete
                 context["interrupted"] = False
+            
+            # Handle tool use events (tool_use_stream)
+            elif event_type == "tool_use_stream":
+                tool_use = event.get("current_tool_use", {})
+                tool_name = tool_use.get("name", "unknown")
+                tool_input = tool_use.get("input", {})
+                print(f"🔧 Tool called: {tool_name} with input: {tool_input}")
+            
+            # Handle tool result events (tool_result)
+            elif event_type == "tool_result":
+                tool_result = event.get("tool_result", {})
+                tool_name = tool_result.get("name", "unknown")
+                result_content = tool_result.get("content", [])
+                result_text = ""
+                for block in result_content:
+                    if isinstance(block, dict) and block.get("type") == "text":
+                        result_text = block.get("text", "")
+                        break
+                print(f"✅ Tool result from {tool_name}: {result_text}")
 
     except asyncio.CancelledError:
         pass
@@ -199,7 +217,7 @@ async def main(duration=180):
 
     # Initialize model and agent
     model = BidiNovaSonicModel(region="us-east-1")
-    agent = BidirectionalAgent(model=model, tools=[calculator], system_prompt="You are a helpful assistant.")
+    agent = BidiAgent(model=model, tools=[calculator], system_prompt="You are a helpful assistant.")
 
     await agent.start()
 
@@ -208,7 +226,7 @@ async def main(duration=180):
         "active": True,
         "audio_in": asyncio.Queue(),
         "audio_out": asyncio.Queue(),
-        "connection": agent._session,
+        "connection": agent._agent_loop,
         "duration": duration,
         "start_time": time.time(),
         "interrupted": False,

src/strands/experimental/bidirectional_streaming/scripts/test_bidi_openai.py

@@ -14,7 +14,7 @@
 import pyaudio
 from strands_tools import calculator
 
-from strands.experimental.bidirectional_streaming.agent.agent import BidirectionalAgent
+from strands.experimental.bidirectional_streaming.agent.agent import BidiAgent
 from strands.experimental.bidirectional_streaming.models.openai import BidiOpenAIRealtimeModel
 
 
@@ -122,18 +122,18 @@ async def receive(agent, context):
             # Get event type
             event_type = event.get("type", "unknown")
 
-            # Handle audio stream events (bidirectional_audio_stream)
-            if event_type == "bidirectional_audio_stream":
+            # Handle audio stream events (bidi_audio_stream)
+            if event_type == "bidi_audio_stream":
                 # Decode base64 audio string to bytes for playback
                 audio_b64 = event["audio"]
                 audio_data = base64.b64decode(audio_b64)
 
                 if not context.get("interrupted", False):
                     await context["audio_out"].put(audio_data)
 
-            # Handle transcript events (bidirectional_transcript_stream)
-            elif event_type == "bidirectional_transcript_stream":
-                source = event.get("source", "assistant")
+            # Handle transcript events (bidi_transcript_stream)
+            elif event_type == "bidi_transcript_stream":
+                source = event.get("role", "assistant")
                 text = event.get("text", "").strip()
 
                 if text:
@@ -142,25 +142,44 @@ async def receive(agent, context):
                     elif source == "assistant":
                         print(f"🔊 Assistant: {text}")
 
-            # Handle interruption events (bidirectional_interruption)
-            elif event_type == "bidirectional_interruption":
+            # Handle interruption events (bidi_interruption)
+            elif event_type == "bidi_interruption":
                 context["interrupted"] = True
                 print("⚠️  Interruption detected")
 
-            # Handle session start events (bidirectional_session_start)
-            elif event_type == "bidirectional_session_start":
+            # Handle connection start events (bidi_connection_start)
+            elif event_type == "bidi_connection_start":
                 print(f"✓ Session started: {event.get('model', 'unknown')}")
 
-            # Handle session end events (bidirectional_session_end)
-            elif event_type == "bidirectional_session_end":
+            # Handle connection close events (bidi_connection_close)
+            elif event_type == "bidi_connection_close":
                 print(f"✓ Session ended: {event.get('reason', 'unknown')}")
                 context["active"] = False
                 break
 
-            # Handle turn complete events (bidirectional_turn_complete)
-            elif event_type == "bidirectional_turn_complete":
+            # Handle response complete events (bidi_response_complete)
+            elif event_type == "bidi_response_complete":
                 # Reset interrupted state since the turn is complete
                 context["interrupted"] = False
+            
+            # Handle tool use events (tool_use_stream)
+            elif event_type == "tool_use_stream":
+                tool_use = event.get("current_tool_use", {})
+                tool_name = tool_use.get("name", "unknown")
+                tool_input = tool_use.get("input", {})
+                print(f"🔧 Tool called: {tool_name} with input: {tool_input}")
+            
+            # Handle tool result events (tool_result)
+            elif event_type == "tool_result":
+                tool_result = event.get("tool_result", {})
+                tool_name = tool_result.get("name", "unknown")
+                result_content = tool_result.get("content", [])
+                result_text = ""
+                for block in result_content:
+                    if isinstance(block, dict) and block.get("type") == "text":
+                        result_text = block.get("text", "")
+                        break
+                print(f"✅ Tool result from {tool_name}: {result_text}")
 
     except asyncio.CancelledError:
         pass
@@ -246,7 +265,7 @@ async def main():
     )
 
     # Create agent
-    agent = BidirectionalAgent(
+    agent = BidiAgent(
         model=model,
         tools=[calculator],
         system_prompt="You are a helpful voice assistant. Keep your responses brief and natural. Say hello when you first connect."

src/strands/experimental/bidirectional_streaming/scripts/test_gemini_live.py

@@ -37,13 +37,11 @@
 import pyaudio
 from strands_tools import calculator
 
-from strands.experimental.bidirectional_streaming.agent.agent import BidirectionalAgent
+from strands.experimental.bidirectional_streaming.agent.agent import BidiAgent
 from strands.experimental.bidirectional_streaming.models.gemini_live import BidiGeminiLiveModel
 
-# Configure logging - debug only for Gemini Live, info for everything else
+# Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-gemini_logger = logging.getLogger('strands.experimental.bidirectional_streaming.models.gemini_live')
-gemini_logger.setLevel(logging.DEBUG)
 logger = logging.getLogger(__name__)
 
 
@@ -145,58 +143,57 @@ async def receive(agent, context):
     """Receive and process events from agent."""
     try:
         async for event in agent.receive():
-            # Debug: Log event type and keys
             event_type = event.get("type", "unknown")
-            event_keys = list(event.keys())
-            logger.debug(f"Received event type: {event_type}, keys: {event_keys}")
 
-            # Handle audio stream events (bidirectional_audio_stream)
-            if event_type == "bidirectional_audio_stream":
+            # Handle audio stream events (bidi_audio_stream)
+            if event_type == "bidi_audio_stream":
                 if not context.get("interrupted", False):
                     # Decode base64 audio string to bytes for playback
                     audio_b64 = event["audio"]
                     audio_data = base64.b64decode(audio_b64)
                     context["audio_out"].put_nowait(audio_data)
-                    logger.info(f"🔊 Audio queued for playback: {len(audio_data)} bytes")
 
-            # Handle interruption events (bidirectional_interruption)
-            elif event_type == "bidirectional_interruption":
+            # Handle interruption events (bidi_interruption)
+            elif event_type == "bidi_interruption":
                 context["interrupted"] = True
-                logger.info("Interruption detected")
+                print("⚠️  Interruption detected")
 
-            # Handle transcript events (bidirectional_transcript_stream)
-            elif event_type == "bidirectional_transcript_stream":
+            # Handle transcript events (bidi_transcript_stream)
+            elif event_type == "bidi_transcript_stream":
                 transcript_text = event.get("text", "")
-                transcript_source = event.get("source", "unknown")
+                transcript_role = event.get("role", "unknown")
                 is_final = event.get("is_final", False)
 
                 # Print transcripts with special formatting
-                if transcript_source == "user":
+                if transcript_role == "user":
                     print(f"🎤 User: {transcript_text}")
-                elif transcript_source == "assistant":
+                elif transcript_role == "assistant":
                     print(f"🔊 Assistant: {transcript_text}")
 
-            # Handle turn complete events (bidirectional_turn_complete)
-            elif event_type == "bidirectional_turn_complete":
-                logger.debug("Turn complete - model ready for next input")
-                # Reset interrupted state since the turn is complete
+            # Handle response complete events (bidi_response_complete)
+            elif event_type == "bidi_response_complete":
+                # Reset interrupted state since the response is complete
                 context["interrupted"] = False
 
-            # Handle session start events (bidirectional_session_start)
-            elif event_type == "bidirectional_session_start":
-                logger.info(f"Session started: {event.get('model', 'unknown')}")
+            # Handle tool use events (tool_use_stream)
+            elif event_type == "tool_use_stream":
+                tool_use = event.get("current_tool_use", {})
+                tool_name = tool_use.get("name", "unknown")
+                tool_input = tool_use.get("input", {})
+                print(f"🔧 Tool called: {tool_name} with input: {tool_input}")
 
-            # Handle session end events (bidirectional_session_end)
-            elif event_type == "bidirectional_session_end":
-                logger.info(f"Session ended: {event.get('reason', 'unknown')}")
-            
-            # Handle error events (bidirectional_error)
-            elif event_type == "bidirectional_error":
-                logger.error(f"Error: {event.get('error_message', 'unknown')}")
-            
-            # Handle turn start events (bidirectional_turn_start)
-            elif event_type == "bidirectional_turn_start":
-                logger.debug(f"Turn started: {event.get('response_id', 'unknown')}")
+            # Handle tool result events (tool_result)
+            elif event_type == "tool_result":
+                tool_result = event.get("tool_result", {})
+                tool_name = tool_result.get("name", "unknown")
+                result_content = tool_result.get("content", [])
+                # Extract text from content blocks
+                result_text = ""
+                for block in result_content:
+                    if isinstance(block, dict) and block.get("type") == "text":
+                        result_text = block.get("text", "")
+                        break
+                print(f"✅ Tool result from {tool_name}: {result_text}")
 
     except asyncio.CancelledError:
         pass
@@ -325,7 +322,7 @@ async def main(duration=180):
     logger.info("Gemini Live model initialized successfully")
     print("Using Gemini Live model")
 
-    agent = BidirectionalAgent(
+    agent = BidiAgent(
         model=model, 
         tools=[calculator], 
         system_prompt="You are a helpful assistant."
@@ -338,7 +335,7 @@ async def main(duration=180):
         "active": True,
         "audio_in": asyncio.Queue(),
         "audio_out": asyncio.Queue(),
-        "connection": agent._session,
+        "connection": agent._agent_loop,
         "duration": duration,
         "start_time": time.time(),
         "interrupted": False,

src/strands/types/_events.py

@@ -145,7 +145,7 @@ class ToolUseStreamEvent(ModelStreamEvent):
 
     def __init__(self, delta: ContentBlockDelta, current_tool_use: dict[str, Any]) -> None:
         """Initialize with delta and current tool use state."""
-        super().__init__({"delta": delta, "current_tool_use": current_tool_use})
+        super().__init__({"type": "tool_use_stream", "delta": delta, "current_tool_use": current_tool_use})
 
 
 class TextStreamEvent(ModelStreamEvent):
@@ -281,7 +281,7 @@ def __init__(self, tool_result: ToolResult) -> None:
         Args:
             tool_result: Final result from the tool execution
         """
-        super().__init__({"tool_result": tool_result})
+        super().__init__({"type": "tool_result", "tool_result": tool_result})
 
     @property
     def tool_use_id(self) -> str:
@@ -309,7 +309,7 @@ def __init__(self, tool_use: ToolUse, tool_stream_data: Any) -> None:
             tool_use: The tool invocation producing the stream
             tool_stream_data: The yielded event from the tool execution
         """
-        super().__init__({"tool_stream_event": {"tool_use": tool_use, "data": tool_stream_data}})
+        super().__init__({"type": "tool_stream", "tool_stream_event": {"tool_use": tool_use, "data": tool_stream_data}})
 
     @property
     def tool_use_id(self) -> str: