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

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

@@ -1,26 +1,33 @@
 import typing as ta
 
 from omlish import check
-from omlish import lang
+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.anthropic.protocol import types as pt
+from .....backends.anthropic.protocol.sse.events import AnthropicSseDecoderEvents
 from ....chat.choices.services import ChatChoicesOutputs
-from ....chat.messages import SystemMessage
-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.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 ....chat.stream.types import PartialToolUseAiDelta
+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 AnthropicChatChoicesService
 from .names import MODEL_NAMES
+from .protocol import build_protocol_chat_messages
+from .protocol import build_protocol_tool
 
 
 ##
@@ -32,42 +39,46 @@ from .names import MODEL_NAMES
 # )
 @static_check_is_chat_choices_stream_service
 class AnthropicChatChoicesStreamService:
-    def __init__(self, *configs: Config) -> None:
+    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(AnthropicChatChoicesService.DEFAULT_MODEL_NAME)
             self._api_key = check.not_none(ApiKey.pop_secret(cc, env='ANTHROPIC_API_KEY'))
 
-    READ_CHUNK_SIZE = 64 * 1024
+    READ_CHUNK_SIZE: ta.ClassVar[int] = -1
 
-    def invoke(
+    async def invoke(
             self,
             request: ChatChoicesStreamRequest,
             *,
             max_tokens: int = 4096,  # FIXME: ChatOption
     ) -> ChatChoicesStreamResponse:
-        messages = []
-        system: str | None = None
-        for i, m in enumerate(request.v):
-            if isinstance(m, SystemMessage):
-                if i != 0 or system is not None:
-                    raise Exception('Only supports one system message and must be first')
-                system = AnthropicChatChoicesService._get_msg_content(m)  # noqa
-            else:
-                messages.append(dict(
-                    role=AnthropicChatChoicesService.ROLES_MAP[type(m)],  # noqa
-                    content=check.isinstance(AnthropicChatChoicesService._get_msg_content(m), str),  # noqa
-                ))
-
-        raw_request = dict(
+        messages, system = build_protocol_chat_messages(request.v)
+
+        tools: list[pt.ToolSpec] = []
+        with tv.TypedValues(*request.options).consume() as oc:
+            t: Tool
+            for t in oc.pop(Tool, []):
+                tools.append(build_protocol_tool(t))
+
+        a_req = pt.MessagesRequest(
             model=MODEL_NAMES.resolve(self._model_name.v),
-            **lang.opt_kw(system=system),
+            system=system,
             messages=messages,
+            tools=tools or None,
             max_tokens=max_tokens,
             stream=True,
         )
 
+        raw_request = msh.marshal(a_req)
+
         http_request = http.HttpRequest(
             'https://api.anthropic.com/v1/messages',
             headers={
@@ -78,16 +89,19 @@ class AnthropicChatChoicesStreamService:
             data=json.dumps(raw_request).encode('utf-8'),
         )
 
-        with UseResources.or_new(request.options) as rs:
-            http_client = rs.enter_context(http.client())
-            http_response = rs.enter_context(http_client.stream_request(http_request))
+        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:
+                msg_start: AnthropicSseDecoderEvents.MessageStart | None = None
+                cbk_start: AnthropicSseDecoderEvents.ContentBlockStart | None = None
+                msg_stop: AnthropicSseDecoderEvents.MessageStop | None = None
 
-            def yield_choices() -> ta.Generator[AiChoiceDeltas, None, ta.Sequence[ChatChoicesOutputs] | None]:
                 db = DelimitingBuffer([b'\r', b'\n', b'\r\n'])
                 sd = sse.SseDecoder()
                 while True:
-                    # FIXME: read1 not on response stream protocol
-                    b = http_response.stream.read1(self.READ_CHUNK_SIZE)  # type: ignore[attr-defined]
+                    b = await http_response.stream.read1(self.READ_CHUNK_SIZE)
                     for l in db.feed(b):
                         if isinstance(l, DelimitingBuffer.Incomplete):
                             # FIXME: handle
@@ -95,29 +109,83 @@ class AnthropicChatChoicesStreamService:
 
                         # FIXME: https://docs.anthropic.com/en/docs/build-with-claude/streaming
                         for so in sd.process_line(l):
-                            # FIXME: AnthropicSseMessageAssembler lol
-                            if isinstance(so, sse.SseEvent) and so.type == b'message':
+                            if isinstance(so, sse.SseEvent):
                                 ss = so.data.decode('utf-8')
                                 if ss == '[DONE]':
                                     return []
 
-                                sj = json.loads(ss)  # ChatCompletionChunk
+                                dct = json.loads(ss)
+                                check.equal(dct['type'], so.type.decode('utf-8'))
+                                ae = msh.unmarshal(dct, AnthropicSseDecoderEvents.Event)
+
+                                match ae:
+                                    case AnthropicSseDecoderEvents.MessageStart():
+                                        check.none(msg_start)
+                                        msg_start = ae
+                                        if msg_start.message.content:
+                                            raise NotImplementedError
+
+                                    case AnthropicSseDecoderEvents.ContentBlockStart():
+                                        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([ContentAiDelta(
+                                                ae.content_block.text,
+                                            )])]))
+
+                                        elif isinstance(ae.content_block, AnthropicSseDecoderEvents.ContentBlockStart.ToolUse):  # noqa
+                                            await sink.emit(AiChoicesDeltas([AiChoiceDeltas([PartialToolUseAiDelta(  # 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([ContentAiDelta(
+                                                ae.delta.text,
+                                            )])]))
+
+                                        elif isinstance(ae.delta, AnthropicSseDecoderEvents.ContentBlockDelta.InputJsonDelta):  # noqa
+                                            await sink.emit(AiChoicesDeltas([AiChoiceDeltas([PartialToolUseAiDelta(  # noqa
+                                                raw_args=ae.delta.partial_json,
+                                            )])]))
+
+                                        else:
+                                            raise TypeError(ae.delta)
+
+                                    case AnthropicSseDecoderEvents.ContentBlockStop():
+                                        check.not_none(cbk_start)
+                                        cbk_start = None
+
+                                    case AnthropicSseDecoderEvents.MessageDelta():
+                                        check.not_none(msg_start)
+                                        check.none(cbk_start)
 
