grasp_agents 0.2.11__py3-none-any.whl → 0.3.1__py3-none-any.whl

grasp_agents/prompt_builder.py

@@ -1,251 +1,234 @@
-from collections.abc import Sequence
-from copy import deepcopy
-from typing import ClassVar, Generic, Protocol, cast
+import json
+from collections.abc import Mapping, Sequence
+from typing import ClassVar, Generic, Protocol, TypeAlias
 
 from pydantic import BaseModel, TypeAdapter
 
 from .generics_utils import AutoInstanceAttributesMixin
-from .run_context import CtxT, RunContextWrapper
-from .typing.content import ImageData
-from .typing.io import (
-    InT,
-    LLMFormattedArgs,
-    LLMFormattedSystemArgs,
-    LLMPrompt,
-    LLMPromptArgs,
-)
+from .run_context import CtxT, RunContext
+from .typing.content import Content, ImageData
+from .typing.io import InT_contra, LLMPrompt, LLMPromptArgs
 from .typing.message import UserMessage
 
 
-class DummySchema(BaseModel):
-    pass
-
-
-class FormatSystemArgsHandler(Protocol[CtxT]):
+class MakeSystemPromptHandler(Protocol[CtxT]):
     def __call__(
         self,
-        sys_args: LLMPromptArgs,
+        sys_args: LLMPromptArgs | None,
         *,
-        ctx: RunContextWrapper[CtxT] | None,
-    ) -> LLMFormattedSystemArgs: ...
+        ctx: RunContext[CtxT] | None,
+    ) -> str: ...
 
 
-class FormatInputArgsHandler(Protocol[InT, CtxT]):
+class MakeInputContentHandler(Protocol[InT_contra, CtxT]):
     def __call__(
         self,
         *,
-        usr_args: LLMPromptArgs,
-        in_args: InT,
+        in_args: InT_contra | None,
+        usr_args: LLMPromptArgs | None,
         batch_idx: int,
-        ctx: RunContextWrapper[CtxT] | None,
-    ) -> LLMFormattedArgs: ...
+        ctx: RunContext[CtxT] | None,
+    ) -> Content: ...
+
 
+PromptArgumentType: TypeAlias = str | bool | int | ImageData
 
-class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
+
+class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT_contra, CtxT]):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {0: "_in_type"}
 
     def __init__(
         self,
-        agent_id: str,
-        sys_prompt: LLMPrompt | None,
-        in_prompt: LLMPrompt | None,
-        sys_args_schema: type[LLMPromptArgs],
-        usr_args_schema: type[LLMPromptArgs],
+        agent_name: str,
+        sys_prompt_template: LLMPrompt | None,
+        in_prompt_template: LLMPrompt | None,
+        sys_args_schema: type[LLMPromptArgs] | None = None,
+        usr_args_schema: type[LLMPromptArgs] | None = None,
     ):
-        self._in_type: type[InT]
+        self._in_type: type[InT_contra]
         super().__init__()
 
-        self._agent_id = agent_id
-        self.sys_prompt = sys_prompt
-        self.in_prompt = in_prompt
+        self._agent_name = agent_name
+        self.sys_prompt_template = sys_prompt_template
+        self.in_prompt_template = in_prompt_template
         self.sys_args_schema = sys_args_schema
         self.usr_args_schema = usr_args_schema
-        self.format_sys_args_impl: FormatSystemArgsHandler[CtxT] | None = None
-        self.format_in_args_impl: FormatInputArgsHandler[InT, CtxT] | None = None
-
-        self._in_args_type_adapter: TypeAdapter[InT] = TypeAdapter(self._in_type)
-
-    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_in_args(
-        self,
-        *,
-        usr_args: LLMPromptArgs,
-        in_args: InT,
-        batch_idx: int = 0,
-        ctx: RunContextWrapper[CtxT] | None = None,
-    ) -> LLMFormattedArgs:
-        if self.format_in_args_impl:
-            return self.format_in_args_impl(
-                usr_args=usr_args, in_args=in_args, batch_idx=batch_idx, ctx=ctx
-            )
-
-        if not isinstance(in_args, BaseModel) and in_args is not None:
-            raise TypeError(
-                "Cannot apply default formatting to non-BaseModel received arguments."
-            )
-
-        usr_args_ = usr_args
-        in_args_ = DummySchema() if in_args is None else in_args
-
-        usr_args_dump = usr_args_.model_dump(exclude_unset=True)
-        in_args_dump = in_args_.model_dump(exclude={"selected_recipient_ids"})
+        self.make_sys_prompt_impl: MakeSystemPromptHandler[CtxT] | None = None
+        self.make_in_content_impl: MakeInputContentHandler[InT_contra, CtxT] | None = (
+            None
+        )
 
