[Feat] Multi model support

llama_cpp/llama.py

@@ -759,6 +759,10 @@ def __init__(
         # Chat Format Params
         chat_format: str = "llama-2",
         chat_handler: Optional[llama_chat_format.LlamaChatCompletionHandler] = None,
+        # Cache
+        cache: bool = False,
+        cache_type: str = "ram",
+        cache_size: int = 2 << 30,
         # Misc
         verbose: bool = True,
         # Extra Params
@@ -816,6 +820,9 @@ def __init__(
             numa: Enable NUMA support. (NOTE: The initial value of this parameter is used for the remainder of the program as this value is set in llama_backend_init)
             chat_format: String specifying the chat format to use when calling create_chat_completion.
             chat_handler: Optional chat handler to use when calling create_chat_completion.
+            cache: Optional if true enables caching.
+            cache_type: String can be "ram" or "disk".
+            cache_size: Number of bytes to cache, defaults to 2GB
             verbose: Print verbose output to stderr.
 
         Raises:
@@ -940,7 +947,7 @@ def __init__(
 
         if self.verbose:
             print(llama_cpp.llama_print_system_info().decode("utf-8"), file=sys.stderr)
-
+        
         self.chat_format = chat_format
         self.chat_handler = chat_handler
 
@@ -958,6 +965,17 @@ def __init__(
             (n_ctx, self._n_vocab), dtype=np.single
         )
 
+        if cache:
+            if cache_type == "disk":
+                if verbose:
+                    print(f"Using disk cache with size {cache_size}")
+                cache = LlamaDiskCache(capacity_bytes=cache_size)
+            else:
+                if verbose:
+                    print(f"Using ram cache with size {cache_size}")
+                cache = LlamaRAMCache(capacity_bytes=cache_size)
+            self.set_cache(cache)
+
     @property
     def ctx(self) -> llama_cpp.llama_context_p:
         assert self._ctx.ctx is not None

llama_cpp/server/__main__.py

@@ -22,12 +22,16 @@
 
 """
 import os
+import sys
 import argparse
 from typing import List, Literal, Union
 
 import uvicorn
 
-from llama_cpp.server.app import create_app, Settings
+from llama_cpp.server.app import create_app
+from llama_cpp.server.settings import Settings, ServerSettings, set_settings
+
+EXE_NAME = 'llama_server'
 
 def get_base_type(annotation):
     if getattr(annotation, '__origin__', None) is Literal:
@@ -67,11 +71,12 @@ def parse_bool_arg(arg):
     else:
         raise ValueError(f'Invalid boolean argument: {arg}')
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    for name, field in Settings.model_fields.items():
+def main():
+    description = "🦙 Llama.cpp python server. Host your own LLMs!🚀"
+    parser = argparse.ArgumentParser(EXE_NAME, description=description)
+    for name, field in (ServerSettings.model_fields|Settings.model_fields).items():
         description = field.description
-        if field.default is not None and description is not None:
+        if field.default and description and not field.is_required():
             description += f" (default: {field.default})"
         base_type = get_base_type(field.annotation) if field.annotation is not None else str
         list_type = contains_list_type(field.annotation)
@@ -90,11 +95,25 @@ def parse_bool_arg(arg):
                 type=parse_bool_arg,
                 help=f"{description}",
             )
-
-    args = parser.parse_args()
-    settings = Settings(**{k: v for k, v in vars(args).items() if v is not None})
-    app = create_app(settings=settings)
+
+    try:
+        args = parser.parse_args()
+        server_settings = ServerSettings(**{k: v for k, v in vars(args).items() if v is not None})
+        set_settings(server_settings)
+        if server_settings.config and os.path.exists(server_settings.config):
+            with open(server_settings.config, 'rb') as f:
+                llama_settings = Settings.model_validate_json(f.read())
+        else:
+            llama_settings = Settings(**{k: v for k, v in vars(args).items() if v is not None})
+        app = create_app(settings=llama_settings)
+    except Exception as e:
+        print(e, file=sys.stderr)
+        parser.print_help()
+        sys.exit(1)
 
     uvicorn.run(
-        app, host=os.getenv("HOST", settings.host), port=int(os.getenv("PORT", settings.port))
+        app, host=server_settings.host, port=server_settings.port
     )
+
+if __name__ == "__main__":
+    main()