# CoT 后端实现方案

## 改动文件清单

| 文件 | 改动性质 |
|------|---------|
| `app/schemas.py` | `model` 字段加入 `auto` 选项 |
| `app/graph/builder.py` | MODEL_PARAMS 加 `auto`；CoT prompt 注入逻辑 |
| `app/graph/nodes.py` | `call_model` 注入含 CoT 的 system prompt |
| `app/api/chat.py` | `ThinkTagParser` + `_stream_response` 扩展 |
| `app/graph/thinking.py` | **新建**：`ThinkTagParser` 状态机 |

---

## 1. schemas.py

```python
# 原来
model: str = Field(default="flash", pattern="^(flash|pro)$")

# 改为
model: str = Field(default="flash", pattern="^(flash|auto|pro)$")
```

---

## 2. graph/thinking.py（新建）

```python
"""Streaming <think> tag parser for CoT extraction."""
from __future__ import annotations


class ThinkTagParser:
    """Parse streaming tokens and separate <think>...</think> from answer.

    Yields (event_type, content) tuples where event_type is
    "thinking" (inside <think> block) or "token" (final answer).

    Handles token boundary issues: tags may arrive split across tokens.
    """

    _OPEN_TAG = "<think>"
    _CLOSE_TAG = "</think>"

    def __init__(self) -> None:
        self._buffer = ""
        self._in_think = False
        self._think_done = False

    def feed(self, token: str) -> list[tuple[str, str]]:
        """Feed one streaming token. Returns list of (type, content) pairs."""
        self._buffer += token
        events: list[tuple[str, str]] = []

        while self._buffer:
            if not self._in_think and not self._think_done:
                # Waiting for <think>
                idx = self._buffer.find(self._OPEN_TAG)
                if idx == -1:
                    # No opening tag found; check for partial tag at end
                    cut = self._safe_cut(self._buffer, "<")
                    if cut > 0:
                        events.append(("token", self._buffer[:cut]))
                        self._buffer = self._buffer[cut:]
                    elif cut == 0:
                        break  # Entire buffer might be a partial tag
                    else:
                        events.append(("token", self._buffer))
                        self._buffer = ""
                else:
                    # Flush any content before <think> as token
                    if idx > 0:
                        events.append(("token", self._buffer[:idx]))
                    self._buffer = self._buffer[idx + len(self._OPEN_TAG):]
                    self._in_think = True

            elif self._in_think:
                # Inside <think>, looking for </think>
                idx = self._buffer.find(self._CLOSE_TAG)
                if idx == -1:
                    cut = self._safe_cut(self._buffer, "<")
                    if cut > 0:
                        events.append(("thinking", self._buffer[:cut]))
                        self._buffer = self._buffer[cut:]
                    elif cut == 0:
                        break
                    else:
                        events.append(("thinking", self._buffer))
                        self._buffer = ""
                else:
                    if idx > 0:
                        events.append(("thinking", self._buffer[:idx]))
                    self._buffer = self._buffer[idx + len(self._CLOSE_TAG):]
                    self._in_think = False
                    self._think_done = True

            else:
                # After </think>: everything is the final answer
                events.append(("token", self._buffer))
                self._buffer = ""

        return events

    def flush(self) -> list[tuple[str, str]]:
        """Flush remaining buffer at stream end."""
        if not self._buffer:
            return []
        kind = "thinking" if self._in_think else "token"
        result = [(kind, self._buffer)]
        self._buffer = ""
        return result

    @staticmethod
    def _safe_cut(text: str, char: str) -> int:
        """Return index of last occurrence of char, or -1 if not found.

        Returns 0 if char is at position 0 (entire string is potential tag).
        """
        idx = text.rfind(char)
        return idx  # -1 if not found, 0 if at start
```

---

## 3. graph/builder.py

