Skip to content

Commit 615c1d5

Browse files
assadyousufassadyousuf
authored
Merge branch 'main' into use-click-for-cli
2 parents 8abf2bd + dc3c9e9 commit 615c1d5

File tree

10 files changed

+165
-78
lines changed

10 files changed

+165
-78
lines changed

‎pydantic_ai_slim/pydantic_ai/_cli.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -260,15 +260,15 @@ async def run_chat(
260260
prog_name: str,
261261
config_dir: Path | None = None,
262262
deps: AgentDepsT = None,
263-
message_history: list[ModelMessage] | None = None,
263+
message_history: Sequence[ModelMessage] | None = None,
264264
) -> int:
265265
prompt_history_path = (config_dir or PYDANTIC_AI_HOME) / PROMPT_HISTORY_FILENAME
266266
prompt_history_path.parent.mkdir(parents=True, exist_ok=True)
267267
prompt_history_path.touch(exist_ok=True)
268268
session: PromptSession[Any] = PromptSession(history=FileHistory(str(prompt_history_path)))
269269

270270
multiline = False
271-
messages: list[ModelMessage] = message_history[:] if message_history else []
271+
messages: list[ModelMessage] = list(message_history) if message_history else []
272272

273273
while True:
274274
try:
@@ -304,7 +304,7 @@ async def ask_agent(
304304
console: Console,
305305
code_theme: str,
306306
deps: AgentDepsT = None,
307-
messages: list[ModelMessage] | None = None,
307+
messages: Sequence[ModelMessage] | None = None,
308308
) -> list[ModelMessage]:
309309
status = Status('[dim]Working on it…[/dim]', console=console)
310310

