fix: 세션 관련 이슈 해결

Author: hanjuhn

app/dto/request.py

@@ -2,3 +2,4 @@
 
 class Request(BaseModel):
     message: str
+    session_id: str = None

app/dto/response.py

@@ -10,6 +10,7 @@ class Response(BaseModel):
     success: Optional[bool] = None
     progress_details: Optional[List[ProgressDetail]] = None
     error_reason: Optional[str] = None
+    session_id: Optional[str] = None
 
     @staticmethod
     def cargo_progres_result_to_response(result: CargoProgressResult) -> "Response":

app/routes.py

@@ -1,7 +1,8 @@
-from flask import Blueprint, request, jsonify
+from flask import Blueprint, request, jsonify, make_response
 from app.dto.request import Request
 from flasgger import swag_from
 from .service import run_model
+import uuid
 
 api_blueprint = Blueprint("api", __name__)
 
@@ -19,6 +20,11 @@
                     'question': {
                         'type': 'string',
                         'example': '이 물건의 세금이 얼마나 나올까?'
+                    },
+                    'session_id': {
+                        'type': 'string',
+                        'example': 'uuid-1234',
+                        'required': False
                     }
                 },
                 'required': ['question']
@@ -34,6 +40,10 @@
                     'answer': {
                         'type': 'string',
                         'example': '이 물건은 8%의 부가세가 부과됩니다.'
+                    },
+                    'session_id': {
+                        'type': 'string',
+                        'example': 'uuid-1234'
                     }
                 }
             }
@@ -42,5 +52,15 @@
 })
 def predict():
     request_data = Request(**request.get_json())
-    answer = run_model(question=request_data.message)
-    return jsonify(answer.model_dump())
+    session_id = request_data.session_id or request.cookies.get("session_id")
+    answer = run_model(question=request_data.message, session_id=session_id)
+    # 응답 dict에서 session_id 제거
+    answer_dict = answer.model_dump()
+    answer_dict.pop("session_id", None)
+    response = make_response(jsonify(answer_dict))
+    # 쿠키는 그대로 관리
+    if answer.session_id:
+        response.set_cookie("session_id", answer.session_id)
+    else:
+        response.delete_cookie("session_id")
+    return response

app/service.py

@@ -1,17 +1,18 @@
 from app.dto.response import Response
 from core.graphs.runner import run_customs_agent
 
-def run_model(question: str) -> "Response":
+def run_model(question: str, session_id: str = None) -> "Response":
     """
         중앙 관리 모델을 통해 각 요청을 적절한 모델로 라우팅
         아무 모델과도 관계없는 경우 중앙 모델에서 적절한 응답을 생성해야 합니다.
         ex : "안녕, 너는 뭘 할 수 있어?"
         예시 코드는 아래와 같습니다.
     """
-    state = run_customs_agent(question)
+    state = run_customs_agent(question, session_id=session_id)
     return Response(
         reply=state.get("final_response"),
         progress_details=state.get("progress_details"),
         error_reason=state.get("error_reason"),
-        success=True
+        success=True,
+        session_id=state.get("session_id")
     )

core/graphs/runner.py

@@ -4,7 +4,7 @@
 from core.shared.states.states import CustomsAgentState
 
 
-def run_customs_agent(query: str) -> CustomsAgentState:
+def run_customs_agent(query: str, session_id: str = None) -> CustomsAgentState:
     """관세청 에이전트를 실행합니다."""
 
     # 그래프 생성
@@ -18,7 +18,8 @@ def run_customs_agent(query: str) -> CustomsAgentState:
         final_response="",
         intermediate_results={},
         error_reason=None,
-        progress_details=None
+        progress_details=None,
+        session_id=session_id
     )
 
     # 그래프 실행

core/shared/router/intent_router.py

@@ -1,6 +1,7 @@
 from langchain_core.messages import SystemMessage, AIMessage
 import re
 from typing import List, Optional
+import uuid
 
 from core.shared.states.states import CustomsAgentState
 from core.shared.utils.llm import get_llm
@@ -89,57 +90,76 @@ def _classify_with_llm(query: str) -> str:
 
 
 def _add_classification_message(state: CustomsAgentState, intent: str, reason: str) -> None:
-    """의도 분류 완료 메시지를 추가합니다."""
-    state["messages"].append(
-        AIMessage(content=f"의도 분류 완료: {intent} ({reason})")
-    )
+    # 의도 분류 메시지를 messages에 남기지 않음
+    pass
 
 
 def intent_router(state: CustomsAgentState) -> CustomsAgentState:
-    """사용자 쿼리의 의도를 분류합니다."""
-    
+    prev_intent = state.get("intent")
     current_query = state["query"].strip()
     if not current_query:
         state["intent"] = DEFAULT_INTENT
         _add_classification_message(state, DEFAULT_INTENT, "빈 쿼리")
         return state
-    
     # 세션 연속성 확인
     is_in_tariff_session = _is_in_tariff_session(state)
-    
     # 관세 예측 세션 중이면 무조건 tariff_prediction으로 분류
     if is_in_tariff_session:
         state["intent"] = "tariff_prediction"
         _add_classification_message(state, "tariff_prediction", "관세 예측 세션 연속성 유지")
+        if not state.get("session_id") or state.get("session_id") in [None, '']:
+            state["session_id"] = str(uuid.uuid4())
         return state
-    
     # 패턴 기반 분류
     is_number_selection = _is_number_selection(current_query)
     is_question = _is_question(current_query)
