[WIP] Webagent updates

tapeagents/agent.py

@@ -238,11 +238,15 @@ def model_post_init(self, __context: Any) -> None:
                 raise ValueError(
                     f"Node {node.name} references unknown LLM {node.llm}. Known LLMs: {list(self.llms.keys())}"
                 )
-            if hasattr(node, "add_known_actions"):
-                node.add_known_actions(self.known_actions)
             node_names.add(node.name)
         return super().model_post_init(__context)
 
+    def update_subagents(self):
+        for subagent in self.subagents:
+            subagent.tools_description = self.tools_description
+            subagent.known_actions = self.known_actions
+            subagent.update_subagents()
+
     @property
     def manager(self):
         """

tapeagents/core.py

@@ -6,7 +6,7 @@
 
 import datetime
 import json
-from typing import Any, Generic, Iterable, Iterator, List, Literal, TypeAlias, TypeVar
+from typing import Any, Generic, Iterable, Iterator, List, Literal, Type, TypeAlias, TypeVar
 from uuid import uuid4
 
 import litellm
@@ -335,6 +335,7 @@ class Prompt(BaseModel):
     tools: list[dict] | None = None
     messages: list[dict] = Field(default_factory=list)
     token_ids: list[int] = Field(default_factory=list)
+    response_format: dict | Type[BaseModel] | None = None
 
     @staticmethod
     def from_user_message(content: str) -> Prompt:

tapeagents/llms/__init__.py

@@ -2,7 +2,7 @@
 
 from .base import LLM, LLMEvent, LLMStream
 from .claude import Claude
-from .litellm import LiteLLM
+from .lite import LiteLLM
 from .mock import MockLLM
 from .replay import ReplayLLM
 from .trainable import TrainableLLM

tapeagents/llms/claude.py

@@ -1,3 +1,4 @@
+import logging
 import time
 from typing import Generator
 
@@ -7,9 +8,10 @@
 from tapeagents.core import Prompt
 from tapeagents.llms.base import LLMEvent, LLMOutput
 from tapeagents.llms.cached import CachedLLM
-from tapeagents.llms.litellm import logger
 from tapeagents.utils import get_step_schemas_from_union_type, resize_base64_message
 
+logger = logging.getLogger(__name__)
+
 
 class Claude(CachedLLM):
     max_tokens: int = 4096

tapeagents/llms/litellm.py → tapeagents/llms/lite.py

@@ -75,6 +75,7 @@ def _generate(
                     messages=prompt.messages,
                     tools=prompt.tools,
                     stream=self.stream,
+                    response_format=prompt.response_format,
                     **kwargs,
                 )
                 break

tapeagents/nodes.py

@@ -5,7 +5,7 @@
 import json
 import logging
 import re
-from typing import Annotated, Any, Generator, Type, Union
+from typing import Annotated, Any, Callable, Generator, Type, Union
 
 from litellm import ChatCompletionMessageToolCall
 from pydantic import Field, TypeAdapter, ValidationError
@@ -28,7 +28,7 @@
 from tapeagents.steps import BranchStep, ReasoningThought
 from tapeagents.tool_calling import as_openai_tool
 from tapeagents.tools.code_executor import PythonCodeAction
-from tapeagents.utils import FatalError, class_for_name, sanitize_json_completion
+from tapeagents.utils import FatalError, class_for_name, sanitize_json_completion, step_schema
 from tapeagents.view import Call, Respond, TapeViewStack
 
 logger = logging.getLogger(__name__)
@@ -71,30 +71,32 @@ class StandardNode(Node):
     trim_obs_except_last_n: int = 2
     use_function_calls: bool = False
     allow_code_blocks: bool = False
+    structured_output: bool = False
     _steps_type: Any = None
     _step_classes: list[type[Step]] | None = None
+    _tools: dict[str, dict] | None = None
+    _tool_name_to_cls: dict[str, type[Step]] | None = None
 
     def model_post_init(self, __context: Any) -> None:
-        self.prepare_step_types()
         super().model_post_init(__context)
 
-    def prepare_step_types(self, actions: list[type[Step]] = None):
-        actions = actions or []
+    def prepare_step_types(self, agent: Agent):
+        actions = agent.known_actions if self.use_known_actions else []
         step_classes_or_str = actions + (self.steps if isinstance(self.steps, list) else [self.steps])
         if not step_classes_or_str:
             return
         self._step_classes = [class_for_name(step) if isinstance(step, str) else step for step in step_classes_or_str]
         if self.allow_code_blocks:
             # remove PythonCodeAction from the list of step classes
             self._step_classes = [c for c in self._step_classes if c != PythonCodeAction]
-        self._name_to_cls = {c.__name__: c for c in self._step_classes}
+        if self.structured_output:
+            assert len(self._step_classes) == 1, "Structured output requires exactly one output step class"
         self._steps_type = Annotated[Union[tuple(self._step_classes)], Field(discriminator="kind")]
+        if self.use_function_calls:
+            self._tools = {step_cls: as_openai_tool(step_cls) for step_cls in self._step_classes}
+            self._tool_name_to_cls = {tool["function"]["name"]: step_cls for step_cls, tool in self._tools.items()}
 
-    def add_known_actions(self, actions: list[type[Step]]):
-        if self.use_known_actions:
-            self.prepare_step_types(actions)
-
-    def make_prompt(self, agent: Any, tape: Tape) -> Prompt:
+    def make_prompt(self, agent: Agent, tape: Tape) -> Prompt:
         """Create a prompt from tape interactions.
 
         This method constructs a prompt by processing the tape content and agent steps description
