import os import json from threading import Lock from typing import List, Optional, Union, Iterator, Dict from typing_extensions import TypedDict, Literal import llama_cpp from fastapi import Depends, FastAPI from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, BaseSettings, Field, create_model_from_typeddict from sse_starlette.sse import EventSourceResponse class Settings(BaseSettings): model: str = os.environ.get("MODEL", "null") n_ctx: int = 2048 n_batch: int = 512 n_threads: int = max((os.cpu_count() or 2) // 2, 1) f16_kv: bool = True use_mlock: bool = False # This causes a silent failure on platforms that don't support mlock (e.g. Windows) took forever to figure out... use_mmap: bool = True embedding: bool = True last_n_tokens_size: int = 64 logits_all: bool = False cache: bool = False # WARNING: This is an experimental feature vocab_only: bool = False app = FastAPI( title="🦙 llama.cpp Python API", version="0.0.1", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) llama: llama_cpp.Llama = None def init_llama(settings: Settings = None): if settings is None: settings = Settings() global llama llama = llama_cpp.Llama( settings.model, f16_kv=settings.f16_kv, use_mlock=settings.use_mlock, use_mmap=settings.use_mmap, embedding=settings.embedding, logits_all=settings.logits_all, n_threads=settings.n_threads, n_batch=settings.n_batch, n_ctx=settings.n_ctx, last_n_tokens_size=settings.last_n_tokens_size, vocab_only=settings.vocab_only, ) if settings.cache: cache = llama_cpp.LlamaCache() llama.set_cache(cache) llama_lock = Lock() def get_llama(): with llama_lock: yield llama model_field = Field( description="The model to use for generating completions." ) max_tokens_field = Field( default=16, ge=1, le=2048, description="The maximum number of tokens to generate." ) temperature_field = Field( default=0.8, ge=0.0, le=2.0, description="Adjust the randomness of the generated text.\n\n" + "Temperature is a hyperparameter that controls the randomness of the generated text. It affects the probability distribution of the model's output tokens. A higher temperature (e.g., 1.5) makes the output more random and creative, while a lower temperature (e.g., 0.5) makes the output more focused, deterministic, and conservative. The default value is 0.8, which provides a balance between randomness and determinism. At the extreme, a temperature of 0 will always pick the most likely next token, leading to identical outputs in each run." ) top_p_field = Field( default=0.95, ge=0.0, le=1.0, description="Limit the next token selection to a subset of tokens with a cumulative probability above a threshold P.\n\n" + "Top-p sampling, also known as nucleus sampling, is another text generation method that selects the next token from a subset of tokens that together have a cumulative probability of at least p. This method provides a balance between diversity and quality by considering both the probabilities of tokens and the number of tokens to sample from. A higher value for top_p (e.g., 0.95) will lead to more diverse text, while a lower value (e.g., 0.5) will generate more focused and conservative text." ) stop_field = Field( default=None, description="A list of tokens at which to stop generation. If None, no stop tokens are used." ) stream_field = Field( default=False, description="Whether to stream the results as they are generated. Useful for chatbots." ) top_k_field = Field( default=40, ge=0, description="Limit the next token selection to the K most probable tokens.\n\n" + "Top-k sampling is a text generation method that selects the next token only from the top k most likely tokens predicted by the model. It helps reduce the risk of generating low-probability or nonsensical tokens, but it may also limit the diversity of the output. A higher value for top_k (e.g., 100) will consider more tokens and lead to more diverse text, while a lower value (e.g., 10) will focus on the most probable tokens and generate more conservative text." ) repeat_penalty_field = Field( default=1.0, ge=0.0, description="A penalty applied to each token that is already generated. This helps prevent the model from repeating itself.\n\n" + "Repeat penalty is a hyperparameter used to penalize the repetition of token sequences during text generation. It helps prevent the model from generating repetitive or monotonous text. A higher value (e.g., 1.5) will penalize repetitions more strongly, while a lower value (e.g., 0.9) will be more lenient." ) class CreateCompletionRequest(BaseModel): prompt: Union[str, List[str]] = Field( default="", description="The prompt to generate completions for." ) suffix: Optional[str] = Field( default=None, description="A suffix to append to the generated text. If None, no suffix is appended. Useful for chatbots." ) max_tokens: int = max_tokens_field temperature: float = temperature_field top_p: float = top_p_field echo: bool = Field( default=False, description="Whether to echo the prompt in the generated text. Useful for chatbots." ) stop: Optional[List[str]] = stop_field stream: bool = stream_field logprobs: Optional[int] = Field( default=None, ge=0, description="The number of logprobs to generate. If None, no logprobs are generated." ) # ignored, but marked as required for the sake of compatibility with openai's api model: str = model_field # llama.cpp specific parameters top_k: int = top_k_field repeat_penalty: float = repeat_penalty_field class Config: schema_extra = { "example": { "prompt": "\n\n### Instructions:\nWhat is the capital of France?\n\n### Response:\n", "stop": ["\n", "###"], } } CreateCompletionResponse = create_model_from_typeddict(llama_cpp.Completion) @app.post( "/v1/completions", response_model=CreateCompletionResponse, ) def create_completion( request: CreateCompletionRequest, llama: llama_cpp.Llama = Depends(get_llama) ): if isinstance(request.prompt, list): request.prompt = "".join(request.prompt) completion_or_chunks = llama( **request.dict( exclude={ "model" } ) ) if request.stream: chunks: Iterator[llama_cpp.CompletionChunk] = completion_or_chunks # type: ignore return EventSourceResponse(dict(data=json.dumps(chunk)) for chunk in chunks) completion: llama_cpp.Completion = completion_or_chunks # type: ignore return completion class CreateEmbeddingRequest(BaseModel): # ignored, but marked as required for the sake of compatibility with openai's api model: str = model_field input: str class Config: schema_extra = { "example": { "input": "The food was delicious and the waiter...", } } CreateEmbeddingResponse = create_model_from_typeddict(llama_cpp.Embedding) @app.post( "/v1/embeddings", response_model=CreateEmbeddingResponse, ) def create_embedding( request: CreateEmbeddingRequest, llama: llama_cpp.Llama = Depends(get_llama) ): return llama.create_embedding(**request.dict(exclude={"model"})) class ChatCompletionRequestMessage(BaseModel): role: Union[Literal["system"], Literal["user"], Literal["assistant"]] = Field( default=Literal["user"], description="The role of the message." ) content: str = Field(default="", description="The content of the message.") class CreateChatCompletionRequest(BaseModel): messages: List[ChatCompletionRequestMessage] = Field( default=[], description="A list of messages to generate completions for." ) max_tokens: int = max_tokens_field temperature: float = temperature_field top_p: float = top_p_field stop: Optional[List[str]] = stop_field stream: bool = stream_field # ignored, but marked as required for the sake of compatibility with openai's api model: str = model_field # llama.cpp specific parameters top_k: int = top_k_field repeat_penalty: float = repeat_penalty_field class Config: schema_extra = { "example": { "messages": [ ChatCompletionRequestMessage( role="system", content="You are a helpful assistant." ), ChatCompletionRequestMessage( role="user", content="What is the capital of France?" ), ] } } CreateChatCompletionResponse = create_model_from_typeddict(llama_cpp.ChatCompletion) @app.post( "/v1/chat/completions", response_model=CreateChatCompletionResponse, ) def create_chat_completion( request: CreateChatCompletionRequest, llama: llama_cpp.Llama = Depends(get_llama), ) -> Union[llama_cpp.ChatCompletion, EventSourceResponse]: completion_or_chunks = llama.create_chat_completion( **request.dict( exclude={ "model" } ), ) if request.stream: async def server_sent_events( chat_chunks: Iterator[llama_cpp.ChatCompletionChunk], ): for chat_chunk in chat_chunks: yield dict(data=json.dumps(chat_chunk)) yield dict(data="[DONE]") chunks: Iterator[llama_cpp.ChatCompletionChunk] = completion_or_chunks # type: ignore return EventSourceResponse( server_sent_events(chunks), ) completion: llama_cpp.ChatCompletion = completion_or_chunks # type: ignore return completion class ModelData(TypedDict): id: str object: Literal["model"] owned_by: str permissions: List[str] class ModelList(TypedDict): object: Literal["list"] data: List[ModelData] GetModelResponse = create_model_from_typeddict(ModelList) @app.get("/v1/models", response_model=GetModelResponse) def get_models() -> ModelList: return { "object": "list", "data": [ { "id": llama.model_path, "object": "model", "owned_by": "me", "permissions": [], } ], }