refactor: 관세 예측 라우터 리펙토링

Author: hanjuhn

app/service.py

@@ -1,18 +1,34 @@
 from app.dto.response import Response
 from core.graphs.runner import run_customs_agent
+from langchain_core.messages import HumanMessage, AIMessage
+
+# 세션 맵: {session_id: messages}
+session_map = {}
 
 def run_model(question: str, session_id: str = None) -> "Response":
-    """
-        중앙 관리 모델을 통해 각 요청을 적절한 모델로 라우팅
-        아무 모델과도 관계없는 경우 중앙 모델에서 적절한 응답을 생성해야 합니다.
-        ex : "안녕, 너는 뭘 할 수 있어?"
-        예시 코드는 아래와 같습니다.
-    """
-    state = run_customs_agent(question, session_id=session_id)
+    # session_id가 없으면 대화내역 없이 동작
+    if not session_id:
+        messages = [HumanMessage(content=question)]
+    else:
+        # session_map에 없으면 빈 리스트 할당
+        if session_id not in session_map:
+            session_map[session_id] = []
+        messages = session_map[session_id]
+        messages.append(HumanMessage(content=question))
+
+    # agent 실행 (messages를 넘김)
+    state = run_customs_agent(question, session_id=session_id, messages=messages)
+
+    # AI 응답 메시지 추가 (state["final_response"] 기준)
+    ai_reply = state.get("final_response")
+    if session_id and ai_reply:
+        messages.append(AIMessage(content=ai_reply))
+        session_map[session_id] = messages
+
     return Response(
-        reply=state.get("final_response"),
+        reply=ai_reply,
         progress_details=state.get("progress_details"),
         error_reason=state.get("error_reason"),
         success=True,
-        session_id=state.get("session_id")
+        session_id=session_id  # session_id를 응답에 포함
     )

core/graphs/runner.py

@@ -4,15 +4,19 @@
 from core.shared.states.states import CustomsAgentState
 
 
-def run_customs_agent(query: str, session_id: str = None) -> CustomsAgentState:
+def run_customs_agent(query: str, session_id: str = None, messages=None) -> CustomsAgentState:
     """관세청 에이전트를 실행합니다."""
 
     # 그래프 생성
     app = create_customs_graph()
 
+    # messages가 없으면 query만 포함
+    if messages is None:
+        messages = [HumanMessage(content=query)]
+
     # 초기 상태 설정
     initial_state = CustomsAgentState(
-        messages=[HumanMessage(content=query)],
+        messages=messages,
         query=query,
         intent=None,
         final_response="",

core/shared/router/intent_router.py

@@ -107,8 +107,6 @@ def intent_router(state: CustomsAgentState) -> CustomsAgentState:
     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)
@@ -117,49 +115,36 @@ def intent_router(state: CustomsAgentState) -> CustomsAgentState:
     if is_question:
         state["intent"] = "qna"
         _add_classification_message(state, "qna", "질문 형태 감지(우선)")
-        if prev_intent == "tariff_prediction" and state.get("session_id"):
+        if prev_intent == "tariff_prediction":
             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"):
+        if prev_intent == "tariff_prediction" and keyword_intent != "tariff_prediction":
             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
     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
     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"):
+    if prev_intent == "tariff_prediction" and intent != "tariff_prediction":
         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