Skip to content

Commit

Permalink
fixing agent docs and refactoring retriever return types
Browse files Browse the repository at this point in the history
  • Loading branch information
@samuelcolvin
samuelcolvin committed Nov 18, 2024
1 parent 6577fde commit 33903b4
Show file tree
Hide file tree
Showing 14 changed files with 607 additions and 246 deletions.
323 changes: 307 additions & 16 deletions docs/agents.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -120,8 +120,8 @@ System prompts might seem simple at first glance since they're just strings (or

Generally, system prompts fall into two categories:

1. **Static system prompts**: These are known when writing the code and can be defined via the `system_prompt` parameter of the `Agent` constructor.
2. **Dynamic system prompts**: These aren't known until runtime and should be defined via functions decorated with `@agent.system_prompt`.
1. **Static system prompts**: These are known when writing the code and can be defined via the `system_prompt` parameter of the [`Agent` constructor][pydantic_ai.Agent.__init__].
2. **Dynamic system prompts**: These aren't known until runtime and should be defined via functions decorated with [`@agent.system_prompt`][pydantic_ai.Agent.system_prompt].

You can add both to a single agent; they're concatenated in the order they're defined at runtime.

Expand Down Expand Up @@ -161,26 +161,317 @@ print(result.data)

## Retrievers

* two different retriever decorators (`retriver_plain` and `retriever_context`) depending on whether you want to use the context or not, show an example using both
* retriever parameters are extracted and used to build the schema for the tool, then validated with pydantic
* if a retriever has a single "model like" parameter (e.g. pydantic mode, dataclass, typed dict), the schema for the tool will but just that type
* docstrings are parsed to get the tool description, thanks to griffe docs for each parameter are extracting using Google, numpy or sphinx docstring styling
* You can raise `ModelRetry` from within a retriever to suggest to the model it should retry
* the return type of retriever can either be `str` or a JSON object typed as `dict[str, Any]` as some models (e.g. Gemini) support structured return values, some expect text (OpenAI) but seem to be just as good at extracting meaning from the data
Retrievers provide a mechanism for models to request extra information to help them generate a response.

They're useful when it is impractical or impossible to put all the context an agent might need into the system prompt, or when you want to make agents' behavior more deterministic by deferring some of the logic required to generate a response to another tool.

!!! info "Retrievers vs. RAG"
Retrievers are basically the "R" of RAG (Retrieval-Augmented Generation) — they augment what the model can do by letting it request extra information.

