ommlds 0.0.0.dev502__py3-none-any.whl → 0.0.0.dev504__py3-none-any.whl

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

@@ -6,21 +6,18 @@ from omlish import typedvalues as tv
 from omlish.formats import json
 from omlish.http import all as http
 from omlish.http import sse
-from omlish.io.buffers import DelimitingBuffer
 
 from .....backends.groq import protocol as pt
 from .....backends.groq.clients import REQUIRED_HTTP_HEADERS
-from ....chat.choices.services import ChatChoicesOutputs
 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 AiChoicesDeltas
 from ....chat.tools.types import Tool
 from ....configs import Config
-from ....resources import UseResources
+from ....http.stream import BytesHttpStreamResponseBuilder
+from ....http.stream import SimpleSseLinesHttpStreamResponseHandler
 from ....standard import ApiKey
-from ....stream.services import StreamResponseSink
-from ....stream.services import new_stream_response
 from .chat import GroqChatChoicesService
 from .names import MODEL_NAMES
 from .protocol import build_gq_request_messages
@@ -50,6 +47,35 @@ class GroqChatChoicesStreamService:
             self._model_name = cc.pop(GroqChatChoicesService.DEFAULT_MODEL_NAME)
             self._api_key = ApiKey.pop_secret(cc, env='GROQ_API_KEY')
 
+    URL: ta.ClassVar[str] = 'https://api.groq.com/openai/v1/chat/completions'
+
+    def _process_sse(self, so: sse.SseDecoderOutput) -> ta.Sequence[AiChoicesDeltas | None]:
+        if not (isinstance(so, sse.SseEvent) and so.type == b'message'):
+            return []
+
+        ss = so.data.decode('utf-8')
+        if ss == '[DONE]':
+            return [None]
+
+        sj = json.loads(ss)  # ChatCompletionChunk
+
+        check.state(sj['object'] == 'chat.completion.chunk')
+
+        ccc = msh.unmarshal(sj, pt.ChatCompletionChunk)
+
+        # FIXME: stop reason
+        if not ccc.choices:
+            return []
+
+        if any(choice.finish_reason for choice in ccc.choices):
+            check.state(all(choice.finish_reason for choice in ccc.choices))
+            return [None]
+
+        return [AiChoicesDeltas([
+            build_mc_ai_choice_deltas(choice.delta)
+            for choice in ccc.choices
+        ])]
+
     READ_CHUNK_SIZE: ta.ClassVar[int] = -1
 
     async def invoke(self, request: ChatChoicesStreamRequest) -> ChatChoicesStreamResponse:
