ommlds 0.0.0.dev467__py3-none-any.whl → 0.0.0.dev469__py3-none-any.whl

ommlds/backends/transformers/filecache.py

@@ -0,0 +1,109 @@
+import contextlib
+import dataclasses as dc
+import os
+import threading
+import typing as ta
+
+import transformers as tfm
+
+from omlish import lang
+
+from ..._hacks.funcs import create_detour
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class _FileCachePatchContext:
+    local_first: bool = False
+    local_config_present_is_authoritative: bool = False
+
+
+_FILE_CACHE_PATCH_CONTEXT_TLS = threading.local()
+
+
+def _get_file_cache_patch_context() -> _FileCachePatchContext:
+    try:
+        return _FILE_CACHE_PATCH_CONTEXT_TLS.context
+    except AttributeError:
+        ctx = _FILE_CACHE_PATCH_CONTEXT_TLS.context = _FileCachePatchContext()
+        return ctx
+
+
+_FILE_CACHE_PATCH_LOCK = threading.Lock()
+
+
+@lang.cached_function(lock=_FILE_CACHE_PATCH_LOCK)
+def patch_file_cache() -> None:
+    """
+    I tried to make a `local_first_pipeline` function to be called instead of `tfm.pipeline`, I really did, but the
+    transformers code is such a disgusting rat's nest full of direct static calls to the caching code strewn about at
+    every layer with no concern whatsoever for forwarding kwargs where they need to go.
+    """
+
+    from transformers.utils.hub import cached_files
+
+    orig_cached_files: ta.Callable[..., str | None] = lang.copy_function(cached_files)  # type: ignore
+
+    get_file_cache_patch_context = _get_file_cache_patch_context
+
+    def new_cached_files(
+            path_or_repo_id: str | os.PathLike,
+            filenames: list[str],
+            **kwargs: ta.Any,
+    ) -> str | None:
+        ctx = get_file_cache_patch_context()
+
+        if ctx.local_first and not kwargs.get('local_files_only'):
+            try:
+                local = orig_cached_files(
+                    path_or_repo_id,
+                    filenames,
+                    **{**kwargs, 'local_files_only': True},
+                )
+            except OSError as e:  # noqa
+                pass
+            else:
+                return local
+
+            if ctx.local_config_present_is_authoritative:
+                try:
+                    local_config = orig_cached_files(
+                        path_or_repo_id,
+                        [tfm.CONFIG_NAME],
+                        **{**kwargs, 'local_files_only': True},
+                    )
+                except OSError as e:  # noqa
+                    pass
+                else:
+                    raise OSError(
+                        f'Files {filenames!r} requested under local_first '
+                        f'but local_config present at {local_config!r}, '
+                        f'assuming files do not exist.',
+                    )
+
+        return orig_cached_files(path_or_repo_id, filenames, **kwargs)
+
+    cached_files.__code__ = create_detour(cached_files, new_cached_files, as_kwargs=True)
+
+
+@contextlib.contextmanager
+def file_cache_patch_context(
+        *,
+        local_first: bool = False,
+        local_config_present_is_authoritative: bool = False,
+) -> ta.Generator[None]:
+    patch_file_cache()
+
+    new_ctx = dc.replace(
+        old_ctx := _get_file_cache_patch_context(),
+        local_first=local_first,
+        local_config_present_is_authoritative=local_config_present_is_authoritative,
+    )
+
+    _FILE_CACHE_PATCH_CONTEXT_TLS.context = new_ctx
+    try:
+        yield
+    finally:
+        _FILE_CACHE_PATCH_CONTEXT_TLS.context = old_ctx

ommlds/backends/transformers/streamers.py

