ommlds 0.0.0.dev466__py3-none-any.whl → 0.0.0.dev512__py3-none-any.whl

ommlds/backends/groq/_marshal.py

@@ -0,0 +1,23 @@
+from omlish import lang
+from omlish import marshal as msh
+
+from .protocol import ChatCompletionRequest
+
+
+##
+
+
+@lang.static_init
+def _install_standard_marshaling() -> None:
+    for root_cls, tag_field in [
+        (ChatCompletionRequest.Message, 'role'),
+    ]:
+        msh.install_standard_factories(*msh.standard_polymorphism_factories(
+            msh.polymorphism_from_subclasses(
+                root_cls,
+                naming=msh.Naming.SNAKE,
+                strip_suffix=msh.AutoStripSuffix,
+            ),
+            msh.FieldTypeTagging(tag_field),
+            unions='partial',
+        ))

ommlds/backends/groq/clients.py

@@ -0,0 +1,9 @@
+import typing as ta
+
+
+##
+
+
+REQUIRED_HTTP_HEADERS: ta.Mapping[bytes, bytes] = {
+    b'User-Agent': b'python-httpx/0.28.1',  # required or it 403's lol
+}

ommlds/backends/groq/protocol.py

@@ -0,0 +1,247 @@
+"""
+https://console.groq.com/docs/api-reference#chat-create
+"""
+import typing as ta
+
+from omlish import dataclasses as dc
+from omlish import lang
+from omlish import marshal as msh
+
+
+##
+
+
+def _set_class_marshal_options(cls):
+    msh.update_object_options(
+        cls,
+        field_defaults=msh.FieldOptions(
+            omit_if=lang.is_none,
+        ),
+    )
+
+    return cls
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class ChatCompletionRequest(lang.Final):
+    @dc.dataclass(frozen=True, kw_only=True)
+    class Message(lang.Sealed, lang.Abstract):
+        pass
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class SystemMessage(Message, lang.Final):
+        content: str | ta.Sequence[str]
+        name: str | None = None
+        role: ta.Literal['system'] = 'system'
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class UserMessage(Message, lang.Final):
+        content: str | ta.Sequence[str]
+        name: str | None = None
+        role: ta.Literal['user'] = 'user'
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class AssistantMessage(Message, lang.Final):
+        content: str | ta.Sequence[str] | None = None
+        name: str | None = None
+        reasoning: str | None = None
+        role: ta.Literal['assistant'] = 'assistant'
+
+        @dc.dataclass(frozen=True, kw_only=True)
+        @_set_class_marshal_options
+        class ToolCall(lang.Final):
+            @dc.dataclass(frozen=True, kw_only=True)
+            @_set_class_marshal_options
+            class Function(lang.Final):
+                arguments: str
+                name: str
+
+            function: Function
+            id: str
+            type: ta.Literal['function'] = 'function'
+
+        tool_calls: ta.Sequence[ToolCall] | None = None
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class ToolMessage(Message, lang.Final):
+        content: str | ta.Sequence[str]
+        role: ta.Literal['tool'] = 'tool'
+        tool_call_id: str
+
+    messages: ta.Sequence[Message]
+    model: str
+    citation_options: ta.Literal['enabled', 'disabled'] | None = None
+    compound_custom: ta.Mapping[str, ta.Any] | None = None
+    disable_tool_validation: bool | None = None
+    documents: ta.Sequence[ta.Mapping[str, ta.Any]] | None = None
+    frequency_penalty: float | None = None
+    include_reasoning: bool | None = None
+    logit_bias: ta.Mapping[str, ta.Any] | None = None
+    logprobs: bool | None = None
+    max_completion_tokens: int | None = None
+    n: int | None = None
+    parallel_tool_calls: bool | None = None
+    presence_penalty: float | None = None
+    reasoning_effort: ta.Literal['none', 'default', 'low', 'medium', 'high'] | None = None
+    reasoning_format: ta.Literal['hidden', 'raw', 'parsed'] | None = None
+    response_format: ta.Any | None = None
+    search_settings: ta.Mapping[str, ta.Any] | None = None
+    seed: int | None = None
+    service_tier: ta.Literal['auto', 'on_demand', 'flex', 'performance', 'null'] | None = None
+    stop: str | ta.Sequence[str] | None = None
+    store: bool | None = None
+    stream: bool | None = None
+    stream_options: ta.Mapping[str, ta.Any] | None = None
+    temperature: float | None = None
+    ool_choice: str | None = None
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class Tool(lang.Final):
+        @dc.dataclass(frozen=True, kw_only=True)
+        @_set_class_marshal_options
+        class Function(lang.Final):
+            description: str | None = None
+            name: str
+            parameters: ta.Mapping[str, ta.Any] | None = None  # json schema
+            strict: bool | None = None
+
+        function: Function
+        type: ta.Literal['function', 'browser_search', 'code_interpreter'] = 'function'
+
+    tools: ta.Sequence[Tool] | None = None
+
+    top_logprobs: int | None = None
+    top_p: float | None = None
+    user: str | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class ExecutedTool(lang.Final):
+    arguments: str
+    index: int
+    type: str
+    browser_results: ta.Sequence[ta.Any] | None = None
+    code_results: ta.Sequence[ta.Any] | None = None
+    output: str | None = None
+    search_results: ta.Any | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class ChatCompletionResponse(lang.Final):
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class Choice(lang.Final):
+        finish_reason: ta.Literal['stop', 'length', 'tool_calls', 'function_call']
+        index: int
+        logprobs: ta.Mapping[str, ta.Any] | None = None
+
+        @dc.dataclass(frozen=True, kw_only=True)
+        @_set_class_marshal_options
+        class Message(lang.Final):
+            annotations: ta.Sequence[ta.Mapping[str, ta.Any]] | None = None
+            content: str | None = None
+
+            executed_tools: ta.Sequence[ExecutedTool] | None = None
+
+            reasoning: str | None = None
+            role: ta.Literal['assistant'] = 'assistant'
+
+            @dc.dataclass(frozen=True, kw_only=True)
+            @_set_class_marshal_options
+            class ToolCall(lang.Final):
+                id: str
+
+                @dc.dataclass(frozen=True, kw_only=True)
+                @_set_class_marshal_options
+                class Function(lang.Final):
+                    arguments: str
+                    name: str
+
+                function: Function
+                type: ta.Literal['function'] = 'function'
+
+            tool_calls: ta.Sequence[ToolCall] | None = None
+
+        message: Message
+
+    choices: ta.Sequence[Choice]
+    created: int
+    id: str
+    model: str
+    object: ta.Literal['chat.completion'] = 'chat.completion'
+    system_fingerprint: str
+    usage: ta.Mapping[str, ta.Any] | None = None
+    usage_breakdown: ta.Mapping[str, ta.Any] | None = None
+    x_groq: ta.Mapping[str, ta.Any] | None = None
+    service_tier: str | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class ChatCompletionChunk(lang.Final):
+    id: str
+    object: ta.Literal['chat.completion.chunk'] = 'chat.completion.chunk'
+    created: int
+    model: str
+    system_fingerprint: str
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class Choice(lang.Final):
+        index: int
+
+        @dc.dataclass(frozen=True, kw_only=True)
+        @_set_class_marshal_options
+        class Delta(lang.Final):
+            role: str | None = None
+            content: str | None = None
+
+            channel: str | None = None
+            reasoning: str | None = None
+
+            @dc.dataclass(frozen=True, kw_only=True)
+            @_set_class_marshal_options
+            class ToolCall(lang.Final):
+                index: int
+                id: str | None = None
+
+                @dc.dataclass(frozen=True, kw_only=True)
+                @_set_class_marshal_options
+                class Function(lang.Final):
+                    arguments: str | None = None
+                    name: str | None = None
+
+                function: Function | None = None
+
+                type: ta.Literal['function'] = 'function'
+
+            tool_calls: ta.Sequence[ToolCall] | None = None
+
+            executed_tools: ta.Sequence[ExecutedTool] | None = None
+
+        delta: Delta
+        logprobs: ta.Mapping[str, ta.Any] | None = None
+        finish_reason: ta.Literal['stop', 'length', 'tool_calls', 'function_call'] | None = None
+
+    choices: ta.Sequence[Choice]
+
+    x_groq: ta.Mapping[str, ta.Any] | None = None
+    service_tier: str | None = None
+    usage: ta.Mapping[str, ta.Any] | None = None
+
+
+##
+
+
+msh.register_global_module_import('._marshal', __package__)

