Update app.py

app.py

@@ -1,18 +1,53 @@
-from fastapi import FastAPI, Body
+
+from fastapi import FastAPI, Body, HTTPException, Request
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+from typing import List, Dict, Any
+from pathlib import Path
+from datetime import datetime
 import subprocess
+import uuid
+import pandas as pd
 
 app = FastAPI(title="Agent Sandbox API")
 
+class IncommingMail(BaseModel):
+    data: List[Dict[str, Any]]   # requires a "data" field in JSON
+    timestamp: str = datetime.utcnow().strftime("%Y%m%dT%H%M%SZ")
+    id: str = uuid.uuid4())[:5]
+    
+    @property 
+    def file_path(self):
+        return f"{self.timestamp}_{self.id}_sample.parquet"
+        
+    def store_mail(self):
+        try:
+            import pandas as pd 
+            data = pd.DataFrame(self.data)
+            data.to_parquet(self.file_path)
+        except Exception as e:
+            raise e
+
+class MailResponse(BaseModel):
+    status: str
+    timestamp: str
+    pigeon_hole: str
+    pigeon_location: str
+
+class CodeScript(BaseModel):
+    code: str
+    
 @app.post("/run")
-def run_code(payload: dict = Body(...)):
+def run_code(payload: CodeScript):
     """
     Execute Python code in a sandboxed subprocess.
     Input JSON: {"code": "print(2+2)"} or {"code": "2+2"}
     Output JSON: {"output": "4"} or {"error": "..."}
     """
-    code = payload.get("code", "")
+    code = payload.code
+    
     if not code:
-        return {"error": "No code provided"}
+        raise HTTPException(status_code=400, detail="Received Payload in incorrect format. Must pass key: 'code' with value of the script to run python code.")
 
     try:
         # Run the code in a subprocess with a 5s timeout
@@ -39,6 +74,21 @@ def run_code(payload: dict = Body(...)):
         return {"output": output}
 
     except subprocess.TimeoutExpired:
-        return {"error": "Execution timed out"}
-    except Exception as e:
-        return {"error": str(e)}
+        raise HTTPException(status_code=408, detail="Execution timed out")
+
+@app.post("/inbox", response_model=MailResponse)
+def agent_mailbox(mail: IncommingMail):
+    """
+    Save incoming data with a timestamp + unique ID.
+    """
+    # Prepare / fetch directory
+    inbox_dir = Path("inbox")
+    inbox_dir.mkdir(exist_ok=True)
+    mail.store_mail()
+    
+    return MailResponse(
+        status="Data Saved!",
+        timestamp=mail.timestamp,
+        pigeon_hole=mail.id,
+        pigeon_location=mail.file_path
+    )