ommlds 0.0.0.dev466__py3-none-any.whl → 0.0.0.dev512__py3-none-any.whl

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

@@ -4,32 +4,57 @@ TODO:
  - https://huggingface.co/blog/aifeifei798/transformers-streaming-output
 """
 import sys
+import threading
 import typing as ta
 
-import transformers as tfm
-
 from omlish import check
 from omlish import lang
 from omlish import typedvalues as tv
+from omlish.asyncs.asyncio.sync import AsyncioBufferRelay
 
 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  # noqa
+from ....chat.choices.stream.types import AiChoicesDeltas  # noqa
 from ....chat.choices.types import AiChoice
+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  # noqa
 from ....completion import CompletionRequest
 from ....completion import CompletionResponse
 from ....completion import static_check_is_completion_service
 from ....configs import Config
 from ....models.configs import ModelPath
+from ....resources import UseResources
+from ....stream.services import StreamResponseSink
+from ....stream.services import new_stream_response
 from ...impls.huggingface.configs import HuggingfaceHubToken
 
 
+with lang.auto_proxy_import(globals()):
+    import transformers as tfm
+
+    from .....backends import transformers as tfm_u
+
+
+##
+
+
+# @omlish-manifest $.minichain.backends.strings.manifests.BackendStringsManifest(
+#     ['ChatChoicesService', 'ChatChoicesStreamService'],
+#     'transformers',
+# )
+
+
 ##
 
 
@@ -128,13 +153,10 @@ def build_chat_message(m: Message) -> ta.Mapping[str, ta.Any]:
         raise TypeError(m)
 
 
-# @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
-#     name='transformers',
-#     aliases=['tfm'],
-#     type='ChatChoicesService',
-# )
-@static_check_is_chat_choices_service
-class TransformersChatChoicesService(lang.ExitStacked):
+##
+
+
+class BaseTransformersChatChoicesService(lang.ExitStacked):
     DEFAULT_MODEL: ta.ClassVar[str] = (
         'meta-llama/Llama-3.2-1B-Instruct'
     )
@@ -148,7 +170,7 @@ class TransformersChatChoicesService(lang.ExitStacked):
             self._huggingface_hub_token = HuggingfaceHubToken.pop_secret(cc, env='HUGGINGFACE_HUB_TOKEN')
 
     @lang.cached_function(transient=True)
-    def _load_pipeline(self) -> tfm.Pipeline:
+    def _load_pipeline(self) -> 'tfm.Pipeline':
         # FIXME: unload
         check.not_none(self._exit_stack)
 
@@ -161,21 +183,118 @@ class TransformersChatChoicesService(lang.ExitStacked):
         for pkw_cfg in self._pipeline_kwargs:
             pkw.update(pkw_cfg.v)
 
-        return tfm.pipeline(
-            'text-generation',
-            **pkw,
-        )
+        with tfm_u.file_cache_patch_context(
+                local_first=True,
+                local_config_present_is_authoritative=True,
+        ):
+            return tfm.pipeline(
+                'text-generation',
+                **pkw,
+            )
 
+
+##
+
+
+# @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
+#     name='transformers',
+#     aliases=['tfm'],
+#     type='ChatChoicesService',
+# )
+@static_check_is_chat_choices_service
+class TransformersChatChoicesService(BaseTransformersChatChoicesService):
     async def invoke(self, request: ChatChoicesRequest) -> ChatChoicesResponse:
         check.empty(request.options)
 
         pipeline = self._load_pipeline()
 
-        output = pipeline(
-            [
-                build_chat_message(m)
-                for m in request.v
-            ],
+        inputs = [
+            build_chat_message(m)
+            for m in request.v
+        ]
+
+        outputs = pipeline(inputs)
+
+        gts = check.single(outputs)['generated_text']
+        ugt, agt = gts
+        check.state(ugt['role'] == 'user')
+        check.state(agt['role'] == 'assistant')
+
+        return ChatChoicesResponse([AiChoice([AiMessage(agt['content'])])])
+
+
+##
+
+
+# @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
+#     name='transformers',
+#     type='ChatChoicesStreamService',
+# )
+@static_check_is_chat_choices_stream_service
+class TransformersChatChoicesStreamService(BaseTransformersChatChoicesService):
+    async def invoke(self, request: ChatChoicesStreamRequest) -> ChatChoicesStreamResponse:
+        check.empty(request.options)
+
+        pipeline = self._load_pipeline()  # noqa
+
+        inputs = [  # noqa
+            build_chat_message(m)
+            for m in request.v
+        ]
+
+        relay: AsyncioBufferRelay = AsyncioBufferRelay()
+
+        def streamer_callback(text: str, *, stream_end: bool) -> None:
+            if text or stream_end:
+                relay.push(text, *([None] if stream_end else []))
+
+        streamer = tfm_u.CancellableTextStreamer(
+            check.not_none(pipeline.tokenizer),  # type: ignore[arg-type]
+            streamer_callback,  # noqa
+            skip_prompt=True,
+            skip_special_tokens=True,
         )
 
-        return ChatChoicesResponse([AiChoice([output])])
+        async with UseResources.or_new(request.options) as rs:
+            thread = threading.Thread(
+                target=tfm_u.CancellableTextStreamer.ignoring_cancelled(pipeline),
+                args=(
+                    inputs,
+                ),
+                kwargs=dict(
+                    streamer=streamer,
+                ),
+            )
+
+            def stop_thread() -> None:
+                streamer.cancel()
+                # thread.join()
+
+            rs.enter_context(lang.defer(stop_thread))
+
+            thread.start()
+
+            async def inner(sink: StreamResponseSink[AiChoicesDeltas]) -> ta.Sequence[ChatChoicesOutputs] | None:
+                while True:
+                    await relay.wait()
+                    got = relay.swap()
+
+                    if not got:
+                        raise RuntimeError
+
+                    if got[-1] is None:
+                        out = ''.join(got[:-1])
+                        end = True
+                    else:
+                        out = ''.join(got)
+                        end = False
+
+                    if out:
+                        await sink.emit(AiChoicesDeltas([AiChoiceDeltas([ContentAiDelta(out)])]))
+
+                    if end:
+                        break
+
+                return []
+
+            return await new_stream_response(rs, inner)

ommlds/minichain/backends/strings/parsing.py

@@ -11,7 +11,7 @@ from omlish import lang
 
 @dc.dataclass(frozen=True, kw_only=True)
 class ParsedBackendString:
-    backend: str | None = dc.xfield(None, repr_fn=dc.opt_repr)
+    backend: str | None = dc.xfield(None, repr_fn=lang.opt_repr)
 
     #
 
@@ -56,7 +56,7 @@ _REPO_MODEL_PAT = re.compile(
 def parse_backend_string(s: str) -> ParsedBackendString:
     backend: str | None
     if ':' in s:
-        backend, s = s.split(':')
+        backend, _, s = s.partition(':')
     else:
         backend = None
 

ommlds/minichain/backends/strings/resolving.py

@@ -4,9 +4,9 @@ TODO:
   - interop with registry somehow, probably? or is it strictly a different concern?
 """
 import abc