ommlds/{huggingface.py → backends/huggingface/cache.py}

@@ -11,15 +11,10 @@ from omlish import check
 from omlish import lang
 
 
-if ta.TYPE_CHECKING:
+with lang.auto_proxy_import(globals()):
     import huggingface_hub as hf
     import huggingface_hub.errors  # noqa
     import huggingface_hub.utils  # noqa
-else:
-    hf = lang.proxy_import('huggingface_hub', extras=[
-        'errors',
-        'utils',
-    ])
 
 
 ##

ommlds/backends/huggingface/cli.py

@@ -0,0 +1,208 @@
+import datetime
+import os
+import sys
+import typing as ta
+
+from omdev.cli import CliModule
+from omlish import lang
+from omlish.argparse import all as ap
+from omlish.formats import json
+from omlish.logs import all as logs
+from omlish.term.confirm import confirm_action
+
+
+with lang.auto_proxy_import(globals()):
+    import huggingface_hub as hf
+    import huggingface_hub.errors  # noqa
+    import huggingface_hub.utils  # noqa
+
+
+log = logs.get_module_logger(globals())
+
+
+##
+
+
+def fmt_ts(f: float) -> ta.Any:
+    dt = datetime.datetime.fromtimestamp(f)  # noqa
+    return dt.isoformat()
+
+
+class Cli(ap.Cli):
+    def _passthrough_args_cmd(
+            self,
+            exe: str,
+            pre_args: ta.Sequence[str] = (),
+            post_args: ta.Sequence[str] = (),
+    ) -> ta.NoReturn:
+        os.execvp(
+            exe,
+            [
+                sys.executable,
+                *pre_args,
+                *self.unknown_args,
+                *self.args.args,
+                *post_args,
+            ],
+        )
+
+    @ap.cmd(
+        ap.arg('args', nargs=ap.REMAINDER),
+        name='cli',
+        accepts_unknown=True,
+    )
+    def cli_cmd(self) -> None:
+        self._passthrough_args_cmd(sys.executable, ['-m', 'huggingface_hub.cli.hf'])
+
+    #
+
+    @ap.cmd(
+        ap.arg('--dir'),
+    )
+    def scan(self) -> None:
+        hf_cache_info = hf.utils.scan_cache_dir(self.args.dir)
+
+        repo_dcts = [
+            {
+                'repo_id': repo.repo_id,
+                'repo_type': repo.repo_type,
+
+                'repo_path': str(repo.repo_path),
+
+                'size_on_disk': repo.size_on_disk,
+                'size_on_disk_str': repo.size_on_disk_str,
+
+                'nb_files': repo.nb_files,
+
+                'revisions': [
+                    {
+                        'commit_hash': rev.commit_hash,
+
+                        'snapshot_path': str(rev.snapshot_path),
+
+                        'size_on_disk': rev.size_on_disk,
+
+                        'files': [
+                            {
+                                'file_name': file.file_name,
+                                'file_path': str(file.file_path),
+                                'blob_path': str(file.blob_path),
+
+                                'size_on_disk': file.size_on_disk,
+                                'size_on_disk_str': file.size_on_disk_str,
+
+                                'blob_last_modified': fmt_ts(file.blob_last_modified),
+                                'blob_last_modified_str': file.blob_last_modified_str,
+                                'blob_last_accessed': fmt_ts(file.blob_last_accessed),
+                                'blob_last_accessed_str': file.blob_last_accessed_str,
+                            }
+                            for file in sorted(rev.files, key=lambda file: file.blob_last_accessed)
+                        ],
+
+                        'refs': sorted(rev.refs),
+
+                        'last_modified': fmt_ts(rev.last_modified),
+                        'last_modified_str': rev.last_modified_str,
+
+                    }
+                    for rev in sorted(repo.revisions, key=lambda rev: rev.last_modified)
+                ],
+
+                'last_modified': fmt_ts(repo.last_modified),
+                'last_modified_str': repo.last_modified_str,
+                'last_accessed': fmt_ts(repo.last_accessed),
+                'last_accessed_str': repo.last_accessed_str,
+
+                'refs': sorted(repo.refs),
+            }
+            for repo in sorted(hf_cache_info.repos, key=lambda repo: repo.last_accessed)
+        ]
+
+        print(json.dumps_pretty(repo_dcts))
+
+    @ap.cmd(
+        ap.arg('--dir'),
+    )
+    def list(self) -> None:
+        hf_cache_info = hf.utils.scan_cache_dir(self.args.dir)
+
+        repos = [
+            repo
+            for repo in hf_cache_info.repos
+            if repo.repo_type == 'model'
+            and repo.nb_files
+        ]
+
+        repo_dcts = [
+            {
+                'repo_id': repo.repo_id,
+                'repo_type': repo.repo_type,
+
+                'repo_path': str(repo.repo_path),
+
+                'size_on_disk': repo.size_on_disk,
+                'size_on_disk_str': repo.size_on_disk_str,
+
+                'nb_files': repo.nb_files,
+
+                'last_modified': fmt_ts(repo.last_modified),
+                'last_modified_str': repo.last_modified_str,
+                'last_accessed': fmt_ts(repo.last_accessed),
+                'last_accessed_str': repo.last_accessed_str,
+            }
+            for repo in sorted(repos, key=lambda repo: repo.last_accessed)
+        ]
+
+        print(json.dumps_pretty(repo_dcts))
+
+    @ap.cmd(
+        ap.arg('key', action='append'),
+        ap.arg('--dir'),
+        ap.arg('--dry-run', action='store_true'),
+        ap.arg('--no-confirm', action='store_true'),
+    )
+    def rm(self) -> None:
+        if not self.args.key:
+            raise ValueError('key is required')
+
+        hf_cache_info = hf.utils.scan_cache_dir(self.args.dir)
+
+        repos_by_id = {repo.repo_id: repo for repo in hf_cache_info.repos}
+        repos_by_rev = {rev.commit_hash: repo for repo in hf_cache_info.repos for rev in repo.revisions}
+
+        rm_revs: dict[str, None] = {}
+
+        for key in self.args.key:
+            if key in repos_by_id:
+                rm_revs.update({rev.commit_hash: None for rev in repos_by_id[key].revisions})
+            elif key in repos_by_rev:
+                rm_revs.update({key: None})
+            else:
+                raise ValueError(f'key {key} not found')
+
+        for rm_rev in rm_revs:
+            rm_repo = repos_by_rev[rm_rev]
+
+            if not self.args.no_confirm:
+                if not confirm_action(f'Delete {rm_repo.repo_id}@{rm_rev}?'):
+                    return
+
+            if not self.args.dry_run:
+                strategy = hf_cache_info.delete_revisions(rm_rev)
+                strategy.execute()
+
+
+##
+
+
+def _main() -> None:
+    logs.configure_standard_logging('INFO')
+    Cli()()
+
+
+# @omlish-manifest
+_CLI_MODULE = CliModule('hf', __name__)
+
+
+if __name__ == '__main__':
+    _main()

