From 38ab4a703d43dfe5ff8f93b142059f6855e65934 Mon Sep 17 00:00:00 2001 From: rhyru9 Date: Wed, 31 Dec 2025 10:15:40 +0700 Subject: [PATCH] Fix 500 error on invalid profile in get_custom_question Added validation to check if profile exists before accessing it. Now returns 404 instead of crashing with KeyError. --- application/api/main.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/application/api/main.py b/application/api/main.py index 7a26fa8..87c14db 100644 --- a/application/api/main.py +++ b/application/api/main.py @@ -1,6 +1,6 @@ import json import traceback -from fastapi import APIRouter, WebSocket, WebSocketDisconnect, status +from fastapi import APIRouter, WebSocket, WebSocketDisconnect, status, HTTPException from nlq.business.log_store import LogManagement from nlq.business.profile import ProfileManagement @@ -33,7 +33,15 @@ def option(id: str=None): @router.get("/get_custom_question", response_model=CustomQuestion) def get_custom_question(data_profile: str): all_profiles = ProfileManagement.get_all_profiles_with_info() - comments = all_profiles[data_profile]['comments'] + + # FIX: Check if profile exists before accessing + if data_profile not in all_profiles: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Profile '{data_profile}' not found." + ) + + comments = all_profiles[data_profile].get('comments', '') comments_questions = [] if len(comments.split("Examples:")) > 1: comments_questions_txt = comments.split("Examples:")[1] @@ -233,4 +241,4 @@ async def response_websocket(websocket: WebSocket, session_id: str, content, } logger.info(content_obj) final_content = json.dumps(content_obj, default=serialize_timestamp) - await websocket.send_text(final_content) + await websocket.send_text(final_content) \ No newline at end of file