CodeWeaver/src/agent/state.py

from typing import Any, Dict, List, Optional, Literal, Tuple, Annotated
from operator import add

from pydantic import BaseModel, Field
from langchain_core.messages import BaseMessage
from langgraph.graph import add_messages


_STEPS_RESET_TOKEN = "__RESET_STEPS__"
_MULTI_ANS_RESET_TOKEN = "__RESET_MULTI_ANS__"


def merge_intermediate_steps(old: List[str], new: List[str]) -> List[str]:
    """
    intermediate_steps reducer.

    - 기본 동작: old + new (병렬 노드에서 동시에 step을 추가 가능)
    - 리셋 동작: new의 첫 원소가 _STEPS_RESET_TOKEN 이면 old를 버리고 new[1:]로 교체
      (체크포인팅으로 누적된 step을 '이번 실행(run)' 기준으로 초기화하기 위함)
    """
    if not new:
        return old
    if new[0] == _STEPS_RESET_TOKEN:
        return new[1:]
    return old + new


def merge_multi_answers(old: List[Dict[str, Any]], new: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    multi_answers reducer.

    - 기본 동작: old + new (병렬 worker에서 답변을 동시에 append 가능)
    - 리셋 동작: new의 첫 원소가 {"__token__": _MULTI_ANS_RESET_TOKEN} 이면
      old를 버리고 new[1:]로 교체
      (체크포인팅/스레드 유지로 인해 이전 턴의 multi_answers가 누적되는 문제 방지)
    """
    if not new:
        return old
    head = new[0]
    if isinstance(head, dict) and head.get("__token__") == _MULTI_ANS_RESET_TOKEN:
        return new[1:]
    return old + new


class SearchResult(BaseModel):
    """검색 도메인에서 공통으로 사용하는 단일 검색 결과 모델."""

    source: str = Field(
        ...,
        description="검색 출처 (예: Stack Overflow, 공식 문서, GitHub Issues 등)",
    )
    content: str = Field(
        ...,
        description="검색 결과의 핵심 내용 또는 발췌 텍스트",
    )
    url: Optional[str] = Field(
        default=None,
        description="검색 결과의 원본 출처 URL (존재하는 경우에만 설정)",
    )
    relevance_score: Optional[float] = Field(
        default=None,
        description="검색 쿼리와의 관련도 점수 (0.0–1.0 범위, 클수록 더 관련 있음)",
    )


class AgentState(BaseModel):
    """CodeWeaver LangGraph 에이전트의 전체 상태를 나타내는 Pydantic 모델.
    
    LangGraph 공식 가이드라인:
    - Pydantic BaseModel 사용 (타입 안전성)
    - messages 필드에 add_messages reducer 적용
    - 모든 필드에 기본값 제공
    """

    # Core fields
    user_question: str = Field(default="", description="사용자의 원본 질문")
    messages: Annotated[List[BaseMessage], add_messages] = Field(
        default_factory=list,
        description="대화 메시지 히스토리 (add_messages reducer 사용)"
    )
    
    # Legacy conversation history (유지하되 messages 우선)
    conversation_history: Optional[List[Tuple[str, str]]] = Field(
        default=None,
        description="레거시 대화 내역 (messages 우선 사용)"
    )
    
    # Intent classification
    detected_intent: Optional[Literal["debugging", "learning", "code_review"]] = Field(
        default=None,
        description="분류된 질문 의도"
    )
    
    # Cache-related
    cached_result: Optional[str] = Field(
        default=None,
        description="벡터 DB 캐시에서 조회된 답변"
    )
    
    # Search results (Send API를 위한 reducer 사용)
    search_results: Annotated[List[SearchResult], add] = Field(
        default_factory=list,
        description="병렬 검색으로 수집된 결과 리스트 (Send API로 병렬 업데이트)"
    )
    
    # Intermediate processing
    subtask_results: Dict[str, Any] = Field(
        default_factory=dict,
        description="서브태스크 실행 결과 저장소"
    )
    
    # Final output
    final_answer: Optional[str] = Field(
        default=None,
        description="최종 생성된 답변"
    )
    
    # Debugging/tracing (병렬 노드 + 실행 단위 리셋 지원 reducer 사용)
    intermediate_steps: Annotated[List[str], merge_intermediate_steps] = Field(
        default_factory=list,
        description="실행 단계별 로그 (병렬 노드에서 동시 업데이트 가능)"
    )
    
    # Question analysis & cache eligibility
    question_type: Optional[Literal["clarification", "new_topic", "independent"]] = Field(
        default=None,
        description="질문 유형 분류 결과"
    )
    analysis_reasoning: Optional[str] = Field(
        default=None,
        description="질문 분석 이유"
    )
    should_cache: Optional[bool] = Field(
        default=None,
        description="캐시 저장 여부"
    )
    canonical_question: Optional[str] = Field(
        default=None,
        description="정규화된 질문 (캐시용)"
    )
    
    # Planning & Refinement (Phase 3: Open Deep Research pattern)
    plan: Optional[Dict[str, Any]] = Field(
        default=None,
        description="질문 분해 계획: {'sub_questions': [...], 'reasoning': '...'}"
    )
    needs_refinement: bool = Field(
        default=False,
        description="검색 결과가 부족하여 쿼리 개선 필요 여부"
    )
    refinement_count: int = Field(
        default=0,
        description="검색 쿼리 개선 시도 횟수 (최대 1회)"
    )
    original_question: Optional[str] = Field(
        default=None,
        description="쿼리 개선 전 원본 질문 (최종 답변 생성 시 참조)"
    )
    
    # Phase 4: Dynamic Parallel Search for Multiple Questions
    is_multi_question: bool = Field(
        default=False,
        description="현재 다중 질문 처리 중인지 여부"
    )
    sub_question_index: int = Field(
        default=0,
        description="서브 질문 인덱스 (0부터 시작)"
    )
    sub_question_text: Optional[str] = Field(
        default=None,
        description="현재 처리 중인 서브 질문 텍스트"
    )
    original_multi_question: Optional[str] = Field(
        default=None,
        description="다중 질문의 원본 질문 (통합 답변 생성 시 참조)"
    )
    multi_answers: Annotated[List[Dict[str, Any]], merge_multi_answers] = Field(
        default_factory=list,
        description="다중 질문의 각 답변 리스트 (reducer로 자동 병합)"
    )
    
    class Config:
        arbitrary_types_allowed = True