‎pydantic_ai_slim/pydantic_ai/agent/__init__.py‎

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -344,6 +344,7 @@ def __init__(
344344

345345
self._event_stream_handler = event_stream_handler
346346

347+
self._override_name: ContextVar[_utils.Option[str]] = ContextVar('_override_name', default=None)
347348
self._override_deps: ContextVar[_utils.Option[AgentDepsT]] = ContextVar('_override_deps', default=None)
348349
self._override_model: ContextVar[_utils.Option[models.Model]] = ContextVar('_override_model', default=None)
349350
self._override_toolsets: ContextVar[_utils.Option[Sequence[AbstractToolset[AgentDepsT]]]] = ContextVar(
@@ -384,7 +385,8 @@ def name(self) -> str | None:
384385
385386
If `None`, we try to infer the agent name from the call frame when the agent is first run.
386387
"""
387-
return self._name
388+
name_ = self._override_name.get()
389+
return name_.value if name_ else self._name
388390

389391
@name.setter
390392
def name(self, value: str | None) -> None:
@@ -415,7 +417,7 @@ def iter(
415417
user_prompt: str | Sequence[_messages.UserContent] | None = None,
416418
*,
417419
output_type: None = None,
418-
message_history: list[_messages.ModelMessage] | None = None,
420+
message_history: Sequence[_messages.ModelMessage] | None = None,
419421
deferred_tool_results: DeferredToolResults | None = None,
420422
model: models.Model | models.KnownModelName | str | None = None,
421423
deps: AgentDepsT = None,
@@ -432,7 +434,7 @@ def iter(
432434
user_prompt: str | Sequence[_messages.UserContent] | None = None,
433435
*,
434436
output_type: OutputSpec[RunOutputDataT],
435-
message_history: list[_messages.ModelMessage] | None = None,
437+
message_history: Sequence[_messages.ModelMessage] | None = None,
436438
deferred_tool_results: DeferredToolResults | None = None,
437439
model: models.Model | models.KnownModelName | str | None = None,
438440
deps: AgentDepsT = None,
@@ -449,7 +451,7 @@ async def iter(
449451
user_prompt: str | Sequence[_messages.UserContent] | None = None,
450452
*,
451453
output_type: OutputSpec[RunOutputDataT] | None = None,
452-
message_history: list[_messages.ModelMessage] | None = None,
454+
message_history: Sequence[_messages.ModelMessage] | None = None,
453455
deferred_tool_results: DeferredToolResults | None = None,
454456
model: models.Model | models.KnownModelName | str | None = None,
455457
deps: AgentDepsT = None,
@@ -566,7 +568,7 @@ async def main():
566568
# Build the initial state
567569
usage = usage or _usage.RunUsage()
568570
state = _agent_graph.GraphAgentState(
569-
message_history=message_history[:] if message_history else [],
571+
message_history=list(message_history) if message_history else [],
570572
usage=usage,
571573
retries=0,
572574
run_step=0,
@@ -690,7 +692,7 @@ def _run_span_end_attributes(
690692
}
691693
else:
692694
attrs = {
693-
'pydantic_ai.all_messages': json.dumps(settings.messages_to_otel_messages(state.message_history)),
695+
'pydantic_ai.all_messages': json.dumps(settings.messages_to_otel_messages(list(state.message_history))),
694696
**settings.system_instructions_attributes(literal_instructions),
695697
}
696698

@@ -712,24 +714,31 @@ def _run_span_end_attributes(
712714
def override(
713715
self,
714716
*,
717+
name: str | _utils.Unset = _utils.UNSET,
715718
deps: AgentDepsT | _utils.Unset = _utils.UNSET,
716719
model: models.Model | models.KnownModelName | str | _utils.Unset = _utils.UNSET,
717720
toolsets: Sequence[AbstractToolset[AgentDepsT]] | _utils.Unset = _utils.UNSET,
718721
tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] | _utils.Unset = _utils.UNSET,
719722
instructions: Instructions[AgentDepsT] | _utils.Unset = _utils.UNSET,
720723
) -> Iterator[None]:
721-
"""Context manager to temporarily override agent dependencies, model, toolsets, tools, or instructions.
724+
"""Context manager to temporarily override agent name, dependencies, model, toolsets, tools, or instructions.
722725
723726
This is particularly useful when testing.
724727
You can find an example of this [here](../testing.md#overriding-model-via-pytest-fixtures).
725728
726729
Args:
730+
name: The name to use instead of the name passed to the agent constructor and agent run.
727731
deps: The dependencies to use instead of the dependencies passed to the agent run.
728732
model: The model to use instead of the model passed to the agent run.
729733
toolsets: The toolsets to use instead of the toolsets passed to the agent constructor and agent run.
730734
tools: The tools to use instead of the tools registered with the agent.
731735
instructions: The instructions to use instead of the instructions registered with the agent.
732736
"""
737+
if _utils.is_set(name):
738+
name_token = self._override_name.set(_utils.Some(name))
739+
else:
740+
name_token = None
741+
733742
if _utils.is_set(deps):
734743
deps_token = self._override_deps.set(_utils.Some(deps))
735744
else:
@@ -759,6 +768,8 @@ def override(
759768
try:
760769
yield
761770
finally:
771+
if name_token is not None:
772+
self._override_name.reset(name_token)
762773
if deps_token is not None:
763774
self._override_deps.reset(deps_token)
764775
if model_token is not None:

‎pydantic_ai_slim/pydantic_ai/agent/abstract.py‎

Lines changed: 21 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ async def run(
128128
user_prompt: str | Sequence[_messages.UserContent] | None = None,
129129
*,
130130
output_type: None = None,
131-
message_history: list[_messages.ModelMessage] | None = None,
131+
message_history: Sequence[_messages.ModelMessage] | None = None,
132132
deferred_tool_results: DeferredToolResults | None = None,
133133
model: models.Model | models.KnownModelName | str | None = None,
134134
deps: AgentDepsT = None,
@@ -146,7 +146,7 @@ async def run(
146146
user_prompt: str | Sequence[_messages.UserContent] | None = None,
147147
*,
148148
output_type: OutputSpec[RunOutputDataT],
149-
message_history: list[_messages.ModelMessage] | None = None,
149+
message_history: Sequence[_messages.ModelMessage] | None = None,
150150
deferred_tool_results: DeferredToolResults | None = None,
151151
model: models.Model | models.KnownModelName | str | None = None,
152152
deps: AgentDepsT = None,
@@ -163,7 +163,7 @@ async def run(
163163
user_prompt: str | Sequence[_messages.UserContent] | None = None,
164164
*,
165165
output_type: OutputSpec[RunOutputDataT] | None = None,
166-
message_history: list[_messages.ModelMessage] | None = None,
166+
message_history: Sequence[_messages.ModelMessage] | None = None,
167167
deferred_tool_results: DeferredToolResults | None = None,
168168
model: models.Model | models.KnownModelName | str | None = None,
169169
deps: AgentDepsT = None,
@@ -242,7 +242,7 @@ def run_sync(
242242
user_prompt: str | Sequence[_messages.UserContent] | None = None,
243243
*,
244244
output_type: None = None,
245-
message_history: list[_messages.ModelMessage] | None = None,
245+
message_history: Sequence[_messages.ModelMessage] | None = None,
246246
deferred_tool_results: DeferredToolResults | None = None,
247247
model: models.Model | models.KnownModelName | str | None = None,
248248
deps: AgentDepsT = None,
@@ -260,7 +260,7 @@ def run_sync(
260260
user_prompt: str | Sequence[_messages.UserContent] | None = None,
261261
*,
262262
output_type: OutputSpec[RunOutputDataT],
263-
message_history: list[_messages.ModelMessage] | None = None,
263+
message_history: Sequence[_messages.ModelMessage] | None = None,
264264
deferred_tool_results: DeferredToolResults | None = None,
265265
model: models.Model | models.KnownModelName | str | None = None,
266266
deps: AgentDepsT = None,
@@ -277,7 +277,7 @@ def run_sync(
277277
user_prompt: str | Sequence[_messages.UserContent] | None = None,
278278
*,
279279
output_type: OutputSpec[RunOutputDataT] | None = None,
280-
message_history: list[_messages.ModelMessage] | None = None,
280+
message_history: Sequence[_messages.ModelMessage] | None = None,
281281
deferred_tool_results: DeferredToolResults | None = None,
282282
model: models.Model | models.KnownModelName | str | None = None,
283283
deps: AgentDepsT = None,
@@ -348,7 +348,7 @@ def run_stream(
348348
user_prompt: str | Sequence[_messages.UserContent] | None = None,
349349
*,
350350
output_type: None = None,
351-
message_history: list[_messages.ModelMessage] | None = None,
351+
message_history: Sequence[_messages.ModelMessage] | None = None,
352352
deferred_tool_results: DeferredToolResults | None = None,
353353
model: models.Model | models.KnownModelName | str | None = None,
354354
deps: AgentDepsT = None,
@@ -366,7 +366,7 @@ def run_stream(
366366
user_prompt: str | Sequence[_messages.UserContent] | None = None,
367367
*,
368368
output_type: OutputSpec[RunOutputDataT],
369-
message_history: list[_messages.ModelMessage] | None = None,
369+
message_history: Sequence[_messages.ModelMessage] | None = None,
370370
deferred_tool_results: DeferredToolResults | None = None,
371371
model: models.Model | models.KnownModelName | str | None = None,
372372
deps: AgentDepsT = None,
@@ -384,7 +384,7 @@ async def run_stream( # noqa C901
384384
user_prompt: str | Sequence[_messages.UserContent] | None = None,
385385
*,
386386
output_type: OutputSpec[RunOutputDataT] | None = None,
387-
message_history: list[_messages.ModelMessage] | None = None,
387+
message_history: Sequence[_messages.ModelMessage] | None = None,
388388
deferred_tool_results: DeferredToolResults | None = None,
389389
model: models.Model | models.KnownModelName | str | None = None,
390390
deps: AgentDepsT = None,
@@ -560,7 +560,7 @@ def run_stream_events(
560560
user_prompt: str | Sequence[_messages.UserContent] | None = None,
561561
*,
562562
output_type: None = None,
563-
message_history: list[_messages.ModelMessage] | None = None,
563+
message_history: Sequence[_messages.ModelMessage] | None = None,
564564
deferred_tool_results: DeferredToolResults | None = None,
565565
model: models.Model | models.KnownModelName | str | None = None,
566566
deps: AgentDepsT = None,
@@ -577,7 +577,7 @@ def run_stream_events(
577577
user_prompt: str | Sequence[_messages.UserContent] | None = None,
578578
*,
579579
output_type: OutputSpec[RunOutputDataT],
580-
message_history: list[_messages.ModelMessage] | None = None,
580+
message_history: Sequence[_messages.ModelMessage] | None = None,
581581
deferred_tool_results: DeferredToolResults | None = None,
582582
model: models.Model | models.KnownModelName | str | None = None,
583583
deps: AgentDepsT = None,
@@ -593,7 +593,7 @@ def run_stream_events(
593593
user_prompt: str | Sequence[_messages.UserContent] | None = None,
594594
*,
595595
output_type: OutputSpec[RunOutputDataT] | None = None,
596-
message_history: list[_messages.ModelMessage] | None = None,
596+
message_history: Sequence[_messages.ModelMessage] | None = None,
597597
deferred_tool_results: DeferredToolResults | None = None,
598598
model: models.Model | models.KnownModelName | str | None = None,
599599
deps: AgentDepsT = None,
@@ -674,7 +674,7 @@ async def _run_stream_events(
674674
user_prompt: str | Sequence[_messages.UserContent] | None = None,
675675
*,
676676
output_type: OutputSpec[RunOutputDataT] | None = None,
677-
message_history: list[_messages.ModelMessage] | None = None,
677+
message_history: Sequence[_messages.ModelMessage] | None = None,
678678
deferred_tool_results: DeferredToolResults | None = None,
679679
model: models.Model | models.KnownModelName | str | None = None,
680680
deps: AgentDepsT = None,
@@ -726,7 +726,7 @@ def iter(
726726
user_prompt: str | Sequence[_messages.UserContent] | None = None,
727727
*,
728728
output_type: None = None,
729-
message_history: list[_messages.ModelMessage] | None = None,
729+
message_history: Sequence[_messages.ModelMessage] | None = None,
730730
deferred_tool_results: DeferredToolResults | None = None,
731731
model: models.Model | models.KnownModelName | str | None = None,
732732
deps: AgentDepsT = None,
@@ -743,7 +743,7 @@ def iter(
743743
user_prompt: str | Sequence[_messages.UserContent] | None = None,
744744
*,
745745
output_type: OutputSpec[RunOutputDataT],
746-
message_history: list[_messages.ModelMessage] | None = None,
746+
message_history: Sequence[_messages.ModelMessage] | None = None,
747747
deferred_tool_results: DeferredToolResults | None = None,
748748
model: models.Model | models.KnownModelName | str | None = None,
749749
deps: AgentDepsT = None,
@@ -761,7 +761,7 @@ async def iter(
761761
user_prompt: str | Sequence[_messages.UserContent] | None = None,
762762
*,
763763
output_type: OutputSpec[RunOutputDataT] | None = None,
764-
message_history: list[_messages.ModelMessage] | None = None,
764+
message_history: Sequence[_messages.ModelMessage] | None = None,
765765
deferred_tool_results: DeferredToolResults | None = None,
766766
model: models.Model | models.KnownModelName | str | None = None,
767767
deps: AgentDepsT = None,
@@ -854,18 +854,20 @@ async def main():
854854
def override(
855855
self,
856856
*,
857+
name: str | _utils.Unset = _utils.UNSET,
857858
deps: AgentDepsT | _utils.Unset = _utils.UNSET,
858859
model: models.Model | models.KnownModelName | str | _utils.Unset = _utils.UNSET,
859860
toolsets: Sequence[AbstractToolset[AgentDepsT]] | _utils.Unset = _utils.UNSET,
860861
tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] | _utils.Unset = _utils.UNSET,
861862
instructions: Instructions[AgentDepsT] | _utils.Unset = _utils.UNSET,
862863
) -> Iterator[None]:
863-
"""Context manager to temporarily override agent dependencies, model, toolsets, tools, or instructions.
864+
"""Context manager to temporarily override agent name, dependencies, model, toolsets, tools, or instructions.
864865
865866
This is particularly useful when testing.
866867
You can find an example of this [here](../testing.md#overriding-model-via-pytest-fixtures).
867868
868869
Args:
870+
name: The name to use instead of the name passed to the agent constructor and agent run.
869871
deps: The dependencies to use instead of the dependencies passed to the agent run.
870872
model: The model to use instead of the model passed to the agent run.
871873
toolsets: The toolsets to use instead of the toolsets passed to the agent constructor and agent run.
@@ -1112,7 +1114,7 @@ async def to_cli(
11121114
self: Self,
11131115
deps: AgentDepsT = None,
11141116
prog_name: str = 'pydantic-ai',
1115-
message_history: list[_messages.ModelMessage] | None = None,
1117+
message_history: Sequence[_messages.ModelMessage] | None = None,
11161118
) -> None:
11171119
"""Run the agent in a CLI chat interface.
11181120
@@ -1149,7 +1151,7 @@ def to_cli_sync(
11491151
self: Self,
11501152
deps: AgentDepsT = None,
11511153
prog_name: str = 'pydantic-ai',
1152-
message_history: list[_messages.ModelMessage] | None = None,
1154+
message_history: Sequence[_messages.ModelMessage] | None = None,
11531155
) -> None:
11541156
"""Run the agent in a CLI chat interface with the non-async interface.
11551157

‎pydantic_ai_slim/pydantic_ai/agent/wrapper.py‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ def iter(
7272
user_prompt: str | Sequence[_messages.UserContent] | None = None,
7373
*,
7474
output_type: None = None,
75-
message_history: list[_messages.ModelMessage] | None = None,
75+
message_history: Sequence[_messages.ModelMessage] | None = None,
7676
deferred_tool_results: DeferredToolResults | None = None,
7777
model: models.Model | models.KnownModelName | str | None = None,
7878
deps: AgentDepsT = None,
@@ -89,7 +89,7 @@ def iter(
8989
user_prompt: str | Sequence[_messages.UserContent] | None = None,
9090
*,
9191
output_type: OutputSpec[RunOutputDataT],
92-
message_history: list[_messages.ModelMessage] | None = None,
92+
message_history: Sequence[_messages.ModelMessage] | None = None,
9393
deferred_tool_results: DeferredToolResults | None = None,
9494
model: models.Model | models.KnownModelName | str | None = None,
9595
deps: AgentDepsT = None,
@@ -106,7 +106,7 @@ async def iter(
106106
user_prompt: str | Sequence[_messages.UserContent] | None = None,
107107
*,
108108
output_type: OutputSpec[RunOutputDataT] | None = None,
109-
message_history: list[_messages.ModelMessage] | None = None,
109+
message_history: Sequence[_messages.ModelMessage] | None = None,
110110
deferred_tool_results: DeferredToolResults | None = None,
111111
model: models.Model | models.KnownModelName | str | None = None,
112112
deps: AgentDepsT = None,
@@ -210,25 +210,28 @@ async def main():
210210
def override(
211211
self,
212212
*,
213+
name: str | _utils.Unset = _utils.UNSET,
213214
deps: AgentDepsT | _utils.Unset = _utils.UNSET,
214215
model: models.Model | models.KnownModelName | str | _utils.Unset = _utils.UNSET,
215216
toolsets: Sequence[AbstractToolset[AgentDepsT]] | _utils.Unset = _utils.UNSET,
216217
tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] | _utils.Unset = _utils.UNSET,
217218
instructions: Instructions[AgentDepsT] | _utils.Unset = _utils.UNSET,
218219
) -> Iterator[None]:
219-
"""Context manager to temporarily override agent dependencies, model, toolsets, tools, or instructions.
220+
"""Context manager to temporarily override agent name, dependencies, model, toolsets, tools, or instructions.
220221
221222
This is particularly useful when testing.
222223
You can find an example of this [here](../testing.md#overriding-model-via-pytest-fixtures).
223224
224225
Args:
226+
name: The name to use instead of the name passed to the agent constructor and agent run.
225227
deps: The dependencies to use instead of the dependencies passed to the agent run.
226228
model: The model to use instead of the model passed to the agent run.
227229
toolsets: The toolsets to use instead of the toolsets passed to the agent constructor and agent run.
228230
tools: The tools to use instead of the tools registered with the agent.
229231
instructions: The instructions to use instead of the instructions registered with the agent.
230232
"""
231233
with self.wrapped.override(
234+
name=name,
232235
deps=deps,
233236
model=model,
234237
toolsets=toolsets,

0 commit comments

Comments
 (0)