@@ -117,20 +119,22 @@ def make_prompt(self, agent: Any, tape: Tape) -> Prompt:
             4. Checks token count and trims if needed
             5. Reconstructs messages if trimming occurred
         """
-        cleaned_tape = self.prepare_tape(tape)
-        steps_description = self.get_steps_description(tape, agent)
-        messages = self.tape_to_messages(cleaned_tape, steps_description)
+        self.prepare_step_types(agent)
+        steps = self.get_steps(tape, agent)
+        steps_description = self.get_steps_description(agent)
+        messages = self.steps_to_messages(steps, steps_description)
         if agent.llms[self.llm].count_tokens(messages) > (agent.llms[self.llm].context_size - 500):
             old_trim = self.trim_obs_except_last_n
             self.trim_obs_except_last_n = 1
-            messages = self.tape_to_messages(cleaned_tape, steps_description)
+            messages = self.steps_to_messages(steps, steps_description)
             self.trim_obs_except_last_n = old_trim
-        prompt = Prompt(messages=messages)
-        if self.use_function_calls:
-            prompt.tools = [as_openai_tool(s) for s in self._step_classes]
+
+        response_format = self._step_classes[0] if self.structured_output else None
+        tools = list(self._tools.values()) if self.use_function_calls else None
+        prompt = Prompt(messages=messages, tools=tools, response_format=response_format)
         return prompt
 
-    def prepare_tape(self, tape: Tape) -> Tape:
+    def get_steps(self, tape: Tape, agent: Agent) -> list[Step]:
         """
         Prepares tape by filtering out control flow steps.
 
@@ -143,8 +147,9 @@ def prepare_tape(self, tape: Tape) -> Tape:
         Returns:
             Tape: A new tape instance containing only non-control flow steps.
         """
-        steps_without_control_flow = [step for step in tape.steps if not isinstance(step, SetNextNode)]
-        return tape.model_copy(update=dict(steps=steps_without_control_flow))
+        steps = agent.compute_view(tape).top.steps
+        steps_without_control_flow = [step for step in steps if not isinstance(step, (SetNextNode, Call, Respond))]
+        return steps_without_control_flow
 
     def make_llm_output(self, agent: Any, tape: Tape, index: int) -> LLMOutput:
         """
@@ -179,7 +184,7 @@ def make_llm_output(self, agent: Any, tape: Tape, index: int) -> LLMOutput:
         content = [step.llm_dict() for step in steps] if len(steps) > 1 else steps[0].llm_dict()
         return LLMOutput(role="assistant", content=json.dumps(content, indent=2, ensure_ascii=False))
 
