635 lines
26 KiB
Python
635 lines
26 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
import uuid
|
|
from contextlib import asynccontextmanager
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Literal
|
|
|
|
from fastapi import Depends, FastAPI, Header, HTTPException, Request
|
|
from fastapi.responses import PlainTextResponse, StreamingResponse
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
from .config import AppConfig, ModelConfig, load_config
|
|
from .manager import ModelManager
|
|
from .modelfile import model_config_from_modelfile
|
|
from .types import BackendError, ChatParams
|
|
|
|
|
|
class Message(BaseModel):
|
|
model_config = ConfigDict(extra="allow")
|
|
|
|
role: Literal["system", "user", "assistant", "tool", "function", "developer"]
|
|
content: Any
|
|
|
|
def payload(self) -> dict[str, Any]:
|
|
return self.model_dump(exclude_none=True)
|
|
|
|
|
|
class ChatCompletionRequest(BaseModel):
|
|
model_config = ConfigDict(extra="allow")
|
|
|
|
model: str | None = None
|
|
messages: list[Message]
|
|
stream: bool = False
|
|
max_tokens: int | None = Field(default=None, ge=1, le=131072)
|
|
max_completion_tokens: int | None = Field(default=None, ge=1, le=131072)
|
|
temperature: float | None = Field(default=None, ge=0, le=2)
|
|
top_p: float | None = Field(default=None, ge=0, le=1)
|
|
top_k: int | None = Field(default=None, ge=0, le=100000)
|
|
min_p: float | None = Field(default=None, ge=0, le=1)
|
|
repeat_penalty: float | None = Field(default=None, ge=0)
|
|
seed: int | None = None
|
|
stop: str | list[str] | None = None
|
|
enable_thinking: bool | None = None
|
|
response_format: str | dict[str, Any] | None = None
|
|
tools: list[dict[str, Any]] = Field(default_factory=list)
|
|
tool_choice: Any = None
|
|
keep_alive: str | int | float | None = None
|
|
extra_body: dict[str, Any] = Field(default_factory=dict)
|
|
|
|
def params(self, config: AppConfig) -> ChatParams:
|
|
del config
|
|
default = ChatParams()
|
|
stop = [self.stop] if isinstance(self.stop, str) else (self.stop or [])
|
|
extra = dict(self.extra_body)
|
|
known = {
|
|
"model", "messages", "stream", "max_tokens", "max_completion_tokens", "temperature",
|
|
"top_p", "top_k", "min_p", "repeat_penalty", "seed", "stop", "enable_thinking",
|
|
"response_format", "tools", "tool_choice", "keep_alive", "extra_body",
|
|
}
|
|
extra.update({key: value for key, value in (self.model_extra or {}).items() if key not in known})
|
|
return ChatParams(
|
|
max_tokens=self.max_completion_tokens or self.max_tokens or default.max_tokens,
|
|
temperature=default.temperature if self.temperature is None else self.temperature,
|
|
top_p=default.top_p if self.top_p is None else self.top_p,
|
|
top_k=default.top_k if self.top_k is None else self.top_k,
|
|
min_p=default.min_p if self.min_p is None else self.min_p,
|
|
repeat_penalty=default.repeat_penalty if self.repeat_penalty is None else self.repeat_penalty,
|
|
seed=self.seed,
|
|
stop=stop,
|
|
enable_thinking=self.enable_thinking,
|
|
response_format=self.response_format,
|
|
tools=self.tools,
|
|
tool_choice=self.tool_choice,
|
|
extra=extra,
|
|
)
|
|
|
|
|
|
class OllamaChatRequest(BaseModel):
|
|
model_config = ConfigDict(extra="allow")
|
|
|
|
model: str
|
|
messages: list[Message]
|
|
stream: bool = True
|
|
format: str | dict[str, Any] | None = None
|
|
options: dict[str, Any] = Field(default_factory=dict)
|
|
keep_alive: str | int | float | None = "5m"
|
|
think: bool | None = None
|
|
tools: list[dict[str, Any]] = Field(default_factory=list)
|
|
|
|
|
|
class OllamaGenerateRequest(BaseModel):
|
|
model_config = ConfigDict(extra="allow")
|
|
|
|
model: str
|
|
prompt: str = ""
|
|
suffix: str | None = None
|
|
system: str | None = None
|
|
template: str | None = None
|
|
context: list[int] | None = None
|
|
stream: bool = True
|
|
raw: bool = False
|
|
format: str | dict[str, Any] | None = None
|
|
options: dict[str, Any] = Field(default_factory=dict)
|
|
keep_alive: str | int | float | None = "5m"
|
|
think: bool | None = None
|
|
|
|
|
|
class OllamaEmbedRequest(BaseModel):
|
|
model: str
|
|
input: str | list[str]
|
|
keep_alive: str | int | float | None = "5m"
|
|
|
|
|
|
class OpenAIEmbeddingRequest(BaseModel):
|
|
model: str
|
|
input: str | list[str]
|
|
encoding_format: str | None = None
|
|
|
|
|
|
class RegisterModelRequest(BaseModel):
|
|
id: str
|
|
path: str | None = None
|
|
kind: str = "auto"
|
|
enabled: bool = True
|
|
description: str = ""
|
|
cloud_profile: str | None = None
|
|
remote_model: str | None = None
|
|
template: str | None = None
|
|
system: str | None = None
|
|
messages: list[dict[str, Any]] = Field(default_factory=list)
|
|
parameters: dict[str, Any] = Field(default_factory=dict)
|
|
options: dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class CreateModelRequest(BaseModel):
|
|
model: str
|
|
modelfile: str
|
|
stream: bool = False
|
|
|
|
|
|
class CopyModelRequest(BaseModel):
|
|
source: str
|
|
destination: str
|
|
|
|
|
|
class PullModelRequest(BaseModel):
|
|
model: str
|
|
path: str | None = None
|
|
kind: str = "auto"
|
|
stream: bool = True
|
|
repo_id: str | None = None
|
|
revision: str | None = None
|
|
destination: str | None = None
|
|
allow_patterns: list[str] | None = None
|
|
alias: str | None = None
|
|
|
|
|
|
def _content_to_text(content: Any) -> str:
|
|
if isinstance(content, str):
|
|
return content
|
|
if isinstance(content, list):
|
|
return "".join(
|
|
str(item.get("text", "")) if isinstance(item, dict) and item.get("type") == "text" else "[image]"
|
|
for item in content
|
|
)
|
|
return str(content)
|
|
|
|
|
|
def _sse(data: dict[str, Any]) -> str:
|
|
return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
|
|
|
|
|
|
def _iso_now() -> str:
|
|
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
|
|
|
|
|
def _ollama_params(options: dict[str, Any], fmt: str | dict[str, Any] | None, think: bool | None, tools: list[dict[str, Any]]) -> ChatParams:
|
|
stop = options.get("stop", [])
|
|
if isinstance(stop, str):
|
|
stop = [stop]
|
|
standard = {
|
|
"num_predict": "max_tokens",
|
|
"temperature": "temperature",
|
|
"top_p": "top_p",
|
|
"top_k": "top_k",
|
|
"min_p": "min_p",
|
|
"repeat_penalty": "repeat_penalty",
|
|
"seed": "seed",
|
|
}
|
|
values: dict[str, Any] = {key: options[source] for source, key in standard.items() if source in options}
|
|
extra = {key: value for key, value in options.items() if key not in standard and key not in {"stop"}}
|
|
return ChatParams(
|
|
max_tokens=int(values.get("max_tokens", 512)),
|
|
temperature=float(values.get("temperature", 0.7)),
|
|
top_p=float(values.get("top_p", 0.95)),
|
|
top_k=int(values.get("top_k", 40)),
|
|
min_p=float(values.get("min_p", 0.05)),
|
|
repeat_penalty=float(values.get("repeat_penalty", 1.05)),
|
|
seed=values.get("seed"),
|
|
stop=stop,
|
|
enable_thinking=think,
|
|
response_format=fmt,
|
|
tools=tools,
|
|
extra=extra,
|
|
)
|
|
|
|
|
|
def _usage_from_metadata(metadata: dict[str, Any]) -> dict[str, int]:
|
|
usage = metadata.get("usage") or {}
|
|
if usage:
|
|
prompt = int(usage.get("prompt_tokens", usage.get("prompt_eval_count", 0)) or 0)
|
|
completion = int(usage.get("completion_tokens", usage.get("eval_count", 0)) or 0)
|
|
return {
|
|
"prompt_tokens": prompt,
|
|
"completion_tokens": completion,
|
|
"total_tokens": int(usage.get("total_tokens", prompt + completion)),
|
|
}
|
|
prompt = int(metadata.get("prompt_eval_count", metadata.get("prompt_tokens", 0)) or 0)
|
|
completion = int(metadata.get("eval_count", metadata.get("completion_tokens", 0)) or 0)
|
|
return {"prompt_tokens": prompt, "completion_tokens": completion, "total_tokens": prompt + completion}
|
|
|
|
|
|
def _ollama_usage(metadata: dict[str, Any]) -> dict[str, int]:
|
|
usage = metadata.get("usage") or {}
|
|
prompt = int(usage.get("prompt_eval_count", usage.get("prompt_tokens", metadata.get("prompt_eval_count", 0))) or 0)
|
|
completion = int(usage.get("eval_count", usage.get("completion_tokens", metadata.get("eval_count", 0))) or 0)
|
|
return {"prompt_eval_count": prompt, "eval_count": completion, "eval_duration": 0, "prompt_eval_duration": 0}
|
|
|
|
|
|
def _ollama_runtime_fields(manager: ModelManager, model_id: str) -> dict[str, int]:
|
|
state = manager.states.get(model_id)
|
|
if not state:
|
|
return {}
|
|
return {
|
|
"total_duration": int(state.last_duration_seconds * 1_000_000_000),
|
|
"load_duration": int(state.last_load_duration_seconds * 1_000_000_000),
|
|
}
|
|
|
|
|
|
def create_app(config_path: str | Path | None = None) -> FastAPI:
|
|
resolved = Path(config_path or r"E:\llm-backend\config\config.json")
|
|
config = load_config(resolved)
|
|
manager = ModelManager(config, resolved)
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_: FastAPI):
|
|
if config.runtime.auto_load and config.runtime.default_model:
|
|
try:
|
|
await manager.load(config.runtime.default_model)
|
|
except Exception:
|
|
pass
|
|
yield
|
|
await manager.shutdown()
|
|
|
|
app = FastAPI(title="LocalPilot", version="0.2.0", lifespan=lifespan)
|
|
app.state.config = config
|
|
app.state.manager = manager
|
|
|
|
async def authorize(request: Request, authorization: str | None = Header(default=None)) -> None:
|
|
del request
|
|
expected = config.runtime.api_key
|
|
if expected and (authorization or "").removeprefix("Bearer ") != expected:
|
|
raise HTTPException(status_code=401, detail="invalid API key")
|
|
|
|
@app.get("/health")
|
|
async def health() -> dict[str, Any]:
|
|
return {
|
|
"status": "ok",
|
|
"version": "0.2.0",
|
|
"active_model": manager.active_model_id,
|
|
"running": manager.running_models(),
|
|
}
|
|
|
|
@app.get("/v1/models")
|
|
async def models(_: None = Depends(authorize)) -> dict[str, Any]:
|
|
data = []
|
|
for item in manager.list_models():
|
|
data.append({**item, "created": int(item.get("modified_at") or time.time())})
|
|
return {"object": "list", "data": data}
|
|
|
|
@app.get("/api/version")
|
|
async def api_version() -> dict[str, str]:
|
|
return {"version": "0.2.0"}
|
|
|
|
@app.get("/api/tags")
|
|
async def tags(_: None = Depends(authorize)) -> dict[str, Any]:
|
|
models_data = []
|
|
for item in manager.list_models():
|
|
details = manager.show_model(item["id"])["details"]
|
|
models_data.append({
|
|
"name": item["name"],
|
|
"model": item["model"],
|
|
"modified_at": datetime.fromtimestamp(item["modified_at"], timezone.utc).isoformat() if item["modified_at"] else _iso_now(),
|
|
"size": item["size"],
|
|
"digest": item["digest"],
|
|
"details": details,
|
|
"capabilities": item["capabilities"],
|
|
})
|
|
return {"models": models_data}
|
|
|
|
@app.get("/api/ps")
|
|
async def ps(_: None = Depends(authorize)) -> dict[str, Any]:
|
|
models_data = []
|
|
for item in manager.running_models():
|
|
state = manager.states[item["id"]]
|
|
expires_at = None
|
|
if state.keep_alive_seconds != float("inf"):
|
|
expires_at = datetime.fromtimestamp(
|
|
state.last_used + state.keep_alive_seconds, timezone.utc
|
|
).isoformat()
|
|
models_data.append({
|
|
"name": item["name"],
|
|
"model": item["model"],
|
|
"size": item["size"],
|
|
"size_vram": item.get("vram_allocated_bytes", 0),
|
|
"expires_at": expires_at,
|
|
"active_requests": item["active_requests"],
|
|
"kind": item["kind"],
|
|
})
|
|
return {"models": models_data}
|
|
|
|
@app.get("/metrics", response_class=PlainTextResponse)
|
|
async def metrics(_: None = Depends(authorize)) -> str:
|
|
records = manager.list_models()
|
|
loaded = sum(1 for item in records if item["loaded"])
|
|
active = sum(int(item["active_requests"]) for item in records)
|
|
total = sum(int(item["total_requests"]) for item in records)
|
|
lines = [
|
|
"# TYPE localpilot_models_loaded gauge",
|
|
f"localpilot_models_loaded {loaded}",
|
|
"# TYPE localpilot_active_requests gauge",
|
|
f"localpilot_active_requests {active}",
|
|
"# TYPE localpilot_requests_total counter",
|
|
f"localpilot_requests_total {total}",
|
|
]
|
|
return "\n".join(lines) + "\n"
|
|
|
|
@app.post("/api/show")
|
|
async def show(body: dict[str, Any], _: None = Depends(authorize)) -> dict[str, Any]:
|
|
model_id = body.get("model") or body.get("name")
|
|
if not model_id:
|
|
raise HTTPException(status_code=400, detail="missing model/name")
|
|
try:
|
|
return manager.show_model(model_id)
|
|
except KeyError as exc:
|
|
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
|
|
|
@app.post("/api/models/register")
|
|
async def register(body: RegisterModelRequest, _: None = Depends(authorize)) -> dict[str, Any]:
|
|
try:
|
|
spec = manager.register_model(ModelConfig(**body.model_dump()))
|
|
return {"status": "registered", "model": spec.model_dump(exclude_none=True)}
|
|
except BackendError as exc:
|
|
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
|
|
|
@app.delete("/api/models/{model_id}")
|
|
async def delete_model(model_id: str, _: None = Depends(authorize)) -> dict[str, Any]:
|
|
try:
|
|
await manager.remove_model(model_id)
|
|
return {"status": "deleted", "model": model_id, "files_removed": False}
|
|
except BackendError as exc:
|
|
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
|
|
|
@app.post("/api/create")
|
|
async def create_model(body: CreateModelRequest, _: None = Depends(authorize)) -> Any:
|
|
try:
|
|
modelfile_text = body.modelfile
|
|
modelfile_path = Path(modelfile_text)
|
|
if modelfile_path.is_file():
|
|
modelfile_text = modelfile_path.read_text(encoding="utf-8")
|
|
spec = model_config_from_modelfile(body.model, modelfile_text)
|
|
manager.register_model(spec)
|
|
except (BackendError, ValueError) as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
if body.stream:
|
|
async def create_stream():
|
|
yield json.dumps({"status": "success", "model": body.model}, ensure_ascii=False) + "\n"
|
|
|
|
return StreamingResponse(create_stream(), media_type="application/x-ndjson")
|
|
return {"status": "success", "model": body.model}
|
|
|
|
@app.post("/api/copy")
|
|
async def copy_model(body: CopyModelRequest, _: None = Depends(authorize)) -> dict[str, Any]:
|
|
try:
|
|
spec = manager.copy_model(body.source, body.destination)
|
|
return {"status": "success", "model": spec.id}
|
|
except (BackendError, KeyError) as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
|
|
@app.post("/api/pull")
|
|
async def pull_model(body: PullModelRequest, _: None = Depends(authorize)) -> Any:
|
|
try:
|
|
if body.path:
|
|
spec = manager.register_model(ModelConfig(id=body.model, kind=body.kind, path=body.path))
|
|
elif body.repo_id or "/" in body.model:
|
|
repo_id = body.repo_id or body.model
|
|
model_id = body.alias or repo_id.replace("/", "-")
|
|
spec = await manager.pull_huggingface(
|
|
model_id=model_id,
|
|
repo_id=repo_id,
|
|
destination=body.destination,
|
|
revision=body.revision,
|
|
allow_patterns=body.allow_patterns,
|
|
)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=501,
|
|
detail="请提供 path 导入本地模型，或使用 repo_id/repo 触发 Hugging Face 下载。",
|
|
)
|
|
except BackendError as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
if body.stream:
|
|
async def pull_stream():
|
|
yield json.dumps({"status": "success", "model": spec.id, "path": spec.path}, ensure_ascii=False) + "\n"
|
|
|
|
return StreamingResponse(pull_stream(), media_type="application/x-ndjson")
|
|
return {"status": "success", "model": spec.id, "path": spec.path}
|
|
|
|
@app.post("/api/models/{model_id}/load")
|
|
async def load_model(model_id: str, _: None = Depends(authorize)) -> dict[str, Any]:
|
|
try:
|
|
backend = await manager.load(model_id)
|
|
return backend.info()
|
|
except (BackendError, KeyError) as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
|
|
@app.post("/api/models/{model_id}/unload")
|
|
async def unload_model(model_id: str, _: None = Depends(authorize)) -> dict[str, Any]:
|
|
await manager.unload(model_id)
|
|
return {"status": "unloaded", "model": model_id}
|
|
|
|
@app.post("/v1/chat/completions")
|
|
async def chat_completions(body: ChatCompletionRequest, _: None = Depends(authorize)) -> Any:
|
|
messages = [item.payload() for item in body.messages]
|
|
params = body.params(config)
|
|
request_id = f"chatcmpl-localpilot-{uuid.uuid4().hex[:12]}"
|
|
model_id = body.model or config.runtime.default_model or manager.active_model_id or "localpilot"
|
|
try:
|
|
if body.stream:
|
|
async def event_stream():
|
|
yield _sse({
|
|
"id": request_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": int(time.time()),
|
|
"model": model_id,
|
|
"choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
|
|
})
|
|
async for chunk in manager.chat(model_id, messages, params, stream=True, keep_alive=body.keep_alive):
|
|
delta = dict(chunk.delta)
|
|
if chunk.text and "content" not in delta:
|
|
delta["content"] = chunk.text
|
|
if delta:
|
|
yield _sse({
|
|
"id": request_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": int(time.time()),
|
|
"model": model_id,
|
|
"choices": [{"index": 0, "delta": delta, "finish_reason": None}],
|
|
})
|
|
if chunk.done:
|
|
yield _sse({
|
|
"id": request_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": int(time.time()),
|
|
"model": model_id,
|
|
"choices": [{"index": 0, "delta": {}, "finish_reason": chunk.finish_reason or "stop"}],
|
|
})
|
|
yield "data: [DONE]\n\n"
|
|
|
|
return StreamingResponse(event_stream(), media_type="text/event-stream")
|
|
|
|
parts: list[str] = []
|
|
finish_reason = "stop"
|
|
metadata: dict[str, Any] = {}
|
|
async for chunk in manager.chat(model_id, messages, params, stream=False, keep_alive=body.keep_alive):
|
|
parts.append(chunk.text)
|
|
finish_reason = chunk.finish_reason or finish_reason
|
|
metadata = chunk.metadata or metadata
|
|
message: dict[str, Any] = {"role": "assistant", "content": "".join(parts)}
|
|
if metadata:
|
|
original_message = ((metadata.get("choices") or [{}])[0].get("message") or {})
|
|
for key in ("tool_calls", "function_call", "refusal"):
|
|
if key in original_message:
|
|
message[key] = original_message[key]
|
|
return {
|
|
"id": request_id,
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": model_id,
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": message,
|
|
"finish_reason": finish_reason,
|
|
}],
|
|
"usage": _usage_from_metadata(metadata),
|
|
}
|
|
except (BackendError, KeyError) as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
|
|
@app.post("/v1/completions")
|
|
async def completions(body: dict[str, Any], _: None = Depends(authorize)) -> Any:
|
|
prompt = body.get("prompt", "")
|
|
if isinstance(prompt, list):
|
|
prompt = "\n".join(str(item) for item in prompt)
|
|
payload = ChatCompletionRequest(
|
|
model=body.get("model"),
|
|
messages=[Message(role="user", content=str(prompt))],
|
|
stream=bool(body.get("stream", False)),
|
|
max_tokens=body.get("max_tokens"),
|
|
temperature=body.get("temperature"),
|
|
top_p=body.get("top_p"),
|
|
stop=body.get("stop"),
|
|
keep_alive=body.get("keep_alive"),
|
|
)
|
|
return await chat_completions(payload)
|
|
|
|
@app.post("/v1/embeddings")
|
|
async def openai_embeddings(body: OpenAIEmbeddingRequest, _: None = Depends(authorize)) -> dict[str, Any]:
|
|
inputs = [body.input] if isinstance(body.input, str) else body.input
|
|
try:
|
|
vectors = await manager.embed(body.model, inputs)
|
|
except (BackendError, KeyError) as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
return {
|
|
"object": "list",
|
|
"data": [{"object": "embedding", "index": index, "embedding": vector} for index, vector in enumerate(vectors)],
|
|
"model": body.model,
|
|
"usage": {"prompt_tokens": 0, "total_tokens": 0},
|
|
}
|
|
|
|
@app.post("/api/embed")
|
|
async def ollama_embed(body: OllamaEmbedRequest, _: None = Depends(authorize)) -> dict[str, Any]:
|
|
inputs = [body.input] if isinstance(body.input, str) else body.input
|
|
try:
|
|
vectors = await manager.embed(body.model, inputs, body.keep_alive)
|
|
except (BackendError, KeyError) as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
return {"model": body.model, "embeddings": vectors}
|
|
|
|
@app.post("/api/embeddings")
|
|
async def ollama_embeddings_legacy(body: OllamaEmbedRequest, _: None = Depends(authorize)) -> dict[str, Any]:
|
|
return await ollama_embed(body)
|
|
|
|
@app.post("/api/chat")
|
|
async def ollama_chat(body: OllamaChatRequest, _: None = Depends(authorize)) -> Any:
|
|
messages = [item.payload() for item in body.messages]
|
|
params = _ollama_params(body.options, body.format, body.think, body.tools)
|
|
try:
|
|
if body.stream:
|
|
async def event_stream():
|
|
answer: list[str] = []
|
|
async for chunk in manager.chat(body.model, messages, params, stream=True, keep_alive=body.keep_alive):
|
|
answer.append(chunk.text)
|
|
message = {"role": "assistant", "content": chunk.text, **chunk.delta}
|
|
yield json.dumps({
|
|
"model": body.model,
|
|
"created_at": _iso_now(),
|
|
"message": message,
|
|
"done": False,
|
|
}, ensure_ascii=False) + "\n"
|
|
yield json.dumps({
|
|
"model": body.model,
|
|
"created_at": _iso_now(),
|
|
"message": {"role": "assistant", "content": ""},
|
|
"done": True,
|
|
"done_reason": "stop",
|
|
**_ollama_runtime_fields(manager, body.model),
|
|
}, ensure_ascii=False) + "\n"
|
|
|
|
return StreamingResponse(event_stream(), media_type="application/x-ndjson")
|
|
|
|
parts: list[str] = []
|
|
metadata: dict[str, Any] = {}
|
|
async for chunk in manager.chat(body.model, messages, params, stream=False, keep_alive=body.keep_alive):
|
|
parts.append(chunk.text)
|
|
metadata = chunk.metadata or metadata
|
|
return {
|
|
"model": body.model,
|
|
"created_at": _iso_now(),
|
|
"message": {"role": "assistant", "content": "".join(parts)},
|
|
"done": True,
|
|
"done_reason": "stop",
|
|
**_ollama_runtime_fields(manager, body.model),
|
|
**_ollama_usage(metadata),
|
|
}
|
|
except (BackendError, KeyError) as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
|
|
@app.post("/api/generate")
|
|
async def ollama_generate(body: OllamaGenerateRequest, _: None = Depends(authorize)) -> Any:
|
|
del body.suffix, body.template, body.context, body.raw
|
|
messages: list[dict[str, Any]] = []
|
|
if body.system:
|
|
messages.append({"role": "system", "content": body.system})
|
|
messages.append({"role": "user", "content": body.prompt})
|
|
params = _ollama_params(body.options, body.format, body.think, [])
|
|
try:
|
|
if body.stream:
|
|
async def event_stream():
|
|
async for chunk in manager.chat(body.model, messages, params, stream=True, keep_alive=body.keep_alive):
|
|
yield json.dumps({
|
|
"model": body.model,
|
|
"created_at": _iso_now(),
|
|
"response": chunk.text,
|
|
"done": False,
|
|
}, ensure_ascii=False) + "\n"
|
|
yield json.dumps({
|
|
"model": body.model,
|
|
"created_at": _iso_now(),
|
|
"response": "",
|
|
"done": True,
|
|
"done_reason": "stop",
|
|
**_ollama_runtime_fields(manager, body.model),
|
|
}, ensure_ascii=False) + "\n"
|
|
|
|
return StreamingResponse(event_stream(), media_type="application/x-ndjson")
|
|
|
|
parts: list[str] = []
|
|
async for chunk in manager.chat(body.model, messages, params, stream=False, keep_alive=body.keep_alive):
|
|
parts.append(chunk.text)
|
|
return {
|
|
"model": body.model,
|
|
"created_at": _iso_now(),
|
|
"response": "".join(parts),
|
|
"done": True,
|
|
**_ollama_runtime_fields(manager, body.model),
|
|
}
|
|
except (BackendError, KeyError) as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
|
|
return app
|