grasp_agents 0.1.5__py3-none-any.whl

grasp_agents/tool_orchestrator.py

@@ -0,0 +1,181 @@
+import asyncio
+import json
+from collections.abc import Coroutine, Sequence
+from logging import getLogger
+from typing import Any, Generic, Protocol
+
+from pydantic import BaseModel
+
+from .llm import LLM, LLMSettings
+from .llm_agent_state import LLMAgentState
+from .run_context import CtxT, RunContextWrapper
+from .typing.converters import Converters
+from .typing.message import AssistantMessage, Conversation, Message, ToolMessage
+from .typing.tool import BaseTool, ToolCall, ToolChoice
+
+logger = getLogger(__name__)
+
+
+class ToolCallLoopExitHandler(Protocol[CtxT]):
+    def __call__(
+        self,
+        conversation: Conversation,
+        *,
+        ctx: RunContextWrapper[CtxT] | None,
+        **kwargs: Any,
+    ) -> bool: ...
+
+
+class ToolOrchestrator(Generic[CtxT]):
+    def __init__(
+        self,
+        agent_id: str,
+        llm: LLM[LLMSettings, Converters],
+        tools: list[BaseTool[BaseModel, BaseModel, CtxT]] | None,
+        max_turns: int,
+        react_mode: bool = False,
+    ) -> None:
+        self._agent_id = agent_id
+
+        self._llm = llm
+        self._tools = tools
+        self.llm.tools = tools
+
+        self._max_turns = max_turns
+        self._react_mode = react_mode
+
+        self.tool_call_loop_exit_impl: ToolCallLoopExitHandler[CtxT] | None = None
+
+    @property
+    def agent_id(self) -> str:
+        return self._agent_id
+
+    @property
+    def llm(self) -> LLM[LLMSettings, Converters]:
+        return self._llm
+
+    @property
+    def tools(self) -> dict[str, BaseTool[BaseModel, BaseModel, CtxT]]:
+        return self._llm.tools or {}
+
+    @property
+    def max_turns(self) -> int:
+        return self._max_turns
+
+    def _tool_call_loop_exit(
+        self,
+        conversation: Conversation,
+        *,
+        ctx: RunContextWrapper[CtxT] | None = None,
+        **kwargs: Any,
+    ) -> bool:
+        if self.tool_call_loop_exit_impl:
+            return self.tool_call_loop_exit_impl(
+                conversation=conversation, ctx=ctx, **kwargs
+            )
+
+        assert conversation, "Conversation must not be empty"
+        assert isinstance(conversation[-1], AssistantMessage), (
+            "Last message in conversation must be an AssistantMessage"
+        )
+
+        return not bool(conversation[-1].tool_calls)
+
+    async def generate_once(
+        self,
+        agent_state: LLMAgentState,
+        tool_choice: ToolChoice | None = None,
+        ctx: RunContextWrapper[CtxT] | None = None,
+    ) -> Sequence[AssistantMessage]:
+        message_history = agent_state.message_history
+        message_batch = await self.llm.generate_message_batch(
+            message_history, tool_choice=tool_choice
+        )
+        message_history.add_message_batch(message_batch)
+
+        self._print_messages_and_track_usage(message_batch, ctx=ctx)
+
+        return message_batch
+
+    async def run_loop(
+        self,
+        agent_state: LLMAgentState,
+        ctx: RunContextWrapper[CtxT] | None = None,
+    ) -> None:
+        message_history = agent_state.message_history
+        assert message_history.batch_size == 1, (
+            "Batch size must be 1 for tool call loop"
+        )
+
+        tool_choice: ToolChoice
+
+        tool_choice = "none" if self._react_mode else "auto"
+        gen_message_batch = await self.generate_once(
+            agent_state, tool_choice=tool_choice, ctx=ctx
+        )
+
+        turns = 0
+
+        while True:
+            if self._tool_call_loop_exit(
+                message_history.batched_conversations[0], ctx=ctx, num_turns=turns
+            ):
+                return
+            if turns >= self.max_turns:
+                logger.info(
+                    f"Max turns reached: {self.max_turns}. Stopping tool call loop."
+                )
+                return
+
+            msg = gen_message_batch[0]
+            if msg.tool_calls:
+                tool_messages = await self.call_tools(msg.tool_calls, ctx=ctx)
+                message_history.add_messages(tool_messages)
+
+            tool_choice = "none" if (self._react_mode and msg.tool_calls) else "auto"
+            gen_message_batch = await self.generate_once(
+                agent_state, tool_choice=tool_choice, ctx=ctx
+            )
+
+            turns += 1
+
+    async def call_tools(
+        self,
+        calls: Sequence[ToolCall],
+        ctx: RunContextWrapper[CtxT] | None = None,
+    ) -> Sequence[ToolMessage]:
+        corouts: list[Coroutine[Any, Any, BaseModel]] = []
+        for call in calls:
+            tool = self.tools[call.tool_name]
+            args = json.loads(call.tool_arguments)
+            corouts.append(tool(ctx=ctx, **args))
+
+        outs = await asyncio.gather(*corouts)
+
+        tool_messages = [
+            ToolMessage.from_tool_output(out, call, model_id=self.agent_id)
+            for out, call in zip(outs, calls, strict=False)
+        ]
+
+        self._print_messages(tool_messages, ctx=ctx)
+
+        return tool_messages
+
+    def _print_messages(
+        self,
+        message_batch: Sequence[Message],
+        ctx: RunContextWrapper[CtxT] | None = None,
+    ) -> None:
+        if ctx:
+            ctx.printer.print_llm_messages(message_batch, agent_id=self.agent_id)
+
+    def _print_messages_and_track_usage(
+        self,
+        message_batch: Sequence[AssistantMessage],
+        ctx: RunContextWrapper[CtxT] | None = None,
+    ) -> None:
+        if ctx:
+            self._print_messages(message_batch, ctx=ctx)
+            ctx.usage_tracker.update(
+                messages=message_batch, model_name=self.llm.model_name
+            )

