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

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

@@ -2,23 +2,24 @@ import typing as ta
 
 from omlish import cached
 from omlish import check
-from omlish import lang
 from omlish import typedvalues as tv
 from omlish.formats import json
 
+from .....backends.openai import protocol as pt
 from ....chat.choices.services import ChatChoicesResponse
 from ....chat.choices.types import AiChoice
+from ....chat.choices.types import AiChoices
 from ....chat.choices.types import ChatChoicesOptions
-from ....chat.messages import AiChat
 from ....chat.messages import AiMessage
 from ....chat.messages import AnyAiMessage
 from ....chat.messages import Chat
-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.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
@@ -26,7 +27,7 @@ from ....llms.types import MaxTokens
 from ....llms.types import Temperature
 from ....llms.types import TokenUsage
 from ....llms.types import TokenUsageOutput
-from ....tools.jsonschema import build_tool_spec_json_schema
+from ....tools.jsonschema import build_tool_spec_params_json_schema
 from ....tools.types import ToolSpec
 from ....tools.types import ToolUse
 from ....types import Option
@@ -35,61 +36,115 @@ from ....types import Option
 ##
 
 
-def build_request_messages(chat: Chat) -> ta.Sequence[ta.Mapping[str, ta.Any]]:
-    out: list[dict[str, ta.Any]] = []
+def build_oai_request_msgs(mc_chat: Chat) -> ta.Sequence[pt.ChatCompletionMessage]:
+    oai_msgs: list[pt.ChatCompletionMessage] = []
 
