hymenjj/llama-cpp-python-prebuilt
0
1from __future__ import annotations2 3import os4import json5import typing6import contextlib7 8from anyio import Lock9from functools import partial10from typing import List, Optional, Union, Dict11 12import llama_cpp13 14import anyio15from anyio.streams.memory import MemoryObjectSendStream16from starlette.concurrency import run_in_threadpool, iterate_in_threadpool17from fastapi import Depends, FastAPI, APIRouter, Request, HTTPException, status, Body18from fastapi.middleware import Middleware19from fastapi.middleware.cors import CORSMiddleware20from fastapi.security import HTTPBearer21from sse_starlette.sse import EventSourceResponse22from starlette_context.plugins import RequestIdPlugin # type: ignore23from starlette_context.middleware import RawContextMiddleware24 25from llama_cpp.server.model import (26 LlamaProxy,27)28from llama_cpp.server.settings import (29 ConfigFileSettings,30 Settings,31 ModelSettings,32 ServerSettings,33)34from llama_cpp.server.types import (35 CreateCompletionRequest,36 CreateEmbeddingRequest,37 CreateChatCompletionRequest,38 ModelList,39 TokenizeInputRequest,40 TokenizeInputResponse,41 TokenizeInputCountResponse,42 DetokenizeInputRequest,43 DetokenizeInputResponse,44)45from llama_cpp.server.errors import RouteErrorHandler46 47 48router = APIRouter(route_class=RouteErrorHandler)49 50_server_settings: Optional[ServerSettings] = None51 52 53def set_server_settings(server_settings: ServerSettings):54 global _server_settings55 _server_settings = server_settings56 57 58def get_server_settings():59 yield _server_settings60 61 62_llama_proxy: Optional[LlamaProxy] = None63 64llama_outer_lock = Lock()65llama_inner_lock = Lock()66 67 68def set_llama_proxy(model_settings: List[ModelSettings]):69 global _llama_proxy70 _llama_proxy = LlamaProxy(models=model_settings)71 72 73async def get_llama_proxy():74 # NOTE: This double lock allows the currently streaming llama model to75 # check if any other requests are pending in the same thread and cancel76 # the stream if so.77 await llama_outer_lock.acquire()78 release_outer_lock = True79 try:80 await llama_inner_lock.acquire()81 try:82 llama_outer_lock.release()83 release_outer_lock = False84 yield _llama_proxy85 finally:86 llama_inner_lock.release()87 finally:88 if release_outer_lock:89 llama_outer_lock.release()90 91 92_ping_message_factory: typing.Optional[typing.Callable[[], bytes]] = None93 94 95def set_ping_message_factory(factory: typing.Callable[[], bytes]):96 global _ping_message_factory97 _ping_message_factory = factory98 99 100def create_app(101 settings: Settings | None = None,102 server_settings: ServerSettings | None = None,103 model_settings: List[ModelSettings] | None = None,104):105 config_file = os.environ.get("CONFIG_FILE", None)106 if config_file is not None:107 if not os.path.exists(config_file):108 raise ValueError(f"Config file {config_file} not found!")109 with open(config_file, "rb") as f:110 # Check if yaml file111 if config_file.endswith(".yaml") or config_file.endswith(".yml"):112 import yaml113 114 config_file_settings = ConfigFileSettings.model_validate_json(115 json.dumps(yaml.safe_load(f))116 )117 else:118 config_file_settings = ConfigFileSettings.model_validate_json(f.read())119 server_settings = ServerSettings.model_validate(config_file_settings)120 model_settings = config_file_settings.models121 122 if server_settings is None and model_settings is None:123 if settings is None:124 settings = Settings()125 server_settings = ServerSettings.model_validate(settings)126 model_settings = [ModelSettings.model_validate(settings)]127 128 assert (129 server_settings is not None and model_settings is not None130 ), "server_settings and model_settings must be provided together"131 132 set_server_settings(server_settings)133 middleware = [Middleware(RawContextMiddleware, plugins=(RequestIdPlugin(),))]134 app = FastAPI(135 middleware=middleware,136 title="🦙 llama.cpp Python API",137 version=llama_cpp.__version__,138 root_path=server_settings.root_path,139 )140 app.add_middleware(141 CORSMiddleware,142 allow_origins=["*"],143 allow_credentials=True,144 allow_methods=["*"],145 allow_headers=["*"],146 )147 app.include_router(router)148 149 assert model_settings is not None150 set_llama_proxy(model_settings=model_settings)151 152 if server_settings.disable_ping_events:153 set_ping_message_factory(lambda: bytes())154 155 return app156 157 158def prepare_request_resources(159 body: CreateCompletionRequest | CreateChatCompletionRequest,160 llama_proxy: LlamaProxy,161 body_model: str | None,162 kwargs,163) -> llama_cpp.Llama:164 if llama_proxy is None:165 raise HTTPException(166 status_code=status.HTTP_503_SERVICE_UNAVAILABLE,167 detail="Service is not available",168 )169 llama = llama_proxy(body_model)170 if body.logit_bias is not None:171 kwargs["logit_bias"] = (172 _logit_bias_tokens_to_input_ids(llama, body.logit_bias)173 if body.logit_bias_type == "tokens"174 else body.logit_bias175 )176 177 if body.grammar is not None:178 kwargs["grammar"] = llama_cpp.LlamaGrammar.from_string(body.grammar)179 180 if body.min_tokens > 0:181 _min_tokens_logits_processor = llama_cpp.LogitsProcessorList(182 [llama_cpp.MinTokensLogitsProcessor(body.min_tokens, llama.token_eos())]183 )184 if "logits_processor" not in kwargs:185 kwargs["logits_processor"] = _min_tokens_logits_processor186 else:187 kwargs["logits_processor"].extend(_min_tokens_logits_processor)188 return llama189 190 191async def get_event_publisher(192 request: Request,193 inner_send_chan: MemoryObjectSendStream[typing.Any],194 body: CreateCompletionRequest | CreateChatCompletionRequest,195 body_model: str | None,196 llama_call,197 kwargs,198):199 server_settings = next(get_server_settings())200 interrupt_requests = (201 server_settings.interrupt_requests if server_settings else False202 )203 async with contextlib.asynccontextmanager(get_llama_proxy)() as llama_proxy:204 llama = prepare_request_resources(body, llama_proxy, body_model, kwargs)205 async with inner_send_chan:206 try:207 iterator = await run_in_threadpool(llama_call, llama, **kwargs)208 async for chunk in iterate_in_threadpool(iterator):209 await inner_send_chan.send(dict(data=json.dumps(chunk)))210 if await request.is_disconnected():211 raise anyio.get_cancelled_exc_class()()212 if interrupt_requests and llama_outer_lock.locked():213 await inner_send_chan.send(dict(data="[DONE]"))214 raise anyio.get_cancelled_exc_class()()215 await inner_send_chan.send(dict(data="[DONE]"))216 except anyio.get_cancelled_exc_class() as e:217 print("disconnected")218 with anyio.move_on_after(1, shield=True):219 print(220 f"Disconnected from client (via refresh/close) {request.client}"221 )222 raise e223 224 225def _logit_bias_tokens_to_input_ids(226 llama: llama_cpp.Llama,227 logit_bias: Dict[str, float],228) -> Dict[str, float]:229 to_bias: Dict[str, float] = {}230 for token, score in logit_bias.items():231 token = token.encode("utf-8")232 for input_id in llama.tokenize(token, add_bos=False, special=True):233 to_bias[str(input_id)] = score234 return to_bias235 236 237# Setup Bearer authentication scheme238bearer_scheme = HTTPBearer(auto_error=False)239 240 241async def authenticate(242 settings: Settings = Depends(get_server_settings),243 authorization: Optional[str] = Depends(bearer_scheme),244):245 # Skip API key check if it's not set in settings246 if settings.api_key is None:247 return True248 249 # check bearer credentials against the api_key250 if authorization and authorization.credentials == settings.api_key:251 # api key is valid252 return authorization.credentials253 254 # raise http error 401255 raise HTTPException(256 status_code=status.HTTP_401_UNAUTHORIZED,257 detail="Invalid API key",258 )259 260 261openai_v1_tag = "OpenAI V1"262 263 264@router.post(265 "/v1/completions",266 summary="Completion",267 dependencies=[Depends(authenticate)],268 response_model=Union[269 llama_cpp.CreateCompletionResponse,270 str,271 ],272 responses={273 "200": {274 "description": "Successful Response",275 "content": {276 "application/json": {277 "schema": {278 "anyOf": [279 {"$ref": "#/components/schemas/CreateCompletionResponse"}280 ],281 "title": "Completion response, when stream=False",282 }283 },284 "text/event-stream": {285 "schema": {286 "type": "string",287 "title": "Server Side Streaming response, when stream=True. "288 + "See SSE format: https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#Event_stream_format", # noqa: E501289 "example": """data: {... see CreateCompletionResponse ...} \\n\\n data: ... \\n\\n ... data: [DONE]""",290 }291 },292 },293 }294 },295 tags=[openai_v1_tag],296)297@router.post(298 "/v1/engines/copilot-codex/completions",299 include_in_schema=False,300 dependencies=[Depends(authenticate)],301 tags=[openai_v1_tag],302)303async def create_completion(304 request: Request,305 body: CreateCompletionRequest,306) -> llama_cpp.Completion:307 if isinstance(body.prompt, list):308 assert len(body.prompt) <= 1309 body.prompt = body.prompt[0] if len(body.prompt) > 0 else ""310 311 body_model = (312 body.model313 if request.url.path != "/v1/engines/copilot-codex/completions"314 else "copilot-codex"315 )316 317 exclude = {318 "n",319 "best_of",320 "logit_bias_type",321 "user",322 "min_tokens",323 }324 kwargs = body.model_dump(exclude=exclude)325 326 # handle streaming request327 if kwargs.get("stream", False):328 send_chan, recv_chan = anyio.create_memory_object_stream(10)329 return EventSourceResponse(330 recv_chan,331 data_sender_callable=partial( # type: ignore332 get_event_publisher,333 request=request,334 inner_send_chan=send_chan,335 body=body,336 body_model=body_model,337 llama_call=llama_cpp.Llama.__call__,338 kwargs=kwargs,339 ),340 sep="\n",341 ping_message_factory=_ping_message_factory,342 )343 344 # handle regular request345 async with contextlib.asynccontextmanager(get_llama_proxy)() as llama_proxy:346 llama = prepare_request_resources(body, llama_proxy, body_model, kwargs)347 348 if await request.is_disconnected():349 print(350 f"Disconnected from client (via refresh/close) before llm invoked {request.client}"351 )352 raise HTTPException(353 status_code=status.HTTP_400_BAD_REQUEST,354 detail="Client closed request",355 )356 357 return await run_in_threadpool(llama, **kwargs)358 359 360@router.post(361 "/v1/embeddings",362 summary="Embedding",363 dependencies=[Depends(authenticate)],364 tags=[openai_v1_tag],365)366async def create_embedding(367 request: CreateEmbeddingRequest,368 llama_proxy: LlamaProxy = Depends(get_llama_proxy),369):370 return await run_in_threadpool(371 llama_proxy(request.model).create_embedding,372 **request.model_dump(exclude={"user"}),373 )374 375 376@router.post(377 "/v1/chat/completions",378 summary="Chat",379 dependencies=[Depends(authenticate)],380 response_model=Union[llama_cpp.ChatCompletion, str],381 responses={382 "200": {383 "description": "Successful Response",384 "content": {385 "application/json": {386 "schema": {387 "anyOf": [388 {389 "$ref": "#/components/schemas/CreateChatCompletionResponse"390 }391 ],392 "title": "Completion response, when stream=False",393 }394 },395 "text/event-stream": {396 "schema": {397 "type": "string",398 "title": "Server Side Streaming response, when stream=True"399 + "See SSE format: https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#Event_stream_format", # noqa: E501400 "example": """data: {... see CreateChatCompletionResponse ...} \\n\\n data: ... \\n\\n ... data: [DONE]""",401 }402 },403 },404 }405 },406 tags=[openai_v1_tag],407)408async def create_chat_completion(409 request: Request,410 body: CreateChatCompletionRequest = Body(411 openapi_examples={412 "normal": {413 "summary": "Chat Completion",414 "value": {415 "model": "gpt-3.5-turbo",416 "messages": [417 {"role": "system", "content": "You are a helpful assistant."},418 {"role": "user", "content": "What is the capital of France?"},419 ],420 },421 },422 "json_mode": {423 "summary": "JSON Mode",424 "value": {425 "model": "gpt-3.5-turbo",426 "messages": [427 {"role": "system", "content": "You are a helpful assistant."},428 {"role": "user", "content": "Who won the world series in 2020"},429 ],430 "response_format": {"type": "json_object"},431 },432 },433 "tool_calling": {434 "summary": "Tool Calling",435 "value": {436 "model": "gpt-3.5-turbo",437 "messages": [438 {"role": "system", "content": "You are a helpful assistant."},439 {"role": "user", "content": "Extract Jason is 30 years old."},440 ],441 "tools": [442 {443 "type": "function",444 "function": {445 "name": "User",446 "description": "User record",447 "parameters": {448 "type": "object",449 "properties": {450 "name": {"type": "string"},451 "age": {"type": "number"},452 },453 "required": ["name", "age"],454 },455 },456 }457 ],458 "tool_choice": {459 "type": "function",460 "function": {461 "name": "User",462 },463 },464 },465 },466 "logprobs": {467 "summary": "Logprobs",468 "value": {469 "model": "gpt-3.5-turbo",470 "messages": [471 {"role": "system", "content": "You are a helpful assistant."},472 {"role": "user", "content": "What is the capital of France?"},473 ],474 "logprobs": True,475 "top_logprobs": 10,476 },477 },478 }479 ),480) -> llama_cpp.ChatCompletion:481 # This is a workaround for an issue in FastAPI dependencies482 # where the dependency is cleaned up before a StreamingResponse483 # is complete.484 # https://github.com/tiangolo/fastapi/issues/11143485 486 body_model = body.model487 exclude = {488 "n",489 "logit_bias_type",490 "user",491 "min_tokens",492 }493 kwargs = body.model_dump(exclude=exclude)494 495 # handle streaming request496 if kwargs.get("stream", False):497 send_chan, recv_chan = anyio.create_memory_object_stream(10)498 return EventSourceResponse(499 recv_chan,500 data_sender_callable=partial( # type: ignore501 get_event_publisher,502 request=request,503 inner_send_chan=send_chan,504 body=body,505 body_model=body_model,506 llama_call=llama_cpp.Llama.create_chat_completion,507 kwargs=kwargs,508 ),509 sep="\n",510 ping_message_factory=_ping_message_factory,511 )512 513 # handle regular request514 async with contextlib.asynccontextmanager(get_llama_proxy)() as llama_proxy:515 llama = prepare_request_resources(body, llama_proxy, body_model, kwargs)516 517 if await request.is_disconnected():518 print(519 f"Disconnected from client (via refresh/close) before llm invoked {request.client}"520 )521 raise HTTPException(522 status_code=status.HTTP_400_BAD_REQUEST,523 detail="Client closed request",524 )525 526 return await run_in_threadpool(llama.create_chat_completion, **kwargs)527 528 529@router.get(530 "/v1/models",531 summary="Models",532 dependencies=[Depends(authenticate)],533 tags=[openai_v1_tag],534)535async def get_models(536 llama_proxy: LlamaProxy = Depends(get_llama_proxy),537) -> ModelList:538 return {539 "object": "list",540 "data": [541 {542 "id": model_alias,543 "object": "model",544 "owned_by": "me",545 "permissions": [],546 }547 for model_alias in llama_proxy548 ],549 }550 551 552extras_tag = "Extras"553 554 555@router.post(556 "/extras/tokenize",557 summary="Tokenize",558 dependencies=[Depends(authenticate)],559 tags=[extras_tag],560)561async def tokenize(562 body: TokenizeInputRequest,563 llama_proxy: LlamaProxy = Depends(get_llama_proxy),564) -> TokenizeInputResponse:565 tokens = llama_proxy(body.model).tokenize(body.input.encode("utf-8"), special=True)566 567 return TokenizeInputResponse(tokens=tokens)568 569 570@router.post(571 "/extras/tokenize/count",572 summary="Tokenize Count",573 dependencies=[Depends(authenticate)],574 tags=[extras_tag],575)576async def count_query_tokens(577 body: TokenizeInputRequest,578 llama_proxy: LlamaProxy = Depends(get_llama_proxy),579) -> TokenizeInputCountResponse:580 tokens = llama_proxy(body.model).tokenize(body.input.encode("utf-8"), special=True)581 582 return TokenizeInputCountResponse(count=len(tokens))583 584 585@router.post(586 "/extras/detokenize",587 summary="Detokenize",588 dependencies=[Depends(authenticate)],589 tags=[extras_tag],590)591async def detokenize(592 body: DetokenizeInputRequest,593 llama_proxy: LlamaProxy = Depends(get_llama_proxy),594) -> DetokenizeInputResponse:595 text = llama_proxy(body.model).detokenize(body.tokens).decode("utf-8")596 597 return DetokenizeInputResponse(text=text)598