XaiModel with tests and stock analysis agent #3400

brightsparc · 2025-11-12T01:04:03Z

This PR includes support for GrokModel which uses the xai-sdk-python library, which supports chat completion, tool calling, and built-in server side tool calls.

I have included a series of tests, as well as an example for a stock analysis agent that uses search and code execution, and returns results in the form of a local tool call.

python examples/pydantic_ai_examples/stock_analysis_agent.py

🔍 Starting stock analysis...

Query: What was the hottest performing stock on NASDAQ yesterday?

18:32:37.095 stock_analysis_agent run
18:32:37.099   chat grok-4-fast
18:32:37.119     chat.stream grok-4-fast
🔧 Server-side tool: web_search
🔧 Server-side tool: browse_page

✅ Analysis complete!

📊 Top Stock: ATGL
💰 Current Price: $20.00

📈 Buy Analysis:
ATGL surged 139.84% on November 18, 2025, making it the hottest performer on NASDAQ, but pulled back 7.45% today to $20.00. The extreme volatility suggests high risk; without clear fundamental drivers, it's not recommended as a buy for most investors. Suitable only for speculative short-term trades.

📊 Usage Statistics:
   Requests: 1
   Input Tokens: 19150
   Output Tokens: 526
   Total Tokens: 19676
   Server-Side Tools: 7

Below is what the output looks like in logfire. It includes custom spans emmited from the xai-sdk.

Cost is coming back as Unknown so it would be good to understand how we can populate this.

DouweM · 2025-11-12T15:01:15Z

@brightsparc Thanks Julian! Let me know when this is ready for review or if you have any questions.

brightsparc · 2025-11-19T02:41:55Z

Due to using otel in our own SDK, I did notice an error when running the pydantic ai code in async.io loop:

Failed to detach context
Traceback (most recent call last):
  File "/Users/julian/workspace/poc/pydantic-ai/.venv/lib/python3.12/site-packages/opentelemetry/context/__init__.py", line 155, in detach
    _RUNTIME_CONTEXT.detach(token)
  File "/Users/julian/workspace/poc/pydantic-ai/.venv/lib/python3.12/site-packages/opentelemetry/context/contextvars_context.py", line 53, in detach
    self._current_context.reset(token)
ValueError: <Token var=<ContextVar name='current_context' default={} at 0x1055d33d0> at 0x109ca6e00> was created in a different Context

We could disable our tracing with XAI_SDK_DISABLE_TRACING=1, but would be good to understand why this happens.

DouweM · 2025-11-19T20:29:58Z

Cost is coming back as Unknown so it would be good to understand how we can populate this.

@brightsparc Looks like grok-4-fast is not on https://github.com/pydantic/genai-prices/blob/main/prices/providers/x_ai.yml yet. Contribution welcome! (See contrib docs)

Failing tests:

│  tests/test_examples.py             │  test_docs_examples[docs/models/grok.md:34]  │  113            │  256         │  ValueError      │

I believe we just need to add the xAI env var to the list here:

pydantic-ai/tests/test_examples.py

Lines 179 to 182 in c096d99

    
           env.set('MOONSHOTAI_API_KEY', 'testing') 
        
           env.set('DEEPSEEK_API_KEY', 'testing') 
        
           env.set('OVHCLOUD_API_KEY', 'testing') 
        
           env.set('PYDANTIC_AI_GATEWAY_API_KEY', 'testing')

│  tests/test_examples.py             │  test_docs_examples[docs/models/grok.md:56]  │  113            │              │                  │

This one should be fixed by running uv run pytest --update-examples tests/test_examples.py -k grok.md

│  tests/models/test_instrumented.py  │  test_instrumented_model_stream_break        │  484            │  512         │  AssertionError  │

Looks like we need the new logfire.exception.fingerprint there as well. Not 100% sure why we're suddenly seeing that...

Due to using otel in our own SDK, I did notice an error when running the pydantic ai code in async.io loop:

Do you have example code for me that triggers it?

brightsparc · 2025-11-19T20:46:41Z