@@ -69,7 +95,7 @@ class GroqChatChoicesStreamService:
         raw_request = msh.marshal(gq_request)
 
         http_request = http.HttpRequest(
-            'https://api.groq.com/openai/v1/chat/completions',
+            self.URL,
             headers={
                 http.consts.HEADER_CONTENT_TYPE: http.consts.CONTENT_TYPE_JSON,
                 http.consts.HEADER_AUTH: http.consts.format_bearer_auth_header(check.not_none(self._api_key).reveal()),
@@ -78,50 +104,11 @@ class GroqChatChoicesStreamService:
             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]:
-                db = DelimitingBuffer([b'\r', b'\n', b'\r\n'])
-                sd = sse.SseDecoder()
-                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 []
-
-                        # FIXME: https://platform.openai.com/docs/guides/function-calling?api-mode=responses#streaming
-                        for so in sd.process_line(l):
-                            if isinstance(so, sse.SseEvent) and so.type == b'message':
-                                ss = so.data.decode('utf-8')
-                                if ss == '[DONE]':
-                                    return []
-
-                                sj = json.loads(ss)  # ChatCompletionChunk
-
-                                check.state(sj['object'] == 'chat.completion.chunk')
-
-                                ccc = msh.unmarshal(sj, pt.ChatCompletionChunk)
-
-                                # FIXME: stop reason
-                                if not ccc.choices:
-                                    continue
-
-                                if any(choice.finish_reason for choice in ccc.choices):
-                                    check.state(all(choice.finish_reason for choice in ccc.choices))
-                                    break
-
-                                await sink.emit(AiChoicesDeltas([
-                                    build_mc_ai_choice_deltas(choice.delta)
-                                    for choice in ccc.choices
-                                ]))
-
-                    if not b:
-                        return []
-
-            # raw_response = json.loads(check.not_none(http_response.data).decode('utf-8'))
-            # return rh.build_response(raw_response)
-
-            return await new_stream_response(rs, inner)
+        return await BytesHttpStreamResponseBuilder(
+            self._http_client,
+            lambda http_response: SimpleSseLinesHttpStreamResponseHandler(self._process_sse).as_lines().as_bytes(),
+            read_chunk_size=self.READ_CHUNK_SIZE,
+        ).new_stream_response(
+            http_request,
+            request.options,
+        )

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

@@ -179,7 +179,7 @@ class OllamaChatChoicesStreamService(BaseOllamaChatChoicesService):
                     for l in db.feed(b):
                         if isinstance(l, DelimitingBuffer.Incomplete):
                             # FIXME: handle
-                            return []
+                            raise TypeError(l)
 
                         lj = json.loads(l.decode('utf-8'))
                         lp: pt.ChatResponse = msh.unmarshal(lj, pt.ChatResponse)

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

@@ -145,6 +145,7 @@ def build_mc_ai_delta(delta: pt.ChatCompletionChunkChoiceDelta) -> AiDelta:
         )
 
     else:
+        # FIXME: no
         return ContentAiDelta('')
 
 

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

@@ -9,10 +9,8 @@ from omlish import typedvalues as tv
 from omlish.formats import json
 from omlish.http import all as http
 from omlish.http import sse
-from omlish.io.buffers import DelimitingBuffer
 
 from .....backends.openai import protocol as pt
-from ....chat.choices.services import ChatChoicesOutputs
 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
@@ -20,12 +18,11 @@ from ....chat.choices.stream.types import AiChoiceDeltas
 from ....chat.choices.stream.types import AiChoicesDeltas
 from ....chat.choices.stream.types import ChatChoicesStreamOption
 from ....configs import Config
+from ....http.stream import BytesHttpStreamResponseBuilder
+from ....http.stream import SimpleSseLinesHttpStreamResponseHandler
 from ....resources import ResourcesOption
-from ....resources import UseResources
 from ....standard import ApiKey
 from ....stream.services import StreamOption
-from ....stream.services import StreamResponseSink
-from ....stream.services import new_stream_response
 from .chat import OpenaiChatChoicesService
 from .format import OpenaiChatRequestHandler
 from .format import build_mc_ai_delta
@@ -54,11 +51,38 @@ class OpenaiChatChoicesStreamService:
             self._model_name = cc.pop(OpenaiChatChoicesService.DEFAULT_MODEL_NAME)
             self._api_key = ApiKey.pop_secret(cc, env='OPENAI_API_KEY')
 
+    URL: ta.ClassVar[str] = 'https://api.openai.com/v1/chat/completions'
+
+    def _process_sse(self, so: sse.SseDecoderOutput) -> ta.Sequence[AiChoicesDeltas | None]:
+        if not (isinstance(so, sse.SseEvent) and so.type == b'message'):
+            return []
+
+        ss = so.data.decode('utf-8')
+        if ss == '[DONE]':
+            return [None]
+
+        sj = json.loads(ss)  # ChatCompletionChunk
+
+        check.state(sj['object'] == 'chat.completion.chunk')
+
+        ccc = msh.unmarshal(sj, pt.ChatCompletionChunk)
+
+        # FIXME: stop reason
+        if not ccc.choices:
+            return []
+
+        if any(choice.finish_reason for choice in ccc.choices):
+            check.state(all(choice.finish_reason for choice in ccc.choices))
+            return [None]
+
+        return [AiChoicesDeltas([
+            AiChoiceDeltas([build_mc_ai_delta(choice.delta)])
+            for choice in ccc.choices
+        ])]
+
     READ_CHUNK_SIZE: ta.ClassVar[int] = -1
 
     async def invoke(self, request: ChatChoicesStreamRequest) -> ChatChoicesStreamResponse:
-        # check.isinstance(request, ChatRequest)
-
         rh = OpenaiChatRequestHandler(
             request.v,
             *[
@@ -78,7 +102,7 @@ class OpenaiChatChoicesStreamService:
         raw_request = msh.marshal(rh.oai_request())
 
         http_request = http.HttpRequest(
-            'https://api.openai.com/v1/chat/completions',
+            self.URL,
             headers={
                 http.consts.HEADER_CONTENT_TYPE: http.consts.CONTENT_TYPE_JSON,
                 http.consts.HEADER_AUTH: http.consts.format_bearer_auth_header(check.not_none(self._api_key).reveal()),
@@ -86,50 +110,11 @@ class OpenaiChatChoicesStreamService:
             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]:
-                db = DelimitingBuffer([b'\r', b'\n', b'\r\n'])
-                sd = sse.SseDecoder()
-                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 []
-
-                        # FIXME: https://platform.openai.com/docs/guides/function-calling?api-mode=responses#streaming
-                        for so in sd.process_line(l):
-                            if isinstance(so, sse.SseEvent) and so.type == b'message':
-                                ss = so.data.decode('utf-8')
-                                if ss == '[DONE]':
-                                    return []
-
-                                sj = json.loads(ss)  # ChatCompletionChunk
-
-                                check.state(sj['object'] == 'chat.completion.chunk')
-
-                                ccc = msh.unmarshal(sj, pt.ChatCompletionChunk)
-
-                                # FIXME: stop reason
-                                if not ccc.choices:
-                                    continue
-
-                                if any(choice.finish_reason for choice in ccc.choices):
-                                    check.state(all(choice.finish_reason for choice in ccc.choices))
-                                    break
-
-                                await sink.emit(AiChoicesDeltas([
-                                    AiChoiceDeltas([build_mc_ai_delta(choice.delta)])
-                                    for choice in ccc.choices
-                                ]))
-
-                    if not b:
-                        return []
-
-            # raw_response = json.loads(check.not_none(http_response.data).decode('utf-8'))
-            # return rh.build_response(raw_response)
-
-            return await new_stream_response(rs, inner)
+        return await BytesHttpStreamResponseBuilder(
+            self._http_client,
+            lambda http_response: SimpleSseLinesHttpStreamResponseHandler(self._process_sse).as_lines().as_bytes(),
+            read_chunk_size=self.READ_CHUNK_SIZE,
+        ).new_stream_response(
+            http_request,
+            request.options,
+        )

ommlds/minichain/http/stream.py

@@ -0,0 +1,195 @@
+"""
+TODO:
+ - better pipeline composition lol
+"""
+import typing as ta
+
+from omlish import check
+from omlish import dataclasses as dc
+from omlish import lang
+from omlish.http import all as http
+from omlish.http import sse
+from omlish.io.buffers import DelimitingBuffer
+
+from ..resources import UseResources
+from ..stream.services import StreamResponse
+from ..stream.services import StreamResponseSink
+from ..stream.services import new_stream_response
+from ..types import Option
+from ..types import Output
+
+
+##
+
+
+@dc.dataclass()
+@dc.extra_class_params(default_repr_fn=lang.opt_repr)
+class HttpStreamResponseError(Exception):
+    response: http.BaseHttpResponse
+
+    data: bytes | None = None
+    data_exception: Exception | None = None
+
+    @classmethod
+    async def from_response(cls, response: http.AsyncStreamHttpResponse) -> 'HttpStreamResponseError':
+        data: bytes | None = None
+        data_exception: Exception | None = None
+
+        try:
+            data = await response.stream.readall()
+        except Exception as de:  # noqa
+            data_exception = de
+
+        return HttpStreamResponseError(
+            response,
+            data=data,
+            data_exception=data_exception,
+        )
+
+
+##
+
+
+class HttpStreamResponseHandler(lang.Abstract):
+    def start(self) -> ta.Sequence[Output]:
+        return ()
+
+    def finish(self) -> ta.Sequence[Output]:
+        return ()
+
+
+##
+
+
+class BytesHttpStreamResponseHandler(HttpStreamResponseHandler, lang.Abstract):
+    def process_bytes(self, data: bytes) -> ta.Iterable:
+        return ()
+
+
+class BytesHttpStreamResponseBuilder:
+    def __init__(
+            self,
+            http_client: http.AsyncHttpClient | None,
+            handling: ta.Callable[[http.AsyncStreamHttpResponse], BytesHttpStreamResponseHandler],
+            *,
+            read_chunk_size: int = -1,
+    ) -> None:
+        super().__init__()
+
+        self._http_client = http_client
+        self._handling = handling
+        self._read_chunk_size = read_chunk_size
+
+    async def new_stream_response(
+            self,
+            http_request: http.HttpRequest,
+            options: ta.Sequence[Option],
+    ) -> StreamResponse:
+        async with UseResources.or_new(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))
+
+            if http_response.status != 200:
+                raise await HttpStreamResponseError.from_response(http_response)
+
+            handler = self._handling(http_response)
+
+            async def inner(sink: StreamResponseSink) -> ta.Sequence | None:
+                while True:
+                    b = await http_response.stream.read1(self._read_chunk_size)
+
+                    for v in handler.process_bytes(b):
+                        if v is None:
+                            break
+
+                        await sink.emit(v)
+
+                    if not b:
+                        break
+
+                return handler.finish()
+
+            return await new_stream_response(
+                rs,
+                inner,
+                handler.start(),
+            )
+
+
+##
+
+
+class LinesHttpStreamResponseHandler(HttpStreamResponseHandler, lang.Abstract):
+    def process_line(self, line: bytes) -> ta.Iterable:
+        return ()
+
+    def as_bytes(self) -> BytesHttpStreamResponseHandler:
+        return LinesBytesHttpStreamResponseHandler(self)
+
+
+class LinesBytesHttpStreamResponseHandler(BytesHttpStreamResponseHandler):
+    def __init__(self, handler: LinesHttpStreamResponseHandler) -> None:
+        super().__init__()
+
+        self._handler = handler
+
+        self._db = DelimitingBuffer([b'\r', b'\n', b'\r\n'])
+
+    def start(self) -> ta.Sequence[Output]:
+        return self._handler.start()
+
+    def process_bytes(self, data: bytes) -> ta.Iterable:
+        for o in self._db.feed(data):
+            if isinstance(o, bytes):
+                yield from self._handler.process_line(o)
+
+            else:
+                raise TypeError(o)
+
+    def finish(self) -> ta.Sequence[Output]:
+        check.state(self._db.is_closed)
+
+        return self._handler.finish()
+
+
+##
+
+
+class SseHttpStreamResponseHandler(HttpStreamResponseHandler, lang.Abstract):
+    def process_sse(self, so: sse.SseDecoderOutput) -> ta.Iterable:
+        return ()
+
+    def as_lines(self) -> LinesHttpStreamResponseHandler:
+        return SseLinesHttpStreamResponseHandler(self)
+
+
+class SseLinesHttpStreamResponseHandler(LinesHttpStreamResponseHandler):
+    def __init__(self, handler: SseHttpStreamResponseHandler) -> None:
+        super().__init__()
+
+        self._handler = handler
+
+        self._sd = sse.SseDecoder()
+
+    def start(self) -> ta.Sequence[Output]:
+        return self._handler.start()
+
+    def process_line(self, line: bytes) -> ta.Iterable:
+        for so in self._sd.process_line(line):
+            yield from self._handler.process_sse(so)
+
+    def finish(self) -> ta.Sequence[Output]:
+        return self._handler.finish()
+
+
+#
+
+
+class SimpleSseLinesHttpStreamResponseHandler(SseHttpStreamResponseHandler):
+    def __init__(self, fn: ta.Callable[[sse.SseDecoderOutput], ta.Iterable]) -> None:
+        super().__init__()
+
+        self._fn = fn
+
+    def process_sse(self, so: sse.SseDecoderOutput) -> ta.Iterable:
+        return self._fn(so)

ommlds/minichain/resources.py

@@ -33,6 +33,10 @@ class ResourcesRefNotRegisteredError(Exception):
 
 @ta.final
 class Resources(lang.Final, lang.NotPicklable):
+    """
+    Essentially a reference-tracked AsyncContextManager.
+    """
+
     def __init__(
             self,
             *,
@@ -80,10 +84,14 @@ class Resources(lang.Final, lang.NotPicklable):
         @contextlib.asynccontextmanager
         async def inner():
             init_ref = Resources._InitRef()
+
             res = Resources(init_ref=init_ref, **kwargs)
+
             await res.init()
+
             try:
                 yield res
+
             finally:
                 await res.remove_ref(init_ref)
 
@@ -94,6 +102,7 @@ class Resources(lang.Final, lang.NotPicklable):
     def add_ref(self, ref: ResourcesRef) -> None:
         check.isinstance(ref, ResourcesRef)
         check.state(not self._closed)
+
         self._refs.add(ref)
 
     def has_ref(self, ref: ResourcesRef) -> bool:
@@ -101,10 +110,13 @@ class Resources(lang.Final, lang.NotPicklable):
 
     async def remove_ref(self, ref: ResourcesRef) -> None:
         check.isinstance(ref, ResourcesRef)
+
         try:
             self._refs.remove(ref)
+
         except KeyError:
             raise ResourcesRefNotRegisteredError(ref) from None
+
         if not self._no_autoclose and not self._refs:
             await self.aclose()
 
@@ -112,10 +124,12 @@ class Resources(lang.Final, lang.NotPicklable):
 
     def enter_context(self, cm: ta.ContextManager[T]) -> T:
         check.state(not self._closed)
+
         return self._aes.enter_context(cm)
 
     async def enter_async_context(self, cm: ta.AsyncContextManager[T]) -> T:
         check.state(not self._closed)
+
         return await self._aes.enter_async_context(cm)
 
     #
@@ -150,7 +164,11 @@ class Resources(lang.Final, lang.NotPicklable):
 class ResourceManaged(ResourcesRef, lang.Final, lang.NotPicklable, ta.Generic[T]):
     """
     A class to 'handoff' a ref to a `Resources`, allowing the `Resources` to temporarily survive being passed from
-    instantiation within a callee to being `__aenter__`'d in the caller.
+    instantiation within a callee.
+
+    This class wraps an arbitrary value, likely an object referencing resources managed by the `Resources`, which is
+    accessed by `__aenter__`'ing. However, as the point of this class is handoff of a `Resources`, not necessarily some
+    arbitrary value, the value needn't necessarily be related to the `Resources`, or may even be `None`.
 
     The ref to the `Resources` is allocated in the ctor, so the contract is that an instance of this must be immediately
     `__aenter__`'d before doing anything else with the return value of the call. Failure to do so leaks the `Resources`.
@@ -172,11 +190,13 @@ class ResourceManaged(ResourcesRef, lang.Final, lang.NotPicklable, ta.Generic[T]
     async def __aenter__(self) -> T:
         check.state(self.__state == 'new')
         self.__state = 'entered'
+
         return self.__v
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         check.state(self.__state == 'entered')
         self.__state = 'exited'
+
         await self.__resources.remove_ref(self)
 
     def __del__(self) -> None:
@@ -203,6 +223,7 @@ class UseResources(tv.UniqueScalarTypedValue[Resources], ResourcesOption, lang.F
         if (ur := tv.as_collection(options).get(UseResources)) is not None:
             async with ResourceManaged(ur.v, ur.v) as rs:
                 yield rs
+
         else:
             async with Resources.new() as rs:
                 yield rs

ommlds/minichain/stream/services.py

@@ -1,3 +1,7 @@
+"""
+TODO:
+ - new_stream_response carry TypeAlias __orig_type__ somehow
+"""
 import abc
 import itertools
 import types
@@ -104,17 +108,23 @@ class _StreamServiceResponse(StreamResponseIterator[V, OutputT]):
             return _StreamServiceResponse._Emit(self._ssr, item)
 
     _state: ta.Literal['new', 'running', 'closed'] = 'new'
+
     _sink: _Sink[V]
-    _a: ta.Any
+
     _cr: ta.Any
+    _a: ta.Any
+    _g: ta.Any
 
     async def __aenter__(self) -> ta.Self:
         check.state(self._state == 'new')
         self._state = 'running'
+
         self._sink = _StreamServiceResponse._Sink(self)
+
         self._cr = self._fn(self._sink)
         self._a = self._cr.__await__()
         self._g = iter(self._a)
+
         return self
 
     @types.coroutine
@@ -123,8 +133,10 @@ class _StreamServiceResponse(StreamResponseIterator[V, OutputT]):
         self._state = 'closed'
         if old_state != 'running':
             return
+
         if self._cr.cr_running or self._cr.cr_suspended:
             cex = StreamServiceCancelledError()
+
             i = None
             for n in itertools.count():
                 try:
@@ -132,13 +144,18 @@ class _StreamServiceResponse(StreamResponseIterator[V, OutputT]):
                         x = self._g.throw(cex)
                     else:
                         x = self._g.send(i)
+
                 except StreamServiceCancelledError as cex2:
                     if cex2 is cex:
                         break
+
                     raise
+
                 i = yield x
+
         if self._cr.cr_running:
             raise RuntimeError(f'Coroutine {self._cr!r} not terminated')
+
         if self._g is not self._a:
             self._g.close()
         self._a.close()
@@ -156,15 +173,18 @@ class _StreamServiceResponse(StreamResponseIterator[V, OutputT]):
     @types.coroutine
     def _anext(self):
         check.state(self._state == 'running')
+
         i = None
         while True:
             try:
                 x = self._g.send(i)
+
             except StopIteration as e:
                 if e.value is not None:
                     self._outputs = tv.TypedValues(*check.isinstance(e.value, ta.Sequence))
                 else:
                     self._outputs = tv.TypedValues()
+
                 raise StopAsyncIteration from None
 
             if isinstance(x, _StreamServiceResponse._Emit) and x.ssr is self:
@@ -200,11 +220,14 @@ async def new_stream_response(
     ssr = _StreamServiceResponse(fn)
 
     v = rs.new_managed(await rs.enter_async_context(ssr))
+
     try:
         return StreamResponse(v, outputs or [])
+
     except BaseException:  # noqa
         # The StreamResponse ctor can raise - for example in `_tv_field_coercer` - in which case we need to clean up the
         # resources ref we have already allocated before reraising.
         async with v:
             pass
+
         raise

ommlds/minichain/wrappers/firstinwins.py

@@ -34,7 +34,7 @@ AnyFirstInWinsService: ta.TypeAlias = ta.Union[
 ##
 
 
-@dc.dataclass(frozen=True)
+@dc.dataclass()
 class FirstInWinsServiceCancelledError(Exception):
     e: BaseException
 

ommlds/minichain/wrappers/instrument.py

@@ -1,6 +1,6 @@
 """
 TODO:
- - final sream outputs
+ - final stream outputs
 """
 import datetime
 import typing as ta