Skip to main content
EN · FR Start for free
Menu
Lesson 9 of 16 · Production: Batch, Optimization, and Observability

Implement Citations in Python

Lesson 9 of 16 Updated November 3, 2026 Advanced
On this page
  1. The Complete Citation Pipeline
  2. Define the Model Contract in the System Message
  3. Declare the Search Tool
  4. Return Results in Reference Format
  5. Separate Text from References
  6. Prepare Frontend Display
  7. Key Takeaways

Updated November 3, 2026 based on the Mistral Citations & References guide: the tool result follows the documented format—an object indexed by source number, without type or reference_id fields—and the reference_ids in the response resolve via str(ref_id) within that object. Source: Citations & References.

The Complete Citation Pipeline

You will now assemble a RAG pipeline that cites its sources, from the first API call to the data structure your interface consumes. Keep this principle in mind: the model never fetches documents itself. It signals that it needs a search, you execute that search in your own database, and then you return the excerpts along with their identifiers. This back-and-forth produces a response where every claim remains tied to a real document.

Specifically, a full cycle requires two calls to chat.complete. The first triggers the tool call; the second generates the cited response once the sources are injected into the history. We will walk through this cycle using a simple case: an internal documentation assistant asked about best practices for deploying an LLM.

Define the Model Contract in the System Message

The system message sets the contract: answer based on the provided sources, cite systematically, and above all, admit when information is missing rather than filling the gap. This last instruction makes the difference in production, because responses tend to drift on questions poorly covered by the documentation base.

from mistralai.client import Mistral
import json

client = Mistral(api_key="your-api-key")

# System message encouraging citations
system_message = {
    "role": "system",
    "content": (
        "You are a documentation assistant. "
        "Answer questions using the provided sources. "
        "Cite your sources systematically. "
        "If the information is not in any source, state this explicitly."
    )
}

user_message = {
    "role": "user",
    "content": "What are the best practices for deploying an LLM in production?"
}

messages = [system_message, user_message]

Declare the Search Tool

The tool describes what a document search looks like in your system: here, a simple search_documentation function taking a text query. The first call does not expect a written answer; it expects the model to formulate the query it deems relevant, saving you from writing a reformulation layer.

reference_tool = {
    "type": "function",
    "function": {
        "name": "search_documentation",
        "description": "Searches the documentation database",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query"
                }
            },
            "required": ["query"]
        }
    }
}

# First call: the model requests to use the tool
response = client.chat.complete(
    model="mistral-large-latest",
    messages=messages,
    tools=[reference_tool]
)

Return Results in Reference Format

This is the structural step. Your search engine results come back in a ToolMessage whose content is the reference object from the previous lesson: each source has a number ("0", "1"…), a URL, a title, and the selected excerpts. These are the numbers the model reuses to mark its citations, and which your frontend will later resolve into links.

from mistralai.client.models import ToolMessage

# Documentation search results, in Mistral reference format
references = {
    "0": {
        "url": "https://docs.example.com/deployment-guide",
        "title": "LLM Deployment Guide",
        "snippets": [[
            "Always configure rate limits to protect the service.",
            "Monitor P99 latency and error rates continuously."
        ]],
        "description": None,
        "date": "2026-09-01",
        "source": "docs"
    },
    "1": {
        "url": "https://docs.example.com/security",
        "title": "AI Application Security",
        "snippets": [[
            "Filter user inputs to prevent injections.",
            "Never expose API keys on the client side."
        ]],
        "description": None,
        "date": "2026-08-15",
        "source": "docs"
    },
    "2": {
        "url": "https://docs.example.com/scaling",
        "title": "Scalability and Performance",
        "snippets": [[
            "Use semantic caching to reduce API calls.",
            "Size instances based on peak traffic."
        ]],
        "description": None,
        "date": "2026-07-20",
        "source": "docs"
    }
}

# Add the tool call and response to the history
tool_call = response.choices[0].message.tool_calls[0]
messages.append(response.choices[0].message)

messages.append(
    ToolMessage(
        tool_call_id=tool_call.id,
        name="search_documentation",
        content=json.dumps(references, ensure_ascii=False)
    )
)

Note that the model’s original message is reinjected into the history before the ToolMessage: without this tool call / tool result pair, the API refuses the continuation. The second call no longer needs the tool definitions, since the search has already taken place.

# Second call: the model generates a response with citations
final_response = client.chat.complete(
    model="mistral-large-latest",
    messages=messages
)

Separate Text from References

The cited response is no longer a simple string but a sequence of chunks alternating prose and citation marks. Parsing involves traversing this sequence and sorting: TextChunk objects reconstruct the readable text, while ReferenceChunk objects deliver the reference_ids actually used. This distinction gives you valuable information: exactly which sources the model used among those you provided.

from mistralai.client.models import TextChunk, ReferenceChunk

def parse_cited_response(response):
    """Extracts text and references from a cited response."""
    content_parts = response.choices[0].message.content
    text_parts = []
    cited_refs = set()

    for chunk in content_parts:
        if isinstance(chunk, TextChunk):
            text_parts.append(chunk.text)
        elif isinstance(chunk, ReferenceChunk):
            for ref_id in chunk.reference_ids:
                cited_refs.add(ref_id)

    return {
        "text": "".join(text_parts),
        "references_used": list(cited_refs)
    }

result = parse_cited_response(final_response)
print("=== Response ===")
print(result["text"])
print("\n=== Cited Sources ===")
for ref_id in result["references_used"]:
    source = references[str(ref_id)]
    print(f"  - {source['title']} : {source['url']}")

Prepare Frontend Display

Flattening the response into plain text would lose the position of citations, and thus their value. It is better to keep the sequence and enrich it: resolve each ReferenceChunk against a reference table to carry the title and URL of each source. Your interface thus receives an ordered list of blocks it can render as clickable links, tooltips, or interactive citation numbers, without ever having to reinterpret the text.

def format_for_frontend(response, reference_map):
    """Formats the response with citations for the frontend."""
    content_parts = response.choices[0].message.content
    formatted = []

    for chunk in content_parts:
        if isinstance(chunk, TextChunk):
            formatted.append({
                "type": "text",
                "content": chunk.text
            })
        elif isinstance(chunk, ReferenceChunk):
            refs = []
            for ref_id in chunk.reference_ids:
                source = reference_map.get(str(ref_id))
                if source:
                    refs.append({
                        "id": ref_id,
                        "title": source["title"],
                        "url": source["url"]
                    })
            formatted.append({
                "type": "citation",
                "references": refs
            })

    return formatted

Before opening this pipeline to your full documentation base, run it with just three to five sources, as in the example above. You will immediately see if the model cites at the right place, if it leverages multiple documents, or if it systematically clings to the first one, allowing you to correct the system message before the behavior gets lost in the volume.

Key Takeaways

  • The pipeline follows a two-call flow: tool call then cited response
  • References are passed via ToolMessage, in an object indexed by source number ("0", "1"…)
  • The response contains a mix of TextChunk and ReferenceChunk to parse
  • Structure the output so your frontend can display interactive citations
  • Test with 3–5 sources to validate behavior before moving to production
← Previous lessonCitations with Sources: Reference Object and RAGNext lesson →Ground Responses in Sources and Reduce Hallucinations