-    
     # 질문 형태이면서 관세 예측 세션이 아닌 경우 QnA로 분류
     if is_question:
         state["intent"] = "qna"
         _add_classification_message(state, "qna", "질문 형태 감지(우선)")
+        if prev_intent == "tariff_prediction" and state.get("session_id"):
+            try:
+                from core.tariff_prediction.agent.tariff_prediction_agent import workflow_manager
+                workflow_manager.cleanup_session(state["session_id"])
+            except Exception:
+                pass
+        state["session_id"] = None
         return state
-    
     # 키워드 기반 분류
     keyword_intent = _classify_by_keywords(current_query)
     if keyword_intent:
         state["intent"] = keyword_intent
         _add_classification_message(state, keyword_intent, "키워드 기반 분류")
+        if prev_intent == "tariff_prediction" and keyword_intent != "tariff_prediction" and state.get("session_id"):
+            try:
+                from core.tariff_prediction.agent.tariff_prediction_agent import workflow_manager
+                workflow_manager.cleanup_session(state["session_id"])
+            except Exception:
+                pass
+        if keyword_intent == "tariff_prediction":
+            if not state.get("session_id") or state.get("session_id") in [None, '']:
+                state["session_id"] = str(uuid.uuid4())
+        else:
+            state["session_id"] = None
         return state
-    
-    # 숫자 선택이지만 관세 예측 세션이 아닌 경우에도 tariff_prediction으로 분류
-    # (HS 코드 직접 입력 등의 경우)
     if is_number_selection:
         state["intent"] = "tariff_prediction"
         _add_classification_message(state, "tariff_prediction", "숫자 선택 감지")
+        if not state.get("session_id") or state.get("session_id") in [None, '']:
+            state["session_id"] = str(uuid.uuid4())
         return state
-    
-    # LLM 기반 의도 분류를 수행
     intent = _classify_with_llm(current_query)
     state["intent"] = intent
     _add_classification_message(state, intent, "LLM 분류")
-    
+    if prev_intent == "tariff_prediction" and intent != "tariff_prediction" and state.get("session_id"):
+        try:
+            from core.tariff_prediction.agent.tariff_prediction_agent import workflow_manager
+            workflow_manager.cleanup_session(state["session_id"])
+        except Exception:
+            pass
+    if intent == "tariff_prediction":
+        if not state.get("session_id") or state.get("session_id") in [None, '']:
+            state["session_id"] = str(uuid.uuid4())
+    else:
+        state["session_id"] = None
     return state

core/shared/states/states.py

@@ -15,4 +15,5 @@ class CustomsAgentState(TypedDict):
     final_response: str
     intermediate_results: Dict[str, Any]
     progress_details: Optional[List[ProgressDetail]]
-    error_reason: Optional[str]
+    error_reason: Optional[str]
+    session_id: Optional[str]  # 세션 ID 추가

core/tariff_prediction/agent/tariff_prediction_agent.py

@@ -23,19 +23,14 @@
 class WorkflowManager:
     def __init__(self):
         self.sessions = {}
-    
     def get_session(self, session_id: str) -> 'TariffPredictionWorkflow':
-        """세션을 가져오거나 새로 생성합니다."""
         if session_id not in self.sessions:
             self.sessions[session_id] = TariffPredictionWorkflow()
         return self.sessions[session_id]
-    
     def cleanup_session(self, session_id: str):
-        """세션을 정리합니다."""
         if session_id in self.sessions:
             del self.sessions[session_id]
 
-# 전역 매니저 인스턴스
 workflow_manager = WorkflowManager()
 
 class TariffPredictionWorkflow:
@@ -524,17 +519,16 @@ def parse_tariff_result(self, tariff_result: str) -> Dict[str, Any]:
 
 
 def tariff_prediction_agent(state: CustomsAgentState) -> CustomsAgentState:
-    session_id = DEFAULT_SESSION_ID 
-    
-    workflow = workflow_manager.get_session(session_id)
-    
+    session_id = state.get("session_id")
+    if session_id:
+        workflow = workflow_manager.get_session(session_id)
+    else:
+        workflow = None
     messages = state.get("messages", [])
     context = ""
     previous_llm_responses = []
-    
     if messages:
         recent_messages = messages[-10:]
-        
         user_messages = []
         for msg in recent_messages:
             if hasattr(msg, 'type'):
@@ -545,19 +539,19 @@ def tariff_prediction_agent(state: CustomsAgentState) -> CustomsAgentState:
                         content = msg.content
                         if any(keyword in content for keyword in ['HS6 코드 후보', 'HS10 코드 후보', '번호를 선택']):
                             previous_llm_responses.append(content)
-        
         if user_messages:
-            context = " ".join(user_messages[-5:]) 
-    
+            context = " ".join(user_messages[-5:])
     enhanced_context = context or ""
     if previous_llm_responses:
         enhanced_context += f"\n\n{RESPONSE_MESSAGES['previous_llm_response']}\n" + "\n".join(previous_llm_responses[-2:])
-    
-    if enhanced_context:
-        enhanced_query = f"{enhanced_context}\n\n{state['query']}"
-        response = workflow.process_user_input(enhanced_query)
+    if workflow:
+        if enhanced_context:
+            enhanced_query = f"{enhanced_context}\n\n{state['query']}"
+            response = workflow.process_user_input(enhanced_query)
+        else:
+            response = workflow.process_user_input(state["query"])
     else:
-        response = workflow.process_user_input(state["query"])
-    
+        response = ERROR_MESSAGES.get('session_error', '세션이 유효하지 않습니다.')
     state["final_response"] = response
+    state["session_id"] = session_id
     return state