LLM utilities

LLM related utils

File: klea_rag/llm.py

Copyright 2026 Ankur Sinha Author: Ankur Sinha <sanjay DOT ankur AT gmail DOT com>

klea_utils.llm.DEFAULT_MAX_OUTPUT_TOKENS = 4096

Fallback max output tokens used when no node/role default provides a value.

class klea_utils.llm.LLMModel(*, model_name: str = '', instance: Any, role_defaults: dict[str, Any] = {}, provider_defaults: dict[str, dict[str, Any]] = {}, modifiable: bool = True)

Bases: BaseModel

Container for a single LLM model instance and its runtime configuration.

instance holds the model object (typically a _ConfigurableModel returned by init_chat_model). role_defaults stores role-wide default parameters (e.g. max_tokens, temperature) that apply to every node sharing this role, unless overridden by node or user config.

build_config() performs a five-layer merge:

Layer 0 — role_defaults: role-wide parameters (e.g. {"max_tokens": 4096}).

Layer 1 — model_name: the default model identifier from the graph config.

Layer 2 — context_overrides: per-request fields from the API (model, api_key, etc.). Only applied when modifiable=True, and skipping any keys frozen by node defaults.

Layer 3 — node_defaults: frozen per-node defaults (always win).

Layer 4 — provider_defaults: per-provider defaults from the graph config (e.g. HuggingFace role budgets), applied after the model string is parsed so the resolved provider is known. Applied with setdefault so explicit role/context/node values always win.

modifiable controls whether the model can be changed at runtime (both the API and web UI reject modifications to locked roles). Set to False to lock a role (e.g. guard) against user overrides in managed deployments.

build_config(context_overrides: dict[str, Any] | None = None, node_defaults: dict[str, Any] | None = None) → langgraph.types.RunnableConfig

Merge up to five layers of model configuration into a RunnableConfig.

Layer order (lowest -> highest priority):

0. self.role_defaults — role-wide parameters

1. self.model_name — role model from graph config

2. context_overrides — per-request user overrides

3. node_defaults — frozen per-node defaults

4. self.provider_defaults — per-provider defaults (setdefault)

Parameters:

• context_overrides – Per-request fields from the API (e.g. model, api_key). Only applied when self.modifiable is True, and skipping any keys present in node_defaults.

• node_defaults – Frozen per-node defaults (e.g. {"temperature": 0.3}). Always win.

Returns:

A RunnableConfig with the configurable key populated.

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class klea_utils.llm.ParsedModelName(provider: str | None, model_name: str, suffix: str | None)

Bases: NamedTuple

Parsed components of a model name string.

model_name: str

Alias for field number 1

provider: str | None

Alias for field number 0

suffix: str | None

Alias for field number 2

klea_utils.llm.add_memory_to_prompt(context_summary: str, messages, num_history_messages) → str

Add memory to system prompt.

Adds the context summary and recent conversation

Parameters:

state – agent state

Returns:

“memory” string to add to the system prompt

klea_utils.llm.check_model_works(model, timeout=30, retries=5)

Check if a model works since it is not tested when loaded

klea_utils.llm.check_ollama_model(logger, model, exit=False)

Check if ollama model is available

Parameters:

• logger (logging) – logger instance

• model (str) – ollama model name

• exit (bool) – if we should call sys.exit if check fails

Returns:

None

Throws ollama.ResponseError:

if model is not available

Throws ConnectionError:

if cannot connect to an Ollama server

klea_utils.llm.classify_llm_invocation_error(exc: BaseException) → LLMInvocationErrorCategory

Classify an LLM invocation exception into a category.

Providers report failures inconsistently, so this uses tolerant regex matching over the exception message and its cause chain. Categories are checked in order of specificity (context overflow first), so a message matching several heuristics lands in the most actionable bucket.

Parameters:

exc – The exception raised by an LLM invocation.

Returns:

The best-effort klea_utils.errors.LLMInvocationErrorCategory.

