import os from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect, HTTPException, Request from sqlalchemy.orm import Session from typing import List from src.utils.db import get_db from src.menu_design.apps.AI_chat.models import Module from pydantic import BaseModel from src.menu_design.tasks.tasks import sample_task_menu_design from src.menu_design.core.menu_design import ChatRequest, MenuDesignGenerator from src.utils.pusher_config import get_pusher_client import json from src.utils.db import SessionLocal from src.menu_design.apps.AI_chat.controller import generate_image from src.menu_design.apps.AI_chat.models import GeneratedImage from src.menu_design.apps.AI_chat.schemas import ImageRequest,ImageResponse,ModuleOut from src.apps.auth.controller import get_current_user from src.apps.stores.models import Branch router = APIRouter() @router.websocket("/chat") async def websocket_chat(websocket: WebSocket): """ WebSocket endpoint for live chat between frontend and backend. """ await websocket.accept() try: print("WebSocket connection established.") while True: try: # Receive message from frontend data = await websocket.receive_text() message = json.loads(data) session_id = message.get("session_id") user_message = message.get("message") # Check for valid input if not session_id or not user_message: await websocket.send_text(json.dumps({"error": "Invalid input"})) continue print(f"Received from client (Session {session_id}): {user_message}") # Process the message using the AI chat agent agent = MenuDesignGenerator(session_id=session_id) req = ChatRequest(session_id=session_id, message=user_message) response = agent.chat(req) # Send AI response back to the client #response_data = {"response": response} print(f"Sending to client (Session {session_id}): {response}") # await websocket.send_text(response.model_dump_json()) await websocket.send_text(json.dumps(response)) except WebSocketDisconnect: print("WebSocket connection closed by the client.") break except Exception as e: error_msg = {"error": str(e)} await websocket.send_text(json.dumps(error_msg)) finally: await websocket.close() print("WebSocket connection terminated.") @router.post("/generate-image", response_model=ImageResponse, dependencies=[Depends(get_current_user)]) async def create_image( request: Request, image_request: ImageRequest, db: Session = Depends(get_db), user=Depends(get_current_user) # Authenticate user ): try: # Query to fetch branch along with user_id and store_id branch = db.query(Branch).filter( Branch.branch_id == image_request.branch_id, Branch.user_id == user.user_id # Ensure the branch belongs to the logged-in user ).first() if not branch: raise HTTPException( status_code=403, detail="You do not have permission to generate images for this branch." ) # Extract user_id and store_id from the branch object user_id = branch.user_id store_id = branch.store_id # Default to 1 image if not provided image_number = image_request.image_number if image_request.image_number else 1 # Generate multiple images image_urls = generate_image( desc=image_request.desc, layout=image_request.layout, style=image_request.style, image_number=image_number ) # Save each image URL to the database along with user_id and store_id for image_url in image_urls: new_image = GeneratedImage( branch_id=image_request.branch_id, user_id=user_id, # Save the fetched user_id store_id=store_id, # Save the fetched store_id image_url=image_url ) db.add(new_image) db.commit() # Return the list of image URLs in the correct format return ImageResponse(image_urls=image_urls) except Exception as e: raise HTTPException(status_code=500, detail=f"Error generating images: {str(e)}") @router.get("/generated-images", response_model=ImageResponse, dependencies=[Depends(get_current_user)]) async def get_generated_images( db: Session = Depends(get_db), user=Depends(get_current_user) # Authenticated user ): try: # Fetch all images where user_id matches the logged-in user images = db.query(GeneratedImage).filter( GeneratedImage.user_id == user.user_id ).all() # Extract just the image URLs image_urls = [image.image_url for image in images] return ImageResponse(image_urls=image_urls) except Exception as e: raise HTTPException(status_code=500, detail=f"Error fetching images: {str(e)}") @router.get("/modules", response_model=List[ModuleOut]) def get_modules(db: Session = Depends(get_db)): return db.query(Module).all() @router.get("/celery-health") def check_celery_health(): """Check if Celery and Redis are available""" try: # Try to connect to Redis via Celery from src.utils.celery_worker import celery_app i = celery_app.control.inspect() stats = i.stats() return { "status": "healthy", "message": "Celery and Redis are available", "stats": stats } except Exception as e: return { "status": "unhealthy", "error": "Celery or Redis is not available", "details": str(e) } @router.get("/run-task") def run_sample_task(name: str = "CeleryUser"): try: task = sample_task_menu_design.delay(name) return {"task_id": task.id, "status": "Task submitted"} except Exception as e: # Handle Celery connection errors gracefully error_msg = str(e) if "Connection refused" in error_msg or "No hostname was supplied" in error_msg: return { "error": "Celery worker is not available. Please ensure Redis is running and Celery worker is started.", "details": error_msg, "status": "failed" } else: return { "error": "Task submission failed", "details": error_msg, "status": "failed" } from pydantic import BaseModel class TestChatRequest(BaseModel): session_id: str message: str @router.post("/test-chat") def test_chat_endpoint(request: TestChatRequest): """ Simple POST endpoint to test menu_agent chat functionality without WebSocket. """ try: print(f"Testing chat session {request.session_id} with message: {request.message}") # Run menu_agent and get response menu_agent = MenuDesignGenerator(session_id=request.session_id) chat_request = ChatRequest(session_id=request.session_id, message=request.message) response = menu_agent.run(chat_request) # Return the full response return response except Exception as e: print(f"Error during chat testing: {str(e)}") raise HTTPException(status_code=500, detail=str(e))