ommlds 0.0.0.dev451__py3-none-any.whl → 0.0.0.dev453__py3-none-any.whl

ommlds/minichain/backends/impls/mistral.py

@@ -90,6 +90,6 @@ class MistralChatChoicesService:
         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

@@ -137,5 +137,5 @@ class MlxChatChoicesService(lang.ExitStacked):
         )
 
         return ChatChoicesResponse([
-            AiChoice(AiMessage(response))  # noqa
+            AiChoice([AiMessage(response)])  # noqa
         ])

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

@@ -14,17 +14,20 @@ TODO:
 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 .....backends.openai import protocol as pt
 from ....chat.choices.services import ChatChoicesRequest
 from ....chat.choices.services import ChatChoicesResponse
 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 .format import OpenaiChatRequestHandler
+from .format2 import OpenaiChatRequestHandler
+from .format2 import build_mc_choices_response
 from .names import MODEL_NAMES
 
 
@@ -63,7 +66,7 @@ class OpenaiChatChoicesService:
             ),
         )
 
-        raw_request = rh.raw_request()
+        raw_request = msh.marshal(rh.oai_request())
 
         http_response = http.request(
             'https://api.openai.com/v1/chat/completions',
@@ -76,4 +79,4 @@ class OpenaiChatChoicesService:
 
         raw_response = json.loads(check.not_none(http_response.data).decode('utf-8'))
 
-        return rh.build_response(raw_response)
+        return build_mc_choices_response(msh.unmarshal(raw_response, pt.ChatCompletionResponse))

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

@@ -9,67 +9,87 @@ from omlish.formats import json
 from ....chat.choices.services import ChatChoicesResponse
 from ....chat.choices.types import AiChoice
 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 ToolExecResultMessage
+from ....chat.messages import ToolUseMessage
+from ....chat.messages import ToolUseResultMessage
 from ....chat.messages import UserMessage
-from ....chat.stream.types import AiMessageDelta
+from ....chat.stream.types import ContentAiChoiceDelta
 from ....chat.tools.types import Tool
+from ....content.json import JsonContent
 from ....content.prepare import prepare_content_str
 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.types import ToolExecRequest
 from ....tools.types import ToolSpec
+from ....tools.types import ToolUse
 from ....types import Option
 
 
 ##
 
 
-def build_request_message(m: Message) -> ta.Mapping[str, ta.Any]:
-    if isinstance(m, SystemMessage):
-        return dict(
-            role='system',
-            content=m.c,
-        )
-
-    elif isinstance(m, AiMessage):
-        return dict(
-            role='assistant',
-            content=check.isinstance(m.c, (str, None)),
-            **(dict(tool_calls=[
-                dict(
-                    id=te.id,
-                    function=dict(
-                        arguments=check.not_none(te.raw_args),
-                        name=te.name,
+def build_request_messages(chat: Chat) -> ta.Sequence[ta.Mapping[str, ta.Any]]:
+    out: list[dict[str, ta.Any]] = []
+
+    for m in chat:
+        if isinstance(m, SystemMessage):
+            out.append(dict(
+                role='system',
+                content=m.c,
+            ))
+
+        elif isinstance(m, AiMessage):
+            out.append(dict(
+                role='assistant',
+                content=check.isinstance(m.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',
                     ),
-                    type='function',
-                )
-                for te in m.tool_exec_requests
-            ]) if m.tool_exec_requests else {}),
-        )
-
-    elif isinstance(m, UserMessage):
-        return dict(
-            role='user',
-            content=prepare_content_str(m.c),
-        )
-
-    elif isinstance(m, ToolExecResultMessage):
-        return dict(
-            role='tool',
-            tool_call_id=m.id,
-            content=check.isinstance(m.c, str),
-        )
-
-    else:
-        raise TypeError(m)
+                ],
+            ))
+
+        elif isinstance(m, UserMessage):
+            out.append(dict(
+                role='user',
+                content=prepare_content_str(m.c),
+            ))
+
+        elif isinstance(m, 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)
+            else:
+                raise TypeError(m.tur.c)
+            out.append(dict(
+                role='tool',
+                tool_call_id=m.tur.id,
+                content=tc,
+            ))
+
+        else:
+            raise TypeError(m)
+
+    return out
 
 
 ##
@@ -94,10 +114,11 @@ class OpenaiChatRequestHandler:
         SystemMessage: 'system',
         UserMessage: 'user',
         AiMessage: 'assistant',
-        ToolExecResultMessage: 'tool',
+        ToolUseMessage: 'assistant',
+        ToolUseResultMessage: 'tool',
     }
 
