|
|
import os |
|
|
|
|
|
if os.environ.get("MODELSCOPE_ENVIRONMENT") == "studio": |
|
|
from modelscope import patch_hub |
|
|
patch_hub() |
|
|
|
|
|
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:256" |
|
|
|
|
|
from config import CONFIG, ModelConfig |
|
|
from utils import ( |
|
|
cleanMessages, |
|
|
parse_think_response, |
|
|
remove_nested_think_tags_stack, |
|
|
format_bytes, |
|
|
log, |
|
|
detect_tools_and_reasoning, |
|
|
universal_tool, |
|
|
) |
|
|
|
|
|
import copy, types, gc, sys, re, time, collections, asyncio |
|
|
from huggingface_hub import hf_hub_download |
|
|
from loguru import logger |
|
|
from rich import print |
|
|
|
|
|
from snowflake import SnowflakeGenerator |
|
|
|
|
|
CompletionIdGenerator = SnowflakeGenerator(42, timestamp=1741101491595) |
|
|
|
|
|
from typing import List, Optional, Union, Any, Dict |
|
|
import uuid |
|
|
from pydantic import BaseModel, Field, model_validator |
|
|
from pydantic_settings import BaseSettings |
|
|
|
|
|
import numpy as np |
|
|
import torch |
|
|
|
|
|
if "cuda" in CONFIG.STRATEGY.lower() and not torch.cuda.is_available(): |
|
|
logger.info(f"CUDA not found, fall back to cpu") |
|
|
CONFIG.STRATEGY = "cpu fp16" |
|
|
|
|
|
_s = CONFIG.STRATEGY.lower() |
|
|
if ("cpu" in _s or "cuda" in _s) and not ("fp16" in _s or "fp32" in _s): |
|
|
logger.info(f"STRATEGY missing precision, appending 'fp16' to `{CONFIG.STRATEGY}`") |
|
|
CONFIG.STRATEGY = CONFIG.STRATEGY + " fp16" |
|
|
|
|
|
try: |
|
|
from pynvml import nvmlInit, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo |
|
|
except Exception: |
|
|
nvmlInit = None |
|
|
nvmlDeviceGetHandleByIndex = None |
|
|
nvmlDeviceGetMemoryInfo = None |
|
|
|
|
|
if "cuda" in CONFIG.STRATEGY.lower() and nvmlInit is not None and nvmlDeviceGetHandleByIndex is not None: |
|
|
nvmlInit() |
|
|
gpu_h = nvmlDeviceGetHandleByIndex(0) |
|
|
|
|
|
def logGPUState(): |
|
|
if "cuda" in CONFIG.STRATEGY and nvmlDeviceGetMemoryInfo is not None: |
|
|
gpu_info = nvmlDeviceGetMemoryInfo(gpu_h) |
|
|
logger.info( |
|
|
f"[STATUS] Torch - {format_bytes(torch.cuda.memory_allocated())} - NVML - vram {format_bytes(gpu_info.total)} used {format_bytes(gpu_info.used)} free {format_bytes(gpu_info.free)}" |
|
|
) |
|
|
|
|
|
torch.backends.cudnn.benchmark = True |
|
|
torch.backends.cudnn.allow_tf32 = True |
|
|
torch.backends.cuda.matmul.allow_tf32 = True |
|
|
os.environ["RWKV_V7_ON"] = "1" |
|
|
os.environ["RWKV_JIT_ON"] = "1" |
|
|
os.environ["RWKV_CUDA_ON"] = ( |
|
|
"1" if CONFIG.RWKV_CUDA_ON and "cuda" in CONFIG.STRATEGY.lower() else "0" |
|
|
) |
|
|
|
|
|
from rwkv.model import RWKV |
|
|
from rwkv.utils import PIPELINE, PIPELINE_ARGS |
|
|
|
|
|
from fastapi import FastAPI, HTTPException, UploadFile, File |
|
|
from starlette.background import BackgroundTask |
|
|
from fastapi.responses import StreamingResponse |
|
|
from fastapi.middleware.cors import CORSMiddleware |
|
|
from fastapi.staticfiles import StaticFiles |
|
|
from fastapi.middleware.gzip import GZipMiddleware |
|
|
|
|
|
from api_types import ( |
|
|
ChatMessage, |
|
|
ChatCompletion, |
|
|
ChatCompletionChunk, |
|
|
Usage, |
|
|
PromptTokensDetails, |
|
|
ChatCompletionChoice, |
|
|
ChatCompletionMessage, |
|
|
SamplerConfig, |
|
|
UploadedFile, |
|
|
FileUploadResponse, |
|
|
) |
|
|
|
|
|
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 |
|
|
|
|
|
STATE_STORE: Dict[tuple, Any] = {} |
|
|
|
|
|
_STATE_STORE_PATH = getattr(CONFIG, 'STATE_STORE_PATH', './state_store.json') |
|
|
_LAST_STATE_STORE_WRITE = 0 |
|
|
|
|
|
TOOL_CALL_RE = re.compile(r"<tool-call>\s*(\{.*?\})\s*</tool-call>", re.S) |
|
|
|
|
|
UPLOADED_FILES: Dict[str, dict] = {} |
|
|
|
|
|
def _serialize_state_store() -> dict: |
|
|
serial = {} |
|
|
for (model_name, state_name), entry in STATE_STORE.items(): |
|
|
try: |
|
|
mt = entry.get('model_tokens') if isinstance(entry, dict) else None |
|
|
if mt is None: |
|
|
continue |
|
|
serial[f"{model_name}|{state_name}"] = { |
|
|
'model': model_name, |
|
|
'state_name': state_name, |
|
|
'model_tokens': mt, |
|
|
} |
|
|
except Exception: |
|
|
continue |
|
|
return serial |
|
|
|
|
|
def _load_state_store_from_disk(): |
|
|
global STATE_STORE |
|
|
try: |
|
|
if os.path.exists(_STATE_STORE_PATH): |
|
|
import json |
|
|
|
|
|
with open(_STATE_STORE_PATH, 'r', encoding='utf-8') as f: |
|
|
data = json.load(f) |
|
|
for k, v in data.items(): |
|
|
model = v.get('model') |
|
|
state_name = v.get('state_name') |
|
|
model_tokens = v.get('model_tokens') |
|
|
if model and state_name and isinstance(model_tokens, list): |
|
|
STATE_STORE[(model, state_name)] = { |
|
|
'state': None, |
|
|
'model_tokens': model_tokens, |
|
|
} |
|
|
logger.info(f"Loaded {len(STATE_STORE)} entries from state store file {_STATE_STORE_PATH}") |
|
|
except Exception as e: |
|
|
logger.info(f"Failed to load state store from disk: {e}") |
|
|
|
|
|
def _save_state_store_to_disk(force=False): |
|
|
global _LAST_STATE_STORE_WRITE |
|
|
now = time.time() |
|
|
if not force and now - _LAST_STATE_STORE_WRITE < getattr(CONFIG, 'STATE_STORE_FLUSH_INTERVAL', 5): |
|
|
return |
|
|
try: |
|
|
serial = _serialize_state_store() |
|
|
if not serial: |
|
|
return |
|
|
import json |
|
|
tmp = _STATE_STORE_PATH + ".tmp" |
|
|
with open(tmp, 'w', encoding='utf-8') as f: |
|
|
json.dump(serial, f) |
|
|
os.replace(tmp, _STATE_STORE_PATH) |
|
|
_LAST_STATE_STORE_WRITE = now |
|
|
except Exception as e: |
|
|
logger.info(f"Write state store to disk failed: {e}") |
|
|
|
|
|
def _recompute_out_and_state_from_tokens(model_name: str, model_tokens: List[int]): |
|
|
ms = MODEL_STORAGE.get(model_name) |
|
|
if not ms or not ms.model: |
|
|
return None, None |
|
|
model_state = None |
|
|
out = None |
|
|
tokens = list(model_tokens) if isinstance(model_tokens, list) else [0] |
|
|
while len(tokens) > 0: |
|
|
out, model_state = ms.model.forward(tokens[: CONFIG.CHUNK_LEN], model_state) |
|
|
tokens = tokens[CONFIG.CHUNK_LEN :] |
|
|
return out, model_state |
|
|
|
|
|
logger.info(f"STRATEGY - {CONFIG.STRATEGY}") |
|
|
|
|
|
logGPUState() |
|
|
|
|
|
logger.info(f"Configured {len(CONFIG.MODELS)} model(s) in ROOT config") |
|
|
|
|
|
for model_config in CONFIG.MODELS: |
|
|
logger.info(f"Load Model - {model_config.SERVICE_NAME}") |
|
|
|
|
|
if model_config.MODEL_FILE_PATH == None: |
|
|
model_config.MODEL_FILE_PATH = hf_hub_download( |
|
|
repo_id=str(model_config.DOWNLOAD_MODEL_REPO_ID), |
|
|
filename=str(model_config.DOWNLOAD_MODEL_FILE_NAME), |
|
|
local_dir=str(model_config.DOWNLOAD_MODEL_DIR), |
|
|
) |
|
|
logger.info(f"Load Model - Path - {model_config.MODEL_FILE_PATH}") |
|
|
|
|
|
if model_config.DEFAULT_CHAT: |
|
|
if DEFALUT_MODEL_NAME != None: |
|
|
logger.info( |
|
|
f"Load Model - Replace `DEFALUT_MODEL_NAME` from `{DEFALUT_MODEL_NAME}` to `{model_config.SERVICE_NAME}`" |
|
|
) |
|
|
DEFALUT_MODEL_NAME = model_config.SERVICE_NAME |
|
|
|
|
|
if model_config.DEFAULT_REASONING: |
|
|
if DEFAULT_REASONING_MODEL_NAME != None: |
|
|
logger.info( |
|
|
f"Load Model - Replace `DEFAULT_REASONING_MODEL_NAME` from `{DEFAULT_REASONING_MODEL_NAME}` to `{model_config.SERVICE_NAME}`" |
|
|
) |
|
|
DEFAULT_REASONING_MODEL_NAME = model_config.SERVICE_NAME |
|
|
|
|
|
logger.info(f"Load Model - Loading `{model_config.SERVICE_NAME}`") |
|
|
print(model_config.DEFAULT_SAMPLER) |
|
|
|
|
|
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() |
|
|
logGPUState() |
|
|
|
|
|
logger.info(f"Load Model - DEFALUT_MODEL_NAME is `{DEFALUT_MODEL_NAME}`") |
|
|
logger.info( |
|
|
f"Load Model - DEFAULT_REASONING_MODEL_NAME is `{DEFAULT_REASONING_MODEL_NAME}`" |
|
|
) |
|
|
|
|
|
if len(MODEL_STORAGE) == 1: |
|
|
single_name = list(MODEL_STORAGE.keys())[0] |
|
|
if DEFALUT_MODEL_NAME != single_name: |
|
|
DEFALUT_MODEL_NAME = single_name |
|
|
logger.info(f"Load Model - Only one model present; DEFALUT_MODEL_NAME set to `{DEFALUT_MODEL_NAME}`") |
|
|
if DEFAULT_REASONING_MODEL_NAME != single_name: |
|
|
DEFAULT_REASONING_MODEL_NAME = single_name |
|
|
logger.info(f"Load Model - Only one model present; DEFAULT_REASONING_MODEL_NAME set to `{DEFAULT_REASONING_MODEL_NAME}`") |
|
|
|
|
|
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=True) |
|
|
state_name: Optional[str] = Field(default=None) |
|
|
include_usage: Optional[bool] = Field(default=False) |
|
|
stop: Optional[list[str]] = Field(["\n\n"]) |
|
|
stop_tokens: Optional[list[int]] = Field([0]) |
|
|
web_search: Optional[bool] = Field(default=True) |
|
|
enable_web_search: Optional[bool] = Field(default=True) |
|
|
auto_web_search: Optional[bool] = Field(default=True) |
|
|
enable_tools: Optional[bool] = Field(default=None) |
|
|
auto_tools: Optional[bool] = Field(default=True) |
|
|
enable_reasoning: Optional[bool] = Field(default=True) |
|
|
auto_reasoning: Optional[bool] = Field(default=True) |
|
|
enable_universal: Optional[bool] = Field(default=None) |
|
|
auto_universal: Optional[bool] = Field(default=True) |
|
|
search_top_k: Optional[int] = Field(default=3) |
|
|
tools: Optional[List[Dict[str, Any]]] = Field(default=None) |
|
|
sampler_allow_web_search: Optional[bool] = Field(default=None) |
|
|
sampler_allow_tools: Optional[bool] = Field(default=None) |
|
|
sampler_allow_reasoning: Optional[bool] = Field(default=None) |
|
|
sampler: Optional[SamplerConfig] = Field(default=None) |
|
|
file_ids: Optional[List[str]] = Field(default=None) |
|
|
enable_file_tool: Optional[bool] = Field(default=None) |
|
|
auto_file_tool: Optional[bool] = Field(default=None) |
|
|
sampler_allow_file_tool: Optional[bool] = Field(default=None) |
|
|
|
|
|
@model_validator(mode="before") |
|
|
@classmethod |
|
|
def validate_mutual_exclusivity(cls, data: Any) -> Any: |
|
|
if not isinstance(data, dict): |
|
|
return data |
|
|
|
|
|
messages_provided = "messages" in data and data["messages"] != None |
|
|
prompt_provided = "prompt" in data and data["prompt"] != None |
|
|
|
|
|
if messages_provided and prompt_provided: |
|
|
raise ValueError("messages and prompt cannot coexist. Choose one.") |
|
|
if not messages_provided and not prompt_provided: |
|
|
raise ValueError("Either messages or prompt must be provided.") |
|
|
return data |
|
|
|
|
|
app = FastAPI(title="RWKV OpenAI-Compatible API") |
|
|
|
|
|
app.add_middleware( |
|
|
CORSMiddleware, |
|
|
allow_origins=["*"], |
|
|
allow_credentials=True, |
|
|
allow_methods=["*"], |
|
|
allow_headers=["*"], |
|
|
) |
|
|
app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=5) |
|
|
|
|
|
@app.on_event("startup") |
|
|
async def _startup_state_load_and_persist_loop(): |
|
|
_load_state_store_from_disk() |
|
|
|
|
|
async def _persist_loop(): |
|
|
while True: |
|
|
try: |
|
|
_save_state_store_to_disk(force=False) |
|
|
except Exception: |
|
|
pass |
|
|
await asyncio.sleep(getattr(CONFIG, 'STATE_STORE_FLUSH_INTERVAL', 5)) |
|
|
|
|
|
try: |
|
|
asyncio.create_task(_persist_loop()) |
|
|
except Exception: |
|
|
pass |
|
|
|
|
|
async def runPrefill( |
|
|
request: ChatCompletionRequest, ctx: str, model_tokens: List[int], model_state |
|
|
): |
|
|
ctx = ctx.replace("\r\n", "\n") |
|
|
out = None |
|
|
|
|
|
ms = MODEL_STORAGE.get(request.model) |
|
|
if not ms or not ms.pipeline or not ms.model: |
|
|
raise HTTPException(500, f"Model {request.model} not loaded or pipeline missing") |
|
|
tokens = ms.pipeline.encode(ctx) |
|
|
tokens = [int(x) for x in tokens] |
|
|
model_tokens += tokens |
|
|
|
|
|
while len(tokens) > 0: |
|
|
out, model_state = ms.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, |
|
|
): |
|
|
ms = MODEL_STORAGE.get(request.model) |
|
|
if not ms or not ms.pipeline or not ms.model: |
|
|
raise HTTPException(500, f"Model {request.model} not loaded or pipeline missing") |
|
|
|
|
|
temperature = request.temperature if request.temperature is not None else 0.2 |
|
|
top_p = request.top_p if request.top_p is not None else 0.9 |
|
|
alpha_frequency = request.count_penalty if request.count_penalty is not None else 0.0 |
|
|
alpha_presence = request.presence_penalty if request.presence_penalty is not None else 0.0 |
|
|
penalty_decay = request.penalty_decay if request.penalty_decay is not None else 0.5 |
|
|
|
|
|
args = PIPELINE_ARGS( |
|
|
temperature=max(0.2, temperature), |
|
|
top_p=top_p, |
|
|
alpha_frequency=alpha_frequency, |
|
|
alpha_presence=alpha_presence, |
|
|
token_ban=[], |
|
|
token_stop=[0], |
|
|
) |
|
|
|
|
|
occurrence = {} |
|
|
out_tokens: List[int] = [] |
|
|
out_last = 0 |
|
|
|
|
|
for i in range(max_tokens): |
|
|
for n in occurrence: |
|
|
out[n] -= args.alpha_presence + occurrence[n] * args.alpha_frequency |
|
|
|
|
|
token = ms.pipeline.sample_logits( |
|
|
out, temperature=args.temperature, top_p=args.top_p |
|
|
) |
|
|
|
|
|
if token == 0 and request.stop_tokens and token in request.stop_tokens: |
|
|
yield { |
|
|
"content": "", |
|
|
"tokens": out_tokens[out_last:], |
|
|
"finish_reason": "stop:token:0", |
|
|
"state": model_state, |
|
|
} |
|
|
|
|
|
del out |
|
|
gc.collect() |
|
|
return |
|
|
|
|
|
out, model_state = ms.model.forward([token], model_state) |
|
|
model_tokens.append(token) |
|
|
out_tokens.append(token) |
|
|
|
|
|
if request.stop_tokens and token in request.stop_tokens: |
|
|
yield { |
|
|
"content": "", |
|
|
"tokens": out_tokens[out_last:], |
|
|
"finish_reason": f"stop:token:{token}", |
|
|
"state": model_state, |
|
|
} |
|
|
|
|
|
del out |
|
|
gc.collect() |
|
|
return |
|
|
|
|
|
for xxx in list(occurrence.keys()): |
|
|
occurrence[xxx] *= penalty_decay |
|
|
occurrence[token] = 1 + (occurrence[token] if token in occurrence else 0) |
|
|
|
|
|
decoded = ms.pipeline.decode([token]) |
|
|
if "\ufffd" in decoded: |
|
|
continue |
|
|
|
|
|
yield { |
|
|
"content": decoded, |
|
|
"tokens": [token], |
|
|
"finish_reason": None, |
|
|
"state": model_state, |
|
|
} |
|
|
out_last = i + 1 |
|
|
|
|
|
else: |
|
|
yield { |
|
|
"content": "", |
|
|
"tokens": [], |
|
|
"finish_reason": "length", |
|
|
} |
|
|
|
|
|
async def chatResponse( |
|
|
request: ChatCompletionRequest, |
|
|
model_state: Any, |
|
|
completionId: str, |
|
|
enableReasoning: bool, |
|
|
) -> ChatCompletion: |
|
|
createTimestamp = time.time() |
|
|
|
|
|
raw_prompt = request.prompt.strip() if request.prompt is not None else cleanMessages(request.messages or []) |
|
|
detection = detect_tools_and_reasoning(raw_prompt) |
|
|
prompt = raw_prompt if request.prompt is not None else f"{cleanMessages(request.messages or [])}\n\nAssistant:{' <think' if enableReasoning else ''}" |
|
|
|
|
|
web_search_enabled = ( |
|
|
True |
|
|
if (request.enable_web_search is not None and request.enable_web_search) |
|
|
else ( |
|
|
request.web_search |
|
|
or (request.auto_web_search if request.auto_web_search is not None else CONFIG.AUTO_ENABLE_WEB_SEARCH and detection.get('need_web_search')) |
|
|
) |
|
|
) |
|
|
if not getattr(CONFIG, 'ENABLE_WEB_SEARCH_BY_DEFAULT', True) and request.enable_web_search is None and not request.web_search: |
|
|
web_search_enabled = False |
|
|
if not getattr(CONFIG, 'ENABLE_WEB_SEARCH_BY_DEFAULT', True) and request.enable_web_search is None and not request.web_search: |
|
|
web_search_enabled = False |
|
|
try: |
|
|
if request.sampler and getattr(request.sampler, 'ALLOW_WEB_SEARCH', None) is not None: |
|
|
web_search_enabled = bool(request.sampler.ALLOW_WEB_SEARCH) |
|
|
elif hasattr(request, 'sampler_allow_web_search') and request.sampler_allow_web_search is not None: |
|
|
web_search_enabled = bool(request.sampler_allow_web_search) |
|
|
else: |
|
|
ms = MODEL_STORAGE.get(request.model) |
|
|
if ms and ms.MODEL_CONFIG: |
|
|
if hasattr(ms.MODEL_CONFIG, 'DEFAULT_SAMPLER') and getattr(ms.MODEL_CONFIG.DEFAULT_SAMPLER, 'ALLOW_WEB_SEARCH', None) is not None: |
|
|
web_search_enabled = bool(ms.MODEL_CONFIG.DEFAULT_SAMPLER.ALLOW_WEB_SEARCH) |
|
|
elif hasattr(ms.MODEL_CONFIG, 'ALLOW_WEB_SEARCH') and not ms.MODEL_CONFIG.ALLOW_WEB_SEARCH: |
|
|
web_search_enabled = False |
|
|
except Exception: |
|
|
pass |
|
|
|
|
|
if request.enable_file_tool is not None: |
|
|
file_tool_enabled = bool(request.enable_file_tool) |
|
|
else: |
|
|
auto_file_flag = request.auto_file_tool if request.auto_file_tool is not None else CONFIG.AUTO_ENABLE_TOOLS |
|
|
file_tool_enabled = bool((request.file_ids and len(request.file_ids) > 0) or (auto_file_flag and request.file_ids)) |
|
|
if not getattr(CONFIG, 'ALLOW_FILE_TOOL_BY_DEFAULT', True) and request.enable_file_tool is None: |
|
|
file_tool_enabled = False |
|
|
try: |
|
|
if request.sampler and getattr(request.sampler, 'ALLOW_FILE_TOOL', None) is not None: |
|
|
file_tool_enabled = bool(request.sampler.ALLOW_FILE_TOOL) |
|
|
elif hasattr(request, 'sampler_allow_file_tool') and request.sampler_allow_file_tool is not None: |
|
|
file_tool_enabled = bool(request.sampler_allow_file_tool) |
|
|
else: |
|
|
ms = MODEL_STORAGE.get(request.model) |
|
|
if ms and ms.MODEL_CONFIG: |
|
|
if hasattr(ms.MODEL_CONFIG, 'DEFAULT_SAMPLER') and getattr(ms.MODEL_CONFIG.DEFAULT_SAMPLER, 'ALLOW_FILE_TOOL', None) is not None: |
|
|
file_tool_enabled = bool(ms.MODEL_CONFIG.DEFAULT_SAMPLER.ALLOW_FILE_TOOL) |
|
|
elif hasattr(ms.MODEL_CONFIG, 'ALLOW_FILE_TOOL') and not ms.MODEL_CONFIG.ALLOW_FILE_TOOL: |
|
|
file_tool_enabled = False |
|
|
except Exception: |
|
|
pass |
|
|
|
|
|
if request.enable_tools is not None: |
|
|
tools_enabled = bool(request.enable_tools) |
|
|
else: |
|
|
auto_tools_flag = request.auto_tools if request.auto_tools is not None else CONFIG.AUTO_ENABLE_TOOLS |
|
|
tools_enabled = bool(request.tools) or CONFIG.ENABLE_TOOLS_BY_DEFAULT or (auto_tools_flag and (detection.get('need_calc') or detection.get('need_web_search'))) |
|
|
try: |
|
|
if request.sampler and getattr(request.sampler, 'ALLOW_TOOLS', None) is not None: |
|
|
tools_enabled = bool(request.sampler.ALLOW_TOOLS) |
|
|
elif hasattr(request, 'sampler_allow_tools') and request.sampler_allow_tools is not None: |
|
|
tools_enabled = bool(request.sampler_allow_tools) |
|
|
else: |
|
|
ms = MODEL_STORAGE.get(request.model) |
|
|
if ms and ms.MODEL_CONFIG: |
|
|
if hasattr(ms.MODEL_CONFIG, 'DEFAULT_SAMPLER') and getattr(ms.MODEL_CONFIG.DEFAULT_SAMPLER, 'ALLOW_TOOLS', None) is not None: |
|
|
if not ms.MODEL_CONFIG.DEFAULT_SAMPLER.ALLOW_TOOLS: |
|
|
tools_enabled = False |
|
|
elif hasattr(ms.MODEL_CONFIG, 'ALLOW_TOOLS') and not ms.MODEL_CONFIG.ALLOW_TOOLS: |
|
|
tools_enabled = False |
|
|
except Exception: |
|
|
pass |
|
|
|
|
|
reasoning_enabled = bool( |
|
|
True |
|
|
if (request.enable_reasoning is not None and request.enable_reasoning) |
|
|
else ( |
|
|
bool(enableReasoning) or bool(request.auto_reasoning if request.auto_reasoning is not None else (CONFIG.AUTO_ENABLE_REASONING and bool(detection.get('need_reasoning')))) |
|
|
) |
|
|
) |
|
|
if not getattr(CONFIG, 'ENABLE_REASONING_BY_DEFAULT', True) and request.enable_reasoning is None: |
|
|
reasoning_enabled = False |
|
|
try: |
|
|
if request.sampler and getattr(request.sampler, 'ALLOW_REASONING', None) is not None: |
|
|
reasoning_enabled = bool(request.sampler.ALLOW_REASONING) |
|
|
elif hasattr(request, 'sampler_allow_reasoning') and request.sampler_allow_reasoning is not None: |
|
|
reasoning_enabled = bool(request.sampler_allow_reasoning) |
|
|
else: |
|
|
ms = MODEL_STORAGE.get(request.model) |
|
|
if ms and ms.MODEL_CONFIG: |
|
|
if hasattr(ms.MODEL_CONFIG, 'DEFAULT_SAMPLER') and getattr(ms.MODEL_CONFIG.DEFAULT_SAMPLER, 'ALLOW_REASONING', None) is not None: |
|
|
if not ms.MODEL_CONFIG.DEFAULT_SAMPLER.ALLOW_REASONING: |
|
|
reasoning_enabled = False |
|
|
elif hasattr(ms.MODEL_CONFIG, 'ALLOW_REASONING') and not ms.MODEL_CONFIG.ALLOW_REASONING: |
|
|
reasoning_enabled = False |
|
|
except Exception: |
|
|
pass |
|
|
|
|
|
enableReasoning = reasoning_enabled |
|
|
try: |
|
|
ms = MODEL_STORAGE.get(request.model) |
|
|
if ms and ms.MODEL_CONFIG and hasattr(ms.MODEL_CONFIG, 'ALLOW_REASONING') and not ms.MODEL_CONFIG.ALLOW_REASONING: |
|
|
enableReasoning = False |
|
|
except Exception: |
|
|
pass |
|
|
|
|
|
if request.enable_web_search is None: |
|
|
request.web_search = web_search_enabled |
|
|
if tools_enabled and not request.tools: |
|
|
if detection.get('detected_tools'): |
|
|
request.tools = detection.get('detected_tools') |
|
|
if (request.enable_universal is True) or ( |
|
|
request.enable_universal is None and (request.auto_universal if request.auto_universal is not None else CONFIG.AUTO_ENABLE_TOOLS and detection.get('need_universal')) |
|
|
): |
|
|
if not request.tools: |
|
|
request.tools = [{"name": "universal", "args": {"query": raw_prompt}}] |
|
|
|
|
|
executed_tool_calls = [] |
|
|
if file_tool_enabled and request.file_ids: |
|
|
for fid in request.file_ids: |
|
|
try: |
|
|
if fid not in UPLOADED_FILES: |
|
|
continue |
|
|
meta = UPLOADED_FILES.get(fid) |
|
|
if not meta: |
|
|
continue |
|
|
from utils import file_read_from_path |
|
|
fpath = meta.get('path') |
|
|
if not fpath or not os.path.exists(fpath): |
|
|
continue |
|
|
file_content = file_read_from_path(fpath, 200000) |
|
|
if file_content: |
|
|
exec_entry = {"name": "file_inject", "args": {"file_id": fid}, "result": {"action": "file_inject", "result": "injected", "metadata": {"file_id": fid, "filename": meta.get('filename')}}} |
|
|
executed_tool_calls.append(exec_entry) |
|
|
prompt = (f"AttachedFile: {meta.get('filename')} (id:{fid})\n{file_content}\n\n" + prompt) |
|
|
except Exception as e: |
|
|
logger.info(f"File injection error: {e}") |
|
|
if file_tool_enabled and request.file_ids: |
|
|
for fid in request.file_ids: |
|
|
try: |
|
|
if fid not in UPLOADED_FILES: |
|
|
continue |
|
|
meta = UPLOADED_FILES.get(fid) |
|
|
if not meta: |
|
|
continue |
|
|
from utils import file_read_from_path |
|
|
fpath = meta.get('path') |
|
|
if not fpath or not os.path.exists(fpath): |
|
|
continue |
|
|
file_content = file_read_from_path(fpath, 200000) |
|
|
if file_content: |
|
|
exec_entry = {"name": "file_inject", "args": {"file_id": fid}, "result": {"action": "file_inject", "result": "injected", "metadata": {"file_id": fid, "filename": meta.get('filename')}}} |
|
|
executed_tool_calls.append(exec_entry) |
|
|
prompt = (f"AttachedFile: {meta.get('filename')} (id:{fid})\n{file_content}\n\n" + prompt) |
|
|
except Exception as e: |
|
|
logger.info(f"File injection error: {e}") |
|
|
if request.tools: |
|
|
try: |
|
|
for tool in request.tools: |
|
|
name = tool.get('name') |
|
|
args = tool.get('args', {}) |
|
|
if name == 'web_search': |
|
|
from utils import web_search |
|
|
|
|
|
search_q = args.get('query') or (request.prompt if request.prompt else cleanMessages(request.messages or [])) |
|
|
search_top_k = int(args.get('top_k') or request.search_top_k or 3) |
|
|
search_str = web_search(search_q, search_top_k) |
|
|
if search_str: |
|
|
search_res_struct = {"action": "web_search", "result": str(search_str), "metadata": {"query": search_q, "top_k": search_top_k, "confidence": 0.9}} |
|
|
executed_tool_calls.append({"name": "web_search", "args": {"query": search_q, "top_k": search_top_k}, "result": search_res_struct}) |
|
|
prompt = (f"ToolResults:\n{search_res_struct.get('result')}\n\nUse these results to answer the prompt.\n\n" + prompt) |
|
|
elif name == 'calc' or name == 'calculator': |
|
|
from utils import calc |
|
|
|
|
|
expr = args.get('expression') |
|
|
if expr: |
|
|
calc_res = calc(expr) |
|
|
calc_res_struct = {"action": "calc", "result": str(calc_res), "metadata": {"expression": expr, "confidence": 0.98}} |
|
|
executed_tool_calls.append({"name": "calc", "args": {"expression": expr}, "result": calc_res_struct}) |
|
|
prompt = (f"ToolResults:\nCalcResult:{expr} = {calc_res_struct.get('result')}\n\nUse this result to answer the prompt.\n\n" + prompt) |
|
|
elif name == 'universal': |
|
|
try: |
|
|
res = universal_tool(args or {"query": raw_prompt}, allow_web_search=bool(web_search_enabled), allow_tools=bool(tools_enabled), allow_file_tool=bool(file_tool_enabled)) |
|
|
if isinstance(res, dict): |
|
|
result_text = res.get('result') if res.get('result') is not None else '' |
|
|
else: |
|
|
result_text = str(res) |
|
|
executed_tool_calls.append({"name": "universal", "args": args, "result": res}) |
|
|
prompt = (f"ToolResults:\n{result_text}\n\nUse this result to answer the prompt.\n\n" + prompt) |
|
|
except Exception as e: |
|
|
logger.info(f"Universal tool execution error: {e}") |
|
|
else: |
|
|
logger.info(f"Unsupported tool requested: {name}") |
|
|
if name == 'file_read': |
|
|
try: |
|
|
fid = args.get('file_id') or args.get('id') or (request.file_ids[0] if request.file_ids else None) |
|
|
if not fid: |
|
|
continue |
|
|
if fid not in UPLOADED_FILES: |
|
|
continue |
|
|
meta = UPLOADED_FILES.get(fid) |
|
|
if not meta: |
|
|
continue |
|
|
from utils import file_read_from_path |
|
|
fpath = meta.get('path') |
|
|
if not fpath or not os.path.exists(fpath): |
|
|
continue |
|
|
file_content = file_read_from_path(fpath, int(args.get('max_bytes') or 100000)) |
|
|
exec_entry = {"name": "file_read", "args": {"file_id": fid, "max_bytes": int(args.get('max_bytes') or 100000)}, "result": {"action": "file_read", "result": file_content, "metadata": {"file_id": fid, "filename": meta.get('filename')}}} |
|
|
executed_tool_calls.append(exec_entry) |
|
|
_res = exec_entry.get('result') if isinstance(exec_entry, dict) else None |
|
|
_res_text = '' |
|
|
if isinstance(_res, dict): |
|
|
_res_text = _res.get('result') or '' |
|
|
elif _res is not None: |
|
|
_res_text = str(_res) |
|
|
prompt = (f"ToolResults:\n{_res_text}\n\nUse these file contents to answer the prompt.\n\n" + prompt) |
|
|
except Exception as e: |
|
|
logger.info(f"file_read tool error: {e}") |
|
|
except Exception as e: |
|
|
logger.info(f"Tool processing error: {e}") |
|
|
elif request.web_search or web_search_enabled: |
|
|
try: |
|
|
from utils import web_search |
|
|
|
|
|
search_q = request.prompt if request.prompt else cleanMessages(request.messages or []) |
|
|
search_res = web_search(search_q, int(request.search_top_k or 3)) |
|
|
if search_res: |
|
|
search_res_struct = {"action": "web_search", "result": str(search_res), "metadata": {"query": search_q, "top_k": int(request.search_top_k or 3), "confidence": 0.9}} |
|
|
executed_tool_calls.append({"name": "web_search", "args": {"query": search_q, "top_k": int(request.search_top_k or 3)}, "result": search_res_struct}) |
|
|
prompt = f"WebSearchResults:\n{search_res_struct.get('result')}\n\n" + prompt |
|
|
except Exception: |
|
|
pass |
|
|
logger.info(f"[REQ] {completionId} - prompt - {prompt}") |
|
|
|
|
|
if request.state_name: |
|
|
state_key = (request.model, request.state_name) |
|
|
if state_key in STATE_STORE: |
|
|
stored = STATE_STORE[state_key] |
|
|
model_state = stored.get('state', None) |
|
|
model_tokens = stored.get('model_tokens', [0]) |
|
|
if model_state is None: |
|
|
out, model_state = _recompute_out_and_state_from_tokens(request.model, model_tokens) |
|
|
else: |
|
|
out, _ = _recompute_out_and_state_from_tokens(request.model, model_tokens[-CONFIG.CHUNK_LEN :]) |
|
|
else: |
|
|
out, model_tokens, model_state = await runPrefill(request, prompt, [0], model_state) |
|
|
else: |
|
|
out, model_tokens, model_state = await runPrefill(request, prompt, [0], model_state) |
|
|
|
|
|
prefillTime = time.time() |
|
|
promptTokenCount = len(model_tokens) |
|
|
|
|
|
fullResponse = " <think" if enableReasoning else "" |
|
|
completionTokenCount = 0 |
|
|
finishReason = None |
|
|
model_initiated_tool_calls = 0 |
|
|
MODEL_MAX_TOOL_CALLS = 3 |
|
|
should_restart = True |
|
|
while should_restart: |
|
|
should_restart = False |
|
|
gen = generate( |
|
|
request, |
|
|
out, |
|
|
model_tokens, |
|
|
model_state, |
|
|
max_tokens=( |
|
|
64000 |
|
|
if "max_tokens" not in request.model_fields_set and enableReasoning |
|
|
else (request.max_tokens or 2048) |
|
|
), |
|
|
) |
|
|
for chunk in gen: |
|
|
fullResponse += chunk["content"] |
|
|
if model_initiated_tool_calls < MODEL_MAX_TOOL_CALLS: |
|
|
m = TOOL_CALL_RE.search(fullResponse) |
|
|
if m: |
|
|
try: |
|
|
payload_raw = m.group(1) |
|
|
import json |
|
|
|
|
|
payload = json.loads(payload_raw) |
|
|
tool_name = payload.get('name') |
|
|
tool_args = payload.get('args', {}) |
|
|
tool_res = None |
|
|
if tool_name == 'web_search': |
|
|
from utils import web_search |
|
|
|
|
|
q = tool_args.get('query') or (request.prompt if request.prompt else cleanMessages(request.messages or [])) |
|
|
k = int(tool_args.get('top_k') or request.search_top_k or 3) |
|
|
tool_res = web_search(q, k) |
|
|
elif tool_name in ('calc', 'calculator'): |
|
|
from utils import calc |
|
|
|
|
|
expr = tool_args.get('expression') |
|
|
if expr: |
|
|
tool_res = calc(expr) |
|
|
else: |
|
|
try: |
|
|
tool_res = universal_tool({'query': tool_args.get('query') or payload.get('query') or ''}, allow_web_search=bool(web_search_enabled), allow_tools=bool(tools_enabled), allow_file_tool=bool(file_tool_enabled)) |
|
|
except Exception: |
|
|
tool_res = None |
|
|
|
|
|
if tool_res: |
|
|
if not isinstance(tool_res, dict): |
|
|
if tool_name in ('calc', 'calculator'): |
|
|
tool_res_struct = {"action": "calc", "result": str(tool_res), "metadata": {"expression": tool_args.get('expression'), "confidence": 0.98}} |
|
|
elif tool_name == 'web_search': |
|
|
tool_res_struct = {"action": "web_search", "result": str(tool_res), "metadata": {"query": tool_args.get('query'), "top_k": tool_args.get('top_k') or request.search_top_k or 3, "confidence": 0.9}} |
|
|
else: |
|
|
tool_res_struct = {"action": tool_name, "result": str(tool_res), "metadata": {"confidence": 0.6}} |
|
|
else: |
|
|
tool_res_struct = tool_res |
|
|
exec_entry = {"name": tool_name, "args": tool_args, "result": tool_res_struct, 'initiated_by_model': True} |
|
|
executed_tool_calls.append(exec_entry) |
|
|
delta_text = f"ToolResults:\n{tool_res_struct.get('result')}\n\n" |
|
|
prompt = delta_text + prompt |
|
|
fullResponse = TOOL_CALL_RE.sub('', fullResponse) |
|
|
buffer = [fullResponse] |
|
|
out, model_tokens, model_state = await runPrefill(request, delta_text, model_tokens, model_state) |
|
|
model_initiated_tool_calls += 1 |
|
|
except Exception as e: |
|
|
logger.info(f"Model-initiated tool handling error: {e}") |
|
|
for stop_words in request.stop or []: |
|
|
if stop_words in fullResponse: |
|
|
finishReason = f"stop:words:{stop_words}" |
|
|
break |
|
|
completionTokenCount += 1 |
|
|
|
|
|
if chunk["finish_reason"]: |
|
|
finishReason = chunk["finish_reason"] |
|
|
|
|
|
generateTime = time.time() |
|
|
|
|
|
responseLog = { |
|
|
"content": fullResponse, |
|
|
"finish": finishReason, |
|
|
"prefill_len": promptTokenCount, |
|
|
"prefill_tps": round(promptTokenCount / (prefillTime - createTimestamp), 2), |
|
|
"gen_len": completionTokenCount, |
|
|
"gen_tps": round(completionTokenCount / (generateTime - prefillTime) if generateTime!=prefillTime else 0, 2), |
|
|
} |
|
|
logger.info(f"[RES] {completionId} - {responseLog}") |
|
|
|
|
|
reasoning_content, content = parse_think_response(fullResponse) |
|
|
|
|
|
response = ChatCompletion( |
|
|
id=completionId, |
|
|
created=int(createTimestamp), |
|
|
model=request.model, |
|
|
usage=Usage( |
|
|
prompt_tokens=promptTokenCount, |
|
|
completion_tokens=completionTokenCount, |
|
|
total_tokens=promptTokenCount + completionTokenCount, |
|
|
prompt_tokens_details=PromptTokensDetails(cached_tokens=0), |
|
|
), |
|
|
choices=[ |
|
|
ChatCompletionChoice( |
|
|
index=0, |
|
|
message=ChatCompletionMessage( |
|
|
role="Assistant", |
|
|
content=content, |
|
|
reasoning_content=reasoning_content if reasoning_content else None, |
|
|
tool_calls=executed_tool_calls if executed_tool_calls else None, |
|
|
), |
|
|
logprobs=None, |
|
|
finish_reason=finishReason, |
|
|
) |
|
|
], |
|
|
) |
|
|
|
|
|
try: |
|
|
if request.state_name: |
|
|
STATE_STORE[(request.model, request.state_name)] = { |
|
|
'state': model_state, |
|
|
'model_tokens': model_tokens, |
|
|
} |
|
|
if getattr(CONFIG, 'STATE_STORE_SAVE_ON_UPDATE', False): |
|
|
try: |
|
|
_save_state_store_to_disk(force=True) |
|
|
except Exception: |
|
|
pass |
|
|
except Exception: |
|
|
pass |
|
|
|
|
|
return response |
|
|
|
|
|
async def chatResponseStream( |
|
|
request: ChatCompletionRequest, |
|
|
model_state: Any, |
|
|
completionId: str, |
|
|
enableReasoning: bool, |
|
|
): |
|
|
createTimestamp = int(time.time()) |
|
|
|
|
|
raw_prompt = request.prompt.strip() if request.prompt is not None else cleanMessages(request.messages or [], False) |
|
|
detection = detect_tools_and_reasoning(raw_prompt) |
|
|
|
|
|
web_search_enabled = ( |
|
|
True |
|
|
if (request.enable_web_search is not None and request.enable_web_search) |
|
|
else ( |
|
|
request.web_search |
|
|
or (request.auto_web_search if request.auto_web_search is not None else CONFIG.AUTO_ENABLE_WEB_SEARCH and detection.get('need_web_search')) |
|
|
) |
|
|
) |
|
|
|
|
|
if request.enable_tools is not None: |
|
|
tools_enabled = bool(request.enable_tools) |
|
|
else: |
|
|
auto_tools_flag = request.auto_tools if request.auto_tools is not None else CONFIG.AUTO_ENABLE_TOOLS |
|
|
tools_enabled = bool(request.tools) or CONFIG.ENABLE_TOOLS_BY_DEFAULT or (auto_tools_flag and (detection.get('need_calc') or detection.get('need_web_search'))) |
|
|
|
|
|
reasoning_enabled = bool( |
|
|
True |
|
|
if (request.enable_reasoning is not None and request.enable_reasoning) |
|
|
else ( |
|
|
bool(enableReasoning) or bool(request.auto_reasoning if request.auto_reasoning is not None else (CONFIG.AUTO_ENABLE_REASONING and bool(detection.get('need_reasoning')))) |
|
|
) |
|
|
) |
|
|
enableReasoning = reasoning_enabled |
|
|
try: |
|
|
ms_cfg = MODEL_STORAGE.get(request.model) |
|
|
if ms_cfg and ms_cfg.MODEL_CONFIG and hasattr(ms_cfg.MODEL_CONFIG, 'ALLOW_REASONING') and not ms_cfg.MODEL_CONFIG.ALLOW_REASONING: |
|
|
enableReasoning = False |
|
|
except Exception: |
|
|
pass |
|
|
if request.enable_file_tool is not None: |
|
|
file_tool_enabled = bool(request.enable_file_tool) |
|
|
else: |
|
|
auto_file_flag = request.auto_file_tool if request.auto_file_tool is not None else CONFIG.AUTO_ENABLE_TOOLS |
|
|
file_tool_enabled = bool((request.file_ids and len(request.file_ids) > 0) or (auto_file_flag and request.file_ids)) |
|
|
if not getattr(CONFIG, 'ALLOW_FILE_TOOL_BY_DEFAULT', True) and request.enable_file_tool is None: |
|
|
file_tool_enabled = False |
|
|
try: |
|
|
if request.sampler and getattr(request.sampler, 'ALLOW_FILE_TOOL', None) is not None: |
|
|
file_tool_enabled = bool(request.sampler.ALLOW_FILE_TOOL) |
|
|
elif hasattr(request, 'sampler_allow_file_tool') and request.sampler_allow_file_tool is not None: |
|
|
file_tool_enabled = bool(request.sampler_allow_file_tool) |
|
|
else: |
|
|
ms2 = MODEL_STORAGE.get(request.model) |
|
|
if ms2 and ms2.MODEL_CONFIG: |
|
|
if hasattr(ms2.MODEL_CONFIG, 'DEFAULT_SAMPLER') and getattr(ms2.MODEL_CONFIG.DEFAULT_SAMPLER, 'ALLOW_FILE_TOOL', None) is not None: |
|
|
file_tool_enabled = bool(ms2.MODEL_CONFIG.DEFAULT_SAMPLER.ALLOW_FILE_TOOL) |
|
|
elif hasattr(ms2.MODEL_CONFIG, 'ALLOW_FILE_TOOL') and not ms2.MODEL_CONFIG.ALLOW_FILE_TOOL: |
|
|
file_tool_enabled = False |
|
|
except Exception: |
|
|
pass |
|
|
prompt = raw_prompt if request.prompt is not None else f"{cleanMessages(request.messages or [], enableReasoning)}\n\nAssistant:{' <think' if enableReasoning else ''}" |
|
|
|
|
|
if tools_enabled and not request.tools: |
|
|
if detection.get('detected_tools'): |
|
|
request.tools = detection.get('detected_tools') |
|
|
executed_tool_calls = [] |
|
|
if request.tools: |
|
|
try: |
|
|
for tool in request.tools: |
|
|
name = tool.get('name') |
|
|
args = tool.get('args', {}) |
|
|
if name == 'web_search': |
|
|
from utils import web_search |
|
|
|
|
|
search_q = args.get('query') or (request.prompt if request.prompt else cleanMessages(request.messages or [])) |
|
|
search_top_k = int(args.get('top_k') or request.search_top_k or 3) |
|
|
search_str = web_search(search_q, search_top_k) |
|
|
if search_str: |
|
|
search_res_struct = {"action": "web_search", "result": str(search_str), "metadata": {"query": search_q, "top_k": search_top_k, "confidence": 0.9}} |
|
|
executed_tool_calls.append({"name": "web_search", "args": {"query": search_q, "top_k": search_top_k}, "result": search_res_struct}) |
|
|
prompt = (f"WebSearchResults:\n{search_res_struct.get('result')}\n\n" + prompt) |
|
|
elif name == 'calc' or name == 'calculator': |
|
|
from utils import calc |
|
|
|
|
|
expr = args.get('expression') |
|
|
if expr: |
|
|
calc_res = calc(expr) |
|
|
calc_res_struct = {"action": "calc", "result": str(calc_res), "metadata": {"expression": expr, "confidence": 0.98}} |
|
|
executed_tool_calls.append({"name": "calc", "args": {"expression": expr}, "result": calc_res_struct}) |
|
|
prompt = (f"CalcResult:{expr} = {calc_res_struct.get('result')}\n\n" + prompt) |
|
|
elif name == 'universal': |
|
|
try: |
|
|
res = universal_tool(args or {"query": raw_prompt}, allow_web_search=bool(web_search_enabled), allow_tools=bool(tools_enabled), allow_file_tool=bool(file_tool_enabled)) |
|
|
if isinstance(res, dict): |
|
|
result_text = res.get('result') if res.get('result') is not None else '' |
|
|
else: |
|
|
result_text = str(res) |
|
|
executed_tool_calls.append({"name": "universal", "args": args, "result": res}) |
|
|
prompt = (f"ToolResults:\n{result_text}\n\n" + prompt) |
|
|
except Exception as e: |
|
|
logger.info(f"Universal tool execution error: {e}") |
|
|
else: |
|
|
logger.info(f"Unsupported tool requested: {name}") |
|
|
except Exception as e: |
|
|
logger.info(f"Tool processing error: {e}") |
|
|
elif request.web_search or web_search_enabled: |
|
|
try: |
|
|
from utils import web_search |
|
|
|
|
|
search_q = request.prompt if request.prompt else cleanMessages(request.messages or []) |
|
|
search_res = web_search(search_q, int(request.search_top_k or 3)) |
|
|
if search_res: |
|
|
search_res_struct = {"action": "web_search", "result": str(search_res), "metadata": {"query": search_q, "top_k": int(request.search_top_k or 3), "confidence": 0.9}} |
|
|
executed_tool_calls.append({"name": "web_search", "args": {"query": search_q, "top_k": int(request.search_top_k or 3)}, "result": search_res_struct}) |
|
|
prompt = f"WebSearchResults:\n{search_res_struct.get('result')}\n\n" + prompt |
|
|
except Exception: |
|
|
pass |
|
|
|
|
|
logger.info(f"[REQ] {completionId} - context\n```{prompt}```") |
|
|
|
|
|
if request.state_name: |
|
|
state_key = (request.model, request.state_name) |
|
|
if state_key in STATE_STORE: |
|
|
stored = STATE_STORE[state_key] |
|
|
model_state = stored.get('state', None) |
|
|
model_tokens = stored.get('model_tokens', [0]) |
|
|
if model_state is None: |
|
|
out, model_state = _recompute_out_and_state_from_tokens(request.model, model_tokens) |
|
|
else: |
|
|
out, _ = _recompute_out_and_state_from_tokens(request.model, model_tokens[-CONFIG.CHUNK_LEN :]) |
|
|
else: |
|
|
out, model_tokens, model_state = await runPrefill(request, prompt, [0], model_state) |
|
|
else: |
|
|
out, model_tokens, model_state = await runPrefill(request, prompt, [0], model_state) |
|
|
|
|
|
prefillTime = time.time() |
|
|
promptTokenCount = len(model_tokens) |
|
|
|
|
|
completionTokenCount = 0 |
|
|
finishReason = None |
|
|
model_initiated_tool_calls = 0 |
|
|
MODEL_MAX_TOOL_CALLS = 3 |
|
|
|
|
|
response = ChatCompletionChunk( |
|
|
id=completionId, |
|
|
created=createTimestamp, |
|
|
model=request.model, |
|
|
usage=( |
|
|
Usage( |
|
|
prompt_tokens=promptTokenCount, |
|
|
completion_tokens=completionTokenCount, |
|
|
total_tokens=promptTokenCount + completionTokenCount, |
|
|
prompt_tokens_details=PromptTokensDetails(cached_tokens=0), |
|
|
) |
|
|
if request.include_usage |
|
|
else None |
|
|
), |
|
|
choices=[ |
|
|
ChatCompletionChoice( |
|
|
index=0, |
|
|
delta=ChatCompletionMessage( |
|
|
role="Assistant", |
|
|
content="", |
|
|
reasoning_content="" if enableReasoning else None, |
|
|
tool_calls=None, |
|
|
), |
|
|
logprobs=None, |
|
|
finish_reason=finishReason, |
|
|
) |
|
|
], |
|
|
) |
|
|
if response.choices and response.choices[0].delta is None: |
|
|
response.choices[0].delta = ChatCompletionMessage(role="Assistant", content="", reasoning_content=None, tool_calls=None) |
|
|
r_dict = response.model_dump() |
|
|
r_dict['state_name'] = request.state_name |
|
|
if executed_tool_calls: |
|
|
r_dict['tool_calls'] = executed_tool_calls |
|
|
try: |
|
|
if r_dict.get('choices') and len(r_dict['choices']) > 0 and r_dict['choices'][0].get('delta') is not None: |
|
|
r_dict['choices'][0]['delta']['tool_calls'] = executed_tool_calls |
|
|
except Exception: |
|
|
pass |
|
|
yield f"data: {r_dict}\n\n" |
|
|
|
|
|
buffer = [] |
|
|
|
|
|
if enableReasoning: |
|
|
buffer.append("<think") |
|
|
|
|
|
streamConfig = { |
|
|
"isChecking": False, |
|
|
"fullTextCursor": 0, |
|
|
"in_think": False, |
|
|
"cacheStr": "", |
|
|
} |
|
|
|
|
|
for chunk in generate( |
|
|
request, |
|
|
out, |
|
|
model_tokens, |
|
|
model_state, |
|
|
max_tokens=( |
|
|
64000 |
|
|
if "max_tokens" not in request.model_fields_set and enableReasoning |
|
|
else (request.max_tokens or 2048) |
|
|
), |
|
|
): |
|
|
completionTokenCount += 1 |
|
|
chunkContent: str = chunk["content"] |
|
|
buffer.append(chunkContent) |
|
|
|
|
|
fullText = "".join(buffer) |
|
|
|
|
|
if chunk["finish_reason"]: |
|
|
finishReason = chunk["finish_reason"] |
|
|
|
|
|
response = ChatCompletionChunk( |
|
|
id=completionId, |
|
|
created=createTimestamp, |
|
|
model=request.model, |
|
|
usage=( |
|
|
Usage( |
|
|
prompt_tokens=promptTokenCount, |
|
|
completion_tokens=completionTokenCount, |
|
|
total_tokens=promptTokenCount + completionTokenCount, |
|
|
prompt_tokens_details=PromptTokensDetails(cached_tokens=0), |
|
|
) |
|
|
if request.include_usage |
|
|
else None |
|
|
), |
|
|
choices=[ |
|
|
ChatCompletionChoice( |
|
|
index=0, |
|
|
delta=ChatCompletionMessage( |
|
|
role="Assistant", |
|
|
content=None, |
|
|
reasoning_content=None, |
|
|
tool_calls=None, |
|
|
), |
|
|
logprobs=None, |
|
|
finish_reason=finishReason, |
|
|
) |
|
|
], |
|
|
) |
|
|
if response.choices and response.choices[0].delta is None: |
|
|
response.choices[0].delta = ChatCompletionMessage(role="Assistant", content="", reasoning_content=None, tool_calls=None) |
|
|
|
|
|
markStart = fullText.find("<", streamConfig["fullTextCursor"]) |
|
|
if not streamConfig["isChecking"] and markStart != -1: |
|
|
streamConfig["isChecking"] = True |
|
|
|
|
|
if streamConfig["in_think"]: |
|
|
delta = response.choices[0].delta |
|
|
if delta is None: |
|
|
delta = ChatCompletionMessage(role="Assistant", content="", reasoning_content=None, tool_calls=None) |
|
|
response.choices[0].delta = delta |
|
|
delta.reasoning_content = fullText[streamConfig["fullTextCursor"] : markStart] |
|
|
else: |
|
|
delta = response.choices[0].delta |
|
|
if delta is None: |
|
|
delta = ChatCompletionMessage(role="Assistant", content="", reasoning_content=None, tool_calls=None) |
|
|
response.choices[0].delta = delta |
|
|
delta.content = fullText[streamConfig["fullTextCursor"] : markStart] |
|
|
|
|
|
streamConfig["cacheStr"] = "" |
|
|
streamConfig["fullTextCursor"] = markStart |
|
|
|
|
|
if streamConfig["isChecking"]: |
|
|
streamConfig["cacheStr"] = fullText[streamConfig["fullTextCursor"] :] |
|
|
else: |
|
|
if streamConfig["in_think"]: |
|
|
delta = response.choices[0].delta |
|
|
if delta is None: |
|
|
delta = ChatCompletionMessage(role="Assistant", content="", reasoning_content=None, tool_calls=None) |
|
|
response.choices[0].delta = delta |
|
|
delta.reasoning_content = chunkContent |
|
|
else: |
|
|
delta = response.choices[0].delta |
|
|
if delta is None: |
|
|
delta = ChatCompletionMessage(role="Assistant", content="", reasoning_content=None, tool_calls=None) |
|
|
response.choices[0].delta = delta |
|
|
delta.content = chunkContent |
|
|
streamConfig["fullTextCursor"] = len(fullText) |
|
|
|
|
|
markEnd = fullText.find(">", streamConfig["fullTextCursor"]) |
|
|
if (streamConfig["isChecking"] and markEnd != -1) or finishReason != None: |
|
|
streamConfig["isChecking"] = False |
|
|
|
|
|
if ( |
|
|
not streamConfig["in_think"] |
|
|
and streamConfig["cacheStr"].find("<think>") != -1 |
|
|
): |
|
|
streamConfig["in_think"] = True |
|
|
|
|
|
delta = response.choices[0].delta |
|
|
if delta is None: |
|
|
delta = ChatCompletionMessage(role="Assistant", content="", reasoning_content=None, tool_calls=None) |
|
|
response.choices[0].delta = delta |
|
|
delta.reasoning_content = ( |
|
|
delta.reasoning_content |
|
|
if delta.reasoning_content != None |
|
|
else "" + streamConfig["cacheStr"].replace("<think>", "") |
|
|
) |
|
|
|
|
|
elif ( |
|
|
streamConfig["in_think"] |
|
|
and streamConfig["cacheStr"].find("</think>") != -1 |
|
|
): |
|
|
streamConfig["in_think"] = False |
|
|
|
|
|
delta = response.choices[0].delta |
|
|
if delta is None: |
|
|
delta = ChatCompletionMessage(role="Assistant", content="", reasoning_content=None, tool_calls=None) |
|
|
response.choices[0].delta = delta |
|
|
delta.content = ( |
|
|
delta.content |
|
|
if delta.content != None |
|
|
else "" + streamConfig["cacheStr"].replace("</think>", "") |
|
|
) |
|
|
else: |
|
|
if streamConfig["in_think"]: |
|
|
delta = response.choices[0].delta |
|
|
if delta is None: |
|
|
delta = ChatCompletionMessage(role="Assistant", content="", reasoning_content=None, tool_calls=None) |
|
|
response.choices[0].delta = delta |
|
|
delta.reasoning_content = ( |
|
|
delta.reasoning_content |
|
|
if delta.reasoning_content != None |
|
|
else "" + streamConfig["cacheStr"] |
|
|
) |
|
|
else: |
|
|
delta = response.choices[0].delta |
|
|
if delta is None: |
|
|
delta = ChatCompletionMessage(role="Assistant", content="", reasoning_content=None, tool_calls=None) |
|
|
response.choices[0].delta = delta |
|
|
delta.content = ( |
|
|
delta.content |
|
|
if delta.content != None |
|
|
else "" + streamConfig["cacheStr"] |
|
|
) |
|
|
streamConfig["fullTextCursor"] = len(fullText) |
|
|
|
|
|
delta = response.choices[0].delta |
|
|
if delta is None: |
|
|
delta = ChatCompletionMessage(role="Assistant", content="", reasoning_content=None, tool_calls=None) |
|
|
response.choices[0].delta = delta |
|
|
if delta.content != None or delta.reasoning_content != None: |
|
|
try: |
|
|
if request.state_name: |
|
|
STATE_STORE[(request.model, request.state_name)] = { |
|
|
'state': model_state, |
|
|
'model_tokens': model_tokens, |
|
|
} |
|
|
if getattr(CONFIG, 'STATE_STORE_SAVE_ON_UPDATE', False): |
|
|
try: |
|
|
_save_state_store_to_disk(force=True) |
|
|
except Exception: |
|
|
pass |
|
|
except Exception: |
|
|
pass |
|
|
if model_initiated_tool_calls < MODEL_MAX_TOOL_CALLS: |
|
|
m = TOOL_CALL_RE.search(fullText) |
|
|
if m: |
|
|
try: |
|
|
payload_raw = m.group(1) |
|
|
import json |
|
|
|
|
|
payload = json.loads(payload_raw) |
|
|
tool_name = payload.get('name') |
|
|
tool_args = payload.get('args', {}) |
|
|
tool_res = None |
|
|
if tool_name == 'web_search': |
|
|
from utils import web_search |
|
|
|
|
|
q = tool_args.get('query') or (request.prompt if request.prompt else cleanMessages(request.messages or [])) |
|
|
k = int(tool_args.get('top_k') or request.search_top_k or 3) |
|
|
tool_res = web_search(q, k) |
|
|
elif tool_name in ('calc', 'calculator'): |
|
|
from utils import calc |
|
|
|
|
|
expr = tool_args.get('expression') |
|
|
if expr: |
|
|
tool_res = calc(expr) |
|
|
else: |
|
|
try: |
|
|
tool_res = universal_tool({'query': tool_args.get('query') or payload.get('query') or ''}, allow_web_search=bool(web_search_enabled), allow_tools=bool(tools_enabled), allow_file_tool=bool(file_tool_enabled)) |
|
|
except Exception: |
|
|
tool_res = None |
|
|
|
|
|
if tool_res: |
|
|
if not isinstance(tool_res, dict): |
|
|
if tool_name in ('calc', 'calculator'): |
|
|
tool_res_struct = {"action": "calc", "result": str(tool_res), "metadata": {"expression": tool_args.get('expression'), "confidence": 0.98}} |
|
|
elif tool_name == 'web_search': |
|
|
tool_res_struct = {"action": "web_search", "result": str(tool_res), "metadata": {"query": tool_args.get('query'), "top_k": tool_args.get('top_k') or request.search_top_k or 3, "confidence": 0.9}} |
|
|
else: |
|
|
tool_res_struct = {"action": tool_name, "result": str(tool_res), "metadata": {"confidence": 0.6}} |
|
|
else: |
|
|
tool_res_struct = tool_res |
|
|
exec_entry = {"name": tool_name, "args": tool_args, "result": tool_res_struct, 'initiated_by_model': True} |
|
|
executed_tool_calls.append(exec_entry) |
|
|
delta_text = f"ToolResults:\n{tool_res_struct.get('result')}\n\n" |
|
|
prompt = delta_text + prompt |
|
|
fullText = TOOL_CALL_RE.sub('', fullText) |
|
|
buffer = [fullText] |
|
|
out, model_tokens, model_state = await runPrefill(request, delta_text, model_tokens, model_state) |
|
|
try: |
|
|
meta_resp = ChatCompletionChunk( |
|
|
id=completionId, |
|
|
created=createTimestamp, |
|
|
model=request.model, |
|
|
usage=( |
|
|
Usage( |
|
|
prompt_tokens=promptTokenCount, |
|
|
completion_tokens=completionTokenCount, |
|
|
total_tokens=promptTokenCount + completionTokenCount, |
|
|
prompt_tokens_details=PromptTokensDetails(cached_tokens=0), |
|
|
) |
|
|
if request.include_usage |
|
|
else None |
|
|
), |
|
|
choices=[ |
|
|
ChatCompletionChoice( |
|
|
index=0, |
|
|
delta=ChatCompletionMessage(role="Assistant", content=None, reasoning_content=None, tool_calls=executed_tool_calls), |
|
|
logprobs=None, |
|
|
finish_reason=None, |
|
|
) |
|
|
], |
|
|
) |
|
|
yield f"data: {meta_resp.model_dump_json()}\n\n" |
|
|
except Exception: |
|
|
pass |
|
|
model_initiated_tool_calls += 1 |
|
|
should_restart = True |
|
|
break |
|
|
except Exception as e: |
|
|
logger.info(f"Model-initiated tool handling error: {e}") |
|
|
yield f"data: {response.model_dump_json()}\n\n" |
|
|
for stop_words in request.stop or []: |
|
|
if stop_words in ''.join(buffer): |
|
|
finishReason = f"stop:words:{stop_words}" |
|
|
return |
|
|
|
|
|
await asyncio.sleep(0) |
|
|
|
|
|
del streamConfig |
|
|
else: |
|
|
should_restart = True |
|
|
while should_restart: |
|
|
should_restart = False |
|
|
gen = generate(request, out, model_tokens, model_state) |
|
|
for chunk in gen: |
|
|
completionTokenCount += 1 |
|
|
buffer.append(chunk["content"]) |
|
|
|
|
|
if chunk["finish_reason"]: |
|
|
finishReason = chunk["finish_reason"] |
|
|
|
|
|
try: |
|
|
if request.state_name: |
|
|
STATE_STORE[(request.model, request.state_name)] = { |
|
|
'state': model_state, |
|
|
'model_tokens': model_tokens, |
|
|
} |
|
|
if getattr(CONFIG, 'STATE_STORE_SAVE_ON_UPDATE', False): |
|
|
try: |
|
|
_save_state_store_to_disk(force=True) |
|
|
except Exception: |
|
|
pass |
|
|
except Exception: |
|
|
pass |
|
|
|
|
|
if model_initiated_tool_calls < MODEL_MAX_TOOL_CALLS: |
|
|
fullText = ''.join(buffer) |
|
|
m = TOOL_CALL_RE.search(fullText) |
|
|
if m: |
|
|
try: |
|
|
payload_raw = m.group(1) |
|
|
import json |
|
|
|
|
|
payload = json.loads(payload_raw) |
|
|
tool_name = payload.get('name') |
|
|
tool_args = payload.get('args', {}) |
|
|
tool_res = None |
|
|
if tool_name == 'web_search': |
|
|
from utils import web_search |
|
|
|
|
|
q = tool_args.get('query') or (request.prompt if request.prompt else cleanMessages(request.messages or [])) |
|
|
k = int(tool_args.get('top_k') or request.search_top_k or 3) |
|
|
tool_res = web_search(q, k) |
|
|
elif tool_name in ('calc', 'calculator'): |
|
|
from utils import calc |
|
|
|
|
|
expr = tool_args.get('expression') |
|
|
if expr: |
|
|
tool_res = calc(expr) |
|
|
else: |
|
|
try: |
|
|
tool_res = universal_tool({'query': tool_args.get('query') or payload.get('query') or ''}, allow_web_search=bool(web_search_enabled), allow_tools=bool(tools_enabled), allow_file_tool=bool(file_tool_enabled)) |
|
|
except Exception: |
|
|
tool_res = None |
|
|
|
|
|
if tool_res: |
|
|
if not isinstance(tool_res, dict): |
|
|
if tool_name in ('calc', 'calculator'): |
|
|
tool_res_struct = {"action": "calc", "result": str(tool_res), "metadata": {"expression": tool_args.get('expression'), "confidence": 0.98}} |
|
|
elif tool_name == 'web_search': |
|
|
tool_res_struct = {"action": "web_search", "result": str(tool_res), "metadata": {"query": tool_args.get('query'), "top_k": tool_args.get('top_k') or request.search_top_k or 3, "confidence": 0.9}} |
|
|
else: |
|
|
tool_res_struct = {"action": tool_name, "result": str(tool_res), "metadata": {"confidence": 0.6}} |
|
|
else: |
|
|
tool_res_struct = tool_res |
|
|
exec_entry = {"name": tool_name, "args": tool_args, "result": tool_res_struct, 'initiated_by_model': True} |
|
|
executed_tool_calls.append(exec_entry) |
|
|
delta_text = f"ToolResults:\n{tool_res_struct.get('result')}\n\n" |
|
|
prompt = delta_text + prompt |
|
|
fullText = TOOL_CALL_RE.sub('', fullText) |
|
|
buffer = [fullText] |
|
|
out, model_tokens, model_state = await runPrefill(request, delta_text, model_tokens, model_state) |
|
|
try: |
|
|
meta_resp = ChatCompletionChunk( |
|
|
id=completionId, |
|
|
created=createTimestamp, |
|
|
model=request.model, |
|
|
usage=( |
|
|
Usage( |
|
|
prompt_tokens=promptTokenCount, |
|
|
completion_tokens=completionTokenCount, |
|
|
total_tokens=promptTokenCount + completionTokenCount, |
|
|
prompt_tokens_details=PromptTokensDetails(cached_tokens=0), |
|
|
) |
|
|
if request.include_usage |
|
|
else None |
|
|
), |
|
|
choices=[ |
|
|
ChatCompletionChoice( |
|
|
index=0, |
|
|
delta=ChatCompletionMessage(role="Assistant", content=None, reasoning_content=None, tool_calls=executed_tool_calls), |
|
|
logprobs=None, |
|
|
finish_reason=None, |
|
|
) |
|
|
], |
|
|
) |
|
|
yield f"data: {meta_resp.model_dump_json()}\n\n" |
|
|
except Exception: |
|
|
pass |
|
|
model_initiated_tool_calls += 1 |
|
|
should_restart = True |
|
|
break |
|
|
except Exception as e: |
|
|
logger.info(f"Model-initiated tool handling error: {e}") |
|
|
|
|
|
response = ChatCompletionChunk( |
|
|
id=completionId, |
|
|
created=createTimestamp, |
|
|
model=request.model, |
|
|
usage=( |
|
|
Usage( |
|
|
prompt_tokens=promptTokenCount, |
|
|
completion_tokens=completionTokenCount, |
|
|
total_tokens=promptTokenCount + completionTokenCount, |
|
|
prompt_tokens_details=PromptTokensDetails(cached_tokens=0), |
|
|
) |
|
|
if request.include_usage |
|
|
else None |
|
|
), |
|
|
choices=[ |
|
|
ChatCompletionChoice( |
|
|
index=0, |
|
|
delta=ChatCompletionMessage(role="Assistant", content=chunk["content"], reasoning_content=None, tool_calls=None), |
|
|
logprobs=None, |
|
|
finish_reason=finishReason, |
|
|
) |
|
|
], |
|
|
) |
|
|
yield f"data: {response.model_dump_json()}\n\n" |
|
|
await asyncio.sleep(0) |
|
|
|
|
|
genenrateTime = time.time() |
|
|
|
|
|
responseLog = { |
|
|
"content": "".join(buffer), |
|
|
"finish": finishReason, |
|
|
"prefill_len": promptTokenCount, |
|
|
"prefill_tps": round(promptTokenCount / (prefillTime - createTimestamp), 2), |
|
|
"gen_len": completionTokenCount, |
|
|
"gen_tps": round(completionTokenCount / (genenrateTime - prefillTime), 2), |
|
|
} |
|
|
logger.info(f"[RES] {completionId} - {responseLog}") |
|
|
if request.messages is None: |
|
|
request.messages = [] |
|
|
request.messages.append(ChatMessage(role="Assistant", content=responseLog["content"])) |
|
|
log( |
|
|
{ |
|
|
**request.model_dump(), |
|
|
**responseLog, |
|
|
"completionId": completionId, |
|
|
"machineLabel": os.environ.get("MACHINE_LABEL"), |
|
|
} |
|
|
) |
|
|
|
|
|
del buffer |
|
|
|
|
|
yield "data: [DONE]\n\n" |
|
|
|
|
|
@app.post("/api/v1/chat/completions") |
|
|
async def chat_completions(request: ChatCompletionRequest): |
|
|
completionId = str(next(CompletionIdGenerator)) |
|
|
logger.info(f"[REQ] {completionId} - {request.model_dump()}") |
|
|
|
|
|
if "rwkv-latest" in request.model: |
|
|
if DEFALUT_MODEL_NAME == None: |
|
|
raise HTTPException(404, "DEFALUT_MODEL_NAME not set") |
|
|
ms_def = MODEL_STORAGE.get(DEFALUT_MODEL_NAME) |
|
|
if not ms_def or not ms_def.MODEL_CONFIG: |
|
|
raise HTTPException(500, "Default sampler config missing for default model") |
|
|
defaultSamplerConfig = ms_def.MODEL_CONFIG.DEFAULT_SAMPLER |
|
|
request.model = DEFALUT_MODEL_NAME |
|
|
elif request.model in MODEL_STORAGE: |
|
|
ms_sel = MODEL_STORAGE.get(request.model) |
|
|
if not ms_sel or not ms_sel.MODEL_CONFIG: |
|
|
raise HTTPException(500, f"Default sampler config missing for model {request.model}") |
|
|
defaultSamplerConfig = ms_sel.MODEL_CONFIG.DEFAULT_SAMPLER |
|
|
else: |
|
|
raise HTTPException(404, f"Can not find `{request.model}`") |
|
|
|
|
|
enableReasoning = request.enable_reasoning if request.enable_reasoning is not None else False |
|
|
|
|
|
async def chatResponseStreamDisconnect(): |
|
|
logGPUState() |
|
|
|
|
|
model_state = None |
|
|
model_tokens_for_resume = [0] |
|
|
state_name = request.state_name |
|
|
if state_name is None: |
|
|
state_name = str(uuid.uuid4()) |
|
|
request.state_name = state_name |
|
|
state_key = (request.model, state_name) |
|
|
if state_key in STATE_STORE: |
|
|
stored = STATE_STORE[state_key] |
|
|
model_state = stored.get('state', None) |
|
|
model_tokens_for_resume = stored.get('model_tokens', [0]) |
|
|
request_dict = request.model_dump() |
|
|
|
|
|
sampler_overrides = request_dict.get('sampler') or {} |
|
|
for k, v in defaultSamplerConfig.model_dump().items(): |
|
|
if sampler_overrides and k in sampler_overrides and sampler_overrides.get(k) is not None: |
|
|
request_dict[k] = sampler_overrides.get(k) |
|
|
continue |
|
|
if k in request_dict and request_dict[k] is None: |
|
|
request_dict[k] = v |
|
|
realRequest = ChatCompletionRequest(**request_dict) |
|
|
if realRequest.stream is None: |
|
|
realRequest.stream = CONFIG.DEFAULT_STREAM |
|
|
|
|
|
logger.info(f"[REQ] {completionId} - Real - {request.model_dump()}") |
|
|
|
|
|
if realRequest.stream: |
|
|
r = StreamingResponse( |
|
|
chatResponseStream(realRequest, model_state, completionId, enableReasoning), |
|
|
media_type="text/event-stream", |
|
|
background=BackgroundTask(chatResponseStreamDisconnect), |
|
|
) |
|
|
else: |
|
|
r = await chatResponse(realRequest, model_state, completionId, enableReasoning) |
|
|
try: |
|
|
import json |
|
|
|
|
|
if isinstance(r, ChatCompletion): |
|
|
d = r.model_dump() |
|
|
d['state_name'] = state_name |
|
|
return d |
|
|
except Exception: |
|
|
pass |
|
|
|
|
|
return r |
|
|
|
|
|
logger.info("Static frontend mount removed for Python-only deploy; use API endpoints for integration") |
|
|
|
|
|
@app.get('/api/v1/models') |
|
|
def list_models(): |
|
|
out = [] |
|
|
root_defaults = { |
|
|
'ALLOW_FILE_TOOL_BY_DEFAULT': getattr(CONFIG, 'ALLOW_FILE_TOOL_BY_DEFAULT', True), |
|
|
'ENABLE_WEB_SEARCH_BY_DEFAULT': getattr(CONFIG, 'ENABLE_WEB_SEARCH_BY_DEFAULT', True), |
|
|
'ENABLE_REASONING_BY_DEFAULT': getattr(CONFIG, 'ENABLE_REASONING_BY_DEFAULT', True), |
|
|
'SHOW_WEB_SEARCH_BUTTON_BY_DEFAULT': getattr(CONFIG, 'SHOW_WEB_SEARCH_BUTTON_BY_DEFAULT', True), |
|
|
'SHOW_FILE_UPLOAD_BUTTON_BY_DEFAULT': getattr(CONFIG, 'SHOW_FILE_UPLOAD_BUTTON_BY_DEFAULT', True), |
|
|
'SHOW_REASONING_TOGGLE_BY_DEFAULT': getattr(CONFIG, 'SHOW_REASONING_TOGGLE_BY_DEFAULT', True), |
|
|
'UPLOAD_URL': '/api/v1/files', |
|
|
} |
|
|
for m in CONFIG.MODELS: |
|
|
out.append( |
|
|
{ |
|
|
'SERVICE_NAME': m.SERVICE_NAME, |
|
|
'DEFAULT_CHAT': m.DEFAULT_CHAT, |
|
|
'DEFAULT_REASONING': m.DEFAULT_REASONING, |
|
|
'ALLOW_WEB_SEARCH': getattr(m, 'ALLOW_WEB_SEARCH', True), |
|
|
'ALLOW_TOOLS': getattr(m, 'ALLOW_TOOLS', True), |
|
|
'ALLOW_REASONING': getattr(m, 'ALLOW_REASONING', True), |
|
|
'ALLOW_FILE_TOOL': getattr(m, 'ALLOW_FILE_TOOL', True), |
|
|
'SHOW_WEB_SEARCH_BUTTON': getattr(m, 'SHOW_WEB_SEARCH_BUTTON', True), |
|
|
'SHOW_FILE_UPLOAD_BUTTON': getattr(m, 'SHOW_FILE_UPLOAD_BUTTON', True), |
|
|
'SHOW_REASONING_TOGGLE': getattr(m, 'SHOW_REASONING_TOGGLE', True), |
|
|
'DEFAULT_SAMPLER': m.DEFAULT_SAMPLER.model_dump() if hasattr(m, 'DEFAULT_SAMPLER') else None, |
|
|
'UPLOAD_URL': '/api/v1/files', |
|
|
'UPLOAD_ALLOWED_BY_DEFAULT': getattr(CONFIG, 'ALLOW_FILE_TOOL_BY_DEFAULT', True), |
|
|
} |
|
|
) |
|
|
return {'root_defaults': root_defaults, 'models': out} |
|
|
|
|
|
@app.post('/api/v1/files', response_model=FileUploadResponse) |
|
|
async def upload_file(file: UploadFile = File(...), model: Optional[str] = None): |
|
|
try: |
|
|
if not getattr(CONFIG, 'ALLOW_FILE_TOOL_BY_DEFAULT', True): |
|
|
raise HTTPException(403, 'File uploads are disabled by server configuration') |
|
|
if model: |
|
|
if model not in MODEL_STORAGE: |
|
|
raise HTTPException(404, f"Model {model} not found") |
|
|
ms = MODEL_STORAGE[model] |
|
|
if ms and ms.MODEL_CONFIG and not getattr(ms.MODEL_CONFIG, 'ALLOW_FILE_TOOL', True): |
|
|
raise HTTPException(403, f"Model {model} does not allow file uploads") |
|
|
from utils import save_bytes_to_upload |
|
|
|
|
|
content = await file.read() |
|
|
fname = file.filename if getattr(file, 'filename', None) else 'uploaded_file' |
|
|
meta = save_bytes_to_upload(fname, content) |
|
|
if meta.get('error'): |
|
|
raise HTTPException(500, f"Could not save file: {meta.get('error')}") |
|
|
UPLOADED_FILES[meta['file_id']] = meta |
|
|
return FileUploadResponse(success=True, file=UploadedFile(**meta)) |
|
|
except Exception as e: |
|
|
raise HTTPException(500, str(e)) |
|
|
|
|
|
@app.get('/api/v1/files') |
|
|
def list_files(): |
|
|
return [UploadedFile(**v).model_dump() for v in UPLOADED_FILES.values()] |
|
|
|
|
|
@app.get('/api/v1/files/{file_id}') |
|
|
def get_file(file_id: str, download: bool = False): |
|
|
if file_id not in UPLOADED_FILES: |
|
|
raise HTTPException(404, 'File not found') |
|
|
meta = UPLOADED_FILES[file_id] |
|
|
if download: |
|
|
try: |
|
|
with open(meta['path'], 'rb') as f: |
|
|
return StreamingResponse(f, media_type='application/octet-stream') |
|
|
except Exception as e: |
|
|
raise HTTPException(500, str(e)) |
|
|
return UploadedFile(**meta) |
|
|
|
|
|
@app.delete('/api/v1/files/{file_id}') |
|
|
def delete_file(file_id: str): |
|
|
if file_id not in UPLOADED_FILES: |
|
|
raise HTTPException(404, 'File not found') |
|
|
meta = UPLOADED_FILES.pop(file_id) |
|
|
try: |
|
|
if os.path.exists(meta['path']): |
|
|
os.remove(meta['path']) |
|
|
except Exception: |
|
|
pass |
|
|
return {'success': True} |
|
|
|
|
|
if __name__ == "__main__": |
|
|
import uvicorn |
|
|
|
|
|
host = CONFIG.HOST or "127.0.0.1" |
|
|
port = CONFIG.PORT or 7860 |
|
|
uvicorn.run(app, host=host, port=port) |