grasp_agents/typing/completion.py

@@ -0,0 +1,30 @@
+from abc import ABC
+
+from pydantic import BaseModel
+
+from .message import AssistantMessage
+
+
+class CompletionChoice(BaseModel):
+    # TODO: add fields
+    message: AssistantMessage
+    finish_reason: str | None
+
+
+class CompletionError(BaseModel):
+    message: str
+    metadata: dict[str, str | None] | None = None
+    code: int
+
+
+class Completion(BaseModel, ABC):
+    # TODO: add fields
+    choices: list[CompletionChoice]
+    model_id: str | None = None
+    error: CompletionError | None = None
+
+
+class CompletionChunk(BaseModel):
+    # TODO: add more fields and tool use support (and choices?)
+    delta: str | None = None
+    model_id: str | None = None

grasp_agents/typing/content.py

@@ -0,0 +1,116 @@
+import base64
+import re
+from collections.abc import Iterable
+from enum import StrEnum
+from pathlib import Path
+from typing import Annotated, Any, Literal, TypeAlias
+
+from pydantic import AnyUrl, BaseModel, Field
+
+
+class ContentType(StrEnum):
+    TEXT = "text"
+    IMAGE = "image"
+
+
+ImageDetail: TypeAlias = Literal["low", "high", "auto"]
+
+
+class ImageData(BaseModel):
+    type: Literal["url", "base64"]
+    url: AnyUrl | None = None
+    base64: str | None = None
+
+    # Supported by OpenAI API
+    detail: ImageDetail = "high"
+
+    @classmethod
+    def from_base64(cls, base64_encoding: str, **kwargs: Any) -> "ImageData":
+        return cls(type="base64", base64=base64_encoding, **kwargs)
+
+    @classmethod
+    def from_path(cls, img_path: str | Path, **kwargs: Any) -> "ImageData":
+        image_bytes = Path(img_path).read_bytes()
+        base64_encoding = base64.b64encode(image_bytes).decode("utf-8")
+        return cls(type="base64", base64=base64_encoding, **kwargs)
+
+    @classmethod
+    def from_url(cls, img_url: str, **kwargs: Any) -> "ImageData":
+        return cls(type="url", url=img_url, **kwargs)  # type: ignore
+
+    def to_str(self) -> str:
+        if self.type == "url":
+            return str(self.url)
+        if self.type == "base64":
+            return str(self.base64)
+        raise ValueError(f"Unsupported image data type: {self.type}")
+
+
+class ContentPartText(BaseModel):
+    type: Literal[ContentType.TEXT] = ContentType.TEXT
+    data: str
+
+
+class ContentPartImage(BaseModel):
+    type: Literal[ContentType.IMAGE] = ContentType.IMAGE
+    data: ImageData
+
+
+ContentPart = Annotated[ContentPartText | ContentPartImage, Field(discriminator="type")]
+
+
+class Content(BaseModel):
+    parts: list[ContentPart]
+
+    @classmethod
+    def from_formatted_prompt(
+        cls,
+        prompt_template: str,
+        prompt_args: dict[str, str | ImageData] | None = None,
+    ) -> "Content":
+        prompt_args = prompt_args or {}
+        image_args = {
+            arg_name: arg_val
+            for arg_name, arg_val in prompt_args.items()
+            if isinstance(arg_val, ImageData)
+        }
+        text_args = {
+            arg_name: arg_val
+            for arg_name, arg_val in prompt_args.items()
+            if isinstance(arg_val, (str, int, float))
+        }
+
+        if not image_args:
+            prompt_with_args = prompt_template.format(**text_args)
+            return cls(parts=[ContentPartText(data=prompt_with_args)])
+
+        pattern = r"({})".format("|".join([r"\{" + s + r"\}" for s in image_args]))
+        input_prompt_chunks = re.split(pattern, prompt_template)
+
+        content_parts: list[ContentPart] = []
+        for chunk in input_prompt_chunks:
+            stripped_chunk = chunk.strip(" \n")
+            if re.match(pattern, stripped_chunk):
+                image_data = image_args[stripped_chunk[1:-1]]
+                content_part = ContentPartImage(data=image_data)
+            else:
+                text_data = stripped_chunk.format(**text_args)
+                content_part = ContentPartText(data=text_data)
+            content_parts.append(content_part)
+
+        return cls(parts=content_parts)
+
+    @classmethod
+    def from_text(cls, text: str) -> "Content":
+        return cls(parts=[ContentPartText(data=text)])
+
+    @classmethod
+    def from_content_parts(cls, content_parts: Iterable[str | ImageData]) -> "Content":
+        parts: list[ContentPart] = []
+        for part in content_parts:
+            if isinstance(part, str):
+                parts.append(ContentPartText(data=part))
+            else:
+                parts.append(ContentPartImage(data=part))
+
+        return cls(parts=parts)

