Debugging shit

app.py

@@ -10,6 +10,7 @@ from typing import Dict, Any
 import autopep8
 import textwrap
 import json
+import asyncio
 
 from datasets import load_dataset
 from fastapi.responses import StreamingResponse
@@ -116,7 +117,7 @@ Here's the complete Python function implementation:
     formatted_code = extract_and_format_code(generated_text)
     return {"solution": formatted_code}
 
-def stream_solution(instruction: str, token: str):
+async def stream_solution(instruction: str, token: str):
     if not verify_token(token):
         return {"error": "Invalid token"}
     
@@ -133,29 +134,35 @@ Here's the complete Python function implementation:
 
 ```python
 """
-    
+    generated_text = ""
     async def generate():
-        generated_text = ""
+        
         try:
-            for chunk in llm(full_prompt, stream=True, **generation_kwargs):
-                token = chunk["choices"][0]["text"]
-                generated_text += token
-                logger.info(f"Generated text: {generated_text}")
-                yield token
-
-                # Optionally send progress updates:
-                progress = len(generated_text) / len(full_prompt)
-                yield {"progress": progress}  # Send a progress update
+            # Set a timeout of 2 minutes
+            timeout = asyncio.TimeoutError("LLM generation timed out")
+            async with asyncio.TimeoutManager(timeout):
+                for chunk in llm(full_prompt, stream=True, **generation_kwargs):
+                    token = chunk["choices"][0]["text"]
+                    generated_text += token
+                    logger.info(f"Generated text: {generated_text}")
+                    yield token
+
+                    # Optionally send progress updates:
+                    progress = len(generated_text) / len(full_prompt)
+                    yield {"progress": progress}  # Send a progress update
         
-            formatted_code = extract_and_format_code(generated_text)
-            logger.info(f"Formatted code: {formatted_code}")
-            yield formatted_code
+        except asyncio.TimeoutError:
+            logger.error("LLM generation timed out")
+            yield {"error": "LLM generation timed out"}
         except Exception as e:
             logger.error(f"Error generating solution: {e}")
             yield {"error": "Error generating solution"}
 
+    formatted_code = extract_and_format_code(generated_text)
+    logger.info(f"Formatted code: {formatted_code}")
+    yield formatted_code
 
-    generated_response = generate()
+    generated_response = await generate()
     logger.info(f"Streaming response {generated_response}")
     return StreamingResponse(generate(), media_type="text/plain")