nat.profiler.prediction_trie.trie_lookup

Classes

PredictionTrieLookup

Looks up predictions in a prediction trie with graceful fallback.

Module Contents

class PredictionTrieLookup(

root: nat.profiler.prediction_trie.data_models.PredictionTrieNode,

)

Looks up predictions in a prediction trie with graceful fallback.

_root

find(

path: list[str],

call_index: int,

) → nat.profiler.prediction_trie.data_models.LLMCallPrediction | None

Find the best matching prediction for the given path and call index.

Walks the trie as far as possible along the path, then returns the deepest match. Falls back to aggregated predictions when exact call_index isn’t found.

Args:

path: Function ancestry path (e.g., [“my_workflow”, “react_agent”]) call_index: The Nth LLM call within the current parent function

Returns:

Best matching prediction, or None if trie is empty

_get_prediction(

node: nat.profiler.prediction_trie.data_models.PredictionTrieNode,

call_index: int,

) → nat.profiler.prediction_trie.data_models.LLMCallPrediction | None

Get prediction from node, preferring exact call_index, falling back to aggregated.