"""
Adaptateur smolagents reliant le moteur LLM MLX ou MTPLX a CodeAgent.
Partage les instances de modeles deja chargees, pour eviter tout rechargement
de poids ou surcout de memoire unifiee.
"""
from __future__ import annotations
import json
import time
import urllib.error
import urllib.request
from collections.abc import Callable
from typing import Any
from mlx_lm.generate import stream_generate as streamGenerateMlxLm
from mlx_vlm import stream_generate as streamGenerateMlxVlm
from smolagents.models import ChatMessage, MessageRole, Model
from smolagents.monitoring import TokenUsage
from ...core.logger import getMirzaLogger
from .reactParser import cleanActionOutput
logger = getMirzaLogger("mirza.agent.adapter")
[docs]
class MirzaModelAdapter(Model):
"""
Adaptateur smolagents reliant le moteur LLM MLX ou MTPLX a CodeAgent.
Parameters
----------
llmEngine : BaseLlm
Instance du moteur LLM (MlxQwenLlm ou MtplxLlm).
model_id : str, optional
Identifiant du modele smolagents (defaut: "qwen-local-mlx").
max_tokens : int, optional
Plafond de tokens generes par tour (defaut: 1024).
temperature : float, optional
Temperature d'echantillonnage (defaut: 0.6).
top_p : float, optional
Seuil de troncature nucleus (defaut: 0.95).
**kwargs
Parametres supplementaires transmis a Model.
"""
[docs]
def __init__(
self,
llmEngine: Any,
model_id: str = "qwen-local-mlx",
max_tokens: int = 1024,
temperature: float = 0.6,
top_p: float = 0.95,
**kwargs: Any,
) -> None:
super().__init__(
model_id=model_id,
flatten_messages_as_text=True,
**kwargs,
)
self.llmEngine = llmEngine
self.max_tokens = max_tokens
self.temperature = temperature
self.top_p = top_p
self.on_token_chunk: Callable[[str], None] | None = None
[docs]
def generate(
self,
messages: list[ChatMessage | dict],
stop_sequences: list[str] | None = None,
response_format: dict[str, str] | None = None,
tools_to_call_from: list[Any] | None = None,
**kwargs: Any,
) -> ChatMessage:
"""
Genere une reponse compatible smolagents via MTPLX ou MLX.
Parameters
----------
messages : list[ChatMessage or dict]
Historique conversationnel smolagents.
stop_sequences : list[str], optional
Chaines d'arret declenchant l'interruption de generation.
response_format : dict[str, str], optional
Format specifique demande pour la sortie.
tools_to_call_from : list[Any], optional
Outils smolagents mis a disposition du modele.
**kwargs
Parametres d'inference additionnels.
Returns
-------
ChatMessage
Message assistant contenant le texte genere et les statistiques de tokens.
"""
if getattr(self.llmEngine, "baseUrl", None):
return self._generateMtplx(
messages=messages,
stop_sequences=stop_sequences,
response_format=response_format,
tools_to_call_from=tools_to_call_from,
**kwargs,
)
return self._generateMlx(
messages=messages,
stop_sequences=stop_sequences,
response_format=response_format,
tools_to_call_from=tools_to_call_from,
**kwargs,
)
def _generateMtplx(
self,
messages: list[ChatMessage | dict],
stop_sequences: list[str] | None = None,
response_format: dict[str, str] | None = None,
tools_to_call_from: list[Any] | None = None,
**kwargs: Any,
) -> ChatMessage:
if not getattr(self.llmEngine, "isLoaded", False):
self.llmEngine.loadModel()
completion_kwargs = self._prepare_completion_kwargs(
messages=messages,
stop_sequences=stop_sequences,
tools_to_call_from=tools_to_call_from,
**kwargs,
)
prepared_messages = completion_kwargs.pop("messages", [])
stops = completion_kwargs.pop("stop", [])
completion_kwargs.pop("tools", None)
completion_kwargs.pop("tool_choice", None)
all_stops = list(stops) if stops else []
for s in ["</code>", "<end_code>", "<|im_end|>", "<|endoftext|>", "Observation:", "Calling tools:"]:
if s not in all_stops:
all_stops.append(s)
clean_messages = []
approx_prompt_chars = 0
for msg in prepared_messages:
if isinstance(msg, dict):
role = msg.get("role")
content = msg.get("content", "")
else:
role = getattr(msg, "role", "user")
content = getattr(msg, "content", "")
role = getattr(role, "value", role)
if not isinstance(role, str):
role = str(role)
if isinstance(content, list):
text_parts = []
for part in content:
if isinstance(part, dict) and "text" in part:
text_parts.append(part["text"])
elif isinstance(part, str):
text_parts.append(part)
content = "\n".join(text_parts)
content_str = str(content)
approx_prompt_chars += len(content_str)
clean_messages.append({"role": role, "content": content_str})
max_toks = kwargs.get("max_tokens", self.max_tokens)
temp = kwargs.get("temperature", self.temperature)
top_p = kwargs.get("top_p", self.top_p)
payload = {
"model": getattr(self.llmEngine, "modelRepo", "pooyakn73/Qwen3.8-9B-Distill-MTPLX-Q6G64"),
"messages": clean_messages,
"max_tokens": max_toks,
"temperature": temp,
"top_p": top_p,
"stream": True,
"stream_options": {"include_usage": True},
}
# Injection des en-tetes specifiques au fournisseur (ex: OpenRouter Auth)
headers = {"Content-Type": "application/json", "User-Agent": "Mirza-AgentAdapter/1.0"}
fournisseur = getattr(self.llmEngine, "_headers", None)
if callable(fournisseur):
headers = {**fournisseur(), **headers}
# Desactivation du raisonnement pour l'agent (MTPLX / OptiQ / OpenRouter)
payload["enable_thinking"] = False
payload["chat_template_kwargs"] = {"enable_thinking": False}
if "Authorization" in headers:
payload["reasoning"] = {"enabled": False}
if all_stops:
payload["stop"] = all_stops
req = urllib.request.Request(
f"{self.llmEngine.baseUrl}/chat/completions",
data=json.dumps(payload).encode("utf-8"),
headers=headers,
)
startTime = time.time()
firstTokenTime = None
output_tokens = 0
text = ""
nativeSpeedTokPerSec = 0.0
nativeTtftMs = 0.0
prompt_tokens_count = max(1, approx_prompt_chars // 4)
try:
with urllib.request.urlopen(req, timeout=90) as resp:
for rawLine in resp:
line = rawLine.decode("utf-8").strip()
if not line or not line.startswith("data: "):
continue
if line == "data: [DONE]":
break
dataStr = line[6:].strip()
try:
chunkJson = json.loads(dataStr)
if chunkJson.get("usage"):
u = chunkJson["usage"]
if "prompt_tokens" in u and u["prompt_tokens"] is not None:
prompt_tokens_count = int(u["prompt_tokens"])
if "completion_tokens" in u and u["completion_tokens"] is not None:
output_tokens = max(output_tokens, int(u["completion_tokens"]))
if chunkJson.get("timings"):
t = chunkJson["timings"]
if "predicted_per_second" in t and t["predicted_per_second"] is not None:
nativeSpeedTokPerSec = float(t["predicted_per_second"])
if "prompt_ms" in t and t["prompt_ms"] is not None:
nativeTtftMs = float(t["prompt_ms"])
if chunkJson.get("mtplx_stats"):
s = chunkJson["mtplx_stats"]
if "decode_tok_s" in s and s["decode_tok_s"] is not None:
nativeSpeedTokPerSec = float(s["decode_tok_s"])
elif "tok_s" in s and s["tok_s"] is not None and nativeSpeedTokPerSec <= 0.0:
nativeSpeedTokPerSec = float(s["tok_s"])
if "ttft_s" in s and s["ttft_s"] is not None:
nativeTtftMs = float(s["ttft_s"]) * 1000.0
choices = chunkJson.get("choices", [])
if not choices:
continue
delta = choices[0].get("delta", {})
contentDelta = delta.get("content", "")
if not contentDelta:
contentDelta = delta.get("reasoning_content", "")
if contentDelta:
if firstTokenTime is None:
firstTokenTime = time.time()
text += contentDelta
output_tokens += 1
if self.on_token_chunk:
try:
self.on_token_chunk(contentDelta)
except Exception as cb_err:
logger.debug("Callback token chunk exception: %s", cb_err)
if all_stops:
for stop in all_stops:
stop_idx = text.rfind(stop)
if stop_idx != -1:
text = text[:stop_idx]
break
except json.JSONDecodeError:
continue
except Exception as e:
logger.error("Erreur HTTP streaming vers le serveur MTPLX : %s", e)
return ChatMessage(
role=MessageRole.ASSISTANT,
content=f"Erreur de communication avec le serveur d'inférence MTPLX : {e}",
)
totalDuration = time.time() - startTime
if nativeTtftMs > 0.0:
self.llmEngine.lastTtftMs = round(nativeTtftMs, 1)
elif firstTokenTime is not None:
self.llmEngine.lastTtftMs = round((firstTokenTime - startTime) * 1000.0, 1)
else:
self.llmEngine.lastTtftMs = round(totalDuration * 1000.0, 1)
if nativeSpeedTokPerSec > 0.0:
currentSpeed = nativeSpeedTokPerSec
else:
decodeDuration = (time.time() - firstTokenTime) if firstTokenTime else totalDuration
currentSpeed = (output_tokens / decodeDuration) if decodeDuration > 0 else 0.0
self.llmEngine.currentSpeedTokPerSec = round(currentSpeed, 1)
self.llmEngine.lastTokensPerSec = self.llmEngine.currentSpeedTokPerSec
self.llmEngine.lastTokenCount = output_tokens
self.llmEngine.estimatedKvTokens = prompt_tokens_count + output_tokens
is_step_stream_action = bool(stop_sequences and any("code" in str(s) for s in stop_sequences))
text = cleanActionOutput(text, is_step_stream_action)
return ChatMessage(
role=MessageRole.ASSISTANT,
content=text,
raw={"out": text},
token_usage=TokenUsage(
input_tokens=prompt_tokens_count,
output_tokens=output_tokens,
),
)
def _generationEntryPoint(self):
"""
La fonction de generation in-process, et son second argument.
Deux runtimes cohabitent : `mlx-vlm` pour les conversions vision-langage,
`mlx-lm` pour le texte seul. Leur point d'entree n'est pas le meme : celui
de `mlx-vlm` lit `model.config.model_type`, que le chargeur `mlx-lm`
attache jamais.
Returns
-------
tuple
La fonction, et l'objet a lui passer en second parametre (processor
pour mlx-vlm, tokenizer pour mlx-lm).
"""
processor = getattr(self.llmEngine, "processor", None)
if processor is not None:
return streamGenerateMlxVlm, processor
return streamGenerateMlxLm, self._inProcessTokenizer()
def _inProcessTokenizer(self):
"""
Le tokenizer du moteur in-process, quel que soit son type.
Ce chemin a ete ecrit pour un moteur vision-langage, qui expose un
`processor` englobant son `tokenizer`. `MlxLmLlm` est texte seul : il n'a
qu'un `tokenizer`.
Returns
-------
object | None
Le tokenizer, ou None si le moteur n'en expose aucun.
"""
processor = getattr(self.llmEngine, "processor", None)
nested = getattr(processor, "tokenizer", None) if processor is not None else None
return nested or getattr(self.llmEngine, "tokenizer", None)
def _generateMlx(
self,
messages: list[ChatMessage | dict],
stop_sequences: list[str] | None = None,
response_format: dict[str, str] | None = None,
tools_to_call_from: list[Any] | None = None,
**kwargs: Any,
) -> ChatMessage:
if getattr(self.llmEngine, "model", None) is None or self._inProcessTokenizer() is None:
self.llmEngine.loadModel()
completion_kwargs = self._prepare_completion_kwargs(
messages=messages,
stop_sequences=stop_sequences,
tools_to_call_from=tools_to_call_from,
**kwargs,
)
prepared_messages = completion_kwargs.pop("messages", [])
stops = completion_kwargs.pop("stop", [])
tools = completion_kwargs.pop("tools", None)
completion_kwargs.pop("tool_choice", None)
all_stops = list(stops) if stops else []
for s in ["</code>", "<end_code>", "<|im_end|>", "<|endoftext|>", "Observation:", "Calling tools:"]:
if s not in all_stops:
all_stops.append(s)
tokenizer = self._inProcessTokenizer()
processor = getattr(self.llmEngine, "processor", None)
chat_template_kwargs = {
"tools": tools,
"add_generation_prompt": True,
"enable_thinking": False,
}
try:
if processor is not None and hasattr(processor, "apply_chat_template"):
prompt_ids = processor.apply_chat_template(
prepared_messages,
**chat_template_kwargs,
)
else:
prompt_ids = tokenizer.apply_chat_template(
prepared_messages,
**chat_template_kwargs,
)
except Exception:
# Les replis retirent les outils puis l'invite de generation, jamais
# la reflexion : un gabarit qui echouait la reactivait silencieusement,
# et l'agent passait de 5,3 s a plus de 27 s par passe.
for dropped in ("tools", "add_generation_prompt"):
try:
chat_template_kwargs.pop(dropped, None)
prompt_ids = tokenizer.apply_chat_template(
prepared_messages,
**chat_template_kwargs,
)
break
except Exception:
continue
else:
raise RuntimeError(
"Le gabarit de chat refuse la requete de l'agent, meme sans "
"outils ni invite de generation. `enable_thinking` n'a pas ete "
"abandonne : la reflexion reste coupee par construction."
)
temperature = kwargs.get("temperature", self.temperature)
topP = kwargs.get("top_p", self.top_p)
maxTokens = kwargs.get("max_tokens", self.max_tokens)
isVlm = getattr(self.llmEngine, "processor", None) is not None
gen_params: dict[str, Any] = {
"max_tokens": maxTokens,
"prefill_step_size": self.llmEngine.config.prefillStepSize,
}
if isVlm:
gen_params["temperature"] = temperature
gen_params["top_p"] = topP
else:
# `mlx-lm` attend un echantillonneur construit, pas des valeurs
# separees : `generate_step` ne connait ni temperature ni top_p.
from mlx_lm.sample_utils import make_sampler
gen_params["sampler"] = make_sampler(temperature, top_p=topP or 0.0)
# Seuls les moteurs vision-langage ont un cache de prefixe automatique
apcManager = getattr(self.llmEngine, "apcManager", None)
if apcManager is not None:
gen_params["apc_manager"] = apcManager
if getattr(self.llmEngine.config, "kvBits", None) is not None:
gen_params["kv_bits"] = self.llmEngine.config.kvBits
draftModel = getattr(self.llmEngine, "draftModel", None)
if draftModel is not None:
gen_params["draft_model"] = draftModel
if isVlm:
gen_params["draft_kind"] = self.llmEngine.draftKind
gen_params["draft_block_size"] = getattr(
self.llmEngine.config, "draftBlockSize", 2
)
output_tokens = 0
text = ""
lastResult = None
startTime = time.time()
firstTokenTime = None
generationTic = None
generate, handler = self._generationEntryPoint()
for responseChunk in generate(
self.llmEngine.model,
handler,
prompt_ids,
**gen_params,
):
output_tokens += 1
lastResult = responseChunk
if firstTokenTime is None:
firstTokenTime = time.time()
generationTic = time.perf_counter()
ttftMs = (firstTokenTime - startTime) * 1000.0
self.llmEngine.lastTtftMs = round(ttftMs, 1)
text += responseChunk.text
if self.on_token_chunk:
try:
self.on_token_chunk(responseChunk.text)
except Exception as exc:
logger.debug("Callback token chunk exception : %s", exc)
if all_stops and any((stop_index := text.rfind(stop)) != -1 for stop in all_stops):
text = text[:stop_index]
break
generationDuration = (time.perf_counter() - generationTic) if generationTic else (time.time() - (firstTokenTime or startTime))
mlxTps = getattr(lastResult, "generation_tps", None)
if mlxTps is not None and mlxTps > 0:
currentSpeed = float(mlxTps)
else:
currentSpeed = (output_tokens / generationDuration) if generationDuration > 0 else 0.0
self.llmEngine.currentSpeedTokPerSec = round(currentSpeed, 1)
self.llmEngine.lastTokensPerSec = self.llmEngine.currentSpeedTokPerSec
self.llmEngine.lastTokenCount = output_tokens
input_token_count = len(prompt_ids) if isinstance(prompt_ids, list) else len(tokenizer.encode(str(prompt_ids)))
self.llmEngine.estimatedKvTokens = input_token_count + output_tokens
is_step_stream_action = bool(stop_sequences and any("code" in str(s) for s in stop_sequences))
text = cleanActionOutput(text, is_step_stream_action)
return ChatMessage(
role=MessageRole.ASSISTANT,
content=text,
raw={"out": text},
token_usage=TokenUsage(
input_tokens=input_token_count,
output_tokens=output_tokens,
),
)