Web Search with References

!pip install mistralai wikipedia

from mistralai.client import Mistral
from mistralai.client.models import  UserMessage, SystemMessage
import getpass
import os

if not os.environ.get("MISTRAL_API_KEY"):
    os.environ["MISTRAL_API_KEY"] = getpass.getpass("Mistral API key: ")

client = Mistral(os.environ["MISTRAL_API_KEY"])

query = "Who won the Nobel Peace Prize in 2024 ?"

#Add the user message to the chat_history
chat_history = [
    SystemMessage(content="You are a helpful assistant that can search the web for information. Use context to answer the question."),
    UserMessage(content=query),
]
print(chat_history)

web_search_tool = {
    "type": "function",
    "function": {
        "name": "web_search",
        "description": "Search the web for a query for which you do not know the answer",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Query to search the web in keyword form.",
                }
            },
            "required": ["query"],
        },
    },
}

chat_response = client.chat.complete(
    model="mistral-medium-latest",
    messages=chat_history,
    tools=[web_search_tool],
)

if hasattr(chat_response.choices[0].message, 'tool_calls'):
    tool_call = chat_response.choices[0].message.tool_calls[0]
    chat_history.append(chat_response.choices[0].message)
    print(tool_call)
else:
    print("No tool call found in the response")

import json
from mistralai.client.models import ToolMessage

query = json.loads(tool_call.function.arguments)["query"]
wb_result = get_wikipedia_search(query)

tool_call_result = ToolMessage(
    content=wb_result,
    tool_call_id=tool_call.id,
    name=tool_call.function.name,
)

# Append the tool call message to the chat_history
chat_history.append(tool_call_result)

#See chunks in the response
print(json.dumps(json.loads(wb_result), indent=2))

for msg in chat_history:
    print(msg,end='\n')

from mistralai.client.models import TextChunk, ReferenceChunk

def format_response(chat_response: list, wb_result:dict):
    print("\n🤖 Answer:\n")
    refs_used = []

    # Print the main response
    for chunk in chat_response.choices[0].message.content:
        if isinstance(chunk, TextChunk):
            print(chunk.text, end="")
        elif isinstance(chunk, ReferenceChunk):
            refs_used += chunk.reference_ids

    # Print references
    if refs_used:
        print("\n\n📚 Sources:")
        for i, ref in enumerate(set(refs_used), 1):
            reference = json.loads(wb_result)[str(ref)]
            print(f"\n{i}. {reference['title']}: {reference['url']}")

# Use the formatter
chat_response = client.chat.complete(
    model="mistral-medium-latest",
    messages=chat_history,
    tools=[web_search_tool],
)
format_response(chat_response, wb_result)

stream_response = client.chat.stream(
    model="mistral-medium-latest",
    messages=chat_history,
    tools=[web_search_tool],
)

last_reference_index = 0
if stream_response is not None:
    for event in stream_response:
        chunk = event.data.choices[0]
        if chunk.delta.content:
            if isinstance(chunk.delta.content, list):
                    # Check if TYPE of chunk is a reference
                    print(chunk.delta.content)
                    references_ids = [
                        ref_id
                        for chunk_elem in chunk.delta.content
                        if chunk_elem.type == "reference"
                        for ref_id in chunk_elem.reference_ids
                    ]
                    last_reference_index += len(references_ids)

                    # Map the references ids to the references data stored in the chat history
                    references_data = [json.loads(wb_result)[str(ref_id)] for ref_id in references_ids]
                    urls = " " + ", ".join(
                        [
                            f"[{i}]({reference['url']})"
                            for i, reference in enumerate(
                                references_data,
                                start=last_reference_index - len(references_ids) + 1,
                            )
                        ]
                    )
                    print(urls, end="")
            else:
                print(chunk.delta.content, end="")

# Start a conversation without creating an Agent
response = client.beta.conversations.start(
    model="mistral-medium-3.5",
    inputs=[{"role": "user", "content": "Who won the Nobel Peace Prize in 2024 ?"}],
    tools=[{"type": "web_search"}],
)

# Display raw outputs
for output in response.outputs:
    print(output.model_dump_json()[:256])