grasp_agents 0.5.10__py3-none-any.whl → 0.5.12__py3-none-any.whl

grasp_agents/printer.py

@@ -72,50 +72,61 @@ ColoringMode: TypeAlias = Literal["agent", "role"]
 CompletionBlockType: TypeAlias = Literal["response", "thinking", "tool_call"]
 
 
-class Printer:
-    def __init__(
-        self, color_by: ColoringMode = "role", msg_trunc_len: int = 20000
-    ) -> None:
-        self.color_by = color_by
-        self.msg_trunc_len = msg_trunc_len
-        self._current_message: str = ""
+def stream_colored_text(new_colored_text: str) -> None:
+    sys.stdout.write(new_colored_text)
+    sys.stdout.flush()
 
-    @staticmethod
-    def get_role_color(role: Role) -> Color:
-        return ROLE_TO_COLOR[role]
 
-    @staticmethod
-    def get_agent_color(agent_name: str) -> Color:
+def get_color(
+    agent_name: str = "", role: Role = Role.ASSISTANT, color_by: ColoringMode = "role"
+) -> Color:
+    if color_by == "agent":
         idx = int(
             hashlib.md5(agent_name.encode()).hexdigest(),  # noqa :S324
             16,
         ) % len(AVAILABLE_COLORS)
 
         return AVAILABLE_COLORS[idx]
+    return ROLE_TO_COLOR[role]
+
 
-    @staticmethod
-    def content_to_str(content: Content | str | None, 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)
+def content_to_str(content: Content | str | None, 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 | None)
+    assert isinstance(content, str | None)
 
