feat(routing): inject native ML complexity Router scoring into telegram wrapper header

This commit is contained in:
Tim Z. 2026-04-07 13:56:19 +02:00
parent f521aefa58
commit 3881f69aa0
3 changed files with 30 additions and 14 deletions

View file

@ -1624,11 +1624,9 @@ func (al *AgentLoop) runAgentLoop(
}
if result.finalContent != "" {
var totalTokens, promptTokens, compTokens int
var totalTokens int
if usage := ts.GetLastUsage(); usage != nil {
totalTokens = usage.TotalTokens
promptTokens = usage.PromptTokens
compTokens = usage.CompletionTokens
}
version := al.GetConfig().BuildInfo.Version
@ -1644,7 +1642,7 @@ func (al *AgentLoop) runAgentLoop(
Agent: agent.ID,
Route: providerName,
Model: modelName,
Complexity: calculateComplexity(promptTokens, compTokens),
Complexity: int(ts.getComplexityScore() * 100),
Tokens: totalTokens,
ToolCalls: ts.currentIteration(),
Processing: time.Since(ts.startedAt),
@ -1850,7 +1848,8 @@ func (al *AgentLoop) runTurn(ctx context.Context, ts *turnState) (turnResult, er
ts.ingestMessage(turnCtx, al, rootMsg)
}
activeCandidates, activeModel, usedLight := al.selectCandidates(ts.agent, ts.userMessage, messages)
activeCandidates, activeModel, usedLight, score := al.selectCandidates(ts.agent, ts.userMessage, messages)
ts.setComplexityScore(score)
activeProvider := ts.agent.Provider
if usedLight && ts.agent.LightProvider != nil {
activeProvider = ts.agent.LightProvider
@ -2977,12 +2976,18 @@ func (al *AgentLoop) selectCandidates(
agent *AgentInstance,
userMsg string,
history []providers.Message,
) (candidates []providers.FallbackCandidate, model string, usedLight bool) {
) (candidates []providers.FallbackCandidate, model string, usedLight bool, score float64) {
if agent.Router == nil || len(agent.LightCandidates) == 0 {
return agent.Candidates, resolvedCandidateModel(agent.Candidates, agent.Model), false
return agent.Candidates, resolvedCandidateModel(agent.Candidates, agent.Model), false, 0
}
_, usedLight, score := agent.Router.SelectModel(userMsg, history, agent.Model)
_, usedLight, score = agent.Router.SelectModel(userMsg, history, agent.Model)
// Cache the native rule-based complexity score to the primary turnState
if ts := turnStateFromContext(context.Background()); ts != nil {
// Note: Since we don't have direct access to ts here, we actually pass it up natively in runTurn
}
if !usedLight {
logger.DebugCF("agent", "Model routing: primary model selected",
map[string]any{
@ -2990,7 +2995,7 @@ func (al *AgentLoop) selectCandidates(
"score": score,
"threshold": agent.Router.Threshold(),
})
return agent.Candidates, resolvedCandidateModel(agent.Candidates, agent.Model), false
return agent.Candidates, resolvedCandidateModel(agent.Candidates, agent.Model), false, score
}
logger.InfoCF("agent", "Model routing: light model selected",
@ -3000,7 +3005,7 @@ func (al *AgentLoop) selectCandidates(
"score": score,
"threshold": agent.Router.Threshold(),
})
return agent.LightCandidates, resolvedCandidateModel(agent.LightCandidates, agent.Router.LightModel()), true
return agent.LightCandidates, resolvedCandidateModel(agent.LightCandidates, agent.Router.LightModel()), true, score
}
// resolveContextManager selects the ContextManager implementation based on config.

View file

@ -23,7 +23,3 @@ func WrapResponse(rawOutput string, m Metrics) string {
)
return header + rawOutput
}
func calculateComplexity(promptTokens, completionTokens int) int {
return int(float64(promptTokens)*0.1 + float64(completionTokens)*0.5)
}

View file

@ -103,6 +103,9 @@ type turnState struct {
tokenBudget *atomic.Int64 // Shared token budget counter
lastFinishReason string // Last LLM finish_reason
lastUsage *providers.UsageInfo // Last LLM usage info
// Telemetry tracking
complexityScore float64 // Cached heuristic difficulty score generated by the router
// Back-reference to the owning AgentLoop (set for SubTurns only, used for hard abort cascade)
al *AgentLoop
@ -214,6 +217,18 @@ func (ts *turnState) setIteration(iteration int) {
ts.iteration = iteration
}
func (ts *turnState) setComplexityScore(score float64) {
ts.mu.Lock()
defer ts.mu.Unlock()
ts.complexityScore = score
}
func (ts *turnState) getComplexityScore() float64 {
ts.mu.RLock()
defer ts.mu.RUnlock()
return ts.complexityScore
}
func (ts *turnState) currentIteration() int {
ts.mu.RLock()
defer ts.mu.RUnlock()