klea_utils.llm.content_to_str(content: str | list[dict | str] | None) → str

Normalise an AIMessage.content value to a plain string.

AIMessage.content can be a plain string, a list of content blocks (when the LLM returns tool calls or structured output), or None. This helper always returns a string suitable for downstream text processing (regex, in checks, prompt interpolation, etc.).

Parameters:

content – The raw .content value from an AIMessage.

Returns:

A plain string.

klea_utils.llm.create_configurable_model(logger: Logger)

Set up a configurable chat model.

Creates a _ConfigurableModel with no default model. Model, provider, and all other parameters (base_url, api_key, temperature, etc.) are specified per-invoke via the config["configurable"] dict passed to ainvoke().

This enables runtime model switching — each ainvoke() call creates a fresh underlying model instance for the given provider, so there is no stale configuration leakage between calls.

The lookup function check_model_works is deliberately not called here — we prefer a “leap before you look” approach so that startup is fast and model availability is checked only at query time.

klea_utils.llm.estimate_input_tokens(input_chars: int) → int

Rough token estimate for a prompt’s character count.

Used only to keep the reserved output window within a model’s total budget (input + output <= context). ~4 characters per token is a reasonable average for mixed English/code text; an exact count would require provider-specific tokenizers.

Parameters:

input_chars – Number of characters in the prompt.

Returns:

Estimated token count.

klea_utils.llm.extract_llm_output_content(output: AIMessage | dict) → str

Extract plain-text content from an LLM output.

Handles both AIMessage (non-structured output) and dict (structured output with raw / parsed keys):

• AIMessage – returns content_to_str(message.content).

• dict – extracts the raw AIMessage from a structured output response and returns its content; falls back to output["parsed"] and finally str(output).

Parameters:

output – The raw output from llm.invoke().

Returns:

A plain-text string.

klea_utils.llm.format_alert(text: str, level: str = 'warning') → str

Wrap text as a GitHub-style markdown alert (e.g. > [!WARNING]).

Multi-line text is prefixed per line so the whole thing stays inside the blockquote. Renderers with the markdown2 alerts extra (the NiceGUI speech bubbles) show it as a styled callout; others fall back to a plain blockquote.

Parameters:

• text – Alert body text

• level – Alert level (note, tip, important, warning, caution)

Returns:

Markdown alert blockquote

klea_utils.llm.get_last_n_conversations(all_messages, start: int = 0, stop: int | None = None) → tuple[str, list[langchain_core.messages.HumanMessage], list[langchain_core.messages.AIMessage]]

Get recent converstations between start and stop indices

Parameters:

• all_messages – all the messages

• start – start index

• stop – stop index

Returns:

(conversation, list of human messages, list of ai messages)

klea_utils.llm.get_provider_allowed_fields(provider: str) → set[str]

Return the set of init-param names accepted by a given provider’s model class.

Uses LangChain’s internal provider registry to look up the Pydantic model class and introspect its fields (including aliases so that both api_key and openai_api_key pass through).

Falls back to an empty set if the provider is not registered in LangChain’s built-in providers. Raises ImportError if the provider’s integration package is not installed — callers should handle this at configuration time, not silently fall through.

The caller should always include {"model", "model_provider"} on top of the returned set since those are consumed by _ConfigurableModel before reaching the model constructor.

klea_utils.llm.get_token_limit_param(provider: str) → str

Return the max-output token parameter name for a provider.

Providers disagree on the parameter name for the maximum number of output tokens: Ollama uses num_predict, while HuggingFace’s ChatHuggingFace (which internally maps it to max_new_tokens) and other OpenAI-compatible providers all use max_tokens.

Note

Known benign warning

When Klea resolves max_tokens for HuggingFace, the inner HuggingFaceEndpoint constructed by ChatHuggingFace.from_model_id (which declares max_new_tokens, not max_tokens) logs WARNING! max_tokens is not default parameter and shuffles it into model_kwargs. This is a false positive: the limit is still delivered correctly as max_tokens to InferenceClient.chat_completion via the outer ChatHuggingFace, which is the parameter the HuggingFace Inference API actually accepts. Do not “fix” it by switching to max_new_tokens here.

