Code generation and debugging with the Grok 4 API

Table of contents

What is Grok 4?

Grok 4's hybrid neural architecture and attention mechanism

Grok 4 compared to other AI models

Performance comparison of Grok 4 with other AI models

Grok 4 pricing

Tutorial: Code generation and observability with Grok 4 and W&B Weave

1. Setting Up Your Environment: Grok 4, OpenRouter, and W&B Weave

1.1 Create OpenRouter and W\&B Accounts & API Keys:
OpenRouter: Sign up on the OpenRouter website (openrouter.ai). Once logged in, navigate to the API keys section and create a new API key. This key will authenticate your requests to the OpenRouter API. Keep it secure.
Weights & Biases: Create a free Weights & Biases account. After signing up, navigate to https://wandb.ai/authorize to find your Weights & Biases API key. This key is essential for logging data to your W&B dashboard.
1.2 Enable Access to xAI’s Grok Model on OpenRouter: On the OpenRouter platform, locate xAI’s Grok 4 in the model list (labeled as x-ai/grok-4). You might need to agree to specific terms of service or ensure you have a payment method on file, as Grok 4 is a premium model.
1.3 Install Necessary Python Libraries: Ensure you have the requests library (for OpenRouter API calls) and weave (for W&B observability) installed. Open your terminal or command prompt and run:

pip install requests weave

1.4 Configure API Access and Initialize Weave in Your Script: Now, let's set up the Python script where all our interactions will happen. We'll define two key functions wrapped with weave.op(): call_grok4 to interact with the model and evaluate_prime_function to test its output. Wrapping these with @weave.op() automatically logs their inputs, outputs, and execution details to your W&B dashboard, creating detailed 'traces'.
Replace "YOUR_OPENROUTER_API_KEY" with your actual key.

import weave
import requests
import math
import os
import re

# --- W&B Weave Initialization ---
project_name = "Grok4-CodeGen-Tutorial"
weave.init(project_name)

# --- OpenRouter API Configuration ---
# Secure way to input API key
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "YOUR_OPENROUTER_API_KEY") # Replace with your key if not using env var
OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
GROK4_MODEL = "x-ai/grok-4"

# Define a Weave operation for calling Grok 4
@weave.op()
def call_grok4(prompt: str, context_messages: list = None) -> str:
    """Sends a prompt to Grok 4 via OpenRouter and returns the response."""
    if context_messages is None:
        context_messages = []
    
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://localhost:8888",  # Required for OpenRouter
        "X-Title": "Grok4-CodeGen-Tutorial"  # Optional but recommended
    }
    
    messages = context_messages + [{"role": "user", "content": prompt}]
    data = {
        "model": GROK4_MODEL,
        "messages": messages,
        "max_tokens": 1000,
        "temperature": 0.7
    }
    
    try:
        response = requests.post(OPENROUTER_URL, json=data, headers=headers)
        response.raise_for_status()
        result = response.json()
        
        # Check if response has the expected structure
        if 'choices' in result and len(result['choices']) > 0:
            return result['choices'][0]['message']['content']
        else:
            return f"Unexpected response format: {result}"
            
    except requests.exceptions.RequestException as e:
        return f"API request failed: {str(e)}"
    except KeyError as e:
        return f"Unexpected response structure: {str(e)}"