-    def tape_to_messages(self, tape: Tape, steps_description: str) -> list[dict]:
+    def steps_to_messages(self, steps: list[Step], steps_description: str) -> list[dict]:
         """
         Converts a Tape object and steps description into a list of messages for LLM conversation.
 
@@ -201,8 +206,8 @@ def tape_to_messages(self, tape: Tape, steps_description: str) -> list[dict]:
             messages.append({"role": "system", "content": self.system_prompt})
         if steps_description:
             messages.append({"role": "user", "content": steps_description})
-        for i, step in enumerate(tape):
-            steps_after_current = len(tape) - i - 1
+        for i, step in enumerate(steps):
+            steps_after_current = len(steps) - i - 1
             role = "assistant" if isinstance(step, AgentStep) else "user"
             if isinstance(step, Observation) and steps_after_current >= self.trim_obs_except_last_n:
                 view = step.short_view()
@@ -217,7 +222,7 @@ def tape_to_messages(self, tape: Tape, steps_description: str) -> list[dict]:
             messages.append({"role": "user", "content": self.guidance})
         return messages
 
-    def get_steps_description(self, tape: Tape, agent: Agent) -> str:
+    def get_steps_description(self, agent: Agent) -> str:
         """
         Get the steps description for the agent's task.
 
@@ -231,10 +236,9 @@ def get_steps_description(self, tape: Tape, agent: Agent) -> str:
         Returns:
             str: The steps prompt describing the sequence of actions.
         """
-        if self.use_function_calls:
-            allowed_steps = ""
-        else:
-            allowed_steps = agent.llms[self.llm].get_step_schema(self._steps_type) if self._steps_type else ""
+        allowed_steps = ""
+        if self._steps_type and not self.use_function_calls:
+            allowed_steps = agent.llms[self.llm].get_step_schema(self._steps_type)
         return self.steps_prompt.format(allowed_steps=allowed_steps, tools_description=agent.tools_description)
 
     def generate_steps(
@@ -284,10 +288,10 @@ def generate_steps(
             yield SetNextNode(next_node=self.next_node)
 
     def tool_call_to_step(self, tool_call: ChatCompletionMessageToolCall) -> Step:
-        step_cls = self._name_to_cls.get(tool_call.function.name)
+        step_cls = self._tool_name_to_cls.get(tool_call.function.name)
         if step_cls is None:
             return LLMOutputParsingFailureAction(
-                error=f"Unknown tool call: {tool_call.function.name}", llm_output=tool_call
+                error=f"Unknown tool call: {tool_call.function.name}", llm_output=tool_call.model_dump_json(indent=2)
             )
         args = tool_call.function.arguments
         return step_cls.model_validate_json(args) if args else step_cls()
@@ -411,6 +415,30 @@ def trim_tape(self, tape: Tape) -> Tape:
         return tape
 
 
+class ViewNode(StandardNode):
+    system_prompt: str
+    view: Any = None
+    prompt: str
+
+    def get_steps(self, tape: Tape, agent: Agent) -> list[Step]:
+        view_cls = class_for_name(self.view)
+        kwargs = view_cls(tape).as_dict() if self.view else {}
+        content = self.prompt.format(**kwargs)
+        return [UserStep(content=content)]
+
+
+class AsStep(StandardNode):
+    def make_prompt(self, agent: Agent, tape: Tape) -> Prompt:
+        self.prepare_step_types(agent)
+        text = tape[-1].reasoning
+        schema = step_schema(self._step_classes[0])
+        response_format = self._step_classes[0] if self.structured_output else None
+        msg = f"Convert the following paragraph into a structured JSON object:\n\n{text}"
+        if not self.structured_output:
+            msg += f"\n\nThe JSON object should match the following schema:\n\n{schema}"
+        return Prompt(messages=[{"role": "user", "content": msg}], response_format=response_format)
+
+
 class ControlFlowNode(Node):
     """
     A node that controls the flow of execution by selecting the next node based on tape content.
@@ -474,6 +502,14 @@ def select_node(self, tape: Tape) -> str:
         return self.next_node if isinstance(tape[-1], self.step_class) else None
 
 