-    DEFAULT_OPTIONS: ta.ClassVar[tv.TypedValues[Option]] = tv.TypedValues(
+    DEFAULT_OPTIONS: ta.ClassVar[tv.TypedValues[Option]] = tv.TypedValues[Option](
         Temperature(0.),
         MaxTokens(1024),
     )
@@ -152,10 +173,7 @@ class OpenaiChatRequestHandler:
 
         return dict(
             model=self._model,
-            messages=[
-                build_request_message(m)
-                for m in self._chat
-            ],
+            messages=build_request_messages(self._chat),
             top_p=1,
             **lang.opt_kw(tools=tools),
             frequency_penalty=0.0,
@@ -163,24 +181,25 @@ class OpenaiChatRequestHandler:
             **po.kwargs,
         )
 
-    def build_ai_message(self, message: ta.Mapping[str, ta.Any]) -> AiMessage:
-        return AiMessage(
-            message.get('content'),
-            tool_exec_requests=[
-                ToolExecRequest(
+    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'],
-                )
-                for tc in message.get('tool_calls', [])
-            ] or None,
-        )
+                ),
+            ))
+        return out
 
     def build_response(self, raw_response: ta.Mapping[str, ta.Any]) -> ChatChoicesResponse:
         return ChatChoicesResponse(
             [
-                AiChoice(self.build_ai_message(choice['message']))
+                AiChoice(self.build_ai_chat(choice['message']))
                 for choice in raw_response['choices']
             ],
 
@@ -193,12 +212,12 @@ class OpenaiChatRequestHandler:
             ),
         )
 
-    def build_ai_message_delta(self, delta: ta.Mapping[str, ta.Any]) -> AiMessageDelta:
-        return AiMessageDelta(
-            delta.get('content'),
+    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=[
-            #     ToolExecRequest(
+            #     ToolUse(
             #         id=tc['id'],
             #         spec=self._process_options().tools_by_name[tc['function']['name']],
             #         args=json.loads(tc['function']['arguments'] or '{}'),

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

@@ -0,0 +1,210 @@
+import typing as ta
+
+from omlish import cached
+from omlish import check
+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 AiMessage
+from ....chat.messages import AnyAiMessage
+from ....chat.messages import Chat
+from ....chat.messages import SystemMessage
+from ....chat.messages import ToolUseMessage
+from ....chat.messages import ToolUseResultMessage
+from ....chat.messages import UserMessage
+from ....chat.tools.types import Tool
+from ....content.json import JsonContent
+from ....content.prepare import prepare_content_str
+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_params_json_schema
+from ....tools.types import ToolSpec
+from ....tools.types import ToolUse
+from ....types import Option
+
+
+##
+
+
+def build_oai_request_msgs(mc_chat: Chat) -> ta.Sequence[pt.ChatCompletionMessage]:
+    oai_msgs: list[pt.ChatCompletionMessage] = []
+
+    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(mc_msg, AiMessage):
+            oai_msgs.append(pt.AssistantChatCompletionMessage(
+                content=check.isinstance(mc_msg.c, (str, None)),
+            ))
+
+        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(mc_msg, UserMessage):
+            oai_msgs.append(pt.UserChatCompletionMessage(
+                content=prepare_content_str(mc_msg.c),
+            ))
+
+        elif isinstance(mc_msg, ToolUseResultMessage):
+            tc: str
+            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(mc_msg.tur.c)
+            oai_msgs.append(pt.ToolChatCompletionMessage(
+                tool_call_id=check.not_none(mc_msg.tur.id),
+                content=tc,
+            ))
+
+        else:
+            raise TypeError(mc_msg)
+
+    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 []),
+        ),
+    )
+
+
+##
+
+
+class OpenaiChatRequestHandler:
+    def __init__(
+            self,
+            chat: Chat,
+            *options: ChatChoicesOptions,
+            model: str,
+            mandatory_kwargs: ta.Mapping[str, ta.Any] | None = None,
+    ) -> None:
+        super().__init__()
+
+        self._chat = chat
+        self._options = options
+        self._model = model
+        self._mandatory_kwargs = mandatory_kwargs
+
+    DEFAULT_OPTIONS: ta.ClassVar[tv.TypedValues[Option]] = tv.TypedValues[Option](
+        Temperature(0.),
+        MaxTokens(1024),
+    )
+
+    _OPTION_KWARG_NAMES_MAP: ta.ClassVar[ta.Mapping[str, type[ChatChoicesOptions]]] = dict(
+        temperature=Temperature,
+        max_tokens=MaxTokens,
+    )
+
+    class _ProcessedOptions(ta.NamedTuple):
+        kwargs: dict[str, ta.Any]
+        tools_by_name: dict[str, ToolSpec]
+
+    @cached.function
+    def _process_options(self) -> _ProcessedOptions:
+        kwargs: dict = dict(
+            temperature=0,
+            max_tokens=1024,
+        )
+
+        tools_by_name: dict[str, ToolSpec] = {}
+
+        with tv.TypedValues(*self._options).consume() as oc:
+            kwargs.update(oc.pop_scalar_kwargs(**self._OPTION_KWARG_NAMES_MAP))
+
+            for t in oc.pop(Tool, []):
+                if t.spec.name in tools_by_name:
+                    raise NameError(t.spec.name)
+                tools_by_name[check.non_empty_str(t.spec.name)] = t.spec
+
+        if (mk := self._mandatory_kwargs):
+            for k, v in mk.items():
+                check.not_in(k, kwargs)
+                kwargs[k] = v
+
+        return self._ProcessedOptions(
+            kwargs=kwargs,
+            tools_by_name=tools_by_name,
+        )
+
+    @cached.function
+    def oai_request(self) -> pt.ChatCompletionRequest:
+        po = self._process_options()
+
+        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 pt.ChatCompletionRequest(
+            model=self._model,
+            messages=build_oai_request_msgs(self._chat),
+            top_p=1,
+            tools=tools or None,
+            frequency_penalty=0.0,
+            presence_penalty=0.0,
+            **po.kwargs,
+        )

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

