grasp_agents 0.1.5__py3-none-any.whl

grasp_agents/openai/openai_llm.py

@@ -0,0 +1,179 @@
+import logging
+from collections.abc import Iterable
+from copy import deepcopy
+from typing import Any, Literal
+
+from openai import AsyncOpenAI
+from openai._types import NOT_GIVEN  # noqa: PLC2701 # type: ignore[import]
+from pydantic import BaseModel
+
+from ..data_retrieval.rate_limiter_chunked import RateLimiterC
+
+from ..cloud_llm import APIProvider, CloudLLM, CloudLLMSettings
+from ..http_client import AsyncHTTPClientParams
+from ..typing.message import AssistantMessage, Conversation
+from ..typing.tool import BaseTool
+from . import (
+    ChatCompletion,
+    ChatCompletionAsyncStream,  # type: ignore[import]
+    ChatCompletionChunk,
+    ChatCompletionMessageParam,
+    ChatCompletionPredictionContentParam,
+    ChatCompletionStreamOptionsParam,
+    ChatCompletionToolChoiceOptionParam,
+    ChatCompletionToolParam,
+    ParsedChatCompletion,
+    # ResponseFormatJSONObject,
+    # ResponseFormatJSONSchema,
+    # ResponseFormatText,
+)
+from .converters import OpenAIConverters
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAILLMSettings(CloudLLMSettings, total=False):
+    reasoning_effort: Literal["low", "medium", "high"] | None
+
+    parallel_tool_calls: bool
+
+    # response_format: (
+    #     ResponseFormatText | ResponseFormatJSONSchema | ResponseFormatJSONObject
+    # )
+
+    modalities: list[Literal["text", "audio"]] | None
+
+    frequency_penalty: float | None
+    presence_penalty: float | None
+    logit_bias: dict[str, int] | None
+    stop: str | list[str] | None
+    logprobs: bool | None
+    top_logprobs: int | None
+    n: int | None
+
+    prediction: ChatCompletionPredictionContentParam | None
+
+    stream_options: ChatCompletionStreamOptionsParam | None
+
+    metadata: dict[str, str] | None
+    store: bool | None
+    user: str
+
+
+class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
+    def __init__(
+        self,
+        # Base LLM args
+        model_name: str,
+        model_id: str | None = None,
+        llm_settings: OpenAILLMSettings | None = None,
+        tools: list[BaseTool[BaseModel, BaseModel, Any]] | None = None,
+        response_format: type | None = None,
+        # Connection settings
+        api_provider: APIProvider = "openai",
+        async_http_client_params: (
+            dict[str, Any] | AsyncHTTPClientParams | None
+        ) = None,
+        async_openai_client_params: dict[str, Any] | None = None,
+        # Rate limiting
+        rate_limiter: (RateLimiterC[Conversation, AssistantMessage] | None) = None,
+        rate_limiter_rpm: float | None = None,
+        rate_limiter_chunk_size: int = 1000,
+        rate_limiter_max_concurrency: int = 300,
+        # Retries
+        num_generation_retries: int = 0,
+        # Disable tqdm for batch processing
+        no_tqdm: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            model_name=model_name,
+            model_id=model_id,
+            llm_settings=llm_settings,
+            converters=OpenAIConverters(),
+            tools=tools,
+            response_format=response_format,
+            api_provider=api_provider,
+            async_http_client_params=async_http_client_params,
+            rate_limiter=rate_limiter,
+            rate_limiter_rpm=rate_limiter_rpm,
+            rate_limiter_chunk_size=rate_limiter_chunk_size,
+            rate_limiter_max_concurrency=rate_limiter_max_concurrency,
+            num_generation_retries=num_generation_retries,
+            no_tqdm=no_tqdm,
+            **kwargs,
+        )
+
+        async_openai_client_params_ = deepcopy(async_openai_client_params or {})
+        if self._async_http_client is not None:
+            async_openai_client_params_["http_client"] = self._async_http_client
+
+        # TODO: context manager for async client
+        self._client: AsyncOpenAI = AsyncOpenAI(
+            base_url=self._base_url,
+            api_key=self._api_key,
+            **async_openai_client_params_,
+            # timeout=10.0,
+            # max_retries=3,
+        )
+
+    async def _get_completion(
+        self,
+        api_messages: Iterable[ChatCompletionMessageParam],
+        api_tools: list[ChatCompletionToolParam] | None = None,
+        api_tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
+        **api_llm_settings: Any,
+    ) -> ChatCompletion:
+        tools = api_tools or NOT_GIVEN
+        tool_choice = api_tool_choice or NOT_GIVEN
+
+        return await self._client.chat.completions.create(
+            model=self._model_name,
+            messages=api_messages,
+            tools=tools,
+            tool_choice=tool_choice,
+            stream=False,
+            **api_llm_settings,
+        )
+
+    async def _get_parsed_completion(
+        self,
+        api_messages: Iterable[ChatCompletionMessageParam],
+        api_tools: list[ChatCompletionToolParam] | None = None,
+        api_tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
+        api_response_format: type | None = None,
+        **api_llm_settings: Any,
+    ) -> ParsedChatCompletion[Any]:
+        tools = api_tools or NOT_GIVEN
+        tool_choice = api_tool_choice or NOT_GIVEN
+        response_format = api_response_format or NOT_GIVEN
+
+        return await self._client.beta.chat.completions.parse(
+            model=self._model_name,
+            messages=api_messages,
+            tools=tools,
+            tool_choice=tool_choice,
+            response_format=response_format,  # type: ignore[arg-type]
+            **api_llm_settings,
+        )
+
+    async def _get_completion_stream(
+        self,
+        api_messages: Iterable[ChatCompletionMessageParam],
+        api_tools: list[ChatCompletionToolParam] | None = None,
+        api_tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
+        **api_llm_settings: Any,
+    ) -> ChatCompletionAsyncStream[ChatCompletionChunk]:
+        assert not api_tools, "Tool use is not supported in streaming mode"
+
+        tools = api_tools or NOT_GIVEN
+        tool_choice = api_tool_choice or NOT_GIVEN
+
+        return await self._client.chat.completions.create(
+            model=self._model_name,
+            messages=api_messages,
+            tools=tools,
+            tool_choice=tool_choice,
+            stream=True,
+            **api_llm_settings,
+        )

