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

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

@@ -0,0 +1,125 @@
+import typing as ta
+
+from omlish import check
+from omlish import marshal as msh
+from omlish import typedvalues as tv
+from omlish.formats import json
+from omlish.http import all as http
+from omlish.http import sse
+from omlish.io.buffers import DelimitingBuffer
+
+from .....backends.groq import protocol as pt
+from ....chat.choices.services import ChatChoicesOutputs
+from ....chat.choices.stream.services import ChatChoicesStreamRequest
+from ....chat.choices.stream.services import ChatChoicesStreamResponse
+from ....chat.choices.stream.services import static_check_is_chat_choices_stream_service
+from ....chat.choices.stream.types import AiChoicesDeltas
+from ....chat.tools.types import Tool
+from ....configs import Config
+from ....resources import UseResources
+from ....standard import ApiKey
+from ....stream.services import StreamResponseSink
+from ....stream.services import new_stream_response
+from .chat import GroqChatChoicesService
+from .names import MODEL_NAMES
+from .protocol import build_gq_request_messages
+from .protocol import build_gq_request_tool
+from .protocol import build_mc_ai_choice_deltas
+
+
+##
+
+
+# @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
+#     name='groq',
+#     type='ChatChoicesStreamService',
+# )
+@static_check_is_chat_choices_stream_service
+class GroqChatChoicesStreamService:
+    def __init__(
+            self,
+            *configs: Config,
+            http_client: http.AsyncHttpClient | None = None,
+    ) -> None:
+        super().__init__()
+
+        self._http_client = http_client
+
+        with tv.consume(*configs) as cc:
+            self._model_name = cc.pop(GroqChatChoicesService.DEFAULT_MODEL_NAME)
+            self._api_key = ApiKey.pop_secret(cc, env='GROQ_API_KEY')
+
+    READ_CHUNK_SIZE: ta.ClassVar[int] = -1
+
+    async def invoke(self, request: ChatChoicesStreamRequest) -> ChatChoicesStreamResponse:
+        tools: list[pt.ChatCompletionRequest.Tool] = []
+        with tv.TypedValues(*request.options).consume() as oc:
+            t: Tool
+            for t in oc.pop(Tool, []):
+                tools.append(build_gq_request_tool(t))
+
+        gq_request = pt.ChatCompletionRequest(
+            messages=build_gq_request_messages(request.v),
+            model=MODEL_NAMES.resolve(self._model_name.v),
+            tools=tools or None,
+            stream=True,
+        )
+
+        raw_request = msh.marshal(gq_request)
+
+        http_request = http.HttpRequest(
+            'https://api.groq.com/openai/v1/chat/completions',
+            headers={
+                http.consts.HEADER_CONTENT_TYPE: http.consts.CONTENT_TYPE_JSON,
+                http.consts.HEADER_AUTH: http.consts.format_bearer_auth_header(check.not_none(self._api_key).reveal()),
+            },
+            data=json.dumps(raw_request).encode('utf-8'),
+        )
+
+        async with UseResources.or_new(request.options) as rs:
+            http_client = await rs.enter_async_context(http.manage_async_client(self._http_client))
+            http_response = await rs.enter_async_context(await http_client.stream_request(http_request))
+
+            async def inner(sink: StreamResponseSink[AiChoicesDeltas]) -> ta.Sequence[ChatChoicesOutputs]:
+                db = DelimitingBuffer([b'\r', b'\n', b'\r\n'])
+                sd = sse.SseDecoder()
+                while True:
+                    b = await http_response.stream.read1(self.READ_CHUNK_SIZE)
+                    for l in db.feed(b):
+                        if isinstance(l, DelimitingBuffer.Incomplete):
+                            # FIXME: handle
+                            return []
+
+                        # FIXME: https://platform.openai.com/docs/guides/function-calling?api-mode=responses#streaming
+                        for so in sd.process_line(l):
+                            if isinstance(so, sse.SseEvent) and so.type == b'message':
+                                ss = so.data.decode('utf-8')
+                                if ss == '[DONE]':
+                                    return []
+
+                                sj = json.loads(ss)  # ChatCompletionChunk
+
+                                check.state(sj['object'] == 'chat.completion.chunk')
+
+                                ccc = msh.unmarshal(sj, pt.ChatCompletionChunk)
+
+                                # FIXME: stop reason
+                                if not ccc.choices:
+                                    continue
+
+                                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([
+                                    build_mc_ai_choice_deltas(choice.delta)
+                                    for choice in ccc.choices
+                                ]))
+
+                    if not b:
+                        return []
+
+            # raw_response = json.loads(check.not_none(http_response.data).decode('utf-8'))
+            # return rh.build_response(raw_response)
+
+            return await new_stream_response(rs, inner)

