Skip to content
177 changes: 177 additions & 0 deletions examples/ollama_chat_automatic_function_calling.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,177 @@
from ollama import ChatResponse, chat


def ollama_automatic_function_calling(
client_fn: chat,
messages: list[dict],
options: dict = None,
**kwargs,
) -> tuple[ChatResponse, list[dict]]:
"""
Automatically handles tool/function calls in a chat loop.

max_turns is reached.

Args:
client_fn: Ollama chat client
messages: List of message dicts for the conversation
max_turns: Maximum number of chat iterations (default: 20)
model: Model name to use (default: "glm-5:cloud")
tools: List of available tool definitions
tool_map: Dict mapping tool names to callable functions
chat_fn: Optional callable to use for chat (default: ollama.chat)

Returns:
(last_response, messages)
"""
model = kwargs.get("model")
if model is None:
raise ValueError("model must be specified")

max_turns = kwargs.get("max_turns", 20)
tools = kwargs.get("tools")
tool_map = kwargs.get("tool_map")
last_response = None

def _to_msg_dict(msg):
if hasattr(msg, "model_dump"):
return msg.model_dump(exclude_none=True)
if isinstance(msg, dict):
return msg
return {
"role": getattr(msg, "role", "assistant"),
"content": getattr(msg, "content", ""),
}

def _extract_tool_call_parts(tool_call):
function = getattr(tool_call, "function", None)
if function is None and isinstance(tool_call, dict):
function = tool_call.get("function", {})

if isinstance(function, dict):
name = function.get("name")
arguments = function.get("arguments") or {}
else:
name = getattr(function, "name", None)
arguments = getattr(function, "arguments", None) or {}

return name, arguments

def _infer_tool_name(t):
if t is None:
return None
if isinstance(t, dict):
fn = t.get("function") or {}
if isinstance(fn, dict):
return fn.get("name")
return None
return getattr(t, "name", None) or getattr(t, "__name__", None)

def _resolve_callable(fn):
wrapped = getattr(fn, "__wrapped__", None)
return wrapped if callable(wrapped) else fn

if tool_map is None and tools:
tool_map = {}
for t in tools:
name = _infer_tool_name(t)
if name:
tool_map[name] = t

if not callable(client_fn):
raise ValueError("client_fn must be a callable")

reserved = {"model", "tools", "messages", "options", "max_turns", "tool_map"}
chat_kwargs = {k: v for k, v in kwargs.items() if k not in reserved}

for _ in range(max_turns):
response = client_fn(
model=model,
tools=tools,
messages=messages,
options=options,
**chat_kwargs,
)
last_response = response
assistant_msg = response.message
messages.append(_to_msg_dict(assistant_msg))

tool_calls = getattr(assistant_msg, "tool_calls", None) or []
if not tool_calls:
return last_response, messages

for tool_call in tool_calls:
tool_name, tool_args = _extract_tool_call_parts(tool_call)
tool_fn = tool_map.get(tool_name) if tool_map else None

if not callable(tool_fn):
messages.append(
{
"role": "tool",
"content": f"Tool '{tool_name}' not found.",
"tool_name": tool_name or "unknown_tool",
}
)
continue

if isinstance(tool_args, str):
try:
import json

tool_args = json.loads(tool_args)
except Exception:
tool_args = {}

try:
result = _resolve_callable(tool_fn)(**(tool_args or {}))
content = str(result)
except Exception as e:
content = f"Tool '{tool_name}' execution failed: {e}"

messages.append(
{
"role": "tool",
"content": content,
"tool_name": tool_name,
}
)

# Safety exit: do not raise, just return the last response.
return last_response, messages


def add(a: float, b: float) -> float:
"""Apply addition"""
return a + b


def sub(a: float, b: float) -> float:
"""Apply subtraction"""
return a - b


def mul(a: float, b: float) -> float:
"""Apply multiplication"""
return a * b


def div(a: float, b: float) -> float:
"""Apply division"""
return a / b


if __name__ == "__main__":
from ollama import chat

response, messages = ollama_automatic_function_calling(
chat,
model="glm-5:cloud",
tools=[add, sub, mul, div],
messages=[{"role": "user", "content": "这个数学题:4+7*9/6是多少?"}],
think="low", # or true, false
options={
"temperature": 1,
"top_p": 1.0,
"top_k": 60,
},
)