Leetmonkey In Action via Inference

app.py

@@ -3,14 +3,11 @@ import re
 import logging
 import textwrap
 import autopep8
+import gradio as gr
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
 import jwt
-from typing import AsyncGenerator
-from fastapi import FastAPI, HTTPException, Depends
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-from pydantic import BaseModel
-from fastapi.responses import StreamingResponse
+from typing import Dict, Any
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -79,32 +76,22 @@ Here's the complete Python function implementation:
     return response["choices"][0]["text"]
 
 def extract_and_format_code(text: str) -> str:
-    # Extract code between triple backticks
     code_match = re.search(r'```python\s*(.*?)\s*```', text, re.DOTALL)
     if code_match:
         code = code_match.group(1)
     else:
         code = text
 
-    # Remove any text before the function definition
     code = re.sub(r'^.*?(?=def\s+\w+\s*\()', '', code, flags=re.DOTALL)
-
-    # Dedent the code to remove any common leading whitespace
     code = textwrap.dedent(code)
-
-    # Split the code into lines
     lines = code.split('\n')
-
-    # Find the function definition line
     func_def_index = next((i for i, line in enumerate(lines) if line.strip().startswith('def ')), 0)
-
-    # Ensure proper indentation
-    indented_lines = [lines[func_def_index]]  # Keep the function definition as is
+    indented_lines = [lines[func_def_index]]
     for line in lines[func_def_index + 1:]:
-        if line.strip():  # If the line is not empty
-            indented_lines.append('    ' + line)  # Add 4 spaces of indentation
+        if line.strip():
+            indented_lines.append('    ' + line)
         else:
-            indented_lines.append(line)  # Keep empty lines as is
+            indented_lines.append(line)
 
     formatted_code = '\n'.join(indented_lines)
 
@@ -113,57 +100,38 @@ def extract_and_format_code(text: str) -> str:
     except:
         return formatted_code
 
-security = HTTPBearer()
-app = FastAPI()
-
-class ProblemRequest(BaseModel):
-    instruction: str
-
-def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
+def verify_token(token: str) -> bool:
     try:
-        jwt.decode(credentials.credentials, JWT_SECRET, algorithms=[JWT_ALGORITHM])
+        jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
         return True
     except jwt.PyJWTError:
-        raise HTTPException(status_code=401, detail="Invalid token")
+        return False
 
-@app.post("/generate_solution")
-async def generate_solution_api(request: ProblemRequest, authorized: bool = Depends(verify_token)):
+def generate_code(instruction: str, token: str) -> Dict[str, Any]:
+    if not verify_token(token):
+        return {"error": "Invalid token"}
+    
     logger.info("Generating solution")
-    generated_output = generate_solution(request.instruction)
+    generated_output = generate_solution(instruction)
     formatted_code = extract_and_format_code(generated_output)
     logger.info("Solution generated successfully")
     return {"solution": formatted_code}
 
-@app.post("/stream_solution")
-async def stream_solution_api(request: ProblemRequest, authorized: bool = Depends(verify_token)):
-    async def generate() -> AsyncGenerator[str, None]:
-        logger.info("Streaming solution")
-        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:
-
-{request.instruction}
-
-### Response:
-Here's the complete Python function implementation:
-
-```python
-"""
-        
-        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)
-        logger.info("Solution generated successfully")
-        yield formatted_code
-
-    return StreamingResponse(generate(), media_type="text/plain")
+# Gradio API
+api = gr.Interface(
+    fn=generate_code,
+    inputs=[
+        gr.Textbox(label="LeetCode Problem Instruction"),
+        gr.Textbox(label="JWT Token")
+    ],
+    outputs=gr.JSON(),
+    title="LeetCode Problem Solver API",
+    description="Provide a LeetCode problem instruction and a valid JWT token to generate a solution.",
+    examples=[
+        ["Implement a function to reverse a linked list", "your_jwt_token_here"],
+        ["Write a function to find the maximum subarray sum", "your_jwt_token_here"]
+    ]
+)
 
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    api.launch(server_name="0.0.0.0", server_port=7860)