""" title: NVIDIA R1 description: 在OpenWebUI中显示NVIDIA R1模型的思维链 - 支持流式<think>标签解析 version: 1.0 license: MIT """ import json import httpx from typing import AsyncGenerator, Callable, Awaitable from pydantic import BaseModel, Field import re class Pipe: class Valves(BaseModel): NVIDIA_API_BASE_URL: str = Field( default="https://integrate.api.nvidia.com/v1", description="NVIDIA API基础地址", ) NVIDIA_API_KEY: str = Field( default="", description="NVIDIA API密钥（以nvapi-开头）" ) R1_MODEL_NAME: str = Field( default="deepseek-ai/deepseek-r1", description="要使用的模型名称" ) def __init__(self): self.valves = self.Valves() self.data_prefix = "data:" self.think_start = re.compile(r"<think>") self.think_end = re.compile(r"</think>") def pipes(self): return [{"id": self.valves.R1_MODEL_NAME, "name": "NVIDIA R1 - DeepSeek"}] async def pipe( self, body: dict, __event_emitter__: Callable[[dict], Awaitable[None]] = None ) -> AsyncGenerator[str, None]: buffer = "" in_think = False # 验证API密钥 if not self.valves.NVIDIA_API_KEY.startswith("nvapi-"): yield json.dumps({"error": "无效的NVIDIA API密钥格式"}, ensure_ascii=False) return headers = { "Authorization": f"Bearer {self.valves.NVIDIA_API_KEY}", "Content-Type": "application/json", } try: # 调整请求体格式 payload = { "model": self.valves.R1_MODEL_NAME, "messages": body["messages"], "temperature": body.get("temperature", 0.7), "stream": True, } async with httpx.AsyncClient(http2=True) as client: async with client.stream( "POST", f"{self.valves.NVIDIA_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 json_str = line[len(self.data_prefix) :].strip() if not json_str: continue try: data = json.loads(json_str) except json.JSONDecodeError: continue if "choices" not in data: continue delta = data["choices"][0].get("delta", {}) content = delta.get("content", "") # 流式内容处理 buffer += content # 处理思维标签状态机 while True: if in_think: end_match = self.think_end.search(buffer) if end_match: pos = end_match.start() yield buffer[:pos] # 发送思考内容 yield "</think>" buffer = buffer[end_match.end() :] in_think = False else: yield buffer # 发送未闭合的思考内容 buffer = "" break else: start_match = self.think_start.search(buffer) if start_match: pos = start_match.start() if pos > 0: yield buffer[:pos] # 发送前置内容 yield "<think>" buffer = buffer[start_match.end() :] in_think = True else: if buffer: yield buffer # 发送普通内容 buffer = "" break except Exception as e: yield self._format_exception(e) def _format_error(self, status_code: int, error: bytes) -> str: try: err_info = json.loads(error.decode()) msg = err_info.get("message", "Unknown error") except: msg = error.decode(errors="ignore") return json.dumps( {"error": f"API Error ({status_code}): {msg[:200]}"}, ensure_ascii=False ) def _format_exception(self, e: Exception) -> str: return json.dumps( {"error": f"Runtime Error: {str(e)[:200]}"}, ensure_ascii=False )