ommlds 0.0.0.dev465__py3-none-any.whl → 0.0.0.dev467__py3-none-any.whl

ommlds/.omlish-manifests.json

@@ -63,7 +63,7 @@
     "module": ".minichain.backends.impls.anthropic.stream",
     "attr": null,
     "file": "ommlds/minichain/backends/impls/anthropic/stream.py",
-    "line": 35,
+    "line": 36,
     "value": {
       "!.minichain.registries.manifests.RegistryManifest": {
         "module": "ommlds.minichain.backends.impls.anthropic.stream",
@@ -137,7 +137,7 @@
     "module": ".minichain.backends.impls.google.stream",
     "attr": null,
     "file": "ommlds/minichain/backends/impls/google/stream.py",
-    "line": 36,
+    "line": 41,
     "value": {
       "!.minichain.registries.manifests.RegistryManifest": {
         "module": "ommlds.minichain.backends.impls.google.stream",
@@ -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/backends/google/protocol/types.py

@@ -148,11 +148,14 @@ class Part(lang.Final):
     video_metadata: VideoMetadata | None = None
 
 
+ContentRole: ta.TypeAlias = ta.Literal['user', 'model']
+
+
 @dc.dataclass(frozen=True, kw_only=True)
 @_set_class_marshal_options
 class Content(lang.Final):
     parts: ta.Sequence[Part] | None = None
-    role: ta.Literal['user', 'model'] | None = None
+    role: ContentRole | None = None
 
 
 ##

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/cli/sessions/chat/chat/ai/services.py

@@ -63,19 +63,17 @@ class ChatChoicesStreamServiceStreamAiChatGenerator(StreamAiChatGenerator):
     ) -> mc.AiChat:
         opts = self._options() if self._options is not None else []
 
-        lst: list[str] = []
-
         async with self._service_provider.provide_backend() as service:
+            joiner = mc.AiChoiceDeltaJoiner()
+
             async with (await service.invoke(mc.ChatChoicesStreamRequest(chat, opts))).v as st_resp:
                 async for o in st_resp:
+                    joiner.add(o.choices)
+
                     choice = check.single(o.choices)
 
                     for delta in choice.deltas:
                         if delta_callback is not None:
                             await delta_callback(delta)
 
-                    c = check.isinstance(delta, mc.ContentAiChoiceDelta).c  # noqa
-                    if c is not None:
-                        lst.append(check.isinstance(c, str))
-
-        return [mc.AiMessage(''.join(lst))]
+        return check.single(joiner.build())

ommlds/minichain/__init__.py

@@ -92,6 +92,10 @@ with _lang.auto_proxy_init(
         ChatChoicesStreamServiceChatChoicesService,
     )
 
+    from .chat.stream.joining import (  # noqa
+        AiChoiceDeltaJoiner,
+    )
+
     from .chat.stream.services import (  # noqa
         ChatChoicesStreamRequest,
         ChatChoicesStreamResponse,
@@ -109,7 +113,9 @@ with _lang.auto_proxy_init(
 
         AiChoiceDelta,
         ContentAiChoiceDelta,
+        AnyToolUseAiChoiceDelta,
         ToolUseAiChoiceDelta,
+        PartialToolUseAiChoiceDelta,
 
         AiChoiceDeltas,
         AiChoicesDeltas,

ommlds/minichain/backends/impls/anthropic/stream.py

@@ -17,6 +17,7 @@ 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 ....chat.stream.types import PartialToolUseAiChoiceDelta
 from ....chat.tools.types import Tool
 from ....configs import Config
 from ....resources import UseResources
@@ -123,19 +124,35 @@ class AnthropicChatChoicesStreamService:
                                         check.not_none(msg_start)
                                         check.none(cbk_start)
                                         cbk_start = ae
+
                                         if isinstance(ae.content_block, AnthropicSseDecoderEvents.ContentBlockStart.Text):  # noqa
                                             await sink.emit(AiChoicesDeltas([AiChoiceDeltas([ContentAiChoiceDelta(
                                                 ae.content_block.text,
                                             )])]))
+
+                                        elif isinstance(ae.content_block, AnthropicSseDecoderEvents.ContentBlockStart.ToolUse):  # noqa
+                                            await sink.emit(AiChoicesDeltas([AiChoiceDeltas([PartialToolUseAiChoiceDelta(  # noqa
+                                                id=ae.content_block.id,
+                                                name=ae.content_block.name,
+                                                raw_args=ae.content_block.input,
+                                            )])]))
+
                                         else:
                                             raise TypeError(ae.content_block)
 
                                     case AnthropicSseDecoderEvents.ContentBlockDelta():
                                         check.not_none(cbk_start)
+
                                         if isinstance(ae.delta, AnthropicSseDecoderEvents.ContentBlockDelta.TextDelta):
                                             await sink.emit(AiChoicesDeltas([AiChoiceDeltas([ContentAiChoiceDelta(
                                                 ae.delta.text,
                                             )])]))
+
+                                        elif isinstance(ae.delta, AnthropicSseDecoderEvents.ContentBlockDelta.InputJsonDelta):  # noqa
+                                            await sink.emit(AiChoicesDeltas([AiChoiceDeltas([PartialToolUseAiChoiceDelta(  # noqa
+                                                raw_args=ae.delta.partial_json,
+                                            )])]))
+
                                         else:
                                             raise TypeError(ae.delta)
 

ommlds/minichain/backends/impls/google/stream.py

@@ -15,6 +15,8 @@ from ....chat.choices.types import ChatChoicesOutputs
 from ....chat.messages import AiMessage
 from ....chat.messages import Message
 from ....chat.messages import SystemMessage
+from ....chat.messages import ToolUseMessage
+from ....chat.messages import ToolUseResultMessage
 from ....chat.messages import UserMessage
 from ....chat.stream.services import ChatChoicesStreamRequest
 from ....chat.stream.services import ChatChoicesStreamResponse
@@ -22,12 +24,15 @@ 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 ....chat.stream.types import ToolUseAiChoiceDelta
+from ....chat.tools.types import Tool
 from ....models.configs import ModelName
 from ....resources import UseResources
 from ....standard import ApiKey
 from ....stream.services import StreamResponseSink
 from ....stream.services import new_stream_response
 from .names import MODEL_NAMES
+from .tools import build_tool_spec_schema
 
 
 ##
@@ -48,22 +53,70 @@ class GoogleChatChoicesStreamService:
             self._model_name = cc.pop(self.DEFAULT_MODEL_NAME)
             self._api_key = ApiKey.pop_secret(cc, env='GEMINI_API_KEY')
 
-    def _get_msg_content(self, m: Message) -> str | None:
-        if isinstance(m, AiMessage):
-            return check.isinstance(m.c, str)
+    def _make_str_content(
+            self,
+            s: str | None,
+            *,
+            role: pt.ContentRole | None = None,
+    ) -> pt.Content | None:
+        if s is None:
+            return None
+
+        return pt.Content(
+            parts=[pt.Part(
+                text=check.not_none(s),
+            )],
+            role=role,
+        )
 
-        elif isinstance(m, (SystemMessage, UserMessage)):
-            return check.isinstance(m.c, str)
+    def _make_msg_content(self, m: Message) -> pt.Content:
+        if isinstance(m, (AiMessage, SystemMessage, UserMessage)):
+            return check.not_none(self._make_str_content(
+                check.isinstance(m.c, str),
+                role=self.ROLES_MAP[type(m)],
+            ))
+
+        elif isinstance(m, ToolUseResultMessage):
+            tr_resp_val: pt.Value
+            if m.tur.c is None:
+                tr_resp_val = pt.NullValue()  # type: ignore[unreachable]
+            elif isinstance(m.tur.c, str):
+                tr_resp_val = pt.StringValue(m.tur.c)
+            else:
+                raise TypeError(m.tur.c)
+            return pt.Content(
+                parts=[pt.Part(
+                    function_response=pt.FunctionResponse(
+                        id=m.tur.id,
+                        name=m.tur.name,
+                        response={
+                            'value': tr_resp_val,
+                        },
+                    ),
+                )],
+            )
+
+        elif isinstance(m, ToolUseMessage):
+            return pt.Content(
+                parts=[pt.Part(
+                    function_call=pt.FunctionCall(
+                        id=m.tu.id,
+                        name=m.tu.name,
+                        args=m.tu.args,
+                    ),
+                )],
+                role='model',
+            )
 
         else:
             raise TypeError(m)
 
     BASE_URL: ta.ClassVar[str] = 'https://generativelanguage.googleapis.com/v1beta/models'
 
-    ROLES_MAP: ta.ClassVar[ta.Mapping[type[Message], str]] = {
-        SystemMessage: 'system',
+    ROLES_MAP: ta.ClassVar[ta.Mapping[type[Message], pt.ContentRole | None]] = {  # noqa
+        SystemMessage: None,
         UserMessage: 'user',
-        AiMessage: 'assistant',
+        AiMessage: 'model',
     }
 
     READ_CHUNK_SIZE = 64 * 1024
@@ -74,16 +127,28 @@ class GoogleChatChoicesStreamService:
     ) -> ChatChoicesStreamResponse:
         key = check.not_none(self._api_key).reveal()
 
+        msgs = list(request.v)
+
+        system_inst: pt.Content | None = None
+        if msgs and isinstance(m0 := msgs[0], SystemMessage):
+            system_inst = self._make_msg_content(m0)
+            msgs.pop(0)
+
+        g_tools: list[pt.Tool] = []
+        with tv.TypedValues(*request.options).consume() as oc:
+            t: Tool
+            for t in oc.pop(Tool, []):
+                g_tools.append(pt.Tool(
+                    function_declarations=[build_tool_spec_schema(t.spec)],
+                ))
+
         g_req = pt.GenerateContentRequest(
             contents=[
-                pt.Content(
-                    parts=[pt.Part(
-                        text=check.not_none(self._get_msg_content(m)),
-                    )],
-                    role=self.ROLES_MAP[type(m)],  # type: ignore[arg-type]
-                )
-                for m in request.v
+                self._make_msg_content(m)
+                for m in msgs
             ],
+            tools=g_tools or None,
+            system_instruction=system_inst,
         )
 
         req_dct = msh.marshal(g_req)
@@ -110,18 +175,38 @@ class GoogleChatChoicesStreamService:
                         if isinstance(bl, DelimitingBuffer.Incomplete):
                             # FIXME: handle
                             return []
+
                         l = bl.decode('utf-8')
                         if not l:
                             continue
+
                         if l.startswith('data: '):
                             gcr = msh.unmarshal(json.loads(l[6:]), pt.GenerateContentResponse)  # noqa
                             cnd = check.single(check.not_none(gcr.candidates))
+
                             for p in check.not_none(cnd.content).parts or []:
-                                await sink.emit(AiChoicesDeltas([
-                                    AiChoiceDeltas([
-                                        ContentAiChoiceDelta(check.not_none(p.text)),
-                                    ]),
-                                ]))
+                                if (txt := p.text) is not None:
+                                    check.none(p.function_call)
+                                    await sink.emit(AiChoicesDeltas([
+                                        AiChoiceDeltas([
+                                            ContentAiChoiceDelta(check.not_none(txt)),
+                                        ]),
+                                    ]))
+
+                                elif (fc := p.function_call) is not None:
+                                    check.none(p.text)
+                                    await sink.emit(AiChoicesDeltas([
+                                        AiChoiceDeltas([
+                                            ToolUseAiChoiceDelta(
+                                                id=fc.id,
+                                                name=fc.name,
+                                                args=fc.args,
+                                            ),
+                                        ]),
+                                    ]))
+
+                                else:
+                                    raise ValueError(p)
 
                     if not b:
                         return []

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/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