lkn
/
ollie-core
Archived
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

backend: add ContextLength, Models, and auto-compaction 

- Add ContextLength(ctx) int to Backend interface; each backend fetches
  from its API (OpenAI /v1/models, Ollama /api/show, CodeWhisperer
  ListAvailableModels) with caching. Falls back to 128k default.
- Add Models(ctx) []string to Backend interface for model discovery.
- Add /models command to list available models.
- Auto-compact before each turn when estimated tokens >= 75% of context length.
This commit is contained in:
Levi Neely 2026-04-10 19:35:44 +02:00
parent 57a92281cf
commit 45a094f46e
8 changed files with 314 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
pkg/agent/core.go
View File

@ -253,6 +253,22 @@ func (s *agentCore) prompt() string {
// Prompt returns the display prompt string for the current session state.
func (s *agentCore) Prompt() string { return s.prompt() }
// defaultContextLength is used when the backend cannot report the model's
// actual context window (e.g. CodeWhisperer). 128k tokens is a safe default
// for modern models.
const defaultContextLength = 128000
// autoCompactLimit returns the token threshold for auto-compaction.
// Uses 75% of the model's context length, reserving room for output and the
// compaction prompt.
func (s *agentCore) autoCompactLimit(ctx context.Context) int {
ctxLen := s.loopcfg.Backend.ContextLength(ctx)
if ctxLen <= 0 {
ctxLen = defaultContextLength
}
return ctxLen * 3 / 4
}
func (s *agentCore) saveSession() {
if s.session == nil || s.sessionID == "" || s.sessionsDir == "" {
return
@ -302,6 +318,14 @@ func (s *agentCore) Submit(ctx context.Context, input string, handler EventHandl
s.loopcfg.Output = handler
// Auto-compact before the turn if approaching the context limit.
if s.session != nil {
if limit := s.autoCompactLimit(ctx); limit > 0 && s.session.estimateTokens() >= limit {
emit(s.loopcfg, Event{Role: "info", Content: "auto-compacting context...\n"})
s.session.compact(ctx, s.loopcfg.Backend) //nolint:errcheck
}
}
s.hooks.Run(HookAgentSpawn)
err := run(actCtx, s.loopcfg, s.session)
@ -364,6 +388,22 @@ func (s *agentCore) handleCommand(ctx context.Context, input string, handler Eve
handler(infoEvent(fmt.Sprintf("switched backend to: %s (model: %s)", be.Name(), be.Model())))
return true
case "/models":
models := s.loopcfg.Backend.Models(ctx)
if len(models) == 0 {
handler(infoEvent("no models available"))
return true
}
current := s.loopcfg.Backend.Model()
for _, m := range models {
marker := " "
if m == current {
marker = "* "
}
handler(infoEvent(marker + m))
}
return true
case "/model":
if len(args) == 0 {
handler(infoEvent("error: /model requires an argument (e.g., /model qwen3:8b)"))
@ -536,6 +576,7 @@ func (s *agentCore) handleCommand(ctx context.Context, input string, handler Eve
" /agent [name] - show or switch active agent",
" /backend <type> - switch backend (ollama, openai)",
" /model <name> - switch model",
" /models - list available models",
" /queued [pop|clear] - manage queued prompts",
" /compact - summarize conversation and compact context",
" /context - show context window debug info",

7
pkg/agent/session.go
View File

@ -243,6 +243,11 @@ func (s *Session) appendUserMessage(content string) {
// contextStatsString returns a one-line human-readable context summary.
func (s *Session) contextStatsString() string {
return fmt.Sprintf("context: ~%d tokens (%d msgs)", s.estimateTokens(), len(s.messages))
}
// estimateTokens returns a rough token count (~4 chars per token).
func (s *Session) estimateTokens() int {
chars := 0
for _, m := range s.messages {
chars += len(m.Content)
@ -250,7 +255,7 @@ func (s *Session) contextStatsString() string {
chars += len(tc.Name) + len(tc.Arguments)
}
}
return fmt.Sprintf("context: ~%d tokens (%d msgs)", chars/4, len(s.messages))
return chars / 4
}
// contextDebug returns a multi-line breakdown of the history.

10
pkg/backend/anthropic.go
View File

@ -31,6 +31,16 @@ func (b *AnthropicBackend) DefaultModel() string { return "claude-sonnet-4-5" }
func (b *AnthropicBackend) Model() string { return b.model }
func (b *AnthropicBackend) SetModel(m string) { b.model = m }
func (b *AnthropicBackend) ContextLength(_ context.Context) int { return 200000 }
func (b *AnthropicBackend) Models(_ context.Context) []string {
return []string{
"claude-sonnet-4-5",
"claude-opus-4",
"claude-3-5-haiku-latest",
}
}
// -- wire types --
type anthropicRequest struct {

5
pkg/backend/backend.go
View File

@ -86,4 +86,9 @@ type Backend interface {
Model() string
// SetModel changes the active model.
SetModel(model string)
// ContextLength returns the context window size in tokens for the active model.
// Returns 0 if unknown. May make an API call; results should be cached by the implementation.
ContextLength(ctx context.Context) int
// Models returns the list of available model IDs from the provider.
Models(ctx context.Context) []string
}

54
pkg/backend/codewhisperer.go
View File

@ -35,6 +35,7 @@ type CodeWhispererBackend struct {
authSource kiroAuthSource
authInitErr error
httpClient *http.Client
ctxLength int // cached
}
// NewCodeWhisperer returns a CodeWhisperer backend. See package doc for apiKey semantics.
@ -56,6 +57,59 @@ func (b *CodeWhispererBackend) Name() string { return "kiro" }
func (b *CodeWhispererBackend) DefaultModel() string { return "auto" }
func (b *CodeWhispererBackend) Model() string { return b.model }
func (b *CodeWhispererBackend) SetModel(m string) { b.model = m }
func (b *CodeWhispererBackend) Models(ctx context.Context) []string {
resp := b.fetchModels(ctx)
if resp == nil {
return nil
}
ids := make([]string, len(resp.Models))
for i, m := range resp.Models {
ids[i] = m.ModelID
}
return ids
}
func (b *CodeWhispererBackend) fetchModels(ctx context.Context) *kiroListModelsResponse {
if b.authInitErr != nil {
return nil
}
token, err := b.authSource.AccessToken(ctx)
if err != nil {
return nil
}
endpoint, err := b.resolveEndpoint(ctx)
if err != nil {
return nil
}
profileARN, _ := b.authSource.ProfileARN(ctx)
client := newKiroAPIClient(endpoint, token, b.extraHeaders, b.httpClient)
resp, err := client.ListModels(ctx, profileARN)
if err != nil {
return nil
}
return resp
}
func (b *CodeWhispererBackend) ContextLength(ctx context.Context) int {
if b.ctxLength > 0 {
return b.ctxLength
}
resp := b.fetchModels(ctx)
if resp == nil {
return 0
}
for _, m := range resp.Models {
if m.TokenLimits != nil && (b.model == "auto" || m.ModelID == b.model) {
b.ctxLength = m.TokenLimits.MaxInputTokens
return b.ctxLength
}
}
if resp.DefaultModel != nil && resp.DefaultModel.TokenLimits != nil {
b.ctxLength = resp.DefaultModel.TokenLimits.MaxInputTokens
return b.ctxLength
}
return 0
}
// ── Backend interface ─────────────────────────────────────────────────────────

67
pkg/backend/codewhisperer_internal.go
View File

@ -608,3 +608,70 @@ func (c *kiroOIDCClient) RefreshToken(ctx context.Context, req kiroRefreshTokenR
}
return &result, nil
}
// ── ListAvailableModels ───────────────────────────────────────────────────────
const (
kiroListModelsTarget = "AmazonCodeWhispererService.ListAvailableModels"
kiroListModelsUserAgent = "aws-sdk-rust/1.3.14 ua/2.1 api/codewhispererruntime/0.1.14474 os/linux lang/rust/1.92.0 md/appVersion-1.27.2 app/AmazonQ-For-CLI"
kiroListModelsAmzUserAgent = "aws-sdk-rust/1.3.14 ua/2.1 api/codewhispererruntime/0.1.14474 os/linux lang/rust/1.92.0 m/F,C app/AmazonQ-For-CLI"
)
type kiroListModelsRequest struct {
Origin string `json:"origin"`
ProfileARN string `json:"profileArn,omitempty"`
}
type kiroListModelsResponse struct {
DefaultModel *kiroAvailableModel `json:"defaultModel,omitempty"`
Models []kiroAvailableModel `json:"models,omitempty"`
}
type kiroAvailableModel struct {
ModelID string `json:"modelId,omitempty"`
TokenLimits *kiroTokenLimits `json:"tokenLimits,omitempty"`
}
type kiroTokenLimits struct {
MaxInputTokens int `json:"maxInputTokens,omitempty"`
MaxOutputTokens int `json:"maxOutputTokens,omitempty"`
}
func (c *kiroAPIClient) ListModels(ctx context.Context, profileARN string) (*kiroListModelsResponse, error) {
body, err := json.Marshal(kiroListModelsRequest{
Origin: kiroDefaultOrigin,
ProfileARN: profileARN,
})
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint, bytes.NewReader(body))
if err != nil {
return nil, err
}
invocationID, _ := kiroNewUUID()
req.Header.Set("Content-Type", "application/x-amz-json-1.0")
req.Header.Set("X-Amz-Target", kiroListModelsTarget)
req.Header.Set("Authorization", "Bearer "+c.token)
req.Header.Set("Accept", "*/*")
req.Header.Set("User-Agent", kiroListModelsUserAgent)
req.Header.Set("X-Amz-User-Agent", kiroListModelsAmzUserAgent)
req.Header.Set("Amz-Sdk-Request", kiroAmzSDKRequest)
req.Header.Set("Amz-Sdk-Invocation-Id", invocationID)
for k, v := range c.extraHeaders {
req.Header.Set(k, v)
}
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return nil, kiroDecodeHTTPError(resp)
}
var result kiroListModelsResponse
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
return nil, err
}
return &result, nil
}

74
pkg/backend/ollama.go
View File

@ -8,13 +8,16 @@ import (
"fmt"
"io"
"net/http"
"strings"
)
// OllamaBackend speaks the Ollama /api/chat wire format.
type OllamaBackend struct {
baseURL string
model string
client *http.Client
baseURL string
model string
client *http.Client
ctxLength int
ctxModel string
}
func NewOllama(baseURL string) *OllamaBackend {
@ -29,7 +32,70 @@ func NewOllama(baseURL string) *OllamaBackend {
func (b *OllamaBackend) Name() string { return "ollama" }
func (b *OllamaBackend) DefaultModel() string { return "qwen3.5:9b" }
func (b *OllamaBackend) Model() string { return b.model }
func (b *OllamaBackend) SetModel(m string) { b.model = m }
func (b *OllamaBackend) SetModel(m string) { b.model = m; b.ctxLength = 0 }
func (b *OllamaBackend) Models(ctx context.Context) []string {
req, err := http.NewRequestWithContext(ctx, "GET", b.baseURL+"/api/tags", nil)
if err != nil {
return nil
}
resp, err := b.client.Do(req)
if err != nil || resp.StatusCode != 200 {
if resp != nil {
resp.Body.Close()
}
return nil
}
defer resp.Body.Close()
var result struct {
Models []struct {
Name string `json:"name"`
} `json:"models"`
}
if json.NewDecoder(resp.Body).Decode(&result) != nil {
return nil
}
names := make([]string, len(result.Models))
for i, m := range result.Models {
names[i] = m.Name
}
return names
}
func (b *OllamaBackend) ContextLength(ctx context.Context) int {
if b.ctxLength > 0 && b.ctxModel == b.model {
return b.ctxLength
}
body, _ := json.Marshal(map[string]string{"name": b.model})
req, err := http.NewRequestWithContext(ctx, "POST", b.baseURL+"/api/show", bytes.NewReader(body))
if err != nil {
return 0
}
resp, err := b.client.Do(req)
if err != nil || resp.StatusCode != 200 {
if resp != nil {
resp.Body.Close()
}
return 0
}
defer resp.Body.Close()
var result struct {
ModelInfo map[string]any `json:"model_info"`
}
if json.NewDecoder(resp.Body).Decode(&result) != nil {
return 0
}
for k, v := range result.ModelInfo {
if strings.HasSuffix(k, ".context_length") {
if f, ok := v.(float64); ok && f > 0 {
b.ctxLength = int(f)
b.ctxModel = b.model
return b.ctxLength
}
}
}
return 0
}
// -- wire types --

62
pkg/backend/openai.go
View File

@ -22,6 +22,14 @@ type OpenAIBackend struct {
model string
client *http.Client
extraHeaders map[string]string // optional; applied after Authorization
ctxLength int // cached; 0 = not yet fetched
ctxModel string // model the cached value is for
cachedModels []openAIModelInfo // cached model list
}
type openAIModelInfo struct {
ID string `json:"id"`
ContextLength int `json:"context_length"`
}
func NewOpenAI(name, baseURL, apiKey string) *OpenAIBackend {
@ -35,7 +43,59 @@ func NewOpenAI(name, baseURL, apiKey string) *OpenAIBackend {
func (b *OpenAIBackend) Name() string { return b.name }
func (b *OpenAIBackend) Model() string { return b.model }
func (b *OpenAIBackend) SetModel(m string) { b.model = m }
func (b *OpenAIBackend) SetModel(m string) { b.model = m; b.ctxLength = 0 }
func (b *OpenAIBackend) fetchModels(ctx context.Context) []openAIModelInfo {
if len(b.cachedModels) > 0 {
return b.cachedModels
}
req, err := http.NewRequestWithContext(ctx, "GET", b.baseURL+"/v1/models", nil)
if err != nil {
return nil
}
if b.apiKey != "" {
req.Header.Set("Authorization", "Bearer "+b.apiKey)
}
resp, err := b.client.Do(req)
if err != nil || resp.StatusCode != 200 {
if resp != nil {
resp.Body.Close()
}
return nil
}
defer resp.Body.Close()
var result struct {
Data []openAIModelInfo `json:"data"`
}
if json.NewDecoder(resp.Body).Decode(&result) != nil {
return nil
}
b.cachedModels = result.Data
return b.cachedModels
}
func (b *OpenAIBackend) ContextLength(ctx context.Context) int {
if b.ctxLength > 0 && b.ctxModel == b.model {
return b.ctxLength
}
for _, m := range b.fetchModels(ctx) {
if m.ID == b.model {
b.ctxLength = m.ContextLength
b.ctxModel = b.model
return b.ctxLength
}
}
return 0
}
func (b *OpenAIBackend) Models(ctx context.Context) []string {
models := b.fetchModels(ctx)
ids := make([]string, len(models))
for i, m := range models {
ids[i] = m.ID
}
return ids
}
func (b *OpenAIBackend) DefaultModel() string {
switch b.name {
case "openrouter":