-        return (content or "").strip(" \n")
+    return (content or "").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 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
+
+
+class Printer:
+    def __init__(
+        self,
+        color_by: ColoringMode = "role",
+        msg_trunc_len: int = 20000,
+        output_to: Literal["print", "log"] = "print",
+        logging_level: Literal["info", "debug", "warning", "error"] = "info",
+    ) -> None:
+        self.color_by: ColoringMode = color_by
+        self.msg_trunc_len = msg_trunc_len
+        self._current_message: str = ""
+        self._logging_level = logging_level
+        self._output_to = output_to
 
     def print_message(
         self,
@@ -128,11 +139,8 @@ class Printer:
             raise ValueError(
                 "Usage information can only be printed for AssistantMessage"
             )
-
-        color = (
-            self.get_agent_color(agent_name)
-            if self.color_by == "agent"
-            else self.get_role_color(message.role)
+        color = get_color(
+            agent_name=agent_name, role=message.role, color_by=self.color_by
         )
         log_kwargs = {"extra": {"color": color}}
 
@@ -144,13 +152,13 @@ class Printer:
             out += f"<thinking>\n{thinking}\n</thinking>\n"
 
         # Content
-        content = self.content_to_str(message.content or "", message.role)
+        content = content_to_str(message.content or "", message.role)
         if content:
             try:
                 content = json.dumps(json.loads(content), indent=2)
             except Exception:
                 pass
-            content = self.truncate_content_str(content, trunc_len=self.msg_trunc_len)
+            content = truncate_content_str(content, trunc_len=self.msg_trunc_len)
             if isinstance(message, SystemMessage):
                 out += f"<system>\n{content}\n</system>\n"
             elif isinstance(message, UserMessage):
@@ -176,7 +184,17 @@ class Printer:
 
             out += f"\n------------------------------------\n{usage_str}\n"
 
-        logger.debug(out, **log_kwargs)  # type: ignore
+        if self._output_to == "log":
+            if self._logging_level == "debug":
+                logger.debug(out, **log_kwargs)  # type: ignore
+            elif self._logging_level == "info":
+                logger.info(out, **log_kwargs)  # type: ignore
+            elif self._logging_level == "warning":
+                logger.warning(out, **log_kwargs)  # type: ignore
+            else:
+                logger.error(out, **log_kwargs)  # type: ignore
+        else:
+            stream_colored_text(colored(out + "\n", color))
 
     def print_messages(
         self,
@@ -193,28 +211,101 @@ class Printer:
             )
 
 
-def stream_text(new_text: str, color: Color) -> None:
-    sys.stdout.write(colored(new_text, color))
-    sys.stdout.flush()
-
-
 async def print_event_stream(
     event_generator: AsyncIterator[Event[Any]],
     color_by: ColoringMode = "role",
     trunc_len: int = 10000,
 ) -> AsyncIterator[Event[Any]]:
-    color = Printer.get_role_color(Role.ASSISTANT)
+    def _make_chunk_text(event: CompletionChunkEvent[CompletionChunk]) -> str:
+        color = get_color(
+            agent_name=event.proc_name or "", role=Role.ASSISTANT, color_by=color_by
+        )
+        text = ""
+
+        if isinstance(event, CompletionStartEvent):
+            text += f"\n<{event.proc_name}> [{event.call_id}]\n"
+        elif isinstance(event, ThinkingStartEvent):
+            text += "<thinking>\n"
+        elif isinstance(event, ResponseStartEvent):
+            text += "<response>\n"
+        elif isinstance(event, ToolCallStartEvent):
+            tc = event.data.tool_call
+            text += f"<tool call> {tc.tool_name} [{tc.id}]\n"
+        elif isinstance(event, AnnotationsStartEvent):
+            text += "<annotations>\n"
+
+        # if isinstance(event, CompletionEndEvent):
+        #     text += f"\n</{event.proc_name}>\n"
+        if isinstance(event, ThinkingEndEvent):
+            text += "\n</thinking>\n"
+        elif isinstance(event, ResponseEndEvent):
+            text += "\n</response>\n"
+        elif isinstance(event, ToolCallEndEvent):
+            text += "\n</tool call>\n"
+        elif isinstance(event, AnnotationsEndEvent):
+            text += "\n</annotations>\n"
+
+        if isinstance(event, ThinkingChunkEvent):
+            thinking = event.data.thinking
+            if isinstance(thinking, str):
+                text += thinking
+            else:
+                text = "\n".join(
+                    [block.get("thinking", "[redacted]") for block in thinking]
+                )
 
-    def _get_color(event: Event[Any], role: Role = Role.ASSISTANT) -> Color:
-        if color_by == "agent":
-            return Printer.get_agent_color(event.proc_name or "")
-        return Printer.get_role_color(role)
+        if isinstance(event, ResponseChunkEvent):
+            text += event.data.response
 
-    def _print_packet(
-        event: ProcPacketOutputEvent | WorkflowResultEvent | RunResultEvent,
-    ) -> None:
-        color = _get_color(event, Role.ASSISTANT)
+        if isinstance(event, ToolCallChunkEvent):
+            text += event.data.tool_call.tool_arguments or ""
 
+        if isinstance(event, AnnotationsChunkEvent):
+            text += "\n".join(
+                [
+                    json.dumps(annotation, indent=2)
+                    for annotation in event.data.annotations
+                ]
+            )
+
+        return colored(text, color)
+
+    def _make_message_text(
+        event: MessageEvent[SystemMessage | UserMessage | ToolMessage],
+    ) -> str:
+        message = event.data
+        role = message.role
+        content = content_to_str(message.content, role=role)
+
+        color = get_color(
+            agent_name=event.proc_name or "", role=role, color_by=color_by
+        )
+        text = f"\n<{event.proc_name}> [{event.call_id}]\n"
+
+        if isinstance(event, (SystemMessageEvent, UserMessageEvent)):
+            content = truncate_content_str(content, trunc_len=trunc_len)
+
+        if isinstance(event, SystemMessageEvent):
+            text += f"<system>\n{content}\n</system>\n"
+
+        elif isinstance(event, UserMessageEvent):
+            text += f"<input>\n{content}\n</input>\n"
+
+        elif isinstance(event, ToolMessageEvent):
+            message = event.data
+            try:
+                content = json.dumps(json.loads(content), indent=2)
+            except Exception:
+                pass
+            text += (
+                f"<tool result> [{message.tool_call_id}]\n{content}\n</tool result>\n"
+            )
+
+        return colored(text, color)
+
+    def _make_packet_text(
+        event: ProcPacketOutputEvent | WorkflowResultEvent | RunResultEvent,
+    ) -> str:
         if isinstance(event, WorkflowResultEvent):
             src = "workflow"
         elif isinstance(event, RunResultEvent):
@@ -222,6 +313,9 @@ async def print_event_stream(
         else:
             src = "processor"
 
+        color = get_color(
+            agent_name=event.proc_name or "", role=Role.ASSISTANT, color_by=color_by
+        )
         text = f"\n<{event.proc_name}> [{event.call_id}]\n"
 
         if event.data.payloads:
@@ -237,7 +331,9 @@ async def print_event_stream(
                 text += f"{p_str}\n"
             text += f"</{src} output>\n"
 
-        stream_text(text, color)
+        return colored(text, color)
+
+    # ------ Wrap event generator -------
 
     async for event in event_generator:
         yield event
@@ -245,91 +341,12 @@ async def print_event_stream(
         if isinstance(event, CompletionChunkEvent) and isinstance(
             event.data, CompletionChunk
         ):
-            color = _get_color(event, Role.ASSISTANT)
-
-            text = ""
-
-            if isinstance(event, CompletionStartEvent):
-                text += f"\n<{event.proc_name}> [{event.call_id}]\n"
-            elif isinstance(event, ThinkingStartEvent):
-                text += "<thinking>\n"
-            elif isinstance(event, ResponseStartEvent):
-                text += "<response>\n"
-            elif isinstance(event, ToolCallStartEvent):
-                tc = event.data.tool_call
-                text += f"<tool call> {tc.tool_name} [{tc.id}]\n"
-            elif isinstance(event, AnnotationsStartEvent):
-                text += "<annotations>\n"
-
-            # if isinstance(event, CompletionEndEvent):
-            #     text += f"\n</{event.proc_name}>\n"
-            if isinstance(event, ThinkingEndEvent):
-                text += "\n</thinking>\n"
-            elif isinstance(event, ResponseEndEvent):
-                text += "\n</response>\n"
-            elif isinstance(event, ToolCallEndEvent):
-                text += "\n</tool call>\n"
-            elif isinstance(event, AnnotationsEndEvent):
-                text += "\n</annotations>\n"
-
-            if isinstance(event, ThinkingChunkEvent):
-                thinking = event.data.thinking
-                if isinstance(thinking, str):
-                    text += thinking
-                else:
-                    text = "\n".join(
-                        [block.get("thinking", "[redacted]") for block in thinking]
-                    )
-
-            if isinstance(event, ResponseChunkEvent):
-                text += event.data.response
-
-            if isinstance(event, ToolCallChunkEvent):
-                text += event.data.tool_call.tool_arguments or ""
-
-            if isinstance(event, AnnotationsChunkEvent):
-                text += "\n".join(
-                    [
-                        json.dumps(annotation, indent=2)
-                        for annotation in event.data.annotations
-                    ]
-                )
-
-            stream_text(text, color)
+            stream_colored_text(_make_chunk_text(event))
 
         if isinstance(event, MessageEvent) and not isinstance(event, GenMessageEvent):
-            assert isinstance(event.data, (SystemMessage | UserMessage | ToolMessage))
-
-            message = event.data
-            role = message.role
-            content = Printer.content_to_str(message.content, role=role)
-            color = _get_color(event, role)
-
-            text = f"\n<{event.proc_name}> [{event.call_id}]\n"
-
-            if isinstance(event, (SystemMessageEvent, UserMessageEvent)):
-                content = Printer.truncate_content_str(content, trunc_len=trunc_len)
-
-            if isinstance(event, SystemMessageEvent):
-                text += f"<system>\n{content}\n</system>\n"
-
-            elif isinstance(event, UserMessageEvent):
-                text += f"<input>\n{content}\n</input>\n"
-
-            elif isinstance(event, ToolMessageEvent):
-                message = event.data
-                try:
-                    content = json.dumps(json.loads(content), indent=2)
-                except Exception:
-                    pass
-                text += (
-                    f"<tool result> [{message.tool_call_id}]\n"
-                    f"{content}\n</tool result>\n"
-                )
-
-            stream_text(text, color)
+            stream_colored_text(_make_message_text(event))
 
         if isinstance(
             event, (ProcPacketOutputEvent, WorkflowResultEvent, RunResultEvent)
         ):
-            _print_packet(event)
+            stream_colored_text(_make_packet_text(event))

grasp_agents/processors/base_processor.py

@@ -60,11 +60,13 @@ def with_retry(func: F) -> F:
                         logger.warning(f"{err_message}:\n{err}")
                     else:
                         logger.warning(f"{err_message} after retrying:\n{err}")
-                    raise ProcRunError(proc_name=self.name, call_id=call_id) from err
+                    # raise ProcRunError(proc_name=self.name, call_id=call_id) from err
+                    return None  # type: ignore[return]
 
                 logger.warning(f"{err_message} (retry attempt {n_attempt + 1}):\n{err}")
         # This part should not be reachable due to the raise in the loop
-        raise ProcRunError(proc_name=self.name, call_id=call_id)
+        # raise ProcRunError(proc_name=self.name, call_id=call_id)
+        return None  # type: ignore[return]
 
     return cast("F", wrapper)
 
@@ -106,7 +108,7 @@ _OutT_contra = TypeVar("_OutT_contra", contravariant=True)
 
 class RecipientSelector(Protocol[_OutT_contra, CtxT]):
     def __call__(
-        self, output: _OutT_contra, ctx: RunContext[CtxT]
+        self, output: _OutT_contra, *, ctx: RunContext[CtxT]
     ) -> Sequence[ProcName] | None: ...
 
 

grasp_agents/processors/parallel_processor.py

@@ -114,7 +114,7 @@ class ParallelProcessor(
         ctx: RunContext[CtxT] | None = None,
     ) -> Packet[OutT]:
         call_id = self._generate_call_id(call_id)
-        ctx = RunContext[CtxT](state=None) if ctx is None else ctx  # type: ignore
+        ctx = ctx or RunContext[CtxT](state=None)  # type: ignore
 
         val_in_args = self._validate_inputs(
             call_id=call_id,
@@ -223,7 +223,7 @@ class ParallelProcessor(
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
         call_id = self._generate_call_id(call_id)
-        ctx = RunContext[CtxT](state=None) if ctx is None else ctx  # type: ignore
+        ctx = ctx or RunContext[CtxT](state=None)  # type: ignore
 
         val_in_args = self._validate_inputs(
             call_id=call_id,

grasp_agents/processors/processor.py

@@ -105,7 +105,7 @@ class Processor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT, MemT, C
         call_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> Packet[OutT]:
-        ctx = RunContext[CtxT](state=None) if ctx is None else ctx  # type: ignore
+        ctx = ctx or RunContext[CtxT](state=None)  # type: ignore
 
         val_in_args, memory, call_id = self._preprocess(
             chat_inputs=chat_inputs,
@@ -136,7 +136,7 @@ class Processor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT, MemT, C
         call_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
-        ctx = RunContext[CtxT](state=None) if ctx is None else ctx  # type: ignore
+        ctx = ctx or RunContext[CtxT](state=None)  # type: ignore
 
         val_in_args, memory, call_id = self._preprocess(
             chat_inputs=chat_inputs,

grasp_agents/prompt_builder.py

@@ -15,11 +15,22 @@ _InT_contra = TypeVar("_InT_contra", contravariant=True)
 
 
 class SystemPromptBuilder(Protocol[CtxT]):
-    def __call__(self, ctx: RunContext[CtxT]) -> str | None: ...
+    def __call__(
+        self,
+        *,
+        ctx: RunContext[CtxT],
+        call_id: str,
+    ) -> str | None: ...
 
 
 class InputContentBuilder(Protocol[_InT_contra, CtxT]):
-    def __call__(self, in_args: _InT_contra, ctx: RunContext[CtxT]) -> Content: ...
+    def __call__(
+        self,
+        in_args: _InT_contra,
+        *,
+        ctx: RunContext[CtxT],
+        call_id: str,
+    ) -> Content: ...
 
 
 PromptArgumentType: TypeAlias = str | bool | int | ImageData
@@ -43,9 +54,9 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
         self._in_args_type_adapter: TypeAdapter[InT] = TypeAdapter(self._in_type)
 
     @final
-    def build_system_prompt(self, ctx: RunContext[CtxT]) -> str | None:
+    def build_system_prompt(self, ctx: RunContext[CtxT], call_id: str) -> str | None:
         if self.system_prompt_builder:
-            return self.system_prompt_builder(ctx=ctx)
+            return self.system_prompt_builder(ctx=ctx, call_id=call_id)
 
         return self.sys_prompt
 
@@ -71,7 +82,7 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
 
     @final
     def _build_input_content(
-        self, in_args: InT | None, ctx: RunContext[CtxT]
+        self, in_args: InT | None, ctx: RunContext[CtxT], call_id: str
     ) -> Content:
         if in_args is None and self._in_type is not type(None):
             raise InputPromptBuilderError(
@@ -83,7 +94,9 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
 
         val_in_args = self._validate_input_args(in_args)
         if self.input_content_builder:
-            return self.input_content_builder(in_args=val_in_args, ctx=ctx)
+            return self.input_content_builder(
+                in_args=val_in_args, ctx=ctx, call_id=call_id
+            )
 
         if issubclass(self._in_type, BaseModel) and isinstance(val_in_args, BaseModel):
             val_in_args_map = self._format_pydantic_prompt_args(val_in_args)
@@ -102,6 +115,7 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
         chat_inputs: LLMPrompt | Sequence[str | ImageData] | None = None,
         *,
         in_args: InT | None = None,
+        call_id: str,
         ctx: RunContext[CtxT],
     ) -> UserMessage | None:
         if chat_inputs is not None:
@@ -116,7 +130,9 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
             return UserMessage.from_content_parts(chat_inputs, name=self._agent_name)
 
         return UserMessage(
-            content=self._build_input_content(in_args=in_args, ctx=ctx),
+            content=self._build_input_content(
+                in_args=in_args, ctx=ctx, call_id=call_id
+            ),
             name=self._agent_name,
         )
 

grasp_agents/run_context.py

@@ -1,11 +1,11 @@
 from collections import defaultdict
-from typing import Any, Generic, TypeVar
+from typing import Generic, TypeVar
 
 from pydantic import BaseModel, ConfigDict, Field
 
 from grasp_agents.typing.completion import Completion
 
-from .printer import ColoringMode, Printer
+from .printer import Printer
 from .typing.io import ProcName
 from .usage_tracker import UsageTracker
 
@@ -19,13 +19,6 @@ class RunContext(BaseModel, Generic[CtxT]):
         default_factory=lambda: defaultdict(list)
     )
     usage_tracker: UsageTracker = Field(default_factory=UsageTracker)
-
     printer: Printer | None = None
-    log_messages: bool = False
-    color_messages_by: ColoringMode = "role"
-
-    def model_post_init(self, context: Any) -> None:  # noqa: ARG002
-        if self.log_messages:
-            self.printer = Printer(color_by=self.color_messages_by)
 
     model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)

grasp_agents/typing/tool.py

@@ -64,20 +64,22 @@ class BaseTool(
         self,
         inp: _InT,
         *,
-        call_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
+        call_id: str | None = None,
     ) -> _OutT_co:
         pass
 
     async def __call__(
         self,
         *,
-        call_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
+        call_id: str | None = None,
         **kwargs: Any,
     ) -> _OutT_co:
+        # NOTE: validation is probably redundant here when tool inputs have been
+        # validated by the LLM already
         input_args = TypeAdapter(self._in_type).validate_python(kwargs)
-        output = await self.run(input_args, call_id=call_id, ctx=ctx)
+        output = await self.run(input_args, ctx=ctx, call_id=call_id)
 
         return TypeAdapter(self._out_type).validate_python(output)
 

{grasp_agents-0.5.10.dist-info → grasp_agents-0.5.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: grasp_agents
-Version: 0.5.10
+Version: 0.5.12
 Summary: Grasp Agents Library
 License-File: LICENSE.md
 Requires-Python: <4,>=3.11.4
@@ -110,24 +110,16 @@ Create a script, e.g., `problem_recommender.py`:
 
 ```python
 import asyncio
-from pathlib import Path
 from typing import Any
 
 from dotenv import load_dotenv
 from pydantic import BaseModel, Field
 
-from grasp_agents.grasp_logging import setup_logging
+from grasp_agents import LLMAgent, BaseTool, RunContext, Printer
 from grasp_agents.litellm import LiteLLM, LiteLLMSettings
-from grasp_agents import LLMAgent, BaseTool, RunContext
-
-load_dotenv()
 
 
-# Configure the logger to output to the console and/or a file
-setup_logging(
-    logs_file_path="grasp_agents_demo.log",
-    logs_config_path=Path().cwd() / "configs/logging/default.yaml",
-)
+load_dotenv()
 
 sys_prompt_react = """
 Your task is to suggest an exciting stats problem to the student. 
@@ -162,13 +154,11 @@ Returns:
 """
 
 
-class AskStudentTool(BaseTool[TeacherQuestion, StudentReply, Any]):
+class AskStudentTool(BaseTool[TeacherQuestion, StudentReply, None]):
     name: str = "ask_student"
     description: str = ask_student_tool_description
 
-    async def run(
-        self, inp: TeacherQuestion, ctx: RunContext[Any] | None = None
-    ) -> StudentReply:
+    async def run(self, inp: TeacherQuestion, **kwargs: Any) -> StudentReply:
         return input(inp.question)
 
 
@@ -178,10 +168,7 @@ class Problem(BaseModel):
 
 teacher = LLMAgent[None, Problem, None](
     name="teacher",
-    llm=LiteLLM(
-        model_name="gpt-4.1",
-        llm_settings=LiteLLMSettings(temperature=0.5),
-    ),
+    llm=LiteLLM(model_name="gpt-4.1"),
     tools=[AskStudentTool()],
     react_mode=True,
     final_answer_as_tool_call=True,
@@ -189,7 +176,7 @@ teacher = LLMAgent[None, Problem, None](
 )
 
 async def main():
-    ctx = RunContext[None](log_messages=True)
+    ctx = RunContext[None](printer=Printer())
     out = await teacher.run("start", ctx=ctx)
     print(out.payloads[0])
     print(ctx.usage_tracker.total_usage)