-        return usr_args_dump | in_args_dump
+        self._in_args_type_adapter: TypeAdapter[InT_contra] = TypeAdapter(self._in_type)
 
     def make_sys_prompt(
-        self,
-        sys_args: LLMPromptArgs,
-        *,
-        ctx: RunContextWrapper[CtxT] | None,
-    ) -> LLMPrompt | None:
-        if self.sys_prompt is None:
+        self, sys_args: LLMPromptArgs | None = None, ctx: RunContext[CtxT] | None = None
+    ) -> str | None:
+        if self.sys_prompt_template 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)
+        val_sys_args = sys_args
+        if sys_args is not None:
+            if self.sys_args_schema is not None:
+                val_sys_args = self.sys_args_schema.model_validate(sys_args)
+            else:
+                raise TypeError(
+                    "System prompt template is set, but system arguments schema is not "
+                    "provided."
+                )
 
-    def _usr_messages_from_text(self, text: str) -> list[UserMessage]:
-        return [UserMessage.from_text(text, model_id=self._agent_id)]
+        if self.make_sys_prompt_impl:
+            return self.make_sys_prompt_impl(sys_args=val_sys_args, ctx=ctx)
 
-    def _usr_messages_from_content_parts(
-        self, content_parts: Sequence[str | ImageData]
-    ) -> Sequence[UserMessage]:
-        return [UserMessage.from_content_parts(content_parts, model_id=self._agent_id)]
+        sys_args_dict = (
+            val_sys_args.model_dump(exclude_unset=True) if val_sys_args else {}
+        )
 
-    def _usr_messages_from_in_args(
-        self, in_args_batch: Sequence[InT]
-    ) -> Sequence[UserMessage]:
-        return [
-            UserMessage.from_text(
-                self._in_args_type_adapter.dump_json(
-                    inp,
-                    exclude_unset=True,
-                    indent=2,
-                    exclude={"selected_recipient_ids"},
-                    warnings="error",
-                ).decode("utf-8"),
-                model_id=self._agent_id,
-            )
-            for inp in in_args_batch
-        ]
+        return self.sys_prompt_template.format(**sys_args_dict)
 