grasp_agents/typing/converters.py

@@ -0,0 +1,118 @@
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+from typing import Any
+
+from pydantic import BaseModel
+
+from .completion import Completion, CompletionChunk
+from .content import Content
+from .message import (
+    AssistantMessage,
+    Message,
+    SystemMessage,
+    ToolMessage,
+    UserMessage,
+)
+from .tool import BaseTool, ToolChoice
+
+
+class Converters(ABC):
+    @staticmethod
+    @abstractmethod
+    def to_system_message(system_message: SystemMessage, **kwargs: Any) -> Any:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def from_system_message(raw_message: Any, **kwargs: Any) -> SystemMessage:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def to_user_message(user_message: UserMessage, **kwargs: Any) -> Any:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def from_user_message(raw_message: Any, **kwargs: Any) -> UserMessage:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def to_assistant_message(assistant_message: AssistantMessage, **kwargs: Any) -> Any:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def from_assistant_message(
+        raw_message: Any, raw_usage: Any, **kwargs: Any
+    ) -> AssistantMessage:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def to_tool_message(tool_message: ToolMessage, **kwargs: Any) -> Any:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def from_tool_message(raw_message: Any, **kwargs: Any) -> ToolMessage:
+        pass
+
+    @classmethod
+    def to_message(cls, message: Message, **kwargs: Any) -> Any:
+        if isinstance(message, UserMessage):
+            return cls.to_user_message(message, **kwargs)
+        if isinstance(message, AssistantMessage):
+            return cls.to_assistant_message(message, **kwargs)
+        if isinstance(message, ToolMessage):
+            return cls.to_tool_message(message, **kwargs)
+
+        return cls.to_system_message(message, **kwargs)
+
+    @staticmethod
+    @abstractmethod
+    def to_tool(tool: BaseTool[BaseModel, BaseModel, Any], **kwargs: Any) -> Any:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def to_tool_choice(tool_choice: ToolChoice, **kwargs: Any) -> Any:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def to_content(content: Content, **kwargs: Any) -> Any:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def from_content(raw_content: Any, **kwargs: Any) -> Content:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def to_completion(completion: Completion, **kwargs: Any) -> Any:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def from_completion(raw_completion: Any, **kwargs: Any) -> Completion:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def to_completion_chunk(chunk: CompletionChunk, **kwargs: Any) -> Any:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def from_completion_chunk(raw_chunk: Any, **kwargs: Any) -> CompletionChunk:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def from_completion_chunk_iterator(
+        raw_chunk_iterator: AsyncIterator[Any], **kwargs: Any
+    ) -> AsyncIterator[CompletionChunk]:
+        pass

