Update main.py

main.py

@@ -6,7 +6,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from queue import Empty
 import os
 app = FastAPI()
-
+root_dir = os.getcwd()
 # @app.on_event("startup")
 # async def startup():
 #     FastAPICache.init(InMemoryBackend(), prefix="fastapi-cache")
@@ -37,8 +37,6 @@ from typing import Dict
 from datetime import datetime, timedelta
 import psutil
 
-app = FastAPI()
-
 # Model for input data
 class CodeExecutionRequest(BaseModel):
     session_token: str
@@ -70,6 +68,11 @@ async def kill_kernel(session_token: str):
 async def execute_code(session_token: str, code: str):
     setup_code = "%matplotlib inline"
     
+    session_dir = os.path.join(root_dir,output, session_token)
+    
+    if not os.path.exists(session_dir):
+        os.makedirs(session_dir)
+    
     if session_token not in kernel_managers:
         await create_kernel(session_token)
 
@@ -77,6 +80,8 @@ async def execute_code(session_token: str, code: str):
     kc = km.client()
 
     try:
+        os.chdir(session_dir)
+        print("current working directory",os.getcwd())
         # Execute setup code
         kc.execute_interactive(setup_code, store_history=False)
         # Execute the provided code