pydantic-ai-slim 0.1.9__tar.gz → 0.1.11__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of pydantic-ai-slim might be problematic. Click here for more details.

Files changed (53) hide show
  1. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/PKG-INFO +6 -6
  2. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/README.md +2 -2
  3. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/__main__.py +2 -2
  4. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/_cli.py +30 -25
  5. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/exceptions.py +2 -2
  6. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/messages.py +1 -1
  7. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/_json_schema.py +1 -1
  8. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/anthropic.py +4 -2
  9. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/bedrock.py +1 -1
  10. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/cohere.py +1 -1
  11. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/function.py +1 -1
  12. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/gemini.py +28 -3
  13. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/groq.py +4 -2
  14. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/mistral.py +1 -1
  15. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/openai.py +7 -3
  16. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/providers/__init__.py +1 -1
  17. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/providers/azure.py +1 -1
  18. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/settings.py +10 -0
  19. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pyproject.toml +8 -2
  20. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/.gitignore +0 -0
  21. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/__init__.py +0 -0
  22. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/_agent_graph.py +0 -0
  23. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/_griffe.py +0 -0
  24. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/_output.py +0 -0
  25. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/_parts_manager.py +0 -0
  26. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/_pydantic.py +0 -0
  27. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/_system_prompt.py +0 -0
  28. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/_utils.py +0 -0
  29. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/agent.py +0 -0
  30. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/common_tools/__init__.py +0 -0
  31. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/common_tools/duckduckgo.py +0 -0
  32. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/common_tools/tavily.py +0 -0
  33. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/format_as_xml.py +0 -0
  34. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/format_prompt.py +0 -0
  35. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/mcp.py +0 -0
  36. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/__init__.py +0 -0
  37. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/fallback.py +0 -0
  38. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/instrumented.py +0 -0
  39. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/test.py +0 -0
  40. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/models/wrapper.py +0 -0
  41. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/providers/anthropic.py +0 -0
  42. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/providers/bedrock.py +0 -0
  43. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/providers/cohere.py +0 -0
  44. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/providers/deepseek.py +0 -0
  45. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/providers/google_gla.py +0 -0
  46. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/providers/google_vertex.py +0 -0
  47. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/providers/groq.py +0 -0
  48. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/providers/mistral.py +0 -0
  49. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/providers/openai.py +0 -0
  50. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/py.typed +0 -0
  51. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/result.py +0 -0
  52. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/tools.py +0 -0
  53. {pydantic_ai_slim-0.1.9 → pydantic_ai_slim-0.1.11}/pydantic_ai/usage.py +0 -0
@@ -1,8 +1,8 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pydantic-ai-slim
3
- Version: 0.1.9
3
+ Version: 0.1.11
4
4
  Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
5
- Author-email: Samuel Colvin <samuel@pydantic.dev>
5
+ Author-email: Samuel Colvin <samuel@pydantic.dev>, Marcelo Trylesinski <marcelotryle@gmail.com>, David Montague <david@pydantic.dev>, Alex Hall <alex@pydantic.dev>
6
6
  License-Expression: MIT
7
7
  Classifier: Development Status :: 4 - Beta
8
8
  Classifier: Environment :: Console
@@ -29,7 +29,7 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
29
29
  Requires-Dist: griffe>=1.3.2
30
30
  Requires-Dist: httpx>=0.27
31
31
  Requires-Dist: opentelemetry-api>=1.28.0
32
- Requires-Dist: pydantic-graph==0.1.9
32
+ Requires-Dist: pydantic-graph==0.1.11
33
33
  Requires-Dist: pydantic>=2.10
34
34
  Requires-Dist: typing-inspection>=0.4.0
35
35
  Provides-Extra: anthropic
@@ -45,7 +45,7 @@ Requires-Dist: cohere>=5.13.11; (platform_system != 'Emscripten') and extra == '
45
45
  Provides-Extra: duckduckgo
46
46
  Requires-Dist: duckduckgo-search>=7.0.0; extra == 'duckduckgo'
47
47
  Provides-Extra: evals
48
- Requires-Dist: pydantic-evals==0.1.9; extra == 'evals'
48
+ Requires-Dist: pydantic-evals==0.1.11; extra == 'evals'
49
49
  Provides-Extra: groq
50
50
  Requires-Dist: groq>=0.15.0; extra == 'groq'
51
51
  Provides-Extra: logfire
@@ -67,8 +67,8 @@ Description-Content-Type: text/markdown
67
67
 
68
68
  [![CI](https://github.com/pydantic/pydantic-ai/actions/workflows/ci.yml/badge.svg?event=push)](https://github.com/pydantic/pydantic-ai/actions/workflows/ci.yml?query=branch%3Amain)
69
69
  [![Coverage](https://coverage-badge.samuelcolvin.workers.dev/pydantic/pydantic-ai.svg)](https://coverage-badge.samuelcolvin.workers.dev/redirect/pydantic/pydantic-ai)
70
- [![PyPI](https://img.shields.io/pypi/v/pydantic-ai.svg)](https://pypi.python.org/pypi/pydantic-ai)
71
- [![versions](https://img.shields.io/pypi/pyversions/pydantic-ai.svg)](https://github.com/pydantic/pydantic-ai)
70
+ [![PyPI](https://img.shields.io/pypi/v/pydantic-ai-slim.svg)](https://pypi.python.org/pypi/pydantic-ai-slim)
71
+ [![versions](https://img.shields.io/pypi/pyversions/pydantic-ai-slim.svg)](https://github.com/pydantic/pydantic-ai)
72
72
  [![license](https://img.shields.io/github/license/pydantic/pydantic-ai.svg?v)](https://github.com/pydantic/pydantic-ai/blob/main/LICENSE)
73
73
 
74
74
  PydanticAI core logic with minimal required dependencies.
@@ -2,8 +2,8 @@
2
2
 
3
3
  [![CI](https://github.com/pydantic/pydantic-ai/actions/workflows/ci.yml/badge.svg?event=push)](https://github.com/pydantic/pydantic-ai/actions/workflows/ci.yml?query=branch%3Amain)
4
4
  [![Coverage](https://coverage-badge.samuelcolvin.workers.dev/pydantic/pydantic-ai.svg)](https://coverage-badge.samuelcolvin.workers.dev/redirect/pydantic/pydantic-ai)
5
- [![PyPI](https://img.shields.io/pypi/v/pydantic-ai.svg)](https://pypi.python.org/pypi/pydantic-ai)
6
- [![versions](https://img.shields.io/pypi/pyversions/pydantic-ai.svg)](https://github.com/pydantic/pydantic-ai)
5
+ [![PyPI](https://img.shields.io/pypi/v/pydantic-ai-slim.svg)](https://pypi.python.org/pypi/pydantic-ai-slim)
6
+ [![versions](https://img.shields.io/pypi/pyversions/pydantic-ai-slim.svg)](https://github.com/pydantic/pydantic-ai)
7
7
  [![license](https://img.shields.io/github/license/pydantic/pydantic-ai.svg?v)](https://github.com/pydantic/pydantic-ai/blob/main/LICENSE)
8
8
 
9
9
  PydanticAI core logic with minimal required dependencies.
@@ -1,6 +1,6 @@
1
1
  """This means `python -m pydantic_ai` should run the CLI."""
2
2
 
3
- from ._cli import app
3
+ from ._cli import cli_exit
4
4
 
5
5
  if __name__ == '__main__':
6
- app()
6
+ cli_exit()
@@ -7,16 +7,16 @@ from asyncio import CancelledError
7
7
  from collections.abc import Sequence
8
8
  from contextlib import ExitStack
9
9
  from datetime import datetime, timezone
10
- from importlib.metadata import version
11
10
  from pathlib import Path
12
11
  from typing import Any, cast
13
12
 
14
13
  from typing_inspection.introspection import get_literal_values
15
14
 
16
- from pydantic_ai.agent import Agent
17
- from pydantic_ai.exceptions import UserError
18
- from pydantic_ai.messages import ModelMessage
19
- from pydantic_ai.models import KnownModelName, infer_model
15
+ from . import __version__
16
+ from .agent import Agent
17
+ from .exceptions import UserError
18
+ from .messages import ModelMessage
19
+ from .models import KnownModelName, infer_model
20
20
 
21
21
  try:
22
22
  import argcomplete
@@ -39,7 +39,7 @@ except ImportError as _import_error:
39
39
  ) from _import_error
40
40
 
41
41
 
42
- __version__ = version('pydantic-ai-slim')
42
+ __all__ = 'cli', 'cli_exit'
43
43
 
44
44
 
45
45
  class SimpleCodeBlock(CodeBlock):
@@ -83,14 +83,20 @@ The current date and time is {datetime.now()} {tzname}.
83
83
  The user is running {sys.platform}."""
84
84
 
85
85
 
86
- def cli(args_list: Sequence[str] | None = None) -> int:
86
+ def cli_exit(prog_name: str = 'pai'): # pragma: no cover
87
+ """Run the CLI and exit."""
88
+ sys.exit(cli(prog_name=prog_name))
89
+
90
+
91
+ def cli(args_list: Sequence[str] | None = None, *, prog_name: str = 'pai') -> int:
92
+ """Run the CLI and return the exit code for the process."""
87
93
  parser = argparse.ArgumentParser(
88
- prog='pai',
94
+ prog=prog_name,
89
95
  description=f"""\
90
96
  PydanticAI CLI v{__version__}\n\n
91
97
 
92
- Special prompt:
93
- * `/exit` - exit the interactive mode
98
+ Special prompts:
99
+ * `/exit` - exit the interactive mode (ctrl-c and ctrl-d also work)
94
100
  * `/markdown` - show the last markdown output of the last question
95
101
  * `/multiline` - toggle multiline mode
96
102
  """,
@@ -101,7 +107,7 @@ Special prompt:
101
107
  '-m',
102
108
  '--model',
103
109
  nargs='?',
104
- help='Model to use, in format "<provider>:<model>" e.g. "openai:gpt-4o". Defaults to "openai:gpt-4o".',
110
+ help='Model to use, in format "<provider>:<model>" e.g. "openai:gpt-4o" or "anthropic:claude-3-7-sonnet-latest". Defaults to "openai:gpt-4o".',
105
111
  default='openai:gpt-4o',
106
112
  )
107
113
  # we don't want to autocomplete or list models that don't include the provider,
@@ -118,10 +124,10 @@ Special prompt:
118
124
  '-t',
119
125
  '--code-theme',
120
126
  nargs='?',
121
- help='Which colors to use for code, can be "dark", "light" or any theme from pygments.org/styles/. Defaults to "monokai".',
122
- default='monokai',
127
+ help='Which colors to use for code, can be "dark", "light" or any theme from pygments.org/styles/. Defaults to "dark" which works well on dark terminals.',
128
+ default='dark',
123
129
  )
124
- parser.add_argument('--no-stream', action='store_true', help='Whether to stream responses from the model')
130
+ parser.add_argument('--no-stream', action='store_true', help='Disable streaming from the model')
125
131
  parser.add_argument('--version', action='store_true', help='Show version and exit')
126
132
 
127
133
  argcomplete.autocomplete(parser)
@@ -129,7 +135,8 @@ Special prompt:
129
135
 
130
136
  console = Console()
131
137
  console.print(
132
- f'[green]pai - PydanticAI CLI v{__version__} using[/green] [magenta]{args.model}[/magenta]', highlight=False
138
+ f'[green]{prog_name} - PydanticAI CLI v{__version__} using[/green] [magenta]{args.model}[/magenta]',
139
+ highlight=False,
133
140
  )
134
141
  if args.version:
135
142
  return 0
@@ -160,23 +167,25 @@ Special prompt:
160
167
  pass
161
168
  return 0
162
169
 
163
- history = Path.home() / '.pai-prompt-history.txt'
170
+ history = Path.home() / f'.{prog_name}-prompt-history.txt'
164
171
  # doing this instead of `PromptSession[Any](history=` allows mocking of PromptSession in tests
165
172
  session: PromptSession[Any] = PromptSession(history=FileHistory(str(history)))
166
173
  try:
167
- return asyncio.run(run_chat(session, stream, cli_agent, console, code_theme))
174
+ return asyncio.run(run_chat(session, stream, cli_agent, console, code_theme, prog_name))
168
175
  except KeyboardInterrupt: # pragma: no cover
169
176
  return 0
170
177
 
171
178
 
172
- async def run_chat(session: PromptSession[Any], stream: bool, agent: Agent, console: Console, code_theme: str) -> int:
179
+ async def run_chat(
180
+ session: PromptSession[Any], stream: bool, agent: Agent, console: Console, code_theme: str, prog_name: str
181
+ ) -> int:
173
182
  multiline = False
174
183
  messages: list[ModelMessage] = []
175
184
 
176
185
  while True:
177
186
  try:
178
187
  auto_suggest = CustomAutoSuggest(['/markdown', '/multiline', '/exit'])
179
- text = await session.prompt_async('pai ➤ ', auto_suggest=auto_suggest, multiline=multiline)
188
+ text = await session.prompt_async(f'{prog_name} ➤ ', auto_suggest=auto_suggest, multiline=multiline)
180
189
  except (KeyboardInterrupt, EOFError): # pragma: no cover
181
190
  return 0
182
191
 
@@ -214,14 +223,14 @@ async def ask_agent(
214
223
 
215
224
  with status, ExitStack() as stack:
216
225
  async with agent.iter(prompt, message_history=messages) as agent_run:
217
- live = Live('', refresh_per_second=15, console=console, vertical_overflow='visible')
226
+ live = Live('', refresh_per_second=15, console=console, vertical_overflow='ellipsis')
218
227
  async for node in agent_run:
219
228
  if Agent.is_model_request_node(node):
220
229
  async with node.stream(agent_run.ctx) as handle_stream:
221
230
  status.stop() # stopping multiple times is idempotent
222
231
  stack.enter_context(live) # entering multiple times is idempotent
223
232
 
224
- async for content in handle_stream.stream_output():
233
+ async for content in handle_stream.stream_output(debounce_by=None):
225
234
  live.update(Markdown(content, code_theme=code_theme))
226
235
 
227
236
  assert agent_run.result is not None
@@ -282,7 +291,3 @@ def handle_slash_command(
282
291
  else:
283
292
  console.print(f'[red]Unknown command[/red] [magenta]`{ident_prompt}`[/magenta]')
284
293
  return None, multiline
285
-
286
-
287
- def app(): # pragma: no cover
288
- sys.exit(cli())
@@ -3,9 +3,9 @@ from __future__ import annotations as _annotations
3
3
  import json
4
4
  import sys
5
5
 
6
- if sys.version_info < (3, 11): # pragma: no cover
6
+ if sys.version_info < (3, 11):
7
7
  from exceptiongroup import ExceptionGroup
8
- else: # pragma: no cover
8
+ else:
9
9
  ExceptionGroup = ExceptionGroup
10
10
 
11
11
  __all__ = (
@@ -81,7 +81,7 @@ class VideoUrl:
81
81
  """Type identifier, this is available on all parts as a discriminator."""
82
82
 
83
83
  @property
84
- def media_type(self) -> VideoMediaType: # pragma: no cover
84
+ def media_type(self) -> VideoMediaType: # pragma: lax no cover
85
85
  """Return the media type of the video, based on the url."""
86
86
  if self.url.endswith('.mkv'):
87
87
  return 'video/x-matroska'
@@ -99,7 +99,7 @@ class WalkJsonSchema(ABC):
99
99
  if (additional_properties := schema.get('additionalProperties')) is not None:
100
100
  if isinstance(additional_properties, bool):
101
101
  schema['additionalProperties'] = additional_properties
102
- else: # pragma: no cover
102
+ else:
103
103
  schema['additionalProperties'] = self._handle(additional_properties)
104
104
 
105
105
  if (pattern_properties := schema.get('patternProperties')) is not None:
@@ -90,7 +90,7 @@ See [the Anthropic docs](https://docs.anthropic.com/en/docs/about-claude/models)
90
90
  """
91
91
 
92
92
 
93
- class AnthropicModelSettings(ModelSettings):
93
+ class AnthropicModelSettings(ModelSettings, total=False):
94
94
  """Settings used for an Anthropic model request.
95
95
 
96
96
  ALL FIELDS MUST BE `anthropic_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
@@ -221,6 +221,8 @@ class AnthropicModel(Model):
221
221
  system_prompt, anthropic_messages = await self._map_message(messages)
222
222
 
223
223
  try:
224
+ extra_headers = model_settings.get('extra_headers', {})
225
+ extra_headers.setdefault('User-Agent', get_user_agent())
224
226
  return await self.client.messages.create(
225
227
  max_tokens=model_settings.get('max_tokens', 1024),
226
228
  system=system_prompt or NOT_GIVEN,
@@ -234,7 +236,7 @@ class AnthropicModel(Model):
234
236
  top_p=model_settings.get('top_p', NOT_GIVEN),
235
237
  timeout=model_settings.get('timeout', NOT_GIVEN),
236
238
  metadata=model_settings.get('anthropic_metadata', NOT_GIVEN),
237
- extra_headers={'User-Agent': get_user_agent()},
239
+ extra_headers=extra_headers,
238
240
  extra_body=model_settings.get('extra_body'),
239
241
  )
240
242
  except APIStatusError as e:
@@ -355,7 +355,7 @@ class BedrockConverseModel(Model):
355
355
 
356
356
  if max_tokens := model_settings.get('max_tokens'):
357
357
  inference_config['maxTokens'] = max_tokens
358
- if temperature := model_settings.get('temperature'):
358
+ if (temperature := model_settings.get('temperature')) is not None:
359
359
  inference_config['temperature'] = temperature
360
360
  if top_p := model_settings.get('top_p'):
361
361
  inference_config['topP'] = top_p
@@ -78,7 +78,7 @@ See [Cohere's docs](https://docs.cohere.com/v2/docs/models) for a list of all av
78
78
  """
79
79
 
80
80
 
81
- class CohereModelSettings(ModelSettings):
81
+ class CohereModelSettings(ModelSettings, total=False):
82
82
  """Settings used for a Cohere model request.
83
83
 
84
84
  ALL FIELDS MUST BE `cohere_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
@@ -282,7 +282,7 @@ def _estimate_string_tokens(content: str | Sequence[UserContent]) -> int:
282
282
  return 0
283
283
  if isinstance(content, str):
284
284
  return len(re.split(r'[\s",.:]+', content.strip()))
285
- else: # pragma: no cover
285
+ else:
286
286
  tokens = 0
287
287
  for part in content:
288
288
  if isinstance(part, str):
@@ -73,7 +73,7 @@ See [the Gemini API docs](https://ai.google.dev/gemini-api/docs/models/gemini#mo
73
73
  """
74
74
 
75
75
 
76
- class GeminiModelSettings(ModelSettings):
76
+ class GeminiModelSettings(ModelSettings, total=False):
77
77
  """Settings used for a Gemini model request.
78
78
 
79
79
  ALL FIELDS MUST BE `gemini_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
@@ -81,6 +81,18 @@ class GeminiModelSettings(ModelSettings):
81
81
 
82
82
  gemini_safety_settings: list[GeminiSafetySettings]
83
83
 
84
+ gemini_thinking_config: ThinkingConfig
85
+ """Thinking is "on" by default in both the API and AI Studio.
86
+
87
+ Being on by default doesn't mean the model will send back thoughts. For that, you would need to set `include_thoughts`
88
+ to `True`, but since end of January 2025, `thoughts` are not returned anymore, and are only displayed in the Google
89
+ AI Studio. See https://discuss.ai.google.dev/t/thoughts-are-missing-cot-not-included-anymore/63653 for more details.
90
+
91
+ If you want to avoid the model spending any tokens on thinking, you can set `thinking_budget` to `0`.
92
+
93
+ See more about it on <https://ai.google.dev/gemini-api/docs/thinking>.
94
+ """
95
+
84
96
 
85
97
  @dataclass(init=False)
86
98
  class GeminiModel(Model):
@@ -223,7 +235,9 @@ class GeminiModel(Model):
223
235
  generation_config['presence_penalty'] = presence_penalty
224
236
  if (frequency_penalty := model_settings.get('frequency_penalty')) is not None:
225
237
  generation_config['frequency_penalty'] = frequency_penalty
226
- if (gemini_safety_settings := model_settings.get('gemini_safety_settings')) != []:
238
+ if (thinkingConfig := model_settings.get('gemini_thinking_config')) is not None:
239
+ generation_config['thinking_config'] = thinkingConfig # pragma: no cover
240
+ if (gemini_safety_settings := model_settings.get('gemini_safety_settings')) is not None:
227
241
  request_data['safetySettings'] = gemini_safety_settings
228
242
  if generation_config:
229
243
  request_data['generationConfig'] = generation_config
@@ -497,6 +511,16 @@ class GeminiSafetySettings(TypedDict):
497
511
  """
498
512
 
499
513
 
514
+ class ThinkingConfig(TypedDict, total=False):
515
+ """The thinking features configuration."""
516
+
517
+ include_thoughts: Annotated[bool, pydantic.Field(alias='includeThoughts')]
518
+ """Indicates whether to include thoughts in the response. If true, thoughts are returned only if the model supports thought and thoughts are available."""
519
+
520
+ thinking_budget: Annotated[int, pydantic.Field(alias='thinkingBudget')]
521
+ """Indicates the thinking budget in tokens."""
522
+
523
+
500
524
  class _GeminiGenerationConfig(TypedDict, total=False):
501
525
  """Schema for an API request to the Gemini API.
502
526
 
@@ -511,6 +535,7 @@ class _GeminiGenerationConfig(TypedDict, total=False):
511
535
  presence_penalty: float
512
536
  frequency_penalty: float
513
537
  stop_sequences: list[str]
538
+ thinking_config: ThinkingConfig
514
539
 
515
540
 
516
541
  class _GeminiContent(TypedDict):
@@ -780,7 +805,7 @@ class _GeminiJsonSchema(WalkJsonSchema):
780
805
  additional_properties = schema.pop(
781
806
  'additionalProperties', None
782
807
  ) # don't pop yet so it's included in the warning
783
- if additional_properties: # pragma: no cover
808
+ if additional_properties:
784
809
  original_schema = {**schema, 'additionalProperties': additional_properties}
785
810
  warnings.warn(
786
811
  '`additionalProperties` is not supported by Gemini; it will be removed from the tool JSON schema.'
@@ -82,7 +82,7 @@ See <https://console.groq.com/docs/models> for an up to date date list of models
82
82
  """
83
83
 
84
84
 
85
- class GroqModelSettings(ModelSettings):
85
+ class GroqModelSettings(ModelSettings, total=False):
86
86
  """Settings used for a Groq model request.
87
87
 
88
88
  ALL FIELDS MUST BE `groq_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
@@ -200,6 +200,8 @@ class GroqModel(Model):
200
200
  groq_messages = self._map_messages(messages)
201
201
 
202
202
  try:
203
+ extra_headers = model_settings.get('extra_headers', {})
204
+ extra_headers.setdefault('User-Agent', get_user_agent())
203
205
  return await self.client.chat.completions.create(
204
206
  model=str(self._model_name),
205
207
  messages=groq_messages,
@@ -217,7 +219,7 @@ class GroqModel(Model):
217
219
  presence_penalty=model_settings.get('presence_penalty', NOT_GIVEN),
218
220
  frequency_penalty=model_settings.get('frequency_penalty', NOT_GIVEN),
219
221
  logit_bias=model_settings.get('logit_bias', NOT_GIVEN),
220
- extra_headers={'User-Agent': get_user_agent()},
222
+ extra_headers=extra_headers,
221
223
  extra_body=model_settings.get('extra_body'),
222
224
  )
223
225
  except APIStatusError as e:
@@ -91,7 +91,7 @@ Since [the Mistral docs](https://docs.mistral.ai/getting-started/models/models_o
91
91
  """
92
92
 
93
93
 
94
- class MistralModelSettings(ModelSettings):
94
+ class MistralModelSettings(ModelSettings, total=False):
95
95
  """Settings used for a Mistral model request.
96
96
 
97
97
  ALL FIELDS MUST BE `mistral_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
@@ -264,6 +264,8 @@ class OpenAIModel(Model):
264
264
  openai_messages = await self._map_messages(messages)
265
265
 
266
266
  try:
267
+ extra_headers = model_settings.get('extra_headers', {})
268
+ extra_headers.setdefault('User-Agent', get_user_agent())
267
269
  return await self.client.chat.completions.create(
268
270
  model=self._model_name,
269
271
  messages=openai_messages,
@@ -284,7 +286,7 @@ class OpenAIModel(Model):
284
286
  logit_bias=model_settings.get('logit_bias', NOT_GIVEN),
285
287
  reasoning_effort=model_settings.get('openai_reasoning_effort', NOT_GIVEN),
286
288
  user=model_settings.get('openai_user', NOT_GIVEN),
287
- extra_headers={'User-Agent': get_user_agent()},
289
+ extra_headers=extra_headers,
288
290
  extra_body=model_settings.get('extra_body'),
289
291
  )
290
292
  except APIStatusError as e:
@@ -610,6 +612,8 @@ class OpenAIResponsesModel(Model):
610
612
  reasoning = self._get_reasoning(model_settings)
611
613
 
612
614
  try:
615
+ extra_headers = model_settings.get('extra_headers', {})
616
+ extra_headers.setdefault('User-Agent', get_user_agent())
613
617
  return await self.client.responses.create(
614
618
  input=openai_messages,
615
619
  model=self._model_name,
@@ -625,7 +629,7 @@ class OpenAIResponsesModel(Model):
625
629
  timeout=model_settings.get('timeout', NOT_GIVEN),
626
630
  reasoning=reasoning,
627
631
  user=model_settings.get('openai_user', NOT_GIVEN),
628
- extra_headers={'User-Agent': get_user_agent()},
632
+ extra_headers=extra_headers,
629
633
  extra_body=model_settings.get('extra_body'),
630
634
  )
631
635
  except APIStatusError as e:
@@ -766,7 +770,7 @@ class OpenAIResponsesModel(Model):
766
770
  file_data=f'data:{item.media_type};base64,{base64_encoded}',
767
771
  )
768
772
  )
769
- elif isinstance(item, DocumentUrl): # pragma: no cover
773
+ elif isinstance(item, DocumentUrl):
770
774
  client = cached_async_http_client()
771
775
  response = await client.get(item.url)
772
776
  response.raise_for_status()
@@ -65,7 +65,7 @@ def infer_provider(provider: str) -> Provider[Any]:
65
65
 
66
66
  return GoogleGLAProvider()
67
67
  # NOTE: We don't test because there are many ways the `boto3.client` can retrieve the credentials.
68
- elif provider == 'bedrock': # pragma: no cover
68
+ elif provider == 'bedrock':
69
69
  from .bedrock import BedrockProvider
70
70
 
71
71
  return BedrockProvider()
@@ -82,7 +82,7 @@ class AzureProvider(Provider[AsyncOpenAI]):
82
82
  self._client = openai_client
83
83
  else:
84
84
  azure_endpoint = azure_endpoint or os.getenv('AZURE_OPENAI_ENDPOINT')
85
- if not azure_endpoint: # pragma: no cover
85
+ if not azure_endpoint:
86
86
  raise UserError(
87
87
  'Must provide one of the `azure_endpoint` argument or the `AZURE_OPENAI_ENDPOINT` environment variable'
88
88
  )
@@ -141,6 +141,16 @@ class ModelSettings(TypedDict, total=False):
141
141
  * Cohere
142
142
  """
143
143
 
144
+ extra_headers: dict[str, str]
145
+ """Extra headers to send to the model.
146
+
147
+ Supported by:
148
+
149
+ * OpenAI
150
+ * Anthropic
151
+ * Groq
152
+ """
153
+
144
154
  extra_body: object
145
155
  """Extra body to send to the model.
146
156
 
@@ -14,7 +14,12 @@ bump = true
14
14
  name = "pydantic-ai-slim"
15
15
  dynamic = ["version", "dependencies", "optional-dependencies"]
16
16
  description = "Agent Framework / shim to use Pydantic with LLMs, slim package"
17
- authors = [{ name = "Samuel Colvin", email = "samuel@pydantic.dev" }]
17
+ authors = [
18
+ { name = "Samuel Colvin", email = "samuel@pydantic.dev" },
19
+ { name = "Marcelo Trylesinski", email = "marcelotryle@gmail.com" },
20
+ { name = "David Montague", email = "david@pydantic.dev" },
21
+ { name = "Alex Hall", email = "alex@pydantic.dev" },
22
+ ]
18
23
  license = "MIT"
19
24
  readme = "README.md"
20
25
  classifiers = [
@@ -87,13 +92,14 @@ dev = [
87
92
  "pytest-recording>=0.13.2",
88
93
  "diff-cover>=9.2.0",
89
94
  "boto3-stubs[bedrock-runtime]",
95
+ "strict-no-cover>=0.1.1",
90
96
  ]
91
97
 
92
98
  [tool.hatch.metadata]
93
99
  allow-direct-references = true
94
100
 
95
101
  [project.scripts]
96
- pai = "pydantic_ai._cli:app"
102
+ pai = "pydantic_ai._cli:cli_exit" # TODO remove this when clai has been out for a while
97
103
 
98
104
  [tool.hatch.build.targets.wheel]
99
105
  packages = ["pydantic_ai"]