The main semantic difference between PydanticAI Retreivers and RAG is RAG is synonymous with vector search, while PydanticAI retrievers are more general purpose. (Note: we might add support for some vector search functionality in the future, particuarly an API for generating embeddings, see [#58](https://github.com/pydantic/pydantic-ai/issues/58))

There are two different decorator functions to register retrievers:

1. [`@agent.retriever_plain`][pydantic_ai.Agent.retriever_plain] — for retrievers that don't need access to the agent [context][pydantic_ai.dependencies.CallContext]
2. [`@agent.retriever_context`][pydantic_ai.Agent.retriever_context] — for retrievers that do need access to the agent [context][pydantic_ai.dependencies.CallContext]

Here's an example using both:

```py title="dice_game.py"
import random

from pydantic_ai import Agent, CallContext

agent = Agent(
'gemini-1.5-flash', # (1)!
deps_type=str, # (2)!
system_prompt=(
"You're a dice game, you should roll the dice and see if the number "
"you got back matches the user's guess, if so tell them they're a winner. "
"Use the player's name in the response."
),
)


@agent.retriever_plain # (3)!
def roll_dice() -> str:
"""Roll a six-sided dice and return the result."""
return str(random.randint(1, 6))


@agent.retriever_context # (4)!
def get_player_name(ctx: CallContext[str]) -> str:
"""Get the player's name."""
return ctx.deps


dice_result = agent.run_sync('My guess is 4', deps='Adam') # (5)!
print(dice_result.data)
#> Congratulations Adam, you guessed correctly! You're a winner!
```

1. This is a pretty simple task, so we can use the fast and cheap Gemini flash model.
2. We pass the user's name as the dependency, to keep things simple we use just the name as a string as the dependency.
3. This retriever doesn't need any context, it just returns a random number. You could probably use a dynamic system prompt in this case.
4. This retriever needs the player's name, so it uses `CallContext` to access dependencies which are just the player's name.
5. Run the agent, passing the player's name as the dependency.

_(This example is complete, it can be run "as is")_

Let's print the messages from that game to see what happened:

```python title="dice_game_messages.py"
from dice_game import dice_result

print(dice_result.all_messages())
"""
[
SystemPrompt(
content="You're a dice game, you should roll the dice and see if the number you got back matches the user's guess, if so tell them they're a winner. Use the player's name in the response.",
role='system',
),
UserPrompt(
content='My guess is 4',
timestamp=datetime.datetime(...),
role='user',
),
ModelStructuredResponse(
calls=[
ToolCall(
tool_name='roll_dice', args=ArgsObject(args_object={}), tool_id=None
)
],
timestamp=datetime.datetime(...),
role='model-structured-response',
),
ToolReturn(
tool_name='roll_dice',
content='4',
tool_id=None,
timestamp=datetime.datetime(...),
role='tool-return',
),
ModelStructuredResponse(
calls=[
ToolCall(
tool_name='get_player_name',
args=ArgsObject(args_object={}),
tool_id=None,
)
],
timestamp=datetime.datetime(...),
role='model-structured-response',
),
ToolReturn(
tool_name='get_player_name',
content='Adam',
tool_id=None,
timestamp=datetime.datetime(...),
role='tool-return',
),
ModelTextResponse(
content="Congratulations Adam, you guessed correctly! You're a winner!",
timestamp=datetime.datetime(...),
role='model-text-response',
),
]
"""
```

We can represent that as a flow diagram, thus:

![Dice game flow diagram](./img/dice-diagram-light.svg#only-light)
![Dice game flow diagram](./img/dice-diagram-dark.svg#only-dark)

### Retrievers, tools, and schema

Under the hood, retrievers use the model's "tools" or "functions" API to let the model know what retrievers are available to call. Tools or functions are also used to define the schema(s) for structured responses, thus a model might have access to many tools, some of which call retrievers while others end the run and return a result.

Function parameters are extracted from the function signature, and all parameters except `CallContext` are used to build the schema for that tool call.

Even better, PydanticAI extracts the docstring from retriever functions and (thanks to [griffe](https://mkdocstrings.github.io/griffe/)) extracts parameter descriptions from the docstring and add them to the schema.

[Griffe supports](https://mkdocstrings.github.io/griffe/reference/docstrings/#docstrings) extracting parameter descriptions from `google`, `numpy` and `sphinx` style docstrings, PydanticAI will infer the format to use based on the docstring. We'll add support in future to explicitly set the style to use, and warn/error if not all parameters are documented, see [#59](https://github.com/pydantic/pydantic-ai/issues/59).

To demonstrate retriever schema, here we use [`FunctionModel`][pydantic_ai.models.function.FunctionModel] to print the schema a model would receive:

```py title="retriever_schema.py"
from pydantic_ai import Agent
from pydantic_ai.messages import Message, ModelAnyResponse, ModelTextResponse
from pydantic_ai.models.function import AgentInfo, FunctionModel

agent = Agent()


@agent.retriever_plain
def foobar(a: int, b: str, c: dict[str, list[float]]) -> str:
"""Get me foobar.
Args:
a: apple pie
b: banana cake
c: carrot smoothie
"""
return f'{a} {b} {c}'


def print_schema(messages: list[Message], info: AgentInfo) -> ModelAnyResponse:
retriever = info.retrievers['foobar']
print(retriever.description)
#> Get me foobar.
print(retriever.json_schema)
"""
{
'description': 'Get me foobar.',
'properties': {
'a': {'description': 'apple pie', 'title': 'A', 'type': 'integer'},
'b': {'description': 'banana cake', 'title': 'B', 'type': 'string'},
'c': {
'additionalProperties': {'items': {'type': 'number'}, 'type': 'array'},
'description': 'carrot smoothie',
'title': 'C',
'type': 'object',
},
},
'required': ['a', 'b', 'c'],
'type': 'object',
'additionalProperties': False,
}
"""
return ModelTextResponse(content='foobar')


agent.run_sync('hello', model=FunctionModel(print_schema))
```

_(This example is complete, it can be run "as is")_

The return type of retriever can either be `str` or a JSON object typed as `dict[str, Any]` as some models (e.g. Gemini) support structured return values, some expect text (OpenAI) but seem to be just as good at extracting meaning from the data.

## Reflection and self-correction

* validation errors from both retrievers parameter validation and structured result validation can be passed back to the with a request to retry
* as described above, you can also raise `ModelRetry` from within a retriever or result validator to tell the model it should retry
* the default retry count is 1, but can be altered both on a whole agent, or on a per-retriever basis and result validator basis
* you can access the current retry count from within a retriever or result validator via `ctx.retry`
Validation errors from both retriever parameter validation and [structured result validation](results.md#structured-result-validation) can be passed back to the model with a request to retry.

You can also raise [`ModelRetry`][pydantic_ai.exceptions.ModelRetry] from within a [retriever](#retrievers) or [result validators functions](results.md#result-validators-functions) to tell the model it should retry.

- The default retry count is **1** but can be altered for the [entire agent][pydantic_ai.Agent.__init__], a [specific retriever][pydantic_ai.Agent.retriever_context], or a [result validator][pydantic_ai.Agent.__init__].
- You can access the current retry count from within a retriever or result validator via [`ctx.retry`][pydantic_ai.dependencies.CallContext].

Here's an example:

```py title="retriever_retry.py"
from fake_database import DatabaseConn
from pydantic import BaseModel

from pydantic_ai import Agent, CallContext, ModelRetry


class ChatResult(BaseModel):
user_id: int
message: str


agent = Agent(
'openai:gpt-4o',
deps_type=DatabaseConn,
result_type=ChatResult,
)


@agent.retriever_context(retries=2)
def get_user_by_name(ctx: CallContext[DatabaseConn], name: str) -> int:
"""Get a user's ID from their full name."""
print(name)
#> John
#> John Doe
user_id = ctx.deps.users.get(name=name)
if user_id is None:
raise ModelRetry(
f'No user found with name {name!r}, remember to provide their full name'
)
return user_id


result = agent.run_sync(
'Send a message to John Doe asking for coffee next week', deps=DatabaseConn()
)
print(result.data)
"""
user_id=123 message='Hello John, would you be free for coffee sometime next week? Let me know what works for you!'
"""
```

## Model errors

* If models behave unexpectedly, e.g. the retry limit is exceed, agent runs will raise `UnexpectedModelBehaviour` exceptions
* If you use PydanticAI in correctly, we try to raise a `UserError` with a helpful message
* show an except of a `UnexpectedModelBehaviour` being raised
* if a `UnexpectedModelBehaviour` is raised, you may want to access the [`.last_run_messages`][pydantic_ai.Agent.last_run_messages] attribute of an agent to see the messages exchanged that led to the error, show an example of accessing `.last_run_messages` in an except block to get more details
If models behave unexpectedly (e.g., the retry limit is exceeded, or their api returns `503`), agent runs will raise [`UnexpectedModelBehaviour`][pydantic_ai.exceptions.UnexpectedModelBehaviour].

In these cases, [`agent.last_run_messages`][pydantic_ai.Agent.last_run_messages] can be used to access the messages exchanged during the run to help diagnose the issue.

```python
from pydantic_ai import Agent, ModelRetry, UnexpectedModelBehaviour

agent = Agent('openai:gpt-4o')


@agent.retriever_plain
def calc_volume(size: int) -> int: # (1)!
if size == 42:
return size**3
else:
raise ModelRetry('Please try again.')


try:
result = agent.run_sync('Please get me the volume of a box with size 6.')
except UnexpectedModelBehaviour as e:
print('An error occurred:', e)
#> An error occurred: Retriever exceeded max retries count of 1
print('cause:', repr(e.__cause__))
#> cause: ModelRetry('Please try again.')
print('messages:', agent.last_run_messages)
"""
messages:
[
UserPrompt(
content='Please get me the volume of a box with size 6.',
timestamp=datetime.datetime(...),
role='user',
),
ModelStructuredResponse(
calls=[
ToolCall(
tool_name='calc_volume',
args=ArgsObject(args_object={'size': 6}),
tool_id=None,
)
],
timestamp=datetime.datetime(...),
role='model-structured-response',
),
RetryPrompt(
content='Please try again.',
tool_name='calc_volume',
tool_id=None,
timestamp=datetime.datetime(...),
role='retry-prompt',
),
ModelStructuredResponse(
calls=[
ToolCall(
tool_name='calc_volume',
args=ArgsObject(args_object={'size': 6}),
tool_id=None,
)
],
timestamp=datetime.datetime(...),
role='model-structured-response',
),
]
"""
else:
print(result.data)
```
1. Define a retriever that will raise `ModelRetry` repeatedly in this case.

## API Reference

Expand Down
4 changes: 2 additions & 2 deletions docs/dependencies.md
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# Dependencies

PydanticAI uses a dependency injection system to provide data and services to your agent's [system prompts](agents.md#system-prompts), [retrievers](agents.md#retrievers) and [result validators](results.md#result-validators).
PydanticAI uses a dependency injection system to provide data and services to your agent's [system prompts](agents.md#system-prompts), [retrievers](agents.md#retrievers) and [result validators](results.md#result-validators-functions).

Matching PydanticAI's design philosophy, our dependency system tries to use existing best practice in Python development rather than inventing esoteric "magic", this should make dependencies type-safe, understandable easier to test and ultimately easier to deploy in production.

Expand Down Expand Up @@ -159,7 +159,7 @@ _(This example is complete, it can be run "as is")_

## Full Example

As well as system prompts, dependencies can be used in [retrievers](agents.md#retrievers) and [result validators](results.md#result-validators).
As well as system prompts, dependencies can be used in [retrievers](agents.md#retrievers) and [result validators](results.md#result-validators-functions).

```python title="full_example.py" hl_lines="27-35 38-48"
from dataclasses import dataclass
Expand Down
Loading

0 comments on commit 33903b4

Please sign in to comment.