ommlds 0.0.0.dev480__py3-none-any.whl → 0.0.0.dev503__py3-none-any.whl

ommlds/minichain/wrappers/stream.py

@@ -0,0 +1,57 @@
+import typing as ta
+
+from omlish import lang
+
+from ..services.requests import Request
+from ..services.services import Service
+from ..stream.services import StreamOptions
+from ..stream.services import StreamResponse
+from ..types import Output
+from .services import WrappedOptionT
+from .services import WrappedOutputT
+from .services import WrappedRequestV
+from .services import WrappedResponseV
+
+
+WrappedStreamOutputT = ta.TypeVar('WrappedStreamOutputT', bound=Output)
+
+
+WrappedStreamOptions: ta.TypeAlias = WrappedOptionT | StreamOptions
+
+WrappedStreamRequest: ta.TypeAlias = Request[
+    WrappedRequestV,
+    WrappedStreamOptions,
+]
+
+WrappedStreamResponse: ta.TypeAlias = StreamResponse[
+    WrappedResponseV,
+    WrappedOutputT,
+    WrappedStreamOutputT,
+]
+
+WrappedStreamService: ta.TypeAlias = Service[
+    WrappedStreamRequest,
+    WrappedStreamResponse,
+]
+
+
+##
+
+
+class WrapperStreamService(
+    lang.Abstract,
+    ta.Generic[
+        WrappedRequestV,
+        WrappedOptionT,
+        WrappedResponseV,
+        WrappedOutputT,
+        WrappedStreamOutputT,
+    ],
+):
+    def __init__(
+            self,
+            service: WrappedStreamService,
+    ) -> None:
+        super().__init__()
+
+        self._service = service

ommlds/nanochat/rustbpe/README.md

