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

ommlds/cli/main.py

@@ -1,191 +1,331 @@
 """
-See:
- - https://github.com/simonw/llm
- - https://github.com/TheR1D/shell_gpt
- - https://github.com/paul-gauthier/aider
+TODO:
+ - bootstrap lol
 """
-import argparse
+import abc
 import functools
-import os.path
-import sys
 import typing as ta
 
 import anyio
 
-from omdev.home.secrets import load_secrets
 from omlish import check
+from omlish import dataclasses as dc
 from omlish import inject as inj
 from omlish import lang
-from omlish.diag import pycharm
+from omlish.argparse import all as ap
 from omlish.logs import all as logs
-from omlish.subprocesses.editor import edit_text_with_user_editor
-from omlish.subprocesses.sync import subprocesses
 
-from .. import minichain as mc
 from .inject import bind_main
+from .secrets import install_secrets
 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.completion.completion import CompletionSession
-from .sessions.embedding.embedding import EmbeddingSession
-from .tools.config import ToolsConfig
+from .sessions.chat.configs import ChatConfig
+from .sessions.completion.configs import CompletionConfig
+from .sessions.embedding.configs import EmbeddingConfig
 
 
-if ta.TYPE_CHECKING:
-    import PIL.Image as pimg  # noqa
-else:
-    pimg = lang.proxy_import('PIL.Image')
+##
 
 
-##
+MAIN_EXTRA_ARGS: ta.Sequence[ap.Arg] = [
+    ap.arg('-v', '--verbose', action='store_true'),
+]
 
 
-async def _a_main(args: ta.Any = None) -> None:
-    logs.configure_standard_logging('INFO')
+def _process_main_extra_args(args: ap.Namespace) -> None:
+    if args.verbose:
+        logs.configure_standard_logging('DEBUG')
+    else:
+        logs.configure_standard_logging('INFO')
+        logs.silence_noisy_loggers()
 
-    #
 
-    parser = argparse.ArgumentParser()
-    parser.add_argument('prompt', nargs='*')
+##
 
-    parser.add_argument('-b', '--backend', default='openai')
 
-    parser.add_argument('-m', '--model-name')
+class Profile(lang.Abstract):
+    @abc.abstractmethod
+    def run(self, argv: ta.Sequence[str]) -> ta.Awaitable[None]:
+        raise NotImplementedError
 
-    parser.add_argument('-C', '--completion', action='store_true')
 
-    parser.add_argument('-n', '--new', action='store_true')
+##
 
-    parser.add_argument('-e', '--editor', action='store_true')
-    parser.add_argument('-i', '--interactive', action='store_true')
-    parser.add_argument('-c', '--code', action='store_true')
-    parser.add_argument('-s', '--stream', action='store_true')
-    parser.add_argument('-M', '--markdown', action='store_true')
 
-    parser.add_argument('-E', '--embed', action='store_true')
-    parser.add_argument('-j', '--image', action='store_true')
+# class ChatAspect(lang.Abstract):
+#     def get_parser_args(self) -> ta.Sequence[ap.Arg]: ...
+#     def set_args(self, args: ap.Namespace) -> None: ...
+#     def configure(self, cfg: ChatConfig) -> ChatConfig: ...
 
-    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')
-    parser.add_argument('--enable-test-weather-tool', action='store_true')
-    parser.add_argument('--dangerous-no-tool-confirmation', action='store_true')
 
-    args = parser.parse_args(args)
+class ChatProfile(Profile):
+    _args: ap.Namespace
 
     #
 
-    content: mc.Content | None
-
-    if args.image:
-        content = mc.ImageContent(pimg.open(check.non_empty_str(check.single(args.prompt))))
+    BACKEND_ARGS: ta.ClassVar[ta.Sequence[ap.Arg]] = [
+        ap.arg('-b', '--backend', group='backend'),
+    ]
+
+    def configure_backend(self, cfg: ChatConfig) -> ChatConfig:
+        return dc.replace(
+            cfg,
+            backend=dc.replace(
+                cfg.backend,
+                backend=self._args.backend,
+            ),
+        )
 
-    elif args.editor:
-        check.arg(not args.prompt)
-        if (ec := edit_text_with_user_editor('', subprocesses)) is None:
-            return
-        content = ec
+    #
 