Parameters:

provider – Klea provider id (huggingface, ollama, …)

Returns:

The token parameter name to send in the invoke config.

klea_utils.llm.is_output_truncated(output: AIMessage | dict[str, Any]) → bool

Return True if an LLM output was truncated by the max-token limit.

Providers signal truncation via finish_reason == "length" on the message metadata. Handles both plain AIMessage outputs and structured-output dicts ({"raw": AIMessage, ...}), and the list-form finish_reason some providers return.

Parameters:

output – The raw output from llm.invoke().

Returns:

True when the model stopped because it hit the output cap.

klea_utils.llm.load_prompt(prompt_name: str, prompt_registry_location: str)

Load a prompt from file called prompt_name.md

Parameters:

• str – prompt file name

• prompt_registry_location – location of prompts folder/registry

Returns:

loaded prompt text

klea_utils.llm.parse_model_name(raw: str) → ParsedModelName

Split a model name into provider, model identifier, and suffix.

Follows the provider:model_id convention. The provider is expected to be explicitly included; no provider inference is done.

With three segments the third is treated as a suffix (provider hint, model tag, base URL, etc.) unless the provider is ollama, for which the second and third segments form the model name (model_name:tag).

Examples:

• ollama:bge-m3:latest -> provider=ollama, model=bge-m3:latest, suffix=None

• huggingface:org/model:auto -> provider=huggingface, model=org/model, suffix=auto

• custom:model:https://example.com/v1 -> provider=custom, model=model, suffix=https://example.com/v1

• openai:gpt-4o -> provider=openai, model=gpt-4o, suffix=None

• bge-m3 -> provider=None, model=bge-m3, suffix=None

Parameters:

raw – Model name with optional provider prefix

Returns:

Parsed model name components

klea_utils.llm.parse_output_with_thought(message: langchain_core.messages.AIMessage, schema: type[TSchema]) → tuple[TSchema, str]

Parse AI message with thought to a dict based on given schema

klea_utils.llm.prompt_value_to_messages(prompt: langchain_core.prompt_values.PromptValue) → list[dict]

Convert a PromptValue to a clean list of message dicts.

Each dict has role and content keys, suitable for JSON serialisation in the inspector debug panel.

Parameters:

prompt – The LangChain PromptValue (filled, variables already substituted).

Returns:

A list of {"role": "...", "content": "..."} dicts.

klea_utils.llm.resolve_output_token_limit(overrides: dict[str, Any], provider: str, role: str | None = None, input_chars: int | None = None) → None

Ensure a bounded max-output token param is set in overrides.

HuggingFace-style providers apply a total budget: the reserved output window (max_new_tokens) is accounted against the model’s context window alongside the input, and an unset value makes them reserve the entire remaining window (causing spurious usage limits and rate limiting). This helper guarantees a finite, clamped value.

Resolution precedence:

1. An explicit provider token param (max_tokens / max_new_tokens / num_predict) already present in overrides (user/node/role value).

2. The generic max_output_tokens key (provider-agnostic count).

3. The built-in per-role fallback for role.

The resolved value is clamped to min(value, catalog limit.output) and, when the catalog exposes a context window and input_chars is given, to the remaining budget (context - estimated input tokens) so HuggingFace’s total-budget check is never exceeded.

Parameters:

• overrides – The merged configurable dict to update in place.

• provider – Klea provider id (huggingface, ollama, …).

• role – Model role (e.g. "chat"), used for the built-in per-role fallback.

• input_chars – Character count of the prompt, to bound the output within the total budget.

klea_utils.llm.split_output_by_section(text: str, section_start_marker: str, section_end_marker: str | None = None)

Split out thoughts and actual responses from AI responses