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

examples/pydantic_ai_examples/stock_analysis.py

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

github-actions · 2025-12-09T14:00:56Z

This PR is stale, and will be closed in 3 days if no reply is received.

DouweM

@brightsparc Thanks for your patience here! Left some more feedback, and can you please look at the merge conflicts with main?

examples/pydantic_ai_examples/stock_analysis.py

DouweM · 2025-12-11T19:12:50Z

docs/examples/stock-analysis.md

Let's remove this example

DouweM · 2025-12-11T19:15:31Z

pydantic_ai_slim/pydantic_ai/models/xai.py

+from ..profiles import ModelProfileSpec
+from ..profiles.grok import GrokModelProfile
+from ..providers import Provider, infer_provider
+from ..providers.xai import XaiModelName


Let's move this inside try/except ImportError above as it also depends on the SDK

DouweM · 2025-12-11T19:16:54Z

pydantic_ai_slim/pydantic_ai/models/xai.py

+    See [xAI SDK documentation](https://docs.x.ai/docs) for more details on these parameters.
+    """
+
+    xai_logprobs: bool


We have logprobs on the base ModelSettings so no need to include it here.

Also note that we're not currently giving the user a way to access the logprobs. In OpenAI models, we do this using ModelResponse.provider_details['logprobs'] or TextPart.provider_details['logprobs'].

If we list xAI under supported providers for ModelSettings.logprobs, we should make sure we do that.

DouweM · 2025-12-11T19:18:08Z

pydantic_ai_slim/pydantic_ai/models/xai.py

+    Corresponds to the `web_search_call.action.sources` value of the `include` parameter in the Responses API.
+    """
+
+    xai_include_x_search_outputs: bool


x_search isn't supported yet so let's leave this out for now

DouweM · 2025-12-11T19:36:35Z

pydantic_ai_slim/pydantic_ai/models/xai.py

+
+            if tool_call.status == chat_types.chat_pb2.ToolCallStatus.TOOL_CALL_STATUS_COMPLETED:
+                # Tool completed - emit return part with result
+                tool_result_content = XaiModel.get_tool_result_content(response.content)


See above, it confuses me that response.content is both text, and the builtin tool return value

DouweM · 2025-12-11T19:36:51Z

pydantic_ai_slim/pydantic_ai/models/xai.py

+            if tool_call.status == chat_types.chat_pb2.ToolCallStatus.TOOL_CALL_STATUS_COMPLETED:
+                # Tool completed - emit return part with result
+                tool_result_content = XaiModel.get_tool_result_content(response.content)
+                return_part = BuiltinToolReturnPart(


As in non-streaming, we probably need a BuiltinToolCallPart as well

DouweM · 2025-12-11T19:38:04Z

pydantic_ai_slim/pydantic_ai/profiles/grok.py

+        # xAI supports JSON object output
+        supports_json_object_output=True,
+        # Default to 'native' for structured output since xAI supports it well
+        default_structured_output_mode='native',


Let's default to tool in line with other models (meaning we can drop this field as that's already the default)

DouweM · 2025-12-11T19:38:48Z

pydantic_ai_slim/pydantic_ai/providers/xai.py

+
+
+# https://docs.x.ai/docs/models
+XaiModelName = str | AllModels


With other providers, we have this in the model file instead, so let's move it just for consistency, unless there's a particular reason we cant

DouweM · 2025-12-11T19:39:41Z

pydantic_ai_slim/pydantic_ai/settings.py

    * OpenAI (some models, not o1)
    * Groq
    * Anthropic
+    * Grok


We can drop this one right

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 24 commits November 25, 2025 16:01

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

Addressing comments for tool output

1faabbc

Concat text parts without new lines

f91affc

Concat reasoning parts without new lines

f163ad6

Fix xai mappings

cb83948

Update map response parts for mcp

5656149

Refactor tool calling

ecffc4c

Fix to clean up builtin tool calls

5f57853

Merge from main

3d3a533

Remove redundant assertions

4440dac

Update to return text part ahead of tool return partn

5894cc8

Return empty dict for default tool call output

7400a13

Adding encrypted contents to assistant message

f653e3d

Fix pyright

167b44e

github-actions bot added the Stale label Dec 9, 2025

DouweM removed Stale awaiting author revision labels Dec 9, 2025

DouweM requested changes Dec 11, 2025

View reviewed changes

DouweM added the awaiting author revision label Dec 11, 2025



		# https://docs.x.ai/docs/models
		XaiModelName = str \| AllModels

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Dec 9, 2025

Uh oh!

DouweM left a comment

Choose a reason for hiding this comment

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!

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