🔄 Refactor streaming logic and update chat message content in AI test scripts

projects/local-ai-test/main-stream.py

@@ -1,44 +1,55 @@
 import os
+from datetime import datetime
 from pprint import pprint
 
 from openai import OpenAI
 from openai._types import Timeout
-from openai.types.chat import ChatCompletionUserMessageParam, ChatCompletionSystemMessageParam
+from openai.types.chat import ChatCompletionUserMessageParam, ChatCompletionSystemMessageParam, \
+    ChatCompletionStreamOptionsParam
 
 
 def main():
     client = OpenAI(
         api_key=os.environ.get("OPENAI_API_KEY", "sk-dummy"),
         base_url=os.environ.get("OPENAI_API_BASE", "http://localhost:11434/v1"),
-        timeout=Timeout(30, connect=5),
+        timeout=Timeout(10, connect=5),
     )
 
-    response = client.chat.completions.create(
+    print("Start", datetime.now())
+    try:
+        stream = client.chat.completions.create(
             model=os.environ.get("OPENAI_MODEL", "qwen-math-turbo"),
             messages=[
                 ChatCompletionSystemMessageParam(role="system", content="You are a useful assistant."),
                 # ChatCompletionUserMessageParam(**{"role": "user", "content": "Hello!"}),
-            ChatCompletionUserMessageParam(role="user", content="Hello!"),
+                ChatCompletionUserMessageParam(role="user", content="Introduce your self."),
             ],
             stream=True,
-        # stream_options={"include_usage": True},
+            stream_options=ChatCompletionStreamOptionsParam(include_usage=True),
         )
         print()
-    pprint(vars(response))
+        pprint(vars(stream))
         print()
-    # print(response.choices[0].message.content)
 
-    for chunk in response:
+        content_parts = []
-        print('xx')
+        print("Begin chunk", datetime.now())
+        for chunk in stream:
+            print("Chunk", datetime.now())
             if chunk.choices:
                 content = chunk.choices[0].delta.content or ""
-            print('CHUNK:', content, end="", flush=True)
+                print('CHUNK:', content, flush=True)
-            # content_parts.append(content)
+                content_parts.append(content)
             elif chunk.usage:
-            print("\n--- 请求用量 ---")
+                print("\n----- Usage -----")
-            print(f"输入 Tokens: {chunk.usage.prompt_tokens}")
+                print(f"Input  Tokens: {chunk.usage.prompt_tokens}")
-            print(f"输出 Tokens: {chunk.usage.completion_tokens}")
+                print(f"Output Tokens: {chunk.usage.completion_tokens}")
-            print(f"总计 Tokens: {chunk.usage.total_tokens}")
+                print(f"Total  Tokens: {chunk.usage.total_tokens}")
+                print("")
+        print("Finally:", "".join(content_parts))
+    except Exception as e:
+        print("Exception", datetime.now(), e)
+    finally:
+        print("End", datetime.now())
 
 
 if __name__ == "__main__":

projects/local-ai-test/main.py

@@ -18,7 +18,7 @@ def main():
         messages=[
             ChatCompletionSystemMessageParam(role="system", content="You are a useful assistant."),
             # ChatCompletionUserMessageParam(**{"role": "user", "content": "Hello!"}),
-            ChatCompletionUserMessageParam(role="user", content="Hello!"),
+            ChatCompletionUserMessageParam(role="user", content="Introduce your self in short a sentence."),
         ],
     )
     print()