mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2025-09-07 13:09:50 +00:00
281 lines
11 KiB
Python
281 lines
11 KiB
Python
from datetime import datetime
|
|
from http.client import NOT_IMPLEMENTED
|
|
import json
|
|
from time import time
|
|
from uuid import uuid4
|
|
from typing import List, Optional
|
|
|
|
from fastapi import APIRouter, Request
|
|
from pydantic import BaseModel, Field
|
|
|
|
from ktransformers.server.config.config import Config
|
|
from ktransformers.server.utils.create_interface import get_interface
|
|
from ktransformers.server.schemas.assistants.streaming import check_link_response
|
|
from ktransformers.server.backend.base import BackendInterfaceBase
|
|
|
|
from ktransformers.server.schemas.endpoints.chat import RawUsage
|
|
|
|
router = APIRouter(prefix='/api')
|
|
|
|
# https://github.com/ollama/ollama/blob/main/docs/api.md#generate-a-completion
|
|
class OllamaGenerateCompletionRequest(BaseModel):
|
|
model: str = Field(..., description="The model name, which is required.")
|
|
prompt: Optional[str] = Field(
|
|
None, description="The prompt to generate a response for.")
|
|
images: Optional[List[str]] = Field(
|
|
None, description="A list of base64-encoded images for multimodal models such as llava.")
|
|
# Advanced parameters
|
|
format: Optional[str] = Field(
|
|
None, description="The format to return a response in, accepted value is json.")
|
|
options: Optional[dict] = Field(
|
|
None, description="Additional model parameters as listed in the documentation.")
|
|
system: Optional[str] = Field(
|
|
None, description="System message to override what is defined in the Modelfile.")
|
|
template: Optional[str] = Field(
|
|
None, description="The prompt template to use, overriding what is defined in the Modelfile.")
|
|
context: Optional[str] = Field(
|
|
None, description="The context parameter from a previous request to keep a short conversational memory.")
|
|
stream: Optional[bool] = Field(
|
|
None, description="If false, the response will be returned as a single response object.")
|
|
raw: Optional[bool] = Field(
|
|
None, description="If true, no formatting will be applied to the prompt.")
|
|
keep_alive: Optional[str] = Field(
|
|
"5m", description="Controls how long the model will stay loaded into memory following the request.")
|
|
|
|
class OllamaGenerationStreamResponse(BaseModel):
|
|
model: str
|
|
created_at: str
|
|
response: str
|
|
done: bool = Field(...)
|
|
|
|
class OllamaGenerationResponse(BaseModel):
|
|
model: str
|
|
created_at: str
|
|
response: str
|
|
done: bool
|
|
|
|
@router.post("/generate", tags=['ollama'])
|
|
async def generate(request: Request, input: OllamaGenerateCompletionRequest):
|
|
id = str(uuid4())
|
|
interface: BackendInterfaceBase = get_interface()
|
|
print(f'COMPLETION INPUT:----\n{input.prompt}\n----')
|
|
config = Config()
|
|
|
|
if input.stream:
|
|
async def inner():
|
|
async for res in interface.inference(input.prompt, id):
|
|
if isinstance(res, RawUsage):
|
|
raw_usage = res
|
|
else:
|
|
token, finish_reason = res
|
|
d = OllamaGenerationStreamResponse(
|
|
model=config.model_name,
|
|
created_at=str(datetime.now()),
|
|
response=token,
|
|
done=False
|
|
)
|
|
yield d.model_dump_json() + '\n'
|
|
d = OllamaGenerationStreamResponse(
|
|
model=config.model_name,
|
|
created_at=str(datetime.now()),
|
|
response='',
|
|
done=True
|
|
)
|
|
yield d.model_dump_json() + '\n'
|
|
return check_link_response(request, inner())
|
|
else:
|
|
complete_response = ""
|
|
async for res in interface.inference(input.prompt, id):
|
|
if isinstance(res, RawUsage):
|
|
raw_usage = res
|
|
else:
|
|
token, finish_reason = res
|
|
complete_response += token
|
|
response = OllamaGenerationResponse(
|
|
model=config.model_name,
|
|
created_at=str(datetime.now()),
|
|
response=complete_response,
|
|
done=True
|
|
)
|
|
return response
|
|
|
|
# https://github.com/ollama/ollama/blob/main/docs/api.md#generate-a-chat-completion
|
|
class OllamaChatCompletionMessage(BaseModel):
|
|
role: str
|
|
content: str
|
|
|
|
class OllamaChatCompletionRequest(BaseModel):
|
|
model: str = Field(..., description="The model name, which is required.")
|
|
messages: List[OllamaChatCompletionMessage] = Field(
|
|
..., description="A list of messages to generate a response for.")
|
|
stream: bool = Field(True, description="If true, the response will be streamed.")
|
|
|
|
class OllamaChatCompletionStreamResponse(BaseModel):
|
|
model: str
|
|
created_at: str
|
|
message: dict
|
|
done: bool = Field(...)
|
|
total_duration: Optional[int] = Field(None, description="Total time spent in nanoseconds")
|
|
load_duration: Optional[int] = Field(None, description="Time spent loading model in nanoseconds")
|
|
prompt_eval_count: Optional[int] = Field(None, description="Number of tokens in prompt")
|
|
prompt_eval_duration: Optional[int] = Field(None, description="Time spent evaluating prompt in nanoseconds")
|
|
eval_count: Optional[int] = Field(None, description="Number of tokens generated")
|
|
eval_duration: Optional[int] = Field(None, description="Time spent generating response in nanoseconds")
|
|
|
|
class OllamaChatCompletionResponse(BaseModel):
|
|
model: str
|
|
created_at: str
|
|
message: dict
|
|
done: bool
|
|
total_duration: Optional[int] = Field(None, description="Total time spent in nanoseconds")
|
|
load_duration: Optional[int] = Field(None, description="Time spent loading model in nanoseconds")
|
|
prompt_eval_count: Optional[int] = Field(None, description="Number of tokens in prompt")
|
|
prompt_eval_duration: Optional[int] = Field(None, description="Time spent evaluating prompt in nanoseconds")
|
|
eval_count: Optional[int] = Field(None, description="Number of tokens generated")
|
|
eval_duration: Optional[int] = Field(None, description="Time spent generating response in nanoseconds")
|
|
|
|
@router.post("/chat", tags=['ollama'])
|
|
async def chat(request: Request, input: OllamaChatCompletionRequest):
|
|
id = str(uuid4())
|
|
interface: BackendInterfaceBase = get_interface()
|
|
config = Config()
|
|
|
|
# 将消息转换为提示字符串
|
|
prompt = ""
|
|
for msg in input.messages:
|
|
prompt += f"{msg.role}: {msg.content}\n"
|
|
prompt += "assistant:"
|
|
|
|
if input.stream:
|
|
async def inner():
|
|
start_time = time() # 记录开始时间（秒）
|
|
eval_count = 0 # 统计生成的 token 数量
|
|
tokens = []
|
|
|
|
async for res in interface.inference(prompt, id):
|
|
if isinstance(res, RawUsage):
|
|
raw_usage = res
|
|
else:
|
|
token, finish_reason = res
|
|
d = OllamaChatCompletionStreamResponse(
|
|
model=config.model_name,
|
|
created_at=str(datetime.now()),
|
|
message={"role": "assistant", "content": token},
|
|
done=False
|
|
)
|
|
yield d.model_dump_json() + '\n'
|
|
# 计算性能数据
|
|
end_time = time()
|
|
total_duration = int((end_time - start_time) * 1_000_000_000) # 转换为纳秒
|
|
prompt_eval_count = len(prompt.split()) # 简单估算提示词数量
|
|
eval_duration = total_duration # 假设全部时间用于生成（简化）
|
|
prompt_eval_duration = 0 # 假设无单独提示评估时间
|
|
load_duration = 0 # 假设加载时间未知
|
|
|
|
d = OllamaChatCompletionStreamResponse(
|
|
model=config.model_name,
|
|
created_at=str(datetime.now()),
|
|
message={},
|
|
done=True,
|
|
total_duration=total_duration,
|
|
load_duration=load_duration,
|
|
prompt_eval_count=prompt_eval_count,
|
|
prompt_eval_duration=prompt_eval_duration,
|
|
eval_count=eval_count,
|
|
eval_duration=eval_duration
|
|
)
|
|
yield d.model_dump_json() + '\n'
|
|
return check_link_response(request, inner())
|
|
else:
|
|
start_time = time()
|
|
complete_response = ""
|
|
eval_count = 0
|
|
|
|
async for res in interface.inference(prompt, id):
|
|
if isinstance(res, RawUsage):
|
|
raw_usage = res
|
|
else:
|
|
token, finish_reason = res
|
|
complete_response += token
|
|
eval_count += 1
|
|
|
|
end_time = time()
|
|
total_duration = int((end_time - start_time) * 1_000_000_000)
|
|
prompt_eval_count = len(prompt.split())
|
|
eval_duration = total_duration
|
|
prompt_eval_duration = 0
|
|
load_duration = 0
|
|
|
|
response = OllamaChatCompletionResponse(
|
|
model=config.model_name,
|
|
created_at=str(datetime.now()),
|
|
message={"role": "assistant", "content": complete_response},
|
|
done=True,
|
|
total_duration=total_duration,
|
|
load_duration=load_duration,
|
|
prompt_eval_count=prompt_eval_count,
|
|
prompt_eval_duration=prompt_eval_duration,
|
|
eval_count=eval_count,
|
|
eval_duration=eval_duration
|
|
)
|
|
return response
|
|
|
|
# https://github.com/ollama/ollama/blob/main/docs/api.md#list-local-models
|
|
class OllamaModel(BaseModel):
|
|
name: str
|
|
modified_at: str
|
|
size: int
|
|
# TODO: fill the rest correctly
|
|
|
|
# mock ollama
|
|
@router.get("/tags", tags=['ollama'])
|
|
async def tags():
|
|
config = Config()
|
|
# TODO: fill this correctly, although it does not effect Tabby
|
|
return {"models": [OllamaModel(name=config.model_name, modified_at="123", size=123)]}
|
|
|
|
class OllamaModelInfo(BaseModel):
|
|
# TODO: fill this correctly
|
|
pass
|
|
|
|
class OllamaShowRequest(BaseModel):
|
|
name: str = Field(..., description="Name of the model to show")
|
|
verbose: Optional[bool] = Field(
|
|
None, description="If set to true, returns full data for verbose response fields")
|
|
|
|
class OllamaShowDetial(BaseModel):
|
|
parent_model: str
|
|
format: str
|
|
family: str
|
|
families: List[str]
|
|
parameter_size: str
|
|
quantization_level: str
|
|
|
|
class OllamaShowResponse(BaseModel):
|
|
modelfile: str
|
|
parameters: str
|
|
template: str
|
|
details: OllamaShowDetial
|
|
model_info: OllamaModelInfo
|
|
|
|
class Config:
|
|
protected_namespaces = ()
|
|
|
|
@router.post("/show", tags=['ollama'])
|
|
async def show(request: Request, input: OllamaShowRequest):
|
|
config = Config()
|
|
# TODO: Add more info in config to return, although it does not effect Tabby
|
|
return OllamaShowResponse(
|
|
modelfile="# Modelfile generated by ...",
|
|
parameters=" ",
|
|
template=" ",
|
|
details=OllamaShowDetial(
|
|
parent_model=" ",
|
|
format="gguf",
|
|
family=" ",
|
|
families=[" "],
|
|
parameter_size=" ",
|
|
quantization_level=" "
|
|
),
|
|
model_info=OllamaModelInfo()
|
|
)
|