ommlds 0.0.0.dev462__py3-none-any.whl → 0.0.0.dev463__py3-none-any.whl

ommlds/backends/mlx/loading.py

@@ -15,6 +15,63 @@ from .tokenization import load_tokenization
 ##
 
 
+def get_model_path(
+        path_or_hf_repo: str,
+        revision: str | None = None,
+) -> tuple[pathlib.Path, str | None]:
+    """
+    Ensures the model is available locally. If the path does not exist locally,
+    it is downloaded from the Hugging Face Hub.
+
+    Args:
+        path_or_hf_repo (str): The local path or Hugging Face repository ID of the model.
+        revision (str, optional): A revision id which can be a branch name, a tag, or a commit hash.
+
+    Returns:
+        Tuple[Path, str]: A tuple containing the local file path and the Hugging Face repo ID.
+    """
+
+    model_path = pathlib.Path(path_or_hf_repo)
+
+    if not model_path.exists():
+        from huggingface_hub import snapshot_download
+        hf_path = path_or_hf_repo
+        model_path = pathlib.Path(
+            snapshot_download(
+                path_or_hf_repo,
+                revision=revision,
+                allow_patterns=[
+                    '*.jinja',
+                    '*.json',
+                    '*.jsonl',
+                    '*.py',
+                    '*.txt',
+
+                    'model*.safetensors',
+
+                    '*.tiktoken',
+                    'tiktoken.model',
+                    'tokenizer.model',
+                ],
+            ),
+        )
+
+    else:
+        from huggingface_hub import ModelCard
+
+        card_path = model_path / 'README.md'
+        if card_path.is_file():
+            card = ModelCard.load(card_path)
+            hf_path = card.data.base_model
+        else:
+            hf_path = None
+
+    return model_path, hf_path
+
+
+##
+
+
 @dc.dataclass(frozen=True, kw_only=True)
 class LoadedModel:
     path: pathlib.Path
@@ -46,7 +103,7 @@ def load_model(
 ) -> LoadedModel:
     # FIXME: get_model_path return annotation is wrong:
     #   https://github.com/ml-explore/mlx-lm/blob/9ee2b7358f5e258af7b31a8561acfbbe56ad5085/mlx_lm/utils.py#L82
-    model_path_res = ta.cast(ta.Any, mlx_lm.utils.get_model_path(path_or_hf_repo))
+    model_path_res = ta.cast(ta.Any, get_model_path(path_or_hf_repo))
     if isinstance(model_path_res, tuple):
         model_path = check.isinstance(model_path_res[0], pathlib.Path)
     else:

ommlds/cli/main.py

@@ -27,6 +27,7 @@ from .sessions.base import Session
 from .sessions.chat.code import CodeChatSession
 from .sessions.chat.interactive import InteractiveChatSession
 from .sessions.chat.prompt import PromptChatSession
+from .sessions.chat2.session import Chat2Session
 from .sessions.completion.completion import CompletionSession
 from .sessions.embedding.embedding import EmbeddingSession
 from .tools.config import ToolsConfig