@@ -16,8 +16,8 @@ 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 AiChoicesDeltas
 from ....chat.stream.types import ChatChoicesStreamOption
 from ....configs import Config
 from ....resources import ResourcesOption
@@ -83,7 +83,7 @@ class OpenaiChatChoicesStreamService:
             http_client = rs.enter_context(http.client())
             http_response = rs.enter_context(http_client.stream_request(http_request))
 
-            async def inner(sink: StreamResponseSink[AiChoiceDeltas]) -> ta.Sequence[ChatChoicesOutputs]:
+            async def inner(sink: StreamResponseSink[AiChoicesDeltas]) -> ta.Sequence[ChatChoicesOutputs]:
                 db = DelimitingBuffer([b'\r', b'\n', b'\r\n'])
                 sd = sse.SseDecoder()
                 while True:
@@ -112,10 +112,13 @@ class OpenaiChatChoicesStreamService:
                                 if not sj['choices']:
                                     continue
 
-                                await sink.emit([
-                                    AiChoiceDelta(rh.build_ai_message_delta(choice['delta']))
-                                    for choice in sj['choices']
-                                ])
+                                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 not b:
                         return []

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

@@ -18,9 +18,9 @@ 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 AiChoiceDelta
 from ....chat.stream.types import AiChoiceDeltas
-from ....chat.stream.types import AiMessageDelta
+from ....chat.stream.types import AiChoicesDeltas
+from ....chat.stream.types import ContentAiChoiceDelta
 from ....chat.types import ChatOption
 from ....llms.types import LlmOption
 from ....resources import UseResources
@@ -126,7 +126,7 @@ class TinygradLlama3ChatChoicesService(BaseTinygradLlama3ChatService):
         for s in tgl3.run_llm(llm, toks):
             out.append(s)
 
-        return ChatChoicesResponse([AiChoice(AiMessage(''.join(out)))])
+        return ChatChoicesResponse([AiChoice([AiMessage(''.join(out))])])
 
 
 ##
@@ -147,9 +147,14 @@ class TinygradLlama3ChatChoicesStreamService(BaseTinygradLlama3ChatService):
                 request.options.get_any((ChatOption, LlmOption)),  # FIXME  # noqa
             )
 
-            async def inner(sink: StreamResponseSink[AiChoiceDeltas]) -> ta.Sequence[ChatChoicesOutputs]:
+            async def inner(sink: StreamResponseSink[AiChoicesDeltas]) -> ta.Sequence[ChatChoicesOutputs]:
                 for s in tgl3.run_llm(llm, toks):