@@ -0,0 +1,73 @@
+import functools
+import typing as ta
+
+import transformers as tfm
+
+
+T = ta.TypeVar('T')
+P = ta.ParamSpec('P')
+
+
+##
+
+
+class CancellableTextStreamer(tfm.TextStreamer):
+    class Callback(ta.Protocol):
+        def __call__(self, text: str, *, stream_end: bool) -> None: ...
+
+    def __init__(
+            self,
+            tokenizer: tfm.AutoTokenizer,
+            callback: Callback,
+            *,
+            skip_prompt: bool = False,
+            **decode_kwargs: ta.Any,
+    ) -> None:
+        super().__init__(
+            tokenizer,
+            skip_prompt=skip_prompt,
+            **decode_kwargs,
+        )
+
+        self.callback = callback
+
+    _cancelled: bool = False
+
+    #
+
+    @property
+    def cancelled(self) -> bool:
+        return self._cancelled
+
+    def cancel(self) -> None:
+        self._cancelled = True
+
+    class Cancelled(BaseException):  # noqa
+        pass
+
+    @staticmethod
+    def ignoring_cancelled(fn: ta.Callable[P, T]) -> ta.Callable[P, T | None]:
+        @functools.wraps(fn)
+        def inner(*args, **kwargs):
+            try:
+                return fn(*args, **kwargs)
+            except CancellableTextStreamer.Cancelled:
+                pass
+
+        return inner
+
+    def _maybe_raise_cancelled(self) -> None:
+        if self._cancelled:
+            raise CancellableTextStreamer.Cancelled
+
+    #
+
+    def put(self, value: ta.Any) -> None:
+        self._maybe_raise_cancelled()
+        super().put(value)
+        self._maybe_raise_cancelled()
+
+    def on_finalized_text(self, text: str, stream_end: bool = False) -> None:
+        self._maybe_raise_cancelled()
+        self.callback(text, stream_end=stream_end)
+        self._maybe_raise_cancelled()

ommlds/cli/main.py

@@ -39,10 +39,6 @@ else:
 
 
 async def _a_main(args: ta.Any = None) -> None:
-    logs.configure_standard_logging('INFO')
-
-    #
-
     parser = argparse.ArgumentParser()
     parser.add_argument('prompt', nargs='*')
 
@@ -64,6 +60,8 @@ async def _a_main(args: ta.Any = None) -> None:
     parser.add_argument('-E', '--embed', action='store_true')
     parser.add_argument('-j', '--image', action='store_true')
 
+    parser.add_argument('-v', '--verbose', action='store_true')
+
     parser.add_argument('--enable-fs-tools', action='store_true')
     parser.add_argument('--enable-todo-tools', action='store_true')
     parser.add_argument('--enable-unsafe-tools-do-not-use-lol', action='store_true')
@@ -74,6 +72,14 @@ async def _a_main(args: ta.Any = None) -> None:
 
     #
 
+    if args.verbose:
+        logs.configure_standard_logging('DEBUG')
+    else:
+        logs.configure_standard_logging('INFO')
+        logs.silence_noisy_loggers()
+
+    #
+
     content: mc.Content | None
 
     if args.image:
@@ -161,7 +167,7 @@ async def _a_main(args: ta.Any = None) -> None:
                 args.enable_test_weather_tool or
                 args.code
             ),