-    elif args.interactive:
-        if args.prompt:
-            raise ValueError('Must not provide prompt')
+    INPUT_ARGS: ta.ClassVar[ta.Sequence[ap.Arg]] = [
+        ap.arg('message', nargs='*', group='input'),
+        ap.arg('-i', '--interactive', action='store_true', group='input'),
+        ap.arg('-e', '--editor', action='store_true', group='input'),
+    ]
+
+    def configure_input(self, cfg: ChatConfig) -> ChatConfig:
+        if self._args.editor:
+            check.arg(not self._args.interactive)
+            check.arg(not self._args.message)
+            raise NotImplementedError
+
+        elif self._args.interactive:
+            check.arg(not self._args.message)
+            return dc.replace(
+                cfg,
+                user=dc.replace(
+                    cfg.user,
+                    interactive=True,
+                ),
+            )
+
+        elif self._args.message:
+            # TODO: '-' -> stdin
+            return dc.replace(
+                cfg,
+                user=dc.replace(
+                    cfg.user,
+                    initial_user_content=' '.join(self._args.message),
+                ),
+            )
 
-    elif args.code:
-        if args.prompt:
-            content = ' '.join(args.prompt)
         else:
-            content = None
+            raise ValueError('Must specify input')
 
-    elif not args.prompt:
-        raise ValueError('Must provide prompt')
+    #
 
-    else:
-        prompt = ' '.join(args.prompt)
+    STATE_ARGS: ta.ClassVar[ta.Sequence[ap.Arg]] = [
+        ap.arg('-n', '--new', action='store_true', group='state'),
+        ap.arg('--ephemeral', action='store_true', group='state'),
+    ]
+
+    def configure_state(self, cfg: ChatConfig) -> ChatConfig:
+        return dc.replace(
+            cfg,
+            state=dc.replace(
+                cfg.state,
+                state='ephemeral' if self._args.ephemeral else 'new' if self._args.new else 'continue',
+            ),
+        )
 
-        # FIXME: ptk / maysync
-        if not sys.stdin.isatty() and not pycharm.is_pycharm_hosted():
-            stdin_data = sys.stdin.read()
-            prompt = '\n'.join([prompt, stdin_data])
+    #
 
-        content = prompt
+    OUTPUT_ARGS: ta.ClassVar[ta.Sequence[ap.Arg]] = [
+        ap.arg('-s', '--stream', action='store_true', group='output'),
+        ap.arg('-M', '--markdown', action='store_true', group='output'),
+    ]
+
+    def configure_output(self, cfg: ChatConfig) -> ChatConfig:
+        return dc.replace(
+            cfg,
+            ai=dc.replace(
+                cfg.ai,
+                stream=bool(self._args.stream),
+            ),
+            rendering=dc.replace(
+                cfg.rendering,
+                markdown=bool(self._args.markdown),
+            ),
+        )
 
     #
 
-    # FIXME: lol garbage
-    for key in [
-        'OPENAI_API_KEY',
-        'HUGGINGFACE_TOKEN',
-        'TAVILY_API_KEY',
-        'ANTHROPIC_API_KEY',
-        'MISTRAL_API_KEY',
-        'GEMINI_API_KEY',
-    ]:
-        if (sec := load_secrets().try_get(key.lower())) is not None:
-            os.environ[key] = sec.reveal()
+    TOOLS_ARGS: ta.ClassVar[ta.Sequence[ap.Arg]] = [
+        ap.arg('--enable-fs-tools', action='store_true', group='tools'),
+        ap.arg('--enable-todo-tools', action='store_true', group='tools'),
+        # ap.arg('--enable-unsafe-tools-do-not-use-lol', action='store_true', group='tools'),
+        ap.arg('--enable-test-weather-tool', action='store_true', group='tools'),
+    ]
+
+    def configure_tools(self, cfg: ChatConfig) -> ChatConfig:
+        return dc.replace(
+            cfg,
+            ai=dc.replace(
+                cfg.ai,
+                enable_tools=(
+                    self._args.enable_fs_tools or
+                    self._args.enable_todo_tools or
+                    # self._args.enable_unsafe_tools_do_not_use_lol or
+                    self._args.enable_test_weather_tool or
+                    self._args.code
+                ),
+            ),
+            tools=dc.replace(
+                cfg.tools,
+                enabled_tools={  # noqa
+                    *(cfg.tools.enabled_tools or []),
+                    *(['fs'] if self._args.enable_fs_tools else []),
+                    *(['todo'] if self._args.enable_todo_tools else []),
+                    *(['weather'] if self._args.enable_test_weather_tool else []),
+                },
+            ),
+        )
 
     #
 
