Adding and testing explain solution

app.py

@@ -11,7 +11,6 @@ import autopep8
 import textwrap
 
 from datasets import load_dataset
-from fastapi.responses import StreamingResponse
 import random
 import asyncio
 
@@ -42,6 +41,8 @@ model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_NAME, cache_dir="./
 llm = Llama(model_path=model_path, n_ctx=1024, n_threads=8, n_gpu_layers=-1, verbose=False, mlock=True)
 logger.info("8-bit model loaded successfully")
 
+user_data = {}
+
 # Generation parameters
 generation_kwargs = {
     "max_tokens": 512,
@@ -90,36 +91,39 @@ def extract_and_format_code(text):
         return autopep8.fix_code(formatted_code)
     except:
         return formatted_code
-
-def generate_solution(instruction: str, token: str) -> Dict[str, Any]:
+    
+def generate_explanation(problem: str, solution: str, token: str) -> Dict[str, Any]:
     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."
+    system_prompt = "You are a Python coding assistant specialized in explaining LeetCode problem solutions. Provide a clear and concise explanation of the given solution."
     full_prompt = f"""### Instruction:
 {system_prompt}
 
-Implement the following function for the LeetCode problem:
+Problem:
+{problem}
 
-{instruction}
+Solution:
+{solution}
+
+Explain this solution step by step.
 
 ### Response:
-Here's the complete Python function implementation:
+Here's the explanation of the solution:
 
-```python
 """
     
     generated_text = ""
     for chunk in llm(full_prompt, stream=True, **generation_kwargs):
         generated_text += chunk["choices"][0]["text"]
     
-    formatted_code = extract_and_format_code(generated_text)
-    return {"solution": formatted_code}
+    return {"explanation": generated_text}
 
-async def stream_solution(instruction: str, token: str):
-    if not verify_token(token):
-        raise Exception("Invalid token")
 
+def generate_solution(instruction: str, token: str) -> Dict[str, Any]:
+    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}
@@ -133,20 +137,14 @@ Here's the complete Python function implementation:
 
 ```python
 """
-    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  # Yield individual tokens for streaming
-        except Exception as e:
-            logger.error(f"Error generating solution: {e}")
-            yield {"error": "Error generating solution"}
-
-    async for token in generate():
-        yield token
+    
+    generated_text = ""
+    for chunk in llm(full_prompt, stream=True, **generation_kwargs):
+        generated_text += chunk["choices"][0]["text"]
+    
+    formatted_code = extract_and_format_code(generated_text)
+    user_data[token] = {"problem": instruction, "solution": formatted_code}
+    return {"solution": formatted_code}
 
 def random_problem(token: str) -> Dict[str, Any]:
     if not verify_token(token):
@@ -157,6 +155,7 @@ def random_problem(token: str) -> Dict[str, Any]:
     
     # Extract the instruction (problem statement) from the randomly selected item
     problem = random_item['instruction']
+    user_data[token] = {"problem": problem, "solution": None}
     
     return {"problem": problem}
 
@@ -169,14 +168,6 @@ generate_interface = gr.Interface(
     description="Provide a LeetCode problem instruction and a valid JWT token to generate a solution."
 )
 
-stream_interface = gr.Interface(
-    fn=stream_solution,
-    inputs=[gr.Textbox(label="Problem Instruction"), gr.Textbox(label="JWT Token")],
-    outputs=gr.Text(),
-    title="Stream Solution API",
-    description="Provide a LeetCode problem instruction and a valid JWT token to stream a solution."
-)
-
 random_problem_interface = gr.Interface(
     fn=random_problem,
     inputs=gr.Textbox(label="JWT Token"),
@@ -185,10 +176,9 @@ random_problem_interface = gr.Interface(
     description="Provide a valid JWT token to get a random LeetCode problem."
 )
 
-# Combine interfaces
 demo = gr.TabbedInterface(
-    [generate_interface, stream_interface, random_problem_interface],
-    ["Generate Solution", "Stream Solution", "Random Problem"]
+    [generate_interface, explain_interface, random_problem_interface],
+    ["Generate Solution", "Explain Solution", "Random Problem"]
 )
 
 # Launch the Gradio app