ommlds/minichain/backends/impls/llamacpp/chat.py

@@ -15,7 +15,8 @@ from ....chat.choices.services import static_check_is_chat_choices_service
 from ....chat.choices.types import AiChoice
 from ....chat.choices.types import ChatChoicesOptions
 from ....chat.messages import AiMessage
-from ....chat.messages import ToolExecResultMessage
+from ....chat.messages import ToolUseMessage
+from ....chat.messages import ToolUseResultMessage
 from ....chat.tools.types import Tool
 from ....configs import Config
 from ....llms.types import MaxTokens
@@ -29,6 +30,15 @@ from .format import get_msg_content
 ##
 
 
+# @omlish-manifest $.minichain.backends.strings.manifests.BackendStringsManifest(
+#     ['ChatChoicesService'],
+#     'llamacpp',
+# )
+
+
+##
+
+
 # @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
 #     name='llamacpp',
 #     type='ChatChoicesService',
@@ -54,7 +64,7 @@ class LlamacppChatChoicesService:
         temperatur=Temperature,
     )
 
-    def invoke(self, request: ChatChoicesRequest) -> ChatChoicesResponse:
+    async def invoke(self, request: ChatChoicesRequest) -> ChatChoicesResponse:
         kwargs: dict = dict(
             # temperature=0,
             max_tokens=1024,
@@ -100,29 +110,34 @@ class LlamacppChatChoicesService:
 
             ims: list = []
             for rm in request.v:
-                if isinstance(rm, ToolExecResultMessage):
+                if isinstance(rm, ToolUseResultMessage):
                     ims.append(dict(
                         role='tool',
-                        **(dict(id=rm.id) if rm.id is not None else {}),
-                        name=rm.name,
-                        content=check.isinstance(rm.c, str),
+                        **(dict(id=rm.tur.id) if rm.tur.id is not None else {}),
+                        name=rm.tur.name,
+                        content=check.isinstance(rm.tur.c, str),
                     ))
+
                 elif isinstance(rm, AiMessage):
-                    tcs: list[dict] = []
-                    for ter in rm.tool_exec_requests or []:
-                        tcs.append(dict(
-                            id=check.not_none(ter.id),
-                            type='function',
-                            function=dict(
-                                name=ter.name,
-                                arguments=check.isinstance(ter.raw_args, str),
-                            ),
-                        ))
                     ims.append(dict(
                         role=ROLES_MAP[type(rm)],
                         **(dict(content=mc) if (mc := get_msg_content(rm)) is not None else {}),
-                        **(dict(tool_calls=tcs) if tcs else {}),
                     ))
+
+                elif isinstance(rm, ToolUseMessage):
+                    ims.append(dict(
+                        role=ROLES_MAP[type(rm)],
+                        content='',
+                        tool_calls=[dict(
+                            id=check.not_none(rm.tu.id),
+                            type='function',
+                            function=dict(
+                                name=rm.tu.name,
+                                arguments=check.isinstance(rm.tu.raw_args, str),
+                            ),
+                        )],
+                    ))
+
                 else:
                     ims.append(dict(
                         role=ROLES_MAP[type(rm)],
@@ -137,6 +152,6 @@ class LlamacppChatChoicesService:
             out: list[AiChoice] = []
             for c in ta.cast(ta.Any, output)['choices']:
                 m = c['message']
-                out.append(AiChoice(AiMessage(m['content'])))
+                out.append(AiChoice([AiMessage(m['content'])]))
 
             return ChatChoicesResponse(out)

ommlds/minichain/backends/impls/llamacpp/completion.py

@@ -51,7 +51,7 @@ class LlamacppCompletionService:
         temperatur=Temperature,
     )
 
-    def invoke(self, request: CompletionRequest) -> CompletionResponse:
+    async def invoke(self, request: CompletionRequest) -> CompletionResponse:
         kwargs: dict = dict(
             # temperature=0,
             max_tokens=1024,

ommlds/minichain/backends/impls/llamacpp/format.py

@@ -5,7 +5,8 @@ from omlish import check
 from ....chat.messages import AiMessage
 from ....chat.messages import Message
 from ....chat.messages import SystemMessage
-from ....chat.messages import ToolExecResultMessage
+from ....chat.messages import ToolUseMessage
+from ....chat.messages import ToolUseResultMessage
 from ....chat.messages import UserMessage
 
 
@@ -16,7 +17,8 @@ ROLES_MAP: ta.Mapping[type[Message], str] = {
     SystemMessage: 'system',
     UserMessage: 'user',
     AiMessage: 'assistant',
-    ToolExecResultMessage: 'tool',
+    ToolUseMessage: 'assistant',
+    ToolUseResultMessage: 'tool',
 }
 
 

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

@@ -10,15 +10,16 @@ from omlish import typedvalues as tv
 
 from .....backends import llamacpp as lcu
 from ....chat.choices.services import ChatChoicesOutputs
-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 AiChoiceDelta
-from ....chat.stream.types import AiChoiceDeltas
-from ....chat.stream.types import AiMessageDelta
+from ....chat.choices.stream.services import ChatChoicesStreamRequest
+from ....chat.choices.stream.services import ChatChoicesStreamResponse
+from ....chat.choices.stream.services import static_check_is_chat_choices_stream_service
+from ....chat.choices.stream.types import AiChoiceDeltas
+from ....chat.choices.stream.types import AiChoicesDeltas
+from ....chat.stream.types import ContentAiDelta
 from ....configs import Config
 from ....models.configs import ModelPath
 from ....resources import UseResources
+from ....stream.services import StreamResponseSink
 from ....stream.services import new_stream_response
 from .chat import LlamacppChatChoicesService
 from .format import ROLES_MAP
@@ -28,6 +29,15 @@ from .format import get_msg_content
 ##
 
 
+# @omlish-manifest $.minichain.backends.strings.manifests.BackendStringsManifest(
+#     ['ChatChoicesStreamService'],
+#     'llamacpp',
+# )
+
+
+##
+
+
 # @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
 #     name='llamacpp',
 #     type='ChatChoicesStreamService',
@@ -49,10 +59,10 @@ class LlamacppChatChoicesStreamService(lang.ExitStacked):
             verbose=False,
         )))
 