grasp_agents/typing/io.py

@@ -0,0 +1,32 @@
+from collections.abc import Sequence
+from typing import TypeAlias, TypeVar
+
+from pydantic import BaseModel
+from pydantic.json_schema import SkipJsonSchema
+
+from .content import ImageData
+
+AgentID: TypeAlias = str
+
+
+class AgentPayload(BaseModel):
+    # TODO: do we need conversation?
+    selected_recipient_ids: SkipJsonSchema[Sequence[AgentID] | None] = None
+
+
+class AgentState(BaseModel):
+    pass
+
+
+InT = TypeVar("InT", bound=AgentPayload, contravariant=True)  # noqa: PLC0105
+OutT = TypeVar("OutT", bound=AgentPayload, covariant=True)  # noqa: PLC0105
+StateT = TypeVar("StateT", bound=AgentState, covariant=True)  # noqa: PLC0105
+
+
+class LLMPromptArgs(BaseModel):
+    pass
+
+
+LLMPrompt: TypeAlias = str
+LLMFormattedSystemArgs: TypeAlias = dict[str, str]
+LLMFormattedArgs: TypeAlias = dict[str, str | ImageData]

grasp_agents/typing/message.py

@@ -0,0 +1,130 @@
+from collections.abc import Hashable, Sequence
+from enum import StrEnum
+from typing import Annotated, Literal, TypeAlias
+from uuid import uuid4
+
+from pydantic import BaseModel, Field, NonNegativeFloat, NonNegativeInt
+
+from .content import Content, ImageData
+from .tool import ToolCall
+
+
+class Role(StrEnum):
+    USER = "user"
+    SYSTEM = "system"
+    ASSISTANT = "assistant"
+    TOOL = "tool"
+
+
+class Usage(BaseModel):
+    input_tokens: NonNegativeInt = 0
+    output_tokens: NonNegativeInt = 0
+    reasoning_tokens: NonNegativeInt | None = None
+    cached_tokens: NonNegativeInt | None = None
+    cost: NonNegativeFloat | None = None
+
+    def __add__(self, add_usage: "Usage") -> "Usage":
+        input_tokens = self.input_tokens + add_usage.input_tokens
+        output_tokens = self.output_tokens + add_usage.output_tokens
+        if self.reasoning_tokens is not None or add_usage.reasoning_tokens is not None:
+            reasoning_tokens = (self.reasoning_tokens or 0) + (
+                add_usage.reasoning_tokens or 0
+            )
+        else:
+            reasoning_tokens = None
+
+        if self.cached_tokens is not None or add_usage.cached_tokens is not None:
+            cached_tokens = (self.cached_tokens or 0) + (add_usage.cached_tokens or 0)
+        else:
+            cached_tokens = None
+
+        cost = (
+            (self.cost or 0.0) + add_usage.cost
+            if (add_usage.cost is not None)
+            else None
+        )
+        return Usage(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            reasoning_tokens=reasoning_tokens,
+            cached_tokens=cached_tokens,
+            cost=cost,
+        )
+
+
+class MessageBase(BaseModel):
+    message_id: Hashable = Field(default_factory=lambda: str(uuid4())[:8])
+    model_id: str | None = None
+
+
+class AssistantMessage(MessageBase):
+    role: Literal[Role.ASSISTANT] = Role.ASSISTANT
+    content: str
+    usage: Usage | None = None
+    tool_calls: Sequence[ToolCall] | None = None
+    refusal: str | None = None
+
+
+class UserMessage(MessageBase):
+    role: Literal[Role.USER] = Role.USER
+    content: Content
+
+    @classmethod
+    def from_text(cls, text: str, model_id: str | None = None) -> "UserMessage":
+        return cls(content=Content.from_text(text), model_id=model_id)
+
+    @classmethod
+    def from_formatted_prompt(
+        cls,
+        prompt_template: str,
+        prompt_args: dict[str, str | ImageData] | None = None,
+        model_id: str | None = None,
+    ) -> "UserMessage":
+        content = Content.from_formatted_prompt(
+            prompt_template=prompt_template, prompt_args=prompt_args
+        )
+
+        return cls(content=content, model_id=model_id)
+
+    @classmethod
+    def from_content_parts(
+        cls,
+        content_parts: Sequence[str | ImageData],
+        model_id: str | None = None,
+    ) -> "UserMessage":
+        content = Content.from_content_parts(content_parts)
+
+        return cls(content=content, model_id=model_id)
+
+
+class SystemMessage(MessageBase):
+    role: Literal[Role.SYSTEM] = Role.SYSTEM
+    content: str
+
+
+class ToolMessage(MessageBase):
+    role: Literal[Role.TOOL] = Role.TOOL
+    content: str
+    tool_call_id: str
+
+    @classmethod
+    def from_tool_output(
+        cls,
+        tool_output: BaseModel,
+        tool_call: ToolCall,
+        model_id: str | None = None,
+        indent: int = 2,
+    ) -> "ToolMessage":
+        return cls(
+            content=tool_output.model_dump_json(indent=indent),
+            tool_call_id=tool_call.id,
+            model_id=model_id,
+        )
+
+
+Message = Annotated[
+    AssistantMessage | UserMessage | SystemMessage | ToolMessage,
+    Field(discriminator="role"),
+]
+
+Conversation: TypeAlias = list[Message]

