fixing agent docs and refactoring retriever return types (#60)

Co-authored-by: hyperlint-ai[bot] <154288675+hyperlint-ai[bot]@users.noreply.github.com>

Author: samuelcolvin

docs/agents.md

@@ -114,14 +114,16 @@ print(result2.data)
 ```
 1. Continue the conversation, without `message_history` the model would not know who "he" was referring to.
 
+_(This example is complete, it can be run "as is")_
+
 ## System Prompts
 
 System prompts might seem simple at first glance since they're just strings (or sequences of strings that are concatenated), but crafting the right system prompt is key to getting the model to behave as you want.
 
 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.
 
@@ -159,28 +161,323 @@ print(result.data)
 3. Dynamic system prompt defined via a decorator.
 4. Another dynamic system prompt, system prompts don't have to have the `CallContext` parameter.
 
+_(This example is complete, it can be run "as is")_
+
 ## 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 any valid JSON object ([`JsonData`][pydantic_ai.dependencies.JsonData]) as some models (e.g. Gemini) support semi-structured return values, some expect text (OpenAI) but seem to be just as good at extracting meaning from the data, if a Python is returned and the model expects a string, the value will be serialized to JSON
 
 ## 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 validator 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.
+
+_(This example is complete, it can be run "as is")_
 
 ## API Reference
 

docs/dependencies.md

@@ -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.
 
@@ -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