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

ommlds/minichain/tools/reflect.py

@@ -23,7 +23,7 @@ from omlish import metadata as md
 from omlish import reflect as rfl
 from omlish.lite.cached import cached_nullary
 
-from ..content.materialize import CanContent
+from ..content.content import Content
 from .types import EnumToolDtype
 from .types import MappingToolDtype
 from .types import NullableToolDtype
@@ -159,7 +159,7 @@ class ToolReflector:
 
     #
 
-    def _prepare_desc(self, s: CanContent) -> CanContent:
+    def _prepare_desc(self, s: Content | None) -> Content | None:
         if s is None:
             return None
         if not self._raw_descs and isinstance(s, str):
@@ -178,14 +178,14 @@ class ToolReflector:
         for o in md.get_object_metadata(fn, type=_ToolSpecOverride):
             ts_ovr.update({
                 k: v
-                for k, v in dc.asdict(o).items()
+                for k, v in dc.shallow_asdict(o).items()
                 if k != 'params'
                 and v is not None
             })
             for op in (o.params or []):
                 p_ovr_dct.setdefault(check.non_empty_str(op.name), {}).update({
                     k: v
-                    for k, v in dc.asdict(op).items()
+                    for k, v in dc.shallow_asdict(op).items()
                     if v is not None
                 })
 
@@ -240,7 +240,7 @@ class ToolReflector:
                 ds_p = ds_p_dct.get(pn)
                 sig_p = sig_p_dct.get(pn)
 
-                p_desc: CanContent
+                p_desc: Content | None
                 if (p_desc := ovr_p.get('desc')) is None:
                     if ds_p is not None:
                         p_desc = ds_p.description
@@ -286,7 +286,7 @@ class ToolReflector:
             if md_tp is not None:
                 tp_kw.update({
                     k: v
-                    for k, v in dc.asdict(md_tp).items()
+                    for k, v in dc.shallow_asdict(md_tp).items()
                     if v is not None
                 })
 

ommlds/minichain/tools/types.py

@@ -6,14 +6,11 @@ from omlish import cached
 from omlish import check
 from omlish import collections as col
 from omlish import dataclasses as dc
-from omlish import dispatch
 from omlish import lang
 from omlish import marshal as msh
 from omlish import reflect as rfl
 
-from ..content.materialize import CanContent
-from ..content.transforms.base import ContentTransform
-from ..content.types import Content
+from ..content.content import Content
 
 
 msh.register_global_module_import('._marshal', __package__)
@@ -150,7 +147,7 @@ class ToolParam:
 
     _: dc.KW_ONLY
 
-    desc: CanContent = None
+    desc: Content | None = None
 
     type: ToolDtype | None = None
 
@@ -168,12 +165,12 @@ class ToolSpec:
 
     _: dc.KW_ONLY
 
-    desc: CanContent = None
+    desc: Content | None = None
 
     params: ta.Sequence[ToolParam] | None = None
     allow_additional_params: bool | None = None
 
-    returns_desc: CanContent = None
+    returns_desc: Content | None = None
     returns_type: ToolDtype | None = None
 
     @cached.property
@@ -207,11 +204,11 @@ class ToolUseResult(lang.Final):
 ##
 
 
-class _ToolUseContentTransform(ContentTransform, lang.Final, lang.NotInstantiable):
-    @dispatch.install_method(ContentTransform.apply)
-    def apply_tool_use(self, tu: ToolUse) -> ToolUse:
-        return tu  # TODO: args are Content
-
-    @dispatch.install_method(ContentTransform.apply)
-    def apply_tool_use_result(self, tur: ToolUseResult) -> ToolUseResult:
-        return dc.replace(tur, c=self.apply(tur.c))  # noqa
+# class _ToolUseContentTransform(ContentTransform, lang.Final, lang.NotInstantiable):
+#     @dispatch.install_method(ContentTransform.apply)
+#     def apply_tool_use(self, tu: ToolUse) -> ToolUse:
+#         return tu  # TODO: args are Content
+#
+#     @dispatch.install_method(ContentTransform.apply)
+#     def apply_tool_use_result(self, tur: ToolUseResult) -> ToolUseResult:
+#         return dc.replace(tur, c=self.apply(tur.c))  # noqa

