ommlds 0.0.0.dev503__py3-none-any.whl → 0.0.0.dev505__py3-none-any.whl

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

@@ -6,20 +6,17 @@ 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.cerebras 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 ....http.stream import BytesHttpStreamResponseBuilder
+from ....http.stream import SimpleSseLinesHttpStreamResponseHandler
 from ....standard import ApiKey
-from ....stream.services import StreamResponseSink
-from ....stream.services import new_stream_response
 from .chat import CerebrasChatChoicesService
 from .names import MODEL_NAMES
 from .protocol import build_cer_request_messages
@@ -49,6 +46,35 @@ class CerebrasChatChoicesStreamService:
             self._model_name = cc.pop(CerebrasChatChoicesService.DEFAULT_MODEL_NAME)
             self._api_key = ApiKey.pop_secret(cc, env='CEREBRAS_API_KEY')
 
+    URL: ta.ClassVar[str] = 'https://api.cerebras.ai/v1/chat/completions'
+
+    def _process_sse(self, so: sse.SseDecoderOutput) -> ta.Sequence[AiChoicesDeltas | None]:
+        if not (isinstance(so, sse.SseEvent) and so.type == b'message'):
+            return []
+
+        ss = so.data.decode('utf-8')
+        if ss == '[DONE]':
+            return [None]
+
+        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:
+            return []
+
+        if any(choice.finish_reason for choice in ccc.choices):
+            check.state(all(choice.finish_reason for choice in ccc.choices))
+            return [None]
+
+        return [AiChoicesDeltas([
+            build_mc_ai_choice_deltas(choice.delta)
+            for choice in ccc.choices
+        ])]
+
     READ_CHUNK_SIZE: ta.ClassVar[int] = -1
 
     async def invoke(self, request: ChatChoicesStreamRequest) -> ChatChoicesStreamResponse:
@@ -76,50 +102,11 @@ class CerebrasChatChoicesStreamService:
             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
-                            raise TypeError(l)
-
-                        # 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)
+        return await BytesHttpStreamResponseBuilder(
+            self._http_client,
+            lambda http_response: SimpleSseLinesHttpStreamResponseHandler(self._process_sse).as_lines().as_bytes(),
+            read_chunk_size=self.READ_CHUNK_SIZE,
+        ).new_stream_response(
+            http_request,
+            request.options,
+        )

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

@@ -16,16 +16,14 @@ 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 AnyAiMessage
-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.tools.types import Tool
 from ....models.configs import ModelName
 from ....standard import ApiKey
 from ....tools.types import ToolUse
 from .names import MODEL_NAMES
+from .protocol import make_msg_content
+from .protocol import pop_system_instructions
 from .tools import build_tool_spec_schema
 
 
@@ -53,90 +51,14 @@ class GoogleChatChoicesService:
             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)
-
-        elif isinstance(m, (SystemMessage, UserMessage)):
-            return check.isinstance(m.c, str)
-
-        else:
-            raise TypeError(m)
-
     BASE_URL: ta.ClassVar[str] = 'https://generativelanguage.googleapis.com/v1beta/models'
 
-    ROLES_MAP: ta.ClassVar[ta.Mapping[type[Message], str]] = {
-        UserMessage: 'user',
-        AiMessage: 'model',
-        ToolUseMessage: 'model',
-    }
-
     async def invoke(
             self,
             request: ChatChoicesRequest,
     ) -> ChatChoicesResponse:
         key = check.not_none(self._api_key).reveal()
 
