Update main.py

main.py

@@ -15,8 +15,8 @@ WEBSOCKET_URL = "wss://4d24-196-75-8-109.ngrok-free.app/ws"
 
 class Item(BaseModel):
     prompt: str
-    history: list
-    system_prompt: str
+    history: list = []
+    system_prompt: str = ""
     temperature: float = 0.0
     max_new_tokens: int = 1048
     top_p: float = 0.15
@@ -50,19 +50,29 @@ async def generate_stream(item: Item):
         yield response.token.text
 
 async def websocket_client():
+    call_count = 0
     while True:
         try:
             async with websockets.connect(WEBSOCKET_URL) as websocket:
                 logger.info("WebSocket connection established")
                 while True:
+                    # Request a prompt
+                    await websocket.send(json.dumps({"type": "getPrompt"}))
+                    
                     message = await websocket.recv()
                     data = json.loads(message)
                     logger.info(f"Received: {data}")
 
                     if data.get('type') == 'prompt':
+                        call_count += 1
                         prompt_id = data.get('id')
-                        item = Item(**data['prompt'])
+                        prompt_text = data.get('prompt')
                         
+                        logger.info(f"Processing prompt: {prompt_text}")
+                        logger.info(f"Call count: {call_count}")
+
+                        item = Item(prompt=prompt_text)
+
                         async for chunk in generate_stream(item):
                             await websocket.send(json.dumps({
                                 "type": "chunk",
@@ -70,12 +80,18 @@ async def websocket_client():
                                 "chunk": chunk
                             }))
                             logger.info(f"Sent chunk: {chunk}")
-                        
+
                         await websocket.send(json.dumps({
                             "type": "completed",
                             "id": prompt_id
                         }))
                         logger.info("Generation completed")
+                    elif data.get('type') == 'error':
+                        logger.info(f"Received error: {data.get('message')}")
+                        # Wait a bit before requesting a new prompt
+                        await asyncio.sleep(5)
+                    else:
+                        logger.info(f"Received unexpected message type: {data.get('type')}")
 
         except websockets.exceptions.ConnectionClosed:
             logger.error("WebSocket connection closed. Retrying in 5 seconds...")