Merge remote-tracking branch 'upstream/main' into xai-sdk

Author: brightsparc

.github/workflows/ci.yml

@@ -202,7 +202,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10", "3.11", "3.12", "3.13"]
+        # TODO(Marcelo): Enable 3.11 again.
+        python-version: ["3.10", "3.12", "3.13"]
     env:
       CI: true
       COVERAGE_PROCESS_START: ./pyproject.toml

docs/.hooks/main.py

@@ -16,10 +16,12 @@
 
 def on_page_markdown(markdown: str, page: Page, config: Config, files: Files) -> str:
     """Called on each file after it is read and before it is converted to HTML."""
-    markdown = inject_snippets(markdown, (DOCS_ROOT / page.file.src_uri).parent)
+    relative_path = DOCS_ROOT / page.file.src_uri
+    markdown = inject_snippets(markdown, relative_path.parent)
     markdown = replace_uv_python_run(markdown)
     markdown = render_examples(markdown)
     markdown = render_video(markdown)
+    markdown = create_gateway_toggle(markdown, relative_path)
     return markdown
 
 
@@ -39,6 +41,7 @@ def on_env(env: Environment, config: Config, files: Files) -> Environment:
 
 def on_post_build(config: Config) -> None:
     """Inject extra CSS into mermaid styles to avoid titles being the same color as the background in dark mode."""
+    assert bundle_path is not None
     if bundle_path.exists():
         content = bundle_path.read_text()
         content, _ = re.subn(r'}(\.statediagram)', '}.statediagramTitleText{fill:#888}\1', content, count=1)
@@ -115,3 +118,109 @@ def sub_cf_video(m: re.Match[str]) -> str:
   ></iframe>
 </div>
 """
+
+
+def create_gateway_toggle(markdown: str, relative_path: Path) -> str:
+    """Transform Python code blocks with Agent() calls to show both Pydantic AI and Gateway versions."""
+    # Pattern matches Python code blocks with or without attributes, and optional annotation definitions after
+    # Annotation definitions are numbered list items like "1. Some text" that follow the code block
+    return re.sub(
+        r'```py(?:thon)?(?: *\{?([^}\n]*)\}?)?\n(.*?)\n```(\n\n(?:\d+\..+?\n)+?\n)?',
+        lambda m: transform_gateway_code_block(m, relative_path),
+        markdown,
+        flags=re.MULTILINE | re.DOTALL,
+    )
+
+
+# Models that should get gateway transformation
+GATEWAY_MODELS = ('anthropic', 'openai', 'openai-responses', 'openai-chat', 'bedrock', 'google-vertex', 'groq')
+
+
+def transform_gateway_code_block(m: re.Match[str], relative_path: Path) -> str:
+    """Transform a single code block to show both versions if it contains Agent() calls."""
+    attrs = m.group(1) or ''
+    code = m.group(2)
+    annotations = m.group(3) or ''  # Capture annotation definitions if present
+
+    # Simple check: does the code contain both "Agent(" and a quoted string?
+    if 'Agent(' not in code:
+        attrs_str = f' {{{attrs}}}' if attrs else ''
+        return f'```python{attrs_str}\n{code}\n```{annotations}'
+
+    # Check if code contains Agent() with a model that should be transformed
+    # Look for Agent(...'model:...' or Agent(..."model:..."
+    agent_pattern = r'Agent\((?:(?!["\']).)*([\"\'])([^"\']+)\1'
+    agent_match = re.search(agent_pattern, code, flags=re.DOTALL)
+
+    if not agent_match:
+        # No Agent() with string literal found
+        attrs_str = f' {{{attrs}}}' if attrs else ''
+        return f'```python{attrs_str}\n{code}\n```{annotations}'
+
+    model_string = agent_match.group(2)
+    # Check if model starts with one of the gateway-supported models
+    should_transform = any(model_string.startswith(f'{model}:') for model in GATEWAY_MODELS)
+
+    if not should_transform:
+        # Model doesn't match gateway models, return original
+        attrs_str = f' {{{attrs}}}' if attrs else ''
+        return f'```python{attrs_str}\n{code}\n```{annotations}'
+
+    # Transform the code for gateway version
+    def replace_agent_model(match: re.Match[str]) -> str:
+        """Replace model string with gateway/ prefix."""
+        full_match = match.group(0)
+        quote = match.group(1)
+        model = match.group(2)
+
+        # Replace the model string while preserving the rest
+        return full_match.replace(f'{quote}{model}{quote}', f'{quote}gateway/{model}{quote}', 1)
+
+    # This pattern finds: "Agent(" followed by anything (lazy), then the first quoted string
+    gateway_code = re.sub(
+        agent_pattern,
+        replace_agent_model,
+        code,
+        flags=re.DOTALL,
+    )
+
+    # Build attributes string
+    docs_path = DOCS_ROOT / 'gateway'
+
+    relative_path_to_gateway = docs_path.relative_to(relative_path, walk_up=True)
+    link = f"<a href='{relative_path_to_gateway}' style='float: right;'>Learn about Gateway</a>"
+    attrs_str = f' {{{attrs}}}' if attrs else ''
+
+    if 'title="' in attrs:
+        gateway_attrs = attrs.replace('title="', f'title="{link} ', 1)
+    else:
+        gateway_attrs = attrs + f' title="{link}"'
+    gateway_attrs_str = f' {{{gateway_attrs}}}'
+
+    # Indent code lines for proper markdown formatting within tabs
+    # Always add 4 spaces to every line (even empty ones) to preserve annotations
+    code_lines = code.split('\n')
+    indented_code = '\n'.join('    ' + line for line in code_lines)
+
+    gateway_code_lines = gateway_code.split('\n')
+    indented_gateway_code = '\n'.join('    ' + line for line in gateway_code_lines)
+
+    # Indent annotation definitions if present (need to be inside tabs for Material to work)
+    indented_annotations = ''
+    if annotations:
+        # Remove surrounding newlines and indent each line with 4 spaces
+        annotation_lines = annotations.strip().split('\n')
+        indented_annotations = '\n\n' + '\n'.join('    ' + line for line in annotation_lines) + '\n\n'
+
+    return f"""\
+=== "With Pydantic AI Gateway"
+
+    ```python{gateway_attrs_str}
+{indented_gateway_code}
+    ```{indented_annotations}
+
+=== "Directly to Provider API"
+
+    ```python{attrs_str}
+{indented_code}
+    ```{indented_annotations}"""

docs/.overrides/main.html

@@ -0,0 +1,8 @@
+{% extends "base.html" %}
+
+{% block announce %}
+    <strong>
+        <a href="/gateway">Pydantic AI Gateway</a> is now available! 🚀
+        Enterprise-ready AI model routing: One key for all your models with real-time monitoring and budget control that works.
+    </strong>
+{% endblock %}

docs/agents.md

@@ -320,7 +320,8 @@ async def main():
                         content='What is the capital of France?',
                         timestamp=datetime.datetime(...),
                     )
-                ]
+                ],
+                run_id='...',
             )
         ),
         CallToolsNode(
@@ -329,6 +330,7 @@ async def main():
                 usage=RequestUsage(input_tokens=56, output_tokens=7),
                 model_name='gpt-5',
                 timestamp=datetime.datetime(...),
+                run_id='...',
             )
         ),
         End(data=FinalResult(output='The capital of France is Paris.')),
@@ -382,7 +384,8 @@ async def main():
                             content='What is the capital of France?',
                             timestamp=datetime.datetime(...),
                         )
-                    ]
+                    ],
+                    run_id='...',
                 )
             ),
             CallToolsNode(
@@ -391,6 +394,7 @@ async def main():
                     usage=RequestUsage(input_tokens=56, output_tokens=7),
                     model_name='gpt-5',
                     timestamp=datetime.datetime(...),
+                    run_id='...',
                 )
             ),
             End(data=FinalResult(output='The capital of France is Paris.')),
@@ -1044,7 +1048,8 @@ with capture_run_messages() as messages:  # (2)!
                         content='Please get me the volume of a box with size 6.',
                         timestamp=datetime.datetime(...),
                     )
-                ]
+                ],
+                run_id='...',
             ),
             ModelResponse(
                 parts=[
@@ -1057,6 +1062,7 @@ with capture_run_messages() as messages:  # (2)!
                 usage=RequestUsage(input_tokens=62, output_tokens=4),
                 model_name='gpt-5',
                 timestamp=datetime.datetime(...),
+                run_id='...',
             ),
             ModelRequest(
                 parts=[
@@ -1066,7 +1072,8 @@ with capture_run_messages() as messages:  # (2)!
                         tool_call_id='pyd_ai_tool_call_id',
                         timestamp=datetime.datetime(...),
                     )
-                ]
+                ],
+                run_id='...',
             ),
             ModelResponse(
                 parts=[
@@ -1079,6 +1086,7 @@ with capture_run_messages() as messages:  # (2)!
                 usage=RequestUsage(input_tokens=72, output_tokens=8),
                 model_name='gpt-5',
                 timestamp=datetime.datetime(...),
+                run_id='...',
             ),
         ]
         """

docs/api/models/function.md

@@ -29,7 +29,8 @@ async def model_function(
                     content='Testing my agent...',
                     timestamp=datetime.datetime(...),
                 )
-            ]
+            ],
+            run_id='...',
         )
     ]
     """