grasp_agents/openai/tool_converters.py

@@ -0,0 +1,37 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+from ..typing.tool import BaseTool, ToolChoice
+from . import (
+    ChatCompletionFunctionDefinition,
+    ChatCompletionNamedToolChoiceFunction,
+    ChatCompletionNamedToolChoiceParam,
+    ChatCompletionToolChoiceOptionParam,
+    ChatCompletionToolParam,
+)
+
+
+def to_api_tool(
+    tool: BaseTool[BaseModel, BaseModel, Any],
+) -> ChatCompletionToolParam:
+    return ChatCompletionToolParam(
+        type="function",
+        function=ChatCompletionFunctionDefinition(
+            name=tool.name,
+            description=tool.description,
+            parameters=tool.in_schema.model_json_schema(),
+            strict=tool.strict,
+        ),
+    )
+
+
+def to_api_tool_choice(
+    tool_choice: ToolChoice,
+) -> ChatCompletionToolChoiceOptionParam:
+    if isinstance(tool_choice, BaseTool):
+        return ChatCompletionNamedToolChoiceParam(
+            type="function",
+            function=ChatCompletionNamedToolChoiceFunction(name=tool_choice.name),
+        )
+    return tool_choice

grasp_agents/printer.py

@@ -0,0 +1,156 @@
+import hashlib
+import json
+import logging
+from collections.abc import Mapping, Sequence
+from typing import Literal, TypeAlias
+
+from termcolor._types import Color  # type: ignore[import]
+
+from .typing.content import Content, ContentPartText
+from .typing.message import AssistantMessage, Message, Role, ToolMessage
+
+logger = logging.getLogger(__name__)
+
+
+ColoringMode: TypeAlias = Literal["agent_id", "role"]
+
+ROLE_TO_COLOR: Mapping[Role, Color] = {
+    Role.SYSTEM: "magenta",
+    Role.USER: "green",
+    Role.ASSISTANT: "light_blue",
+    Role.TOOL: "light_cyan",
+}
+
+AVAILABLE_COLORS: list[Color] = [
+    "magenta",
+    "green",
+    "light_blue",
+    "light_cyan",
+    "yellow",
+    "blue",
+    "red",
+]
+
+
+class Printer:
+    def __init__(
+        self,
+        source_id: str,
+        color_by: ColoringMode = "role",
+        msg_trunc_len: int = 20000,
+        print_messages: bool = False,
+    ) -> None:
+        self.source_id = source_id
+        self.color_by = color_by
+        self.msg_trunc_len = msg_trunc_len
+        self.print_messages = print_messages
+
+    @staticmethod
+    def get_role_color(role: Role) -> Color:
+        return ROLE_TO_COLOR[role]
+
+    @staticmethod
+    def get_agent_color(agent_id: str) -> Color:
+        idx = int(
+            hashlib.md5(agent_id.encode()).hexdigest(),  # noqa :S324
+            16,
+        ) % len(AVAILABLE_COLORS)
+
+        return AVAILABLE_COLORS[idx]
+
+    @staticmethod
+    def content_to_str(content: Content | str, role: Role) -> str:
+        if role == Role.USER and isinstance(content, Content):
+            content_str_parts: list[str] = []
+            for content_part in content.parts:
+                if isinstance(content_part, ContentPartText):
+                    content_str_parts.append(content_part.data.strip(" \n"))
+                elif content_part.data.type == "url":
+                    content_str_parts.append(str(content_part.data.url))
+                elif content_part.data.type == "base64":
+                    content_str_parts.append("<ENCODED_IMAGE>")
+            return "\n".join(content_str_parts)
+
+        assert isinstance(content, str)
+
+        return content.strip(" \n")
+
+    @staticmethod
+    def truncate_content_str(content_str: str, trunc_len: int = 2000) -> str:
+        if len(content_str) > trunc_len:
+            return content_str[:trunc_len] + "[...]"
+
+        return content_str
+
+    def print_llm_message(self, message: Message, agent_id: str) -> None:
+        if not self.print_messages:
+            return
+
+        role = message.role
+        usage = message.usage if isinstance(message, AssistantMessage) else None
+        content_str = self.content_to_str(message.content, message.role)
+
+        if self.color_by == "agent_id":
+            color = self.get_agent_color(agent_id)
+        elif self.color_by == "role":
+            color = self.get_role_color(role)
+
+        log_kwargs = {"extra": {"color": color}}  # type: ignore
+
+        # Print message title
+
+        out = f"\n<{agent_id}>"
+        out += "[" + role.value.upper() + "]"
+
+        if isinstance(message, ToolMessage):
+            out += f"\nTool call ID: {message.tool_call_id}"
+
+        # Print message content
+
+        if content_str:
+            try:
+                content_str = json.dumps(json.loads(content_str), indent=2)
+            except Exception:
+                pass
+            content_str_truncated = self.truncate_content_str(
+                content_str, trunc_len=self.msg_trunc_len
+            )
+            out += f"\n{content_str_truncated}"
+
+        logger.debug(out, **log_kwargs)  # type: ignore
+
+        # Print tool calls
+
+        if isinstance(message, AssistantMessage) and message.tool_calls is not None:
+            for tool_call in message.tool_calls:
+                if self.color_by == "agent_id":
+                    tool_color = self.get_agent_color(agent_id=agent_id)
+                elif self.color_by == "role":
+                    tool_color = self.get_role_color(role=Role.TOOL)
+                logger.debug(
+                    f"\n[TOOL_CALL]<{agent_id}>\n{tool_call.tool_name} "
+                    f"| {tool_call.id}\n{tool_call.tool_arguments}",
+                    extra={"color": tool_color},  # type: ignore
+                )
+
+        # Print usage
+
+        if usage is not None:
+            usage_str = (
+                f"I/O/(R)/(C) tokens: {usage.input_tokens}/{usage.output_tokens}"
+            )
+            if usage.reasoning_tokens is not None:
+                usage_str += f"/{usage.reasoning_tokens}"
+            if usage.cached_tokens is not None:
+                usage_str += f"/{usage.cached_tokens}"
+            logger.debug(
+                f"\n------------------------------------\n{usage_str}",
+                **log_kwargs,  # type: ignore
+            )
+
+    def print_llm_messages(self, messages: Sequence[Message], agent_id: str) -> None:
+        if not self.print_messages:
+            return
+
+        for message in messages:
+            self.print_llm_message(message, agent_id)

