Leetmonkey In Action via Inference

app.py

@@ -8,6 +8,9 @@ from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
 import jwt
 from typing import Generator
+from fastapi import FastAPI, HTTPException, Depends
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from pydantic import BaseModel
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -15,6 +18,8 @@ logger = logging.getLogger(__name__)
 
 # JWT settings
 JWT_SECRET = os.environ.get("JWT_SECRET")
+if not JWT_SECRET:
+    raise ValueError("JWT_SECRET environment variable is not set")
 JWT_ALGORITHM = "HS256"
 
 # Model settings
@@ -108,73 +113,88 @@ def extract_and_format_code(text: str) -> str:
     except:
         return formatted_code
 
-def verify_token(token: str) -> bool:
+security = HTTPBearer()
+app = FastAPI()
+
+class ProblemRequest(BaseModel):
+    instruction: str
+
+def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
     try:
-        jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
+        jwt.decode(credentials.credentials, JWT_SECRET, algorithms=[JWT_ALGORITHM])
         return True
     except jwt.PyJWTError:
-        return False
+        raise HTTPException(status_code=401, detail="Invalid token")
 
-def generate_solution_api(instruction: str, token: str) -> str:
-    if not verify_token(token):
-        return "Invalid token. Please provide a valid JWT token."
-    
+@app.post("/generate_solution")
+async def generate_solution_api(request: ProblemRequest, authorized: bool = Depends(verify_token)):
     logger.info("Generating solution")
-    generated_output = generate_solution(instruction)
+    generated_output = generate_solution(request.instruction)
     formatted_code = extract_and_format_code(generated_output)
     logger.info("Solution generated successfully")
-    return formatted_code
-
-def stream_solution_api(instruction: str, token: str) -> Generator[str, None, None]:
-    if not verify_token(token):
-        yield "Invalid token. Please provide a valid JWT token."
-        return
-
-    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:
+    return {"solution": formatted_code}
+
+@app.post("/stream_solution")
+async def stream_solution_api(request: ProblemRequest, authorized: bool = Depends(verify_token)):
+    async def generate():
+        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:
 
-{instruction}
+{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"]["text"]
-        generated_text += token
-        yield generated_text
-    
-    formatted_code = extract_and_format_code(generated_text)
-    logger.info("Solution generated successfully")
-    yield formatted_code
-
-# Gradio interface
-def gradio_generate(instruction: str, token: str) -> str:
-    return generate_solution_api(instruction, token)
-
-def gradio_stream(instruction: str, token: str) -> str:
-    return "".join(list(stream_solution_api(instruction, token)))
-
-iface = gr.Interface(
-    fn=[gradio_generate, gradio_stream],
-    inputs=[
-        gr.Textbox(label="LeetCode Problem Instruction"),
-        gr.Textbox(label="JWT Token")
-    ],
-    outputs=[
-        gr.Code(label="Generated Solution"),
-        gr.Code(label="Streamed Solution")
-    ],
-    title="LeetCode Problem Solver",
-    description="Enter a LeetCode problem instruction and your JWT token to generate a solution."
-)
+        
+        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 generate()
+
+# Gradio wrapper for FastAPI
+def gradio_wrapper(app):
+    def inference(instruction, token):
+        import requests
+        url = "http://localhost:8000/generate_solution"
+        headers = {"Authorization": f"Bearer {token}"}
+        response = requests.post(url, json={"instruction": instruction}, headers=headers)
+        if response.status_code == 200:
+            return response.json()["solution"]
+        else:
+            return f"Error: {response.status_code}, {response.text}"
+
+    iface = gr.Interface(
+        fn=inference,
+        inputs=[
+            gr.Textbox(label="LeetCode Problem Instruction"),
+            gr.Textbox(label="JWT Token")
+        ],
+        outputs=gr.Code(label="Generated Solution"),
+        title="LeetCode Problem Solver API",
+        description="Enter a LeetCode problem instruction and your JWT token to generate a solution."
+    )
+    return iface
 
 if __name__ == "__main__":
-    iface.launch(share=True)
+    import uvicorn
+    from threading import Thread
+
+    # Start FastAPI in a separate thread
+    Thread(target=lambda: uvicorn.run(app, host="0.0.0.0", port=8000)).start()
+
+    # Launch Gradio interface
+    iface = gradio_wrapper(app)
+    iface.launch(share=True)

requirements.txt

@@ -2,4 +2,7 @@ gradio
 llama-cpp-python
 huggingface_hub
 pyjwt
-autopep8
+autopep8
+fastapi
+uvicorn
+pydantic