rename retriever_context -> retriever (#68)

Author: samuelcolvin

docs/agents.md

@@ -31,7 +31,7 @@ roulette_agent = Agent(  # (1)!
 )
 
 
-@roulette_agent.retriever_context
+@roulette_agent.retriever
 async def roulette_wheel(ctx: CallContext[int], square: int) -> str:  # (2)!
     """check if the square is a winner"""
     return 'winner' if square == ctx.deps else 'loser'
@@ -179,7 +179,9 @@ They're useful when it is impractical or impossible to put all the context an ag
 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]
+2. [`@agent.retriever`][pydantic_ai.Agent.retriever] — for retrievers that do need access to the agent [context][pydantic_ai.dependencies.CallContext]
+
+`@agent.retriever` is the default since in the majority of cases retrievers will need access to the agent context.
 
 Here's an example using both:
 
@@ -205,7 +207,7 @@ def roll_dice() -> str:
     return str(random.randint(1, 6))
 
 
-@agent.retriever_context  # (4)!
+@agent.retriever  # (4)!
 def get_player_name(ctx: CallContext[str]) -> str:
     """Get the player's name."""
     return ctx.deps
@@ -219,7 +221,7 @@ print(dice_result.data)
 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.
+4. This retriever needs the player's name, so it uses `CallContext` to access dependencies which are just the player's name in this case.
 5. Run the agent, passing the player's name as the dependency.
 
 _(This example is complete, it can be run "as is")_
@@ -362,7 +364,7 @@ Validation errors from both retriever parameter validation and [structured resul
 
 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__].
+- 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], 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:
@@ -386,7 +388,7 @@ agent = Agent(
 )
 
 
-@agent.retriever_context(retries=2)
+@agent.retriever(retries=2)
 def get_user_by_name(ctx: CallContext[DatabaseConn], name: str) -> int:
     """Get a user's ID from their full name."""
     print(name)

docs/api/agent.md

@@ -13,5 +13,5 @@
         - last_run_messages
         - system_prompt
         - retriever_plain
-        - retriever_context
+        - retriever
         - result_validator

docs/dependencies.md

@@ -188,7 +188,7 @@ async def get_system_prompt(ctx: CallContext[MyDeps]) -> str:
     return f'Prompt: {response.text}'
 
 
