""" title: RAGLFOW-COT author: dalisanxianbao description: 适配ragflow思维链 version: 1.1 licence: MIT """ import json import httpx from typing import AsyncGenerator, Callable, Awaitable, Optional from pydantic import BaseModel, Field class Pipe: class Valves(BaseModel): DEEPSEEK_API_BASE_URL: str = Field( default="https://api.siliconflow.com/v1", description="DeepSeek API的基础请求地址", ) DEEPSEEK_API_KEY: str = Field( default="", description="用于身份验证的DeepSeek API密钥" ) DEEPSEEK_R1_MODEL: str = Field( default="deepseek-ai/DeepSeek-R1", description="DeepSeek R1模型的ID", ) def __init__(self): self.valves = self.Valves() self.data_prefix = "data:" self.thinking = -1 # 新增状态：-1=可能有遗留内容，0=思考中，1=已回答 self.buffer = "" self.emitter = None def pipes(self): return [{"id": self.valves.DEEPSEEK_R1_MODEL, "name": "ragflow-cot"}] async def pipe( self, body: dict, __event_emitter__: Callable[[dict], Awaitable[None]] = None ) -> AsyncGenerator[str, None]: self.thinking = -1 self.buffer = "" self.emitter = __event_emitter__ if not self.valves.DEEPSEEK_API_KEY: yield json.dumps({"error": "未配置API密钥"}, ensure_ascii=False) return headers = { "Authorization": f"Bearer {self.valves.DEEPSEEK_API_KEY}", "Content-Type": "application/json", } try: model_id = body["model"].split(".", 1)[-1] payload = {**body, "model": model_id} # 消息处理逻辑保持不变 messages = payload["messages"] i = 0 while i < len(messages) - 1: if messages[i]["role"] == messages[i + 1]["role"]: alternate_role = ( "assistant" if messages[i]["role"] == "user" else "user" ) messages.insert( i + 1, {"role": alternate_role, "content": "[Thinking]"} ) i += 1 async with httpx.AsyncClient(http2=True) as client: async with client.stream( "POST", f"{self.valves.DEEPSEEK_API_BASE_URL}/chat/completions", json=payload, headers=headers, timeout=30, ) as response: if response.status_code != 200: error = await response.aread() yield self._format_error(response.status_code, error) return async for line in response.aiter_lines(): if not line.startswith(self.data_prefix): continue data = json.loads(line[len(self.data_prefix) :]) choice = data.get("choices", [{}])[0] if choice.get("finish_reason"): return content = choice.get("delta", {}).get("content", "") if content: # 处理内容并生成事件 async for chunk in self._process_content(content): yield chunk except Exception as e: yield self._format_exception(e) async def _process_content(self, content: str) -> AsyncGenerator[str, None]: self.buffer += content while True: # 状态0：等待思考开始 if self.thinking == -1: think_start = self.buffer.find("<think>") if think_start != -1: # 处理起始标签前可能存在的非预期内容 if think_start > 0: yield self.buffer[:think_start] self.buffer = self.buffer[think_start:] self.thinking = 0 yield "<think>" self.buffer = self.buffer[7:] # 移除<think> continue else: # 非思考内容直接输出 if self.buffer: yield self.buffer self.buffer = "" break # 状态1：思考进行中 if self.thinking == 0: end_pos = self.buffer.find("</think>") # 情况1：没有闭合标签（思考已自然结束） if end_pos == -1: # 输出全部缓冲内容作为最终思考 yield "</think>\n\n" yield self.buffer self.thinking = 1 self.buffer = "" break # 情况2：存在闭合标签 else: remaining = self.buffer[end_pos + 8 :] # 跳过</think> # 闭合标签后有内容，闭合标签后为最终答案内容，闭合标签前为思考内容 if remaining.strip(): # 输出闭合前的思考内容 yield self.buffer[:end_pos] # 闭合思考并输出后续内容 yield "</think>\n\n" yield remaining self.thinking = 1 self.buffer = "" break # 闭合标签后无内容，直接输出闭合前的思考内容 else: yield self.buffer[:end_pos] self.buffer = "" break # 状态2：已回答阶段 elif self.thinking == 1: if self.buffer: yield self.buffer self.buffer = "" break def _format_error(self, status_code: int, error: bytes) -> str: try: err_msg = json.loads(error).get("message", error.decode(errors="ignore"))[ :200 ] except: err_msg = error.decode(errors="ignore")[:200] return json.dumps( {"error": f"HTTP {status_code}: {err_msg}"}, ensure_ascii=False ) def _format_exception(self, e: Exception) -> str: err_type = type(e).__name__ return json.dumps({"error": f"{err_type}: {str(e)}"}, ensure_ascii=False)