From 61ba8cc12c673c0f2563ffc710a984c220989d36 Mon Sep 17 00:00:00 2001
From: grillazz <the@grillazz.com>
Date: Sat, 3 May 2025 08:47:08 +0200
Subject: [PATCH] refactor chat testing client

---
 tests/chat.py | 27 ++++++++++++++-------------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/tests/chat.py b/tests/chat.py
index a231c6c..f99c819 100644
--- a/tests/chat.py
+++ b/tests/chat.py
@@ -2,29 +2,30 @@ import anyio
 import httpx
 import orjson
 
+API_URL = "http://localhost:8000/chat/"
+
+
 async def chat_with_endpoint():
     async with httpx.AsyncClient() as client:
         while True:
-            # Get user input
             prompt = input("\nYou: ")
             if prompt.lower() == "exit":
                 break
 
-            # Send request to the API
             print("\nModel: ", end="", flush=True)
-            async with client.stream(
-                "POST",
-                "http://localhost:8000/chat/",
-                data={"prompt": prompt},
-                timeout=60
-            ) as response:
-                async for chunk in response.aiter_lines():
-                    if chunk:
+            try:
+                async with client.stream("POST", API_URL, data={"prompt": prompt}, timeout=60) as response:
+                    async for chunk in response.aiter_lines():
+                        if not chunk:
+                            continue
+
                         try:
-                            data = orjson.loads(chunk)
-                            print(data["content"], end="", flush=True)
+                            print(orjson.loads(chunk)["content"], end="", flush=True)
                         except Exception as e:
                             print(f"\nError parsing chunk: {e}")
+            except httpx.RequestError as e:
+                print(f"\nConnection error: {e}")
+
 
 if __name__ == "__main__":
-    anyio.run(chat_with_endpoint)
+    anyio.run(chat_with_endpoint)
\ No newline at end of file