ommlds 0.0.0.dev426__py3-none-any.whl → 0.0.0.dev485__py3-none-any.whl

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

@@ -1,32 +1,46 @@
+import contextlib
 import typing as ta
 
 from omlish import check
 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
+
+
+with lang.auto_proxy_import(globals()):
+    from .....backends import mlx as mlxu
 
 
 ##
 
 
 # @omlish-manifest $.minichain.backends.strings.manifests.BackendStringsManifest(
-#     ['ChatChoicesService'],
+#     ['ChatChoicesService', 'ChatChoicesStreamService'],
 #     'mlx',
 # )
 
@@ -34,12 +48,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 +61,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,17 +79,14 @@ 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:
             raise TypeError(m)
 
     @lang.cached_function(transient=True)
-    def _load_model(self) -> mlxu.LoadedModel:
+    def _load_model(self) -> 'mlxu.LoadedModel':
         # FIXME: walk state, find all mx.arrays, dealloc/set to empty
         check.not_none(self._exit_stack)
 
@@ -96,10 +102,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 +112,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)
+
+    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
 
-        kwargs = dict()
+
+# @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 +160,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)

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

@@ -0,0 +1,199 @@
+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.io.buffers import DelimitingBuffer
+
+from .....backends.ollama import protocol as pt
+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.messages import AiMessage
+from ....chat.messages import AnyAiMessage
+from ....chat.messages import Message
+from ....chat.messages import SystemMessage
+from ....chat.messages import UserMessage
+from ....chat.stream.types import ContentAiDelta
+from ....models.configs import ModelName
+from ....resources import UseResources
+from ....standard import ApiUrl
+from ....stream.services import StreamResponseSink
+from ....stream.services import new_stream_response
+
+
+##
+
+
+# @omlish-manifest $.minichain.backends.strings.manifests.BackendStringsManifest(
+#     [
+#         'ChatChoicesService',
+#         'ChatChoicesStreamService',
+#     ],
+#     'ollama',
+# )
+
+
+##
+
+
+class BaseOllamaChatChoicesService(lang.Abstract):
+    DEFAULT_API_URL: ta.ClassVar[ApiUrl] = ApiUrl('http://localhost:11434/api')
+    DEFAULT_MODEL_NAME: ta.ClassVar[ModelName] = ModelName('llama3.2')
+
+    def __init__(
+            self,
+            *configs: ApiUrl | ModelName,
+            http_client: http.AsyncHttpClient | None = None,
+    ) -> None:
+        super().__init__()
+
+        self._http_client = http_client
+
+        with tv.consume(*configs) as cc:
+            self._api_url = cc.pop(self.DEFAULT_API_URL)
+            self._model_name = cc.pop(self.DEFAULT_MODEL_NAME)
+
+    #
+
+    ROLE_MAP: ta.ClassVar[ta.Mapping[type[Message], pt.Role]] = {  # noqa
+        SystemMessage: 'system',
+        UserMessage: 'user',
+        AiMessage: 'assistant',
+    }
+
+    @classmethod
+    def _get_message_content(cls, m: Message) -> str | None:
+        if isinstance(m, (AiMessage, UserMessage, SystemMessage)):
+            return check.isinstance(m.c, str)
+        else:
+            raise TypeError(m)
+
+    @classmethod
+    def _build_request_messages(cls, mc_msgs: ta.Iterable[Message]) -> ta.Sequence[pt.Message]:
+        messages: list[pt.Message] = []
+        for m in mc_msgs:
+            messages.append(pt.Message(
+                role=cls.ROLE_MAP[type(m)],
+                content=cls._get_message_content(m),
+            ))
+        return messages
+
+
+##
+
+
+# @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
+#     name='ollama',
+#     type='ChatChoicesService',
+# )
+@static_check_is_chat_choices_service
+class OllamaChatChoicesService(BaseOllamaChatChoicesService):
+    async def invoke(
+            self,
+            request: ChatChoicesRequest,
+    ) -> ChatChoicesResponse:
+        messages = self._build_request_messages(request.v)
+
+        a_req = pt.ChatRequest(
+            model=self._model_name.v,
+            messages=messages,
+            # tools=tools or None,
+            stream=False,
+        )
+
+        raw_request = msh.marshal(a_req)
+
+        async with http.manage_async_client(self._http_client) as http_client:
+            raw_response = await http_client.request(http.HttpRequest(
+                self._api_url.v.removesuffix('/') + '/chat',
+                data=json.dumps(raw_request).encode('utf-8'),
+            ))
+
+        json_response = json.loads(check.not_none(raw_response.data).decode('utf-8'))
+
+        resp = msh.unmarshal(json_response, pt.ChatResponse)
+
+        out: list[AnyAiMessage] = []
+        if resp.message.role == 'assistant':
+            out.append(AiMessage(
+                check.not_none(resp.message.content),
+            ))
+        else:
+            raise TypeError(resp.message.role)
+
+        return ChatChoicesResponse([
+            AiChoice(out),
+        ])
+
+
+##
+
+
+# @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
+#     name='ollama',
+#     type='ChatChoicesStreamService',
+# )
+@static_check_is_chat_choices_stream_service
+class OllamaChatChoicesStreamService(BaseOllamaChatChoicesService):
+    READ_CHUNK_SIZE: ta.ClassVar[int] = -1
+
+    async def invoke(
+            self,
+            request: ChatChoicesStreamRequest,
+    ) -> ChatChoicesStreamResponse:
+        messages = self._build_request_messages(request.v)
+
+        a_req = pt.ChatRequest(
+            model=self._model_name.v,
+            messages=messages,
+            # tools=tools or None,
+            stream=True,
+        )
+
+        raw_request = msh.marshal(a_req)
+
+        http_request = http.HttpRequest(
+            self._api_url.v.removesuffix('/') + '/chat',
+            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] | None:
+                db = DelimitingBuffer([b'\r', b'\n', b'\r\n'])
+                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 []
+
+                        lj = json.loads(l.decode('utf-8'))
+                        lp: pt.ChatResponse = msh.unmarshal(lj, pt.ChatResponse)
+
+                        check.state(lp.message.role == 'assistant')
+                        check.none(lp.message.tool_name)
+                        check.state(not lp.message.tool_calls)
+
+                        if (c := lp.message.content):
+                            await sink.emit(AiChoicesDeltas([AiChoiceDeltas([ContentAiDelta(
+                                c,
+                            )])]))
+
+                    if not b:
+                        return []
+
+            return await new_stream_response(rs, inner)

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

