I built a small mathematics reasoning agent last month to check derivations for a simulation project. Instead of trusting a single pass, I wanted explicit steps and a self-audit. In this tutorial, we will rebuild that agent. It takes a plain-text math problem, reasons through it with chain-of-thought logic, and returns a structured derivation you can inspect.
What you'll need
You will need an Oxlo.ai API key from https://portal.oxlo.ai, Python 3.10 or newer, and the OpenAI SDK.
pip install openai- Set your API key in your environment or replace
YOUR_OXLO_API_KEYin the code below.
I recommend Oxlo.ai for this because flat per-request pricing keeps long reasoning traces predictable. Token-based providers make extended chain-of-thought expensive fast. See https://oxlo.ai/pricing for details.
Step 1: Initialize the Oxlo.ai client
First, point the OpenAI SDK at Oxlo.ai. I am using kimi-k2.6 because it handles advanced reasoning and agentic coding tasks well. If you want to experiment, Oxlo.ai also offers deepseek-r1-671b for deep reasoning and qwen-3-32b for agent workflows.
from openai import OpenAI
import json
client = OpenAI(
base_url="https://api.oxlo.ai/v1",
api_key="YOUR_OXLO_API_KEY"
)
MODEL = "kimi-k2.6"Step 2: Lock down the system prompt
Most math errors come from skipped steps or hidden assumptions. We force the model to emit a structured derivation and a separate verification block before it is allowed to return a final answer.
SYSTEM_PROMPT = """You are a mathematics reasoning agent.
Your task is to solve the given problem rigorously.
Follow these rules:
1. Break the solution into clear logical steps.
2. Show all substitutions, rearrangements, and simplifications.
3. After deriving the answer, verify it by substituting back into the original problem or by checking dimensional consistency.
4. Return your entire response as a single JSON object with exactly these keys:
- "steps": a list of strings, each one logical step.
- "verification": a string describing how you checked the result.
- "confidence": an integer from 1 to 10.
- "final_answer": a string with the concise answer.
Do not include markdown formatting inside the JSON."""Step 3: Request structured output with JSON mode
We ask the model to reply in JSON so we can inspect the reasoning programmatically. Oxlo.ai supports JSON mode, so we set response_format to enforce valid output. I keep temperature at 0.2 because precision matters more than creativity in algebra.
def solve_problem(problem_text):
response = client.chat.completions.create(
model=MODEL,
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": problem_text},
],
response_format={"type": "json_object"},
temperature=0.2,
)
raw = response.choices[0].message.content
return json.loads(raw)
result = solve_problem("Solve for x: 2^(x+1) + 2^(x-1) = 40")
print(json.dumps(result, indent=2))Step 4: Add a self-correction loop
In practice, about one in twenty solutions had a minor sign error. Rather than bail out, I check the confidence score and verification string. If either looks weak, I feed the draft back into the model with a critique prompt. This second call costs one more request on Oxlo.ai, but because pricing is per-request, the cost is predictable.
def solve_with_audit(problem_text):
draft = solve_problem(problem_text)
if draft.get("confidence", 0) < 8 or "uncertain" in draft.get("verification", "").lower():
critique = (
f"Previous solution: {json.dumps(draft)}\n"
f"Critique this solution, fix any errors, and return a revised JSON object "
f"with the same schema. Problem: {problem_text}"
)
response = client.chat.completions.create(
model=MODEL,
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": problem_text},
{"role": "assistant", "content": json.dumps(draft)},
{"role": "user", "content": critique},
],
response_format={"type": "json_object"},
temperature=0.2,
)
draft = json.loads(response.choices[0].message.content)
return draftStep 5: Wrap it in a CLI
Finally, we unpack the JSON and print each step with its index. This makes it easy to scan where a mistake happened if the final answer still looks wrong.
def print_solution(problem_text):
data = solve_with_audit(problem_text)
print(f"Problem: {problem_text}\n")
print("Derivation:")
for i, step in enumerate(data.get("steps", []), 1):
print(f" {i}. {step}")
print(f"\nVerification: {data.get('verification', 'N/A')}")
print(f"Confidence: {data.get('confidence', 'N/A')}/10")
print(f"\nFinal Answer: {data.get('final_answer', 'N/A')}")
if __name__ == "__main__":
import sys
query = sys.argv[1] if len(sys.argv) > 1 else "Solve for x: 2^(x+1) + 2^(x-1) = 40"
print_solution(query)Run it
Save the full script as math_agent.py and pass a problem as an argument.
python math_agent.py "Solve for x: 2^(x+1) + 2^(x-1) = 40"Typical output looks like this. The exact wording varies, but the structure remains consistent because JSON mode enforces the schema.
Problem: Solve for x: 2^(x+1) + 2^(x-1) = 40
Derivation:
1. Factor out 2^(x-1) from both terms: 2^(x-1) * (2^2 + 1) = 40.
2. Simplify the constant factor: 2^(x-1) * 5 = 40.
3. Divide both sides by 5: 2^(x-1) = 8.
4. Express 8 as a power of 2: 2^(x-1) = 2^3.
5. Equate exponents: x - 1 = 3, therefore x = 4.
6. Substitute x = 4 back into the original equation to confirm: 2^5 + 2^3 = 32 + 8 = 40.
Verification: Substitution confirms the left-hand side equals the right-hand side.
Confidence: 10/10
Final Answer: x = 4Wrap-up
Two directions I am exploring next. First, piping the steps through SymPy so the verification block executes real symbolic logic instead of natural language. Second, batch-processing an entire problem set overnight. Because Oxlo.ai does not charge by the token, a long reasoning trace on kimi-k2.6 costs the same as a one-line reply. That predictability makes batch experiments easy to budget.
Top comments (0)