refactor: refactor the FastAPI response streaming with Claude3

examples/fastapi-response-streaming/app/main.py

@@ -13,38 +13,44 @@
 
 app.mount("/demo", StaticFiles(directory="static", html=True))
 
+
 @app.get("/")
 async def root():
     return RedirectResponse(url='/demo/')
 
+
 class Story(BaseModel):
-   topic: Optional[str] = None
+    topic: Optional[str] = None
+
 
 @app.post("/api/story")
 def api_story(story: Story):
     if story.topic == None or story.topic == "":
-       return None
+        return None
 
     return StreamingResponse(bedrock_stream(story.topic), media_type="text/html")
 
 
 bedrock = boto3.client('bedrock-runtime')
 
+
 async def bedrock_stream(topic: str):
     instruction = f"""
     You are a world class writer. Please write a sweet bedtime story about {topic}.
     """
-
     body = json.dumps({
-        'prompt': f'Human:{instruction}\n\nAssistant:', 
-        'max_tokens_to_sample': 1028,
-        'temperature': 1,
-        'top_k': 250,
-        'top_p': 0.999,
-        'stop_sequences': ['\n\nHuman:']
+        "anthropic_version": "bedrock-2023-05-31",
+        "max_tokens": 1024,
+        "messages": [
+            {
+                "role": "user",
+                "content": instruction,
+            }
+        ],
     })
+
     response = bedrock.invoke_model_with_response_stream(
-        modelId='anthropic.claude-v2',
+        modelId='anthropic.claude-3-haiku-20240307-v1:0',
         body=body
     )
 
@@ -53,8 +59,12 @@ async def bedrock_stream(topic: str):
         for event in stream:
             chunk = event.get('chunk')
             if chunk:
-                yield json.loads(chunk.get('bytes').decode())['completion']
+                message = json.loads(chunk.get("bytes").decode())
+                if message['type'] == "content_block_delta":
+                    yield message['delta']['text'] or ""
+                elif message['type'] == "message_stop":
+                    yield "\n"
 
 
 if __name__ == "__main__":
-  uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "8080")))
+    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "8080")))