Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from fastapi.responses import Response
 from fastapi.responses import FileResponse
 from pydantic import BaseModel
@@ -46,9 +46,16 @@ async def create_ticket(ticket: Ticket):
 
 
 @app.post("/run_code")
-async def run_code(code: Code):
+async def run_code(code: Code, request: Request):
     # img_buffer = io.BytesIO()
-    exec(code.code)
+
+    pattern = f"""```python([\s\S]*?)```"""
+    codes = re.findall(pattern, code.code)
+
+    if codes:
+        exec("\n".join(codes))
+    else:
+        exec(code.code)
     # img_buffer.seek(0)  # Reset buffer position
 
     file_path = f"graphs_{code.project_id}_{code.chain_id}_{code.session_id}.pdf"
@@ -79,7 +86,11 @@ async def run_code(code: Code):
     
     # return Response(content=img_buffer.getvalue(), media_type="image/png")
     # return FileResponse(file_path, media_type="image/png")
-    return FileResponse(file_path, media_type="application/pdf", filename="graph.pdf")
+    # return FileResponse(file_path, media_type="application/pdf", filename="graph.pdf")
+    return {
+        "message": "Graph created succesfully!",
+        "url": request.url + f"/chart/{code.project_id}/{code.chain_id}/{code.session_id}"
+    }