import os import uvicorn import re import numpy as np import matplotlib matplotlib.use('Agg') # Prevents extra GUI windows from opening on your Ubuntu desktop import matplotlib.pyplot as plt import mpld3 import plotly.graph_objects as go from fastapi import FastAPI, Request from fastapi.responses import HTMLResponse from llama_cpp import Llama app = FastAPI() # 1. GPU INITIALIZATION (RTX 4060 Ti) llm = Llama( model_path="/mnt/ai_data/math_APP/qwen_math_q4_k_m.gguf", n_gpu_layers=-1, n_ctx=2048, verbose=False ) def fix_plotly_colorscales(code): """Intercepts Matplotlib colormap names and swaps them for Plotly equivalents.""" # Maps common Matplotlib names to their closest Plotly counterparts color_map = { "spring": "Viridis", "summer": "Plasma", "autumn": "Inferno", "winter": "Cividis", "magma": "Magma" } for mpl_name, plotly_name in color_map.items(): # Handles both single and double quotes, and case sensitivity code = re.sub(f"['\"]{mpl_name}['\"]", f"'{plotly_name}'", code, flags=re.IGNORECASE) return code @app.get("/", response_class=HTMLResponse) async def index(): return """ Interactive Math Core

Math Visualizer Core

Resilient Engine v2.0

                    

                


                

                    

                        
                        

                            

                                Ready for computation.
                            

                        

                    

                

            

        


        
    
    
    """

@app.post("/process")
async def process_request(request: Request):
    data = await request.json()
    user_prompt = data.get("prompt")
    
    instruction = (
        "Output ONLY raw Python code. Use plotly.graph_objects (go) for 3D/animations "
        "and name the figure 'fig'. Use matplotlib.pyplot (plt) for 2D. "
        "CRITICAL: Never use plt.show() or fig.show(). Use Plotly-safe colorscales."
    )
    
    formatted = f"<|im_start|>system\n{instruction}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
    output = llm(formatted, max_tokens=1536, stop=["<|im_end|>"])
    code = re.sub(r"```python|```", "", output['choices'][0]['text']).strip()
    
    # Apply the colorscale fix
    code = fix_plotly_colorscales(code)

    try:
        plt.clf()
        plt.close('all')
        
        # Scope with "Dummy Show" to prevent Ubuntu window popups
        def dummy_show(*args, **kwargs): pass
        
        exec_scope = {
            "plt": plt, "np": np, "go": go, "mpld3": mpld3, 
            "__builtins__": __builtins__
        }
        plt.show = dummy_show 
        
        exec(code, exec_scope)
        
        if "fig" in exec_scope:
            # Captures Plotly (3D/Animated)
            html_output = exec_scope["fig"].to_html(full_html=False, include_plotlyjs='cdn')
        else:
            # Captures Matplotlib (2D)
            html_output = mpld3.fig_to_html(plt.gcf())

        return {"code": code, "html": html_output}

    except Exception as e:
        return {"code": code, "error": str(e)}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)