Feat/add xai model provider

[Cerrix]

Motivation

xAI's Grok models offer unique capabilities including server-side agentic tools (web_search, x_search, code_execution), reasoning models with visible or encrypted thinking, and high-performance inference. Developers building agents that need real-time web information, X/Twitter search, or advanced reasoning capabilities currently cannot use these features with Strands Agents.

This PR adds native support for xAI's Grok models using the official xai-sdk Python package, enabling developers to leverage the full xAI feature set within the Strands Agents framework.

Public API Changes

New xAIModel class for xAI integration:

from strands import Agent
from strands.models.xai import xAIModel

# Basic usage
model = xAIModel(
    client_args={"api_key": "your-api-key"},
    model_id="grok-4-1-fast-non-reasoning-latest",
)
agent = Agent(model=model)
result = agent("Hello!")

# With reasoning (grok-3-mini)
model = xAIModel(
    client_args={"api_key": "your-api-key"},
    model_id="grok-3-mini",
    reasoning_effort="high",  # "low" or "high"
)

# With encrypted reasoning for multi-turn context (grok-4)
model = xAIModel(
    client_args={"api_key": "your-api-key"},
    model_id="grok-4-fast-reasoning",
    use_encrypted_content=True,
)

# With server-side tools (executed by xAI)
from xai_sdk.tools import x_search, web_search

model = xAIModel(
    client_args={"api_key": "your-api-key"},
    model_id="grok-4-1-fast-non-reasoning-latest",
    xai_tools=[x_search(), web_search()],
)
agent = Agent(model=model)
result = agent("What are people saying about AI on X?")

# Hybrid: server-side + client-side tools
from strands import tool

@tool
def get_weather(city: str) -> str:
    """Get weather for a city."""
    return f"Weather in {city}: Sunny, 22°C"

model = xAIModel(
    client_args={"api_key": "your-api-key"},
    model_id="grok-4-1-fast-non-reasoning-latest",
    xai_tools=[x_search()],
)
agent = Agent(model=model, tools=[get_weather])
result = agent("What's the weather in Paris and what are people tweeting about it?")

Configuration Options

Parameter	Type	Description
model_id	str	Grok model ID (e.g., "grok-4", "grok-3-mini")
client_args	dict	Arguments for xAI client (api_key, timeout, etc.)
params	dict	Model parameters (temperature, max_tokens, etc.)
xai_tools	list	Server-side tools from xai_sdk.tools
reasoning_effort	str	"low" or "high" (grok-3-mini only)
use_encrypted_content	bool	Enable encrypted reasoning for multi-turn
include	list	Optional xAI features (e.g., ["inline_citations", "verbose_streaming"])

Installation

pip install strands-agents[xai]

Use Cases

• Real-time information: Use web_search or x_search server-side tools to get current information without implementing custom search tools
• Social media analysis: Query X/Twitter directly through x_search for sentiment analysis, trend monitoring, or content research
• Reasoning tasks: Leverage grok-3-mini's visible reasoning or grok-4's encrypted reasoning for complex problem-solving with full multi-turn context preservation
• Hybrid tool usage: Combine xAI's server-side tools with Strands' client-side tools in the same agent

Implementation Notes

The xAI SDK uses a gRPC-based API where server-side tools (x_search, web_search, code_execution) are executed by xAI's infrastructure and return encrypted results. These encrypted results cannot be reconstructed from plain text, which presents a challenge for multi-turn conversations.

To solve this, the provider serializes the xAI SDK's internal protobuf messages after each response and stores them in reasoningContent.redactedContent. This field was chosen because:

1. It's designed for encrypted/hidden content
2. It's NOT rendered in AgentResult.__str__(), keeping output clean for users
3. The Strands event loop already handles it properly

On subsequent turns, the provider extracts this state and restores the full conversation context, enabling seamless multi-turn conversations even with server-side tools.

Testing

• 63 unit tests covering all provider functionality
• 15 integration tests with real API calls covering:
  • Basic invocation (sync/async)
  • Streaming
  • Multi-turn conversations
  • Structured output
  • Image understanding
  • Reasoning models (grok-3-mini with visible reasoning)
  • Server-side tools (web_search)
  • Encrypted content preservation (grok-4 multi-turn)

Interactive Test Script

Copy the script below to manually test all xAI features after installing strands-agents[xai]:

# Set your API key
export XAI_API_KEY="your-api-key"

# Run the script (after copying to a file)
python test_xai.py

The script provides a menu with 10 different agent configurations:

1. Simple (non-streaming) - Basic agent without streaming
2. Streaming - Agent with streaming output
3. Streaming (DEBUG) - Shows all callback events for debugging
4. Strands tools - Client-side tools (calculate, weather)
5. X search - Server-side X/Twitter search
6. Hybrid - Both server-side (X search) and client-side tools
7. Reasoning (grok-3-mini) - Visible reasoning with reasoning_effort="high"
8. Reasoning encrypted (grok-4) - Encrypted reasoning for multi-turn context
9. Reasoning encrypted DEBUG - Debug mode for encrypted reasoning
10. Web search with citations - Web search with include=["inline_citations"]

