Skip to content
Merged
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 98 additions & 0 deletions docs/.hooks/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ def on_page_markdown(markdown: str, page: Page, config: Config, files: Files) ->
markdown = replace_uv_python_run(markdown)
markdown = render_examples(markdown)
markdown = render_video(markdown)
markdown = create_gateway_toggle(markdown)
return markdown


Expand All @@ -39,6 +40,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)
Expand Down Expand Up @@ -115,3 +117,99 @@ def sub_cf_video(m: re.Match[str]) -> str:
></iframe>
</div>
"""


def create_gateway_toggle(markdown: str) -> 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'```python(?:\s+\{([^}]*)\})?\n(.*?)\n```(\n\n(?:\d+\..+(?:\n(?: .+|\n))*)+)?',
transform_gateway_code_block,
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]) -> 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
attrs_str = f' {{{attrs}}}' if attrs else ''

# 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 leading 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)

return f"""\
=== "Pydantic AI Gateway"
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this really be the first tab? As much as we want to drive people to the Gateway, most people won't, and this may turn them off.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we add a link to the gateway docs?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Kludex did you manage to add the disclaimers? If not, I'm happy to do it.


```python{attrs_str}
{indented_gateway_code}
```{indented_annotations}

=== "Pydantic AI"

```python{attrs_str}
{indented_code}
```{indented_annotations}"""