diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c34e2729e..799bc91842 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +- feat: add context manager support to Llama +- feat(server): add get model by id endpoint and container methods to llamaproxy +- feat(grammar): expose grammar and root properties on LlamaGrammar +- fix(windows): safely check directory existence before calling add_dll_directory +- fix(server): safely handle non-json and get requests in route error handler +- fix(types): allow tool role in chat request message and deduplicate union +- fix(cache): remove leftover debug print statements from disk cache +- fix(server): fix mutually exclusive argument branch in cli parser + + ## [0.3.36] - feat: update llama.cpp to ggml-org/llama.cpp@0c1e57098 diff --git a/docs/server.md b/docs/server.md index 9c09a1f1cf..fd8ad6dbb4 100644 --- a/docs/server.md +++ b/docs/server.md @@ -46,6 +46,19 @@ CLI arguments and environment variables are available for all of the fields defi Additionally the server supports configuration check out the [configuration section](#configuration-and-multi-model-support) for more information and examples. +### Supported Endpoints + +The server exposes standard OpenAI-compatible and helper endpoints: + +- `GET /v1/models` - List all configured models. +- `GET /v1/models/{model}` - Retrieve metadata for a specific model by alias/ID. +- `POST /v1/completions` - Text completions (supports streaming). +- `POST /v1/chat/completions` - Chat completions (supports streaming, tools, and JSON mode). +- `POST /v1/embeddings` - Embedding creation. +- `POST /extras/tokenize` - Tokenize text into token IDs. +- `POST /extras/tokenize/count` - Count tokens for input. +- `POST /extras/detokenize` - Detokenize token IDs to text. + ## Guides diff --git a/examples/high_level_api/high_level_api_context_manager.py b/examples/high_level_api/high_level_api_context_manager.py new file mode 100644 index 0000000000..d47d1d420c --- /dev/null +++ b/examples/high_level_api/high_level_api_context_manager.py @@ -0,0 +1,27 @@ +import json +import argparse + +from llama_cpp import Llama + +parser = argparse.ArgumentParser( + description="Run inference using Llama as a context manager." +) +parser.add_argument("-m", "--model", type=str, required=True, help="Path to GGUF model") +parser.add_argument( + "-p", + "--prompt", + type=str, + default="Question: What is the speed of light? Answer:", + help="Prompt for completion", +) +args = parser.parse_args() + +# Using Llama in a context manager automatically closes the model and frees resources +with Llama(model_path=args.model, verbose=False) as llm: + output = llm( + args.prompt, + max_tokens=64, + stop=["\n", "Question:"], + echo=True, + ) + print(json.dumps(output, indent=2)) diff --git a/llama_cpp/_ctypes_extensions.py b/llama_cpp/_ctypes_extensions.py index 02cee8a88f..f0237c4064 100644 --- a/llama_cpp/_ctypes_extensions.py +++ b/llama_cpp/_ctypes_extensions.py @@ -55,18 +55,19 @@ def load_shared_library(lib_base_name: str, base_path: pathlib.Path): # Add the library directory to the DLL search path on Windows (if needed) if sys.platform == "win32": - os.add_dll_directory(str(base_path)) - os.environ["PATH"] = str(base_path) + os.pathsep + os.environ["PATH"] - - if sys.platform == "win32" and sys.version_info >= (3, 8): - os.add_dll_directory(str(base_path)) - if "CUDA_PATH" in os.environ: - os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "bin")) - os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "lib")) - if "HIP_PATH" in os.environ: - os.add_dll_directory(os.path.join(os.environ["HIP_PATH"], "bin")) - os.add_dll_directory(os.path.join(os.environ["HIP_PATH"], "lib")) - cdll_args["winmode"] = ctypes.RTLD_GLOBAL + if base_path.exists() and base_path.is_dir(): + if sys.version_info >= (3, 8): + os.add_dll_directory(str(base_path)) + os.environ["PATH"] = str(base_path) + os.pathsep + os.environ["PATH"] + + if sys.version_info >= (3, 8): + for env_var in ("CUDA_PATH", "HIP_PATH"): + if env_var in os.environ: + for sub in ("bin", "lib"): + sub_path = os.path.join(os.environ[env_var], sub) + if os.path.isdir(sub_path): + os.add_dll_directory(sub_path) + cdll_args["winmode"] = ctypes.RTLD_GLOBAL if sys.platform == "emscripten": cdll_args["mode"] = ctypes.RTLD_GLOBAL diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 14e2f8500f..7cca101a6a 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -2296,6 +2296,12 @@ def close(self) -> None: """Explicitly free the model from memory.""" self._stack.close() + def __enter__(self) -> Llama: + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + self.close() + def __del__(self) -> None: self.close() diff --git a/llama_cpp/llama_cache.py b/llama_cpp/llama_cache.py index 5220c79337..88709caca8 100644 --- a/llama_cpp/llama_cache.py +++ b/llama_cpp/llama_cache.py @@ -142,14 +142,10 @@ def __contains__(self, key: Sequence[int]) -> bool: return self._find_longest_prefix_key(tuple(key)) is not None def __setitem__(self, key: Sequence[int], value: "llama_cpp.llama.LlamaState"): - print("LlamaDiskCache.__setitem__: called", file=sys.stderr) key = tuple(key) if key in self.cache: - print("LlamaDiskCache.__setitem__: delete", file=sys.stderr) del self.cache[key] self.cache[key] = value - print("LlamaDiskCache.__setitem__: set", file=sys.stderr) while self.cache_size > self.capacity_bytes and len(self.cache) > 0: key_to_remove = next(iter(self.cache)) del self.cache[key_to_remove] - print("LlamaDiskCache.__setitem__: trim", file=sys.stderr) diff --git a/llama_cpp/llama_grammar.py b/llama_cpp/llama_grammar.py index ba34dda831..2f0936da79 100644 --- a/llama_cpp/llama_grammar.py +++ b/llama_cpp/llama_grammar.py @@ -21,6 +21,16 @@ def __init__(self, *args, _grammar: str, **kwargs): self._grammar = _grammar self._root = LLAMA_GRAMMAR_DEFAULT_ROOT + @property + def grammar(self) -> str: + """Return the GBNF grammar string.""" + return self._grammar + + @property + def root(self) -> str: + """Return the root rule name.""" + return self._root + @classmethod def from_string(cls, grammar: str, verbose: bool = True) -> "LlamaGrammar": return cls(_grammar=grammar) @@ -39,7 +49,7 @@ def from_file(cls, file: Union[str, Path], verbose: bool = True) -> "LlamaGramma return cls.from_string(grammar, verbose=verbose) raise ValueError( - f"{cls.from_file.__name__}: error parsing grammar file: params_grammer is empty" + f"{cls.from_file.__name__}: error parsing grammar file: params_grammar is empty" ) @classmethod diff --git a/llama_cpp/llama_tokenizer.py b/llama_cpp/llama_tokenizer.py index 1375e1392d..91083456a1 100644 --- a/llama_cpp/llama_tokenizer.py +++ b/llama_cpp/llama_tokenizer.py @@ -8,7 +8,6 @@ ) import llama_cpp -from llama_cpp.llama_types import List class BaseLlamaTokenizer(abc.ABC): diff --git a/llama_cpp/llama_types.py b/llama_cpp/llama_types.py index f647822ff5..beab2688a9 100644 --- a/llama_cpp/llama_types.py +++ b/llama_cpp/llama_types.py @@ -255,7 +255,6 @@ class ChatCompletionRequestFunctionMessage(TypedDict): ChatCompletionRequestSystemMessage, ChatCompletionRequestUserMessage, ChatCompletionRequestAssistantMessage, - ChatCompletionRequestUserMessage, ChatCompletionRequestToolMessage, ChatCompletionRequestFunctionMessage, ] diff --git a/llama_cpp/server/app.py b/llama_cpp/server/app.py index f776fe159c..51e306751a 100644 --- a/llama_cpp/server/app.py +++ b/llama_cpp/server/app.py @@ -36,6 +36,7 @@ CreateEmbeddingRequest, CreateChatCompletionRequest, ModelList, + ModelData, TokenizeInputRequest, TokenizeInputResponse, TokenizeInputCountResponse, @@ -549,6 +550,29 @@ async def get_models( } +@router.get( + "/v1/models/{model:path}", + summary="Get Model", + dependencies=[Depends(authenticate)], + tags=[openai_v1_tag], +) +async def get_model( + model: str, + llama_proxy: LlamaProxy = Depends(get_llama_proxy), +) -> ModelData: + if model not in llama_proxy: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Model '{model}' not found", + ) + return { + "id": model, + "object": "model", + "owned_by": "me", + "permissions": [], + } + + extras_tag = "Extras" diff --git a/llama_cpp/server/cli.py b/llama_cpp/server/cli.py index 3a5cde9cd2..0280cc2fb6 100644 --- a/llama_cpp/server/cli.py +++ b/llama_cpp/server/cli.py @@ -98,20 +98,20 @@ def add_args_from_model(parser: argparse.ArgumentParser, model: Type[BaseModel]) type=_parse_json_object_arg, help=description, ) - elif base_type is not bool: + elif base_type is bool: parser.add_argument( f"--{name}", dest=name, - nargs="*" if list_type else None, - type=base_type, - help=description, + type=_parse_bool_arg, + help=f"{description}", ) - if base_type is bool: + else: parser.add_argument( f"--{name}", dest=name, - type=_parse_bool_arg, - help=f"{description}", + nargs="*" if list_type else None, + type=base_type, + help=description, ) diff --git a/llama_cpp/server/errors.py b/llama_cpp/server/errors.py index d0eda5664b..ad8042ef4c 100644 --- a/llama_cpp/server/errors.py +++ b/llama_cpp/server/errors.py @@ -178,25 +178,27 @@ async def custom_route_handler(request: Request) -> Response: # api key check failed raise unauthorized except Exception as exc: - json_body = await request.json() + body: Optional[ + Union[ + CreateChatCompletionRequest, + CreateCompletionRequest, + CreateEmbeddingRequest, + ] + ] = None try: - if "messages" in json_body: - # Chat completion - body: Optional[ - Union[ - CreateChatCompletionRequest, - CreateCompletionRequest, - CreateEmbeddingRequest, - ] - ] = CreateChatCompletionRequest(**json_body) - elif "prompt" in json_body: - # Text completion - body = CreateCompletionRequest(**json_body) - else: - # Embedding - body = CreateEmbeddingRequest(**json_body) + json_body = await request.json() + if isinstance(json_body, dict): + if "messages" in json_body: + # Chat completion + body = CreateChatCompletionRequest(**json_body) + elif "prompt" in json_body: + # Text completion + body = CreateCompletionRequest(**json_body) + else: + # Embedding + body = CreateEmbeddingRequest(**json_body) except Exception: - # Invalid request body + # Invalid or unreadable request body (e.g. GET request or non-JSON) body = None # Get proper error message from the exception diff --git a/llama_cpp/server/model.py b/llama_cpp/server/model.py index 8aa929202c..8acfea8454 100644 --- a/llama_cpp/server/model.py +++ b/llama_cpp/server/model.py @@ -2,7 +2,7 @@ import json -from typing import Dict, Optional, Union, List +from typing import Dict, Optional, Union, List, Any import llama_cpp import llama_cpp.llama_speculative as llama_speculative @@ -65,6 +65,17 @@ def __iter__(self): for model in self._model_settings_dict: yield model + def __contains__(self, model: str) -> bool: + return model in self._model_settings_dict + + def __len__(self) -> int: + return len(self._model_settings_dict) + + def get(self, model: str, default: Any = None) -> Any: + if model in self._model_settings_dict: + return self[model] + return default + def free(self): if self._current_model: self._current_model.close() diff --git a/llama_cpp/server/types.py b/llama_cpp/server/types.py index eda2b92801..88d84cb635 100644 --- a/llama_cpp/server/types.py +++ b/llama_cpp/server/types.py @@ -181,7 +181,7 @@ class CreateEmbeddingRequest(BaseModel): class ChatCompletionRequestMessage(BaseModel): - role: Literal["system", "user", "assistant", "function"] = Field( + role: Literal["system", "user", "assistant", "tool", "function"] = Field( default="user", description="The role of the message." ) content: Optional[str] = Field( @@ -303,7 +303,7 @@ class TokenizeInputCountResponse(BaseModel): class DetokenizeInputRequest(BaseModel): model: Optional[str] = model_field - tokens: List[int] = Field(description="A list of toekns to detokenize.") + tokens: List[int] = Field(description="A list of tokens to detokenize.") model_config = {"json_schema_extra": {"example": [{"tokens": [123, 321, 222]}]}} diff --git a/tests/test_llama.py b/tests/test_llama.py index c1c16e30ca..0e3ea1f67c 100644 --- a/tests/test_llama.py +++ b/tests/test_llama.py @@ -20,6 +20,16 @@ def test_llama_cpp_version(): assert llama_cpp.__version__ +def test_llama_context_manager(): + from unittest.mock import MagicMock + + mock_llama = llama_cpp.Llama.__new__(llama_cpp.Llama) + mock_llama._stack = MagicMock() + with mock_llama as llm: + assert llm is mock_llama + mock_llama._stack.close.assert_called_once() + + def test_llama_cpp_tokenization(): llama = llama_cpp.Llama(model_path=MODEL, vocab_only=True, verbose=False) diff --git a/tests/test_llama_chat_format.py b/tests/test_llama_chat_format.py index 18c7279cf0..c82cccc9a1 100644 --- a/tests/test_llama_chat_format.py +++ b/tests/test_llama_chat_format.py @@ -92,3 +92,64 @@ def test_hf_tokenizer_config_str_to_chat_formatter(): ) assert chat_formatter_respoonse.prompt == ("[INST] Hello, world! [/INST]") + + +def test_format_llama3(): + messages = [ + llama_types.ChatCompletionRequestSystemMessage( + role="system", content="You are a helpful assistant." + ), + llama_types.ChatCompletionRequestUserMessage( + role="user", content="Hello Llama 3" + ), + ] + response = llama_chat_format.format_llama3(messages=messages) + expected_prompt = ( + "<|start_header_id|>system<|end_header_id|>\n\n" + "You are a helpful assistant.<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n" + "Hello Llama 3<|eot_id|>" + "<|start_header_id|>assistant<|end_header_id|>\n\n" + ) + assert response.prompt == expected_prompt + assert response.stop == "<|eot_id|>" + + +def test_format_alpaca(): + messages = [ + llama_types.ChatCompletionRequestUserMessage( + role="user", content="Summarize this text." + ), + ] + response = llama_chat_format.format_alpaca(messages=messages) + assert "### Instruction: Summarize this text." in response.prompt + + +def test_format_qwen(): + messages = [ + llama_types.ChatCompletionRequestUserMessage( + role="user", content="Hello Qwen" + ), + ] + response = llama_chat_format.format_qwen(messages=messages) + assert "<|im_start|>system\nYou are a helpful assistant.<|im_end|>" in response.prompt + assert "<|im_start|>user\nHello Qwen<|im_end|>" in response.prompt + assert "<|im_start|>assistant\n" in response.prompt + assert response.stop == "<|endoftext|>" + + +def test_format_chatml(): + messages = [ + llama_types.ChatCompletionRequestSystemMessage( + role="system", content="System instruction" + ), + llama_types.ChatCompletionRequestUserMessage( + role="user", content="User message" + ), + ] + response = llama_chat_format.format_chatml(messages=messages) + assert "<|im_start|>system\nSystem instruction<|im_end|>" in response.prompt + assert "<|im_start|>user\nUser message<|im_end|>" in response.prompt + assert "<|im_start|>assistant\n" in response.prompt + assert response.stop == "<|im_end|>" + diff --git a/tests/test_llama_grammar.py b/tests/test_llama_grammar.py index 34ef2874df..1eb735b14c 100644 --- a/tests/test_llama_grammar.py +++ b/tests/test_llama_grammar.py @@ -10,9 +10,8 @@ def test_grammar_from_string(): grammar = llama_cpp.LlamaGrammar.from_string(tree) - # assert grammar._n_rules == 3 - # assert grammar._start_rule_index == 2 - # assert grammar.grammar is not None + assert grammar.grammar == tree + assert grammar.root == "root" def test_composed_pydantic_grammar(): diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 0000000000..74250729ef --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,177 @@ +import argparse +import pytest +from unittest.mock import MagicMock + +from llama_cpp.server.settings import ServerSettings, ModelSettings +from llama_cpp.server.cli import add_args_from_model, parse_model_from_args +from llama_cpp.server.types import ( + ChatCompletionRequestMessage, + DetokenizeInputRequest, + TokenizeInputRequest, +) + + +def test_server_types(): + # Verify tool role is allowed in ChatCompletionRequestMessage + msg = ChatCompletionRequestMessage(role="tool", content="tool output") + assert msg.role == "tool" + assert msg.content == "tool output" + + detok = DetokenizeInputRequest(tokens=[1, 2, 3]) + assert detok.tokens == [1, 2, 3] + + tok = TokenizeInputRequest(input="hello") + assert tok.input == "hello" + + +def test_server_cli_args(): + parser = argparse.ArgumentParser() + add_args_from_model(parser, ServerSettings) + args = parser.parse_args( + [ + "--host", + "127.0.0.1", + "--port", + "9000", + "--interrupt_requests", + "false", + ] + ) + settings = parse_model_from_args(ServerSettings, args) + assert settings.host == "127.0.0.1" + assert settings.port == 9000 + assert settings.interrupt_requests is False + + +def test_llama_proxy_container(): + from llama_cpp.server.model import LlamaProxy + + mock_model = MagicMock() + model_settings = [ + ModelSettings(model="model-a.gguf", model_alias="model-a"), + ModelSettings(model="model-b.gguf", model_alias="model-b"), + ] + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + LlamaProxy, "load_llama_from_model_settings", lambda self, s: mock_model + ) + proxy = LlamaProxy(models=model_settings) + + assert len(proxy) == 2 + assert "model-a" in proxy + assert "model-b" in proxy + assert "non-existent" not in proxy + assert proxy.get("model-a") is not None + assert proxy.get("non-existent") is None + assert proxy.get("non-existent", "default") == "default" + + +def test_server_models_endpoints(): + pytest.importorskip("sse_starlette") + pytest.importorskip("starlette_context") + from fastapi.testclient import TestClient + from llama_cpp.server.app import create_app + from llama_cpp.server.model import LlamaProxy + + mock_model = MagicMock() + model_settings = [ + ModelSettings(model="model-a.gguf", model_alias="model-a"), + ] + server_settings = ServerSettings(api_key=None) + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + LlamaProxy, "load_llama_from_model_settings", lambda self, s: mock_model + ) + app = create_app(server_settings=server_settings, model_settings=model_settings) + client = TestClient(app) + + # GET /v1/models + response = client.get("/v1/models") + assert response.status_code == 200 + data = response.json() + assert data["object"] == "list" + assert len(data["data"]) == 1 + assert data["data"][0]["id"] == "model-a" + + # GET /v1/models/{model} for existing model + response = client.get("/v1/models/model-a") + assert response.status_code == 200 + model_data = response.json() + assert model_data["id"] == "model-a" + assert model_data["object"] == "model" + assert model_data["owned_by"] == "me" + + # GET /v1/models/{model} for non-existent model returns 404 + response = client.get("/v1/models/unknown-model") + assert response.status_code == 404 + assert "not found" in response.json()["detail"].lower() + + +def test_server_authentication(): + pytest.importorskip("sse_starlette") + pytest.importorskip("starlette_context") + from fastapi.testclient import TestClient + from llama_cpp.server.app import create_app + from llama_cpp.server.model import LlamaProxy + + mock_model = MagicMock() + model_settings = [ + ModelSettings(model="model-a.gguf", model_alias="model-a"), + ] + server_settings = ServerSettings(api_key="secret-token-123") + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + LlamaProxy, "load_llama_from_model_settings", lambda self, s: mock_model + ) + app = create_app(server_settings=server_settings, model_settings=model_settings) + client = TestClient(app) + + # Missing token -> 401 + res = client.get("/v1/models") + assert res.status_code == 401 + + # Invalid token -> 401 + res = client.get( + "/v1/models", headers={"Authorization": "Bearer wrong-token"} + ) + assert res.status_code == 401 + + # Valid token -> 200 + res = client.get( + "/v1/models", headers={"Authorization": "Bearer secret-token-123"} + ) + assert res.status_code == 200 + + +def test_server_route_error_handler_resilience(): + pytest.importorskip("sse_starlette") + pytest.importorskip("starlette_context") + from fastapi.testclient import TestClient + from llama_cpp.server.app import create_app + from llama_cpp.server.model import LlamaProxy + + mock_model = MagicMock() + mock_model.tokenize.return_value = [1, 2, 3] + mock_model.detokenize.return_value = b"hello" + model_settings = [ + ModelSettings(model="model-a.gguf", model_alias="model-a"), + ] + server_settings = ServerSettings(api_key=None) + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + LlamaProxy, "load_llama_from_model_settings", lambda self, s: mock_model + ) + app = create_app(server_settings=server_settings, model_settings=model_settings) + client = TestClient(app) + + # Post non-json body to tokenize + response = client.post( + "/extras/tokenize", + content=b"not-json-content", + headers={"Content-Type": "application/json"}, + ) + # Should gracefully return 422 unprocessable entity, not 500 JSONDecodeError crash + assert response.status_code in (400, 422)