-    session_cfg: Session.Config
+    CODE_CONFIG: ta.ClassVar[ta.Sequence[ap.Arg]] = [
+        ap.arg('-c', '--code', action='store_true', group='code'),
+    ]
 
-    if args.interactive:
-        session_cfg = InteractiveChatSession.Config(
-            backend=args.backend,
-            model_name=args.model_name,
-            new=bool(args.new),
-            dangerous_no_tool_confirmation=bool(args.dangerous_no_tool_confirmation),
-        )
+    def configure_code(self, cfg: ChatConfig) -> ChatConfig:
+        if not self._args.code:
+            return cfg
 
-    elif args.code:
-        session_cfg = CodeChatSession.Config(
-            backend=args.backend,
-            model_name=args.model_name,
-            new=bool(args.new),
-            dangerous_no_tool_confirmation=bool(args.dangerous_no_tool_confirmation),
-            initial_message=content,  # noqa
-            markdown=bool(args.markdown),
+        cfg = dc.replace(
+            cfg,
+            ai=dc.replace(
+                cfg.ai,
+                enable_tools=True,
+            ),
         )
 
-    elif args.embed:
-        session_cfg = EmbeddingSession.Config(
-            check.not_none(content),  # noqa
-            backend=args.backend,
-        )
+        if self._args.new or self._args.ephemeral:
+            from ..minichain.lib.code.prompts import CODE_AGENT_SYSTEM_PROMPT
+            system_content = CODE_AGENT_SYSTEM_PROMPT
+
+            cfg = dc.replace(
+                cfg,
+                user=dc.replace(
+                    cfg.user,
+                    initial_system_content=system_content,
+                ),
+            )
+
+        return cfg
+
+    #
+
+    async def run(self, argv: ta.Sequence[str]) -> None:
+        parser = ap.ArgumentParser()
+
+        for grp_name, grp_args in [
+            ('backend', self.BACKEND_ARGS),
+            ('input', self.INPUT_ARGS),
+            ('state', self.STATE_ARGS),
+            ('output', self.OUTPUT_ARGS),
+            ('tools', self.TOOLS_ARGS),
+            ('code', self.CODE_CONFIG),
+        ]:
+            grp = parser.add_argument_group(grp_name)
+            for a in grp_args:
+                grp.add_argument(*a.args, **a.kwargs)
+
+        self._args = parser.parse_args(argv)
+
+        cfg = ChatConfig()
+        cfg = self.configure_backend(cfg)
+        cfg = self.configure_input(cfg)
+        cfg = self.configure_state(cfg)
+        cfg = self.configure_output(cfg)
+        cfg = self.configure_tools(cfg)
+        cfg = self.configure_code(cfg)
+
+        with inj.create_managed_injector(bind_main(
+                session_cfg=cfg,
+        )) as injector:
+            await injector[Session].run()
+
+
+##
 
-    elif args.completion:
-        session_cfg = CompletionSession.Config(
-            check.not_none(content),  # noqa
+
+class CompletionProfile(Profile):
+    async def run(self, argv: ta.Sequence[str]) -> None:
+        parser = ap.ArgumentParser()
+        parser.add_argument('prompt', nargs='*')
+        parser.add_argument('-b', '--backend', default='openai')
+        args = parser.parse_args(argv)
+
+        content = ' '.join(args.prompt)
+
+        cfg = CompletionConfig(
+            check.non_empty_str(content),
             backend=args.backend,
         )
 
-    else:
-        session_cfg = PromptChatSession.Config(
-            check.not_none(content),  # noqa
+        with inj.create_managed_injector(bind_main(
+                session_cfg=cfg,
+        )) as injector:
+            await injector[Session].run()
+
+
+##
+
+
+class EmbedProfile(Profile):
+    async def run(self, argv: ta.Sequence[str]) -> None:
+        parser = ap.ArgumentParser()
+        parser.add_argument('prompt', nargs='*')
+        parser.add_argument('-b', '--backend', default='openai')
+        args = parser.parse_args(argv)
+
+        content = ' '.join(args.prompt)
+
+        cfg = EmbeddingConfig(
+            check.non_empty_str(content),
             backend=args.backend,
-            model_name=args.model_name,
-            new=bool(args.new),
-            stream=bool(args.stream),
-            markdown=bool(args.markdown),
-            dangerous_no_tool_confirmation=bool(args.dangerous_no_tool_confirmation),
         )
 
-    #
+        with inj.create_managed_injector(bind_main(
+                session_cfg=cfg,
+        )) as injector:
+            await injector[Session].run()
 
-    tools_config = ToolsConfig(
-        enable_fs_tools=args.enable_fs_tools or args.code,
-        enable_todo_tools=args.enable_todo_tools or args.code,
-        enable_unsafe_tools_do_not_use_lol=args.enable_unsafe_tools_do_not_use_lol,
-        enable_test_weather_tool=args.enable_test_weather_tool,
-    )
 
-    #
+##
+
+
+PROFILE_TYPES: ta.Mapping[str, type[Profile]] = {
+    'chat': ChatProfile,
+    'complete': CompletionProfile,
+    'embed': EmbedProfile,
+}
+
+
+##
+
+
+MAIN_PROFILE_ARGS: ta.Sequence[ap.Arg] = [
+    ap.arg('-p', '--profile', default='chat'),
+    ap.arg('args', nargs=ap.REMAINDER),
+]
+
+
+async def _a_main(argv: ta.Any = None) -> None:
+    parser = ap.ArgumentParser()
+
+    for a in [*MAIN_PROFILE_ARGS, *MAIN_EXTRA_ARGS]:
+        parser.add_argument(*a.args, **a.kwargs)
+
+    args, unk_args = parser.parse_known_args(argv)
+
+    _process_main_extra_args(args)
+
+    install_secrets()
 
-    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)),
-    )) as injector:
-        await injector[Session].run()
+    profile_cls = PROFILE_TYPES[args.profile]
+    profile = profile_cls()
+    await profile.run([*unk_args, *args.args])
 
 
 def _main(args: ta.Any = None) -> None:
