123 lines
3.3 KiB
Python
123 lines
3.3 KiB
Python
"""Mistral-specific utilities.
|
|
|
|
This module contains utilities specific to the Mistral provider,
|
|
including reask functions, response handlers, and message formatting.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from ...mode import Mode
|
|
from ...processing.schema import generate_openai_schema
|
|
from ...utils.core import dump_message
|
|
|
|
|
|
def reask_mistral_structured_outputs(
|
|
kwargs: dict[str, Any],
|
|
response: Any,
|
|
exception: Exception,
|
|
):
|
|
"""
|
|
Handle reask for Mistral structured outputs mode when validation fails.
|
|
|
|
Kwargs modifications:
|
|
- Adds: "messages" (assistant content and user correction request)
|
|
"""
|
|
kwargs = kwargs.copy()
|
|
reask_msgs = [
|
|
{
|
|
"role": "assistant",
|
|
"content": response.choices[0].message.content,
|
|
}
|
|
]
|
|
reask_msgs.append(
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
f"Validation Error found:\n{exception}\nRecall the function correctly, fix the errors"
|
|
),
|
|
}
|
|
)
|
|
kwargs["messages"].extend(reask_msgs)
|
|
return kwargs
|
|
|
|
|
|
def reask_mistral_tools(
|
|
kwargs: dict[str, Any],
|
|
response: Any,
|
|
exception: Exception,
|
|
):
|
|
"""
|
|
Handle reask for Mistral tools mode when validation fails.
|
|
|
|
Kwargs modifications:
|
|
- Adds: "messages" (tool response messages indicating validation errors)
|
|
"""
|
|
kwargs = kwargs.copy()
|
|
reask_msgs = [dump_message(response.choices[0].message)]
|
|
for tool_call in response.choices[0].message.tool_calls:
|
|
reask_msgs.append(
|
|
{
|
|
"role": "tool", # type: ignore
|
|
"tool_call_id": tool_call.id,
|
|
"name": tool_call.function.name,
|
|
"content": (
|
|
f"Validation Error found:\n{exception}\nRecall the function correctly, fix the errors"
|
|
),
|
|
}
|
|
)
|
|
kwargs["messages"].extend(reask_msgs)
|
|
return kwargs
|
|
|
|
|
|
def handle_mistral_tools(
|
|
response_model: type[Any], new_kwargs: dict[str, Any]
|
|
) -> tuple[type[Any], dict[str, Any]]:
|
|
"""
|
|
Handle Mistral tools mode.
|
|
|
|
Kwargs modifications:
|
|
- Adds: "tools" (list with function schema)
|
|
- Adds: "tool_choice" set to "any"
|
|
"""
|
|
new_kwargs["tools"] = [
|
|
{
|
|
"type": "function",
|
|
"function": generate_openai_schema(response_model),
|
|
}
|
|
]
|
|
new_kwargs["tool_choice"] = "any"
|
|
return response_model, new_kwargs
|
|
|
|
|
|
def handle_mistral_structured_outputs(
|
|
response_model: type[Any], new_kwargs: dict[str, Any]
|
|
) -> tuple[type[Any], dict[str, Any]]:
|
|
"""
|
|
Handle Mistral structured outputs mode.
|
|
|
|
Kwargs modifications:
|
|
- Adds: "response_format" derived from the response model
|
|
- Removes: "tools" and "response_model" from kwargs
|
|
"""
|
|
from mistralai.extra import response_format_from_pydantic_model
|
|
|
|
new_kwargs["response_format"] = response_format_from_pydantic_model(response_model)
|
|
new_kwargs.pop("tools", None)
|
|
new_kwargs.pop("response_model", None)
|
|
return response_model, new_kwargs
|
|
|
|
|
|
# Handler registry for Mistral
|
|
MISTRAL_HANDLERS = {
|
|
Mode.MISTRAL_TOOLS: {
|
|
"reask": reask_mistral_tools,
|
|
"response": handle_mistral_tools,
|
|
},
|
|
Mode.MISTRAL_STRUCTURED_OUTPUTS: {
|
|
"reask": reask_mistral_structured_outputs,
|
|
"response": handle_mistral_structured_outputs,
|
|
},
|
|
}
|