-            enabled_tools={
+            enabled_tools={  # noqa
                 *(['fs'] if args.enable_fs_tools else []),
                 *(['todo'] if args.enable_todo_tools else []),
                 *(['weather'] if args.enable_test_weather_tool else []),

ommlds/cli/sessions/chat/backends/catalog.py

@@ -32,7 +32,7 @@ class _CatalogBackendProvider(BackendProvider[ServiceT], lang.Abstract):
     @contextlib.asynccontextmanager
     async def _provide_backend(self, cls: type[ServiceT]) -> ta.AsyncIterator[ServiceT]:
         service: ServiceT
-        async with lang.async_or_sync_maybe_managing(self._catalog.get_backend(
+        async with lang.async_or_sync_maybe_managing(self._catalog.new_backend(
                 cls,
                 self._name,
                 *(self._configs or []),

ommlds/cli/sessions/completion/session.py

@@ -31,7 +31,7 @@ class CompletionSession(Session['CompletionSession.Config']):
         prompt = check.isinstance(self._config.content, str)
 
         mdl: mc.CompletionService
-        async with lang.async_maybe_managing(self._backend_catalog.get_backend(
+        async with lang.async_maybe_managing(self._backend_catalog.new_backend(
                 mc.CompletionService,
                 self._config.backend or DEFAULT_COMPLETION_MODEL_BACKEND,
         )) as mdl:

ommlds/cli/sessions/embedding/session.py

@@ -29,7 +29,7 @@ class EmbeddingSession(Session['EmbeddingSession.Config']):
 
     async def run(self) -> None:
         mdl: mc.EmbeddingService
-        async with lang.async_maybe_managing(self._backend_catalog.get_backend(
+        async with lang.async_maybe_managing(self._backend_catalog.new_backend(
             mc.EmbeddingService,
             self._config.backend or DEFAULT_EMBEDDING_MODEL_BACKEND,
         )) as mdl:

ommlds/minichain/__init__.py

@@ -322,6 +322,7 @@ with _lang.auto_proxy_init(
     ##
 
     from .registries.globals import (  # noqa
+        get_registry_cls,
         register_type,
         registry_new,
         registry_of,
@@ -558,6 +559,10 @@ with _lang.auto_proxy_init(
     )
 
     from .standard import (  # noqa
+        Device,
+
+        ApiUrl,
+
         ApiKey,
 
         DefaultOptions,

ommlds/minichain/backends/catalogs/base.py

@@ -3,15 +3,28 @@ import typing as ta
 
 from omlish import lang
 
+from ...configs import Config
+
+
+T = ta.TypeVar('T')
+
 
 ##
 
 
 class BackendCatalog(lang.Abstract):
+    class Backend(ta.NamedTuple):
+        factory: ta.Callable[..., ta.Any]
+        configs: ta.Sequence[Config] | None
+
     @abc.abstractmethod
-    def get_backend(self, service_cls: ta.Any, name: str, *args: ta.Any, **kwargs: ta.Any) -> ta.Any:
+    def get_backend(self, service_cls: type[T], name: str) -> Backend:
         raise NotImplementedError
 
+    def new_backend(self, service_cls: ta.Any, name: str, *args: ta.Any, **kwargs: ta.Any) -> ta.Any:
+        be = self.get_backend(service_cls, name)
+        return be.factory(*be.configs or [], *args, **kwargs)
+
     # #
     #
     # class Bound(lang.Final, ta.Generic[T]):

ommlds/minichain/backends/catalogs/simple.py

@@ -40,9 +40,9 @@ class SimpleBackendCatalog(BackendCatalog):
             sc_dct[e.name] = e
         self._dct = dct
 
-    def get_backend(self, service_cls: ta.Any, name: str, *args: ta.Any, **kwargs: ta.Any) -> ta.Any:
+    def get_backend(self, service_cls: ta.Any, name: str, *args: ta.Any, **kwargs: ta.Any) -> BackendCatalog.Backend:
         e = self._dct[service_cls][name]
-        return e.factory_fn(*args, **kwargs)
+        return BackendCatalog.Backend(e.factory_fn, None)
 
 
 ##

ommlds/minichain/backends/catalogs/strings.py

@@ -5,7 +5,7 @@ from omlish import check
 from ...models.configs import ModelPath
 from ...models.configs import ModelRepo
 from ...models.repos.resolving import ModelRepoResolver
-from ...registries.globals import registry_new
+from ...registries.globals import get_registry_cls
 from ..strings.parsing import parse_backend_string
 from ..strings.resolving import BackendStringResolver
 from ..strings.resolving import ResolveBackendStringArgs
@@ -30,14 +30,14 @@ class BackendStringBackendCatalog(BackendCatalog):
         self._string_resolver = string_resolver
         self._model_repo_resolver = model_repo_resolver
 
-    def get_backend(self, service_cls: ta.Any, name: str, *args: ta.Any, **kwargs: ta.Any) -> ta.Any:
+    def get_backend(self, service_cls: ta.Any, name: str, *args: ta.Any, **kwargs: ta.Any) -> BackendCatalog.Backend:
         ps = parse_backend_string(name)
         rs = check.not_none(self._string_resolver.resolve_backend_string(ResolveBackendStringArgs(
             service_cls,
             ps,
         )))
 
-        al = list(rs.args or [])
+        al: list = list(rs.args or [])
 
         # FIXME: lol
         if al and isinstance(al[0], ModelRepo):
@@ -46,10 +46,12 @@ class BackendStringBackendCatalog(BackendCatalog):
             mrp = check.not_none(mrr.resolve(mr))
             al = [ModelPath(mrp.path), *al[1:]]
 
-        return registry_new(
+        cls = get_registry_cls(
             service_cls,
             rs.name,
-            *al,
-            *args,
-            **kwargs,
+        )
+
+        return BackendCatalog.Backend(
+            cls,
+            al,
         )

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

@@ -95,8 +95,7 @@ class AnthropicChatChoicesStreamService:
                 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 = http_response.stream.read1(self.READ_CHUNK_SIZE)
                     for l in db.feed(b):
                         if isinstance(l, DelimitingBuffer.Incomplete):
                             # FIXME: handle

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

@@ -169,8 +169,7 @@ class GoogleChatChoicesStreamService:
             async def inner(sink: StreamResponseSink[AiChoicesDeltas]) -> ta.Sequence[ChatChoicesOutputs] | None:
                 db = DelimitingBuffer([b'\r', b'\n', b'\r\n'])
                 while True:
-                    # FIXME: read1 not on response stream protocol
-                    b = http_response.stream.read1(self.READ_CHUNK_SIZE)  # type: ignore[attr-defined]
+                    b = http_response.stream.read1(self.READ_CHUNK_SIZE)
                     for bl in db.feed(b):
                         if isinstance(bl, DelimitingBuffer.Incomplete):
                             # FIXME: handle

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

@@ -30,6 +30,15 @@ from .format import get_msg_content
 ##
 
 
+# @omlish-manifest $.minichain.backends.strings.manifests.BackendStringsManifest(
+#     ['ChatChoicesService'],
+#     'llamacpp',
+# )
+
+
+##
+
+
 # @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
 #     name='llamacpp',
 #     type='ChatChoicesService',

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

@@ -29,6 +29,15 @@ from .format import get_msg_content
 ##
 
 
+# @omlish-manifest $.minichain.backends.strings.manifests.BackendStringsManifest(
+#     ['ChatChoicesStreamService'],
+#     'llamacpp',
+# )
+
+
+##
+
+
 # @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
 #     name='llamacpp',
 #     type='ChatChoicesStreamService',
@@ -76,18 +85,25 @@ class LlamacppChatChoicesStreamService(lang.ExitStacked):
             rs.enter_context(lang.defer(close_output))
 
             async def inner(sink: StreamResponseSink[AiChoicesDeltas]) -> ta.Sequence[ChatChoicesOutputs] | None:
+                last_role: ta.Any = None
+
                 for chunk in output:
                     check.state(chunk['object'] == 'chat.completion.chunk')
-                    l: list[AiChoiceDeltas] = []
-                    for choice in chunk['choices']:
-                        # FIXME: check role is assistant
-                        # FIXME: stop reason
-                        if not (delta := choice.get('delta', {})):
-                            continue
-                        if not (content := delta.get('content', '')):
-                            continue
-                        l.append(AiChoiceDeltas([ContentAiChoiceDelta(content)]))
-                    await sink.emit(AiChoicesDeltas(l))
+
+                    choice = check.single(chunk['choices'])
+
+                    if not (delta := choice.get('delta', {})):
+                        continue
+
+                    # FIXME: check role is assistant
+                    if (role := delta.get('role')) != last_role:
+                        last_role = role
+
+                    # FIXME: stop reason
+
+                    if (content := delta.get('content', '')):
+                        await sink.emit(AiChoicesDeltas([AiChoiceDeltas([ContentAiChoiceDelta(content)])]))
+
                 return None
 
             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.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.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 AiChoiceDeltas
+from ....chat.stream.types import AiChoicesDeltas
+from ....chat.stream.types import ContentAiChoiceDelta
+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]] = {
+        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 = 64 * 1024
+
+    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([ContentAiChoiceDelta(
+                                c,
+                            )])]))
+
+                    if not b:
+                        return []
+
+            return await new_stream_response(rs, inner)

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

@@ -88,8 +88,7 @@ class OpenaiChatChoicesStreamService:
                 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 = http_response.stream.read1(self.READ_CHUNK_SIZE)
                     for l in db.feed(b):
                         if isinstance(l, DelimitingBuffer.Incomplete):
                             # FIXME: handle