-    def invoke(self, request: ChatChoicesStreamRequest) -> ChatChoicesStreamResponse:
+    async def invoke(self, request: ChatChoicesStreamRequest) -> ChatChoicesStreamResponse:
         lcu.install_logging_hook()
 
-        with UseResources.or_new(request.options) as rs:
+        async with UseResources.or_new(request.options) as rs:
             rs.enter_context(self._lock)
 
             model: ta.Any = self._load_model()  # FIXME: the types are awful lol
@@ -74,19 +84,26 @@ class LlamacppChatChoicesStreamService(lang.ExitStacked):
 
             rs.enter_context(lang.defer(close_output))
 
-            def yield_choices() -> ta.Generator[AiChoiceDeltas, None, ta.Sequence[ChatChoicesOutputs] | None]:
+            async def inner(sink: StreamResponseSink[AiChoicesDeltas]) -> ta.Sequence[ChatChoicesOutputs] | None:
+                last_role: ta.Any = None
+
                 for chunk in output:
                     check.state(chunk['object'] == 'chat.completion.chunk')
-                    l: list[AiChoiceDelta] = []
-                    for choice in chunk['choices']:
-                        # FIXME: check role is assistant
-                        # FIXME: stop reason
-                        if not (delta := choice.get('delta', {})):
-                            continue
-                        if not (content := delta.get('content', '')):
-                            continue
-                        l.append(AiChoiceDelta(AiMessageDelta(content)))
-                    yield l
+
+                    choice = check.single(chunk['choices'])
+
+                    if not (delta := choice.get('delta', {})):
+                        continue
+
+                    # FIXME: check role is assistant
+                    if (role := delta.get('role')) != last_role:
+                        last_role = role
+
+                    # FIXME: stop reason
+
+                    if (content := delta.get('content', '')):
+                        await sink.emit(AiChoicesDeltas([AiChoiceDeltas([ContentAiDelta(content)])]))
+
                 return None
 