-import dataclasses as dc
 import typing as ta
 
+from omlish import dataclasses as dc
 from omlish import lang
 from omlish.manifests.globals import GlobalManifestLoader
 
@@ -97,6 +97,8 @@ class ManifestBackendStringResolver(BackendStringResolver):
             for scn in manifest.service_cls_names
         }
 
+    __repr__ = lang.attr_ops('_manifest').repr
+
     def _resolve_name_model(
             self,
             args: ResolveBackendStringArgs,
@@ -108,7 +110,10 @@ class ManifestBackendStringResolver(BackendStringResolver):
 
         mn: str | None = mdl.name
 
-        if mn == m.backend_name:
+        if args.parsed.backend == m.backend_name and mn is not None:
+            pass
+
+        elif mn == m.backend_name:
             if m.model_names is not None:
                 mn = m.model_names.resolved_default
             else:

ommlds/minichain/chat/{stream → choices/stream}/adapters.py

@@ -1,11 +1,11 @@
 from omlish import dataclasses as dc
 
-from ...services import Response
-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 .joining import AiChoiceDeltaJoiner
+from ....services import Response
+from ..services import ChatChoicesRequest
+from ..services import static_check_is_chat_choices_service
+from ..types import AiChoice
+from ..types import AiChoices
+from .joining import AiChoicesDeltaJoiner
 from .services import ChatChoicesOutputs
 from .services import ChatChoicesStreamOutputs
 from .services import ChatChoicesStreamService
@@ -23,7 +23,7 @@ class ChatChoicesStreamServiceChatChoicesService:
         AiChoices,
         ChatChoicesOutputs | ChatChoicesStreamOutputs,
     ]:
-        joiner = AiChoiceDeltaJoiner()
+        joiner = AiChoicesDeltaJoiner()
 
         async with (resp := await self.service.invoke(request)).v as it:  # noqa
             async for cs in it:

ommlds/minichain/chat/choices/stream/joining.py

@@ -0,0 +1,31 @@
+import typing as ta
+
+from omlish import check
+
+from ...messages import AiChat
+from ...stream.joining import AiDeltaJoiner
+from .types import AiChoiceDeltas
+
+
+##
+
+
+class AiChoicesDeltaJoiner:
+    def __init__(self) -> None:
+        super().__init__()
+
+        self._seq = 0
+        self._channels: list[AiDeltaJoiner] = []
+
+    def add(self, choices: ta.Sequence[AiChoiceDeltas]) -> None:
+        if not self._seq:
+            check.empty(self._channels)
+            self._channels.extend(AiDeltaJoiner() for _ in range(len(choices)))
+
+        for chan, c in zip(self._channels, choices, strict=True):
+            chan.add(c.deltas)
+
+        self._seq += 1
+
+    def build(self) -> list[AiChat]:
+        return [list(chan.build()) for chan in self._channels]

ommlds/minichain/chat/choices/stream/services.py

@@ -0,0 +1,45 @@
+import abc
+import typing as ta
+
+from omlish import lang
+
+from ....registries.globals import register_type
+from ....services import Request
+from ....services import Service
+from ....stream.services import StreamResponse
+from ...messages import Chat
+from ..types import ChatChoicesOutputs
+from .types import AiChoicesDeltas
+from .types import ChatChoicesStreamOptions
+from .types import ChatChoicesStreamOutputs
+
+
+##
+
+
+ChatChoicesStreamRequest: ta.TypeAlias = Request[Chat, ChatChoicesStreamOptions]
+
+ChatChoicesStreamResponse: ta.TypeAlias = StreamResponse[
+    AiChoicesDeltas,
+    ChatChoicesOutputs,
+    ChatChoicesStreamOutputs,
+]
+
+# @omlish-manifest $.minichain.registries.manifests.RegistryTypeManifest
+ChatChoicesStreamService: ta.TypeAlias = Service[ChatChoicesStreamRequest, ChatChoicesStreamResponse]
+
+register_type(ChatChoicesStreamService, module=__name__)
+
+
+def static_check_is_chat_choices_stream_service[T: ChatChoicesStreamService](t: type[T]) -> type[T]:
+    return t
+
+
+##
+
+
+@static_check_is_chat_choices_stream_service
+class AbstractChatChoicesStreamService(lang.Abstract):
+    @abc.abstractmethod
+    def invoke(self, request: ChatChoicesStreamRequest) -> ta.Awaitable[ChatChoicesStreamResponse]:
+        raise NotImplementedError

ommlds/minichain/chat/choices/stream/types.py

@@ -0,0 +1,43 @@
+import typing as ta
+
+from omlish import dataclasses as dc
+from omlish import lang
+
+from ....stream.services import StreamOptions
+from ....types import Option
+from ....types import Output
+from ...stream.types import AiDeltas
+from ..types import ChatChoicesOptions
+
+
+##
+
+
+class ChatChoicesStreamOption(Option, lang.Abstract, lang.PackageSealed):
+    pass
+
+
+ChatChoicesStreamOptions: ta.TypeAlias = ChatChoicesStreamOption | StreamOptions | ChatChoicesOptions
+
+
+##
+
+
+class ChatChoicesStreamOutput(Output, lang.Abstract, lang.PackageSealed):
+    pass
+
+
+ChatChoicesStreamOutputs: ta.TypeAlias = ChatChoicesStreamOutput
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class AiChoiceDeltas(lang.Final):
+    deltas: AiDeltas
+
+
+@dc.dataclass(frozen=True)
+class AiChoicesDeltas(lang.Final):
+    choices: ta.Sequence[AiChoiceDeltas]