ommlds/minichain/vectors/_marshal.py

@@ -2,10 +2,10 @@
 TODO:
  - serialize as base64 bytes? at least support deserializing as it
 """
-import dataclasses as dc
 import typing as ta
 
 from omlish import check
+from omlish import dataclasses as dc
 from omlish import lang
 from omlish import marshal as msh
 from omlish import reflect as rfl

ommlds/minichain/vectors/embeddings.py

@@ -2,7 +2,7 @@ import typing as ta
 
 from omlish import lang
 
-from ..content.types import Content
+from ..content.content import Content
 from ..registries.globals import register_type
 from ..services import Request
 from ..services import Response

ommlds/minichain/wrappers/__init__.py

@@ -0,0 +1,7 @@
+"""
+TODO:
+ - ratelimit
+ - option - output injection / removal
+ - 'switch'
+ - ** reflection passthrough **
+"""

ommlds/minichain/wrappers/firstinwins.py

@@ -0,0 +1,144 @@
+"""
+TODO:
+ - stream lol
+  - take first open vs take first chunk
+"""
+import typing as ta
+
+from omlish import dataclasses as dc
+from omlish import lang
+
+from ..services.services import Service
+from ..types import Output
+from .services import MultiWrapperService
+from .services import WrappedOptionT
+from .services import WrappedOutputT
+from .services import WrappedRequest
+from .services import WrappedRequestV
+from .services import WrappedResponse
+from .services import WrappedResponseV
+from .stream import WrappedStreamOutputT
+from .stream import WrapperStreamService
+
+
+with lang.auto_proxy_import(globals()):
+    import asyncio
+
+
+AnyFirstInWinsService: ta.TypeAlias = ta.Union[
+    'FirstInWinsService',
+    'FirstInWinsStreamService',
+]
+
+
+##
+
+
+@dc.dataclass()
+class FirstInWinsServiceCancelledError(Exception):
+    e: BaseException
+
+
+class FirstInWinsServiceExceptionGroup(ExceptionGroup):
+    pass
+
+
+@dc.dataclass(frozen=True)
+class FirstInWinsServiceOutput(Output):
+    first_in_wins_service: AnyFirstInWinsService
+    response_service: Service
+    service_exceptions: ta.Mapping[Service, Exception] | None = None
+
+
+##
+
+
+class FirstInWinsService(
+    MultiWrapperService[
+        WrappedRequestV,
+        WrappedOptionT,
+        WrappedResponseV,
+        WrappedOutputT,
+    ],
+    lang.Abstract,
+):
+    pass
+
+
+##
+
+
+class FirstInWinsStreamService(
+    WrapperStreamService[
+        WrappedRequestV,
+        WrappedOptionT,
+        WrappedResponseV,
+        WrappedOutputT,
+        WrappedStreamOutputT,
+    ],
+    lang.Abstract,
+):
+    pass
+
+
+##
+
+
+class AsyncioFirstInWinsService(
+    FirstInWinsService[
+        WrappedRequestV,
+        WrappedOptionT,
+        WrappedResponseV,
+        WrappedOutputT,
+    ],
+):
+    async def invoke(self, request: WrappedRequest) -> WrappedResponse:
+        tasks: list = []
+        services_by_task: dict = {}
+        for svc in self._services:
+            task: asyncio.Task = asyncio.create_task(svc.invoke(request))  # type: ignore[arg-type]
+            tasks.append(task)
+            services_by_task[task] = svc
+
+        failures_by_service: dict[Service, Exception] = {}
+
+        try:
+            pending = set(tasks)
+
+            while pending:
+                done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
+
+                for t in done:
+                    svc = services_by_task[t]
+
+                    try:
+                        response = t.result()
+                    except asyncio.CancelledError as exc:
+                        failures_by_service[svc] = FirstInWinsServiceCancelledError(exc)
+                        continue
+                    except Exception as exc:  # noqa
+                        failures_by_service[svc] = exc
+                        continue
+
+                    for p in pending:
+                        p.cancel()
+
+                    await asyncio.gather(*pending, return_exceptions=True)
+
+                    return response.with_outputs(FirstInWinsServiceOutput(
+                        self,
+                        svc,
+                        failures_by_service,
+                    ))
+
+            raise FirstInWinsServiceExceptionGroup(
+                'All service calls failed',
+                list(failures_by_service.values()),
+            )
+
+        finally:
+            for t in tasks:
+                if not t.done():
+                    t.cancel()
+
+            await asyncio.gather(*tasks, return_exceptions=True)

ommlds/minichain/wrappers/instrument.py

@@ -0,0 +1,146 @@
+"""
+TODO:
+ - final stream outputs
+"""
+import datetime
+import typing as ta
+
+from omlish import dataclasses as dc
+from omlish import lang
+
+from ..services.requests import Request
+from ..services.responses import Response
+from .services import WrappedOptionT
+from .services import WrappedOutputT
+from .services import WrappedRequest
+from .services import WrappedRequestV
+from .services import WrappedResponse
+from .services import WrappedResponseV
+from .services import WrappedService
+from .services import WrapperService
+from .stream import WrappedStreamOutputT
+from .stream import WrappedStreamResponse
+from .stream import WrappedStreamService
+from .stream import WrapperStreamService
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@dc.extra_class_params(default_repr_fn=lang.opt_or_just_repr)
+class InstrumentedServiceEvent:
+    dt: datetime.datetime = dc.field(default_factory=lang.utcnow)
+
+    req: Request | None = None
+
+    resp: Response | None = None
+    exc: BaseException | None = None
+
+    stream_v: lang.Maybe[ta.Any] = lang.empty()
+
+
+class InstrumentedServiceEventSink(ta.Protocol):
+    def __call__(self, ev: InstrumentedServiceEvent) -> ta.Awaitable[None]: ...
+
+
+class ListInstrumentedServiceEventSink:
+    def __init__(self, lst: list[InstrumentedServiceEvent] | None = None) -> None:
+        super().__init__()
+
+        if lst is None:
+            lst = []
+        self._lst = lst
+
+    @property
+    def events(self) -> ta.Sequence[InstrumentedServiceEvent]:
+        return self._lst
+
+    async def __call__(self, ev: InstrumentedServiceEvent) -> None:
+        self._lst.append(ev)
+
+
+##
+
+
+class InstrumentedService(
+    WrapperService[
+        WrappedRequestV,
+        WrappedOptionT,
+        WrappedResponseV,
+        WrappedOutputT,
+    ],
+):
+    def __init__(
+            self,
+            service: WrappedService,
+            sink: InstrumentedServiceEventSink | None = None,
+    ) -> None:
+        super().__init__(service)
+
+        if sink is None:
+            sink = ListInstrumentedServiceEventSink()
+        self._sink = sink
+
+    async def invoke(self, request: WrappedRequest) -> WrappedResponse:
+        await self._sink(InstrumentedServiceEvent(req=request))
+
+        try:
+            resp = await self._service.invoke(request)
+
+        except Exception as e:  # noqa
+            await self._sink(InstrumentedServiceEvent(req=request, exc=e))
+
+            raise
+
+        await self._sink(InstrumentedServiceEvent(req=request, resp=resp))
+
+        return resp
+
+
+##
+
+
+class InstrumentedStreamService(
+    WrapperStreamService[
+        WrappedRequestV,
+        WrappedOptionT,
+        WrappedResponseV,
+        WrappedOutputT,
+        WrappedStreamOutputT,
+    ],
+):
+    def __init__(
+            self,
+            service: WrappedStreamService,
+            sink: InstrumentedServiceEventSink | None = None,
+    ) -> None:
+        super().__init__(service)
+
+        if sink is None:
+            sink = ListInstrumentedServiceEventSink()
+        self._sink = sink
+
+    async def invoke(self, request: WrappedRequest) -> WrappedStreamResponse:
+        await self._sink(InstrumentedServiceEvent(req=request))
+
+        try:
+            resp = await self._service.invoke(request)
+
+        except Exception as e:  # noqa
+            await self._sink(InstrumentedServiceEvent(req=request, exc=e))
+
+            raise
+
+        await self._sink(InstrumentedServiceEvent(req=request, resp=resp))
+
+        async def inner(sink):  # noqa
+            async with resp.v as st:
+                async for v in st:
+                    await self._sink(InstrumentedServiceEvent(req=request, resp=resp, stream_v=v))
+
+                    await sink(v)
+
+        # return resp.with_v(inner())
+
+        raise NotImplementedError

ommlds/minichain/wrappers/retry.py

@@ -0,0 +1,168 @@
+"""
+TODO:
+ - tenacity shit
+  - exception filter - retryable vs not
+   - explicit RetryableError wrapper?
+  - sleep
+   - jitter
+ - log, on retry / on except callbacks, blah blah
+ - stream retry:
+  - failed to open
+  - failed during stream
+   - buffer and replay??
+  - accept death mid-stream?
+  - ** probably **: cannot sanely impose any nontrivial stream retry strat at this layer -
+"""
+import typing as ta
+
+from omlish import dataclasses as dc
+
+from ..resources import Resources
+from ..stream.services import StreamResponseSink
+from ..stream.services import new_stream_response
+from ..types import Output
+from .services import WrappedOptionT
+from .services import WrappedOutputT
+from .services import WrappedRequest
+from .services import WrappedRequestV
+from .services import WrappedResponse
+from .services import WrappedResponseV
+from .services import WrappedService
+from .services import WrapperService
+from .stream import WrappedStreamOutputT
+from .stream import WrappedStreamResponse
+from .stream import WrappedStreamService
+from .stream import WrapperStreamService
+
+
+AnyRetryService: ta.TypeAlias = ta.Union[
+    'RetryService',
+    'RetryStreamService',
+]
+
+
+##
+
+
+@dc.dataclass()
+class RetryServiceMaxRetriesExceededError(Exception):
+    pass
+
+
+@dc.dataclass(frozen=True)
+class RetryServiceOutput(Output):
+    retry_service: AnyRetryService
+    num_retries: int
+
+
+##
+
+
+class RetryService(
+    WrapperService[
+        WrappedRequestV,
+        WrappedOptionT,
+        WrappedResponseV,
+        WrappedOutputT,
+    ],
+):
+    DEFAULT_MAX_RETRIES: ta.ClassVar[int] = 3
+
+    def __init__(
+            self,
+            service: WrappedService,
+            *,
+            max_retries: int | None = None,
+    ) -> None:
+        super().__init__(service)
+
+        if max_retries is None:
+            max_retries = self.DEFAULT_MAX_RETRIES
+        self._max_retries = max_retries
+
+    async def invoke(self, request: WrappedRequest) -> WrappedResponse:
+        n = 0
+
+        while True:
+            try:
+                resp = await self._service.invoke(request)
+
+            except Exception as e:  # noqa
+                if n < self._max_retries:
+                    n += 1
+                    continue
+
+                raise RetryServiceMaxRetriesExceededError from e
+
+            return resp.with_outputs(RetryServiceOutput(
+                retry_service=self,
+                num_retries=n,
+            ))
+
+        raise RuntimeError  # unreachable
+
+
+##
+
+
+class RetryStreamService(
+    WrapperStreamService[
+        WrappedRequestV,
+        WrappedOptionT,
+        WrappedResponseV,
+        WrappedOutputT,
+        WrappedStreamOutputT,
+    ],
+):
+    DEFAULT_MAX_RETRIES: ta.ClassVar[int] = 3
+
+    def __init__(
+            self,
+            service: WrappedStreamService,
+            *,
+            max_retries: int | None = None,
+    ) -> None:
+        super().__init__(service)
+
+        if max_retries is None:
+            max_retries = self.DEFAULT_MAX_RETRIES
+        self._max_retries = max_retries
+
+    async def invoke(self, request: WrappedRequest) -> WrappedStreamResponse:
+        n = 0
+
+        while True:
+            try:
+                async with Resources.new() as rs:
+                    in_resp = await self._service.invoke(request)
+                    in_vs = await rs.enter_async_context(in_resp.v)
+
+                    async def inner(sink: StreamResponseSink[WrappedResponseV]) -> ta.Sequence[WrappedOutputT] | None:
+                        async for v in in_vs:
+                            await sink.emit(v)
+
+                        return in_vs.outputs
+
+                    outs = [
+                        *in_resp.outputs,
+                        RetryServiceOutput(
+                            retry_service=self,
+                            num_retries=n,
+                        ),
+                    ]
+
+                    # FIXME: ??
+                    # if (ur := tv.as_collection(request.options).get(UseResources)) is not None:
+
+                    return await new_stream_response(
+                        rs,
+                        inner,
+                        outs,
+                    )
+
+            except Exception as e:  # noqa
+                if n < self._max_retries:
+                    n += 1
+                    continue
+
+                raise RetryServiceMaxRetriesExceededError from e

ommlds/minichain/wrappers/services.py

@@ -0,0 +1,98 @@
+import typing as ta
+
+from omlish import check
+from omlish import lang
+
+from ..services.requests import Request
+from ..services.responses import Response
+from ..services.services import Service
+from ..types import Option
+from ..types import Output
+
+
+P = ta.ParamSpec('P')
+
+
+WrappedRequestV = ta.TypeVar('WrappedRequestV')
+WrappedOptionT = ta.TypeVar('WrappedOptionT', bound=Option)
+
+WrappedResponseV = ta.TypeVar('WrappedResponseV')
+WrappedOutputT = ta.TypeVar('WrappedOutputT', bound=Output)
+
+
+WrappedRequest: ta.TypeAlias = Request[
+    WrappedRequestV,
+    WrappedOptionT,
+]
+
+WrappedResponse: ta.TypeAlias = Response[
+    WrappedResponseV,
+    WrappedOutputT,
+]
+
+WrappedService: ta.TypeAlias = Service[
+    WrappedRequest,
+    WrappedResponse,
+]
+
+
+##
+
+
+class WrapperService(
+    lang.Abstract,
+    ta.Generic[
+        WrappedRequestV,
+        WrappedOptionT,
+        WrappedResponseV,
+        WrappedOutputT,
+    ],
+):
+    def __init__(
+            self,
+            service: WrappedService,
+    ) -> None:
+        super().__init__()
+
+        self._service = service
+
+
+class MultiWrapperService(
+    lang.Abstract,
+    ta.Generic[
+        WrappedRequestV,
+        WrappedOptionT,
+        WrappedResponseV,
+        WrappedOutputT,
+    ],
+):
+    def __init__(
+            self,
+            *services: WrappedService,
+    ) -> None:
+        super().__init__()
+
+        self._services = check.not_empty(services)
+
+
+##
+
+
+def wrap_service(
+        wrapped: WrappedService,
+        wrapper: ta.Callable[
+            ta.Concatenate[
+                WrappedService,
+                P,
+            ],
+            WrapperService[
+                WrappedRequestV,
+                WrappedOptionT,
+                WrappedResponseV,
+                WrappedOutputT,
+            ],
+        ],
+        *args: P.args,
+        **kwargs: P.kwargs,
+) -> WrappedService:
+    return ta.cast(ta.Any, wrapper(wrapped, *args, **kwargs))