Merge pull request #7 from brnaba-aws/feat/streaming

added streaming

Author: benlec

agent/strands_agent.py

@@ -21,25 +21,26 @@ def weather():
 agent = Agent(
     model=model,
     tools=[calculator, weather],
-    system_prompt="You're a helpful assistant. You can do simple math calculation, and tell the weather."
+    system_prompt="You're a helpful assistant. You can do simple math calculation, and tell the weather.",
+    callback_handler=None
 )
 
 @app.entrypoint
-def strands_agent_bedrock(payload):
+async def agent_invocation(payload):
     """
     Invoke the agent with a payload
-    
+
     IMPORTANT: Payload structure varies depending on invocation method:
     - Direct invocation (Python SDK, Console, agentcore CLI): {"prompt": "..."}
     - AWS SDK invocation (JS/Java/etc via InvokeAgentRuntimeCommand): {"input": {"prompt": "..."}}
-    
+
     The AWS SDK automatically wraps payloads in an "input" field as part of the API contract.
     This function handles both formats for maximum compatibility.
     """
     # Handle both dict and string payloads
     if isinstance(payload, str):
         payload = json.loads(payload)
-    
+
     # Extract the prompt from the payload
     # Try AWS SDK format first (most common for production): {"input": {"prompt": "..."}}
     # Fall back to direct format: {"prompt": "..."}
@@ -49,14 +50,19 @@ def strands_agent_bedrock(payload):
             user_input = payload["input"].get("prompt")
         else:
             user_input = payload.get("prompt")
-    
+
     if not user_input:
         raise ValueError(f"No prompt found in payload. Expected {{'prompt': '...'}} or {{'input': {{'prompt': '...'}}}}. Received: {payload}")
-    
-    response = agent(user_input)
-    response_text = response.message['content'][0]['text']
-    
-    return response_text
+
+    # response = agent(user_input)
+    # response_text = response.message['content'][0]['text']
+    stream = agent.stream_async(user_input)
+    async for event in stream:
+        if (event.get('event',{}).get('contentBlockDelta',{}).get('delta',{}).get('text')):
+            print(event.get('event',{}).get('contentBlockDelta',{}).get('delta',{}).get('text'))
+            yield (event.get('event',{}).get('contentBlockDelta',{}).get('delta',{}).get('text'))
+
+    # return response_text
 
 if __name__ == "__main__":
     app.run()

frontend/src/App.tsx

@@ -191,20 +191,54 @@ function App() {
     const currentPrompt = prompt;
     setPrompt('');
 
+    // Add a placeholder message for the streaming response
+    const streamingMessageIndex = messages.length + 1;
+    setMessages(prev => [...prev, {
+      type: 'agent',
+      content: '',
+      timestamp: new Date()
+    }]);
+
     try {
-      const data = await invokeAgent({ prompt: currentPrompt });
+      let streamedContent = '';
+
+      const data = await invokeAgent({
+        prompt: currentPrompt,
+        onChunk: (chunk: string) => {
+          // Accumulate the streamed content
+          streamedContent += chunk;
+
+          // Update the last message with the streamed content
+          setMessages(prev => {
+            const updated = [...prev];
+            updated[streamingMessageIndex] = {
+              type: 'agent',
+              content: streamedContent,
+              timestamp: new Date()
+            };
+            return updated;
+          });
+        }
+      });
 
-      const agentMessage: Message = {
-        type: 'agent',
-        content: cleanResponse(data.response || ''),
-        timestamp: new Date()
-      };
+      // Update with the final cleaned response
+      const finalContent = cleanResponse(data.response || streamedContent);
+      setMessages(prev => {
+        const updated = [...prev];
+        updated[streamingMessageIndex] = {
+          type: 'agent',
+          content: finalContent,
+          timestamp: new Date()
+        };
+        return updated;
+      });
 
-      setMessages(prev => [...prev, agentMessage]);
       // Show support prompts after agent responds
       setShowSupportPrompts(true);
     } catch (err: any) {
       setError(err.message);
+      // Remove the placeholder message on error
+      setMessages(prev => prev.slice(0, -1));
     } finally {
       setLoading(false);
     }
@@ -574,4 +608,4 @@ function App() {
   );
 }
 
-export default App;
+export default App;

frontend/src/agentcore.ts

