ommlds 0.0.0.dev440__py3-none-any.whl → 0.0.0.dev480__py3-none-any.whl

ommlds/backends/groq/protocol.py

@@ -0,0 +1,249 @@
+"""
+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_metadata(
+        cls,
+        field_defaults=msh.FieldMetadata(
+            options=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/backends/mlx/generation.py

@@ -254,7 +254,7 @@ def stream_generate(
                 add_special_tokens=add_special_tokens,
             )
 
-        prompt = mx.array(prompt)  # type: ignore[arg-type]
+        prompt = mx.array(ta.cast(ta.Any, prompt))
 
     detokenizer = tokenization.detokenizer
     detokenizer.reset()

ommlds/backends/mlx/loading.py

@@ -15,6 +15,63 @@ from .tokenization import load_tokenization
 ##
 
 
+def get_model_path(
+        path_or_hf_repo: str,
+        revision: str | None = None,
+) -> tuple[pathlib.Path, str | None]:
+    """
+    Ensures the model is available locally. If the path does not exist locally,
+    it is downloaded from the Hugging Face Hub.
+
+    Args:
+        path_or_hf_repo (str): The local path or Hugging Face repository ID of the model.
+        revision (str, optional): A revision id which can be a branch name, a tag, or a commit hash.
+
+    Returns:
+        Tuple[Path, str]: A tuple containing the local file path and the Hugging Face repo ID.
+    """
+
+    model_path = pathlib.Path(path_or_hf_repo)
+
+    if not model_path.exists():
+        from huggingface_hub import snapshot_download
+        hf_path = path_or_hf_repo
+        model_path = pathlib.Path(
+            snapshot_download(
+                path_or_hf_repo,
+                revision=revision,
+                allow_patterns=[
+                    '*.jinja',
+                    '*.json',
+                    '*.jsonl',
+                    '*.py',
+                    '*.txt',
+
+                    'model*.safetensors',
+
+                    '*.tiktoken',
+                    'tiktoken.model',
+                    'tokenizer.model',
+                ],
+            ),
+        )
+
+    else:
+        from huggingface_hub import ModelCard
+
+        card_path = model_path / 'README.md'
+        if card_path.is_file():
+            card = ModelCard.load(card_path)
+            hf_path = card.data.base_model
+        else:
+            hf_path = None
+
+    return model_path, hf_path
+
+
+##
+
+
 @dc.dataclass(frozen=True, kw_only=True)
 class LoadedModel:
     path: pathlib.Path
@@ -46,7 +103,7 @@ def load_model(
 ) -> LoadedModel:
     # FIXME: get_model_path return annotation is wrong:
     #   https://github.com/ml-explore/mlx-lm/blob/9ee2b7358f5e258af7b31a8561acfbbe56ad5085/mlx_lm/utils.py#L82
-    model_path_res = ta.cast(ta.Any, mlx_lm.utils.get_model_path(path_or_hf_repo))
+    model_path_res = ta.cast(ta.Any, get_model_path(path_or_hf_repo))
     if isinstance(model_path_res, tuple):
         model_path = check.isinstance(model_path_res[0], pathlib.Path)
     else:

ommlds/backends/ollama/protocol.py

@@ -0,0 +1,170 @@
+"""
+https://docs.ollama.com/api
+"""
+import typing as ta
+
+from omlish import dataclasses as dc
+from omlish import lang
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@dc.extra_class_params(default_repr_fn=dc.opt_repr)
+class Options:
+    # loading
+    numa: bool | None = None
+    num_ctx: int | None = None
+    num_batch: int | None = None
+    num_gpu: int | None = None
+    main_gpu: int | None = None
+    low_vram: bool | None = None
+    f16_kv: bool | None = None
+    logits_all: bool | None = None
+    vocab_only: bool | None = None
+    use_mmap: bool | None = None
+    use_mlock: bool | None = None
+    embedding_only: bool | None = None
+    num_thread: int | None = None
+
+    # querying
+    num_keep: int | None = None
+    seed: int | None = None
+    num_predict: int | None = None
+    top_k: int | None = None
+    top_p: float | None = None
+    tfs_z: float | None = None
+    typical_p: float | None = None
+    repeat_last_n: int | None = None
+    temperature: float | None = None
+    repeat_penalty: float | None = None
+    presence_penalty: float | None = None
+    frequency_penalty: float | None = None
+    mirostat: int | None = None
+    mirostat_tau: float | None = None
+    mirostat_eta: float | None = None
+    penalize_newline: bool | None = None
+    stop: ta.Sequence[str] | None = None
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class BaseRequest(lang.Abstract):
+    model: str
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class BaseStreamableRequest(BaseRequest, lang.Abstract):
+    stream: bool | None = None
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class BaseGenerateRequest(BaseStreamableRequest, lang.Abstract):
+    options: Options | None = None
+    format: ta.Literal['', 'json'] | None = None  # TODO: jsonschema
+    keep_alive: float | str | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@dc.extra_class_params(default_repr_fn=dc.opt_repr)
+class GenerateRequest(BaseGenerateRequest):
+    prompt: str | None = None
+    suffix: str | None = None
+    system: str | None = None
+    template: str | None = None
+    context: ta.Sequence[int] | None = None
+    raw: bool | None = None
+    images: ta.Sequence[bytes] | None = None
+    think: bool | ta.Literal['low', 'medium', 'high'] | None = None
+
+
+#
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class BaseGenerateResponse(lang.Abstract):
+    model: str | None = None
+    created_at: str | None = None
+    done: bool | None = None
+    done_reason: str | None = None
+    total_duration: int | None = None
+    load_duration: int | None = None
+    prompt_eval_count: int | None = None
+    prompt_eval_duration: int | None = None
+    eval_count: int | None = None
+    eval_duration: int | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@dc.extra_class_params(default_repr_fn=dc.opt_repr)
+class GenerateResponse(BaseGenerateResponse):
+    response: str
+    thinking: str | None = None
+    context: ta.Sequence[int] | None = None
+
+
+##
+
+
+Role: ta.TypeAlias = ta.Literal[
+    'system',
+    'user',
+    'assistant',
+    'tool',
+]
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@dc.extra_class_params(default_repr_fn=dc.opt_repr)
+class Message:
+    role: Role
+    content: str | None = None
+    thinking: str | None = None
+    images: ta.Sequence[bytes] | None = None
+    tool_name: str | None = None
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    class ToolCall:
+        @dc.dataclass(frozen=True, kw_only=True)
+        class Function:
+            name: str
+            arguments: ta.Mapping[str, ta.Any]
+
+        function: Function
+
+    tool_calls: ta.Sequence[ToolCall] | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@dc.extra_class_params(default_repr_fn=dc.opt_repr)
+class Tool:
+    type: str | None = 'function'
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @dc.extra_class_params(default_repr_fn=dc.opt_repr)
+    class Function:
+        name: str | None = None
+        description: str | None = None
+        parameters: ta.Any | None = None
+
+    function: Function | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@dc.extra_class_params(default_repr_fn=dc.opt_repr)
+class ChatRequest(BaseGenerateRequest):
+    messages: ta.Sequence[Message] | None = None
+    tools: ta.Sequence[Tool] | None = None
+    think: bool | ta.Literal['low', 'medium', 'high'] | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@dc.extra_class_params(default_repr_fn=dc.opt_repr)
+class ChatResponse(BaseGenerateResponse):
+    message: Message

ommlds/backends/openai/protocol/__init__.py

@@ -5,11 +5,8 @@ https://github.com/openai/openai-openapi/blob/master/openapi.yaml
 # ruff: noqa: I001
 
 from .chatcompletion.chunk import (  # noqa
-    ChatCompletionChunkChoiceDeltaToolCallFunction,
-    ChatCompletionChunkChoiceDeltaToolCall,
     ChatCompletionChunkChoiceDelta,
 
-    ChatCompletionChunkChoiceLogprobs,
     ChatCompletionChunkChoice,
 
     ChatCompletionChunk,
@@ -18,13 +15,10 @@ from .chatcompletion.chunk import (  # noqa
 from .chatcompletion.contentpart import (  # noqa
     TextChatCompletionContentPart,
 
-    ImageChatCompletionContentPartImageUrl,
-    ImageChatCompletionContentPart,
+    ImageUrlChatCompletionContentPart,
 
-    FileChatCompletionContentPartFileInfo,
     FileChatCompletionContentPart,
 
-    InputAudioChatCompletionContentPartInputAudio,
     InputAudioChatCompletionContentPart,
 
     RefusalChatCompletionContentPart,
@@ -39,9 +33,6 @@ from .chatcompletion.message import (  # noqa
 
     UserChatCompletionMessage,
 
-    AssistantChatCompletionMessageAudio,
-    AssistantChatCompletionMessageToolCallFunction,
-    AssistantChatCompletionMessageToolCall,
     AssistantChatCompletionMessage,
 
     ToolChatCompletionMessage,
@@ -52,18 +43,12 @@ from .chatcompletion.message import (  # noqa
 )
 
 from .chatcompletion.request import (  # noqa
-    ChatCompletionRequestWebSearchOptionsUserLocationApproximate,
-    ChatCompletionRequestWebSearchOptionsUserLocation,
     ChatCompletionRequestWebSearchOptions,
 
     ChatCompletionRequestPrediction,
 
-    ChatCompletionRequestToolFunction,
     ChatCompletionRequestTool,
 
-    ChatCompletionRequestStreamOptions,
-
-    ChatCompletionRequestNamedToolChoiceFunction,
     ChatCompletionRequestNamedToolChoice,
 
     ChatCompletionRequestAudio,
@@ -72,16 +57,8 @@ from .chatcompletion.request import (  # noqa
 )
 
 from .chatcompletion.response import (  # noqa
-    ChatCompletionResponseAnnotationUrlCitation,
-    ChatCompletionResponseAnnotation,
-
-    ChatCompletionResponseAudio,
-
-    ChatCompletionResponseMessageToolCallFunction,
-    ChatCompletionResponseMessageToolCall,
     ChatCompletionResponseMessage,
 
-    ChatCompletionResponseChoiceLogprobs,
     ChatCompletionResponseChoice,
 
     ChatCompletionResponse,
@@ -90,7 +67,6 @@ from .chatcompletion.response import (  # noqa
 from .chatcompletion.responseformat import (  # noqa
     TextChatCompletionResponseFormat,
 
-    JsonSchemaChatCompletionResponseFormatJsonSchema,
     JsonSchemaChatCompletionResponseFormat,
 
     JsonObjectChatCompletionResponseFormat,
@@ -99,12 +75,17 @@ from .chatcompletion.responseformat import (  # noqa
 )
 
 from .chatcompletion.tokenlogprob import (  # noqa
-    ChatCompletionTokenLogprobTopLogprob,
     ChatCompletionTokenLogprob,
 )
 
 from .completionusage import (  # noqa
-    CompletionUsageCompletionTokensDetails,
-    CompletionUsagePromptTokensDetails,
     CompletionUsage,
 )
+
+
+##
+
+
+from omlish import marshal as _msh  # noqa
+
+_msh.register_global_module_import('._marshal', __package__)

ommlds/backends/openai/protocol/_common.py

@@ -0,0 +1,18 @@
+from omlish import lang
+from omlish import marshal as msh
+
+
+##
+
+
+def _set_class_marshal_options(cls):
+    msh.update_object_metadata(
+        cls,
+        field_defaults=msh.FieldMetadata(
+            options=msh.FieldOptions(
+                omit_if=lang.is_none,
+            ),
+        ),
+    )
+
+    return cls

ommlds/backends/openai/protocol/_marshal.py

@@ -0,0 +1,27 @@
+from omlish import lang
+from omlish import marshal as msh
+
+from .chatcompletion.contentpart import ChatCompletionContentPart
+from .chatcompletion.message import ChatCompletionMessage
+from .chatcompletion.responseformat import ChatCompletionResponseFormat
+
+
+##
+
+
+@lang.static_init
+def _install_standard_marshaling() -> None:
+    for root_cls, tag_field in [
+        (ChatCompletionContentPart, 'type'),
+        (ChatCompletionMessage, 'role'),
+        (ChatCompletionResponseFormat, 'type'),
+    ]:
+        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',
+        ))