```python
# Model parameter presets
MODEL_PARAMS: dict[str, dict] = {
    "flash": {"max_tokens": 500,  "temperature": 0.2, "thinking": False},
    "auto":  {"max_tokens": 2048, "temperature": 0.3, "thinking": True},
    "pro":   {"max_tokens": 4096, "temperature": 0.3, "thinking": True},
}

SYSTEM_PROMPT_BASE = (
    "You are SOC Assistant, an enterprise AI assistant. "
    "You help users with knowledge base queries, ticket management, "
    "and general questions. Always respond in the same language the user uses. "
    "Be concise, accurate, and helpful."
)

# Auto: concise thinking (key decisions only)
COT_PROMPT_AUTO = (
    "\n\nBefore answering, briefly think through the key decision points "
    "inside <think> tags, then give your final answer outside the tags.\n"
    "Format:\n<think>\n[key reasoning steps]\n</think>\n\n[final answer]"
)

# Pro: full step-by-step reasoning
COT_PROMPT_PRO = (
    "\n\nBefore answering, think through the problem step by step inside "
    "<think> tags. Analyze the question thoroughly, consider multiple "
    "approaches, then provide your final answer outside the tags.\n"
    "Format:\n<think>\n[detailed step-by-step reasoning]\n</think>\n\n[final answer]"
)


def _get_system_prompt(model: str) -> str:
    if model == "auto":
        return SYSTEM_PROMPT_BASE + COT_PROMPT_AUTO
    if model == "pro":
        return SYSTEM_PROMPT_BASE + COT_PROMPT_PRO
    return SYSTEM_PROMPT_BASE


def _get_llm(model: str) -> AzureChatOpenAI:
    params = MODEL_PARAMS.get(model, MODEL_PARAMS["flash"])
    return AzureChatOpenAI(
        azure_endpoint=settings.azure_openai_endpoint,
        api_key=settings.azure_openai_api_key,
        api_version=settings.azure_openai_api_version,
        azure_deployment=settings.azure_openai_deployment,
        max_tokens=params["max_tokens"],
        temperature=params["temperature"],
        streaming=True,
    )
```

---

## 4. graph/nodes.py

```python
from langchain_core.messages import SystemMessage
from app.graph.builder import _get_llm, _get_system_prompt


async def call_model(state: ChatState) -> dict:
    model = state.get("model", "flash")
    llm = _get_llm(model)

    messages = list(state["messages"])
    system_content = _get_system_prompt(model)
    messages.insert(0, SystemMessage(content=system_content))

    response = await llm.ainvoke(messages)
    return {"messages": [response]}
```

---

## 5. api/chat.py（核心改动）

在 `_stream_response` 中集成 `ThinkTagParser`：

```python
from app.graph.thinking import ThinkTagParser

async def _stream_response(request: ChatRequest) -> AsyncIterator[bytes]:
    # ... 现有初始化代码 ...

    thinking_enabled = request.model in ("auto", "pro")
    parser = ThinkTagParser() if thinking_enabled else None

    try:
        async for event in graph.astream_events(input_data, config=config, version="v2"):
            kind = event.get("event", "")

            if kind == "on_chat_model_stream":
                chunk = event.get("data", {}).get("chunk")
                if chunk and hasattr(chunk, "content") and chunk.content:
                    if isinstance(chunk.content, str):
                        raw_token = chunk.content

                        if parser:
                            for evt_type, evt_content in parser.feed(raw_token):
                                if not evt_content:
                                    continue
                                if evt_type == "thinking":
                                    full_thinking.append(evt_content)
                                else:
                                    full_content.append(evt_content)
                                sse = json.dumps(
                                    {"type": evt_type, "content": evt_content},
                                    ensure_ascii=False,
                                )
                                yield f"data: {sse}\n\n".encode("utf-8")
                        else:
                            # Flash mode: direct token passthrough
                            full_content.append(raw_token)
                            sse = json.dumps(
                                {"type": "token", "content": raw_token},
                                ensure_ascii=False,
                            )
                            yield f"data: {sse}\n\n".encode("utf-8")

            elif kind == "on_tool_start":
                tool_name = event.get("name", "unknown")
                sse = json.dumps({"type": "tool_start", "tool": tool_name}, ensure_ascii=False)
                yield f"data: {sse}\n\n".encode("utf-8")

            elif kind == "on_tool_end":
                tool_name = event.get("name", "unknown")
                sse = json.dumps({"type": "tool_end", "tool": tool_name}, ensure_ascii=False)
                yield f"data: {sse}\n\n".encode("utf-8")

    except Exception as exc:
        # ... 现有错误处理 ...
        pass

    finally:
        # Flush parser buffer
        if parser:
            for evt_type, evt_content in parser.flush():
                if evt_content:
                    sse = json.dumps({"type": evt_type, "content": evt_content}, ensure_ascii=False)
                    yield f"data: {sse}\n\n".encode("utf-8")

        ai_content = "".join(full_content)
        if ai_content:
            await _persist_ai_message(request.conversation_id, ai_content)

        done_data = json.dumps({"type": "done"})
        yield f"data: {done_data}\n\n".encode("utf-8")
```

---

## 实施顺序

1. `app/schemas.py` — 加 `auto`
2. `app/graph/thinking.py` — 新建 `ThinkTagParser`
3. `app/graph/builder.py` — MODEL_PARAMS + prompt 函数
4. `app/graph/nodes.py` — 注入 system prompt
5. `app/api/chat.py` — 集成 parser + 新 SSE 事件