@@ -0,0 +1,9 @@
+# https://github.com/karpathy/nanochat/tree/9467d83cf23dcc9a9b4ca6e35103142f48a55b27
+
+---
+
+# rustbpe
+
+> The missing tiktoken training code
+
+A very lightweight Rust library for training a GPT tokenizer. The issue is that the inference library [tiktoken](https://github.com/openai/tiktoken) is great, but only does inference. Separately, the huggingface [tokenizers](https://github.com/huggingface/tokenizers) library does training, but it is rather bloated and really hard to navigate because it has to support all the different historical baggage of how people dealt with tokenizers over the years. More recently, I also wrote the [minbpe](https://github.com/karpathy/minbpe) library which does both training and inference, but only in inefficient Python. Basically what I really want is a non-fancy, super simple, but still relatively efficient training code for GPT tokenizer (more efficient than minbpe, much cleaner/simpler than tokenizers), and then export the trained vocab for inference with tiktoken. Does that make sense? So here we are. There are more opportunities for optimization here, I just stopped a bit early because unlike minbpe before it, rustbpe is now simple and fast enough, and not a significant bottleneck for nanochat.

ommlds/nanochat/tokenizers.py

@@ -18,7 +18,10 @@ from omlish import lang
 
 with lang.auto_proxy_import(globals()):
     import tiktoken
-    import tokenizers
+    import tokenizers.decoders
+    import tokenizers.models
+    import tokenizers.pre_tokenizers
+    import tokenizers.trainers
 
 
 rustbpe: ta.Any = lang.proxy_import('.rustbpe', __package__)
@@ -27,7 +30,7 @@ rustbpe: ta.Any = lang.proxy_import('.rustbpe', __package__)
 ##
 
 
-SPECIAL_TOKENS = [
+SPECIAL_TOKENS: ta.Sequence[str] = [
     # every document begins with the Beginning of Sequence (BOS) token that delimits documents
     '<|bos|>',
     # tokens below are only used during finetuning to render Conversations into token ids
@@ -45,10 +48,18 @@ SPECIAL_TOKENS = [
 # NOTE: this split pattern deviates from GPT-4 in that we use \p{N}{1,2} instead of \p{N}{1,3}
 # I did this because I didn't want to "waste" too many tokens on numbers for smaller vocab sizes.
 # I haven't validated that this is actually a good idea, TODO.
-SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""  # noqa
+SPLIT_PATTERN = (
+    r"'(?i:[sdmt]|ll|ve|re)|"
+    r"[^\r\n\p{L}\p{N}]?+\p{L}+|"
+    r"\p{N}{1,2}|"
+    r" ?[^\s\p{L}\p{N}]++[\r\n]*|"
+    r"\s*[\r\n]|"
+    r"\s+(?!\S)|"
+    r"\s+"
+)
 
 
-# -----------------------------------------------------------------------------
+##
 # Generic GPT-4-style tokenizer based on HuggingFace Tokenizer
 
 
@@ -87,22 +98,28 @@ class HuggingFaceTokenizer:
             unk_token=None,
             fuse_unk=False,
         ))
+
         # Normalizer: None
         tokenizer.normalizer = None
+
         # Pre-tokenizer: GPT-4 style
         # the regex pattern used by GPT-4 to split text into groups before BPE
         # NOTE: The pattern was changed from \p{N}{1,3} to \p{N}{1,2} because I suspect it is harmful to
         # very small models and smaller vocab sizes, because it is a little bit wasteful in the token space.
         # (but I haven't validated this! TODO)
         gpt4_split_regex = tokenizers.Regex(split_pattern)  # huggingface demands that you wrap it in Regex!!
+
         tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.Sequence([
             tokenizers.pre_tokenizers.Split(pattern=gpt4_split_regex, behavior='isolated', invert=False),
             tokenizers.pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False),
         ])
+
         # Decoder: ByteLevel (it pairs together with the ByteLevel pre-tokenizer)
         tokenizer.decoder = tokenizers.decoders.ByteLevel()
+
         # Post-processor: None
         tokenizer.post_processor = None
+
         # Trainer: BPE
         trainer = tokenizers.trainers.BpeTrainer(
             vocab_size=vocab_size,
@@ -111,8 +128,10 @@ class HuggingFaceTokenizer:
             initial_alphabet=tokenizers.pre_tokenizers.ByteLevel.alphabet(),
             special_tokens=special_tokens,
         )
+
         # Kick off the training
         tokenizer.train_from_iterator(text_iterator, trainer)
+
         return cls(tokenizer)
 
     def encode_ordinary(self, text):
@@ -174,7 +193,7 @@ class HuggingFaceTokenizer:
         print(f'Saved tokenizer to {tokenizer_path}')
 
 
-# -----------------------------------------------------------------------------
+##
 # Tokenizer based on rustbpe + tiktoken combo
 
 
@@ -255,6 +274,7 @@ class RustBPETokenizer:
                 ids.insert(0, prepend_id)  # TODO: slightly inefficient here? :( hmm
             if append is not None:
                 ids.append(append_id)
+
         elif isinstance(text, list):
             ids = self.enc.encode_ordinary_batch(text, num_threads=num_threads)
             if prepend is not None:
@@ -263,6 +283,7 @@ class RustBPETokenizer:
             if append is not None:
                 for ids_row in ids:
                     ids_row.append(append_id)
+
         else:
             raise ValueError(f'Invalid input type: {type(text)}')  # noqa
 
@@ -285,6 +306,7 @@ class RustBPETokenizer:
     def render_conversation(self, conversation, max_tokens=2048):
         """
         Tokenize a single Chat conversation (which we call a "doc" or "document" here).
+
         Returns:
         - ids: list[int] is a list of token ids of this rendered conversation
         - mask: list[int] of same length, mask = 1 for tokens that the Assistant is expected to train on.
@@ -324,7 +346,10 @@ class RustBPETokenizer:
         for i, message in enumerate(messages):
             # some sanity checking here around assumptions, to prevent footguns
             must_be_from = 'user' if i % 2 == 0 else 'assistant'
-            check.state(message['role'] == must_be_from, f"Message {i} is from {message['role']} but should be from {must_be_from}")  # noqa
+            check.state(
+                message['role'] == must_be_from,
+                f"Message {i} is from {message['role']} but should be from {must_be_from}",
+            )
 
             # content can be either a simple string or a list of parts (e.g. containing tool calls)
             content = message['content']
@@ -335,33 +360,42 @@ class RustBPETokenizer:
                 add_tokens(user_start, 0)
                 add_tokens(value_ids, 0)
                 add_tokens(user_end, 0)
+
             elif message['role'] == 'assistant':
                 add_tokens(assistant_start, 0)
+
                 if isinstance(content, str):
                     # simple string => simply add the tokens
                     value_ids = self.encode(content)
                     add_tokens(value_ids, 1)
+
                 elif isinstance(content, list):
                     for part in content:
                         value_ids = self.encode(part['text'])
+
                         if part['type'] == 'text':
                             # string part => simply add the tokens
                             add_tokens(value_ids, 1)
+
                         elif part['type'] == 'python':
                             # python tool call => add the tokens inside <|python_start|> and <|python_end|>
                             add_tokens(python_start, 1)
                             add_tokens(value_ids, 1)
                             add_tokens(python_end, 1)
+
                         elif part['type'] == 'python_output':
                             # python output => add the tokens inside <|output_start|> and <|output_end|>
                             # none of these tokens are supervised because the tokens come from Python at test time
                             add_tokens(output_start, 0)
                             add_tokens(value_ids, 0)
                             add_tokens(output_end, 0)
+
                         else:
                             raise ValueError(f"Unknown part type: {part['type']}")
+
                 else:
                     raise ValueError(f'Unknown content type: {type(content)}')
+
                 add_tokens(assistant_end, 1)
 
         # truncate to max_tokens tokens MAX (helps prevent OOMs)

ommlds/specs/mcp/clients.py

@@ -0,0 +1,146 @@
+import contextlib
+import subprocess
+import typing as ta
+
+import anyio.abc
+
+from omlish import check
+from omlish import dataclasses as dc
+from omlish import marshal as msh
+from omlish.asyncs import anyio as aiu
+from omlish.specs import jsonrpc as jr
+
+from . import protocol as pt
+
+
+##
+
+
+class McpServerConnection:
+    def __init__(
+            self,
+            tg: anyio.abc.TaskGroup,
+            stream: anyio.abc.ByteStream,
+            *,
+            default_timeout: float | None = 30.,
+    ) -> None:
+        super().__init__()
+
+        self._conn = jr.Connection(
+            tg,
+            stream,
+            request_handler=self._handle_client_request,
+            notification_handler=self._handle_client_notification,
+            default_timeout=default_timeout,
+        )
+
+    #
+
+    @classmethod
+    def from_process(
+            cls,
+            tg: anyio.abc.TaskGroup,
+            proc: anyio.abc.Process,
+            **kwargs: ta.Any,
+    ) -> 'McpServerConnection':
+        return cls(
+            tg,
+            aiu.StapledByteStream(
+                check.not_none(proc.stdin),
+                check.not_none(proc.stdout),
+            ),
+            **kwargs,
+        )
+
+    @classmethod
+    def open_process(
+            cls,
+            tg: anyio.abc.TaskGroup,
+            cmd: ta.Sequence[str],
+            open_kwargs: ta.Mapping[str, ta.Any] | None = None,
+            **kwargs: ta.Any,
+    ) -> ta.AsyncContextManager[tuple[anyio.abc.Process, 'McpServerConnection']]:
+        @contextlib.asynccontextmanager
+        async def inner():
+            async with await anyio.open_process(
+                    cmd,
+                    stdin=subprocess.PIPE,
+                    stdout=subprocess.PIPE,
+                    **open_kwargs or {},
+            ) as proc:
+                async with cls.from_process(
+                        tg,
+                        proc,
+                        **kwargs,
+                ) as client:
+                    yield (proc, client)
+
+        return inner()
+
+    #
+
+    async def __aenter__(self) -> 'McpServerConnection':
+        await self._conn.__aenter__()
+        return self
+
+    async def __aexit__(self, et, e, tb) -> None:
+        await self._conn.__aexit__(et, e, tb)
+
+    #
+
+    async def _handle_client_request(self, _client: jr.Connection, req: jr.Request) -> None:
+        pass
+
+    async def _handle_client_notification(self, _client: jr.Connection, no: jr.Request) -> None:
+        pass
+
+    #
+
+    async def request(self, req: pt.ClientRequest[pt.ClientResultT]) -> pt.ClientResultT:
+        res_cls = pt.MESSAGE_TYPES_BY_JSON_RPC_METHOD_NAME[pt.ClientResult][req.json_rpc_method_name]  # type: ignore[type-abstract]  # noqa
+        req_mv = msh.marshal(req)
+        res_mv = await self._conn.request(req.json_rpc_method_name, req_mv)  # type: ignore[arg-type]
+        res = msh.unmarshal(res_mv, res_cls)
+        return ta.cast(pt.ClientResultT, res)
+
+    async def notify(self, no: pt.Notification) -> None:
+        no_mv = msh.marshal(no)
+        await self._conn.notify(no.json_rpc_method_name, no_mv)  # type: ignore[arg-type]
+
+    #
+
+    async def yield_cursor_request(
+            self,
+            req: pt.CursorClientRequest[pt.CursorClientResultT],
+    ) -> ta.AsyncGenerator[pt.CursorClientResultT]:
+        check.none(req.cursor)
+
+        cursor: str | None = None
+        while True:
+            res = await self.request(dc.replace(req, cursor=cursor))  # noqa
+            yield res
+
+            if (cursor := res.next_cursor) is None:
+                break
+
+    async def list_cursor_request(
+            self,
+            req: pt.CursorClientRequest[pt.CursorClientResultT],
+    ) -> list[pt.CursorClientResultT]:
+        return [res async for res in self.yield_cursor_request(req)]
+
+    #
+
+    async def list_tools(self) -> list[pt.Tool]:
+        return [
+            tool
+            async for res in self.yield_cursor_request(pt.ListToolsRequest())
+            for tool in res.tools
+        ]
+
+    async def list_prompts(self) -> list[pt.Prompt]:
+        return [
+            prompt
+            async for res in self.yield_cursor_request(pt.ListPromptsRequest())
+            for prompt in res.prompts
+        ]

ommlds/specs/mcp/protocol.py

@@ -1,6 +1,15 @@
 """
 https://modelcontextprotocol.io/specification/2025-06-18
 https://modelcontextprotocol.io/specification/2025-06-18/schema
+
+TODO:
+ - https://modelcontextprotocol.io/specification/2025-06-18/basic/utilities/cancellation
+ - https://modelcontextprotocol.io/specification/2025-06-18/basic/utilities/progress
+ - https://modelcontextprotocol.io/specification/2025-06-18/client/sampling
+ - https://modelcontextprotocol.io/specification/2025-06-18/client/elicitation
+ - https://modelcontextprotocol.io/specification/2025-06-18/server/prompts
+ - https://modelcontextprotocol.io/specification/2025-06-18/server/resources
+ - https://modelcontextprotocol.io/specification/2025-06-18/server/utilities/logging
 """
 import abc
 import typing as ta
@@ -11,6 +20,16 @@ from omlish import lang
 from omlish import marshal as msh
 
 
+ClientRequestT = ta.TypeVar('ClientRequestT', bound='ClientRequest')
+ClientResultT = ta.TypeVar('ClientResultT', bound='ClientResult')
+
+ServerRequestT = ta.TypeVar('ServerRequestT', bound='ServerRequest')
+ServerResultT = ta.TypeVar('ServerResultT', bound='ServerResult')
+
+CursorClientRequestT = ta.TypeVar('CursorClientRequestT', bound='CursorClientRequest')
+CursorClientResultT = ta.TypeVar('CursorClientResultT', bound='CursorClientResult')
+
+
 msh.register_global_module_import('._marshal', __package__)
 
 
@@ -41,27 +60,36 @@ class Message(lang.Sealed, lang.Abstract):
         raise NotImplementedError
 
 
-class ClientRequest(Message, lang.Abstract):
+#
+
+
+class ClientRequest(Message, lang.Abstract, ta.Generic[ClientResultT]):
     pass
 
 
-class ClientResult(Message, lang.Abstract):
+class ClientResult(Message, lang.Abstract, ta.Generic[ClientRequestT]):
     pass
 
 
-class ServerRequest(Message, lang.Abstract):
+#
+
+
+class ServerRequest(Message, lang.Abstract, ta.Generic[ServerResultT]):
     pass
 
 
-class ServerResult(Message, lang.Abstract):
+class ServerResult(Message, lang.Abstract, ta.Generic[ServerRequestT]):
     pass
 
 
+#
+
+
 class Notification(Message, lang.Abstract):
     pass
 
 
-#
+##
 
 
 @dc.dataclass(frozen=True, kw_only=True)
@@ -84,7 +112,7 @@ DEFAULT_PROTOCOL_VERSION = '2025-06-18'
 
 @dc.dataclass(frozen=True, kw_only=True)
 @_set_class_marshal_options
-class ClientCapabilities:
+class ClientCapabilities(lang.Final):
     elicitation: ta.Any | None = None
 
     experimental: ta.Mapping[str, ta.Any] | None = None
@@ -101,7 +129,7 @@ class ClientCapabilities:
 
 @dc.dataclass(frozen=True, kw_only=True)
 @_set_class_marshal_options
-class InitializeRequest(ClientRequest):
+class InitializeRequest(ClientRequest['InitializeResult']):
     json_rpc_method_name: ta.ClassVar[str] = 'initialize'
 
     client_info: Implementation
@@ -143,12 +171,13 @@ class ServerCapabilities:
 
 @dc.dataclass(frozen=True, kw_only=True)
 @_set_class_marshal_options
-class InitializeResult(ClientResult, WithMeta):
+class InitializeResult(ClientResult[InitializeRequest], WithMeta):
     json_rpc_method_name: ta.ClassVar[str] = 'initialize'
 
     server_info: Implementation
     protocol_version: str
     capabilities: ServerCapabilities
+    instructions: str | None = None
 
 
 @dc.dataclass(frozen=True, kw_only=True)
@@ -160,6 +189,19 @@ class InitializedNotification(Notification):
 ##
 
 
+@dc.dataclass(frozen=True, kw_only=True)
+class CursorClientRequest(ClientRequest[CursorClientResultT], lang.Abstract):
+    cursor: str | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class CursorClientResult(ClientResult[CursorClientRequestT], lang.Abstract):
+    next_cursor: str | None = None
+
+
+##
+
+
 @dc.dataclass(frozen=True, kw_only=True)
 @_set_class_marshal_options
 class ToolAnnotations(lang.Final):
@@ -172,7 +214,7 @@ class ToolAnnotations(lang.Final):
 
 @dc.dataclass(frozen=True, kw_only=True)
 @_set_class_marshal_options
-class Tool(lang.Final):
+class Tool(WithMeta, lang.Final):
     name: str
     title: str | None = None
 
@@ -186,19 +228,16 @@ class Tool(lang.Final):
 
 @dc.dataclass(frozen=True, kw_only=True)
 @_set_class_marshal_options
-class ListToolsRequest(ClientRequest):
+class ListToolsRequest(CursorClientRequest['ListToolsResult']):
     json_rpc_method_name: ta.ClassVar[str] = 'tools/list'
 
-    cursor: str | None = None
-
 
 @dc.dataclass(frozen=True, kw_only=True)
 @_set_class_marshal_options
-class ListToolsResult(ClientResult, WithMeta):
+class ListToolsResult(CursorClientResult[ListToolsRequest], WithMeta):
     json_rpc_method_name: ta.ClassVar[str] = 'tools/list'
 
     tools: ta.Sequence[Tool]
-    next_cursor: str | None = None
 
 
 ##
@@ -219,7 +258,7 @@ class TextContentBlock(ContentBlock, lang.Final):
 
 @dc.dataclass(frozen=True, kw_only=True)
 @_set_class_marshal_options
-class CallToolRequest(ClientRequest):
+class CallToolRequest(ClientRequest['CallToolResult']):
     json_rpc_method_name: ta.ClassVar[str] = 'tools/call'
 
     name: str
@@ -228,7 +267,7 @@ class CallToolRequest(ClientRequest):
 
 @dc.dataclass(frozen=True, kw_only=True)
 @_set_class_marshal_options
-class CallToolResult(ClientResult, WithMeta):
+class CallToolResult(ClientResult[CallToolRequest], WithMeta):
     json_rpc_method_name: ta.ClassVar[str] = 'tools/call'
 
     content: ta.Sequence[ContentBlock]
@@ -241,16 +280,82 @@ class CallToolResult(ClientResult, WithMeta):
 
 @dc.dataclass(frozen=True, kw_only=True)
 @_set_class_marshal_options
-class PingRequest(ClientRequest, WithMeta):
+class PingClientRequest(ClientRequest['PingClientResult'], WithMeta):
     json_rpc_method_name: ta.ClassVar[str] = 'ping'
 
 
 @dc.dataclass(frozen=True, kw_only=True)
 @_set_class_marshal_options
-class PingResult(ClientResult):
+class PingClientResult(ClientResult[PingClientRequest]):
     json_rpc_method_name: ta.ClassVar[str] = 'ping'
 
 
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class PingServerRequest(ServerRequest['PingServerResult'], WithMeta):
+    json_rpc_method_name: ta.ClassVar[str] = 'ping'
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class PingServerResult(ServerResult[PingServerRequest]):
+    json_rpc_method_name: ta.ClassVar[str] = 'ping'
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class PromptArgument(lang.Final):
+    name: str
+    title: str | None = None
+
+    description: str | None = None
+
+    required: bool | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class Prompt(WithMeta, lang.Final):
+    name: str
+    title: str | None = None
+
+    description: str | None = None
+
+    arguments: ta.Sequence[PromptArgument] | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class ListPromptsRequest(CursorClientRequest['ListPromptsResult']):
+    json_rpc_method_name: ta.ClassVar[str] = 'prompts/list'
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class ListPromptsResult(CursorClientResult[ListPromptsRequest], WithMeta):
+    json_rpc_method_name: ta.ClassVar[str] = 'prompts/list'
+
+    prompts: ta.Sequence[Prompt]
+
+
+##
+
+
+LoggingLevel: ta.TypeAlias = ta.Literal[
+    'debug',
+    'info',
+    'notice',
+    'warning',
+    'error',
+    'critical',
+    'alert',
+    'emergency',
+]
+
+
 ##