# Define a Weave operation to evaluate the generated code
@weave.op()
def evaluate_prime_function(code: str, test_cases: list) -> dict:
    """
    Executes the provided code and tests the is_prime function against test cases.
    Returns a dictionary of results for Weave to log.
    """
    results = {}
    total_tests = len(test_cases)
    passed_tests = 0

    try:
        # Create a clean namespace for code execution
        exec_globals = {"math": math}  # Provide math module
        
        # Execute the generated code to define the is_prime function
        exec(code, exec_globals)
        is_prime_func = exec_globals.get('is_prime')

        if not is_prime_func:
            raise ValueError("Generated code does not define 'is_prime' function.")

        # Test each case
        for num, expected in test_cases:
            try:
                actual = is_prime_func(num)
                is_correct = (actual == expected)
                if is_correct:
                    passed_tests += 1
                results[f"Test_{num}"] = {
                    "input": num,
                    "expected": expected,
                    "actual": actual,
                    "correct": is_correct
                }
            except Exception as test_error:
                results[f"Test_{num}"] = {
                    "input": num,
                    "expected": expected,
                    "actual": None,
                    "correct": False,
                    "error": str(test_error)
                }
                
    except Exception as e:
        results["execution_error"] = str(e)
        passed_tests = 0

    return {
        "passed_tests": passed_tests,
        "total_tests": total_tests,
        "accuracy": passed_tests / total_tests if total_tests > 0 else 0,
        "test_details": results
    }

print("Setup complete!")

You've now configured your environment and defined two weave.op functions. Every time call_grok4 or evaluate_prime_function is called, Weave will automatically log their inputs, outputs, and execution details to your W&B dashboard, creating a 'trace' of the process.

2. Generate Initial Code with Grok 4

initial_prompt = """Write a Python function is_prime(n) that returns True if n is a prime number and False otherwise. The function should be efficient for large n and include comments explaining the logic. Only return the function, with no explanations."""

print("Sending initial prompt to Grok 4...")
# Calling our Weave-wrapped function logs this interaction
generated_code = call_grok4(initial_prompt) 
print("\n--- Grok 4 Generated Code ---")
print(generated_code)
print("-----------------------------\n")



```python
import math

def is_prime(n):
    """
    Checks if a number n is prime.
    
    A prime number is a natural number greater than 1 that has no positive divisors other than 1 and itself.
    This function uses an efficient trial division method optimized for performance, checking divisibility
    up to the square root of n, with additional optimizations to skip multiples of 2 and 3.
    
    Time complexity: O(sqrt(n)), which is efficient for n up to around 10^18 on modern hardware.
    
    :param n: An integer to check for primality.
    :return: True if n is prime, False otherwise.
    """
    # Handle small numbers: primes start from 2, so anything <= 1 is not prime.
    if n <= 1:
        return False
    # 2 and 3 are primes.
    if n <= 3:
        return True
    # Eliminate multiples of 2 and 3 early, as they can't be prime (except 2 and 3 themselves).
    if n % 2 == 0 or n % 3 == 0:
        return False
    
    # Now check for factors from 5 onwards, using the 6k ± 1 optimization.
    # All primes greater than 3 are of the form 6k ± 1.
    # We check i and i+2 (which is i+2 = (i+6)-4, but effectively covers the pattern) in steps of 6.
    # We only need to check up to sqrt(n) because if n has a factor larger than sqrt(n), 
    # it must also have a smaller one we've already checked.
    i = 5
    while i * i <= n:
        if n % i == 0 or n % (i + 2) == 0:
            return False
        i += 6
    return True
```

3. Test and evaluate the generated code with Weave

test_cases = [
    (1, False),    # 1 is not prime
    (2, True),     # 2 is prime
    (3, True),     # 3 is prime
    (4, False),    # 4 is not prime
    (15, False),   # 15 is not prime
    (17, True),    # 17 is prime
    (97, True),    # 97 is prime
    (100, False),  # 100 is not prime
    (997, True)    # 997 is prime
]

print("Evaluating the generated code...")
# Calling our Weave-wrapped evaluation function logs the results
evaluation_results = evaluate_prime_function(generated_code, test_cases)
print("\n--- Evaluation Results ---")
print(evaluation_results)
print("--------------------------\n")



4. Debug and refine with Grok 4 (Observed with Weave)

