What's new in agent-framework-mlx: Tool Calling, Agent Framework 1.12.1 and more

Β· 1135 words Β· 6 minutes to read

Back in December I introduced the MLX Integration Library for Agent Framework, a small library that lets you plug local MLX models into Agent Framework applications as a regular chat client, sitting right next to your cloud-backed agents. Agent Framework has been moving quickly since then, and rather than letting the library drift behind, I spent some time catching up and, in the process, added one of the features people kept asking me about: tool calling.

Updated to Agent Framework 1.12.1 πŸ”—

The library now targets Agent Framework 1.12.1, the latest stable release (as of today). Agent Framework is still evolving quickly under the hood, so rather than a loose version bound, agent-framework-mlx now pins its dependencies exactly:

dependencies = [
    "agent-framework-core==1.12.1",
    "mlx-lm==0.31.3",
    "mlx==0.32.0",
    "pydantic>=2.0.0",
]

Bumping to 1.12.1 was not just a version number change. A few core primitives were renamed along the way (ChatMessage became Message, TextContent became Content), and the recommended way to build a chat client also changed. Instead of implementing everything by hand on top of BaseChatClient, the client now composes a few mixins that Agent Framework provides:

class MLXChatClient(
    ChatMiddlewareLayer[MLXChatOptions],
    FunctionInvocationLayer[MLXChatOptions],
    ChatTelemetryLayer[MLXChatOptions],
    BaseChatClient[MLXChatOptions],
):
    ...

FunctionInvocationLayer is what actually gives us the tool-calling loop for free (more on that below), and ChatTelemetryLayer wires up OpenTelemetry tracing, metrics and token usage tracking automatically, without any extra setup on your side.

Latest MLX dependencies πŸ”—

On the MLX side, the library now runs on mlx-lm 0.31.3 and mlx 0.32.0, the latest versions available at the time of writing. As with Agent Framework, these are also pinned exactly rather than left open-ended, since mlx-lm’s generation and sampling APIs have a habit of shifting between releases too.

Tool calling πŸ”—

This is the headline feature. The Phi family of models (which is what most of my examples use) supports function calling through a specific convention: available tools are embedded as JSON in the system message, and the model responds with a <|tool_call|>[...]<|/tool_call|> block instead of prose when it wants to invoke one.

agent-framework-mlx now understands this convention end to end. When you attach tools to a request, the library converts each Agent Framework tool into the flat schema Phi expects, injects it into the prompt, and then parses the completion looking for a tool call, whether it is wrapped in the <|tool_call|> tags or just a bare JSON array, tolerating some of the stray text models occasionally produce around it. If a tool call is found, it comes back as a proper function-call content item, which FunctionInvocationLayer then uses to actually invoke your Python function and feed the result back to the model.

Streaming works too - since a tool call cannot be meaningfully rendered token by token, the library buffers the streamed chunks internally and only decides, once generation has finished, whether what came out was prose (which gets replayed to you as normal streaming updates) or a tool call (which gets surfaced as a function-call update). Either way, you never see raw tool-call syntax leak into your output.

Here is the local_agent.py sample from the repository, showing a small BMI calculator tool wired up to a local agent:

from agent_framework_mlx import MLXChatClient, MLXGenerationConfig

def calculate_bmi(weight_kg: float, height_m: float) -> str:
    """Calculates the Body Mass Index (BMI) given weight in kg and height in meters."""
    bmi = weight_kg / (height_m ** 2)
    return f"{bmi:.2f}"

config = MLXGenerationConfig(
    temp=0.0,  # lower temperature works best for reliable tool use
    max_tokens=1000
)

client = MLXChatClient(
    model_path="mlx-community/Phi-4-mini-instruct-8bit",
    generation_config=config,
)

# as_agent() is a convenience method that ships with Agent Framework 1.12.1
agent = client.as_agent(
    name="HealthAssistant",
    instructions="You are a helpful health assistant. You can calculate BMI.",
    tools=[calculate_bmi]
)

response = await agent.run("My weight is 70kg and my height is 1.75m. What is my BMI?")
print(response)

And the streaming follow-up:

async for update in agent.run("Now calculate it for someone who is 90kg and 1.80m.", stream=True):
    if update.text:
        print(update.text, end="", flush=True)

Running the full sample end to end gives us something like this:

--- πŸš€ Loading MLX Model: mlx-community/Phi-4-mini-instruct-8bit ---

πŸ“ User: My weight is 70kg and my height is 1.75m. What is my BMI?

--- ⚑️ Running Agent (Automatic Tool Use) ---
πŸ€– Assistant: Your Body Mass Index (BMI) is 22.86. 
This is considered to be within the healthy weight range. 
However, keep in mind that BMI is a general guide and 
individual health can vary. It's always a good idea to 
consult with a healthcare professional for personalized advice.

--- 🌊 Running Streaming Agent ---

πŸ“ User: Now calculate it for someone who is 90kg and 1.80m.

The Body Mass Index (BMI) for a person who weighs 90kg 
and is 1.80m tall is approximately 27.78. This value is 
considered overweight according to the World Health 
Organization's BMI categories.

Notice that the agent picks up the tool on its own, calls it with the right arguments extracted from natural language, and folds the result back into a natural language answer, both for the buffered and the streaming run.

Observability out of the box πŸ”—

Thanks to ChatTelemetryLayer, every request made through MLXChatClient is now automatically traced, with token usage and other metrics recorded through OpenTelemetry, the same as you would get from any other Agent Framework provider. There is nothing extra to configure, it comes bundled with the base client.

A few smaller improvements πŸ”—

A handful of other things landed in this round of updates:

  • Environment variable configuration. You can now set MLX_MODEL_PATH (and MLX_ADAPTER_PATH) as environment variables, or load them from a .env file, and skip passing them as constructor arguments entirely:

    export MLX_MODEL_PATH="mlx-community/Phi-4-mini-instruct-4bit"
    
    # no arguments needed if the env vars are set
    client = MLXChatClient()
    
  • A typed MLXChatOptions. MLX-specific sampling knobs, min_p, top_k, xtc_probability, xtc_threshold, repetition_penalty and repetition_context_size, are now first-class fields on MLXChatOptions instead of being smuggled in through additional_properties.

  • A dedicated worker thread for the model. MLX ties model weights and generation state to the thread they were created on, and running generation from a different thread than the one used to load the model used to produce a rather cryptic “There is no Stream(…) in current thread” error. The client now runs model loading and generation on a single dedicated worker thread, which makes this class of error go away.

Getting the update πŸ”—

The new release is on PyPI, so upgrading is just:

pip install --upgrade agent-framework-mlx

Final thoughts πŸ”—

Tool calling was the biggest gap between agent-framework-mlx and the cloud-backed chat clients in Agent Framework, and I am glad to have closed it. Combined with the move to Agent Framework 1.12.1, local MLX agents can now hold their own in more realistic agentic scenarios, not just answer questions.

As before, the code is on GitHub under the MIT license. If you run into models or quantizations where tool calling does not behave well, or you have other feature requests, issues and pull requests are always welcome.

About


Hi! I'm Filip W., a software architect from ZΓΌrich πŸ‡¨πŸ‡­. I like Toronto Maple Leafs πŸ‡¨πŸ‡¦, Rancid and quantum computing. Oh, and I love the Lowlands 🏴󠁧󠁒󠁳󠁣󠁴󠁿.

You can find me on Github, on Mastodon and on Bluesky.

My Introduction to Quantum Computing with Q# and QDK book
Microsoft MVP