ommlds/minichain/chat/content.py

@@ -0,0 +1,42 @@
+import typing as ta
+
+from ..content.content import Content
+from ..content.transform.strings import transform_content_strings
+from .messages import AiMessage
+from .messages import Message
+from .messages import SystemMessage
+from .messages import ToolUseMessage
+from .messages import ToolUseResultMessage
+from .messages import UserMessage
+
+
+MessageT = ta.TypeVar('MessageT', bound=Message)
+
+
+##
+
+
+def transform_message_content(fn: ta.Callable[[Content], Content], m: MessageT) -> MessageT:
+    if isinstance(m, UserMessage):
+        return m.replace(c=fn(m.c))
+
+    elif isinstance(m, AiMessage):
+        return m.replace(c=fn(m.c))
+
+    elif isinstance(m, SystemMessage):
+        return m.replace(c=fn(m.c))
+
+    elif isinstance(m, ToolUseMessage):
+        # TODO: m.tu.args?
+        return m
+
+    elif isinstance(m, ToolUseResultMessage):
+        # TODO: m.tur.c?
+        return m
+
+    else:
+        raise TypeError(m)
+
+
+def transform_message_content_strings(fn: ta.Callable[[str], str], m: MessageT) -> MessageT:
+    return transform_message_content(lambda c: transform_content_strings(fn, c), m)

ommlds/minichain/chat/messages.py

@@ -1,20 +1,20 @@
+"""
+TODO:
+ - channel? reasoning / thinking?
+"""
 import operator
 import typing as ta
 
 from omlish import check
 from omlish import dataclasses as dc
-from omlish import dispatch
 from omlish import lang
 from omlish import marshal as msh
-from omlish import typedvalues as tv
 
-from .._typedvalues import _tv_field_metadata
-from ..content.materialize import CanContent
-from ..content.transforms.base import ContentTransform
-from ..content.types import Content
-from ..metadata import MetadataContainer
+from ..content.content import Content
+from ..metadata import MetadataContainerDataclass
 from ..tools.types import ToolUse
 from ..tools.types import ToolUseResult
+from .metadata import MessageMetadata
 from .metadata import MessageMetadatas
 
 
@@ -26,7 +26,7 @@ msh.register_global_module_import('._marshal', __package__)
 
 @dc.dataclass(frozen=True)
 class Message(  # noqa
-    MetadataContainer[MessageMetadatas],
+    MetadataContainerDataclass[MessageMetadatas],
     lang.Abstract,
     lang.Sealed,
 ):
@@ -34,18 +34,14 @@ class Message(  # noqa
         default=(),
         kw_only=True,
         repr=False,
-        metadata=_tv_field_metadata(
-            MessageMetadatas,
-            marshal_name='metadata',
-        ),
     )
 
-    @property
-    def metadata(self) -> tv.TypedValues[MessageMetadatas]:
-        return check.isinstance(self._metadata, tv.TypedValues)
+    MetadataContainerDataclass._configure_metadata_field(_metadata, MessageMetadatas)  # noqa
 
-    def with_metadata(self, *mds: MessageMetadatas, override: bool = False) -> ta.Self:
-        return dc.replace(self, _metadata=tv.TypedValues(*self._metadata, *mds, override=override))
+    def replace(self, **kwargs: ta.Any) -> ta.Self:
+        if (n := dc.replace_is_not(self, **kwargs)) is self:
+            return self
+        return n.with_metadata(MessageOriginal(self), discard=[MessageOriginal], override=True)
 
 
 Chat: ta.TypeAlias = ta.Sequence[Message]
@@ -54,6 +50,14 @@ Chat: ta.TypeAlias = ta.Sequence[Message]
 ##
 
 
+@dc.dataclass(frozen=True)
+class MessageOriginal(MessageMetadata, lang.Final):
+    c: Message
+
+
+##
+
+
 @dc.dataclass(frozen=True)
 class AnyUserMessage(Message, lang.Abstract):
     pass
@@ -90,18 +94,18 @@ def check_ai_chat(chat: Chat) -> AiChat:
 
 @dc.dataclass(frozen=True)
 class SystemMessage(AnyUserMessage, lang.Final):
-    c: CanContent
+    c: Content
 
 
 #
 
 
 @dc.dataclass(frozen=True)