Full test script code (click to expand)

"""Strands Agents - xAI Model Interactive Chat.

This script provides an interactive multi-turn chat with different xAI agent configurations.
Run with: python test_grok_strands.py
"""

import os
from typing import Any

from strands import Agent, tool
from strands.handlers.callback_handler import PrintingCallbackHandler
from strands.models.xai import xAIModel


def get_api_key() -> str:
    """Get XAI API key from environment."""
    api_key = os.getenv("XAI_API_KEY")
    if not api_key:
        raise ValueError("XAI_API_KEY environment variable not set")
    return api_key


# =============================================================================
# STRANDS TOOLS (CLIENT-SIDE)
# =============================================================================

@tool
def calculate(expression: str) -> str:
    """Evaluate a mathematical expression.
    
    Args:
        expression: A mathematical expression to evaluate (e.g., "2 + 2", "15 * 7")
        
    Returns:
        The result of the calculation
    """
    try:
        allowed = set("0123456789+-*/.() ")
        if all(c in allowed for c in expression):
            result = eval(expression)
            return f"Result: {result}"
        return "Error: Invalid expression"
    except Exception as e:
        return f"Error: {e}"


@tool
def get_weather(city: str) -> str:
    """Get the current weather for a city.
    
    Args:
        city: Name of the city
        
    Returns:
        Weather description for the city
    """
    weather_data = {
        "paris": "Sunny, 22°C",
        "london": "Cloudy, 15°C", 
        "tokyo": "Rainy, 18°C",
        "new york": "Clear, 20°C",
    }
    city_lower = city.lower()
    if city_lower in weather_data:
        return f"Weather in {city}: {weather_data[city_lower]}"
    return f"Weather in {city}: Partly cloudy, 20°C (default)"


# =============================================================================
# CALLBACK HANDLER
# =============================================================================

class StreamingCallbackHandler(PrintingCallbackHandler):
    """Callback handler that shows streaming text output with reasoning support."""
    
    def __init__(self, debug: bool = False, show_reasoning: bool = True):
        super().__init__(verbose_tool_use=True)
        self.debug = debug
        self.show_reasoning = show_reasoning
        self._in_reasoning = False
        self._reasoning_started = False
    
    def __call__(self, **kwargs: Any) -> None:
        # Handle reasoning content with visual styling
        if "reasoningText" in kwargs:
            if self.show_reasoning:
                if not self._reasoning_started:
                    print("\n💭 Thinking...", flush=True)
                    self._reasoning_started = True
                print(kwargs['reasoningText'], end="", flush=True)
            return
        
        # Detect transition from reasoning to final answer
        if "data" in kwargs and self._reasoning_started and not self._in_reasoning:
            print("\n───────────────────────────────────────")
            print("📝 Answer: ", end="", flush=True)
            self._in_reasoning = True
        
        # Reset state on new message
        if "complete" in kwargs and kwargs.get("complete"):
            self._reasoning_started = False
            self._in_reasoning = False
        
        if self.debug:
            filtered = {k: v for k, v in kwargs.items() 
                       if v is not None and v != "" and v != {} and k != "delta"}
            if filtered:
                print(f"\n[DEBUG] {filtered}")
        else:
            super().__call__(**kwargs)


# =============================================================================
# AGENT CONFIGURATIONS
# =============================================================================

def create_simple_agent(api_key: str) -> Agent:
    """Create a simple non-streaming agent."""
    model = xAIModel(
        client_args={"api_key": api_key},
        model_id="grok-4-1-fast-non-reasoning-latest",
    )
    return Agent(model=model, system_prompt="You are a helpful assistant. Be concise.", callback_handler=None)


def create_streaming_agent(api_key: str, debug: bool = False) -> Agent:
    """Create a streaming agent."""
    model = xAIModel(
        client_args={"api_key": api_key},
        model_id="grok-4-1-fast-non-reasoning-latest",
    )
    return Agent(model=model, system_prompt="You are a helpful assistant. Be concise.",
                 callback_handler=StreamingCallbackHandler(debug=debug))


def create_tools_agent(api_key: str, debug: bool = False) -> Agent:
    """Create an agent with Strands tools (client-side)."""
    model = xAIModel(
        client_args={"api_key": api_key},
        model_id="grok-4-1-fast-non-reasoning-latest",
    )
    return Agent(model=model, system_prompt="You are a helpful assistant with access to tools.",
                 tools=[calculate, get_weather], callback_handler=StreamingCallbackHandler(debug=debug))


def create_x_search_agent(api_key: str, debug: bool = False) -> Agent:
    """Create an agent with X search server-side tool."""
    from xai_sdk.tools import x_search
    
    model = xAIModel(
        client_args={"api_key": api_key},
        model_id="grok-4-1-fast-non-reasoning-latest",
        xai_tools=[x_search()],
    )
    return Agent(model=model, system_prompt="You are a helpful assistant that can search X (Twitter).",
                 callback_handler=StreamingCallbackHandler(debug=debug))