-    for m in chat:
-        if isinstance(m, SystemMessage):
-            out.append(dict(
-                role='system',
-                content=m.c,
+    for mc_msg in mc_chat:
+        if isinstance(mc_msg, SystemMessage):
+            oai_msgs.append(pt.SystemChatCompletionMessage(
+                content=check.isinstance(mc_msg.c, str),
             ))
 
-        elif isinstance(m, AiMessage):
-            out.append(dict(
-                role='assistant',
-                content=check.isinstance(m.c, (str, None)),
+        elif isinstance(mc_msg, AiMessage):
+            oai_msgs.append(pt.AssistantChatCompletionMessage(
+                content=check.isinstance(mc_msg.c, (str, None)),
             ))
 
-        elif isinstance(m, ToolUseMessage):
-            out.append(dict(
-                role='assistant',
-                tool_calls=[
-                    dict(
-                        id=m.tu.id,
-                        function=dict(
-                            arguments=check.not_none(m.tu.raw_args),
-                            name=m.tu.name,
-                        ),
-                        type='function',
+        elif isinstance(mc_msg, ToolUseMessage):
+            oai_msgs.append(pt.AssistantChatCompletionMessage(
+                tool_calls=[pt.AssistantChatCompletionMessage.ToolCall(
+                    id=check.not_none(mc_msg.tu.id),
+                    function=pt.AssistantChatCompletionMessage.ToolCall.Function(
+                        arguments=check.not_none(mc_msg.tu.raw_args),
+                        name=mc_msg.tu.name,
                     ),
-                ],
+                )],
             ))
 
-        elif isinstance(m, UserMessage):
-            out.append(dict(
-                role='user',
-                content=prepare_content_str(m.c),
+        elif isinstance(mc_msg, UserMessage):
+            oai_msgs.append(pt.UserChatCompletionMessage(
+                content=prepare_content_str(mc_msg.c),
             ))
 
-        elif isinstance(m, ToolUseResultMessage):
+        elif isinstance(mc_msg, ToolUseResultMessage):
             tc: str
-            if isinstance(m.tur.c, str):
-                tc = m.tur.c
-            elif isinstance(m.tur.c, JsonContent):
-                tc = json.dumps_compact(m.tur.c)
+            if isinstance(mc_msg.tur.c, str):
+                tc = mc_msg.tur.c
+            elif isinstance(mc_msg.tur.c, JsonContent):
+                tc = json.dumps_compact(mc_msg.tur.c)
             else:
-                raise TypeError(m.tur.c)
-            out.append(dict(
-                role='tool',
-                tool_call_id=m.tur.id,
+                raise TypeError(mc_msg.tur.c)
+            oai_msgs.append(pt.ToolChatCompletionMessage(
+                tool_call_id=check.not_none(mc_msg.tur.id),
                 content=tc,
             ))
 
         else:
-            raise TypeError(m)
+            raise TypeError(mc_msg)
 
-    return out
+    return oai_msgs
+
+
+#
+
+
+def build_mc_ai_choice(oai_choice: pt.ChatCompletionResponseChoice) -> AiChoice:
+    cur: list[AnyAiMessage] = []
+
+    oai_msg = oai_choice.message
+
+    if (oai_c := oai_msg.content) is not None:
+        cur.append(AiMessage(check.isinstance(oai_c, str)))
+
+    for oai_tc in oai_msg.tool_calls or []:
+        cur.append(ToolUseMessage(ToolUse(
+            id=oai_tc.id,
+            name=oai_tc.function.name,
+            args=json.loads(oai_tc.function.arguments or '{}'),
+            raw_args=oai_tc.function.arguments,
+        )))
+
+    return AiChoice(cur)
+
+
+def build_mc_ai_choices(oai_resp: pt.ChatCompletionResponse) -> AiChoices:
+    return [
+        build_mc_ai_choice(oai_choice)
+        for oai_choice in oai_resp.choices
+    ]
+
+
+def build_mc_choices_response(oai_resp: pt.ChatCompletionResponse) -> ChatChoicesResponse:
+    return ChatChoicesResponse(
+        build_mc_ai_choices(oai_resp),
+
+        tv.TypedValues(
+            *([TokenUsageOutput(TokenUsage(
+                input=tu.prompt_tokens,
+                output=tu.completion_tokens,
+                total=tu.total_tokens,
+            ))] if (tu := oai_resp.usage) is not None else []),
+        ),
+    )
+
+
+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)
 
 
 ##
@@ -110,14 +165,6 @@ class OpenaiChatRequestHandler:
         self._model = model
         self._mandatory_kwargs = mandatory_kwargs
 
-    ROLES_MAP: ta.ClassVar[ta.Mapping[type[Message], str]] = {
-        SystemMessage: 'system',
-        UserMessage: 'user',
-        AiMessage: 'assistant',
-        ToolUseMessage: 'assistant',
-        ToolUseResultMessage: 'tool',
-    }
-
     DEFAULT_OPTIONS: ta.ClassVar[tv.TypedValues[Option]] = tv.TypedValues[Option](
         Temperature(0.),
         MaxTokens(1024),
@@ -160,69 +207,26 @@ class OpenaiChatRequestHandler:
         )
 
     @cached.function
-    def raw_request(self) -> ta.Mapping[str, ta.Any]:
+    def oai_request(self) -> pt.ChatCompletionRequest:
         po = self._process_options()
 
-        tools = [
-            dict(
-                type='function',
-                function=build_tool_spec_json_schema(ts),
+        tools: list[pt.ChatCompletionRequestTool] = [
+            pt.ChatCompletionRequestTool(
+                function=pt.ChatCompletionRequestTool.Function(
+                    name=check.not_none(ts.name),
+                    description=prepare_content_str(ts.desc),
+                    parameters=build_tool_spec_params_json_schema(ts),
+                ),
             )
             for ts in po.tools_by_name.values()
         ]
 
-        return dict(
+        return pt.ChatCompletionRequest(
             model=self._model,
-            messages=build_request_messages(self._chat),
+            messages=build_oai_request_msgs(self._chat),
             top_p=1,
-            **lang.opt_kw(tools=tools),
+            tools=tools or None,
             frequency_penalty=0.0,
             presence_penalty=0.0,
             **po.kwargs,
         )
-
-    def build_ai_chat(self, message: ta.Mapping[str, ta.Any]) -> AiChat:
-        out: list[AnyAiMessage] = []
-        if (c := message.get('content')) is not None:
-            out.append(AiMessage(c))
-        for tc in message.get('tool_calls', []):
-            out.append(ToolUseMessage(
-                ToolUse(
-                    id=tc['id'],
-                    name=tc['function']['name'],
-                    args=json.loads(tc['function']['arguments'] or '{}'),
-                    raw_args=tc['function']['arguments'],
-                ),
-            ))
-        return out
-
-    def build_response(self, raw_response: ta.Mapping[str, ta.Any]) -> ChatChoicesResponse:
-        return ChatChoicesResponse(
-            [
-                AiChoice(self.build_ai_chat(choice['message']))
-                for choice in raw_response['choices']
-            ],
-
-            tv.TypedValues(
-                *([TokenUsageOutput(TokenUsage(
-                    input=tu['prompt_tokens'],
-                    output=tu['completion_tokens'],
-                    total=tu['total_tokens'],
-                ))] if (tu := raw_response.get('usage')) is not None else []),
-            ),
-        )
-
-    def build_ai_choice_delta(self, delta: ta.Mapping[str, ta.Any]) -> ContentAiChoiceDelta:
-        return ContentAiChoiceDelta(
-            check.not_none(delta.get('content')),
-            # FIXME:
-            # tool_exec_requests=[
-            #     ToolUse(
-            #         id=tc['id'],
-            #         spec=self._process_options().tools_by_name[tc['function']['name']],
-            #         args=json.loads(tc['function']['arguments'] or '{}'),
-            #         raw_args=tc['function']['arguments'],
-            #     )
-            #     for tc in message_or_delta.get('tool_calls', [])
-            # ] or None,
-        )

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/minichain/chat/stream/adapters.py

@@ -1,22 +1,14 @@
-from omlish import check
 from omlish import dataclasses as dc
-from omlish import lang
 
 from ...services import Response
-from ...tools.types import ToolUse
 from ..choices.services import ChatChoicesRequest
 from ..choices.services import static_check_is_chat_choices_service
 from ..choices.types import AiChoice
 from ..choices.types import AiChoices
-from ..messages import AiMessage
-from ..messages import AnyAiMessage
-from ..messages import ToolUseMessage
+from .joining import AiChoiceDeltaJoiner
 from .services import ChatChoicesOutputs
 from .services import ChatChoicesStreamOutputs
 from .services import ChatChoicesStreamService
-from .types import AiChoiceDelta
-from .types import ContentAiChoiceDelta
-from .types import ToolUseAiChoiceDelta
 
 
 ##
@@ -31,50 +23,13 @@ class ChatChoicesStreamServiceChatChoicesService:
         AiChoices,
         ChatChoicesOutputs | ChatChoicesStreamOutputs,
     ]:
-        choice_lsts: list[list[list[str] | ToolUse]] = []
-
-        def add(l: list[list[str] | ToolUse], d: AiChoiceDelta) -> None:
-            if isinstance(d, ContentAiChoiceDelta):
-                s = check.isinstance(d.c, str)
-                if l and isinstance(l[-1], list):
-                    l[-1].append(s)
-                else:
-                    l.append([s])
-
-            elif isinstance(d, ToolUseAiChoiceDelta):
-                l.append(d.tu)
-
-            else:
-                raise TypeError(d)
+        joiner = AiChoiceDeltaJoiner()
 
         async with (resp := await self.service.invoke(request)).v as it:  # noqa
-            i = -1  # noqa
-            l: list[list[str] | ToolUse]
-            async for i, cs in lang.async_enumerate(it):
-                if i == 0:
-                    for c in cs.choices:
-                        choice_lsts.append(l := [])
-                        for d in c.deltas:
-                            add(l, d)
-
-                else:
-                    for l, c in zip(choice_lsts, cs.choices, strict=True):
-                        for d in c.deltas:
-                            add(l, d)
+            async for cs in it:
+                joiner.add(cs.choices)
 
         # check.state(resp_v.is_done)
 
-        ret: list[AiChoice] = []
-        for cl in choice_lsts:
-            cc: list[AnyAiMessage] = []
-            for e in cl:
-                if isinstance(e, list):
-                    cc.append(AiMessage(''.join(e)))
-                elif isinstance(e, ToolUse):
-                    cc.append(ToolUseMessage(e))
-                else:
-                    raise TypeError(e)
-            ret.append(AiChoice(cc))
-
         # FIXME: outputs lol
-        return Response(ret)
+        return Response([AiChoice(ms) for ms in joiner.build()])

ommlds/minichain/chat/stream/joining.py

@@ -0,0 +1,96 @@
+import typing as ta
+
+from omlish import check
+from omlish.formats import json
+
+from ...tools.types import ToolUse
+from ..messages import AiChat
+from ..messages import AiMessage
+from ..messages import AnyAiMessage
+from ..messages import ToolUseMessage
+from .types import AiChoiceDelta
+from .types import AiChoiceDeltas
+from .types import ContentAiChoiceDelta
+from .types import PartialToolUseAiChoiceDelta
+from .types import ToolUseAiChoiceDelta
+
+
+##
+
+
+class AiChoiceDeltaJoiner:
+    def __init__(self) -> None:
+        super().__init__()
+
+        self._seq = 0
+        self._channels: list[AiChoiceDeltaJoiner._Channel] = []
+
+    class _Channel(ta.NamedTuple):
+        deltas: list[AiChoiceDelta]
+        messages: list[AnyAiMessage]
+
+    def _build_joined(self, deltas: ta.Sequence[AiChoiceDelta]) -> AnyAiMessage:
+        dty = check.single(set(map(type, check.not_empty(deltas))))
+
+        if dty is ContentAiChoiceDelta:
+            cds = ta.cast(ta.Sequence[ContentAiChoiceDelta], deltas)
+            return AiMessage(''.join(check.isinstance(cd.c, str) for cd in cds))
+
+        elif dty is ToolUseAiChoiceDelta:
+            raise TypeError(dty)
+
+        elif dty is PartialToolUseAiChoiceDelta:
+            tds = ta.cast(ta.Sequence[PartialToolUseAiChoiceDelta], deltas)
+            for td in ta.cast(ta.Sequence[PartialToolUseAiChoiceDelta], deltas)[1:]:
+                check.none(td.id)
+                check.none(td.name)
+
+            ra = ''.join(filter(None, (td.raw_args for td in tds)))
+
+            return ToolUseMessage(ToolUse(
+                id=tds[0].id,
+                name=check.non_empty_str(tds[0].name),
+                args=json.loads(ra),
+                raw_args=ra,
+            ))
+
+        else:
+            raise TypeError(dty)
+
+    def _join_one(self, chan: _Channel) -> None:
+        if not chan.deltas:
+            return
+
+        chan.messages.append(self._build_joined(chan.deltas))
+        chan.deltas.clear()
+
+    def _add_to(self, chan: _Channel, d: AiChoiceDelta) -> None:
+        if chan.deltas and type(chan.deltas[0]) is not type(d):
+            self._join_one(chan)
+
+        if isinstance(d, ToolUseAiChoiceDelta):
+            chan.messages.append(ToolUseMessage(ToolUse(
+                id=d.id,
+                name=check.not_none(d.name),
+                args=d.args or {},
+            )))
+
+        else:
+            chan.deltas.append(d)
+
+    def add(self, choices: ta.Sequence[AiChoiceDeltas]) -> None:
+        if not self._seq:
+            check.empty(self._channels)
+            self._channels.extend(self._Channel([], []) for _ in range(len(choices)))
+
+        for chan, c in zip(self._channels, choices, strict=True):
+            for d in c.deltas:
+                self._add_to(chan, d)
+
+        self._seq += 1
+
+    def build(self) -> list[AiChat]:
+        for chan in self._channels:
+            self._join_one(chan)
+
+        return [list(chan.messages) for chan in self._channels]

ommlds/minichain/chat/stream/types.py

@@ -6,7 +6,6 @@ from omlish import marshal as msh
 
 from ...content.types import Content
 from ...stream.services import StreamOptions
-from ...tools.types import ToolUse
 from ...types import Option
 from ...types import Output
 from ..choices.types import ChatChoicesOptions
@@ -43,14 +42,31 @@ class AiChoiceDelta(lang.Sealed, lang.Abstract):
     pass
 
 
+#
+
+
 @dc.dataclass(frozen=True)
 class ContentAiChoiceDelta(AiChoiceDelta, lang.Final):
     c: Content
 
 
-@dc.dataclass(frozen=True)
-class ToolUseAiChoiceDelta(AiChoiceDelta, lang.Final):
-    tu: ToolUse
+#
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class AnyToolUseAiChoiceDelta(AiChoiceDelta, lang.Abstract):
+    id: str | None = None
+    name: str | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class ToolUseAiChoiceDelta(AnyToolUseAiChoiceDelta, lang.Final):
+    args: ta.Mapping[str, ta.Any] | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class PartialToolUseAiChoiceDelta(AnyToolUseAiChoiceDelta, lang.Final):
+    raw_args: ta.Any | None = None
 
 
 #

{ommlds-0.0.0.dev465.dist-info → ommlds-0.0.0.dev467.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ommlds
-Version: 0.0.0.dev465
+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.dev465
-Requires-Dist: omlish==0.0.0.dev465
+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"