@@ -66,6 +67,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('-2', '--two', 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')
@@ -128,7 +131,16 @@ async def _a_main(args: ta.Any = None) -> None:
 
     session_cfg: Session.Config
 
-    if args.interactive:
+    if args.two:
+        session_cfg = Chat2Session.Config(
+            backend=args.backend,
+            model_name=args.model_name,
+            state='ephemeral',
+            initial_content=content,  # noqa
+            # dangerous_no_tool_confirmation=bool(args.dangerous_no_tool_confirmation),
+        )
+
+    elif args.interactive:
         session_cfg = InteractiveChatSession.Config(
             backend=args.backend,
             model_name=args.model_name,
@@ -183,7 +195,11 @@ async def _a_main(args: ta.Any = None) -> None:
     with inj.create_managed_injector(bind_main(
             session_cfg=session_cfg,
             tools_config=tools_config,
-            enable_backend_strings=isinstance(session_cfg, (PromptChatSession.Config, CodeChatSession.Config)),
+            enable_backend_strings=isinstance(session_cfg, (
+                Chat2Session.Config,
+                CodeChatSession.Config,
+                PromptChatSession.Config,
+            )),
     )) as injector:
         await injector[Session].run()
 

ommlds/cli/sessions/chat/state.py

@@ -86,7 +86,7 @@ class StateStorageChatStateManager(ChatStateManager):
     def get_state(self) -> ChatState:
         if self._state is not None:
             return self._state
-        state: ChatState | None = self._storage.load_state(self._key, ChatState)
+        state: ChatState | None = lang.sync_await(self._storage.load_state(self._key, ChatState))
         if state is None:
             state = ChatState()
         self._state = state
@@ -94,7 +94,7 @@ class StateStorageChatStateManager(ChatStateManager):
 
     def clear_state(self) -> ChatState:
         state = ChatState()
-        self._storage.save_state(self._key, state, ChatState)
+        lang.sync_await(self._storage.save_state(self._key, state, ChatState))
         self._state = state
         return state
 
@@ -105,6 +105,6 @@ class StateStorageChatStateManager(ChatStateManager):
             chat=[*state.chat, *chat_additions],
             updated_at=lang.utcnow(),
         )
-        self._storage.save_state(self._key, state, ChatState)
+        lang.sync_await(self._storage.save_state(self._key, state, ChatState))
         self._state = state
         return state

ommlds/cli/sessions/chat2/_inject.py

@@ -0,0 +1,105 @@
+from omlish import lang
+
+
+with lang.auto_proxy_init(globals()):
+    from .backends.catalog import (  # noqa
+        CatalogChatChoicesServiceBackendProvider,
+        CatalogChatChoicesStreamServiceBackendProvider,
+    )
+
+    from .backends.types import (  # noqa
+        BackendName,
+        BackendConfigs,
+        BackendProvider,
+        ChatChoicesServiceBackendProvider,
+        ChatChoicesStreamServiceBackendProvider,
+    )
+
+    from .chat.ai.rendering import (  # noqa
+        RenderingAiChatGenerator,
+        RenderingStreamAiChatGenerator,
+    )
+
+    from .chat.ai.services import (  # noqa
+        ChatChoicesServiceOptions,
+        ChatChoicesServiceAiChatGenerator,
+        ChatChoicesStreamServiceStreamAiChatGenerator,
+    )
+
+    from .chat.ai.types import (  # noqa
+        AiChatGenerator,
+        StreamAiChatGenerator,
+    )
+
+    from .chat.state.inmemory import (  # noqa
+        InMemoryChatStateManager,
+    )
+
+    from .chat.state.storage import (  # noqa
+        StateStorageChatStateManager,
+    )
+
+    from .chat.state.types import (  # noqa
+        ChatState,
+        ChatStateManager,
+    )
+
+    from .chat.user.interactive import (  # noqa
+        InteractiveUserChatInput,
+    )
+
+    from .chat.user.oneshot import (  # noqa
+        OneshotUserChatInputInitialChat,
+        OneshotUserChatInput,
+    )
+
+    from .chat.user.types import (  # noqa
+        UserChatInput,
+    )
+
+    from .content.messages import (  # noqa
+        MessageContentExtractor,
+        MessageContentExtractorImpl,
+    )
+
+    from .content.strings import (  # noqa
+        ContentStringifier,
+        ContentStringifierImpl,
+    )
+
+    from .rendering.markdown import (  # noqa
+        MarkdownContentRendering,
+        MarkdownStreamContentRendering,
+    )
+
+    from .rendering.raw import (  # noqa
+        RawContentRendering,
+        RawContentStreamRendering,
+    )
+
+    from .rendering.types import (  # noqa
+        ContentRendering,
+        StreamContentRendering,
+    )
+
+    from .tools.confirmation import (  # noqa
+        ToolExecutionRequestDeniedError,
+        ToolExecutionConfirmation,
+        InteractiveToolExecutionConfirmation,
+    )
+
+    from .tools.execution import (  # noqa
+        ToolUseExecutor,
+        ToolUseExecutorImpl,
+    )
+
+    from .driver import (  # noqa
+        ChatDriver,
+    )
+
+    from .phases import (  # noqa
+        ChatPhase,
+        ChatPhaseCallback,
+        ChatPhaseCallbacks,
+        ChatPhaseManager,
+    )

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

@@ -0,0 +1,56 @@
+import contextlib
+import typing as ta
+
+from omlish import lang
+
+from ..... import minichain as mc
+from .types import BackendConfigs
+from .types import BackendName
+from .types import BackendProvider
+from .types import ChatChoicesServiceBackendProvider
+from .types import ChatChoicesStreamServiceBackendProvider
+from .types import ServiceT
+
+
+##
+
+
+class _CatalogBackendProvider(BackendProvider[ServiceT], lang.Abstract):
+    def __init__(
+            self,
+            *,
+            name: BackendName,
+            catalog: mc.BackendCatalog,
+            configs: BackendConfigs | None = None,
+    ) -> None:
+        super().__init__()
+
+        self._name = name
+        self._catalog = catalog
+        self._configs = configs
+
+    @contextlib.asynccontextmanager
+    async def _provide_backend(self, cls: type[ServiceT]) -> ta.AsyncIterator[ServiceT]:
+        service: ServiceT
+        async with lang.async_maybe_managing(self._catalog.get_backend(
+                cls,
+                self._name,
+                *(self._configs or []),
+        )) as service:
+            yield service
+
+
+class CatalogChatChoicesServiceBackendProvider(
+    _CatalogBackendProvider[mc.ChatChoicesService],
+    ChatChoicesServiceBackendProvider,
+):
+    def provide_backend(self) -> ta.AsyncContextManager[mc.ChatChoicesService]:
+        return self._provide_backend(mc.ChatChoicesService)  # type: ignore[type-abstract]
+
+
+class CatalogChatChoicesStreamServiceBackendProvider(
+    _CatalogBackendProvider[mc.ChatChoicesStreamService],
+    ChatChoicesStreamServiceBackendProvider,
+):
+    def provide_backend(self) -> ta.AsyncContextManager[mc.ChatChoicesStreamService]:
+        return self._provide_backend(mc.ChatChoicesStreamService)  # type: ignore[type-abstract]

ommlds/cli/sessions/chat2/backends/types.py

@@ -0,0 +1,36 @@
+import abc
+import typing as ta
+
+from omlish import lang
+
+from ..... import minichain as mc
+
+
+ServiceT = ta.TypeVar('ServiceT', bound=mc.Service)
+
+
+##
+
+
+BackendName = ta.NewType('BackendName', str)
+BackendConfigs = ta.NewType('BackendConfigs', ta.Sequence[mc.Config])
+
+
+##
+
+
+class BackendProvider(lang.Abstract, ta.Generic[ServiceT]):
+    @abc.abstractmethod
+    def provide_backend(self) -> ta.AsyncContextManager[ServiceT]:
+        raise NotImplementedError
+
+
+##
+
+
+class ChatChoicesServiceBackendProvider(BackendProvider[mc.ChatChoicesService], lang.Abstract):
+    pass
+
+
+class ChatChoicesStreamServiceBackendProvider(BackendProvider[mc.ChatChoicesStreamService], lang.Abstract):
+    pass

ommlds/cli/sessions/chat2/chat/ai/rendering.py

@@ -0,0 +1,67 @@
+import typing as ta
+
+from ...... import minichain as mc
+from ...content.messages import MessageContentExtractor
+from ...content.messages import MessageContentExtractorImpl
+from ...rendering.types import ContentRendering
+from .types import AiChatGenerator
+from .types import StreamAiChatGenerator
+
+
+##
+
+
+class RenderingAiChatGenerator(AiChatGenerator):
+    def __init__(
+            self,
+            *,
+            wrapped: AiChatGenerator,
+            extractor: MessageContentExtractor | None = None,
+            renderer: ContentRendering,
+    ) -> None:
+        super().__init__()
+
+        self._wrapped = wrapped
+        if extractor is None:
+            extractor = MessageContentExtractorImpl()
+        self._extractor = extractor
+        self._renderer = renderer
+
+    async def get_next_ai_messages(self, chat: mc.Chat) -> mc.AiChat:
+        out = await self._wrapped.get_next_ai_messages(chat)
+
+        for msg in out:
+            if (c := self._extractor.extract_message_content(msg)) is not None:
+                await self._renderer.render_content(c)
+
+        return out
+
+
+class RenderingStreamAiChatGenerator(StreamAiChatGenerator):
+    def __init__(
+            self,
+            *,
+            wrapped: StreamAiChatGenerator,
+            extractor: MessageContentExtractor | None = None,
+            renderer: ContentRendering,
+    ) -> None:
+        super().__init__()
+
+        self._wrapped = wrapped
+        if extractor is None:
+            extractor = MessageContentExtractorImpl()
+        self._extractor = extractor
+        self._renderer = renderer
+
+    async def get_next_ai_messages_streamed(
+            self,
+            chat: mc.Chat,
+            delta_callback: ta.Callable[[mc.AiChoiceDelta], ta.Awaitable[None]] | None = None,
+    ) -> mc.AiChat:
+        async def inner(delta: mc.AiChoiceDelta) -> None:
+            # FIXME: render lol
+
+            if delta_callback is not None:
+                await delta_callback(delta)
+
+        return await self._wrapped.get_next_ai_messages_streamed(chat, delta_callback=inner)

ommlds/cli/sessions/chat2/chat/ai/services.py

@@ -0,0 +1,70 @@
+import typing as ta
+
+from omlish import check
+
+from ...... import minichain as mc
+from ...backends.types import ChatChoicesServiceBackendProvider
+from .types import AiChatGenerator
+from .types import StreamAiChatGenerator
+
+
+##
+
+
+ChatChoicesServiceOptions = ta.NewType('ChatChoicesServiceOptions', ta.Sequence[mc.ChatChoicesOptions])
+
+
+##
+
+
+class ChatChoicesServiceAiChatGenerator(AiChatGenerator):
+    def __init__(
+            self,
+            service_provider: ChatChoicesServiceBackendProvider,
+            *,
+            options: ChatChoicesServiceOptions | None = None,
+    ) -> None:
+        super().__init__()
+
+        self._service_provider = service_provider
+        self._options = options
+
+    async def get_next_ai_messages(self, chat: mc.Chat) -> mc.AiChat:
+        async with self._service_provider.provide_backend() as service:
+            resp = await service.invoke(mc.ChatChoicesRequest(chat, self._options or []))
+
+        return check.single(resp.v).ms
+
+
+class ChatChoicesStreamServiceStreamAiChatGenerator(StreamAiChatGenerator):
+    def __init__(
+            self,
+            service: mc.ChatChoicesStreamService,
+            *,
+            options: ChatChoicesServiceOptions | None = None,
+    ) -> None:
+        super().__init__()
+
+        self._service = service
+        self._options = options
+
+    async def get_next_ai_messages_streamed(
+            self,
+            chat: mc.Chat,
+            delta_callback: ta.Callable[[mc.AiChoiceDelta], ta.Awaitable[None]] | None = None,
+    ) -> mc.AiChat:
+        lst: list[str] = []
+
+        async with (await self._service.invoke(mc.ChatChoicesStreamRequest(chat, self._options or []))).v as st_resp:
+            async for o in st_resp:
+                choice = check.single(o.choices)
+
+                for delta in choice.deltas:
+                    if delta_callback is not None:
+                        await delta_callback(delta)
+
+                c = check.isinstance(delta, mc.ContentAiChoiceDelta).c  # noqa
+                if c is not None:
+                    lst.append(check.isinstance(c, str))
+
+        return [mc.AiMessage(''.join(lst))]

ommlds/cli/sessions/chat2/chat/ai/types.py

@@ -0,0 +1,28 @@
+import abc
+import typing as ta
+
+from omlish import lang
+
+from ...... import minichain as mc
+
+
+##
+
+
+class AiChatGenerator(lang.Abstract):
+    @abc.abstractmethod
+    def get_next_ai_messages(self, chat: mc.Chat) -> ta.Awaitable[mc.AiChat]:
+        raise NotImplementedError
+
+
+class StreamAiChatGenerator(AiChatGenerator, lang.Abstract):
+    def get_next_ai_messages(self, chat: mc.Chat) -> ta.Awaitable[mc.AiChat]:
+        return self.get_next_ai_messages_streamed(chat)
+
+    @abc.abstractmethod
+    def get_next_ai_messages_streamed(
+            self,
+            chat: mc.Chat,
+            delta_callback: ta.Callable[[mc.AiChoiceDelta], ta.Awaitable[None]] | None = None,
+    ) -> ta.Awaitable[mc.AiChat]:
+        raise NotImplementedError

ommlds/cli/sessions/chat2/chat/state/inmemory.py

@@ -0,0 +1,34 @@
+import dataclasses as dc
+
+from omlish import lang
+
+from ...... import minichain as mc
+from .types import ChatState
+from .types import ChatStateManager
+
+
+##
+
+
+class InMemoryChatStateManager(ChatStateManager):
+    def __init__(self, initial_state: ChatState | None = None) -> None:
+        super().__init__()
+
+        if initial_state is None:
+            initial_state = ChatState()
+        self._state = initial_state
+
+    async def get_state(self) -> ChatState:
+        return self._state
+
+    async def clear_state(self) -> ChatState:
+        self._state = ChatState()
+        return self._state
+
+    async def extend_chat(self, chat_additions: mc.Chat) -> ChatState:
+        self._state = dc.replace(
+            self._state,
+            chat=[*self._state.chat, *chat_additions],
+            updated_at=lang.utcnow(),
+        )
+        return self._state

ommlds/cli/sessions/chat2/chat/state/storage.py

@@ -0,0 +1,53 @@
+import dataclasses as dc
+
+from omlish import check
+from omlish import lang
+
+from ...... import minichain as mc
+from .....state import StateStorage
+from .types import ChatState
+from .types import ChatStateManager
+
+
+##
+
+
+class StateStorageChatStateManager(ChatStateManager):
+    def __init__(
+            self,
+            *,
+            storage: StateStorage,
+            key: str = 'chat',
+    ) -> None:
+        super().__init__()
+
+        self._storage = storage
+        self._key = check.non_empty_str(key)
+
+        self._state: ChatState | None = None
+
+    async def get_state(self) -> ChatState:
+        if self._state is not None:
+            return self._state
+        state: ChatState | None = await self._storage.load_state(self._key, ChatState)
+        if state is None:
+            state = ChatState()
+        self._state = state
+        return state
+
+    async def clear_state(self) -> ChatState:
+        state = ChatState()
+        await self._storage.save_state(self._key, state, ChatState)
+        self._state = state
+        return state
+
+    async def extend_chat(self, chat_additions: mc.Chat) -> ChatState:
+        state = await self.get_state()
+        state = dc.replace(
+            state,
+            chat=[*state.chat, *chat_additions],
+            updated_at=lang.utcnow(),
+        )
+        await self._storage.save_state(self._key, state, ChatState)
+        self._state = state
+        return state

ommlds/cli/sessions/chat2/chat/state/types.py

@@ -0,0 +1,38 @@
+import abc
+import dataclasses as dc
+import datetime
+import typing as ta
+
+from omlish import lang
+
+from ...... import minichain as mc
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class ChatState:
+    name: str | None = None
+
+    created_at: datetime.datetime = dc.field(default_factory=lang.utcnow)
+    updated_at: datetime.datetime = dc.field(default_factory=lang.utcnow)
+
+    chat: mc.Chat = ()
+
+
+##
+
+
+class ChatStateManager(lang.Abstract):
+    @abc.abstractmethod
+    def get_state(self) -> ta.Awaitable[ChatState]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def clear_state(self) -> ta.Awaitable[ChatState]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def extend_chat(self, chat_additions: mc.Chat) -> ta.Awaitable[ChatState]:
+        raise NotImplementedError

ommlds/cli/sessions/chat2/chat/user/interactive.py

@@ -0,0 +1,29 @@
+import functools
+import typing as ta
+
+from omlish import lang
+
+from ...... import minichain as mc
+from .types import UserChatInput
+
+
+##
+
+
+class InteractiveUserChatInput(UserChatInput):
+    def __init__(
+            self,
+            string_input: ta.Callable[[], ta.Awaitable[str]] | None = None,
+    ) -> None:
+        super().__init__()
+
+        if string_input is None:
+            string_input = lang.as_async(functools.partial(input, '> '))
+        self._string_input = string_input
+
+    async def get_next_user_messages(self) -> mc.UserChat:
+        try:
+            s = await self._string_input()
+        except EOFError:
+            return []
+        return [mc.UserMessage(s)]