def create_hybrid_agent(api_key: str, debug: bool = False) -> Agent:
    """Create an agent with both server-side and client-side tools."""
    from xai_sdk.tools import x_search
    
    model = xAIModel(
        client_args={"api_key": api_key},
        model_id="grok-4-1-fast-non-reasoning-latest",
        xai_tools=[x_search()],
    )
    return Agent(model=model, 
                 system_prompt="You are a helpful assistant with X search, calculator, and weather tools.",
                 tools=[calculate, get_weather], callback_handler=StreamingCallbackHandler(debug=debug))


def create_reasoning_agent(api_key: str, debug: bool = False) -> Agent:
    """Create a reasoning agent (grok-3-mini)."""
    model = xAIModel(
        client_args={"api_key": api_key},
        model_id="grok-3-mini",
        reasoning_effort="high",
    )
    return Agent(model=model, system_prompt="You are a helpful assistant that thinks through problems carefully.",
                 callback_handler=StreamingCallbackHandler(debug=debug))


def create_reasoning_encrypted_agent(api_key: str, debug: bool = False) -> Agent:
    """Create a grok-4 reasoning agent with encrypted content."""
    model = xAIModel(
        client_args={"api_key": api_key},
        model_id="grok-4-fast-reasoning",
        use_encrypted_content=True,
    )
    return Agent(model=model, system_prompt="You are a helpful assistant that thinks through problems carefully.",
                 callback_handler=StreamingCallbackHandler(debug=debug))


def create_web_search_citations_agent(api_key: str, debug: bool = False) -> Agent:
    """Create an agent with web search and inline citations."""
    from xai_sdk.tools import web_search
    
    model = xAIModel(
        client_args={"api_key": api_key},
        model_id="grok-4-1-fast-non-reasoning-latest",
        xai_tools=[web_search()],
        include=["inline_citations"],
    )
    return Agent(model=model, system_prompt="You are a helpful assistant that searches the web. Always cite sources.",
                 callback_handler=StreamingCallbackHandler(debug=debug))


# =============================================================================
# INTERACTIVE CHAT
# =============================================================================

def chat_loop(agent: Agent, agent_name: str, is_streaming: bool = True) -> None:
    """Run an interactive multi-turn chat loop."""
    print(f"\n{'=' * 60}")
    print(f"CHAT: {agent_name}")
    print(f"{'=' * 60}")
    print("Commands: /quit, /clear, /help")
    print("-" * 60)
    
    while True:
        try:
            user_input = input("\nYou: ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\nGoodbye!")
            break
        
        if not user_input:
            continue
        if user_input.lower() in ["/quit", "/q"]:
            break
        elif user_input.lower() in ["/clear", "/c"]:
            agent.messages.clear()
            print("[Conversation cleared]")
            continue
        
        if is_streaming:
            print("\nAssistant: ", end="", flush=True)
        else:
            print("\nAssistant: ", end="")
        
        try:
            result = agent(user_input)
            if not is_streaming:
                print(result)
        except Exception as e:
            print(f"\n[Error]: {e}")


def main() -> None:
    """Main entry point."""
    api_key = get_api_key()
    
    agents = {
        "1": ("Simple (non-streaming)", lambda: create_simple_agent(api_key), False),
        "2": ("Streaming", lambda: create_streaming_agent(api_key), True),
        "3": ("Streaming (DEBUG)", lambda: create_streaming_agent(api_key, debug=True), True),
        "4": ("Strands tools", lambda: create_tools_agent(api_key), True),
        "5": ("X search", lambda: create_x_search_agent(api_key), True),
        "6": ("Hybrid", lambda: create_hybrid_agent(api_key), True),
        "7": ("Reasoning (grok-3-mini)", lambda: create_reasoning_agent(api_key), True),
        "8": ("Reasoning encrypted (grok-4)", lambda: create_reasoning_encrypted_agent(api_key), True),
        "9": ("Reasoning encrypted DEBUG", lambda: create_reasoning_encrypted_agent(api_key, debug=True), True),
        "10": ("Web search + citations", lambda: create_web_search_citations_agent(api_key), True),
    }
    
    while True:
        print("\n" + "=" * 60)
        print("STRANDS AGENTS - xAI INTERACTIVE CHAT")
        print("=" * 60)
        for k, (name, _, _) in agents.items():
            print(f"  {k}. {name}")
        print("  q. Quit\n")
        
        choice = input("Enter choice: ").strip().lower()
        
        if choice == "q":
            break
        elif choice in agents:
            name, create_fn, is_streaming = agents[choice]
            try:
                agent = create_fn()
                chat_loop(agent, name, is_streaming)
            except Exception as e:
                print(f"\n[Error]: {e}")


if __name__ == "__main__":
    main()

[Cerrix]

The PR is large, but 1400 (out of 2230) lines come from unit tests and integration tests. The added model provider is 765 lines.