@@ -14,10 +14,12 @@ 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
@@ -25,7 +27,8 @@ from ....models.configs import ModelName
 from ....standard import ApiKey
 from ....standard import DefaultOptions
 from .format import OpenaiChatRequestHandler
-from .names import MODEL_NAMES
+from .format import build_mc_choices_response
+from .names import CHAT_MODEL_NAMES
 
 
 ##
@@ -37,17 +40,23 @@ from .names import MODEL_NAMES
 # )
 @static_check_is_chat_choices_service
 class OpenaiChatChoicesService:
-    DEFAULT_MODEL_NAME: ta.ClassVar[ModelName] = ModelName(check.not_none(MODEL_NAMES.default))
+    DEFAULT_MODEL_NAME: ta.ClassVar[ModelName] = ModelName(check.not_none(CHAT_MODEL_NAMES.default))
 
-    def __init__(self, *configs: ApiKey | ModelName | DefaultOptions) -> None:
+    def __init__(
+            self,
+            *configs: ApiKey | ModelName | DefaultOptions,
+            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='OPENAI_API_KEY')
             self._default_options: tv.TypedValues = DefaultOptions.pop(cc)
 
-    def invoke(self, request: ChatChoicesRequest) -> ChatChoicesResponse:
+    async def invoke(self, request: ChatChoicesRequest) -> ChatChoicesResponse:
         # check.isinstance(request, ChatRequest)
 
         rh = OpenaiChatRequestHandler(
@@ -57,23 +66,24 @@ class OpenaiChatChoicesService:
                 *request.options,
                 override=True,
             ),
-            model=MODEL_NAMES.resolve(self._model_name.v),
+            model=CHAT_MODEL_NAMES.resolve(self._model_name.v),
             mandatory_kwargs=dict(
                 stream=False,
             ),
         )
 
-        raw_request = rh.raw_request()
+        raw_request = msh.marshal(rh.oai_request())
 
-        http_response = http.request(
+        http_response = await http.async_request(
             'https://api.openai.com/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'),
+            client=self._http_client,
         )
 
         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/completion.py

@@ -23,13 +23,19 @@ from ....standard import ApiKey
 class OpenaiCompletionService:
     DEFAULT_MODEL_NAME: ta.ClassVar[str] = 'gpt-3.5-turbo-instruct'
 
-    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._api_key = ApiKey.pop_secret(cc, env='OPENAI_API_KEY')
 
-    def invoke(self, t: CompletionRequest) -> CompletionResponse:
+    async def invoke(self, t: CompletionRequest) -> CompletionResponse:
         raw_request = dict(
             model=self.DEFAULT_MODEL_NAME,
             prompt=t.v,
@@ -41,13 +47,14 @@ class OpenaiCompletionService:
             stream=False,
         )
 
-        raw_response = http.request(
+        raw_response = await http.async_request(
             'https://api.openai.com/v1/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'),
+            client=self._http_client,
         )
 
         response = json.loads(check.not_none(raw_response.data).decode('utf-8'))

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

@@ -22,25 +22,32 @@ from ....vectors.types import Vector
 class OpenaiEmbeddingService:
     model = 'text-embedding-3-small'
 
-    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._api_key = ApiKey.pop_secret(cc, env='OPENAI_API_KEY')
 
-    def invoke(self, request: EmbeddingRequest) -> EmbeddingResponse:
+    async def invoke(self, request: EmbeddingRequest) -> EmbeddingResponse:
         raw_request = dict(
             model=self.model,
             input=check.isinstance(request.v, str),
         )
 
-        raw_response = http.request(
+        raw_response = await http.async_request(
             'https://api.openai.com/v1/embeddings',
             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'),
+            client=self._http_client,
         )
 
         response = json.loads(check.not_none(raw_response.data).decode('utf-8'))