finish websocket example

cookbook/pocketflow-fastapi-websocket/utils/stream_llm.py

@@ -1,22 +1,27 @@
 import os
-from openai import OpenAI
+from openai import AsyncOpenAI
 
-def stream_llm(messages):
-    client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", "your-api-key"))
+async def stream_llm(messages):
+    client = AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY", "your-api-key"))
     
-    stream = client.chat.completions.create(
+    stream = await client.chat.completions.create(
         model="gpt-4o-mini",
         messages=messages,
         stream=True,
         temperature=0.7
     )
     
-    for chunk in stream:
+    async for chunk in stream:
         if chunk.choices[0].delta.content is not None:
             yield chunk.choices[0].delta.content
 
 if __name__ == "__main__":
-    messages = [{"role": "user", "content": "Hello!"}]
-    for chunk in stream_llm(messages):
-        print(chunk, end="", flush=True)
-    print() 
+    import asyncio
+    
+    async def test():
+        messages = [{"role": "user", "content": "Hello!"}]
+        async for chunk in stream_llm(messages):
+            print(chunk, end="", flush=True)
+        print()
+    
+    asyncio.run(test())