Update api_routes_v2.py

api_routes_v2.py

@@ -403,13 +403,13 @@
 from fastapi import APIRouter, HTTPException, UploadFile, File
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
-from typing import Optional, List, Dict, Any, Generator
+from typing import Optional, List, Dict, Any, Generator, Callable, Tuple
 import json
 import os
 from datetime import datetime
 from urllib.parse import urlparse
-from pathlib import Path
 import tempfile
+from pathlib import Path
 
 # AWS S3 (server-side access, no presigned URLs)
 import boto3
@@ -465,12 +465,34 @@ class ChatResponse(BaseModel):
 # ========================
 
 def _ensure_chat(chat_id: Optional[str]) -> str:
-    """Create a new chat if none provided."""
-    return chat_id or session_manager.create_session()
+    """
+    Ensure a chat exists; if not provided or missing, create a new one.
+    """
+    if chat_id:
+        if session_manager.get_session(chat_id):
+            return chat_id
+    # Create a brand new chat
+    new_id = session_manager.create_session()
+    # Warm/init
+    session_manager.get_session(new_id)
+    return new_id
+
+
+def _get_session_or_init(chat_id: str) -> Dict[str, Any]:
+    """
+    Always return a session dict; initialize if missing.
+    """
+    session = session_manager.get_session(chat_id)
+    if not session:
+        session_manager.update_session(chat_id, {"state": "initial", "messages": []})
+        session = session_manager.get_session(chat_id) or {"state": "initial", "messages": []}
+    return session
 
 
 def _normalize_history_for_api(chat_id: str) -> List[Message]:
-    """Return history in Gradio-like format: [{role, content, timestamp}]"""
+    """
+    Return history in Gradio-like format: [{role, content, timestamp}]
+    """
     session = session_manager.get_session(chat_id) or {}
     raw_messages = session.get("messages") or []
     history: List[Message] = []
@@ -505,7 +527,7 @@ def _assistant_response_payload(
     )
 
 
-def parse_s3_uri(uri: str):
+def parse_s3_uri(uri: str) -> Tuple[str, str]:
     """
     Parse s3://bucket/key to (bucket, key).
     """
@@ -515,9 +537,9 @@ def parse_s3_uri(uri: str):
     return p.netloc, p.path.lstrip("/")
 
 
-def download_to_temp_file(file_ref: Optional[str]) -> (Optional[str], callable):
+def download_to_temp_file(file_ref: Optional[str]) -> Tuple[Optional[str], Callable[[], None]]:
     """
-    If file_ref is an S3 URI, download to a temporary file and return path + cleanup callback.
+    If file_ref is an S3 URI, download to a temporary file and return (path, cleanup).
     If local path or None, return as-is and a no-op cleanup.
     """
     def noop():
@@ -526,7 +548,7 @@ def download_to_temp_file(file_ref: Optional[str]) -> (Optional[str], callable):
     if not file_ref:
         return None, noop
 
-    if file_ref.startswith("s3://"):
+    if isinstance(file_ref, str) and file_ref.startswith("s3://"):
         bucket, key = parse_s3_uri(file_ref)
         suffix = os.path.splitext(key)[1] or ""
         fd, temp_path = tempfile.mkstemp(prefix="masterllm_", suffix=suffix)
@@ -569,12 +591,12 @@ async def smart_chat(request: ChatRequest):
     - assistant_response (LLM-visible reply), plus full history like Gradio UI
     """
     chat_id = _ensure_chat(request.chat_id)
-    session = session_manager.get_session(chat_id) or {}
+    session = _get_session_or_init(chat_id)
 
     # Update file if provided (can be local path or s3://)
     if request.file_path:
         session_manager.update_session(chat_id, {"current_file": request.file_path})
-        session = session_manager.get_session(chat_id)
+        session = _get_session_or_init(chat_id)
 
     # Add user message
     session_manager.add_message(chat_id, "user", request.message)
@@ -665,7 +687,7 @@ async def smart_chat(request: ChatRequest):
                 }
                 return _assistant_response_payload(chat_id, friendly, intent_data, api_data, current_state)
 
-            # Generate pipeline (pass the stored file ref; generator usually doesn't need bytes)
+            # Generate pipeline (no need to download file)
             try:
                 pipeline = generate_pipeline(
                     user_input=request.message,
@@ -807,12 +829,12 @@ def smart_chat_stream(request: ChatRequest):
 
     def gen() -> Generator[bytes, None, None]:
         chat_id = _ensure_chat(request.chat_id)
-        session = session_manager.get_session(chat_id) or {}
+        session = _get_session_or_init(chat_id)
 
         # Update file if provided
         if request.file_path:
             session_manager.update_session(chat_id, {"current_file": request.file_path})
-            session = session_manager.get_session(chat_id)
+            session = _get_session_or_init(chat_id)
 
         # Add user message
         session_manager.add_message(chat_id, "user", request.message)
@@ -824,7 +846,6 @@ def smart_chat_stream(request: ChatRequest):
         def emit(obj: Dict[str, Any]):
             obj.setdefault("chat_id", chat_id)
             obj.setdefault("state", session_manager.get_session(chat_id).get("state", current_state))
-            # Add history only on finals to keep payloads small
             line = json.dumps(obj, ensure_ascii=False).encode("utf-8") + b"\n"
             return line
 
@@ -1045,7 +1066,9 @@ async def execute_pipeline_now(chat_id: str):
     Execute the currently proposed pipeline and return final result (non-streaming).
     Downloads the file to a temp path for execution and deletes it afterward.
     """
