virtus 0.0.3__py3-none-any.whl

virtus/__init__.py

@@ -0,0 +1,109 @@
+from typing import TYPE_CHECKING
+
+from virtus.agent import Agent, agent
+from virtus.capability import Capabilities, context, on_input, on_message
+from virtus.compaction import compaction
+from virtus.journal import Journal, JsonJournal, MemoryJournal
+from virtus.knowledge import Document, Passage, Store, knowledge, memory
+from virtus.lazy import lazy
+from virtus.live import Interrupted, Level, Live
+from virtus.message import (
+    AssistantMessage,
+    Audio,
+    AudioDelta,
+    AudioType,
+    Delta,
+    Event,
+    File,
+    FileType,
+    Image,
+    ImageType,
+    MediaType,
+    Message,
+    TextDelta,
+    ToolCall,
+    ToolCallDelta,
+    ToolResult,
+    Usage,
+    UserMessage,
+)
+from virtus.middleware import Capability, Middleware, Next
+from virtus.observe import observe
+from virtus.output import Snapshot
+from virtus.provider import Embedder, LiveProvider, Provider, Session
+from virtus.request import Request
+from virtus.thread import Conversation, Thread
+from virtus.tool import Tool, tool
+from virtus.vault import secrets
+from virtus.workspace import workspace
+
+if TYPE_CHECKING:
+    from virtus.servers import Http, Stdio, mcp
+
+__all__ = [
+    "Agent",
+    "AssistantMessage",
+    "Audio",
+    "AudioDelta",
+    "AudioType",
+    "Capabilities",
+    "Capability",
+    "Conversation",
+    "Delta",
+    "Document",
+    "Embedder",
+    "Event",
+    "File",
+    "FileType",
+    "Http",
+    "Image",
+    "ImageType",
+    "Interrupted",
+    "Journal",
+    "JsonJournal",
+    "Level",
+    "Live",
+    "LiveProvider",
+    "MediaType",
+    "MemoryJournal",
+    "Message",
+    "Middleware",
+    "Next",
+    "Passage",
+    "Provider",
+    "Request",
+    "Session",
+    "Snapshot",
+    "Stdio",
+    "Store",
+    "TextDelta",
+    "Thread",
+    "Tool",
+    "ToolCall",
+    "ToolCallDelta",
+    "ToolResult",
+    "Usage",
+    "UserMessage",
+    "agent",
+    "compaction",
+    "context",
+    "knowledge",
+    "lazy",
+    "mcp",
+    "memory",
+    "observe",
+    "on_input",
+    "on_message",
+    "secrets",
+    "tool",
+    "workspace",
+]
+
+
+def __getattr__(name: str) -> object:
+    """`mcp`, `Stdio` and `Http` load on demand, so the `mcp` extra is only needed by those who use it."""
+    if name in ("Http", "Stdio", "mcp"):
+        from virtus import servers
+
+        return getattr(servers, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

virtus/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.0.3'
+__version_tuple__ = version_tuple = (0, 0, 3)
+
+__commit_id__ = commit_id = None

virtus/agent.py

@@ -0,0 +1,177 @@
+import inspect
+from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
+from dataclasses import dataclass
+from typing import Final, overload
+from uuid import uuid4
+
+from pydantic import BaseModel
+
+from virtus.capability import flatten
+from virtus.dispatch import Target, dispatching
+from virtus.journal import Journal, MemoryJournal
+from virtus.message import Audio, File, Image, UserMessage, placeholder
+from virtus.middleware import Capability, Middleware
+from virtus.output import Output, output
+from virtus.provider import Provider
+from virtus.thread import Conversation, Thread
+from virtus.tool import parameters_model
+
+
+class Unset:
+    """Marks a `replace` setting to keep."""
+
+
+UNSET: Final = Unset()
+
+
+class Agent[**P, R](Capability):
+    """An LLM call loop described by a template function.
+
+    The template's docstring, formatted with the call arguments, is the user message. Its return
+    annotation decides what a call returns; see `virtus.output`. `R` is that return type, so an
+    agent declared `-> Person` has `R = Awaitable[Person]` and one declared `-> AsyncIterator[Person]`
+    has `R = AsyncIterator[Person]`.
+
+    Used as a capability, the agent becomes a target of the caller's `dispatch` tool (`virtus.dispatch`):
+    the calling model addresses it by name for a one-off request, or as `name:thread` to keep a
+    conversation going, and passes the template's arguments, which format the message exactly as a
+    direct call does. Those threads, one-off requests included, live under the caller's thread in its
+    journal, so separate conversations never share them.
+    """
+
+    def __init__(
+        self,
+        template: Callable[P, object],
+        *,
+        output: Output[R],
+        provider: Provider,
+        name: str | None = None,
+        description: str | None = None,
+        system: str = "",
+        capabilities: Sequence[Capability | Middleware] = (),
+        journal: Journal | None = None,
+    ) -> None:
+        signature = inspect.signature(template)
+
+        self.name = name or template.__name__
+        self.description = description or inspect.getdoc(template) or ""
+        self.provider = provider
+        self.system = system
+        self.journal: Journal = journal or MemoryJournal()
+        """Where every thread of this agent keeps its history; direct one-off calls are never written to it."""
+        self.output: Final[Output[R]] = output
+        self.capabilities: Final[tuple[Capability | Middleware, ...]] = tuple(capabilities)
+        self.chain: tuple[Middleware, ...] = flatten(capabilities)
+        """Middleware applied around each step of a run, outermost first."""
+        self.parameters: type[BaseModel] = parameters_model(template)
+        self._template = template
+        self._signature = signature
+        self._prompt = inspect.getdoc(template) or "\n\n".join(f"{{{p}}}" for p in signature.parameters)
+
+        if not self._prompt:
+            raise TypeError(f"{self.name}: an agent needs parameters or a docstring")
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        return Thread(self)(*args, **kwargs)
+
+    def replace(
+        self,
+        *,
+        provider: Provider | Unset = UNSET,
+        name: str | Unset = UNSET,
+        description: str | Unset = UNSET,
+        system: str | Unset = UNSET,
+        capabilities: Sequence[Capability | Middleware] | Unset = UNSET,
+        journal: Journal | Unset = UNSET,
+    ) -> "Agent[P, R]":
+        """A copy with these settings changed; same template and output mode, and the same journal unless
+        one is given, so both see the same threads.
+
+        `capabilities` replaces the whole list; `extend` appends to it.
+        """
+
+        def pick[T](value: T | Unset, current: T) -> T:
+            return current if isinstance(value, Unset) else value
+
+        return Agent(
+            self._template,
+            output=self.output,
+            provider=pick(provider, self.provider),
+            name=pick(name, self.name),
+            description=pick(description, self.description),
+            system=pick(system, self.system),
+            capabilities=pick(capabilities, self.capabilities),
+            journal=pick(journal, self.journal),
+        )
+
+    def extend(self, *capabilities: Capability | Middleware) -> "Agent[P, R]":
+        """A copy with these capabilities appended after the current ones."""
+        return self.replace(capabilities=(*self.capabilities, *capabilities))
+
+    def thread(self, id: str | None = None) -> Thread[P, R]:
+        """A handle to the conversation with this id in the journal, new or already under way.
+
+        Without an id, start a new one.
+        """
+        return Thread(self, (self.name, id or uuid4().hex), self.journal)
+
+    def input(self, *args: object, **kwargs: object) -> UserMessage:
+        """The user message a call with these arguments sends.
+
+        `Image`, `Audio` and `File` arguments, alone or in a sequence, are attached to the message; in
+        the prompt they read as `[image]`, `[audio]` and `[file]`.
+        """
+        bound = self._signature.bind(*args, **kwargs)
+        bound.apply_defaults()
+
+        values = list(bound.arguments.values())
+        images = tuple(image for value in values for image in Image.found_in(value))
+        audio = tuple(clip for value in values for clip in Audio.found_in(value))
+        files = tuple(file for value in values for file in File.found_in(value))
+        text = {name: placeholder(value) or value for name, value in bound.arguments.items()}
+        return UserMessage(self._prompt.format(**text), images, audio, files)
+
+    def middleware(self) -> tuple[Middleware, ...]:
+        def conversation(parent: Conversation, thread: str | None) -> Conversation:
+            return parent.child((self.name, thread or uuid4().hex), self)
+
+        return (dispatching(Target(self, conversation)),)
+
+
+def agent(
+    *,
+    provider: Provider,
+    name: str | None = None,
+    description: str | None = None,
+    system: str = "",
+    capabilities: Sequence[Capability | Middleware] = (),
+    journal: Journal | None = None,
+) -> "Decorator":
+    """Turn an async template function into an `Agent`; see `virtus.output` for the return types."""
+    return Decorator(provider, name, description, system, capabilities, journal)
+
+
+@dataclass(frozen=True)
+class Decorator:
+    provider: Provider
+    name: str | None
+    description: str | None
+    system: str
+    capabilities: Sequence[Capability | Middleware]
+    journal: Journal | None
+
+    @overload
+    def __call__[**P, T](self, template: Callable[P, Awaitable[AsyncIterator[T]]]) -> Agent[P, AsyncIterator[T]]: ...
+    @overload
+    def __call__[**P, Out](self, template: Callable[P, Awaitable[Out]]) -> Agent[P, Awaitable[Out]]: ...
+    def __call__[**P](self, template: Callable[P, Awaitable[object]]) -> Agent[P, object]:
+        return Agent(
+            template,
+            output=output(template),
+            provider=self.provider,
+            name=self.name,
+            description=self.description,
+            system=self.system,
+            capabilities=self.capabilities,
+            journal=self.journal,
+        )

virtus/capability.py

@@ -0,0 +1,126 @@
+import re
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Sequence
+from dataclasses import dataclass, replace
+from itertools import takewhile
+from types import MethodType
+
+from virtus.message import AudioDelta, Event, Message, TextDelta, ToolCallDelta, UserMessage
+from virtus.middleware import Capability, Middleware, Next
+from virtus.request import Request
+from virtus.tool import Tool
+
+
+class Capabilities(Capability):
+    """Groups tools, context and hooks declared as methods, so they can share state through `self`.
+
+    Tool names are prefixed with the class name in snake case: `Memory.remember` -> `memory__remember`.
+    """
+
+    def middleware(self) -> tuple[Middleware, ...]:
+        prefix = re.sub(r"(?<!^)(?=[A-Z])", "_", type(self).__name__).lower()
+        names = dict.fromkeys(
+            name
+            for klass in reversed(type(self).__mro__)
+            for name, value in vars(klass).items()
+            if isinstance(value, Capability)
+        )
+
+        def member_middleware(name: str) -> Iterator[Middleware]:
+            match getattr(self, name):
+                case Tool() as tool:
+                    yield from replace(tool, name=f"{prefix}__{tool.name}").middleware()
+                case Capability() as capability:
+                    yield from capability.middleware()
+
+        return tuple(f for name in names for f in member_middleware(name))
+
+
+@dataclass(frozen=True)
+class Context(Capability):
+    function: Callable[..., str | Awaitable[str]]
+
+    def __get__(self, instance: object | None, owner: type | None = None) -> "Context":
+        return self if instance is None else replace(self, function=MethodType(self.function, instance))
+
+    def middleware(self) -> tuple[Middleware, ...]:
+        async def add(request: Request, next: Next) -> AsyncIterator[Event]:
+            text = await _resolve(self.function())
+            system = "\n\n".join(part for part in (request.system, text) if part)
+            async for event in next(replace(request, system=system)):
+                yield event
+
+        return (add,)
+
+
+@dataclass(frozen=True)
+class OnInput(Capability):
+    function: Callable[..., UserMessage | Awaitable[UserMessage]]
+
+    def __get__(self, instance: object | None, owner: type | None = None) -> "OnInput":
+        return self if instance is None else replace(self, function=MethodType(self.function, instance))
+
+    def middleware(self) -> tuple[Middleware, ...]:
+        async def rewrite(request: Request, next: Next) -> AsyncIterator[Event]:
+            if request.iteration == 0:
+                request = replace(request, messages=await self._rewrite(request.messages))
+
+            async for event in next(request):
+                yield event
+
+        return (rewrite,)
+
+    async def _rewrite(self, messages: tuple[Message, ...]) -> tuple[Message, ...]:
+        """The new input is the run of user messages at the end of the conversation."""
+        start = len(messages) - sum(1 for _ in takewhile(lambda m: isinstance(m, UserMessage), reversed(messages)))
+        rewritten = [await _resolve(self.function(m)) if isinstance(m, UserMessage) else m for m in messages[start:]]
+        return (*messages[:start], *rewritten)
+
+
+@dataclass(frozen=True)
+class OnMessage(Capability):
+    """Sees every complete message: assistant replies, including tool calls, and tool results.
+
+    Streamed deltas pass through untouched. Raising on an assistant message stops the run
+    before its tool calls execute; returning a modified message does not change which calls run.
+    """
+
+    function: Callable[..., Message | Awaitable[Message]]
+
+    def __get__(self, instance: object | None, owner: type | None = None) -> "OnMessage":
+        return self if instance is None else replace(self, function=MethodType(self.function, instance))
+
+    def middleware(self) -> tuple[Middleware, ...]:
+        async def apply(request: Request, next: Next) -> AsyncIterator[Event]:
+            async for event in next(request):
+                match event:
+                    case TextDelta() | AudioDelta() | ToolCallDelta():
+                        yield event
+                    case message:
+                        result = await _resolve(self.function(message))
+                        if type(result) is not type(message):
+                            raise TypeError(f"on_message must return a {type(message).__name__}, got {result!r}")
+                        yield result
+
+        return (apply,)
+
+
+def context(function: Callable[..., str | Awaitable[str]]) -> Context:
+    """Append the returned text to the system prompt of every provider call."""
+    return Context(function)
+
+
+def on_input(function: Callable[..., UserMessage | Awaitable[UserMessage]]) -> OnInput:
+    """Rewrite, or reject by raising, each new user message once per run."""
+    return OnInput(function)
+
+
+def on_message(function: Callable[..., Message | Awaitable[Message]]) -> OnMessage:
+    return OnMessage(function)
+
+
+def flatten(capabilities: Sequence[Capability | Middleware]) -> tuple[Middleware, ...]:
+    return tuple(f for c in capabilities for f in (c.middleware() if isinstance(c, Capability) else (c,)))
+
+
+async def _resolve[T](value: T | Awaitable[T]) -> T:
+    return await value if isinstance(value, Awaitable) else value

virtus/compaction.py

@@ -0,0 +1,75 @@
+from collections.abc import AsyncIterator, Sequence
+from dataclasses import replace
+
+from virtus.message import AssistantMessage, Event, Message, ToolResult, UserMessage
+from virtus.middleware import Middleware, Next
+from virtus.provider import Provider
+from virtus.request import Request
+
+SUMMARY_PROMPT = (
+    "Summarize the conversation below so the summary can replace it. Keep facts, decisions, results "
+    "of tool calls, open questions and anything the user asked to remember. Be concise."
+)
+
+
+SUMMARY = "Summary of the earlier conversation:"
+
+
+def compaction(
+    max_tokens: int = 32_000,
+    keep: int = 8,
+    *,
+    chars_per_token: int = 4,
+    provider: Provider | None = None,
+    prompt: str = SUMMARY_PROMPT,
+    header: str = SUMMARY,
+) -> Middleware:
+    """Replace older messages with a summary once the conversation exceeds `max_tokens`.
+
+    Tokens are estimated as `chars_per_token` characters each. At least the last `keep` messages stay
+    verbatim; the boundary moves back so a tool result is never separated from the call that produced it.
+    The summary is written by `provider`, by default the request's, following `prompt`, and goes in as a
+    user message opened by `header`. Like any rewrite of the messages it becomes the history, so it is not
+    recomputed until the conversation grows past the limit again.
+    """
+
+    async def summarize(messages: Sequence[Message], request: Request) -> str:
+        summary = replace(request, system=prompt, messages=(UserMessage(_transcript(messages)),), tools=(), output=None)
+        replies = [e async for e in (provider or request.provider).complete(summary) if isinstance(e, AssistantMessage)]
+        return replies[-1].content
+
+    async def compact(request: Request, next: Next) -> AsyncIterator[Event]:
+        messages = request.messages
+        start = max(
+            (i for i in range(1, len(messages) - keep + 1) if not isinstance(messages[i], ToolResult)),
+            default=0,
+        )
+
+        if start and len(_transcript(messages)) // chars_per_token > max_tokens:
+            summary = UserMessage(f"{header}\n\n{await summarize(messages[:start], request)}")
+            request = replace(request, messages=(summary, *messages[start:]))
+
+        async for event in next(request):
+            yield event
+
+    return compact
+
+
+def _transcript(messages: Sequence[Message]) -> str:
+    def lines(message: Message) -> list[str]:
+        match message:
+            case UserMessage(content=content, images=images, audio=audio, files=files):
+                return [f"user: {content}{_media(images, audio, files)}"]
+            case AssistantMessage(content=content, tool_calls=calls, images=images, audio=audio):
+                text = [f"assistant: {content}{_media(images, (audio,) if audio else ())}"] if content or images else []
+                return [*text, *(f"assistant called {c.name}({c.arguments})" for c in calls)]
+            case ToolResult(name=name, content=content, images=images, audio=audio, files=files):
+                return [f"{name} returned: {content}{_media(images, audio, files)}"]
+
+    return "\n".join(line for message in messages for line in lines(message))
+
+
+def _media(images: Sequence[object], audio: Sequence[object], files: Sequence[object] = ()) -> str:
+    counts = ((images, "image(s)"), (audio, "audio clip(s)"), (files, "file(s)"))
+    notes = [f"{len(items)} {kind}" for items, kind in counts if items]
+    return f" [{', '.join(notes)}]" if notes else ""

virtus/dispatch.py

@@ -0,0 +1,144 @@
+"""The one tool through which an agent talks to other agents.
+
+Every agent (or thread) given as a capability adds itself as a target of the same `dispatch` tool
+instead of becoming a tool of its own, so the model learns one tool however many agents it can
+reach. The `arguments` are those of the addressed agent's template, which formats the message as a
+direct call would; the targets and their open conversations are listed in the system prompt.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import re
+from collections.abc import AsyncIterator, Callable, Mapping
+from dataclasses import dataclass, replace
+from functools import reduce
+from operator import or_
+from typing import TYPE_CHECKING, get_args
+
+from pydantic import Field, TypeAdapter, create_model
+
+from virtus.message import Audio, Event, Image
+from virtus.middleware import Middleware, Next
+from virtus.request import Request
+from virtus.tool import Tool
+
+if TYPE_CHECKING:
+    from virtus.agent import Agent
+    from virtus.thread import Conversation
+
+
+@dataclass(frozen=True)
+class Target:
+    """An agent the model can address: `conversation` gives the thread behind `name` or `name:thread`."""
+
+    agent: Agent[..., object]
+    conversation: Callable[[Conversation, str | None], Conversation]
+
+
+@dataclass(frozen=True)
+class Failed:
+    """What a failing address answers with, so the other addresses' replies still come back."""
+
+    error: str
+
+    def as_reply(self) -> dict[str, str]:
+        return {"error": self.error}
+
+
+@dataclass(frozen=True)
+class Dispatch(Tool):
+    targets: tuple[Target, ...]
+
+    @classmethod
+    def to(cls, targets: tuple[Target, ...], parent: Conversation) -> Dispatch:
+        by_name = {target.agent.name: target for target in targets}
+
+        async def dispatch(to: list[str], arguments: dict[str, object]) -> dict[str, object]:
+            if unknown := [address for address in to if address.partition(":")[0] not in by_name]:
+                raise ValueError(f"unknown agents: {', '.join(unknown)}; known: {', '.join(by_name)}")
+
+            async def deliver(address: str) -> object | Failed:
+                """The agent's reply, or `Failed`, so one address failing doesn't lose the others' replies."""
+                name, _, thread = address.partition(":")
+                target = by_name[name]
+                try:
+                    values = dict(target.agent.parameters.model_validate(arguments))
+                    return await target.conversation(parent, thread or None).send(target.agent.input(**values))
+                except Exception as error:
+                    return Failed(f"{type(error).__name__}: {error}")
+
+            replies = dict(zip(to, await asyncio.gather(*(deliver(address) for address in to)), strict=True))
+            failed = {address: reply for address, reply in replies.items() if isinstance(reply, Failed)}
+            if failed and len(failed) == len(replies):
+                raise ValueError("\n".join(f"{address}: {reply.error}" for address, reply in failed.items()))
+            return {address: r.as_reply() if isinstance(r, Failed) else r for address, r in replies.items()}
+
+        return cls("dispatch", DESCRIPTION, PARAMETERS, dispatch, targets)
+
+    @property
+    def schema(self) -> Mapping[str, object]:
+        """The parameters schema with `arguments` spelled out as one of the targets' templates."""
+        union = reduce(or_, (target.agent.parameters for target in self.targets))
+        arguments = TypeAdapter[object](union).json_schema()
+        schema = PARAMETERS.model_json_schema()
+        defs = {**schema.get("$defs", {}), **arguments.pop("$defs", {})}
+        properties = {**schema["properties"], "arguments": {**schema["properties"]["arguments"], **arguments}}
+        return {**schema, "properties": properties, **({"$defs": defs} if defs else {})}
+
+
+def dispatching(target: Target) -> Middleware:
+    """Middleware adding `target` to the request's `dispatch` tool, creating the tool if it is the first.
+
+    An agent with an `Image` or `Audio` parameter is refused: the model has no way to pass media.
+    """
+    fields = target.agent.parameters.model_fields.items()
+    if media := [name for name, field in fields if _media(field.annotation)]:
+        names = ", ".join(media)
+        raise TypeError(f"{target.agent.name}: can't be dispatched to; the model can't fill {names} with media")
+
+    async def add(request: Request, next: Next) -> AsyncIterator[Event]:
+        parent = request.thread
+        found = next_dispatch(request)
+        tool = Dispatch.to((*(found.targets if found else ()), target), parent)
+        tools = (*(t for t in request.tools if t is not found), tool)
+
+        name = target.agent.name
+        open = [thread for agent, thread in parent.children if agent == name and not ONE_OFF.fullmatch(thread)]
+        line = f"- {name}: {target.agent.description}" + (f" Open conversations: {', '.join(open)}" if open else "")
+        if found:
+            system = f"{request.system}\n{line}"
+        else:
+            system = "\n\n".join(part for part in (request.system, f"{HEADER}\n{line}") if part)
+
+        async for event in next(replace(request, system=system, tools=tools)):
+            yield event
+
+    return add
+
+
+def next_dispatch(request: Request) -> Dispatch | None:
+    return next((tool for tool in request.tools if isinstance(tool, Dispatch)), None)
+
+
+def _media(annotation: object) -> bool:
+    return annotation in (Image, Audio) or any(_media(inner) for inner in get_args(annotation))
+
+
+HEADER = "Agents you can `dispatch` to:"
+
+ONE_OFF = re.compile(r"[0-9a-f]{32}")
+"""The generated name of a one-off request's thread, kept in the caller's journal but not offered to the model."""
+
+DESCRIPTION = (
+    "Send a request to one or more agents and get their replies, keyed by address. An address is an agent "
+    "name for a one-off request, or `agent:conversation` to continue (or start) a named conversation with "
+    "that agent, which then remembers earlier requests. `arguments` are the addressed agent's; one call "
+    "sends the same arguments to every address, at once."
+)
+
+PARAMETERS = create_model(
+    "dispatch",
+    to=(list[str], Field(description="Addresses: `agent` or `agent:conversation`.", min_length=1)),
+    arguments=(dict[str, object], Field(description="The arguments of the addressed agent.")),
+)