ommlds 0.0.0.dev466__tar.gz → 0.0.0.dev467__tar.gz

{ommlds-0.0.0.dev466/ommlds.egg-info → ommlds-0.0.0.dev467}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ommlds
-Version: 0.0.0.dev466
+Version: 0.0.0.dev467
 Summary: ommlds
 Author: wrmsr
 License-Expression: BSD-3-Clause
@@ -14,8 +14,8 @@ Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.13
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: omdev==0.0.0.dev466
-Requires-Dist: omlish==0.0.0.dev466
+Requires-Dist: omdev==0.0.0.dev467
+Requires-Dist: omlish==0.0.0.dev467
 Provides-Extra: all
 Requires-Dist: llama-cpp-python~=0.3; extra == "all"
 Requires-Dist: mlx~=0.29; extra == "all"

{ommlds-0.0.0.dev466 → ommlds-0.0.0.dev467}/ommlds/.omlish-manifests.json

@@ -229,11 +229,12 @@
     "module": ".minichain.backends.impls.mlx.chat",
     "attr": null,
     "file": "ommlds/minichain/backends/impls/mlx/chat.py",
-    "line": 28,
+    "line": 39,
     "value": {
       "!.minichain.backends.strings.manifests.BackendStringsManifest": {
         "service_cls_names": [
-          "ChatChoicesService"
+          "ChatChoicesService",
+          "ChatChoicesStreamService"
         ],
         "backend_name": "mlx",
         "model_names": null
@@ -244,7 +245,7 @@
     "module": ".minichain.backends.impls.mlx.chat",
     "attr": null,
     "file": "ommlds/minichain/backends/impls/mlx/chat.py",
-    "line": 37,
+    "line": 133,
     "value": {
       "!.minichain.registries.manifests.RegistryManifest": {
         "module": "ommlds.minichain.backends.impls.mlx.chat",
@@ -255,6 +256,21 @@
       }
     }
   },
+  {
+    "module": ".minichain.backends.impls.mlx.chat",
+    "attr": null,
+    "file": "ommlds/minichain/backends/impls/mlx/chat.py",
+    "line": 164,
+    "value": {
+      "!.minichain.registries.manifests.RegistryManifest": {
+        "module": "ommlds.minichain.backends.impls.mlx.chat",
+        "attr": "MlxChatChoicesStreamService",
+        "name": "mlx",
+        "aliases": null,
+        "type": "ChatChoicesStreamService"
+      }
+    }
+  },
   {
     "module": ".minichain.backends.impls.openai.chat",
     "attr": null,
@@ -356,7 +372,7 @@
     "module": ".minichain.backends.impls.openai.stream",
     "attr": null,
     "file": "ommlds/minichain/backends/impls/openai/stream.py",
-    "line": 37,
+    "line": 38,
     "value": {
       "!.minichain.registries.manifests.RegistryManifest": {
         "module": "ommlds.minichain.backends.impls.openai.stream",

{ommlds-0.0.0.dev466 → ommlds-0.0.0.dev467}/ommlds/cli/sessions/chat/backends/catalog.py

@@ -32,7 +32,7 @@ class _CatalogBackendProvider(BackendProvider[ServiceT], lang.Abstract):
     @contextlib.asynccontextmanager
     async def _provide_backend(self, cls: type[ServiceT]) -> ta.AsyncIterator[ServiceT]:
         service: ServiceT
-        async with lang.async_maybe_managing(self._catalog.get_backend(
+        async with lang.async_or_sync_maybe_managing(self._catalog.get_backend(
                 cls,
                 self._name,
                 *(self._configs or []),

{ommlds-0.0.0.dev466 → ommlds-0.0.0.dev467}/ommlds/minichain/backends/impls/mlx/chat.py

@@ -1,3 +1,4 @@
+import contextlib
 import typing as ta
 
 from omlish import check
@@ -5,6 +6,7 @@ from omlish import lang
 from omlish import typedvalues as tv
 
 from .....backends import mlx as mlxu
+from ....chat.choices.services import ChatChoicesOutputs
 from ....chat.choices.services import ChatChoicesRequest
 from ....chat.choices.services import ChatChoicesResponse
 from ....chat.choices.services import static_check_is_chat_choices_service
@@ -14,19 +16,28 @@ from ....chat.messages import AiMessage
 from ....chat.messages import Message
 from ....chat.messages import SystemMessage
 from ....chat.messages import UserMessage
+from ....chat.stream.services import ChatChoicesStreamRequest
+from ....chat.stream.services import ChatChoicesStreamResponse
+from ....chat.stream.services import static_check_is_chat_choices_stream_service
+from ....chat.stream.types import AiChoiceDeltas
+from ....chat.stream.types import AiChoicesDeltas
+from ....chat.stream.types import ContentAiChoiceDelta
 from ....configs import Config
 from ....llms.types import MaxTokens
 from ....models.configs import ModelPath
 from ....models.configs import ModelRepo
 from ....models.configs import ModelSpecifier
+from ....resources import UseResources
 from ....standard import DefaultOptions
+from ....stream.services import StreamResponseSink
+from ....stream.services import new_stream_response
 
 
 ##
 
 
 # @omlish-manifest $.minichain.backends.strings.manifests.BackendStringsManifest(
-#     ['ChatChoicesService'],
+#     ['ChatChoicesService', 'ChatChoicesStreamService'],
 #     'mlx',
 # )
 
@@ -34,12 +45,7 @@ from ....standard import DefaultOptions
 ##
 
 
-# @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
-#     name='mlx',
-#     type='ChatChoicesService',
-# )
-@static_check_is_chat_choices_service
-class MlxChatChoicesService(lang.ExitStacked):
+class BaseMlxChatChoicesService(lang.ExitStacked):
     DEFAULT_MODEL: ta.ClassVar[ModelSpecifier] = (
         # 'mlx-community/DeepSeek-Coder-V2-Lite-Instruct-8bit'
         # 'mlx-community/Llama-3.3-70B-Instruct-4bit'
@@ -52,8 +58,8 @@ class MlxChatChoicesService(lang.ExitStacked):
         # 'mlx-community/Qwen2.5-0.5B-4bit'
         # 'mlx-community/Qwen2.5-32B-Instruct-8bit'
         # 'mlx-community/Qwen2.5-Coder-32B-Instruct-8bit'
-        # 'mlx-community/mamba-2.8b-hf-f16'
         # 'mlx-community/Qwen3-30B-A3B-6bit'
+        # 'mlx-community/mamba-2.8b-hf-f16'
     )
 
     def __init__(self, *configs: Config) -> None:
@@ -70,10 +76,7 @@ class MlxChatChoicesService(lang.ExitStacked):
     }
 
     def _get_msg_content(self, m: Message) -> str | None:
-        if isinstance(m, AiMessage):
-            return check.isinstance(m.c, str)
-
-        elif isinstance(m, (SystemMessage, UserMessage)):
+        if isinstance(m, (AiMessage, SystemMessage, UserMessage)):
             return check.isinstance(m.c, str)
 
         else:
@@ -96,10 +99,9 @@ class MlxChatChoicesService(lang.ExitStacked):
         max_tokens=MaxTokens,
     )
 
-    async def invoke(self, request: ChatChoicesRequest) -> ChatChoicesResponse:
-        loaded_model = self._load_model()
-
-        tokenizer = loaded_model.tokenization.tokenizer
+    @lang.cached_function(transient=True)
+    def _get_tokenizer(self) -> mlxu.tokenization.Tokenizer:
+        tokenizer = self._load_model().tokenization.tokenizer
 
         if not (
                 hasattr(tokenizer, 'apply_chat_template') and
@@ -107,26 +109,44 @@ class MlxChatChoicesService(lang.ExitStacked):
         ):
             raise RuntimeError(tokenizer)
 
-        prompt = tokenizer.apply_chat_template(
+        return tokenizer
+
+    def _build_prompt(self, messages: ta.Sequence[Message]) -> str:
+        return check.isinstance(self._get_tokenizer().apply_chat_template(
             [  # type: ignore[arg-type]
                 dict(
                     role=self.ROLES_MAP[type(m)],
                     content=self._get_msg_content(m),
                 )
-                for m in request.v
+                for m in messages
             ],
             tokenize=False,
             add_generation_prompt=True,
-        )
+        ), str)
 
-        kwargs = dict()
+    def _build_kwargs(self, oc: tv.TypedValuesConsumer) -> dict[str, ta.Any]:
+        kwargs: dict[str, ta.Any] = {}
+        kwargs.update(oc.pop_scalar_kwargs(**self._OPTION_KWARG_NAMES_MAP))
+        return kwargs
+
+
+# @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
+#     name='mlx',
+#     type='ChatChoicesService',
+# )
+@static_check_is_chat_choices_service
+class MlxChatChoicesService(BaseMlxChatChoicesService):
+    async def invoke(self, request: ChatChoicesRequest) -> ChatChoicesResponse:
+        loaded_model = self._load_model()
+
+        prompt = self._build_prompt(request.v)
 
         with tv.consume(
                 *self._default_options,
                 *request.options,
                 override=True,
         ) as oc:
-            kwargs.update(oc.pop_scalar_kwargs(**self._OPTION_KWARG_NAMES_MAP))
+            kwargs = self._build_kwargs(oc)
 
         response = mlxu.generate(
             loaded_model.model,
@@ -139,3 +159,57 @@ class MlxChatChoicesService(lang.ExitStacked):
         return ChatChoicesResponse([
             AiChoice([AiMessage(response)])  # noqa
         ])
+
+
+# @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
+#     name='mlx',
+#     type='ChatChoicesStreamService',
+# )
+@static_check_is_chat_choices_stream_service
+class MlxChatChoicesStreamService(BaseMlxChatChoicesService):
+    def __init__(self, *configs: Config) -> None:
+        super().__init__()
+
+        with tv.consume(*configs) as cc:
+            self._model = cc.pop(MlxChatChoicesService.DEFAULT_MODEL)
+            self._default_options: tv.TypedValues = DefaultOptions.pop(cc)
+
+    READ_CHUNK_SIZE = 64 * 1024
+
+    async def invoke(
+            self,
+            request: ChatChoicesStreamRequest,
+            *,
+            max_tokens: int = 4096,  # FIXME: ChatOption
+    ) -> ChatChoicesStreamResponse:
+        loaded_model = self._load_model()
+
+        prompt = self._build_prompt(request.v)
+
+        with tv.consume(
+                *self._default_options,
+                *request.options,
+                override=True,
+        ) as oc:
+            oc.pop(UseResources, None)
+            kwargs = self._build_kwargs(oc)
+
+        async with UseResources.or_new(request.options) as rs:
+            gen: ta.Iterator[mlxu.GenerationOutput] = rs.enter_context(contextlib.closing(mlxu.stream_generate(
+                loaded_model.model,
+                loaded_model.tokenization,
+                check.isinstance(prompt, str),
+                mlxu.GenerationParams(**kwargs),
+                # verbose=True,
+            )))
+
+            async def inner(sink: StreamResponseSink[AiChoicesDeltas]) -> ta.Sequence[ChatChoicesOutputs]:
+                for go in gen:
+                    if go.text:
+                        await sink.emit(AiChoicesDeltas([AiChoiceDeltas([
+                            ContentAiChoiceDelta(go.text),
+                        ])]))
+
+                return []
+
+            return await new_stream_response(rs, inner)

{ommlds-0.0.0.dev466 → ommlds-0.0.0.dev467}/ommlds/minichain/backends/impls/openai/chat.py

@@ -26,8 +26,8 @@ from ....chat.choices.services import static_check_is_chat_choices_service
 from ....models.configs import ModelName
 from ....standard import ApiKey
 from ....standard import DefaultOptions
-from .format2 import OpenaiChatRequestHandler
-from .format2 import build_mc_choices_response
+from .format import OpenaiChatRequestHandler
+from .format import build_mc_choices_response
 from .names import MODEL_NAMES
 
 

ommlds-0.0.0.dev466/ommlds/minichain/backends/impls/openai/format2.py → ommlds-0.0.0.dev467/ommlds/minichain/backends/impls/openai/format.py

@@ -17,6 +17,9 @@ from ....chat.messages import SystemMessage
 from ....chat.messages import ToolUseMessage
 from ....chat.messages import ToolUseResultMessage
 from ....chat.messages import UserMessage
+from ....chat.stream.types import AiChoiceDelta
+from ....chat.stream.types import ContentAiChoiceDelta
+from ....chat.stream.types import PartialToolUseAiChoiceDelta
 from ....chat.tools.types import Tool
 from ....content.json import JsonContent
 from ....content.prepare import prepare_content_str
@@ -125,6 +128,25 @@ def build_mc_choices_response(oai_resp: pt.ChatCompletionResponse) -> ChatChoice
     )
 
 
+def build_mc_ai_choice_delta(delta: pt.ChatCompletionChunkChoiceDelta) -> AiChoiceDelta:
+    if delta.content is not None:
+        check.state(not delta.tool_calls)
+        return ContentAiChoiceDelta(delta.content)
+
+    elif delta.tool_calls is not None:
+        check.state(delta.content is None)
+        tc = check.single(delta.tool_calls)
+        tc_fn = check.not_none(tc.function)
+        return PartialToolUseAiChoiceDelta(
+            id=tc.id,
+            name=tc_fn.name,
+            raw_args=tc_fn.arguments,
+        )
+
+    else:
+        raise ValueError(delta)
+
+
 ##
 
 

{ommlds-0.0.0.dev466 → ommlds-0.0.0.dev467}/ommlds/minichain/backends/impls/openai/stream.py

@@ -11,7 +11,7 @@ from omlish.http import all as http
 from omlish.http import sse
 from omlish.io.buffers import DelimitingBuffer
 
-from .....backends.openai.protocol.chatcompletion.chunk import ChatCompletionChunk
+from .....backends.openai import protocol as pt
 from ....chat.choices.services import ChatChoicesOutputs
 from ....chat.stream.services import ChatChoicesStreamRequest
 from ....chat.stream.services import ChatChoicesStreamResponse
@@ -28,6 +28,7 @@ from ....stream.services import StreamResponseSink
 from ....stream.services import new_stream_response
 from .chat import OpenaiChatChoicesService
 from .format import OpenaiChatRequestHandler
+from .format import build_mc_ai_choice_delta
 from .names import MODEL_NAMES
 
 
@@ -62,13 +63,13 @@ class OpenaiChatChoicesStreamService:
             model=MODEL_NAMES.resolve(self._model_name.v),
             mandatory_kwargs=dict(
                 stream=True,
-                stream_options=dict(
+                stream_options=pt.ChatCompletionRequest.StreamOptions(
                     include_usage=True,
                 ),
             ),
         )
 
-        raw_request = rh.raw_request()
+        raw_request = msh.marshal(rh.oai_request())
 
         http_request = http.HttpRequest(
             'https://api.openai.com/v1/chat/completions',
@@ -105,20 +106,20 @@ class OpenaiChatChoicesStreamService:
 
                                 check.state(sj['object'] == 'chat.completion.chunk')
 
-                                ccc = msh.unmarshal(sj, ChatCompletionChunk)  # noqa
-                                # print(ccc)
+                                ccc = msh.unmarshal(sj, pt.ChatCompletionChunk)
 
                                 # FIXME: stop reason
-                                if not sj['choices']:
+                                if not ccc.choices:
                                     continue
 
-                                if any(choice['delta'] for choice in sj['choices']):
-                                    await sink.emit(AiChoicesDeltas([
-                                        AiChoiceDeltas(
-                                            [rh.build_ai_choice_delta(choice['delta'])] if choice['delta'] else [],
-                                        )
-                                        for choice in sj['choices']
-                                    ]))
+                                if any(choice.finish_reason for choice in ccc.choices):
+                                    check.state(all(choice.finish_reason for choice in ccc.choices))
+                                    break
+
+                                await sink.emit(AiChoicesDeltas([
+                                    AiChoiceDeltas([build_mc_ai_choice_delta(choice.delta)])
+                                    for choice in ccc.choices
+                                ]))
 
                     if not b:
                         return []

{ommlds-0.0.0.dev466 → ommlds-0.0.0.dev467}/ommlds/minichain/chat/stream/types.py

@@ -42,6 +42,9 @@ class AiChoiceDelta(lang.Sealed, lang.Abstract):
     pass
 
 
+#
+
+
 @dc.dataclass(frozen=True)
 class ContentAiChoiceDelta(AiChoiceDelta, lang.Final):
     c: Content

{ommlds-0.0.0.dev466 → ommlds-0.0.0.dev467/ommlds.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ommlds
-Version: 0.0.0.dev466
+Version: 0.0.0.dev467
 Summary: ommlds
 Author: wrmsr
 License-Expression: BSD-3-Clause
@@ -14,8 +14,8 @@ Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.13
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: omdev==0.0.0.dev466
-Requires-Dist: omlish==0.0.0.dev466
+Requires-Dist: omdev==0.0.0.dev467
+Requires-Dist: omlish==0.0.0.dev467
 Provides-Extra: all
 Requires-Dist: llama-cpp-python~=0.3; extra == "all"
 Requires-Dist: mlx~=0.29; extra == "all"

{ommlds-0.0.0.dev466 → ommlds-0.0.0.dev467}/ommlds.egg-info/SOURCES.txt

@@ -206,7 +206,6 @@ ommlds/minichain/backends/impls/openai/chat.py
 ommlds/minichain/backends/impls/openai/completion.py
 ommlds/minichain/backends/impls/openai/embedding.py
 ommlds/minichain/backends/impls/openai/format.py
-ommlds/minichain/backends/impls/openai/format2.py
 ommlds/minichain/backends/impls/openai/names.py
 ommlds/minichain/backends/impls/openai/stream.py
 ommlds/minichain/backends/impls/sentencepiece/__init__.py

{ommlds-0.0.0.dev466 → ommlds-0.0.0.dev467}/ommlds.egg-info/requires.txt

@@ -1,5 +1,5 @@
-omdev==0.0.0.dev466
-omlish==0.0.0.dev466
+omdev==0.0.0.dev467
+omlish==0.0.0.dev467
 
 [all]
 llama-cpp-python~=0.3

{ommlds-0.0.0.dev466 → ommlds-0.0.0.dev467}/pyproject.toml

@@ -13,7 +13,7 @@ urls = {source = 'https://github.com/wrmsr/omlish'}
 license = 'BSD-3-Clause'
 readme = 'README.md'
 requires-python = '>=3.13'
-version = '0.0.0.dev466'
+version = '0.0.0.dev467'
 classifiers = [
     'Development Status :: 2 - Pre-Alpha',
     'Intended Audience :: Developers',
@@ -24,8 +24,8 @@ classifiers = [
 ]
 description = 'ommlds'
 dependencies = [
-    'omdev == 0.0.0.dev466',
-    'omlish == 0.0.0.dev466',
+    'omdev == 0.0.0.dev467',
+    'omlish == 0.0.0.dev467',
 ]
 
 [project.optional-dependencies]