-@agent.retriever_context  # (1)!
+@agent.retriever  # (1)!
 async def get_joke_material(ctx: CallContext[MyDeps], subject: str) -> str:
     response = await ctx.deps.http_client.get(
         'https://example.com#jokes',
@@ -220,7 +220,7 @@ async def main():
         #> Did you hear about the toothpaste scandal? They called it Colgate.
 ```
 
-1. To pass `CallContext` and to a retriever, us the [`retriever_context`][pydantic_ai.Agent.retriever_context] decorator.
+1. To pass `CallContext` and to a retriever, us the [`retriever`][pydantic_ai.Agent.retriever] decorator.
 2. `CallContext` may optionally be passed to a [`result_validator`][pydantic_ai.Agent.result_validator] function as the first argument.
 
 _(This example is complete, it can be run "as is")_
@@ -324,7 +324,7 @@ joke_agent = Agent(
 factory_agent = Agent('gemini-1.5-pro', result_type=list[str])
 
 
-@joke_agent.retriever_context
+@joke_agent.retriever
 async def joke_factory(ctx: CallContext[MyDeps], count: int) -> str:
     r = await ctx.deps.factory_agent.run(f'Please generate {count} jokes.')
     return '\n'.join(r.data)

docs/index.md

@@ -94,7 +94,7 @@ async def add_customer_name(ctx: CallContext[SupportDependencies]) -> str:
     return f"The customer's name is {customer_name!r}"
 
 
-@support_agent.retriever_context  # (6)!
+@support_agent.retriever  # (6)!
 async def customer_balance(
     ctx: CallContext[SupportDependencies], include_pending: bool
 ) -> str:

pydantic_ai/_pydantic.py

@@ -78,10 +78,10 @@ def function_schema(either_function: _retriever.RetrieverEitherFunc[AgentDeps, R
 
             if index == 0 and takes_ctx:
                 if not _is_call_ctx(annotation):
-                    errors.append('First argument must be a CallContext instance when using `.retriever_context`')
+                    errors.append('First argument must be a CallContext instance when using `.retriever`')
                 continue
             elif not takes_ctx and _is_call_ctx(annotation):
-                errors.append('CallContext instance can only be used with `.retriever_context`')
+                errors.append('CallContext instance can only be used with `.retriever`')
                 continue
             elif index != 0 and _is_call_ctx(annotation):
                 errors.append('CallContext instance can only be used as the first argument')

pydantic_ai/agent.py

@@ -338,16 +338,16 @@ def result_validator(
         return func
 
     @overload
-    def retriever_context(
+    def retriever(
         self, func: RetrieverContextFunc[AgentDeps, RetrieverParams], /
     ) -> _r.Retriever[AgentDeps, RetrieverParams]: ...
 
     @overload
-    def retriever_context(
+    def retriever(
         self, /, *, retries: int | None = None
     ) -> Callable[[RetrieverContextFunc[AgentDeps, RetrieverParams]], _r.Retriever[AgentDeps, RetrieverParams]]: ...
 
-    def retriever_context(
+    def retriever(
         self,
         func: RetrieverContextFunc[AgentDeps, RetrieverParams] | None = None,
         /,

pydantic_ai_examples/bank_support.py

@@ -62,7 +62,7 @@ async def add_customer_name(ctx: CallContext[SupportDependencies]) -> str:
     return f"The customer's name is {customer_name!r}"
 
 
-@support_agent.retriever_context
+@support_agent.retriever
 async def customer_balance(
     ctx: CallContext[SupportDependencies], include_pending: bool
 ) -> str:

pydantic_ai_examples/rag.py

@@ -51,7 +51,7 @@ class Deps:
 agent = Agent('openai:gpt-4o', deps_type=Deps)
 
 
-@agent.retriever_context
+@agent.retriever
 async def retrieve(context: CallContext[Deps], search_query: str) -> str:
     """Retrieve documentation sections based on a search query.
 

pydantic_ai_examples/weather_agent.py

@@ -41,7 +41,7 @@ class Deps:
 )
 
 
-@weather_agent.retriever_context
+@weather_agent.retriever
 async def get_lat_lng(
     ctx: CallContext[Deps], location_description: str
 ) -> dict[str, float]:
@@ -71,7 +71,7 @@ async def get_lat_lng(
         raise ModelRetry('Could not find the location')
 
 
-@weather_agent.retriever_context
+@weather_agent.retriever
 async def get_weather(ctx: CallContext[Deps], lat: float, lng: float) -> dict[str, Any]:
     """Get the weather at a location.
 

tests/models/test_model_function.py

@@ -119,7 +119,7 @@ async def get_location(location_description: str) -> str:
     return json.dumps(lat_lng)
 
 
-@weather_agent.retriever_context
+@weather_agent.retriever
 async def get_weather(_: CallContext[None], lat: int, lng: int):
     if (lat, lng) == (51, 0):
         # it always rains in London
@@ -199,7 +199,7 @@ async def call_function_model(messages: list[Message], _: AgentInfo) -> ModelAny
 var_args_agent = Agent(FunctionModel(call_function_model), deps_type=int)
 
 
-@var_args_agent.retriever_context
+@var_args_agent.retriever
 def get_var_args(ctx: CallContext[int], *args: int):
     assert ctx.deps == 123
     return json.dumps({'args': args})
@@ -233,7 +233,7 @@ async def call_retriever(messages: list[Message], info: AgentInfo) -> ModelAnyRe
 def test_deps_none():
     agent = Agent(FunctionModel(call_retriever))
 
-    @agent.retriever_context
+    @agent.retriever
     async def get_none(ctx: CallContext[None]):
         nonlocal called
 
@@ -259,7 +259,7 @@ def get_check_foobar(ctx: CallContext[tuple[str, str]]) -> str:
         return ''
 
     agent = Agent(FunctionModel(call_retriever), deps_type=tuple[str, str])
-    agent.retriever_context(get_check_foobar)
+    agent.retriever(get_check_foobar)
     called = False
     agent.run_sync('Hello', deps=('foo', 'bar'))
     assert called
@@ -277,12 +277,12 @@ def test_model_arg():
 agent_all = Agent()
 
 
-@agent_all.retriever_context
+@agent_all.retriever
 async def foo(_: CallContext[None], x: int) -> str:
     return str(x + 1)
 
 
-@agent_all.retriever_context(retries=3)
+@agent_all.retriever(retries=3)
 def bar(ctx, x: int) -> str:  # pyright: ignore[reportUnknownParameterType,reportMissingParameterType]
     return str(x + 2)