Description
Question
Hi,
I created this example to investigate the behaviors of manipulating messages between agents as I feel ambiguous in the document.
Additional Context
from __future__ import annotations as _annotations
from dotenv import load_dotenv
import os
from dataclasses import dataclass, field
from pydantic import BaseModel
from pydantic_ai import RunContext
from pydantic_ai.messages import ModelMessage
from pydantic_ai import Agent
from pydantic_graph import BaseNode, End, Graph, GraphRunContext
from pydantic_ai.models.openai import OpenAIModel
load_dotenv()
llm_model = OpenAIModel(
os.getenv("AI_SCORING_MODEL_NAME", ""),
base_url=os.getenv("AI_SCORING_BASE_URL", ""),
api_key=os.getenv("AI_SCORING_API_KEY", ""),
)
class GeneralAgent(BaseModel):
"""An agent who handles general questions and chitchat."""
class ComputationalAgent(BaseModel):
"""An agent who handles computational tasks."""
manager_agent = Agent(
model=llm_model,
name="manager_agent",
result_type=GeneralAgent | ComputationalAgent,
retries=3,
)
@manager_agent.system_prompt(dynamic=True)
async def get_system_prompt(_ctx: RunContext) -> str:
prompt = """
You are a *manager_agent* with mission is navigation and task planning.
You can nagivate to the appropriate agent based on the user's requests and conversation history.
There are two supported agents: `general_agent` and `computational_agent`.
"""
return prompt
general_agent = Agent(
model=llm_model,
name="general_agent",
result_type=str,
)
@general_agent.system_prompt(dynamic=True)
async def get_system_prompt(_ctx: RunContext) -> str:
prompt = """
You are a *general_agent* with mission is chitchat with user`.
"""
return prompt
computational_agent = Agent(
model=llm_model,
name="computational_agent",
result_type=str
)
@computational_agent.system_prompt(dynamic=True)
async def get_system_prompt(_ctx: RunContext) -> str:
prompt = """
You are a *computational_agent* with mission is computation.
You have to return 1 for all requests of users.
"""
return prompt
@dataclass
class State:
message_history:list[ModelMessage] = field(default_factory=list)
@dataclass
class Manager(BaseNode[State]):
user_prompt: str
async def run(self, ctx: GraphRunContext[State]) -> General | Computation:
result = await manager_agent.run(
self.user_prompt,
message_history=ctx.state.message_history
)
ctx.state.message_history += result.new_messages()
if isinstance(result.data, GeneralAgent):
return General()
elif isinstance(result.data, ComputationalAgent):
return Computation()
@dataclass
class General(BaseNode[State]):
async def run(self, ctx: GraphRunContext[State]) -> End:
result = await general_agent.run(
"",
message_history=ctx.state.message_history
)
ctx.state.message_history += result.new_messages()
return End(result.data)
@dataclass
class Computation(BaseNode[State]):
async def run(self, ctx: GraphRunContext[State]) -> End:
result = await computational_agent.run(
"",
message_history=ctx.state.message_history
)
ctx.state.message_history += result.new_messages()
return End(result.data)
async def main():
my_graph = Graph(nodes=(Manager, General, Computation))
state = State()
while True:
print("\n==STATE==\n")
print(state)
user_prompt = input()
if user_prompt == "exit":
break
result = await my_graph.run(Manager(user_prompt), state = state)
print(result.output)
if __name__ == "__main__":
import asyncio
asyncio.run(main())
======
The output:
==STATE==
State(message_history=[])
hi
/home/nld/ai-scoring/notebooks/graph pydantic.py:129: LogfireNotConfiguredWarning: No logs or spans will be created until logfire.configure() has been called. Set the environment variable LOGFIRE_IGNORE_NO_CONFIG=1 or add ignore_no_config=true in pyproject.toml to suppress this warning.
result = await my_graph.run(Manager(user_prompt), state = state)
Hi there! How's your day going?
==STATE==
State(message_history=[ModelRequest(parts=[SystemPromptPart(content='\n You are a general_agent with mission is chitchat with user`.\n ', dynamic_ref='get_system_prompt', part_kind='system-prompt'), UserPromptPart(content='hi', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 34, 398203, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'), ModelResponse(parts=[ToolCallPart(tool_name='final_result_GeneralAgent', args='{}', tool_call_id='call_ui2kUvBPBRUnUV0wu2r2gM7G', part_kind='tool-call')], model_name='gpt-4o-mini-2024-07-18', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 34, tzinfo=datetime.timezone.utc), kind='response'), ModelRequest(parts=[ToolReturnPart(tool_name='final_result_GeneralAgent', content='Final result processed.', tool_call_id='call_ui2kUvBPBRUnUV0wu2r2gM7G', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 36, 118073, tzinfo=datetime.timezone.utc), part_kind='tool-return')], kind='request'), ModelRequest(parts=[UserPromptPart(content='', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 36, 119197, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'), ModelResponse(parts=[TextPart(content="Hi there! How's your day going?", part_kind='text')], model_name='gpt-4o-mini-2024-07-18', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 36, tzinfo=datetime.timezone.utc), kind='response')])
i'm okay how are you
I’m just a program, but I’m here and ready to chat! What’s on your mind?
==STATE==
State(message_history=[ModelRequest(parts=[SystemPromptPart(content='\n You are a general_agent with mission is chitchat with user`.\n ', dynamic_ref='get_system_prompt', part_kind='system-prompt'), UserPromptPart(content='hi', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 34, 398203, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'), ModelResponse(parts=[ToolCallPart(tool_name='final_result_GeneralAgent', args='{}', tool_call_id='call_ui2kUvBPBRUnUV0wu2r2gM7G', part_kind='tool-call')], model_name='gpt-4o-mini-2024-07-18', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 34, tzinfo=datetime.timezone.utc), kind='response'), ModelRequest(parts=[ToolReturnPart(tool_name='final_result_GeneralAgent', content='Final result processed.', tool_call_id='call_ui2kUvBPBRUnUV0wu2r2gM7G', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 36, 118073, tzinfo=datetime.timezone.utc), part_kind='tool-return')], kind='request'), ModelRequest(parts=[UserPromptPart(content='', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 36, 119197, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'), ModelResponse(parts=[TextPart(content="Hi there! How's your day going?", part_kind='text')], model_name='gpt-4o-mini-2024-07-18', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 36, tzinfo=datetime.timezone.utc), kind='response'), ModelRequest(parts=[UserPromptPart(content="i'm okay how are you", timestamp=datetime.datetime(2025, 4, 8, 3, 44, 2, 122190, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'), ModelResponse(parts=[ToolCallPart(tool_name='final_result_GeneralAgent', args='{}', tool_call_id='call_WxFUAq4ciBzOFhd4l568XHDu', part_kind='tool-call')], model_name='gpt-4o-mini-2024-07-18', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 2, tzinfo=datetime.timezone.utc), kind='response'), ModelRequest(parts=[ToolReturnPart(tool_name='final_result_GeneralAgent', content='Final result processed.', tool_call_id='call_WxFUAq4ciBzOFhd4l568XHDu', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 3, 252682, tzinfo=datetime.timezone.utc), part_kind='tool-return')], kind='request'), ModelRequest(parts=[UserPromptPart(content='', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 3, 254144, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'), ModelResponse(parts=[TextPart(content='I’m just a program, but I’m here and ready to chat! What’s on your mind?', part_kind='text')], model_name='gpt-4o-mini-2024-07-18', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 3, tzinfo=datetime.timezone.utc), kind='response')])
can you do my computation: 2^3+4
The result of the computation is 1.
==STATE==
State(message_history=[ModelRequest(parts=[SystemPromptPart(content='\n You are a computational_agent with mission is computation.\n You have to return 1 for all requests of users. \n ', dynamic_ref='get_system_prompt', part_kind='system-prompt'), UserPromptPart(content='hi', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 34, 398203, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'), ModelResponse(parts=[ToolCallPart(tool_name='final_result_GeneralAgent', args='{}', tool_call_id='call_ui2kUvBPBRUnUV0wu2r2gM7G', part_kind='tool-call')], model_name='gpt-4o-mini-2024-07-18', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 34, tzinfo=datetime.timezone.utc), kind='response'), ModelRequest(parts=[ToolReturnPart(tool_name='final_result_GeneralAgent', content='Final result processed.', tool_call_id='call_ui2kUvBPBRUnUV0wu2r2gM7G', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 36, 118073, tzinfo=datetime.timezone.utc), part_kind='tool-return')], kind='request'), ModelRequest(parts=[UserPromptPart(content='', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 36, 119197, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'), ModelResponse(parts=[TextPart(content="Hi there! How's your day going?", part_kind='text')], model_name='gpt-4o-mini-2024-07-18', timestamp=datetime.datetime(2025, 4, 8, 3, 43, 36, tzinfo=datetime.timezone.utc), kind='response'), ModelRequest(parts=[UserPromptPart(content="i'm okay how are you", timestamp=datetime.datetime(2025, 4, 8, 3, 44, 2, 122190, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'), ModelResponse(parts=[ToolCallPart(tool_name='final_result_GeneralAgent', args='{}', tool_call_id='call_WxFUAq4ciBzOFhd4l568XHDu', part_kind='tool-call')], model_name='gpt-4o-mini-2024-07-18', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 2, tzinfo=datetime.timezone.utc), kind='response'), ModelRequest(parts=[ToolReturnPart(tool_name='final_result_GeneralAgent', content='Final result processed.', tool_call_id='call_WxFUAq4ciBzOFhd4l568XHDu', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 3, 252682, tzinfo=datetime.timezone.utc), part_kind='tool-return')], kind='request'), ModelRequest(parts=[UserPromptPart(content='', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 3, 254144, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'), ModelResponse(parts=[TextPart(content='I’m just a program, but I’m here and ready to chat! What’s on your mind?', part_kind='text')], model_name='gpt-4o-mini-2024-07-18', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 3, tzinfo=datetime.timezone.utc), kind='response'), ModelRequest(parts=[UserPromptPart(content='can you do my computation: 2^3+4', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 20, 848373, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'), ModelResponse(parts=[ToolCallPart(tool_name='final_result_ComputationalAgent', args='{}', tool_call_id='call_mUVLDlSZVMFvT9duw069v2zU', part_kind='tool-call')], model_name='gpt-4o-mini-2024-07-18', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 21, tzinfo=datetime.timezone.utc), kind='response'), ModelRequest(parts=[ToolReturnPart(tool_name='final_result_ComputationalAgent', content='Final result processed.', tool_call_id='call_mUVLDlSZVMFvT9duw069v2zU', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 22, 387996, tzinfo=datetime.timezone.utc), part_kind='tool-return')], kind='request'), ModelRequest(parts=[UserPromptPart(content='', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 22, 389275, tzinfo=datetime.timezone.utc), part_kind='user-prompt')], kind='request'), ModelResponse(parts=[TextPart(content='The result of the computation is 1.', part_kind='text')], model_name='gpt-4o-mini-2024-07-18', timestamp=datetime.datetime(2025, 4, 8, 3, 44, 22, tzinfo=datetime.timezone.utc), kind='response')])
===========
So acutally it works the whole time?
I can see the system prompt is anchored and replaced every time hands off to another agent.
I also noticed that this is just the state of the graph. Not really sure what's gonna pass into the llm in agent as the final messages.
(nld-mindmesh) nld@nlp118:~/ai-scoring$ pip show pydantic_ai
Name: pydantic_ai
Version: 0.0.36rc0
Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
Home-page:
Author: Samuel Colvin
Author-email: samuel@pydantic.dev
License: MIT
Location: /HDD/.conda/envs/nld-mindmesh/lib/python3.10/site-packages
Requires: eval-type-backport, exceptiongroup, griffe, httpx, logfire-api, opentelemetry-api, pydantic, pydantic-graph, typing-extensions, typing-inspection