@@ -194,7 +334,6 @@ def _main(args: ta.Any = None) -> None:
             _a_main,
             args,
         ),
-        backend='asyncio',
     )  # noqa
 
 

ommlds/cli/rendering/configs.py

@@ -0,0 +1,9 @@
+from omlish import dataclasses as dc
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class RenderingConfig:
+    markdown: bool = False

ommlds/cli/rendering/inject.py

@@ -0,0 +1,31 @@
+from omlish import inject as inj
+from omlish import lang
+
+from .configs import RenderingConfig
+
+
+with lang.auto_proxy_import(globals()):
+    from . import markdown as _markdown
+    from . import raw as _raw
+    from . import types as _types
+
+
+##
+
+
+def bind_rendering(cfg: RenderingConfig = RenderingConfig()) -> inj.Elements:
+    els: list[inj.Elemental] = []
+
+    if cfg.markdown:
+        els.extend([
+            inj.bind(_types.ContentRendering, to_ctor=_markdown.MarkdownContentRendering, singleton=True),
+            inj.bind(_types.StreamContentRendering, to_ctor=_markdown.MarkdownStreamContentRendering, singleton=True),
+        ])
+
+    else:
+        els.extend([
+            inj.bind(_types.ContentRendering, to_ctor=_raw.RawContentRendering, singleton=True),
+            inj.bind(_types.StreamContentRendering, to_ctor=_raw.RawStreamContentRendering, singleton=True),
+        ])
+
+    return inj.as_elements(*els)

ommlds/cli/rendering/markdown.py

@@ -0,0 +1,52 @@
+import typing as ta
+
+from omdev.tui import rich
+from omlish import lang
+
+from ... import minichain as mc
+from ..content.strings import ContentStringifier
+from ..content.strings import HasContentStringifier
+from .types import ContentRendering
+from .types import StreamContentRendering
+
+
+##
+
+
+class MarkdownContentRendering(ContentRendering, HasContentStringifier):
+    def __init__(
+            self,
+            *,
+            content_stringifier: ContentStringifier | None = None,
+    ) -> None:
+        super().__init__(content_stringifier=content_stringifier)
+
+    async def render_content(self, content: 'mc.Content') -> None:
+        if (s := self._content_stringifier.stringify_content(content)) is not None and (s := s.strip()):
+            rich.Console().print(rich.Markdown(s))
+
+
+class MarkdownStreamContentRendering(StreamContentRendering, HasContentStringifier):
+    def __init__(
+            self,
+            *,
+            content_stringifier: ContentStringifier | None = None,
+    ) -> None:
+        super().__init__(content_stringifier=content_stringifier)
+
+    @ta.final
+    class _ContextInstance(ContentRendering, lang.AsyncExitStacked):
+        def __init__(self, owner: 'MarkdownStreamContentRendering') -> None:
+            self._owner = owner
+
+        _ir: rich.MarkdownLiveStream
+
+        async def _async_enter_contexts(self) -> None:
+            self._ir = self._enter_context(rich.IncrementalMarkdownLiveStream())
+
+        async def render_content(self, content: 'mc.Content') -> None:
+            if (s := self._owner._content_stringifier.stringify_content(content)) is not None:  # noqa: SLF001
+                self._ir.feed(s)
+
+    def create_context(self) -> ta.AsyncContextManager[ContentRendering]:
+        return MarkdownStreamContentRendering._ContextInstance(self)

