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_agent.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.

DouweM · 2025-11-19T21:03:26Z

pydantic_ai_slim/pydantic_ai/models/__init__.py

        'grok:grok-4',
        'grok:grok-4-0709',
+        'grok:grok-4-fast-non-reasoning',
+        'grok:grok-4-fast-reasoning',


grok-4-fast as well?

grok-4-fast is an alias for grok-4-fast-reasoning - how should we handle those, do we need all aliases defined?

pydantic_ai_slim/pydantic_ai/profiles/grok.py

DouweM · 2025-11-19T21:45:09Z

pydantic_ai_slim/pydantic_ai/models/xai.py

+        usage = self._map_usage(response)
+
+        # Map finish reason
+        finish_reason_map = {


Please move this mapping to a typed constant as we have it on other models

DouweM · 2025-11-19T21:47:32Z

pydantic_ai_slim/pydantic_ai/models/xai.py

+        if not hasattr(response, 'usage'):
+            return RequestUsage()
+
+        usage_obj = getattr(response, 'usage', None)


All this usage of has/getattr confuses me; shouldn't the Response type have these fields so we can read them directly?

DouweM · 2025-11-19T21:48:35Z

tests/models/mock_xai.py

I'd rather not do any mocking and test everything against the (recorded) API.

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

pydantic_ai_slim/pydantic_ai/models/xai.py

+        if model_request_parameters.builtin_tools:
+            tools.extend(self._get_builtin_tools(model_request_parameters))
+        if model_request_parameters.tool_defs:
+            tools.extend(self._map_tools(model_request_parameters))


Does the API support tool_choice? Look at the logic OpenAIChatModel uses to set it; we should do that here as well.

Yes it does, I can add support for that.

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

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 9 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

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

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

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