-                    await sink.emit([AiChoiceDelta(AiMessageDelta(s))])
+                    await sink.emit(AiChoicesDeltas([
+                        AiChoiceDeltas([
+                            ContentAiChoiceDelta(s),
+                        ]),
+                    ]))
+
                 return []
 
             return await new_stream_response(rs, inner)

ommlds/minichain/backends/impls/transformers/transformers.py

@@ -15,10 +15,12 @@ from omlish import typedvalues as tv
 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.types import AiChoice
 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
 from ....completion import CompletionRequest
 from ....completion import CompletionResponse
@@ -94,17 +96,19 @@ def build_chat_message(m: Message) -> ta.Mapping[str, ta.Any]:
         return dict(
             role='assistant',
             content=check.isinstance(m.c, str),
-            **(dict(tool_calls=[
-                dict(
-                    id=te.id,
-                    function=dict(
-                        arguments=te.args,
-                        name=te.name,
-                    ),
-                    type='function',
-                )
-                for te in m.tool_exec_requests
-            ]) if m.tool_exec_requests else {}),
+        )
+
+    elif isinstance(m, ToolUseMessage):
+        return dict(
+            role='assistant',
+            tool_calls=[dict(
+                id=m.tu.id,
+                function=dict(
+                    arguments=m.tu.args,
+                    name=m.tu.name,
+                ),
+                type='function',
+            )],
         )
 
     elif isinstance(m, UserMessage):
@@ -113,11 +117,11 @@ def build_chat_message(m: Message) -> ta.Mapping[str, ta.Any]:
             content=check.isinstance(m.c, str),
         )
 
-    elif isinstance(m, ToolExecResultMessage):
+    elif isinstance(m, ToolUseResultMessage):
         return dict(
             role='tool',
-            tool_call_id=m.id,
-            content=check.isinstance(m.c, str),
+            tool_call_id=m.tur.id,
+            content=check.isinstance(m.tur.c, str),
         )
 
     else:
@@ -174,4 +178,4 @@ class TransformersChatChoicesService(lang.ExitStacked):
             ],
         )
 
-        return ChatChoicesResponse(output)
+        return ChatChoicesResponse([AiChoice([output])])

ommlds/minichain/chat/_marshal.py

@@ -5,6 +5,8 @@ TODO:
 from omlish import lang
 from omlish import marshal as msh
 
+from .messages import AnyAiMessage
+from .messages import AnyUserMessage
 from .messages import Message
 
 
@@ -13,12 +15,17 @@ from .messages import Message
 
 @lang.static_init
 def _install_standard_marshaling() -> None:
-    msgs_poly = msh.polymorphism_from_subclasses(
+    for cls in [
+        AnyAiMessage,
+        AnyUserMessage,
         Message,
-        naming=msh.Naming.SNAKE,
-        strip_suffix=True,
-    )
-    msh.install_standard_factories(
-        msh.PolymorphismMarshalerFactory(msgs_poly),
-        msh.PolymorphismUnmarshalerFactory(msgs_poly),
-    )
+    ]:
+        cls_poly = msh.polymorphism_from_subclasses(
+            cls,
+            naming=msh.Naming.SNAKE,
+            strip_suffix='Message',
+        )
+        msh.install_standard_factories(
+            msh.PolymorphismMarshalerFactory(cls_poly),
+            msh.PolymorphismUnmarshalerFactory(cls_poly),
+        )

ommlds/minichain/chat/choices/adapters.py

@@ -2,7 +2,7 @@ from omlish import check
 from omlish import dataclasses as dc
 
 from ...services import Response
-from ..messages import AiMessage
+from ..messages import AiChat
 from ..services import ChatRequest
 from ..services import static_check_is_chat_service
 from .services import ChatChoicesService
@@ -17,6 +17,6 @@ from .types import ChatChoicesOutputs
 class ChatChoicesServiceChatService:
     service: ChatChoicesService
 
-    async def invoke(self, request: ChatRequest) -> Response[AiMessage, ChatChoicesOutputs]:
+    async def invoke(self, request: ChatRequest) -> Response[AiChat, ChatChoicesOutputs]:
         resp = await self.service.invoke(request)
-        return Response(check.single(resp.v).m, resp.outputs)
+        return Response(check.single(resp.v).ms, resp.outputs)