ommlds/cli/rendering/raw.py

@@ -0,0 +1,73 @@
+import typing as ta
+
+from omlish import lang
+
+from ... import minichain as mc
+from ..content.strings import ContentStringifier
+from ..content.strings import HasContentStringifier
+from .types import ContentRendering
+from .types import StreamContentRendering
+
+
+##
+
+
+class RawContentRendering(ContentRendering, HasContentStringifier):
+    def __init__(
+            self,
+            *,
+            printer: ta.Callable[[str], ta.Awaitable[None]] | None = None,
+            content_stringifier: ContentStringifier | None = None,
+    ) -> None:
+        super().__init__(content_stringifier=content_stringifier)
+
+        if printer is None:
+            printer = lang.as_async(print)
+        self._printer = printer
+
+    async def render_content(self, content: 'mc.Content') -> None:
+        if (s := self._content_stringifier.stringify_content(content)) is not None:
+            await self._printer(s)
+
+
+class RawStreamContentRendering(StreamContentRendering, HasContentStringifier):
+    class Output(ta.Protocol):
+        def write(self, s: str) -> ta.Awaitable[None]: ...
+        def flush(self) -> ta.Awaitable[None]: ...
+
+    class PrintOutput:
+        async def write(self, s: str) -> None:
+            print(s, end='', flush=True)
+
+        async def flush(self) -> None:
+            print(flush=True)
+
+    def __init__(
+            self,
+            *,
+            output: Output | None = None,
+            content_stringifier: ContentStringifier | None = None,
+    ) -> None:
+        super().__init__(content_stringifier=content_stringifier)
+
+        if output is None:
+            output = RawStreamContentRendering.PrintOutput()
+        self._output = output
+
+    @ta.final
+    class _ContextInstance(ContentRendering, ta.AsyncContextManager):
+        def __init__(self, owner: 'RawStreamContentRendering') -> None:
+            self._owner = owner
+
+        async def __aenter__(self) -> ta.Self:
+            return self
+
+        async def __aexit__(self, *exc_info) -> None:
+            await self._owner._output.flush()  # noqa: SLF001
+
+        async def render_content(self, content: 'mc.Content') -> None:
+            if (s := self._owner._content_stringifier.stringify_content(content)) is not None:  # noqa: SLF001
+                await self._owner._output.write(s)  # noqa: SLF001
+
+    def create_context(self) -> ta.AsyncContextManager[ContentRendering]:
+        return RawStreamContentRendering._ContextInstance(self)

ommlds/cli/rendering/types.py

@@ -0,0 +1,21 @@
+import abc
+import typing as ta
+
+from omlish import lang
+
+from ... import minichain as mc
+
+
+##
+
+
+class ContentRendering(lang.Abstract):
+    @abc.abstractmethod
+    def render_content(self, content: 'mc.Content') -> ta.Awaitable[None]:
+        raise NotImplementedError
+
+
+class StreamContentRendering(lang.Abstract):
+    @abc.abstractmethod
+    def create_context(self) -> ta.AsyncContextManager[ContentRendering]:
+        raise NotImplementedError

ommlds/cli/secrets.py

@@ -0,0 +1,21 @@
+import os
+
+from omdev.home.secrets import load_secrets
+
+
+##
+
+
+def install_secrets() -> None:
+    # FIXME: lol garbage
+    for key in [
+        'ANTHROPIC_API_KEY',
+        'GEMINI_API_KEY',
+        'GROQ_API_KEY',
+        'HUGGINGFACE_TOKEN',
+        'MISTRAL_API_KEY',
+        'OPENAI_API_KEY',
+        'TAVILY_API_KEY',
+    ]:
+        if (sec := load_secrets().try_get(key.lower())) is not None:
+            os.environ[key] = sec.reveal()