-        g_sys_content: pt.Content | None = None
-        g_contents: list[pt.Content] = []
-        for i, m in enumerate(request.v):
-            if isinstance(m, SystemMessage):
-                check.arg(i == 0)
-                check.none(g_sys_content)
-                g_sys_content = pt.Content(
-                    parts=[pt.Part(
-                        text=check.not_none(self._get_msg_content(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)
-                g_contents.append(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, AiMessage):
-                g_contents.append(pt.Content(
-                    parts=[pt.Part(
-                        text=check.not_none(self._get_msg_content(m)),
-                    )],
-                    role='model',
-                ))
-
-            elif isinstance(m, ToolUseMessage):
-                g_contents.append(pt.Content(
-                    parts=[pt.Part(
-                        function_call=pt.FunctionCall(
-                            id=m.tu.id,
-                            name=m.tu.name,
-                            args=m.tu.args,
-                        ),
-                    )],
-                    role='model',
-                ))
-
-            else:
-                g_contents.append(pt.Content(
-                    parts=[pt.Part(
-                        text=check.not_none(self._get_msg_content(m)),
-                    )],
-                    role=self.ROLES_MAP[type(m)],  # type: ignore[arg-type]
-                ))
-
         g_tools: list[pt.Tool] = []
         with tv.TypedValues(*request.options).consume() as oc:
             t: Tool
@@ -145,10 +67,17 @@ class GoogleChatChoicesService:
                     function_declarations=[build_tool_spec_schema(t.spec)],
                 ))
 
+        msgs = list(request.v)
+
+        system_inst = pop_system_instructions(msgs)
+
         g_req = pt.GenerateContentRequest(
-            contents=g_contents or None,
+            contents=[
+                make_msg_content(m)
+                for m in msgs
+            ] or None,
             tools=g_tools or None,
-            system_instruction=g_sys_content,
+            system_instruction=system_inst,
         )
 
         req_dct = msh.marshal(g_req)

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

@@ -0,0 +1,105 @@
+import typing as ta
+
+from omlish import check
+
+from .....backends.google.protocol import types as pt
+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
+
+
+##
+
+
+ROLES_MAP: ta.Mapping[type[Message], pt.ContentRole | None] = {  # noqa
+    SystemMessage: None,
+    UserMessage: 'user',
+    AiMessage: 'model',
+    ToolUseMessage: 'model',
+}
+
+
+def make_str_content(
+        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,
+    )
+
+
+def make_msg_content(m: Message) -> pt.Content:
+    if isinstance(m, (AiMessage, SystemMessage, UserMessage)):
+        return check.not_none(make_str_content(
+            check.isinstance(m.c, str),
+            role=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)
+
+
+def pop_system_instructions(msgs: list[Message]) -> pt.Content | None:
+    if not msgs:
+        return None
+
+    m0 = msgs[0]
+    if not isinstance(m0 := msgs[0], SystemMessage):
+        return None
+
+    msgs.pop(0)
+    return make_msg_content(m0)
+
+
+def get_msg_content(m: Message) -> str | None:
+    if isinstance(m, AiMessage):
+        return check.isinstance(m.c, str)
+
+    elif isinstance(m, (SystemMessage, UserMessage)):
+        return check.isinstance(m.c, str)
+
+    else:
+        raise TypeError(m)

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

@@ -8,7 +8,7 @@ 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.io.buffers import DelimitingBuffer
+from omlish.http import sse
 
 from .....backends.google.protocol import types as pt
 from ....chat.choices.stream.services import ChatChoicesStreamRequest
@@ -16,22 +16,16 @@ 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 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.types import ContentAiDelta
 from ....chat.stream.types import ToolUseAiDelta
 from ....chat.tools.types import Tool
+from ....http.stream import BytesHttpStreamResponseBuilder
+from ....http.stream import SimpleSseLinesHttpStreamResponseHandler
 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 .protocol import make_msg_content
+from .protocol import pop_system_instructions
 from .tools import build_tool_spec_schema
 
 
@@ -59,71 +53,38 @@ class GoogleChatChoicesStreamService:
             self._model_name = cc.pop(self.DEFAULT_MODEL_NAME)
             self._api_key = ApiKey.pop_secret(cc, env='GEMINI_API_KEY')
 
-    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,
-        )
-
-    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], pt.ContentRole | None]] = {  # noqa
-        SystemMessage: None,
-        UserMessage: 'user',
-        AiMessage: 'model',
-    }
+    def _process_sse(self, so: sse.SseDecoderOutput) -> ta.Iterable[AiChoicesDeltas | None]:
+        if not (isinstance(so, sse.SseEvent) and so.type == b'message'):
+            return
+
+        gcr = msh.unmarshal(json.loads(so.data.decode('utf-8')), pt.GenerateContentResponse)  # noqa
+        cnd = check.single(check.not_none(gcr.candidates))
+
+        for p in check.not_none(cnd.content).parts or []:
+            if (txt := p.text) is not None:
+                check.none(p.function_call)
+                yield AiChoicesDeltas([
+                    AiChoiceDeltas([
+                        ContentAiDelta(check.not_none(txt)),
+                    ]),
+                ])
+
+            elif (fc := p.function_call) is not None:
+                check.none(p.text)
+                yield AiChoicesDeltas([
+                    AiChoiceDeltas([
+                        ToolUseAiDelta(
+                            id=fc.id,
+                            name=fc.name,
+                            args=fc.args,
+                        ),
+                    ]),
+                ])
+
+            else:
+                raise ValueError(p)
 
     READ_CHUNK_SIZE: ta.ClassVar[int] = -1
 
@@ -133,13 +94,6 @@ 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
@@ -148,11 +102,15 @@ class GoogleChatChoicesStreamService:
                     function_declarations=[build_tool_spec_schema(t.spec)],
                 ))
 
+        msgs = list(request.v)
+
+        system_inst = pop_system_instructions(msgs)
+
         g_req = pt.GenerateContentRequest(
             contents=[
-                self._make_msg_content(m)
+                make_msg_content(m)
                 for m in msgs
-            ],
+            ] or None,
             tools=g_tools or None,
             system_instruction=system_inst,
         )
@@ -168,52 +126,11 @@ class GoogleChatChoicesStreamService:
             method='POST',
         )
 
-        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] | None:
-                db = DelimitingBuffer([b'\r', b'\n', b'\r\n'])
-                while True:
-                    b = await http_response.stream.read1(self.READ_CHUNK_SIZE)
-                    for bl in db.feed(b):
-                        if isinstance(bl, DelimitingBuffer.Incomplete):
-                            # FIXME: handle
-                            raise TypeError(bl)
-
-                        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 []:
-                                if (txt := p.text) is not None:
-                                    check.none(p.function_call)
-                                    await sink.emit(AiChoicesDeltas([
-                                        AiChoiceDeltas([
-                                            ContentAiDelta(check.not_none(txt)),
-                                        ]),
-                                    ]))
-
-                                elif (fc := p.function_call) is not None:
-                                    check.none(p.text)
-                                    await sink.emit(AiChoicesDeltas([
-                                        AiChoiceDeltas([
-                                            ToolUseAiDelta(
-                                                id=fc.id,
-                                                name=fc.name,
-                                                args=fc.args,
-                                            ),
-                                        ]),
-                                    ]))
-
-                                else:
-                                    raise ValueError(p)
-
-                    if not b:
-                        return []
-
-            return await new_stream_response(rs, inner)
+        return await BytesHttpStreamResponseBuilder(
+            self._http_client,
+            lambda http_response: SimpleSseLinesHttpStreamResponseHandler(self._process_sse).as_lines().as_bytes(),
+            read_chunk_size=self.READ_CHUNK_SIZE,
+        ).new_stream_response(
+            http_request,
+            request.options,
+        )