Fixing generate and random problem API

app.py

@@ -9,6 +9,7 @@ import jwt
 from typing import Dict, Any
 import autopep8
 import textwrap
+import json
 
 from datasets import load_dataset
 from fastapi.responses import StreamingResponse
@@ -141,7 +142,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) -> Dict[str, Any]:
+def stream_solution(instruction: str, token: str):
     if not verify_token(token):
         return {"error": "Invalid token"}
     
@@ -160,10 +161,16 @@ Here's the complete Python function implementation:
 """
     
     def generate():
+        generated_text = ""
         for chunk in llm(full_prompt, stream=True, **generation_kwargs):
-            yield chunk["choices"][0]["text"]
-    
-    return generate()
+            token = chunk["choices"][0]["text"]
+            generated_text += token
+            yield json.dumps({"token": token, "generated_text": generated_text}) + "\n"
+        
+        formatted_code = extract_and_format_code(generated_text)
+        yield json.dumps({"complete": True, "formatted_code": formatted_code}) + "\n"
+
+    return StreamingResponse(generate(), media_type="application/x-ndjson")
 
 def random_problem(token: str) -> Dict[str, Any]:
     if not verify_token(token):