(agents)= # `agents`: Agent as Code ```{Note} If you just want to call LLM API. See [models](#models) section first! ``` An agent is defined as an executable control flow of a text generation session using LLMs, tools, and other functions. PromptTrail offers a simple and intuitive DSL to build agents with code. We call this `Agent as Code`. ```{Note} You can explore the `core` module to learn how you can mock, cache and debug your agent. ``` You can use the agent via CLI, API, etc. Therefore, you can build a chatbot on Agent, but you can also build any application that requires multiple-step text generation. If you're only building applications with single-turn text generation, you just need to use prompttrail.models, which allows you to use LLMs with a simple API. ## Introduction ### Template `agent.templates` offers a simple DSL to define the conversation flow. You can write how you would like to interact with the user, LLM, and functions in it. Example: A simple proofreading agent This is actually used in this repository to housekeep README.md, etc. See [examples/dogfooding/fix_markdown.py](https://github.com/combinatrix-ai/PromptTrail/blob/main/examples/dogfooding/fix_markdown.py) for the actual code. ```python from prompttrail.agent.templates import LinearTemplate, AssistantTemplate, MessageTemplate, SystemTemplate, UserTemplate templates = LinearTemplate( [ SystemTemplate( content=""" You're an AI proofreader that helps users fix markdown. You're given markdown content by the user. You only emit the corrected markdown. No explanation, comments, or anything else is needed. Do not remove > in the code section, which represents the prompt. """, ), UserTemplate( content="{{content}}", ), AssistantTemplate(), # Generates response using LLM ], ) ``` The template above is an example of a very simple agent. `LinearTemplate` is a template that runs templates in order. So, let's see child templates. The first `SystemTemplate` is a convenience template that automatically sets the role to "system" following OpenAI's convention. It's used to tell LLM what they are. In this agent, markdown is passed to LLM and LLM returns the corrected markdown. The second `UserTemplate` is a template that takes the user's input. `{{content}}` is a placeholder that will be replaced by `runner`. This template automatically sets the role to "user". This is where the actual markdown is passed. As some of you may have noticed, this is `Jinja2` template syntax. We use Jinja to dynamically generate templates. Finally, `AssistantTemplate` is a template that generates content using LLM. When no content is provided, it automatically calls the LLM to generate a response. This template automatically sets the role to "assistant". OK. You may grasp what's going on here. Let's run this agent. ### Runners PromptTrail provides different runners to execute your agent templates in various environments. Each runner type is designed for specific use cases: #### CommandLineRunner The `CommandLineRunner` is ideal for development, debugging, and CLI applications. It provides a simple command-line interface for interacting with your agent. ```python from prompttrail.agent.runners import CommandLineRunner from prompttrail.agent.user_interface import CLIInterface runner = CommandLineRunner( model=model, # LLM model instance template=templates, # Your template definition user_interface=CLIInterface() ) # Run with optional parameters session = runner.run( session=None, # Optional: existing session to continue max_messages=100, # Optional: limit on number of messages debug_mode=False # Optional: enable debug output ) ``` The CommandLineRunner prints all interactions to the console, making it easy to follow the conversation flow and debug your agent. #### APIRunner The `APIRunner` exposes your agent as a REST API server, perfect for web applications and services. It provides endpoints for session management and user interaction. ```python from prompttrail.agent.runners import APIRunner from prompttrail.agent.user_interface import CLIInterface runner = APIRunner( model=model, template=templates, user_interface=CLIInterface(), host="127.0.0.1", # Optional: server host (default: 127.0.0.1) port=8000 # Optional: server port (default: 8000) ) # Start the API server runner.start_server() ``` The APIRunner provides the following endpoints: - `POST /sessions` - Create a new session - Request body: `{ "metadata": {} }` (optional) - Returns: `{ "session_id": "uuid", "session": {...} }` - `GET /sessions/{session_id}` - Get session status and current state - Returns: ```json { "is_running": bool, "has_event": bool, "session": {...}, "current_event": { "instruction": "string", "default": "string" } } ``` - `POST /sessions/{session_id}/start` - Start template execution for a session - Returns: `{ "status": "started" }` - `POST /sessions/{session_id}/input` - Provide user input when requested - Request body: `{ "input": "user input" }` - Returns: `{ "status": "success" }` For a complete working example of using APIRunner, see [examples/agent/weather_forecast_api.py](examples/agent/weather_forecast_api.py). This example demonstrates: - Setting up an APIRunner with a weather forecast tool - Configuring the API server - Using session metadata to pass user input - Interacting with the API endpoints Example usage with curl: ```bash # Create a new session with weather query curl -X POST http://localhost:8000/sessions \ -H "Content-Type: application/json" \ -d '{"metadata": {"user_input": "What'\''s the weather in Tokyo?"}}' # Start the session curl -X POST http://localhost:8000/sessions/{session_id}/start # Check session status and get weather forecast curl http://localhost:8000/sessions/{session_id} ``` #### Common Runner Features All runners share these core capabilities: - **Session Management**: Continue conversations by passing existing sessions - **Template Search**: Find templates by ID using `runner.search_template()` - **Debug Mode**: Enable detailed logging with `debug_mode=True` - **Message Limits**: Control conversation length with `max_messages` Let's see a basic example using the CommandLineRunner: ```python import os from prompttrail.core import Session from prompttrail.agent.runner import CommandLineRunner from prompttrail.agent.user_interface import CLIInterface from prompttrail.models.openai import ( OpenAIModel, OpenAIConfiguration, OpenAIParam, ) # Setup LLM model # Don't forget to set OPENAI_API_KEY environment variable config = OpenAIConfig( api_key=os.environ.get("OPENAI_API_KEY", ""), model_name="gpt-4o-mini", temperature=0.0, max_tokens=8000 ) model = OpenAIModel(configuration=config) # Define runner runner = CommandLineRunner( model=model, template=templates, user_interface=CLIInterface(), ) ``` OK, we are ready to run the agent. Let's run it! We need to prepare the markdown file to be proofread. ```python markdown = """ # PromptTrail PromptTrail is a library to build a text generation agent with LLMs. """ ``` Then, we need to pass the markdown to the agent. The point here is `session`. `session` is a state that is passed to the templates. In this example, we pass the markdown to the template. `session.metadata` is passed to the Jinja2 processor and impute the template. Each message also has its own metadata that can be accessed via `message.metadata`. The metadata is managed by a dedicated `Metadata` class that provides dictionary-like operations with type safety. You can also update the metadata with LLM outputs, function results, etc. See [examples/agent/fermi_problem.py] for an example. Finally, run the agent! ```python result = runner.run( session=Session( metadata={"content": markdown}, ), ) ``` You will see the following output on your terminal. ```python StatefulMessage( 'content': """\nYou're an AI proofreader that helps users fix markdown.\nYou're given markdown content by the user.\nYou only emit the corrected markdown. No explanation, comments, or anything else is needed.\nDo not remove > in the code section, which represents the prompt.""", 'role': 'system', ), StatefulMessage( 'content': """\n# PromptTrail\n\nPromptTrail is a library to build a text generation agent with LLMs.""", 'role': 'user', ), StatefulMessage( 'content': """'# PromptTrail\n\nPromptTrail is a library to build a text generation agent with LLMs.""", 'role': 'assistant', ) ``` Pretty simple, right? What we want is the last message, which is the corrected markdown. Now we saved the output of the runner in `result`, which is the final `state`. We can extract the conversation as follows: ```python result.messages ``` We just need the last message: ```python corrected_markdown = result.messages[-1].content print(corrected_markdown) ``` The result is (may vary depending on the LLM): ```markdown # PromptTrail PromptTrail is a library to build a text generation agent with LLMs. ``` Great! We have built our first agent! Here we have reviewed the core concepts of `prompttrail.agent`. You may start using `prompttrail.agent` to build your own agent now! ## Subroutines Subroutines allow you to execute templates in an isolated session context with flexible strategies for session initialization and message handling. This is particularly useful when you want to: - Break down complex tasks into smaller, manageable pieces - Isolate certain parts of the conversation - Reuse template logic across different agents - Control message inheritance and propagation ### Session Initialization Strategies When executing a subroutine, you can control how the session is initialized using various strategies: ```python from prompttrail.agent.session_init_strategy import ( CleanSessionStrategy, InheritSystemStrategy, LastNMessagesStrategy, FilteredInheritStrategy, ) # Start with a clean session clean_strategy = CleanSessionStrategy() # Inherit only system messages system_strategy = InheritSystemStrategy() # Keep last N messages from parent last_n_strategy = LastNMessagesStrategy(n=3) # Custom filtering of messages def is_important(msg): return msg.metadata.get("importance", 0) > 0.5 filtered_strategy = FilteredInheritStrategy(filter_fn=is_important) ``` ### Message Squashing Strategies After subroutine execution, you can control which messages are propagated back to the parent session: ```python from prompttrail.agent.squash_strategy import ( LastMessageStrategy, FilterByRoleStrategy, ) # Keep only the last message last_message = LastMessageStrategy() # Keep messages with specific roles assistant_only = FilterByRoleStrategy(roles=["assistant"]) ``` ### Using SubroutineTemplate Here's an example of using subroutines to solve math problems: ```python from prompttrail.agent.subroutine import SubroutineTemplate from prompttrail.agent.templates import Template, LinearTemplate class CalculationTemplate(Template): """Template for performing calculations""" def _render(self, session): yield Message(role="assistant", content="Let me solve this step by step:") yield Message(role="assistant", content="1. First let's identify the key numbers...") yield Message(role="assistant", content="The result is 42") return session class MathProblemTemplate(Template): """Template for solving math problems using subroutines""" def __init__(self): super().__init__() # Create calculation subroutine calculation = CalculationTemplate() self.calculation_subroutine = SubroutineTemplate( template=calculation, session_init_strategy=InheritSystemStrategy(), squash_strategy=FilterByRoleStrategy(roles=["assistant"]), ) def _render(self, session): # Main problem-solving flow yield Message(role="assistant", content="I'll help solve this math problem") # Execute calculation subroutine for message in self.calculation_subroutine.render(session): yield message yield Message(role="assistant", content="Problem solved!") return session # Use in a linear template template = LinearTemplate([ SystemTemplate(content="You are a math teacher."), UserTemplate(content="What is 6 x 7?"), MathProblemTemplate(), ]) ``` This example demonstrates: 1. Isolated execution of the calculation logic 2. Inheritance of system context 3. Filtering of messages to keep only assistant responses 4. Clean separation of concerns between problem setup and calculation See [examples/agent/subroutine_example.py](examples/agent/subroutine_example.py) for a complete working example. ## Hooks Hooks are used to enhance the template. Let's see an excerpt from [examples/agent/fermi_problem.py]: ```python AssistantTemplate( after_transform=[ ExtractMarkdownCodeBlockHook( key="python_segment", lang="python" ), EvaluatePythonCodeHook(key="answer", code="python_segment"), ] ), ``` This template orders the LLM to generate text, extract a Python code block from the generated text, and evaluate the code. `after_transform` is called after the LLM generates text. We passed `ExtractMarkdownCodeBlockHook` and `EvaluatePythonCodeHook`. Let's see what they do. `ExtractMarkdownCodeBlockHook` extracts a code block of the language specified by `lang` from the generated text and stores it in the message's metadata under the key `"python_segment"`. `EvaluatePythonCodeHook` evaluates the code stored in the message's metadata under the key `"python_segment"` and stores the result under the key `"answer"`. As a convention, `key` is used to represent the key in the message's metadata to store the result of the hook. `gather_feedback` and `first.template_id` are template ids, the unique identifier of the template passed to the runner. `template_id` can be set at instantiation like: ```python GenerateTemplate( role="assistant", template_id="gather_feedback", ... ) ``` However, you can omit it. In that case, `template_id` is automatically generated. You can get it by `template.template_id`. Anyway, we omitted the templates in this example, so we will not explain it here. See [examples/agent/fermi_problem.py] for more details. As there are `after_transform` and there are `before_transform`. They are called before the `rendering` of the template. The order of hooks is: - `before_transform` - (rendering) - `after_transform` ## Rendering `rendering` is a process to create a message from a template. Every template has a `render` method. For `MessageTemplate`, it simply renders the template with the message's metadata via Jinja2 and returns the result as a message. For `GenerateTemplate`, it calls the LLM and returns the result as a message. For `InputTemplate`, it asks for user input using `user_interface` and returns the result as a message. You can also add your own template. See [template.py] for more details. ## Session `Session` is a state that is passed to the templates and holds the state of the conversation. If you're going to build an application with `prompttrail.agent`, you just need the following: - `Session.metadata` - This is a `Metadata` instance that holds the session-level metadata. - You can use this to pass data to the templates and store session state. - Each message also has its own metadata that can be accessed via `message.metadata`. - If a key is not found in the metadata, an error will be raised unless you specify a `default` in the template or hooks. - `Session.messages` - This is a list of messages in the conversation. - Each message has `content`, `role`, and `metadata`. ```python class Session(BaseModel): """A session represents a conversation between a user and a model, or API etc...""" messages: List[Message] = Field(default_factory=list) metadata: Metadata = Field(default_factory=Metadata) runner: Optional["Runner"] = Field(default=None, exclude=True) debug_mode: bool = Field(default=False) stack: List["Stack"] = Field(default_factory=list) ``` Other attributes can also be accessed: - `runner`: You can access the runner itself. If you want to search templates passed to the runner, you can use `session.runner.search_template`. - `stack`: You can access the template stack. This is used for template control flow. ## Control Flow PromptTrail provides several templates for controlling the flow of conversation: - `LoopTemplate`: Repeats a sequence of templates until an exit condition is met or maximum iterations reached ```python LoopTemplate( templates=[...], exit_condition=lambda session: session.get_last().content == "END", exit_loop_count=10 # Optional: maximum number of iterations ) ``` - `IfTemplate`: Conditionally executes templates based on a condition ```python IfTemplate( condition=lambda session: "answer" in session.metadata, true_template=AssistantTemplate(...), false_template=BreakTemplate() ) ``` The conditions are defined as lambda functions that take a Session object and return a boolean value. ## Tool (Function Calling) `agent.tool` is a set of tools that can be used by LLMs, especially OpenAI's function calling feature. Using `agent.tool`, functions called by LLMs can be written with a unified interface. The tool system provides type safety and automatic documentation generation from type annotations. Let's see an example from [examples/agent/weather_forecast.py]: ```python from typing import Literal from typing_extensions import TypedDict from prompttrail.agent.tools import Tool, ToolArgument, ToolResult # First, define the structure of the tool's output using TypedDict class WeatherData(TypedDict): """Weather data structure""" temperature: float weather: Literal["sunny", "rainy", "cloudy", "snowy"] # Then define the result class that wraps the output structure class WeatherForecastResult(ToolResult): """Weather forecast result This class defines the structure of the data returned by the weather forecast tool. The content field contains: - temperature: The temperature in the specified unit (float) - weather: The weather condition (one of: sunny, rainy, cloudy, snowy) """ content: WeatherData # Finally, implement the tool itself class WeatherForecastTool(Tool): """Weather forecast tool This tool simulates getting weather forecast data for a location. """ name: str = "get_weather_forecast" description: str = "Get the current weather in a given location" # Define the tool's arguments using ToolArgument arguments: Dict[str, ToolArgument[Any]] = { "location": ToolArgument[str]( name="location", description="The location to get the weather forecast", value_type=str, required=True ), "unit": ToolArgument[str]( name="unit", description="Temperature unit (celsius or fahrenheit)", value_type=str, required=False ) } def _execute(self, args: Dict[str, Any]) -> ToolResult: """Execute the weather forecast tool This is where the actual weather data fetching would happen. For this example, we return mock data. """ return WeatherForecastResult( content={ "temperature": 20.5, "weather": "sunny" } ) ``` This tool definition is automatically converted to OpenAI's function calling format: ```json { "name": "get_weather_forecast", "description": "Get the current weather in a given location", "parameters": { "type": "object", "properties": { "location": { "type": "string", "description": "The location to get the weather forecast" }, "unit": { "type": "string", "description": "Temperature unit (celsius or fahrenheit)" } }, "required": ["location"] } } ``` Then, you can use this tool with OpenAI's function calling feature through the `ToolingTemplate`: ```python from prompttrail.agent.templates import ( LinearTemplate, UserTemplate, ToolingTemplate, ) template = LinearTemplate( [ MessageTemplate( content="You are a helpful weather assistant that provides weather forecasts.", role="system" ), UserTemplate( role="user", content="What's the weather in Tokyo?", ), # This template handles the function calling flow: # 1. Sends the function definition to OpenAI # 2. Gets back which function to call with what arguments # 3. Executes the function with those arguments # 4. Sends the result back to OpenAI for final response ToolingTemplate( role="assistant", functions=[WeatherForecastTool()], ), ] ) ``` The tool system handles all the complexity of function calling for you: - Type-safe argument validation - Automatic documentation generation from type hints and docstrings - Function calling API formatting and execution - Result parsing and conversion ### Subroutine Tool In addition to regular tools, PromptTrail provides `SubroutineTool` that allows you to execute templates as tools. This is particularly useful when you want to: - Execute complex conversation flows as tools - Reuse existing templates in function calling context - Maintain isolation between main conversation and tool execution Here's an example of using SubroutineTool: ```python from prompttrail.agent.tools import SubroutineTool from prompttrail.agent.templates import Template class WeatherTemplate(Template): """Template for weather forecasting""" def _render(self, session): yield Message( role="assistant", content="Let me check the weather forecast..." ) yield Message( role="assistant", content="Based on the data, it will be sunny with 25°C." ) return session # Create weather tool from template weather_tool = SubroutineTool( name="get_weather", description="Get weather forecast for a location", template=WeatherTemplate(), ) # Use like any other tool in ToolingTemplate template = LinearTemplate([ SystemTemplate(content="You are a weather assistant."), UserTemplate(content="What's the weather?"), ToolingTemplate(tools=[weather_tool]) ]) ``` The SubroutineTool provides: - Isolated execution context for templates - Automatic message handling and result formatting - Integration with function calling flow - Access to all template features (hooks, control flow, etc.) See [examples/agent/subroutine_tool_example.py](examples/agent/subroutine_tool_example.py) for a complete working example. This allows you to focus on implementing the actual tool functionality rather than dealing with API integration details. Isn't it great?