Cost is coming back as Unknown so it would be good to understand how we can populate this.

@brightsparc Looks like grok-4-fast is not on https://github.com/pydantic/genai-prices/blob/main/prices/providers/x_ai.yml yet. Contribution welcome! (See contrib docs)

I will make a separate PR here

Failing tests:
│  tests/test_examples.py             │  test_docs_examples[docs/models/grok.md:34]  │  113            │  256         │  ValueError      │
I believe we just need to add the xAI env var to the list here:

pydantic-ai/tests/test_examples.py

Lines 179 to 182 in c096d99

env.set('MOONSHOTAI_API_KEY', 'testing')

env.set('DEEPSEEK_API_KEY', 'testing')

env.set('OVHCLOUD_API_KEY', 'testing')

env.set('PYDANTIC_AI_GATEWAY_API_KEY', 'testing')
│  tests/test_examples.py             │  test_docs_examples[docs/models/grok.md:56]  │  113            │              │                  │

Added this:

This one should be fixed by running uv run pytest --update-examples tests/test_examples.py -k grok.md
│  tests/models/test_instrumented.py  │  test_instrumented_model_stream_break        │  484            │  512         │  AssertionError  │

Not sure I follow this

Looks like we need the new logfire.exception.fingerprint there as well. Not 100% sure why we're suddenly seeing that...

Due to using otel in our own SDK, I did notice an error when running the pydantic ai code in async.io loop:

Do you have example code for me that triggers it?

You can run the following command to repo that error:

XAI_API_KEY="a-valid-key" uv run python examples/pydantic_ai_examples/stock_analysis_agent.py

docs/models/grok.md

DouweM · 2025-11-19T20:57:11Z

examples/pydantic_ai_examples/stock_analysis.py

If we're going to include this example, we should have a doc under docs/examples as well. I'm also OK with not including this example.

I've added a page in the docs for this example. I can remove it if you don't think it's worthwhile.

pydantic_ai_slim/pydantic_ai/models/__init__.py

pydantic_ai_slim/pydantic_ai/profiles/grok.py

pydantic_ai_slim/pydantic_ai/models/xai.py

tests/models/mock_xai.py

pydantic_ai_slim/pydantic_ai/models/xai.py

DouweM · 2025-11-19T21:50:46Z

pydantic_ai_slim/pydantic_ai/models/xai.py

+                xai_settings['frequency_penalty'] = model_settings['frequency_penalty']
+
+        # Create chat instance
+        chat = client.chat.create(model=self._model_name, messages=xai_messages, tools=tools_param, **xai_settings)


Just to confirm that the API does not support structured JSON output? If it does, we need to implement support for that as well: https://ai.pydantic.dev/output/#native-output

Yes it does, I can add support for that.

…d XaiProvider and move api_key and client to this

…level

…nale

…common _create_chat method

… of utc_now

DouweM · 2025-11-25T18:19:36Z

docs/api/models/grok.md

@@ -0,0 +1,7 @@
+# `pydantic_ai.models.grok`


We should add both new docs pages to mkdocs.yml so they show up in the navbar

docs/models/overview.md

docs/models/xai.md

pydantic_ai_slim/pydantic_ai/models/xai.py

tests/models/test_xai.py

DouweM · 2025-11-25T19:56:23Z

tests/models/test_xai.py

+    response_message = messages[-1]
+    assert isinstance(response_message, ModelResponse)
+
+    # Should have at least one BuiltinToolCallPart for MCP tools (prefixed with server_label, e.g. "linear.list_issues")


Instead of these types of tests with multiple assertions, I'd much rather test against the entire message history snapshot, so that we can see immediately whether it looks like what we expect -- not just that the subfields we test for have the right values

tests/models/test_xai.py

Initial commit for xAI grok model

311395d

brightsparc marked this pull request as draft November 12, 2025 01:04

Adding flight booking example for live test

afd8a4e

DouweM self-assigned this Nov 12, 2025

DouweM added the awaiting author revision label Nov 12, 2025

brightsparc added 12 commits November 17, 2025 18:56

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