grasp_agents/prompt_builder.py

@@ -0,0 +1,204 @@
+from collections.abc import Sequence
+from copy import deepcopy
+from typing import Generic, Protocol
+
+from .run_context import CtxT, RunContextWrapper, UserRunArgs
+from .typing.content import ImageData
+from .typing.io import (
+    AgentPayload,
+    InT,
+    LLMFormattedArgs,
+    LLMFormattedSystemArgs,
+    LLMPrompt,
+    LLMPromptArgs,
+)
+from .typing.message import UserMessage
+
+
+class FormatSystemArgsHandler(Protocol[CtxT]):
+    def __call__(
+        self,
+        sys_args: LLMPromptArgs,
+        *,
+        ctx: RunContextWrapper[CtxT] | None,
+    ) -> LLMFormattedSystemArgs: ...
+
+
+class FormatInputArgsHandler(Protocol[InT, CtxT]):
+    def __call__(
+        self,
+        usr_args: LLMPromptArgs,
+        rcv_args: InT,
+        ctx: RunContextWrapper[CtxT] | None,
+    ) -> LLMFormattedArgs: ...
+
+
+class PromptBuilder(Generic[InT, CtxT]):
+    def __init__(
+        self,
+        agent_id: str,
+        sys_prompt: LLMPrompt | None,
+        inp_prompt: LLMPrompt | None,
+        sys_args_schema: type[LLMPromptArgs],
+        usr_args_schema: type[LLMPromptArgs],
+        rcv_args_schema: type[InT],
+    ):
+        self._agent_id = agent_id
+        self.sys_prompt = sys_prompt
+        self.inp_prompt = inp_prompt
+        self.sys_args_schema = sys_args_schema
+        self.usr_args_schema = usr_args_schema
+        self.rcv_args_schema = rcv_args_schema
+        self.format_sys_args_impl: FormatSystemArgsHandler[CtxT] | None = None
+        self.format_inp_args_impl: FormatInputArgsHandler[InT, CtxT] | None = None
+
+    def _format_sys_args(
+        self,
+        sys_args: LLMPromptArgs,
+        ctx: RunContextWrapper[CtxT] | None = None,
+    ) -> LLMFormattedSystemArgs:
+        if self.format_sys_args_impl:
+            return self.format_sys_args_impl(sys_args=sys_args, ctx=ctx)
+
+        return sys_args.model_dump(exclude_unset=True)
+
+    def _format_inp_args(
+        self,
+        usr_args: LLMPromptArgs,
+        rcv_args: InT,
+        ctx: RunContextWrapper[CtxT] | None = None,
+    ) -> LLMFormattedArgs:
+        if self.format_inp_args_impl:
+            return self.format_inp_args_impl(
+                usr_args=usr_args, rcv_args=rcv_args, ctx=ctx
+            )
+
+        return usr_args.model_dump(exclude_unset=True) | rcv_args.model_dump(
+            exclude_unset=True, exclude={"selected_recipient_ids"}
+        )
+
+    def make_sys_prompt(
+        self,
+        sys_args: LLMPromptArgs,
+        ctx: RunContextWrapper[CtxT] | None,
+    ) -> LLMPrompt | None:
+        if self.sys_prompt is None:
+            return None
+        val_sys_args = self.sys_args_schema.model_validate(sys_args)
+        fmt_sys_args = self._format_sys_args(val_sys_args, ctx=ctx)
+
+        return self.sys_prompt.format(**fmt_sys_args)
+
+    def _usr_messages_from_text(self, text: str) -> list[UserMessage]:
+        return [UserMessage.from_text(text, model_id=self._agent_id)]
+
+    def _usr_messages_from_content_parts(
+        self, content_parts: list[str | ImageData]
+    ) -> list[UserMessage]:
+        return [UserMessage.from_content_parts(content_parts, model_id=self._agent_id)]
+
+    def _usr_messages_from_rcv_args(
+        self, rcv_args_batch: Sequence[InT]
+    ) -> list[UserMessage]:
+        val_rcv_args_batch = [
+            self.rcv_args_schema.model_validate(rcv) for rcv in rcv_args_batch
+        ]
+
+        return [
+            UserMessage.from_text(
+                rcv.model_dump_json(
+                    exclude_unset=True,
+                    indent=2,
+                    exclude={"selected_recipient_ids"},
+                ),
+                model_id=self._agent_id,
+            )
+            for rcv in val_rcv_args_batch
+        ]
+
+    def _usr_messages_from_prompt_template(
+        self,
+        inp_prompt: LLMPrompt,
+        usr_args: UserRunArgs | None = None,
+        rcv_args_batch: Sequence[InT] | None = None,
+        ctx: RunContextWrapper[CtxT] | None = None,
+    ) -> Sequence[UserMessage]:
+        usr_args_batch, rcv_args_batch = self._make_batched(usr_args, rcv_args_batch)
+        val_usr_args_batch = [
+            self.usr_args_schema.model_validate(u) for u in usr_args_batch
+        ]
+        val_rcv_args_batch = [
+            self.rcv_args_schema.model_validate(r) for r in rcv_args_batch
+        ]
+        formatted_inp_args_batch = [
+            self._format_inp_args(usr_args=val_usr_args, rcv_args=val_rcv_args, ctx=ctx)
+            for val_usr_args, val_rcv_args in zip(
+                val_usr_args_batch, val_rcv_args_batch, strict=False
+            )
+        ]
+
+        return [
+            UserMessage.from_formatted_prompt(
+                prompt_template=inp_prompt, prompt_args=inp_args
+            )
+            for inp_args in formatted_inp_args_batch
+        ]
+
+    def make_user_messages(
+        self,
+        inp_items: LLMPrompt | list[str | ImageData] | None = None,
+        usr_args: UserRunArgs | None = None,
+        rcv_args_batch: Sequence[InT] | None = None,
+        entry_point: bool = False,
+        ctx: RunContextWrapper[CtxT] | None = None,
+    ) -> Sequence[UserMessage]:
+        # 1) Direct user input (e.g. chat input)
+        if inp_items is not None or entry_point:
+            """
+            * If user inputs are provided, use them instead of the predefined
+                input prompt template
+            * In a multi-agent system, the predefined input prompt is used to
+                construct agent inputs using the combination of received
+                and user arguments.
+                However, the first agent run (entry point) has no received
+                messages, so we use the user inputs directly, if provided.
+            """
+            if isinstance(inp_items, LLMPrompt):
+                return self._usr_messages_from_text(inp_items)
+            if isinstance(inp_items, list) and inp_items:
+                return self._usr_messages_from_content_parts(inp_items)
+            return []
+
+        # 2) No input prompt template + received args → raw JSON messages
+        if self.inp_prompt is None and rcv_args_batch:
+            return self._usr_messages_from_rcv_args(rcv_args_batch)
+
+        # 3) Input prompt template + any args → batch & format
+        if self.inp_prompt is not None:
+            return self._usr_messages_from_prompt_template(
+                inp_prompt=self.inp_prompt,
+                usr_args=usr_args,
+                rcv_args_batch=rcv_args_batch,
+                ctx=ctx,
+            )
+        return []
+
+    def _make_batched(
+        self,
+        usr_args: UserRunArgs | None = None,
+        rcv_args_batch: Sequence[InT] | None = None,
+    ) -> tuple[Sequence[LLMPromptArgs], Sequence[InT]]:
+        usr_args_batch_ = (
+            usr_args if isinstance(usr_args, list) else [usr_args or LLMPromptArgs()]
+        )
+        rcv_args_batch_ = rcv_args_batch or [AgentPayload()]
+
+        # Broadcast singleton → match lengths
+        if len(usr_args_batch_) == 1 and len(rcv_args_batch_) > 1:
+            usr_args_batch_ = [deepcopy(usr_args_batch_[0]) for _ in rcv_args_batch_]
+        if len(rcv_args_batch_) == 1 and len(usr_args_batch_) > 1:
+            rcv_args_batch_ = [deepcopy(rcv_args_batch_[0]) for _ in usr_args_batch_]
+        if len(usr_args_batch_) != len(rcv_args_batch_):
+            raise ValueError("User args and received args must have the same length")
+
+        return usr_args_batch_, rcv_args_batch_  # type: ignore

