LQConstPlanWriterAgent/core/document_chat/schemas.py

# -*- coding: utf-8 -*-
"""Schemas for the document chat module."""

from typing import Any, Dict, List, Literal, Optional

from pydantic import BaseModel, Field


class SelectedSection(BaseModel):
    index: str = Field(..., description="Section index, for example 2.1")
    title: str = Field(..., description="Section title")
    content: str = Field(default="", description="Current section content from the editor")
    code: str = Field(default="", description="Section code")


class DocumentContext(BaseModel):
    before: str = Field(default="", description="Previous context snippet")
    after: str = Field(default="", description="Following context snippet")
    siblings: List[Dict[str, Any]] = Field(default_factory=list)
    references: List[Dict[str, Any]] = Field(default_factory=list)


class DocumentChatRequest(BaseModel):
    user_id: str
    message: str = Field(..., min_length=1, description="User message")
    selected_section: SelectedSection
    conversation_id: Optional[str] = None
    task_id: Optional[str] = None
    project_info: Dict[str, Any] = Field(default_factory=dict)
    document_context: DocumentContext = Field(default_factory=DocumentContext)
    conversation_history: List[Dict[str, Any]] = Field(default_factory=list)
    response_mode: Literal["json", "sse"] = "json"

    class Config:
        extra = "forbid"


class IntentResult(BaseModel):
    intent: Literal["document_modify", "document_answer", "clarify", "unsupported"]
    confidence: float = Field(default=0.0, ge=0.0, le=1.0)
    skill_name: Optional[str] = None
    operation: str = ""
    target_scope: str = "selected_section"
    normalized_instruction: str = ""
    needs_clarification: bool = False
    clarification_question: str = ""
    reason: str = ""
    warnings: List[str] = Field(default_factory=list)


class DocumentChatSkillInput(BaseModel):
    user_id: str
    user_message: str
    selected_section: SelectedSection
    intent_result: IntentResult
    conversation_id: Optional[str] = None
    task_id: Optional[str] = None
    project_info: Dict[str, Any] = Field(default_factory=dict)
    document_context: DocumentContext = Field(default_factory=DocumentContext)
    conversation_history: List[Dict[str, Any]] = Field(default_factory=list)


class DocumentChatSkillOutput(BaseModel):
    skill_name: str
    response_type: Literal["answer", "proposal", "clarify", "unsupported"]
    answer: Optional[str] = None
    old_content: Optional[str] = None
    proposed_content: Optional[str] = None
    change_summary: List[str] = Field(default_factory=list)
    references: List[Dict[str, Any]] = Field(default_factory=list)
    warnings: List[str] = Field(default_factory=list)


class DiffItem(BaseModel):
    type: Literal["equal", "insert", "delete", "replace", "full_content"]
    old_text: str = ""
    new_text: str = ""


class DiffResult(BaseModel):
    old_content_hash: str
    new_content_hash: str
    diff: List[DiffItem] = Field(default_factory=list)
    diff_granularity: Literal["line", "full_content"] = "line"


class DocumentChatData(BaseModel):
    callback_task_id: str
    response_type: Literal["answer", "proposal", "clarify", "unsupported", "error"]
    intent_result: Optional[Dict[str, Any]] = None
    answer: Optional[str] = None
    proposed_content: Optional[str] = None
    old_content_hash: Optional[str] = None
    new_content_hash: Optional[str] = None
    diff: List[Dict[str, Any]] = Field(default_factory=list)
    diff_granularity: Optional[str] = None
    change_summary: List[str] = Field(default_factory=list)
    references: List[Dict[str, Any]] = Field(default_factory=list)
    warnings: List[str] = Field(default_factory=list)
    selected_section: Dict[str, Any] = Field(default_factory=dict)
    error_message: Optional[str] = None


class DocumentChatResponse(BaseModel):
    code: int
    message: str
    data: Optional[DocumentChatData] = None


def model_to_dict(value: Any) -> Dict[str, Any]:
    """Return a dict for Pydantic v1/v2 models."""
    if value is None:
        return {}
    if isinstance(value, dict):
        return value
    if hasattr(value, "model_dump"):
        return value.model_dump()
    if hasattr(value, "dict"):
        return value.dict()
    return dict(value)