ab2ff90

Fix pre-commit issues

4513c8f

Updated grok to support passing in AsyncClient, added initial tests

e60d3ea

Fix pyright

43e6890

Update grok all available grok models

4f15966

Adding image tests that require an active API_KEY

185a6be

Adding MCP added stock analysis agent

a632cbf

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

eb10f47

Updating docs

18b21a5

Set provider response id

b417485

Adding reasoning and usage

0f1c113

Update stock analytics agent to stream responses

2b58f59

brightsparc marked this pull request as ready for review November 19, 2025 02:36

brightsparc changed the title ~~Initial commit for xAI grok model~~ xAI grok model with tests and stock analysis agent Nov 19, 2025

brightsparc added 2 commits November 18, 2025 18:38

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

735da97

Remove new line for tool calls

65bb646

Adding XAI_API_KEY to test_examples.py

b19479e

DouweM requested changes Nov 19, 2025

View reviewed changes

brightsparc added 2 commits November 19, 2025 15:19

Adding grok-4-1-fast-non-reasoning as default

8268267

Update to mention GrokModelName in docs

843f588

brightsparc changed the title ~~xAI grok model with tests and stock analysis agent~~ XaiModel with tests and stock analysis agent Nov 20, 2025

brightsparc added 13 commits November 20, 2025 14:18

Renmae GrokModel to XaiModel, and update dependencies and tests. Adde…

99ad3aa

…d XaiProvider and move api_key and client to this

Remove use of hasattr

d933b2c

Fix RetryPromptPart

fe93dad

Add support for DocumentUrl, and downloading an dummy pdf with casette

5224e56

Implement first pass of response parts using openai chat completions …

128c05a

…level

Adding support for formatting builtin tools

b35bcfd

Update docs for xAI refence

59bf2b6

Minor tweak to tidy process response

e4ebcca

Adding prepare_request

f69a730

Updating live tests to use mocks, with a design document on the ratio…

052086b

…nale

Reduce duplication between request and request_stream, by creating a …

4521909

…common _create_chat method

Adding support for system instructions

c9326e1

Using system for provider_name, passing in response.timestamp instead…

a53badc

… of utc_now

DouweM requested changes Nov 25, 2025

View reviewed changes

brightsparc added 16 commits November 25, 2025 15:02

Add built in tool name and finish reason map

26744cc

Fix tests

ca70b96

Adding tool_choice and response_format arguments to create_chat

999eb81

Updating docs for xai.md

c335ca6

Update case of xAI

bcc7834

Adding tests for tool_choice and native json output

9d06c87

Updating docs link for models

df20aab

Create XaiModelName and deprecated GrokProvider

656eb62

Update builtin tests to verify output parts with snapshot

34035c0

Fix mcp_server tool with prefix and snapshot

5452ec1

Adding streaming tests for web_search and code_execution

bf6d4fd

Update remaining test to use snapshot

da5a14f

Adding tests for run_stream_events and added constants for model names

99c83f6

Merge from main

fb027c2

Updating documentation for stock_analysis

9ceb5df

Updating to unlress xai-sdk version

88f49af

XaiModel with tests and stock analysis agent #3400

Are you sure you want to change the base?

XaiModel with tests and stock analysis agent #3400

Uh oh!

Conversation

brightsparc commented Nov 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

DouweM commented Nov 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

brightsparc commented Nov 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

DouweM commented Nov 19, 2025

Uh oh!

brightsparc commented Nov 19, 2025

Uh oh!

Uh oh!

DouweM Nov 19, 2025

Choose a reason for hiding this comment

Uh oh!

brightsparc Nov 26, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

DouweM Nov 19, 2025

Choose a reason for hiding this comment

Uh oh!

brightsparc Nov 19, 2025

Choose a reason for hiding this comment

Uh oh!

DouweM Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

DouweM Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

brightsparc commented Nov 12, 2025 •

edited

Loading

DouweM commented Nov 12, 2025 •

edited

Loading

brightsparc commented Nov 19, 2025 •

edited

Loading