-    session = session_manager.get_session(chat_id) or {}
+    session = session_manager.get_session(chat_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Chat not found")
     plan = session.get("proposed_pipeline")
     if not plan:
         raise HTTPException(status_code=400, detail="No pipeline proposed for this chat.")
@@ -1084,7 +1107,9 @@ def execute_pipeline_stream_endpoint(chat_id: str):
     Stream the execution of the currently proposed pipeline (NDJSON).
     Downloads the file to a temp path for execution and deletes it afterward.
     """
-    session = session_manager.get_session(chat_id) or {}
+    session = session_manager.get_session(chat_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Chat not found")
     plan = session.get("proposed_pipeline")
     if not plan:
         raise HTTPException(status_code=400, detail="No pipeline proposed for this chat.")
@@ -1188,14 +1213,22 @@ async def upload_file_to_chat(chat_id: str, file: UploadFile = File(...)):
     key = f"{S3_PREFIX}/{chat_id}/{file.filename}"
     config = TransferConfig(multipart_threshold=8 * 1024 * 1024, max_concurrency=4)
 
-    # Stream from request to S3 (no full in-memory read, no local disk)
-    s3.upload_fileobj(
-        Fileobj=file.file,
-        Bucket=S3_BUCKET,
-        Key=key,
-        ExtraArgs={"ContentType": file.content_type or "application/octet-stream"},
-        Config=config
-    )
+    try:
+        # Stream from request to S3 (no full in-memory read, no local disk)
+        s3.upload_fileobj(
+            Fileobj=file.file,
+            Bucket=S3_BUCKET,
+            Key=key,
+            ExtraArgs={"ContentType": file.content_type or "application/octet-stream"},
+            Config=config
+        )
+    except ClientError as e:
+        code = e.response.get("Error", {}).get("Code", "Unknown")
+        msg = f"S3 upload failed: {code}. Check AWS credentials, permissions (s3:PutObject), region and bucket."
+        raise HTTPException(
+            status_code=403 if code in ("AccessDenied", "InvalidAccessKeyId", "SignatureDoesNotMatch") else 500,
+            detail=msg
+        )
 
     s3_uri = f"s3://{S3_BUCKET}/{key}"
     session_manager.update_session(chat_id, {"current_file": s3_uri, "state": "initial"})
@@ -1214,10 +1247,14 @@ async def upload_file_to_chat(chat_id: str, file: UploadFile = File(...)):
 @router.get("/chats/{chat_id}/file")
 def download_chat_file(chat_id: str):
     s = session_manager.get_session(chat_id)
-    if not s or not s.get("current_file", "").startswith("s3://"):
+    if not s:
+        raise HTTPException(status_code=404, detail="Chat not found")
+
+    file_ref = s.get("current_file")
+    if not file_ref or not isinstance(file_ref, str) or not file_ref.startswith("s3://"):
         raise HTTPException(status_code=404, detail="No S3 file attached to this chat")
 
-    bucket, key = parse_s3_uri(s["current_file"])
+    bucket, key = parse_s3_uri(file_ref)
     try:
         obj = s3.get_object(Bucket=bucket, Key=key)
     except ClientError as e:
@@ -1241,4 +1278,4 @@ def download_chat_file(chat_id: str):
 
 @router.get("/health")
 def health_check():
-    return {"status": "ok", "service": "MasterLLM v2.0", "time": datetime.utcnow().isoformat() + "Z"}
+    return {"status": "ok", "service": "MasterLLM v2.0", "time": datetime.utcnow().isoformat() + "Z"}