ommlds/backends/llamacpp/logging.py

@@ -1,4 +1,7 @@
 """
+NOTE: This can't be cleaned up too much - the callback can't be a closure to hide its guts because it needs to be
+      picklable for multiprocessing.
+
 FIXME:
  - it outputs newline-terminated so buffer and chop on newlines - DelimitingBuffer again
 """
@@ -27,4 +30,4 @@ def llama_log_callback(
 
 @lang.cached_function
 def install_logging_hook() -> None:
-    llama_cpp.llama_log_set(llama_log_callback, ct.c_void_p(0))
+    llama_cpp.llama_log_set(llama_log_callback, ct.c_void_p(0))  # noqa

ommlds/backends/mlx/caching.py

@@ -17,7 +17,11 @@
 # https://github.com/ml-explore/mlx-lm/blob/ce2358d297af245b002e690623f00195b6507da0/mlx_lm/generate.py
 import typing as ta
 
-import mlx_lm.models.cache
+from omlish import lang
+
+
+with lang.auto_proxy_import(globals()):
+    import mlx_lm.models.cache as mlx_lm_models_cache
 
 
 ##
@@ -32,13 +36,13 @@ def maybe_quantize_kv_cache(
 ) -> None:
     if not (
             kv_bits is not None and
-            not isinstance(prompt_cache[0], mlx_lm.models.cache.QuantizedKVCache) and
+            not isinstance(prompt_cache[0], mlx_lm_models_cache.QuantizedKVCache) and
             prompt_cache[0].offset > quantized_kv_start
     ):
         return
 
     for i in range(len(prompt_cache)):
-        if isinstance(prompt_cache[i], mlx_lm.models.cache.KVCache):
+        if isinstance(prompt_cache[i], mlx_lm_models_cache.KVCache):
             prompt_cache[i] = prompt_cache[i].to_quantized(
                 bits=kv_bits,
                 group_size=kv_group_size,

ommlds/backends/mlx/cli.py

@@ -20,16 +20,19 @@ import json
 import sys
 import typing as ta
 
-import mlx.core as mx
-import mlx_lm.models.cache
-import mlx_lm.sample_utils
-import mlx_lm.utils
+from omlish import lang
 
 from .generation import GenerationParams
 from .generation import generate
 from .loading import load_model
 
 
+with lang.auto_proxy_import(globals()):
+    import mlx.core as mx
+    import mlx_lm.models.cache as mlx_lm_models_cache
+    import mlx_lm.sample_utils as mlx_lm_sample_utils
+
+
 ##
 
 
@@ -214,11 +217,11 @@ def _main() -> None:
     # Load the prompt cache and metadata if a cache file is provided
     using_cache = args.prompt_cache_file is not None
     if using_cache:
-        prompt_cache, metadata = mlx_lm.models.cache.load_prompt_cache(
+        prompt_cache, metadata = mlx_lm_models_cache.load_prompt_cache(
             args.prompt_cache_file,
             return_metadata=True,
         )
-        if isinstance(prompt_cache[0], mlx_lm.models.cache.QuantizedKVCache):
+        if isinstance(prompt_cache[0], mlx_lm_models_cache.QuantizedKVCache):
             if args.kv_bits is not None and args.kv_bits != prompt_cache[0].bits:
                 raise ValueError('--kv-bits does not match the kv cache loaded from --prompt-cache-file.')
             if args.kv_group_size != prompt_cache[0].group_size:
@@ -293,7 +296,7 @@ def _main() -> None:
     else:
         prompt = tokenizer.encode(prompt)
 
-    sampler = mlx_lm.sample_utils.make_sampler(
+    sampler = mlx_lm_sample_utils.make_sampler(
         args.temp,
         args.top_p,
         args.min_p,

ommlds/backends/mlx/generation.py

@@ -21,10 +21,6 @@ import io
 import sys
 import typing as ta
 
-import mlx.core as mx
-import mlx_lm.models.cache
-from mlx import nn
-
 from omlish import check
 from omlish import lang
 
@@ -33,6 +29,12 @@ from .limits import wired_limit_context
 from .tokenization import Tokenization
 
 
+with lang.auto_proxy_import(globals()):
+    import mlx.core as mx
+    import mlx.nn as mlx_nn
+    import mlx_lm.models.cache as mlx_lm_models_cache
+
+
 ##
 
 
@@ -47,9 +49,9 @@ def _generation_stream():
 class LogitProcessor(ta.Protocol):
     def __call__(
             self,
-            tokens: mx.array,
-            logits: mx.array,
-    ) -> mx.array:
+            tokens: 'mx.array',
+            logits: 'mx.array',
+    ) -> 'mx.array':
         ...
 
 
@@ -99,12 +101,12 @@ class GenerationParams:
 
 class _GenerationStep(ta.NamedTuple):
     token: int
-    logprobs: mx.array
+    logprobs: 'mx.array'
 
 
 def _generate_step(
-        prompt: mx.array,
-        model: nn.Module,
+        prompt: 'mx.array',
+        model: 'mlx_nn.Module',
         params: GenerationParams = GenerationParams(),
 ) -> ta.Generator[_GenerationStep]:
     y = prompt
@@ -113,7 +115,7 @@ def _generate_step(
     # Create the Kv cache for generation
     prompt_cache = params.prompt_cache
     if prompt_cache is None:
-        prompt_cache = mlx_lm.models.cache.make_prompt_cache(
+        prompt_cache = mlx_lm_models_cache.make_prompt_cache(
             model,
             max_kv_size=params.max_kv_size,
         )
@@ -221,7 +223,7 @@ class GenerationOutput:
     token: int
 
     # A vector of log probabilities.
-    logprobs: mx.array
+    logprobs: 'mx.array'
 
     # The number of tokens in the prompt.
     prompt_tokens: int
@@ -234,9 +236,9 @@ class GenerationOutput:
 
 
 def stream_generate(
-        model: nn.Module,
+        model: 'mlx_nn.Module',
         tokenization: Tokenization,
-        prompt: str | mx.array,
+        prompt: ta.Union[str, 'mx.array'],
         params: GenerationParams = GenerationParams(),
 ) -> ta.Generator[GenerationOutput]:
     if not isinstance(prompt, mx.array):
@@ -308,9 +310,9 @@ def stream_generate(
 
 
 def generate(
-        model: nn.Module,
+        model: 'mlx_nn.Module',
         tokenization: Tokenization,
-        prompt: str | mx.array,
+        prompt: ta.Union[str, 'mx.array'],
         params: GenerationParams = GenerationParams(),
         *,
         verbose: bool = False,