-            return new_stream_response(rs, yield_choices())
+            return await new_stream_response(rs, inner)

ommlds/minichain/backends/impls/mistral.py

@@ -21,13 +21,18 @@ from ...chat.messages import UserMessage
 ##
 
 
+# TODO: generalize lol
+class TooManyRequestsMistralError(Exception):
+    pass
+
+
 # @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
 #     name='mistral',
 #     type='ChatChoicesService',
 # )
 @static_check_is_chat_choices_service
 class MistralChatChoicesService:
-    model: ta.ClassVar[str] = 'mistral-large-latest'
+    model: ta.ClassVar[str] = 'mistral-medium-2508'
 
     ROLES_MAP: ta.ClassVar[ta.Mapping[type[Message], str]] = {
         SystemMessage: 'system',
@@ -35,10 +40,16 @@ class MistralChatChoicesService:
         AiMessage: 'assistant',
     }
 
-    def __init__(self, *, api_key: str | None = None) -> None:
+    def __init__(
+            self,
+            *,
+            api_key: str | None = None,
+            http_client: http.AsyncHttpClient | None = None,
+    ) -> None:
         super().__init__()
 
         self._api_key = api_key
+        self._http_client = http_client
 
     def _get_msg_content(self, m: Message) -> str | None:
         if isinstance(m, AiMessage):
@@ -50,7 +61,7 @@ class MistralChatChoicesService:
         else:
             raise TypeError(m)
 
-    def invoke(
+    async def invoke(
             self,
             request: ChatChoicesRequest,
     ) -> ChatChoicesResponse:
@@ -68,7 +79,7 @@ class MistralChatChoicesService:
             ],
         }
 
-        resp = http.request(
+        resp = await http.async_request(
             'https://api.mistral.ai/v1/chat/completions',
             method='POST',
             data=json.dumps_compact(req_dct).encode('utf-8'),
@@ -77,11 +88,15 @@ class MistralChatChoicesService:
                 'Accept': 'application/json',
                 'Authorization': f'Bearer {key}',
             },
+            client=self._http_client,
         )
 
+        if resp.status == 429:
+            raise TooManyRequestsMistralError
+
         resp_dct = json.loads(check.not_none(resp.data).decode('utf-8'))
 
         return ChatChoicesResponse([
-            AiChoice(AiMessage(c['message']['content']))
+            AiChoice([AiMessage(c['message']['content'])])
             for c in resp_dct['choices']
         ])

ommlds/minichain/backends/impls/mlx/chat.py

@@ -1,3 +1,4 @@
+import contextlib
 import typing as ta
 
 from omlish import check
@@ -5,28 +6,38 @@ 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
+from ....chat.choices.stream.services import ChatChoicesStreamRequest
+from ....chat.choices.stream.services import ChatChoicesStreamResponse
+from ....chat.choices.stream.services import static_check_is_chat_choices_stream_service
+from ....chat.choices.stream.types import AiChoiceDeltas
+from ....chat.choices.stream.types import AiChoicesDeltas
 from ....chat.choices.types import AiChoice
 from ....chat.choices.types import ChatChoicesOptions
 from ....chat.messages import AiMessage
 from ....chat.messages import Message
 from ....chat.messages import SystemMessage
 from ....chat.messages import UserMessage
+from ....chat.stream.types import ContentAiDelta
 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,
     )
 
-    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,
@@ -137,5 +157,59 @@ class MlxChatChoicesService(lang.ExitStacked):
         )
 
         return ChatChoicesResponse([
-            AiChoice(AiMessage(response))  # noqa
+            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([
+                            ContentAiDelta(go.text),
+                        ])]))
+
+                return []
+
+            return await new_stream_response(rs, inner)