grasp_agents/typing/tool.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, TypeVar
+
+from pydantic import BaseModel
+
+if TYPE_CHECKING:
+    from ..run_context import CtxT, RunContextWrapper
+else:
+    CtxT = TypeVar("CtxT")
+
+    class RunContextWrapper(Generic[CtxT]):
+        """Runtime placeholder so RunContextWrapper[CtxT] works"""
+
+
+ToolInT = TypeVar("ToolInT", bound=BaseModel, contravariant=True)  # noqa: PLC0105
+ToolOutT = TypeVar("ToolOutT", bound=BaseModel, covariant=True)  # noqa: PLC0105
+
+
+class ToolCall(BaseModel):
+    id: str
+    tool_name: str
+    tool_arguments: str
+
+
+class BaseTool(BaseModel, ABC, Generic[ToolInT, ToolOutT, CtxT]):
+    name: str
+    description: str
+    in_schema: type[ToolInT]
+    out_schema: type[ToolOutT]
+
+    # Supported by OpenAI API
+    strict: bool | None = None
+
+    @abstractmethod
+    async def run(
+        self, inp: ToolInT, ctx: RunContextWrapper[CtxT] | None = None
+    ) -> ToolOutT:
+        pass
+
+    async def __call__(
+        self, ctx: RunContextWrapper[CtxT] | None = None, **kwargs: Any
+    ) -> ToolOutT:
+        result = await self.run(self.in_schema(**kwargs), ctx=ctx)
+
+        return self.out_schema.model_validate(result)
+
+
+ToolChoice: TypeAlias = (
+    Literal["none", "auto", "required"] | BaseTool[BaseModel, BaseModel, Any]
+)