Fixing generate and random problem API

app.py

@@ -8,8 +8,10 @@ import os
 import jwt
 from typing import Dict, Any
 import autopep8
+import textwrap
 
 from datasets import load_dataset
+from fastapi.responses import StreamingResponse
 import random
 
 # Load the dataset (you might want to do this once at the start of your script)
@@ -52,7 +54,36 @@ def verify_token(token: str) -> bool:
         return True
     except jwt.PyJWTError:
         return False
+
+def stream_solution(instruction: str, token: str):
+    if not verify_token(token):
+        return {"error": "Invalid token"}
     
+    system_prompt = "You are a Python coding assistant specialized in solving LeetCode problems. Provide only the complete implementation of the given function. Ensure proper indentation and formatting. Do not include any explanations or multiple solutions."
+    full_prompt = f"""### Instruction:
+{system_prompt}
+
+Implement the following function for the LeetCode problem:
+
+{instruction}
+
+### Response:
+Here's the complete Python function implementation:
+
+```python
+"""
+    
+    def generate():
+        generated_text = ""
+        for chunk in llm(full_prompt, stream=True, **generation_kwargs):
+            token = chunk["choices"][0]["text"]
+            generated_text += token
+            yield token
+        
+        formatted_code = extract_and_format_code(generated_text)
+        yield formatted_code
+
+    return StreamingResponse(generate(), media_type="text/plain")
 
 def extract_and_format_code(text):
     # Extract code between triple backticks