+class If(ControlFlowNode):
+    predicate: Callable[[Tape], bool]
+    next_node: str
+
+    def select_node(self, tape: Tape) -> str:
+        return self.next_node if self.predicate(tape) else None
+
+
 class ObservationControlNode(ControlFlowNode):
     """
     A control flow node that selects the next node based on the last observation in the tape.
@@ -544,6 +580,10 @@ def generate_steps(
             yield step
 
 
+class Return(FixedStepsNode):
+    steps: list[Step] = [Respond(copy_output=True)]
+
+
 class GoTo(Node):
     next_node: str
 

tapeagents/tool_calling.py

@@ -158,6 +158,11 @@ def as_openai_tool(action: Action) -> dict:
     props.pop("kind", None)
     name = schema["title"]
     description = schema.get("description", "")
+    if name.lower().endswith("action"):
+        name = name[:-6]
+    elif name.lower().endswith("thought"):
+        name = f"Produce{name}"
+        description = f"Produce {description}"
     if len(description) > 1024:
         description = description[:1024]
         logger.warning(f"Description of {name} truncated to 1024 characters: {description}")
@@ -198,7 +203,7 @@ def as_function_def(action: Action) -> str:
         ptype = type_aliases.get(param_spec["type"], param_spec["type"])
         fdef += f"{param}: {ptype}, "
     fdef = fdef[:-2] + "):"
-    fdef += f"\n    \"\"\"{tool_spec['function']['description']}\n"
+    fdef += f'\n    """{tool_spec["function"]["description"]}\n'
     for param, param_spec in tool_spec["function"]["parameters"]["properties"].items():
         fdef += f"    {param}: {param_spec['description']}\n"
     fdef += '    """\n'

tapeagents/tools/code_executor.py

@@ -68,17 +68,7 @@ def execute_action(self, action: PythonCodeAction) -> CodeExecutionResult:
         return obs
 
     def prepare_code(self, action: PythonCodeAction) -> str:
-        lines = action.code.splitlines()
-        if len(lines) == 1 and "\\n" in lines[0]:
-            lines = lines[0].split("\\n")
-        lines = [f"# {action.name}"] + lines
-        if "print(" not in lines[-1] and "break" not in lines[-1]:
-            if " = " in lines[-1]:
-                name = lines[-1].split("=")[0].strip()
-                lines.append(f"print({name})")
-            else:
-                lines[-1] = f"print({lines[-1]})"
-        return "\n".join(lines)
+        return f"# {action.name}\n{action.code}"
 
     def trim_output(self, output: str) -> str:
         if len(output) > self.max_output_length:

tapeagents/utils.py

@@ -16,7 +16,7 @@
 
 import jsonref
 from PIL import Image
-from pydantic import TypeAdapter
+from pydantic import BaseModel, TypeAdapter
 from termcolor import colored
 
 
@@ -110,6 +110,18 @@ def get_step_schemas_from_union_type(cls, simplify: bool = True) -> str:
     return json.dumps(clean_schema, ensure_ascii=False)
 
 
+def step_schema(step: BaseModel, simple: bool = True) -> str:
+    schema = step.model_json_schema()
+    step_dict: dict = dict(jsonref.replace_refs(schema, proxies=False))  # type: ignore
+    step_dict["properties"].pop("metadata", None)
+    if simple:
+        step_dict.pop("title", None)
+        for prop in step_dict["properties"]:
+            step_dict["properties"][prop].pop("title", None)
+        step_dict["properties"]["kind"] = {"const": step_dict["properties"]["kind"]["const"]}
+    return json.dumps(step_dict, ensure_ascii=False, indent=2)
+
+
 def image_base64_message(image_path: str) -> dict:
     image_extension = os.path.splitext(image_path)[1][1:]
     content_type = f"image/{image_extension}"

tapeagents/view.py

@@ -68,8 +68,8 @@ def add_step(self, step: StepType):
                 self.next_node = ""
             self.last_prompt_id = step.metadata.prompt_id
             self.last_node = step.metadata.node
-        if isinstance(step, SetNextNode):
-            self.next_node = step.next_node
+            if isinstance(step, SetNextNode):
+                self.next_node = step.next_node
         self.steps_by_kind[kind].append(step)
 
     def get_output(self, subagent_name_or_index: int | str) -> StepType: