import os import re import gc import sys import time import json import queue import random import asyncio import threading import requests import collections import torch import numpy as np from typing import List, Optional, Dict, Any, Literal, Union from pydantic import BaseModel, Field, model_validator from pydantic_settings import BaseSettings from fastapi import FastAPI, HTTPException, Request from fastapi.responses import StreamingResponse from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from fastapi.middleware.gzip import GZipMiddleware from huggingface_hub import hf_hub_download from snowflake import SnowflakeGenerator if os.environ.get("MODELSCOPE_ENVIRONMENT") == "studio": from modelscope import patch_hub patch_hub() os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:64" os.environ["RWKV_V7_ON"] = "1" os.environ["RWKV_JIT_ON"] = "1" os.environ["RWKV_CUDA_ON"] = "1" GPU_LOCK = asyncio.Lock() class ChatMessage(BaseModel): role: str = Field() content: str = Field() name: Optional[str] = Field(None) tool_call_id: Optional[str] = Field(None) class Logprob(BaseModel): token: str logprob: float top_logprobs: Optional[List[Dict[str, Any]]] = None class LogprobsContent(BaseModel): content: Optional[List[Logprob]] = None refusal: Optional[List[Logprob]] = None class ChatCompletionMessage(BaseModel): role: Optional[str] = Field(None) content: Optional[str] = Field(None) reasoning_content: Optional[str] = Field(None) tool_calls: Optional[List[Dict[str, Any]]] = Field(None) class PromptTokensDetails(BaseModel): cached_tokens: int class Usage(BaseModel): prompt_tokens: int completion_tokens: int total_tokens: int prompt_tokens_details: Optional[PromptTokensDetails] = None class ChatCompletionChoice(BaseModel): index: int message: Optional[ChatCompletionMessage] = None delta: Optional[ChatCompletionMessage] = None logprobs: Optional[LogprobsContent] = None finish_reason: Optional[str] = Field(...) class ChatCompletionChunk(BaseModel): id: str = Field(...) object: Literal["chat.completion.chunk"] = "chat.completion.chunk" created: int = Field(...) model: str choices: List[ChatCompletionChoice] usage: Optional[Usage] = None class ToolFunction(BaseModel): name: str description: str parameters: Dict[str, Any] class Tool(BaseModel): type: Literal["function"] = "function" function: ToolFunction def remove_nested_think_tags_stack(text): stack = [] result = "" i = 0 while i < len(text): if text[i : i + 7] == "": stack.append("") i += 7 elif text[i : i + 8] == "": if stack and stack[-1] == "": stack.pop() i += 8 else: result += text[i : i + 8] i += 8 elif not stack: result += text[i] i += 1 else: i += 1 return result def cleanMessages(messages: List[ChatMessage], removeThinkingContent: bool = False): promptStrList = [] # Safety check in case messages is None if not messages: return "" for message in messages: content = message.content.strip() content = re.sub(r"\n+", "\n", content) role_str = message.role.strip().lower().capitalize() if role_str == 'Assistant' and removeThinkingContent: content = remove_nested_think_tags_stack(content) if message.role == "tool": promptStrList.append(f"Tool Output ({message.name}): {content}") elif message.role == "system": promptStrList.append(f"System: {content}") elif message.role == "user": promptStrList.append(f"User: {content}") elif message.role == "assistant": promptStrList.append(f"Assistant: {content}") else: promptStrList.append(f"{role_str}: {content}") return "\n\n".join(promptStrList) class SamplerConfig(BaseModel): max_tokens: int = 4096 temperature: float = 1.0 top_p: float = 0.3 presence_penalty: float = 0.5 count_penalty: float = 0.5 penalty_decay: float = 0.996 stop: List[str] = ["\n\n"] stop_tokens: List[int] = [0] class ModelConfig(BaseModel): SERVICE_NAME: str DOWNLOAD_MODEL_FILE_NAME: str DOWNLOAD_MODEL_REPO_ID: str DOWNLOAD_MODEL_DIR: str = "models" MODEL_FILE_PATH: Optional[str] = None DEFAULT_CHAT: bool = False DEFAULT_REASONING: bool = False REASONING: bool = False VOCAB: str = "rwkv_vocab_v20230424" CTX_LEN: int = 4096 DEFAULT_SAMPLER: SamplerConfig = Field(default_factory=SamplerConfig) class Config(BaseSettings): HOST: str = "0.0.0.0" PORT: int = 7860 STRATEGY: str = "cuda fp16" RWKV_CUDA_ON: bool = True CHUNK_LEN: int = 256 MODELS: List[ModelConfig] = [ ModelConfig( SERVICE_NAME="rwkv7-g1a4-2.9b-20251118-ctx8192", DOWNLOAD_MODEL_FILE_NAME="rwkv7-g1a4-2.9b-20251118-ctx8192.pth", DOWNLOAD_MODEL_REPO_ID="BlinkDL/rwkv7-g1", REASONING=True, CTX_LEN=8192 ), ModelConfig( SERVICE_NAME="rwkv7-g1a3-1.5b-20251015-ctx8192", DOWNLOAD_MODEL_FILE_NAME="rwkv7-g1a3-1.5b-20251015-ctx8192.pth", DOWNLOAD_MODEL_REPO_ID="BlinkDL/rwkv7-g1", REASONING=True, CTX_LEN=8192 ), ModelConfig( SERVICE_NAME="rwkv7-g1a-0.4b-20250905-ctx4096", DOWNLOAD_MODEL_FILE_NAME="rwkv7-g1a-0.4b-20250905-ctx4096.pth", DOWNLOAD_MODEL_REPO_ID="BlinkDL/rwkv7-g1", REASONING=True, CTX_LEN=4096 ), ModelConfig( SERVICE_NAME="rwkv7-g1a-0.1b-20250728-ctx4096", DOWNLOAD_MODEL_FILE_NAME="rwkv7-g1a-0.1b-20250728-ctx4096.pth", DOWNLOAD_MODEL_REPO_ID="BlinkDL/rwkv7-g1", REASONING=True, DEFAULT_CHAT=True, DEFAULT_REASONING=True, CTX_LEN=4096 ), ] CONFIG = Config() try: from duckduckgo_search import DDGS HAS_DDG = True except ImportError: HAS_DDG = False try: from faker import Faker fake = Faker() HAS_FAKER = True except ImportError: HAS_FAKER = False CompletionIdGenerator = SnowflakeGenerator(42, timestamp=1741101491595) if "cuda" in CONFIG.STRATEGY.lower() and not torch.cuda.is_available(): CONFIG.STRATEGY = "cpu fp16" CONFIG.RWKV_CUDA_ON = False if CONFIG.RWKV_CUDA_ON and "cuda" in CONFIG.STRATEGY.lower(): from pynvml import * nvmlInit() os.environ["RWKV_CUDA_ON"] = "1" torch.backends.cudnn.benchmark = True torch.backends.cudnn.allow_tf32 = True torch.backends.cuda.matmul.allow_tf32 = True else: os.environ["RWKV_CUDA_ON"] = "0" from rwkv.model import RWKV from rwkv.utils import PIPELINE, PIPELINE_ARGS class ModelStorage: MODEL_CONFIG: Optional[ModelConfig] = None model: Optional[RWKV] = None pipeline: Optional[PIPELINE] = None MODEL_STORAGE: Dict[str, ModelStorage] = {} DEFALUT_MODEL_NAME = None DEFAULT_REASONING_MODEL_NAME = None for model_config in CONFIG.MODELS: if model_config.MODEL_FILE_PATH is None: model_config.MODEL_FILE_PATH = hf_hub_download( repo_id=model_config.DOWNLOAD_MODEL_REPO_ID, filename=model_config.DOWNLOAD_MODEL_FILE_NAME, local_dir=model_config.DOWNLOAD_MODEL_DIR, ) if model_config.DEFAULT_CHAT: DEFALUT_MODEL_NAME = model_config.SERVICE_NAME if model_config.DEFAULT_REASONING: DEFAULT_REASONING_MODEL_NAME = model_config.SERVICE_NAME MODEL_STORAGE[model_config.SERVICE_NAME] = ModelStorage() MODEL_STORAGE[model_config.SERVICE_NAME].MODEL_CONFIG = model_config MODEL_STORAGE[model_config.SERVICE_NAME].model = RWKV( model=model_config.MODEL_FILE_PATH.replace(".pth", ""), strategy=CONFIG.STRATEGY, ) MODEL_STORAGE[model_config.SERVICE_NAME].pipeline = PIPELINE( MODEL_STORAGE[model_config.SERVICE_NAME].model, model_config.VOCAB ) if "cuda" in CONFIG.STRATEGY: torch.cuda.empty_cache() gc.collect() class ChatCompletionRequest(BaseModel): model: str = Field(default="rwkv-latest") messages: Optional[List[ChatMessage]] = Field(default=None) prompt: Optional[str] = Field(default=None) max_tokens: Optional[int] = Field(default=None) temperature: Optional[float] = Field(default=None) top_p: Optional[float] = Field(default=None) presence_penalty: Optional[float] = Field(default=None) count_penalty: Optional[float] = Field(default=None) penalty_decay: Optional[float] = Field(default=None) stream: Optional[bool] = Field(default=False) stop: Optional[List[str]] = Field(["\n\n"]) stop_tokens: Optional[List[int]] = Field([0]) tools: Optional[List[Tool]] = Field(default=None) tool_choice: Optional[Union[str, Dict]] = Field(default="auto") @model_validator(mode="before") @classmethod def validate_mutual_exclusivity(cls, data: Any) -> Any: if not isinstance(data, dict): return data if "messages" in data and "prompt" in data and data["messages"] and data["prompt"]: raise ValueError("messages and prompt cannot coexist.") return data class ToolEngine: TOOL_SYSTEM_PROMPT = """ CAPABILITY: You have access to real-time tools. INSTRUCTION: To use a tool, output exactly: tool_name("argument") Do not describe the tool, just call it. After the System provides the result, synthesize the answer. AVAILABLE TOOLS: 1. google_search(query): Searches Google and DuckDuckGo for real-time information. 2. visit_page(url): Accesses a specific link, reads the text, and finds sub-links. """.strip() @staticmethod def google_search_request(query: str) -> str: try: headers = {"User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"} resp = requests.get("https://www.google.com/search", params={"q": query, "gl": "us", "hl": "en"}, headers=headers, timeout=6) if resp.status_code != 200: raise Exception("Google blocked request") clean_text = re.sub(r'.*?', '', resp.text, flags=re.DOTALL) clean_text = re.sub(r'.*?', '', clean_text, flags=re.DOTALL) headings = re.findall(r'(.*?)', clean_text) links = re.findall(r'', '', headings[i])} - Link: {links[i]}\n" if not headings: return ToolEngine.duckduckgo_fallback(query) return output except: return ToolEngine.duckduckgo_fallback(query) @staticmethod def duckduckgo_fallback(query: str) -> str: try: if HAS_DDG: res = DDGS().text(query, max_results=5) return "\n".join([f"- {r['title']}: {r['body']} ({r['href']})" for r in res]) resp = requests.get("https://html.duckduckgo.com/html/", params={"q": query}, headers={"User-Agent": "Mozilla/5.0"}, timeout=5) titles = re.findall(r']*>(.*?)', resp.text) snippets = re.findall(r']*>(.*?)', resp.text) limit = min(len(titles), len(snippets), 4) out = "DuckDuckGo HTML Results:\n" for i in range(limit): t = re.sub(r'<.*?>', '', titles[i]).strip() s = re.sub(r'<.*?>', '', snippets[i]).strip() out += f"{i+1}. {t}: {s}\n" return out except Exception as e: return f"Search failed: {str(e)}" @staticmethod def visit_page(url: str) -> str: try: headers = {"User-Agent": "Mozilla/5.0 (compatible; RWKV-Bot/1.0)"} resp = requests.get(url, headers=headers, timeout=8) resp.encoding = resp.apparent_encoding text = re.sub(r'.*?', '', resp.text, flags=re.DOTALL) text = re.sub(r'.*?', '', text, flags=re.DOTALL) text = re.sub(r'.*?', '', text, flags=re.DOTALL) text = re.sub(r'', '', text, flags=re.DOTALL) text = re.sub(r'<[^>]+>', ' ', text) text = re.sub(r'\s+', ' ', text).strip() links = re.findall(r'href=["\'](http[s]?://[^"\']+)["\']', resp.text) unique_links = list(set(links))[:5] content_preview = text[:3000] + ("..." if len(text) > 3000 else "") return f"PAGE CONTENT ({url}):\n{content_preview}\n\nFOUND SUB-LINKS:\n" + "\n".join(unique_links) except Exception as e: return f"Error visiting page: {str(e)}" @staticmethod def execute(call_str: str) -> str: try: match = re.match(r'(\w+)\(["\'](.*?)["\']\)', call_str) if not match: return "Invalid tool call syntax." func, arg = match.groups() if func == "google_search": return ToolEngine.google_search_request(arg) elif func == "visit_page": return ToolEngine.visit_page(arg) else: return f"Unknown tool: {func}" except Exception as e: return f"Tool execution error: {e}" app = FastAPI(title="RWKV Ultimate Agent Server") app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=5) @app.middleware("http") async def privacy_middleware(request: Request, call_next): if HAS_FAKER: request.scope["client"] = (fake.ipv4(), request.client.port if request.client else 80) return await call_next(request) def prune_context(messages: List[ChatMessage], model_name: str, max_gen_tokens: int): storage = MODEL_STORAGE[model_name] limit = storage.MODEL_CONFIG.CTX_LEN pipeline = storage.pipeline current_text = cleanMessages(messages) tokens = pipeline.encode(current_text) if len(tokens) + max_gen_tokens < limit: return messages system_msgs = [m for m in messages if m.role == "System"] other_msgs = [m for m in messages if m.role != "System"] while len(other_msgs) > 1: candidate_text = cleanMessages(system_msgs + other_msgs) if len(pipeline.encode(candidate_text)) + max_gen_tokens < limit: break other_msgs.pop(0) return system_msgs + other_msgs async def runPrefill(request: ChatCompletionRequest, ctx: str, model_tokens: List[int], model_state): ctx = ctx.replace("\r\n", "\n") tokens = MODEL_STORAGE[request.model].pipeline.encode(ctx) model_tokens.extend([int(x) for x in tokens]) while len(tokens) > 0: out, model_state = MODEL_STORAGE[request.model].model.forward(tokens[: CONFIG.CHUNK_LEN], model_state) tokens = tokens[CONFIG.CHUNK_LEN :] await asyncio.sleep(0) return out, model_tokens, model_state def generate(request: ChatCompletionRequest, out, model_tokens: List[int], model_state, max_tokens=2048): args = PIPELINE_ARGS( temperature=request.temperature, top_p=request.top_p, alpha_frequency=request.count_penalty, alpha_presence=request.presence_penalty, token_ban=[], token_stop=[0] ) occurrence = {} out_tokens = [] out_last = 0 cache_word_list = [] stop_sequences = request.stop if request.stop else [] stop_sequences.append("") for i in range(max_tokens): for n in occurrence: out[n] -= args.alpha_presence + occurrence[n] * args.alpha_frequency token = MODEL_STORAGE[request.model].pipeline.sample_logits(out, temperature=args.temperature, top_p=args.top_p) if token == 0: yield {"content": "".join(cache_word_list), "finish_reason": "stop", "state": model_state} del out; gc.collect(); return out, model_state = MODEL_STORAGE[request.model].model.forward([token], model_state) model_tokens.append(token) out_tokens.append(token) for xxx in occurrence: occurrence[xxx] *= request.penalty_decay occurrence[token] = 1 + (occurrence.get(token, 0)) tmp = MODEL_STORAGE[request.model].pipeline.decode(out_tokens[out_last:]) if "\ufffd" in tmp: continue cache_word_list.append(tmp) out_last = i + 1 current_buffer = "".join(cache_word_list) if "" in current_buffer: pre_call = current_buffer.split("")[0] yield {"content": pre_call, "finish_reason": "tool_start", "state": model_state} del out; gc.collect(); return for s in stop_sequences: if s in current_buffer and s != "": final_content = current_buffer.split(s)[0] yield {"content": final_content, "finish_reason": "stop", "state": model_state} del out; gc.collect(); return if len(cache_word_list) > 2: yield {"content": cache_word_list.pop(0), "finish_reason": None} yield {"content": "".join(cache_word_list), "finish_reason": "length"} async def chatResponseStream(request: ChatCompletionRequest, model_state: any, completionId: str, enableReasoning: bool): current_messages = request.messages for step in range(4): clean_msg = cleanMessages(current_messages, enableReasoning) prompt = f"{clean_msg}\n\nAssistant:{' " in current_gen: full_tool_call = current_gen.split("")[0] break finally: pass if full_tool_call: result = ToolEngine.execute(full_tool_call) current_messages.append(ChatMessage(role="assistant", content=f"{full_tool_call}")) current_messages.append(ChatMessage(role="tool", content=result, name="system")) else: break else: break yield "data: [DONE]\n\n" @app.post("/v1/chat/completions") @app.post("/v1/chat/") @app.post("/v1/completions") @app.post("/v1/responses") @app.post("/responses") @app.post("/api/generate") @app.post("/api/v1/chat/completions") async def chat_completions(request: ChatCompletionRequest): completionId = str(next(CompletionIdGenerator)) raw_model = request.model model_key = request.model.split(":")[0].replace(":online", "") is_reasoning = ":thinking" in request.model target_model = model_key if "rwkv-latest" in model_key: if is_reasoning and DEFAULT_REASONING_MODEL_NAME: target_model = DEFAULT_REASONING_MODEL_NAME elif DEFALUT_MODEL_NAME: target_model = DEFALUT_MODEL_NAME if target_model not in MODEL_STORAGE: raise HTTPException(404, f"Model {target_model} not loaded.") request.model = target_model default_sampler = MODEL_STORAGE[target_model].MODEL_CONFIG.DEFAULT_SAMPLER req_data = request.model_dump() for k, v in default_sampler.model_dump().items(): if req_data.get(k) is None: req_data[k] = v realRequest = ChatCompletionRequest(**req_data) # FIX: Handle missing messages (legacy completion API support) if realRequest.messages is None: if realRequest.prompt: realRequest.messages = [ChatMessage(role="user", content=realRequest.prompt)] else: # Fallback to empty list to prevent crashes in prune_context/cleanMessages realRequest.messages = [] enable_tools = ":online" in raw_model or realRequest.tools is not None if enable_tools: sys_msg = ChatMessage(role="System", content=ToolEngine.TOOL_SYSTEM_PROMPT) if realRequest.messages: if realRequest.messages[0].role == "System": realRequest.messages[0].content += f"\n\n{ToolEngine.TOOL_SYSTEM_PROMPT}" else: realRequest.messages.insert(0, sys_msg) else: realRequest.messages.append(sys_msg) realRequest.messages = prune_context(realRequest.messages, target_model, realRequest.max_tokens or 1024) return StreamingResponse(chatResponseStream(realRequest, None, completionId, is_reasoning), media_type="text/event-stream") @app.get("/api/v1/models") @app.get("/v1/models") @app.get("/models") async def list_models(): models_list = [] ts = int(time.time()) for model_id in MODEL_STORAGE.keys(): models_list.append({"id": model_id, "object": "model", "created": ts, "owned_by": "rwkv-server"}) models_list.append({"id": f"{model_id}:online", "object": "model", "created": ts, "owned_by": "rwkv-server"}) if DEFALUT_MODEL_NAME: models_list.append({"id": "rwkv-latest", "object": "model", "created": ts, "owned_by": "rwkv-system"}) models_list.append({"id": "rwkv-latest:online", "object": "model", "created": ts, "owned_by": "rwkv-system"}) if DEFAULT_REASONING_MODEL_NAME: models_list.append({"id": "rwkv-latest:thinking", "object": "model", "created": ts, "owned_by": "rwkv-system"}) models_list.append({"id": "rwkv-latest:thinking:online", "object": "model", "created": ts, "owned_by": "rwkv-system"}) return {"object": "list", "data": models_list} app.mount("/", StaticFiles(directory="dist-frontend", html=True), name="static") if __name__ == "__main__": import uvicorn uvicorn.run(app, host=CONFIG.HOST, port=CONFIG.PORT)