grasp_agents/run_context.py

@@ -0,0 +1,90 @@
+from collections.abc import Sequence
+from typing import Any, Generic, TypeAlias, TypeVar
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+
+from .printer import Printer
+from .typing.content import ImageData
+from .typing.io import (
+    AgentID,
+    AgentPayload,
+    AgentState,
+    InT,
+    LLMPrompt,
+    LLMPromptArgs,
+    OutT,
+    StateT,
+)
+from .usage_tracker import UsageTracker
+
+SystemRunArgs: TypeAlias = LLMPromptArgs
+UserRunArgs: TypeAlias = LLMPromptArgs | list[LLMPromptArgs]
+
+
+class RunArgs(BaseModel):
+    sys: SystemRunArgs = Field(default_factory=LLMPromptArgs)
+    usr: UserRunArgs = Field(default_factory=LLMPromptArgs)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class InteractionRecord(BaseModel, Generic[InT, OutT, StateT]):
+    source_id: str
+    recipient_ids: Sequence[AgentID]
+    state: StateT
+    inp_items: LLMPrompt | list[str | ImageData] | None = None
+    sys_prompt: LLMPrompt | None = None
+    inp_prompt: LLMPrompt | None = None
+    sys_args: SystemRunArgs | None = None
+    usr_args: UserRunArgs | None = None
+    rcv_args: Sequence[InT] | None = None
+    outputs: Sequence[OutT]
+
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+
+InteractionHistory: TypeAlias = list[
+    InteractionRecord[AgentPayload, AgentPayload, AgentState]
+]
+
+
+CtxT = TypeVar("CtxT")
+
+
+class RunContextWrapper(BaseModel, Generic[CtxT]):
+    context: CtxT | None = None
+    run_id: str = Field(default_factory=lambda: str(uuid4())[:8], frozen=True)
+    run_args: dict[AgentID, RunArgs] = Field(default_factory=dict)
+    interaction_history: InteractionHistory = Field(default_factory=list)
+
+    print_messages: bool = False
+
+    _usage_tracker: UsageTracker = PrivateAttr()
+    _printer: Printer = PrivateAttr()
+
+    # usage_tracker: Optional[UsageTracker] = None
+    # printer: Optional[Printer] = None
+
+    # @model_validator(mode="after")
+    # def set_usage_tracker_and_printer(self) -> "RunContextWrapper":
+    #     self.usage_tracker = UsageTracker(source_id=self.run_id)
+    #     self.printer = Printer(source_id=self.run_id)
+
+    #     return self
+
+    def model_post_init(self, context: Any) -> None:  # noqa: ARG002
+        self._usage_tracker = UsageTracker(source_id=self.run_id)
+        self._printer = Printer(
+            source_id=self.run_id, print_messages=self.print_messages
+        )
+
+    @property
+    def usage_tracker(self) -> UsageTracker:
+        return self._usage_tracker
+
+    @property
+    def printer(self) -> Printer:
+        return self._printer
+
+    model_config = ConfigDict(extra="forbid")