@@ -8,6 +8,7 @@ const localAgentUrl = (import.meta as any).env.VITE_AGENT_RUNTIME_URL || '/api';
 
 export interface InvokeAgentRequest {
   prompt: string;
+  onChunk?: (chunk: string) => void;
 }
 
 export interface InvokeAgentResponse {
@@ -20,7 +21,7 @@ export const invokeAgent = async (request: InvokeAgentRequest): Promise<InvokeAg
     if (isLocalDev) {
       console.log('Invoking local AgentCore:', { url: localAgentUrl });
       console.log('Request payload:', { prompt: request.prompt });
-      
+
       const response = await fetch(`${localAgentUrl}/invocations`, {
         method: 'POST',
         headers: {
@@ -30,15 +31,65 @@ export const invokeAgent = async (request: InvokeAgentRequest): Promise<InvokeAg
           prompt: request.prompt
         }),
       });
-      
+
       console.log('Local AgentCore response status:', response.status);
-      
+
       if (!response.ok) {
         const errorText = await response.text();
         console.error('Local AgentCore error response:', errorText);
         throw new Error(`Local AgentCore invocation failed: ${response.status} ${response.statusText} - ${errorText}`);
       }
 
+      // Check if streaming callback is provided
+      if (request.onChunk && response.body) {
+        const reader = response.body.getReader();
+        const decoder = new TextDecoder();
+        let fullResponse = '';
+        let buffer = '';
+
+        try {
+          while (true) {
+            const { done, value } = await reader.read();
+
+            if (done) {
+              break;
+            }
+
+            const chunk = decoder.decode(value, { stream: true });
+            buffer += chunk;
+
+            // Process complete SSE messages in the buffer
+            const lines = buffer.split('\n');
+            // Keep the last incomplete line in the buffer
+            buffer = lines.pop() || '';
+
+            for (const line of lines) {
+              if (line.startsWith('data: ')) {
+                const data = line.slice(6); // Remove 'data: ' prefix
+
+                // Try to parse as JSON string
+                try {
+                  const parsed = JSON.parse(data);
+                  fullResponse += parsed;
+                  // Call the chunk callback with parsed content
+                  request.onChunk(parsed);
+                } catch {
+                  // If not JSON, use the raw data
+                  fullResponse += data;
+                  request.onChunk(data);
+                }
+              }
+            }
+          }
+
+          console.log('Streaming completed. Full response:', fullResponse);
+          return { response: fullResponse };
+        } finally {
+          reader.releaseLock();
+        }
+      }
+
+      // Non-streaming mode (backward compatibility)
       let data;
       try {
         data = await response.json();
@@ -49,7 +100,7 @@ export const invokeAgent = async (request: InvokeAgentRequest): Promise<InvokeAg
         console.log('Raw response text:', textResponse);
         throw new Error(`Invalid JSON response from local AgentCore: ${textResponse}`);
       }
-      
+
       // Handle different response formats from AgentCore
       let responseText = '';
       if (typeof data === 'string') {
@@ -59,9 +110,9 @@ export const invokeAgent = async (request: InvokeAgentRequest): Promise<InvokeAg
       } else {
         responseText = 'No response from agent';
       }
-      
+
       console.log('Final response text:', responseText);
-      
+
       return {
         response: responseText
       };
@@ -82,13 +133,13 @@ export const invokeAgent = async (request: InvokeAgentRequest): Promise<InvokeAg
 
     // URL encode the agent runtime ARN for the API call (as per AWS documentation)
     const encodedAgentRuntimeArn = encodeURIComponent(agentRuntimeArn);
-    
+
     // Use the correct AgentCore endpoint format from AWS documentation
     const url = `https://bedrock-agentcore.${region}.amazonaws.com/runtimes/${encodedAgentRuntimeArn}/invocations?qualifier=DEFAULT`;
-    
+
     console.log('Invoking AgentCore directly:', { url, agentRuntimeArn, region });
     console.log('Request payload:', { prompt: request.prompt });
-    
+
     const response = await fetch(url, {
       method: 'POST',
       headers: {
@@ -101,7 +152,7 @@ export const invokeAgent = async (request: InvokeAgentRequest): Promise<InvokeAg
         prompt: request.prompt
       }),
     });
-    
+
     console.log('AgentCore response status:', response.status);
     console.log('AgentCore response headers:', Object.fromEntries(response.headers.entries()));
 
@@ -111,6 +162,56 @@ export const invokeAgent = async (request: InvokeAgentRequest): Promise<InvokeAg
       throw new Error(`AgentCore invocation failed: ${response.status} ${response.statusText} - ${errorText}`);
     }
 
+    // Check if streaming callback is provided
+    if (request.onChunk && response.body) {
+      const reader = response.body.getReader();
+      const decoder = new TextDecoder();
+      let fullResponse = '';
+      let buffer = '';
+
+      try {
+        while (true) {
+          const { done, value } = await reader.read();
+
+          if (done) {
+            break;
+          }
+
+          const chunk = decoder.decode(value, { stream: true });
+          buffer += chunk;
+
+          // Process complete SSE messages in the buffer
+          const lines = buffer.split('\n');
+          // Keep the last incomplete line in the buffer
+          buffer = lines.pop() || '';
+
+          for (const line of lines) {
+            if (line.startsWith('data: ')) {
+              const data = line.slice(6); // Remove 'data: ' prefix
+
+              // Try to parse as JSON string
+              try {
+                const parsed = JSON.parse(data);
+                fullResponse += parsed;
+                // Call the chunk callback with parsed content
+                request.onChunk(parsed);
+              } catch {
+                // If not JSON, use the raw data
+                fullResponse += data;
+                request.onChunk(data);
+              }
+            }
+          }
+        }
+
+        console.log('Streaming completed. Full response:', fullResponse);
+        return { response: fullResponse };
+      } finally {
+        reader.releaseLock();
+      }
+    }
+
+    // Non-streaming mode (backward compatibility)
     let data;
     try {
       data = await response.json();
@@ -121,7 +222,7 @@ export const invokeAgent = async (request: InvokeAgentRequest): Promise<InvokeAg
       console.log('Raw response text:', textResponse);
       throw new Error(`Invalid JSON response from AgentCore: ${textResponse}`);
     }
-    
+
     // Handle different response formats from AgentCore
     let responseText = '';
     if (typeof data === 'string') {
@@ -131,9 +232,9 @@ export const invokeAgent = async (request: InvokeAgentRequest): Promise<InvokeAg
     } else {
       responseText = 'No response from agent';
     }
-    
+
     console.log('Final response text:', responseText);
-    
+
     return {
       response: responseText
     };
@@ -142,4 +243,4 @@ export const invokeAgent = async (request: InvokeAgentRequest): Promise<InvokeAg
     console.error('AgentCore invocation error:', error);
     throw new Error(`Failed to invoke agent: ${error.message}`);
   }
-};
+};