Implement Predicted Outputs in Python
On this page
From Concept to Production Code
You understand the principle behind Predicted Outputs and their use cases. This lesson provides complete Python code ready for integration into your applications, including error handling and performance measurement.
Basic Implementation
The minimal pattern requires just one parameter: prediction, where you pass the expected content. Everything else in the call remains unchanged—the same model, same messages, same response handling. This makes the feature easy to adopt and just as easy to remove if metrics don’t confirm the gain:
from mistralai.client import Mistral
client = Mistral(api_key="your-api-key")
source_code = """
import pandas as pd
def load_data(filepath):
df = pd.read_csv(filepath)
df['date'] = pd.to_datetime(df['date'])
df = df.dropna(subset=['value'])
return df
def compute_stats(df):
return {
'mean': df['value'].mean(),
'median': df['value'].median(),
'std': df['value'].std()
}
"""
prompt = (
"Add a 'separator' parameter to the load_data function "
"(default ',') and use it in read_csv. "
"Respond only with the complete modified code."
)
response = client.chat.complete(
model="codestral-latest",
messages=[
{"role": "user", "content": prompt},
{"role": "user", "content": source_code}
],
prediction={"type": "content", "content": source_code}
)
modified_code = response.choices[0].message.content
print(modified_code)Utility Class for Production
Encapsulation here offers benefits beyond clean code: it centralizes the decision to enable or disable prediction. This is exactly the kind of setting you want to toggle instantly. If a deployment reveals consistently incorrect predictions, you can disable them at a single point instead of hunting down the parameter throughout your codebase:
from mistralai.client import Mistral
import time
from dataclasses import dataclass
@dataclass
class PredictionResult:
content: str
latency_ms: float
tokens_input: int
tokens_output: int
model: str
class PredictedEditor:
"""Content editor using Predicted Outputs."""
def __init__(self, api_key: str, model: str = "mistral-large-latest"):
self.client = Mistral(api_key=api_key)
self.model = model
def edit(
self,
original: str,
instruction: str,
system_prompt: str | None = None
) -> PredictionResult:
"""Edits content with output prediction."""
messages = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.extend([
{"role": "user", "content": instruction},
{"role": "user", "content": original}
])
start = time.perf_counter()
response = self.client.chat.complete(
model=self.model,
messages=messages,
prediction={"type": "content", "content": original}
)
latency = (time.perf_counter() - start) * 1000
choice = response.choices[0]
return PredictionResult(
content=choice.message.content,
latency_ms=round(latency, 1),
tokens_input=response.usage.prompt_tokens,
tokens_output=response.usage.completion_tokens,
model=self.model
)
def edit_code(self, code: str, instruction: str) -> PredictionResult:
"""Shortcut for code editing."""
return self.edit(
original=code,
instruction=instruction,
system_prompt=(
"You are a programming assistant. "
"Respond only with the modified code, "
"without explanation or markdown."
)
)
def edit_text(self, text: str, instruction: str) -> PredictionResult:
"""Shortcut for text editing."""
return self.edit(
original=text,
instruction=instruction,
system_prompt=(
"You are a professional text editor. "
"Apply the requested modification while preserving "
"the style and structure of the original text. "
"Respond only with the modified text."
)
)Using the Class
editor = PredictedEditor(api_key="your-api-key", model="codestral-latest")
# Edit code
result = editor.edit_code(
code=my_python_file,
instruction="Replace all print() calls with logging.info() calls"
)
print(f"Latency: {result.latency_ms} ms")
print(f"Tokens: {result.tokens_input} in / {result.tokens_output} out")
print(result.content)Benchmark: With vs. Without Prediction
Do not skip this step. The gain from Predicted Outputs depends entirely on the accuracy of your predictions, which depends on your data. No figure from documentation replaces measurement on your own examples. This script compares both modes on the same inputs, giving you the answer in seconds:
import statistics
def benchmark_prediction(
client, model, messages, prediction_content, runs=5
):
"""Compares latency with and without Predicted Output."""
times_normal = []
times_predicted = []
for _ in range(runs):
# Without prediction
start = time.perf_counter()
client.chat.complete(model=model, messages=messages)
times_normal.append((time.perf_counter() - start) * 1000)
# With prediction
start = time.perf_counter()
client.chat.complete(
model=model,
messages=messages,
prediction={"type": "content", "content": prediction_content}
)
times_predicted.append((time.perf_counter() - start) * 1000)
avg_normal = statistics.mean(times_normal)
avg_predicted = statistics.mean(times_predicted)
gain_pct = (1 - avg_predicted / avg_normal) * 100
return {
"without_prediction_ms": round(avg_normal, 1),
"with_prediction_ms": round(avg_predicted, 1),
"gain_percentage": round(gain_pct, 1),
"runs": runs
}Error Handling
Error handling is straightforward here for a structural reason: an incorrect prediction is not an error. The model ignores it and generates normally—you lose the latency gain, never the quality. The real edge cases lie elsewhere:
def safe_edit(editor, original, instruction):
"""Editing with robust error handling."""
try:
result = editor.edit(original, instruction)
# Check that the response is not empty
if not result.content or len(result.content.strip()) < 10:
raise ValueError("Response too short or empty")
return result
except Exception as e:
# Fallback: retry without prediction
print(f"Predicted Output failed ({e}), falling back to standard...")
response = editor.client.chat.complete(
model=editor.model,
messages=[
{"role": "user", "content": instruction},
{"role": "user", "content": original}
]
)
return PredictionResult(
content=response.choices[0].message.content,
latency_ms=0,
tokens_input=response.usage.prompt_tokens,
tokens_output=response.usage.completion_tokens,
model=editor.model
)Constraints to Know
Before deploying, integrate three points into your code review. First, the n parameter (multiple generations) is incompatible with prediction. If your pipeline requests multiple completions, you must choose. Second, the match between prediction and output is handled by the model regardless of the content’s position in the response—you do not need to prepare any alignment. Finally, and this is what makes the feature safe to deploy: a bad prediction never degrades response quality; it only costs you the latency gain. The worst-case scenario is simply a return to normal speed on the two compatible models (mistral-large-latest and codestral-latest).
Key Takeaways
- The
predictionparameter simply adds to thechat.complete()call - Encapsulate the logic in a reusable class for production
- Systematically measure the gain with a benchmark before deploying
- Plan a fallback without prediction in case of errors
- A bad prediction does not degrade quality, only latency