# Simulate a problematic code example for demonstration purposes
# In a real scenario, this would be the actual 'generated_code' from step 2 if it had a flaw.
problematic_code_example = """
def is_prime(n: int) -> bool:
    # Simulating a bug: This version might not handle n < 2 correctly
    if n % 2 == 0:
        return n == 2
    import math
    limit = int(math.sqrt(n)) + 1
    for divisor in range(3, limit, 2):
        if n % divisor == 0:
            return False
    return True
""" 

# Provide feedback to Grok 4 based on the evaluation findings
debugging_prompt = f"""The `is_prime` function you provided:
```python
{problematic_code_example}

Calling call_grok4 again logs this new interaction

print("Sending debugging prompt to Grok 4...")
corrected_code = call_grok4(debugging_prompt)
print("\n--- Grok 4 Corrected Code ---")
print(corrected_code)
print("------------------------------\n")



Re-evaluate the corrected code, logging the new evaluation with Weave

print("Re-evaluating the corrected code...")
corrected_evaluation_results = evaluate_prime_function(corrected_code, test_cases)
print("\n--- Corrected Evaluation Results ---")
print(corrected_evaluation_results)
print("-----------------------------------\n")

5. Leveraging W&B Weave for deeper analysis and workflow improvement

Version Control for Prompts & Models: Every weave.op run is recorded as a trace. You can easily navigate through these traces to see how different prompts, context messages, or even versions of Grok 4 (if you were to compare models) affect the generated code and its subsequent evaluation scores. This is crucial for systematic prompt engineering and understanding the impact of your inputs.
Cost and Performance Monitoring: Weave automatically captures metrics like token usage and latency for each `call_grok4` operation (when provided by the API). You can visualize these trends over time to manage costs effectively, identify any performance bottlenecks, and optimize your prompts for efficiency.
Failure Analysis and Debugging: If Grok 4 ever produces "hallucinations," logically flawed code, or simply doesn't meet expectations, you have the exact prompt, response, and evaluation results recorded. This allows you to go back, understand *why* it failed, and iterate on your prompting strategy to improve future outcomes.
Comparing Model Variants and Prompts: Weave's robust logging capabilities enable you to run experiments comparing different prompts for the same task, or even compare the performance of Grok 4 against other LLMs (e.g., `gpt-4o`, `claude-3-opus`) on specific coding benchmarks within a unified dashboard.
Automated Dashboards and Reports: You can create custom, interactive dashboards in W&B to visualize key metrics (e.g., accuracy over time, cost per successful generation, average debugging iterations). These dashboards can be shared with your team, providing transparent insights into the efficiency of your AI-powered development.

Alternative use cases for Grok 4



Medical imaging analysis: With its multimodal prowess, Grok 4 can analyze and interpret images alongside text. For example, it could examine an X-ray or MRI scan and provide diagnostic suggestions or detailed descriptions of the findings . This could assist doctors by triaging cases or offering a second opinion based on visual data. (Of course, any medical use would require rigorous validation, but the potential is there.)
Creative content generation: Grok 4 isn’t limited to structured tasks – it can also generate creative text such as stories, poetry, or artwork captions. Imagine feeding it a prompt to write a short sci-fi story or a screenplay snippet; Grok’s advanced understanding allows it to maintain coherent plots and characters. Its image understanding might even allow it to write descriptions or alt-text for images, and future updates hint at image generation abilities as well . This makes Grok a tool not just for coders, but for writers and artists looking to brainstorm or automate creative workflows.
Educational tutoring and problem solving: Thanks to its strong reasoning skills, Grok 4 can serve as an AI tutor or educational assistant. It can break down complex concepts in math, science, or programming into step-by-step explanations. For instance, a student could ask Grok to explain a calculus problem or to help debug a piece of code they wrote. Grok’s detailed reasoning (it was trained to be methodical in problem solving) is a huge asset here. It can also generate practice problems or quiz questions on the fly. Its proficiency in mathematics and logic (recall that it can even solve advanced math problems and has proved itself in competitions) makes it a powerful tool for learning environments.

Conclusion