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.

klea_utils.llm.INPUT_TOKEN_ESTIMATE_SAFETY_FACTOR = 1.05

Safety factor applied on top of the char-based token estimate so the reserved output window keeps headroom for tokenizer variance: the ~4 chars/token rule can undercount the real token count (e.g. 9736 chars estimate 2434 tokens while the server tokenizes 2435), which pushes input + output one token over the model’s context window and fails the request.

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, required: 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.

required marks roles that need a default model for the app to function (e.g. chat). At startup, required roles with an empty model trigger a warning (not a failure) listing the environment variables to set. Optional roles (e.g. guard) are skipped when their model is empty.

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) → str

Add the context summary to the system prompt.

Returns a text block framing the previous-context summary. Recent conversation messages are no longer flattened into this block: they are injected as real message objects by the node’s prompt assembly (see klea_utils.nodes.base).

Parameters:

context_summary – Summary of the past conversation.

Returns:

Prompt text block, or "" when there is no summary.

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

Conservative token estimate for a prompt’s character count.

Used 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. The estimate is deliberately inflated by INPUT_TOKEN_ESTIMATE_SAFETY_FACTOR (rounded up) so a tokenizer that counts more tokens than the average never pushes the request over the context window.

Parameters:

input_chars – Number of characters in the prompt.

Returns:

Conservative 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.BaseMessage]]

Get recent conversations between start and stop indices.

Returns the conversation as a single text block (used as prompt/summary input) along with the ordered BaseMessage objects, preserving the interleaved user/assistant order of the original history.

Parameters:

• all_messages – all the messages

• start – start index

• stop – stop index

Returns:

(conversation, ordered list of human/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_recent_messages(messages, max_chars: int, keep_at_least: int = 1) → list[langchain_core.messages.BaseMessage]

Return the most recent human/ai messages bounded by max_chars.

Walks backwards through messages (in their original interleaved order) accumulating pretty_repr() length until max_chars would be exceeded. The first keep_at_least messages are always included, so the latest exchange is never dropped even when it alone exceeds the budget.

Parameters:

• messages – All conversation messages.

• max_chars – Maximum total characters of the returned window.

• keep_at_least – Minimum number of messages to always include.

Returns:

Ordered list of recent human/ai messages.

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_langchain_endpoint(instance: Any, config: langgraph.types.RunnableConfig) → str | None

Resolve the base URL/endpoint of a configurable chat model.

create_configurable_model returns a generic _ConfigurableModel whose concrete provider instance is only built at invoke time (via its private _model(config) method, which is already called on every ainvoke). Native providers (mistral:, anthropic:, deepseek:, …) do not carry a base_url in the configurable dict – the provider resolves its own default endpoint internally – so to probe such an endpoint we materialise the concrete instance and read its resolved endpoint attribute.

Cheap by construction: _model(config) only parses the model string and constructs the provider object (module import is cached per provider); no network or API calls happen at construction.

Parameters:

• instance – The configurable model (_ConfigurableModel) whose concrete instance to materialise.

• config – Per-invoke RunnableConfig carrying the merged configurable dict.

Returns:

The resolved base URL, or None when the concrete model exposes no known endpoint attribute (e.g. plain openai: models, whose default URL lives inside the OpenAI SDK client, not on the model object).

klea_utils.llm.resolve_output_token_limit(overrides: dict[str, Any], provider: str, role: str | None = None, input_chars: int | None = None, *, use_endpoint: bool = False) → 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.

The context source depends on use_endpoint: for endpoints with a known base_url (custom OpenAI-compatible endpoints, or a native provider endpoint resolved via resolve_langchain_endpoint()), max_model_len from the live /models probe is authoritative (models.dev’s values are per-deployment and may under-report). The normal (non-retry) path passes use_endpoint=False so it stays on the fast, offline models.dev value – its purpose is only to produce a finite budget (mainly for the HuggingFace whole-window reservation), not an exact cap. The retry path passes use_endpoint=True to clamp against the real server context, falling back to models.dev when the endpoint probe fails. The output clamp always uses models.dev (the endpoint exposes no separate output cap).

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.

• use_endpoint – When True, prefer the live endpoint’s max_model_len as the context source, falling back to models.dev. When False (default), use models.dev only and never query the endpoint.

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