-@msh.update_fields_metadata(['name'], omit_if=operator.not_)
+@msh.update_fields_options(['name'], omit_if=operator.not_)
 class UserMessage(AnyUserMessage, lang.Final):
-    c: CanContent
+    c: Content
 
-    name: str | None = dc.xfield(None, repr_fn=dc.opt_repr)
+    name: str | None = dc.xfield(None, repr_fn=lang.opt_repr)
 
 
 #
@@ -109,7 +113,7 @@ class UserMessage(AnyUserMessage, lang.Final):
 
 @dc.dataclass(frozen=True)
 class AiMessage(AnyAiMessage, lang.Final):
-    c: Content = dc.xfield(None, repr_fn=dc.opt_repr)  # TODO: non-null?
+    c: Content = dc.xfield(None, repr_fn=lang.opt_repr)  # TODO: non-null?
 
 
 #
@@ -128,23 +132,23 @@ class ToolUseResultMessage(AnyUserMessage, lang.Final):
 ##
 
 
-class _MessageContentTransform(ContentTransform, lang.Final, lang.NotInstantiable):
-    @dispatch.install_method(ContentTransform.apply)
-    def apply_system_message(self, m: SystemMessage) -> SystemMessage:
-        return dc.replace(m, c=self.apply(m.c))
-
-    @dispatch.install_method(ContentTransform.apply)
-    def apply_user_message(self, m: UserMessage) -> UserMessage:
-        return dc.replace(m, c=self.apply(m.c))
-
-    @dispatch.install_method(ContentTransform.apply)
-    def apply_ai_message(self, m: AiMessage) -> AiMessage:
-        return dc.replace(m, c=self.apply(m.c))
-
-    @dispatch.install_method(ContentTransform.apply)
-    def apply_tool_use_message(self, m: ToolUseMessage) -> ToolUseMessage:
-        return dc.replace(m, tu=self.apply(m.tu))
-
-    @dispatch.install_method(ContentTransform.apply)
-    def apply_tool_use_result_message(self, m: ToolUseResultMessage) -> ToolUseResultMessage:
-        return dc.replace(m, tur=self.apply(m.tur))
+# class _MessageContentTransform(ContentTransform, lang.Final, lang.NotInstantiable):
+#     @dispatch.install_method(ContentTransform.apply)
+#     def apply_system_message(self, m: SystemMessage) -> SystemMessage:
+#         return dc.replace(m, c=self.apply(m.c))
+#
+#     @dispatch.install_method(ContentTransform.apply)
+#     def apply_user_message(self, m: UserMessage) -> UserMessage:
+#         return dc.replace(m, c=self.apply(m.c))
+#
+#     @dispatch.install_method(ContentTransform.apply)
+#     def apply_ai_message(self, m: AiMessage) -> AiMessage:
+#         return dc.replace(m, c=self.apply(m.c))
+#
+#     @dispatch.install_method(ContentTransform.apply)
+#     def apply_tool_use_message(self, m: ToolUseMessage) -> ToolUseMessage:
+#         return dc.replace(m, tu=self.apply(m.tu))
+#
+#     @dispatch.install_method(ContentTransform.apply)
+#     def apply_tool_use_result_message(self, m: ToolUseResultMessage) -> ToolUseResultMessage:
+#         return dc.replace(m, tur=self.apply(m.tur))

ommlds/minichain/chat/stream/_marshal.py

@@ -1,7 +1,7 @@
 from omlish import lang
 from omlish import marshal as msh
 
-from .types import AiChoiceDelta
+from .types import AiDelta
 
 
 ##
@@ -9,8 +9,8 @@ from .types import AiChoiceDelta
 
 @lang.static_init
 def _install_standard_marshaling() -> None:
-    acd_poly = msh.polymorphism_from_subclasses(AiChoiceDelta, naming=msh.Naming.SNAKE)
+    ad_poly = msh.polymorphism_from_subclasses(AiDelta, naming=msh.Naming.SNAKE)
     msh.install_standard_factories(
-        msh.PolymorphismMarshalerFactory(acd_poly),
-        msh.PolymorphismUnmarshalerFactory(acd_poly),
+        msh.PolymorphismMarshalerFactory(ad_poly),
+        msh.PolymorphismUnmarshalerFactory(ad_poly),
     )