-                                check.state(sj['object'] == 'chat.completion.chunk')
+                                    case AnthropicSseDecoderEvents.MessageStop():
+                                        check.not_none(msg_start)
+                                        check.none(msg_stop)
+                                        msg_stop = ae
 
-                                # FIXME: stop reason
-                                if not sj['choices']:
-                                    continue
+                                    case AnthropicSseDecoderEvents.Ping():
+                                        pass
 
-                                yield [
-                                    AiChoiceDelta(choice['delta'])
-                                    for choice in sj['choices']
-                                ]
+                                    case _:
+                                        raise TypeError(ae)
 
                     if not b:
+                        check.not_none(msg_stop)
+                        check.none(cbk_start)
                         return []
 
             # raw_response = json.loads(check.not_none(http_response.data).decode('utf-8'))
             # return rh.build_response(raw_response)
 
-            return new_stream_response(rs, yield_choices())
+            return await new_stream_response(rs, inner)

ommlds/minichain/backends/impls/duckduckgo/search.py

@@ -17,7 +17,7 @@ from ....search import static_check_is_search_service
 # )
 @static_check_is_search_service
 class DuckduckgoSearchService:
-    def invoke(self, request: SearchRequest) -> SearchResponse:
+    async def invoke(self, request: SearchRequest) -> SearchResponse:
         dsch = ddgs.DDGS()
         res = dsch.text(request.v)
         return SearchResponse(SearchHits(

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

@@ -0,0 +1,69 @@
+# ruff: noqa: PERF402
+import typing as ta
+
+from omlish.text.lorem import LOREM
+
+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 ChatChoicesOutputs
+from ....chat.messages import AiMessage
+from ....chat.stream.types import ContentAiDelta
+from ....resources import UseResources
+from ....stream.services import StreamResponseSink
+from ....stream.services import new_stream_response
+
+
+##
+
+
+# @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
+#     name='dummy',
+#     type='ChatChoicesService',
+# )
+@static_check_is_chat_choices_service
+class DummyChatChoicesService:
+    async def invoke(self, request: ChatChoicesRequest) -> ChatChoicesResponse:
+        return ChatChoicesResponse([AiChoice([AiMessage(LOREM)])])
+
+
+##
+
+
+# @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
+#     name='dummy',
+#     type='ChatChoicesStreamService',
+# )
+@static_check_is_chat_choices_stream_service
+class DummyChatChoicesStreamService:
+    async def invoke(self, request: ChatChoicesStreamRequest) -> ChatChoicesStreamResponse:
+        async with UseResources.or_new(request.options) as rs:
+            async def inner(sink: StreamResponseSink[AiChoicesDeltas]) -> ta.Sequence[ChatChoicesOutputs]:
+                for s in LOREM:
+                    await sink.emit(AiChoicesDeltas([
+                        AiChoiceDeltas([
+                            ContentAiDelta(s),
+                        ]),
+                    ]))
+
+                return []
+
+            return await new_stream_response(rs, inner)
+
+
+##
+
+
+# @omlish-manifest $.minichain.backends.strings.manifests.BackendStringsManifest(
+#     [
+#         'ChatChoicesService',
+#         'ChatChoicesStreamService',
+#     ],
+#     'dummy',
+# )

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

@@ -4,21 +4,29 @@ https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models
 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.google.protocol import types 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 ....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 .tools import build_tool_spec_schema
 
 
 ##
@@ -32,9 +40,15 @@ from .names import MODEL_NAMES
 class GoogleChatChoicesService:
     DEFAULT_MODEL_NAME: ta.ClassVar[ModelName] = ModelName(check.not_none(MODEL_NAMES.default))
 
-    def __init__(self, *configs: ApiKey | ModelName) -> None:
+    def __init__(
+            self,
+            *configs: ApiKey | ModelName,
+            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(self.DEFAULT_MODEL_NAME)
             self._api_key = ApiKey.pop_secret(cc, env='GEMINI_API_KEY')
@@ -52,43 +66,121 @@ class GoogleChatChoicesService:
     BASE_URL: ta.ClassVar[str] = 'https://generativelanguage.googleapis.com/v1beta/models'
 
     ROLES_MAP: ta.ClassVar[ta.Mapping[type[Message], str]] = {
-        SystemMessage: 'system',
         UserMessage: 'user',
-        AiMessage: 'assistant',
+        AiMessage: 'model',
+        ToolUseMessage: 'model',
     }
 
-    def invoke(
+    async def invoke(
             self,
             request: ChatChoicesRequest,
     ) -> ChatChoicesResponse:
         key = check.not_none(self._api_key).reveal()
 
-        req_dct = {
-            'contents': [
-                {
-                    'role': self.ROLES_MAP[type(m)],
-                    'parts': [
-                        {
-                            'text': self._get_msg_content(m),
-                        },
-                    ],
-                }
-                for m in request.v
-            ],
-        }
+        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
+            for t in oc.pop(Tool, []):
+                g_tools.append(pt.Tool(
+                    function_declarations=[build_tool_spec_schema(t.spec)],
+                ))
+
+        g_req = pt.GenerateContentRequest(
+            contents=g_contents or None,
+            tools=g_tools or None,
+            system_instruction=g_sys_content,
+        )
+
+        req_dct = msh.marshal(g_req)
 
         model_name = MODEL_NAMES.resolve(self._model_name.v)
 
-        resp = http.request(
+        resp = await http.async_request(
             f'{self.BASE_URL.rstrip("/")}/{model_name}:generateContent?key={key}',
             headers={'Content-Type': 'application/json'},
             data=json.dumps_compact(req_dct).encode('utf-8'),
             method='POST',
+            client=self._http_client,
         )
 
         resp_dct = json.loads(check.not_none(resp.data).decode('utf-8'))
 
-        return ChatChoicesResponse([
-            AiChoice(AiMessage(c['content']['parts'][0]['text']))
-            for c in resp_dct['candidates']
-        ])
+        g_resp = msh.unmarshal(resp_dct, pt.GenerateContentResponse)
+
+        ai_choices: list[AiChoice] = []
+        for c in g_resp.candidates or []:
+            out: list[AnyAiMessage] = []
+            for g_resp_part in check.not_none(check.not_none(c.content).parts):
+                if (g_txt := g_resp_part.text) is not None:
+                    out.append(AiMessage(g_txt))
+                elif (g_fc := g_resp_part.function_call) is not None:
+                    out.append(ToolUseMessage(ToolUse(
+                        id=g_fc.id,
+                        name=g_fc.name,
+                        args=g_fc.args or {},
+                    )))
+                else:
+                    raise TypeError(g_resp_part)
+            ai_choices.append(AiChoice(out))
+
+        return ChatChoicesResponse(ai_choices)

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

@@ -82,13 +82,17 @@ class CseSearchService:
             self,
             cse_id: str | None = None,
             cse_api_key: str | None = None,
+            *,
+            http_client: http.AsyncHttpClient | None = None,
     ) -> None:
         super().__init__()
 
         self._cse_id = cse_id
         self._cse_api_key = cse_api_key
 
-    def invoke(
+        self._http_client = http_client
+
+    async def invoke(
             self,
             request: SearchRequest,
     ) -> SearchResponse:
@@ -97,8 +101,9 @@ class CseSearchService:
             cx=check.non_empty_str(self._cse_id),
             q=request.v,
         ))
-        resp = http.request(
+        resp = await http.async_request(
             f'https://www.googleapis.com/customsearch/v1?{qs}',
+            client=self._http_client,
         )
         out = check.not_none(resp.data)