-    def _usr_messages_from_prompt_template(
+    def make_in_content(
         self,
-        in_prompt: LLMPrompt,
-        in_args_batch: Sequence[InT] | None = None,
-        usr_args_batch: Sequence[LLMPromptArgs] | None = None,
-        ctx: RunContextWrapper[CtxT] | None = None,
-    ) -> Sequence[UserMessage]:
-        usr_args_batch_, in_args_batch_ = self._align_in_usr_batches(
-            in_args_batch, usr_args_batch
+        *,
+        in_args: InT_contra | None,
+        usr_args: LLMPromptArgs | None,
+        batch_idx: int = 0,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> Content:
+        val_in_args, val_usr_args = self._validate_prompt_args(
+            in_args=in_args, usr_args=usr_args
         )
 
-        val_usr_args_batch_ = [
-            self.usr_args_schema.model_validate(u) for u in usr_args_batch_
-        ]
-        val_in_args_batch_ = [
-            self._in_args_type_adapter.validate_python(inp) for inp in in_args_batch_
-        ]
-
-        formatted_in_args_batch = [
-            self._format_in_args(
-                usr_args=val_usr_args, in_args=val_in_args, batch_idx=i, ctx=ctx
+        if self.make_in_content_impl:
+            return self.make_in_content_impl(
+                in_args=val_in_args, usr_args=val_usr_args, batch_idx=batch_idx, ctx=ctx
             )
-            for i, (val_usr_args, val_in_args) in enumerate(
-                zip(val_usr_args_batch_, val_in_args_batch_, strict=False)
-            )
-        ]
 
-        return [
-            UserMessage.from_formatted_prompt(
-                prompt_template=in_prompt, prompt_args=in_args
+        combined_args = self._combine_args(in_args=val_in_args, usr_args=val_usr_args)
+        if isinstance(combined_args, str):
+            return Content.from_text(combined_args)
+
+        if self.in_prompt_template is not None:
+            return Content.from_formatted_prompt(
+                self.in_prompt_template, prompt_args=combined_args
             )
-            for in_args in formatted_in_args_batch
-        ]
+
+        return Content.from_text(json.dumps(combined_args, indent=2))
 
     def make_user_messages(
         self,
         chat_inputs: LLMPrompt | Sequence[str | ImageData] | None = None,
-        in_args: InT | Sequence[InT] | None = None,
-        usr_args: LLMPromptArgs | Sequence[LLMPromptArgs] | None = None,
-        entry_point: bool = False,
-        ctx: RunContextWrapper[CtxT] | None = None,
+        in_args_batch: Sequence[InT_contra] | None = None,
+        usr_args: LLMPromptArgs | None = None,
+        ctx: RunContext[CtxT] | None = None,
     ) -> Sequence[UserMessage]:
-        # 1) Direct user input (e.g. chat input)
-        if chat_inputs is not None or entry_point:
-            """
-            * If chat 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 chat inputs directly, if provided.
-            """
+        if chat_inputs:
             if isinstance(chat_inputs, LLMPrompt):
                 return self._usr_messages_from_text(chat_inputs)
+            return self._usr_messages_from_content_parts(chat_inputs)
 
-            if isinstance(chat_inputs, Sequence) and chat_inputs:
-                return self._usr_messages_from_content_parts(chat_inputs)
+        in_content_batch = [
+            self.make_in_content(
+                in_args=in_args, usr_args=usr_args, batch_idx=i, ctx=ctx
+            )
+            for i, in_args in enumerate(in_args_batch or [None])
+        ]
+        return [
+            UserMessage(content=in_content, name=self._agent_name)
+            for in_content in in_content_batch
+        ]
 
-        in_args_batch = cast(
-            "Sequence[InT] | None",
-            in_args if (isinstance(in_args, Sequence) or not in_args) else [in_args],
-        )
+    def _usr_messages_from_text(self, text: str) -> list[UserMessage]:
+        return [UserMessage.from_text(text, name=self._agent_name)]
 
-        # 2) No input prompt template + received args → raw JSON messages
-        if self.in_prompt is None and in_args_batch:
-            return self._usr_messages_from_in_args(in_args_batch)
-
-        usr_args_batch = cast(
-            "Sequence[LLMPromptArgs] | None",
-            (
-                usr_args
-                if (isinstance(usr_args, Sequence) or not usr_args)
-                else [usr_args]
-            ),
-        )
+    def _usr_messages_from_content_parts(
+        self, content_parts: Sequence[str | ImageData]
+    ) -> list[UserMessage]:
+        return [UserMessage.from_content_parts(content_parts, name=self._agent_name)]
 
-        # 3) Input prompt template + any args → batch & format
-        if self.in_prompt is not None:
-            if in_args_batch and not isinstance(in_args_batch[0], BaseModel):
+    def _validate_prompt_args(
+        self,
+        *,
+        in_args: InT_contra | None,
+        usr_args: LLMPromptArgs | None,
+    ) -> tuple[InT_contra | None, LLMPromptArgs | None]:
+        val_usr_args = usr_args
+        if usr_args is not None:
+            if self.in_prompt_template is None:
+                raise TypeError(
+                    "Input prompt template is not set, but user arguments are provided."
+                )
+            if self.usr_args_schema is None:
+                raise TypeError(
+                    "User arguments schema is not provided, but user arguments are "
+                    "given."
+                )
+            val_usr_args = self.usr_args_schema.model_validate(usr_args)
+
+        val_in_args = in_args
+        if in_args is not None:
+            val_in_args = self._in_args_type_adapter.validate_python(in_args)
+            if isinstance(val_in_args, BaseModel):
+                _, has_image = self._format_pydantic_prompt_args(val_in_args)
+                if has_image and self.in_prompt_template is None:
+                    raise TypeError(
+                        "BaseModel input arguments contain ImageData, but input prompt "
+                        "template is not set. Cannot format input arguments."
+                    )
+            elif self.in_prompt_template is not None:
                 raise TypeError(
                     "Cannot use the input prompt template with "
-                    "non-BaseModel received arguments."
+                    "non-BaseModel input arguments."
                 )
-            return self._usr_messages_from_prompt_template(
-                in_prompt=self.in_prompt,
-                usr_args_batch=usr_args_batch,
-                in_args_batch=in_args_batch,
-                ctx=ctx,
-            )
 
-        return []
+        return val_in_args, val_usr_args
+
+    @staticmethod
+    def _format_pydantic_prompt_args(
+        inp: BaseModel,
+    ) -> tuple[dict[str, PromptArgumentType], bool]:
+        formatted_args: dict[str, PromptArgumentType] = {}
+        contains_image_data = False
+        for field in type(inp).model_fields:
+            if field == "selected_recipients":
+                continue
+
+            val = getattr(inp, field)
+            if isinstance(val, (int, str, bool)):
+                formatted_args[field] = val
+            elif isinstance(val, ImageData):
+                formatted_args[field] = val
+                contains_image_data = True
+            elif isinstance(val, BaseModel):
+                formatted_args[field] = val.model_dump_json(indent=2, warnings="error")
+            else:
+                raise TypeError(
+                    f"Field '{field}' in prompt arguments must be of type "
+                    "int, str, bool, BaseModel, or ImageData."
+                )
+
+        return formatted_args, contains_image_data
 
-    def _align_in_usr_batches(
+    def _combine_args(
         self,
-        in_args_batch: Sequence[InT] | None = None,
-        usr_args_batch: Sequence[LLMPromptArgs] | None = None,
-    ) -> tuple[Sequence[LLMPromptArgs | DummySchema], Sequence[InT | DummySchema]]:
-        usr_args_batch_ = usr_args_batch or [DummySchema()]
-        in_args_batch_ = in_args_batch or [DummySchema()]
-
-        # Broadcast singleton → match lengths
-        if len(usr_args_batch_) == 1 and len(in_args_batch_) > 1:
-            usr_args_batch_ = [deepcopy(usr_args_batch_[0]) for _ in in_args_batch_]
-        if len(in_args_batch_) == 1 and len(usr_args_batch_) > 1:
-            in_args_batch_ = [deepcopy(in_args_batch_[0]) for _ in usr_args_batch_]
-        if len(usr_args_batch_) != len(in_args_batch_):
-            raise ValueError("User args and received args must have the same length")
-
-        return usr_args_batch_, in_args_batch_
+        *,
+        in_args: InT_contra | None,
+        usr_args: LLMPromptArgs | None,
+    ) -> Mapping[str, PromptArgumentType] | str:
+        fmt_usr_args, _ = (
+            self._format_pydantic_prompt_args(usr_args) if usr_args else ({}, False)
+        )
+
+        if in_args is None:
+            return fmt_usr_args
+
+        if isinstance(in_args, BaseModel):
+            fmt_in_args, _ = self._format_pydantic_prompt_args(in_args)
+            return fmt_in_args | fmt_usr_args
+
+        combined_args_str = self._in_args_type_adapter.dump_json(
+            in_args, indent=2, warnings="error"
+        ).decode("utf-8")
+        if usr_args is not None:
+            fmt_usr_args_str = usr_args.model_dump_json(indent=2, warnings="error")
+            combined_args_str += "\n" + fmt_usr_args_str
+
+        return combined_args_str

grasp_agents/run_context.py

@@ -1,58 +1,39 @@
-from collections.abc import Sequence
-from typing import Any, Generic, TypeAlias, TypeVar
+from collections import defaultdict
+from collections.abc import Mapping
+from typing import Any, Generic, 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,
-    AgentState,
-    InT,
-    LLMPrompt,
-    LLMPromptArgs,
-    OutT,
-    StateT,
-)
+from grasp_agents.typing.completion import Completion
+
+from .printer import ColoringMode, Printer
+from .typing.io import LLMPromptArgs, ProcName
 from .usage_tracker import UsageTracker
 
 
 class RunArgs(BaseModel):
     sys: LLMPromptArgs = Field(default_factory=LLMPromptArgs)
-    usr: LLMPromptArgs | Sequence[LLMPromptArgs] = Field(default_factory=LLMPromptArgs)
+    usr: LLMPromptArgs = 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
-    chat_inputs: LLMPrompt | Sequence[str | ImageData] | None = None
-    sys_prompt: LLMPrompt | None = None
-    in_prompt: LLMPrompt | None = None
-    sys_args: LLMPromptArgs | None = None
-    usr_args: LLMPromptArgs | Sequence[LLMPromptArgs] | None = None
-    in_args: InT | Sequence[InT] | None = None
-    outputs: Sequence[OutT]
-
-    model_config = ConfigDict(extra="forbid", frozen=True)
-
-
-InteractionHistory: TypeAlias = list[InteractionRecord[Any, Any, AgentState]]
-
-
 CtxT = TypeVar("CtxT")
 
 
-class RunContextWrapper(BaseModel, Generic[CtxT]):
-    context: CtxT | None = None
+class RunContext(BaseModel, Generic[CtxT]):
     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)  # type: ignore[valid-type]
+
+    state: CtxT | None = None
+
+    run_args: dict[ProcName, RunArgs] = Field(default_factory=dict)
+    completions: Mapping[ProcName, list[Completion]] = Field(
+        default_factory=lambda: defaultdict(list)
+    )
 
     print_messages: bool = False
+    color_messages_by: ColoringMode = "role"
 
     _usage_tracker: UsageTracker = PrivateAttr()
     _printer: Printer = PrivateAttr()
@@ -60,7 +41,9 @@ class RunContextWrapper(BaseModel, Generic[CtxT]):
     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
+            source_id=self.run_id,
+            print_messages=self.print_messages,
+            color_by=self.color_messages_by,
         )
 
     @property

grasp_agents/typing/completion.py

@@ -1,14 +1,58 @@
-from abc import ABC
+import time
+from typing import Literal, TypeAlias
+from uuid import uuid4
 
-from pydantic import BaseModel
+from openai.types.chat.chat_completion import ChoiceLogprobs as CompletionChoiceLogprobs
+from pydantic import BaseModel, Field, NonNegativeFloat, NonNegativeInt
 
 from .message import AssistantMessage
 
+FinishReason: TypeAlias = Literal[
+    "stop", "length", "tool_calls", "content_filter", "function_call"
+]
+
+
+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 CompletionChoice(BaseModel):
-    # TODO: add fields
     message: AssistantMessage
-    finish_reason: str | None
+    finish_reason: FinishReason | None
+    index: int
+    logprobs: CompletionChoiceLogprobs | None = None
 
 
 class CompletionError(BaseModel):
@@ -17,14 +61,16 @@ class CompletionError(BaseModel):
     code: int
 
 
-class Completion(BaseModel, ABC):
-    # TODO: add fields
+class Completion(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4())[:8])
+    created: int = Field(default_factory=lambda: int(time.time()))
+    model: str
+    name: str | None = None
+    system_fingerprint: str | None = None
     choices: list[CompletionChoice]
-    model_id: str | None = None
+    usage: Usage | 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
+    @property
+    def messages(self) -> list[AssistantMessage]:
+        return [choice.message for choice in self.choices if choice.message]