"""Ollama client wrapper with tool-calling chat loop and streaming.""" from __future__ import annotations import json from collections.abc import AsyncIterator, Callable, Iterator from typing import Any, cast import ollama from ollama import Message from ollama._types import Image as OllamaImage from .config import settings from .i18n import t from .mcp import redact_secrets from .prompts import get_prompt from .tools.helper.image import resize_image from .tools.helper.pagination import parse_paged_result, parse_token from .tools.registry import ToolExecutionResult, ToolRegistry, default_registry # Maximum successive tool-call rounds before we force a final answer. # Some models request many tools per round; we keep this bounded but large # enough to let iterative research complete (e.g. search, then fetch). MAX_TOOL_ROUNDS = 5 # Upper bound on the cumulative size of tool results sent back to ollama. # Larger values can cause ollama to spend a very long time processing the # context between rounds. MAX_TOOL_CONTEXT_CHARS = 24000 # Type alias matching ollama's expected tools signature. ToolCallable = Callable[..., Any] # Marker that separates a paged tool result's body from its page-info footer. # Must match ``ollama_web.tools.helper.pagination.PAGE_INFO_MARKER``. _PAGE_INFO_MARKER = "\n---page-info---" def _safe_tool_arguments(arguments: Any) -> Any: return redact_secrets(arguments) def _safe_tool_result(result: str) -> str: redacted = redact_secrets(result) return str(redacted) def get_client(host: str | None = None) -> ollama.Client: """Return an ``ollama.Client`` configured for the configured host.""" import httpx timeout = httpx.Timeout( settings.ollama_timeout, connect=10.0, read=settings.ollama_timeout, write=10.0, pool=10.0, ) return ollama.Client(host=host or settings.ollama_host, timeout=timeout) def list_models(host: str | None = None) -> list[str]: """Return the list of locally available model names.""" client = get_client(host) try: resp = client.list() except Exception: # noqa: BLE001 return [] models: list[str] = [] for item in resp.get("models", []) if isinstance(resp, dict) else getattr(resp, "models", []): name = item.get("model") if isinstance(item, dict) else getattr(item, "model", None) if name: models.append(str(name)) return models def get_model_capabilities(model: str, host: str | None = None) -> set[str]: """Return the capability set for a single model via ``ollama.show()``. Unknown or unreachable models are treated as having no special capabilities, which keeps the UI conservative and avoids sending unsupported parameters. """ try: client = get_client(host) show_response = client.show(model) except Exception: # noqa: BLE001 return set() caps = getattr(show_response, "capabilities", None) if caps is None and isinstance(show_response, dict): caps = show_response.get("capabilities") if not caps: return set() return {str(c).lower() for c in caps} def _is_image_message(message: Message) -> bool: """Check whether a message has image attachments.""" images = getattr(message, "images", None) return bool(images) def _remove_images_with_notice(message: Message, language: str | None = None) -> Message: """Remove unsupported image bytes and add a localized fixed notice.""" content = str(getattr(message, "content", "") or "") notice = get_prompt("vision_unsupported_attachment", lang=language) content = f"{content}\n\n{notice}".strip() kwargs: dict[str, Any] = {"role": getattr(message, "role", "user"), "content": content} for attr in ("name", "tool_calls", "tool_name"): value = getattr(message, attr, None) if value is not None: kwargs[attr] = value if getattr(message, "thinking", None): kwargs["thinking"] = message.thinking return Message(**kwargs) def _tool_calls_from_message(msg: Any) -> list[dict[str, Any]]: """Extract tool calls from a message object/dict.""" if msg is None: return [] calls = getattr(msg, "tool_calls", None) if not isinstance(msg, dict) else msg.get("tool_calls") if not calls: return [] out: list[dict[str, Any]] = [] for c in calls: if isinstance(c, dict): fn = c.get("function", {}) name = fn.get("name") or c.get("name", "") args = fn.get("arguments") or c.get("arguments", {}) out.append({"name": name, "arguments": args}) else: fn = getattr(c, "function", None) if fn is not None: name = getattr(fn, "name", "") args = getattr(fn, "arguments", {}) else: name = getattr(c, "name", "") args = getattr(c, "arguments", {}) out.append({"name": name, "arguments": args}) return out def _token_counts(response: Any) -> dict[str, int] | None: """Extract prompt_eval_count and eval_count from an ollama response.""" if response is None: return None if isinstance(response, dict): prompt_eval = response.get("prompt_eval_count") eval_count = response.get("eval_count") else: prompt_eval = getattr(response, "prompt_eval_count", None) eval_count = getattr(response, "eval_count", None) if prompt_eval is None or eval_count is None: return None return {"prompt_eval_count": int(prompt_eval), "eval_count": int(eval_count)} def _trim_tool_result(result: str, limit: int) -> str: """Truncate ``result`` to ``limit`` chars while keeping it readable. Paged tools must fit their own complete envelope before advancing a cursor. An oversized paged result is therefore rejected atomically instead of being truncated into an unrecoverable content gap. """ if len(result) <= limit: return result # A paged tool must size its own body before advancing its cursor. Trimming # it here would create an unrecoverable gap between visible text and the # next token, so fail the call atomically and let the model repeat it. if _PAGE_INFO_MARKER in result: return ( "Paged tool result exceeded the context-safe result limit. " "Repeat the same tool call with a smaller max_chars value; no part " "of this page was consumed." ) # No footer: legacy truncation behaviour. cutoff = result[:limit].rsplit("\n\n", 1)[0] if len(cutoff) < limit * 0.5: # No good paragraph boundary; cut at a space instead. cutoff = result[:limit].rsplit(" ", 1)[0] return cutoff + "\n\n…[tool result truncated]" def _paged_source_key(name: str, arguments: Any) -> tuple[str, str]: """Identify one paged source without exposing its content.""" if not isinstance(arguments, dict): return name, "unknown" if name == "search_and_fetch": return name, f"{arguments.get('query', '')}|{arguments.get('file_type', '')}" token = arguments.get("continue_from") if isinstance(token, str) and token: try: payload = parse_token(token, name) return name, str(payload.get("url", "unknown")) except ValueError: pass source = arguments.get("url_or_path", arguments.get("url", "unknown")) return name, str(source) def stream_chat_with_tools( messages: list[Message], model: str, *, think: bool | str | None = None, registry: ToolRegistry | None = None, host: str | None = None, capabilities: set[str] | None = None, options: dict[str, Any] | None = None, language: str | None = None, ) -> Iterator[dict[str, Any]]: """Run a tool-calling chat loop and yield real streaming events from ollama. The final assistant answer is streamed token-by-token. Tool and thinking events are emitted inline so the UI can render them inside the assistant message. """ import logging import time logger = logging.getLogger("ollama_web.llm") caps = capabilities or set() raw_client = get_client(host) reg = registry or default_registry() tools: list[ToolCallable] = [cast(Any, t) for t in reg.ollama_tools] effective_tools = tools if "tools" in caps else [] effective_think = think if "thinking" in caps else None # Non-vision models receive a localized notice, never base64 image text. if "vision" not in caps: messages = [ _remove_images_with_notice(m, language) if _is_image_message(m) else m for m in messages ] cumulative_tool_chars = 0 outstanding_paged_sources: set[tuple[str, str]] = set() MAX_EMPTY_RETRIES = 2 empty_retries = 0 rounds = 0 while rounds < MAX_TOOL_ROUNDS: rounds += 1 full_message: dict[str, Any] = {"role": "assistant", "content": "", "thinking": ""} tool_calls: list[dict[str, Any]] = [] logger.info( "tool round %d/%d start messages=%d model=%s cumulative_tool_chars=%d", rounds, MAX_TOOL_ROUNDS, len(messages), model, cumulative_tool_chars, ) t0 = time.time() try: stream = cast(Any, raw_client).chat( model=model, messages=messages, tools=effective_tools, think=effective_think, options=options or None, stream=True, ) except Exception as exc: # noqa: BLE001 logger.warning("ollama chat call failed in %.2fs: %s", time.time() - t0, exc) yield {"type": "error", "message": str(exc)} return chunk_count = 0 last_chunk = None for chunk in stream: chunk_count += 1 last_chunk = chunk thinking = _assistant_thinking(chunk) if thinking: full_message.setdefault("thinking", "") full_message["thinking"] += thinking yield {"type": "thinking", "content": thinking} content = _assistant_content(chunk) if content: full_message["content"] += content yield {"type": "delta", "content": content} calls = _tool_calls_from_message(getattr(chunk, "message", None)) for c in calls: # Avoid duplicates from repeated chunks containing the same call. if c not in tool_calls: tool_calls.append(c) logger.info( "ollama stream finished round=%d chunks=%d elapsed=%.2fs", rounds, chunk_count, time.time() - t0, ) # Detect empty responses. Some models / proxies occasionally return a # stream with almost no meaningful content. Log it and retry a few # times before giving up. is_empty = ( not full_message.get("content") and not full_message.get("thinking") and not tool_calls ) if is_empty and empty_retries < MAX_EMPTY_RETRIES: empty_retries += 1 logger.warning( "ollama returned an empty stream (chunks=%d); retry %d/%d", chunk_count, empty_retries, MAX_EMPTY_RETRIES, ) yield { "type": "status", "message": t("status.empty_response_retry", lang=language), } # Retry the same round without consuming the tool round budget. rounds -= 1 # Small pause to avoid hammering the upstream. time.sleep(0.5) continue # Reset retry budget on any non-empty round so isolated hiccups are fine. if not is_empty: empty_retries = 0 # Reconstruct the assistant message including any tool_calls so the # next round can see them. if tool_calls: full_message["tool_calls"] = tool_calls msg = _assistant_message_chunk(full_message) if msg is not None: messages.append(msg) if not tool_calls: if outstanding_paged_sources: yield { "type": "warning", "message": t("status.pagination_incomplete", lang=language), } if is_empty: # Ran out of tool rounds while still empty; surface a warning. yield { "type": "warning", "message": t("status.empty_response", lang=language), } done_event: dict[str, Any] = {"type": "done"} tokens = _token_counts(last_chunk) if tokens: done_event.update(tokens) yield done_event return logger.info("executing %d tool call(s) round=%d", len(tool_calls), rounds) tool_results: list[tuple[str, ToolExecutionResult]] = [] total_result_len = 0 for call in tool_calls: name = call.get("name", "") args = call.get("arguments", {}) paged_source = _paged_source_key(name, args) yield {"type": "tool_start", "name": name, "arguments": _safe_tool_arguments(args)} t1 = time.time() remaining_context = MAX_TOOL_CONTEXT_CHARS - (cumulative_tool_chars + total_result_len) is_builtin_paged = name in {"scrape_url", "pdf_to_text", "search_and_fetch"} if is_builtin_paged and remaining_context < settings.max_tool_result_chars: detailed_result = ToolExecutionResult( "Tool context budget exhausted before this paged call was executed. " "Unread content remains; do not claim that the source was read in full." ) outstanding_paged_sources.add(paged_source) logger.info("paged tool skipped name=%s reason=context_budget", name) else: try: detailed_result = reg.execute_detailed( name, args if isinstance(args, dict) else json.dumps(args) ) except Exception as exc: # noqa: BLE001 logger.warning( "tool %s execution failed in %.2fs: %s", name, time.time() - t1, exc ) detailed_result = ToolExecutionResult(f"Tool execution failed: {exc}") result = detailed_result.text raw_is_paged = _PAGE_INFO_MARKER in result if raw_is_paged: _, _, next_token = parse_paged_result(result) if next_token and isinstance(args, dict): paged_source = _paged_source_key( name, {**args, "continue_from": next_token} ) outstanding_paged_sources.add(paged_source) elif is_builtin_paged: outstanding_paged_sources.discard(paged_source) # Always truncate per-tool result immediately so it does not bloat # the context even before we reach the cumulative limit. result = _safe_tool_result( _trim_tool_result(detailed_result.text, settings.max_tool_result_chars) ) detailed_result.text = result logger.info( "tool %s executed len=%d elapsed=%.2fs", name, len(result), time.time() - t1, ) media = [ artifact.public_dict() for artifact in detailed_result.artifacts if artifact.for_user and artifact.file_id is not None ] event: dict[str, Any] = {"type": "tool_end", "name": name, "result": result} if media: event["media"] = media yield event tool_results.append((name, detailed_result)) total_result_len += len(result) cumulative_tool_chars += sum(len(r.text) for _, r in tool_results) for name, detailed_result in tool_results: content = detailed_result.text image_values: list[OllamaImage] = [] has_hidden_image = False has_audio = False for artifact in detailed_result.artifacts: if not artifact.for_assistant: continue if artifact.kind == "image": if "vision" in caps: processed = resize_image(artifact.data, name=artifact.name) image_values.append(OllamaImage(value=processed)) else: has_hidden_image = True elif artifact.kind == "audio": has_audio = True notices: list[str] = [] if has_hidden_image: notices.append(get_prompt("mcp_image_artifact_notice", lang=language)) if has_audio: notices.append(get_prompt("mcp_audio_artifact_notice", lang=language)) if notices: content = f"{content}\n\n" + "\n\n".join(notices) kwargs: dict[str, Any] = {"role": "tool", "tool_name": name, "content": content} if image_values: kwargs["images"] = image_values messages.append(Message(**kwargs)) # Exceeded tool rounds: ask for a final answer without tools. logger.info( ( "max tool rounds exceeded, requesting final answer without tools " "(cumulative_tool_chars=%d)" ), cumulative_tool_chars, ) yield { "type": "status", "message": t("status.tool_limit_reached", lang=language), } limit_instruction = get_prompt("tool_limit_instruction", lang=language) if outstanding_paged_sources: limit_instruction += "\n\n" + get_prompt("pagination_incomplete_instruction", lang=language) yield { "type": "warning", "message": t("status.pagination_incomplete", lang=language), } messages.append( Message( role="user", content=limit_instruction, ) ) t0 = time.time() try: stream = cast(Any, raw_client).chat( model=model, messages=messages, tools=[], think=effective_think, options=options or None, stream=True, ) except Exception as exc: # noqa: BLE001 logger.warning("final ollama chat call failed in %.2fs: %s", time.time() - t0, exc) yield {"type": "error", "message": str(exc)} return logger.info("streaming final answer") last_chunk = None for chunk in stream: last_chunk = chunk thinking = _assistant_thinking(chunk) if thinking: yield {"type": "thinking", "content": thinking} content = _assistant_content(chunk) if content: yield {"type": "delta", "content": content} logger.info("final answer stream finished elapsed=%.2fs", time.time() - t0) done_event_final: dict[str, Any] = {"type": "done"} tokens = _token_counts(last_chunk) if tokens: done_event_final.update(tokens) yield done_event_final async def astream_chat_with_tools( messages: list[Message], model: str, *, think: bool | str | None = None, registry: ToolRegistry | None = None, host: str | None = None, capabilities: set[str] | None = None, options: dict[str, Any] | None = None, language: str | None = None, ) -> AsyncIterator[dict[str, Any]]: """Async wrapper that yields events from ``stream_chat_with_tools``.""" import anyio gen = stream_chat_with_tools( messages, model, think=think, registry=registry, host=host, capabilities=capabilities, options=options, language=language, ) def _next() -> dict[str, Any] | None: try: return next(gen) except StopIteration: return None while True: event = await anyio.to_thread.run_sync(_next, abandon_on_cancel=True) if event is None: break yield event def _assistant_content(response: Any) -> str: msg = getattr(response, "message", None) or ( response.get("message") if isinstance(response, dict) else None ) if msg is None: return "" if isinstance(msg, dict): return str(msg.get("content", "") or "") return str(getattr(msg, "content", "") or "") def _assistant_thinking(response: Any) -> str: msg = getattr(response, "message", None) or ( response.get("message") if isinstance(response, dict) else None ) if msg is None: return "" if isinstance(msg, dict): return str(msg.get("thinking", "") or "") return str(getattr(msg, "thinking", "") or "") def _assistant_message_chunk(full_message: dict[str, Any]) -> Message | None: """Build an assistant ``Message`` from the accumulated streamed message. ``tool_calls`` are expected to be already accumulated into ``full_message``. If the message is empty, ``None`` is returned. """ content = full_message.get("content", "") thinking = full_message.get("thinking", "") calls = full_message.get("tool_calls", []) if not content and not thinking and not calls: return None kwargs: dict[str, Any] = {"role": "assistant", "content": content} if thinking: kwargs["thinking"] = thinking if calls: # Convert plain dicts into ollama ToolCall objects. tc_list: list[Message.ToolCall] = [] for c in calls: name = c.get("name", "") args = c.get("arguments", {}) tc_list.append( Message.ToolCall(function=Message.ToolCall.Function(name=name, arguments=args)) ) kwargs["tool_calls"] = tc_list return Message(**kwargs)