Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 9abd5c3b5f | |||
| 3d403bd685 | |||
| 7b8bad5e62 | |||
| f235b8ce90 | |||
| 950c63adaa | |||
| e4c75c7c0e | |||
| 8dbcb1147f |
@@ -2,3 +2,4 @@ agentu*.yaml
|
|||||||
/agentu
|
/agentu
|
||||||
/coverage.out
|
/coverage.out
|
||||||
*.swp
|
*.swp
|
||||||
|
chat-history.txt
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ model/provider switching, and local tools.
|
|||||||
- Multiple providers from `~/.agentu/config.yaml`.
|
- Multiple providers from `~/.agentu/config.yaml`.
|
||||||
- Session-only provider, model, and thinking-level switching.
|
- Session-only provider, model, and thinking-level switching.
|
||||||
- Read-only project tools enabled by default.
|
- Read-only project tools enabled by default.
|
||||||
|
- Built-in web search and URL fetching.
|
||||||
- `--yolo` mode for file writes and shell execution.
|
- `--yolo` mode for file writes and shell execution.
|
||||||
|
|
||||||
## Quick Start
|
## Quick Start
|
||||||
@@ -19,13 +20,19 @@ model/provider switching, and local tools.
|
|||||||
mkdir -p ~/.agentu
|
mkdir -p ~/.agentu
|
||||||
cp agentu.example.yaml ~/.agentu/config.yaml
|
cp agentu.example.yaml ~/.agentu/config.yaml
|
||||||
export AGENTU_API_KEY='your-api-key'
|
export AGENTU_API_KEY='your-api-key'
|
||||||
go run ./cmd/agentu
|
go run .
|
||||||
|
```
|
||||||
|
|
||||||
|
Use dark mode when your terminal is dark:
|
||||||
|
|
||||||
|
```sh
|
||||||
|
go run . --theme dark
|
||||||
```
|
```
|
||||||
|
|
||||||
Use `--yolo` only when you want agentu to write files or run shell commands:
|
Use `--yolo` only when you want agentu to write files or run shell commands:
|
||||||
|
|
||||||
```sh
|
```sh
|
||||||
go run ./cmd/agentu --yolo
|
go run . --yolo
|
||||||
```
|
```
|
||||||
|
|
||||||
## Config
|
## Config
|
||||||
@@ -43,48 +50,52 @@ providers:
|
|||||||
loveuer:
|
loveuer:
|
||||||
base_url: https://ai.loveuer.com
|
base_url: https://ai.loveuer.com
|
||||||
api_key: ${AGENTU_API_KEY}
|
api_key: ${AGENTU_API_KEY}
|
||||||
model: deepseek-v4-flash
|
|
||||||
models:
|
models:
|
||||||
- deepseek-v4-flash
|
- id: deepseek-v4-flash
|
||||||
thinking: none
|
thinking: none
|
||||||
thinking_param: thinking
|
context: 256000
|
||||||
|
|
||||||
agent:
|
|
||||||
system_prompt: |
|
|
||||||
You are agentu, a concise local agent. Answer simple conversation directly.
|
|
||||||
Use tools only when the user's request needs local project context or local
|
|
||||||
command execution.
|
|
||||||
working_dir: .
|
|
||||||
tool_timeout: 30s
|
|
||||||
```
|
```
|
||||||
|
|
||||||
|
Web tools are enabled by default with a built-in no-key backend. agentu exposes
|
||||||
|
`web_search` for current or external web information and `web_fetch` for reading
|
||||||
|
a known URL.
|
||||||
|
|
||||||
Provider names are the keys under `providers`. At startup, agentu selects the
|
Provider names are the keys under `providers`. At startup, agentu selects the
|
||||||
first provider by name. Use `/model provider <name>` to switch providers for the
|
first provider by name. Use `/model provider <name>` to switch providers for the
|
||||||
current session.
|
current session.
|
||||||
|
|
||||||
`models` is optional. If present, `/model model <name>` is restricted to that
|
`models` is required. `/model model <name>` is restricted to that list. Each
|
||||||
list. If omitted, any model name is accepted for that provider.
|
entry is an object with:
|
||||||
|
|
||||||
`thinking` is optional. Supported values are:
|
- `id`: model name sent to the provider.
|
||||||
|
- `thinking`: optional default thinking level for that model.
|
||||||
|
- `context`: optional context-window token count used for auto-compaction.
|
||||||
|
|
||||||
|
If `model` is omitted, agentu uses the first `models` entry. The `agent` block
|
||||||
|
is optional: `working_dir`, `tool_timeout`, `system_prompt`, and
|
||||||
|
`max_context_tokens` all have defaults. `max_context_tokens` overrides the
|
||||||
|
selected model's `context` when set.
|
||||||
|
|
||||||
|
Supported thinking values are:
|
||||||
|
|
||||||
```text
|
```text
|
||||||
none, middle, high, xhigh, max
|
none, middle, high, xhigh, max
|
||||||
```
|
```
|
||||||
|
|
||||||
When configured or changed with `/model thinking ...`, the value is sent as a
|
When configured or changed with `/model thinking ...`, the value is sent as a
|
||||||
top-level OpenAI-compatible request field. `thinking_param` controls that field
|
top-level OpenAI-compatible request field named `thinking`.
|
||||||
name and defaults to `thinking`.
|
|
||||||
|
|
||||||
## CLI Flags
|
## CLI Flags
|
||||||
|
|
||||||
```sh
|
```sh
|
||||||
go run ./cmd/agentu [flags]
|
go run . [flags]
|
||||||
```
|
```
|
||||||
|
|
||||||
- `--config <path>` loads a YAML config file. Default: `~/.agentu/config.yaml`.
|
- `--config <path>` loads a YAML config file. Default: `~/.agentu/config.yaml`.
|
||||||
- `--theme light|dark` selects the TUI theme. Default: `light`.
|
- `--theme light|dark` selects the TUI theme. Default: `light`.
|
||||||
- `--plain` uses the simple line-based REPL instead of the TUI.
|
- `--plain` uses the simple line-based REPL instead of the TUI.
|
||||||
- `--yolo` enables automatic file writes and shell execution.
|
- `--yolo` enables automatic file writes and shell execution.
|
||||||
|
- `--resume <id>` resumes a saved session from `~/.agentu/sessions`.
|
||||||
|
|
||||||
## TUI Controls
|
## TUI Controls
|
||||||
|
|
||||||
@@ -98,10 +109,18 @@ go run ./cmd/agentu [flags]
|
|||||||
Slash commands:
|
Slash commands:
|
||||||
|
|
||||||
- `/clear` resets conversation history.
|
- `/clear` resets conversation history.
|
||||||
|
- `/compact` summarizes older conversation history and keeps recent context.
|
||||||
|
- `/status` shows changed files with `git status --short`.
|
||||||
|
- `/diff [--stat] [path...]` shows the unstaged git diff.
|
||||||
|
- `/test [command...]` runs a project test command; defaults to `go test ./...` for Go modules or `npm test` for Node projects.
|
||||||
- `/model` shows current provider, model, thinking, and switch usage.
|
- `/model` shows current provider, model, thinking, and switch usage.
|
||||||
- `/model provider <name>` switches provider for the current session.
|
- `/model provider <name>` switches provider for the current session.
|
||||||
- `/model model <name>` or `/model <name>` switches model for the current session.
|
- `/model model <name>` or `/model <name>` switches model for the current session.
|
||||||
- `/model thinking <none|middle|high|xhigh|max>` changes thinking for the current session.
|
- `/model thinking <none|middle|high|xhigh|max>` changes thinking for the current session.
|
||||||
|
- `/sessions` lists saved sessions.
|
||||||
|
- `/resume [id]` resumes by ID or opens an interactive picker.
|
||||||
|
- `/rename <name>` renames the current session.
|
||||||
|
- `/new` saves the current session and starts a fresh one.
|
||||||
- `/exit` or `/quit` exits.
|
- `/exit` or `/quit` exits.
|
||||||
|
|
||||||
## Local Tools
|
## Local Tools
|
||||||
@@ -110,12 +129,22 @@ Read-only tools are available by default:
|
|||||||
|
|
||||||
- `file_read`
|
- `file_read`
|
||||||
- `file_list`
|
- `file_list`
|
||||||
- `file_search`
|
- `file_search` with optional `context_lines` and `max_results`
|
||||||
|
- `code_symbols` for Go packages, imports, types, functions, and methods
|
||||||
|
- `git_status` shows changed files
|
||||||
|
- `git_diff` shows unstaged or staged diff output
|
||||||
|
- `git_log` shows recent commit history
|
||||||
|
|
||||||
|
Web tools are also available in read-only mode:
|
||||||
|
|
||||||
|
- `web_search`
|
||||||
|
- `web_fetch`
|
||||||
|
|
||||||
`--yolo` additionally enables:
|
`--yolo` additionally enables:
|
||||||
|
|
||||||
- `file_write`
|
- `file_write` for creating or replacing whole files
|
||||||
- `shell_run`
|
- `file_edit` for surgical line-range or pattern edits to existing files
|
||||||
|
- `shell_run` for non-interactive shell commands
|
||||||
|
|
||||||
Use `--yolo` for prompts that ask agentu to run or inspect local commands, for
|
Use `--yolo` for prompts that ask agentu to run or inspect local commands, for
|
||||||
example:
|
example:
|
||||||
@@ -124,6 +153,8 @@ example:
|
|||||||
请帮我运行 `go test ./...` 并总结结果
|
请帮我运行 `go test ./...` 并总结结果
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
## Development
|
## Development
|
||||||
|
|
||||||
Run the verification suite:
|
Run the verification suite:
|
||||||
@@ -131,7 +162,7 @@ Run the verification suite:
|
|||||||
```sh
|
```sh
|
||||||
go test ./...
|
go test ./...
|
||||||
go vet ./...
|
go vet ./...
|
||||||
go build ./cmd/agentu
|
go build .
|
||||||
```
|
```
|
||||||
|
|
||||||
Local config and build outputs are intentionally ignored by git.
|
Local config and build outputs are intentionally ignored by git.
|
||||||
|
|||||||
+3
-12
@@ -2,16 +2,7 @@ providers:
|
|||||||
loveuer:
|
loveuer:
|
||||||
base_url: https://ai.loveuer.com
|
base_url: https://ai.loveuer.com
|
||||||
api_key: ${AGENTU_API_KEY}
|
api_key: ${AGENTU_API_KEY}
|
||||||
model: deepseek-v4-flash
|
|
||||||
models:
|
models:
|
||||||
- deepseek-v4-flash
|
- id: deepseek-v4-flash
|
||||||
thinking: none
|
thinking: none
|
||||||
thinking_param: thinking
|
context: 256000
|
||||||
|
|
||||||
agent:
|
|
||||||
system_prompt: |
|
|
||||||
You are agentu, a concise local agent. Answer simple conversation directly.
|
|
||||||
Use tools only when the user's request needs local project context or local
|
|
||||||
command execution.
|
|
||||||
working_dir: .
|
|
||||||
tool_timeout: 30s
|
|
||||||
|
|||||||
@@ -1,129 +0,0 @@
|
|||||||
package main
|
|
||||||
|
|
||||||
import (
|
|
||||||
"bufio"
|
|
||||||
"context"
|
|
||||||
"errors"
|
|
||||||
"flag"
|
|
||||||
"fmt"
|
|
||||||
"net/http"
|
|
||||||
"os"
|
|
||||||
"os/signal"
|
|
||||||
"strings"
|
|
||||||
|
|
||||||
"agentu/internal/agent"
|
|
||||||
"agentu/internal/config"
|
|
||||||
"agentu/internal/llm"
|
|
||||||
"agentu/internal/session"
|
|
||||||
"agentu/internal/tools"
|
|
||||||
"agentu/internal/tui"
|
|
||||||
)
|
|
||||||
|
|
||||||
func main() {
|
|
||||||
if err := run(); err != nil {
|
|
||||||
fmt.Fprintln(os.Stderr, "agentu:", err)
|
|
||||||
os.Exit(1)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func run() error {
|
|
||||||
configPath := flag.String("config", config.DefaultPath, "path to YAML config")
|
|
||||||
yolo := flag.Bool("yolo", false, "also enable automatic file writes and shell execution")
|
|
||||||
plain := flag.Bool("plain", false, "use the simple line-based REPL instead of the TUI")
|
|
||||||
themeName := flag.String("theme", "light", "TUI theme: light or dark")
|
|
||||||
flag.Parse()
|
|
||||||
|
|
||||||
themeMode, err := tui.ParseThemeMode(*themeName)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
cfg, err := config.Load(*configPath)
|
|
||||||
if err != nil {
|
|
||||||
if errors.Is(err, os.ErrNotExist) {
|
|
||||||
resolvedPath, resolveErr := config.ResolvePath(*configPath)
|
|
||||||
if resolveErr != nil {
|
|
||||||
return resolveErr
|
|
||||||
}
|
|
||||||
return fmt.Errorf("config not found at %s; create ~/.agentu/config.yaml from agentu.example.yaml and set AGENTU_API_KEY", resolvedPath)
|
|
||||||
}
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
providerName := cfg.DefaultProviderName()
|
|
||||||
providerConfig := cfg.Providers[providerName]
|
|
||||||
provider := llm.NewOpenAICompatibleClient(providerConfig.BaseURL, providerConfig.APIKey, http.DefaultClient)
|
|
||||||
|
|
||||||
var registry *tools.Registry
|
|
||||||
if *yolo {
|
|
||||||
registry = tools.Builtins(cfg.Agent.WorkingDir)
|
|
||||||
} else {
|
|
||||||
registry = tools.ReadOnlyBuiltins(cfg.Agent.WorkingDir)
|
|
||||||
}
|
|
||||||
systemPrompt := strings.TrimSpace(cfg.Agent.SystemPrompt + "\n\n" + tools.AgentInstructions(cfg.Agent.WorkingDir, *yolo))
|
|
||||||
|
|
||||||
assistant := agent.New(agent.Options{
|
|
||||||
Provider: provider,
|
|
||||||
ProviderName: providerName,
|
|
||||||
Model: providerConfig.Model,
|
|
||||||
Thinking: providerConfig.Thinking,
|
|
||||||
ThinkingParam: providerConfig.ThinkingParam,
|
|
||||||
ThinkingEnabled: providerConfig.ThinkingConfigured,
|
|
||||||
SystemPrompt: systemPrompt,
|
|
||||||
ToolRegistry: registry,
|
|
||||||
ToolTimeout: cfg.Agent.ToolTimeout,
|
|
||||||
})
|
|
||||||
modelManager, err := session.NewManager(cfg, assistant, http.DefaultClient)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
|
|
||||||
defer stop()
|
|
||||||
|
|
||||||
if !*plain {
|
|
||||||
return tui.Run(ctx, assistant, tui.Options{
|
|
||||||
ModelName: providerConfig.Model,
|
|
||||||
Yolo: *yolo,
|
|
||||||
ThemeMode: themeMode,
|
|
||||||
ModelManager: modelManager,
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
return repl(ctx, assistant)
|
|
||||||
}
|
|
||||||
|
|
||||||
func repl(ctx context.Context, assistant *agent.Agent) error {
|
|
||||||
scanner := bufio.NewScanner(os.Stdin)
|
|
||||||
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
|
|
||||||
|
|
||||||
fmt.Println("agentu")
|
|
||||||
fmt.Println("Type /exit to quit, /clear to reset context.")
|
|
||||||
for {
|
|
||||||
fmt.Print("> ")
|
|
||||||
if !scanner.Scan() {
|
|
||||||
if err := scanner.Err(); err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
fmt.Println()
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
input := strings.TrimSpace(scanner.Text())
|
|
||||||
switch input {
|
|
||||||
case "":
|
|
||||||
continue
|
|
||||||
case "/exit", "/quit":
|
|
||||||
return nil
|
|
||||||
case "/clear":
|
|
||||||
assistant.Clear()
|
|
||||||
fmt.Println("context cleared")
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := assistant.RunTurn(ctx, input, os.Stdout, os.Stderr); err != nil {
|
|
||||||
fmt.Fprintln(os.Stderr, "\nerror:", err)
|
|
||||||
}
|
|
||||||
fmt.Println()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -7,6 +7,7 @@ require (
|
|||||||
github.com/charmbracelet/bubbletea v1.3.10
|
github.com/charmbracelet/bubbletea v1.3.10
|
||||||
github.com/charmbracelet/glamour v1.0.0
|
github.com/charmbracelet/glamour v1.0.0
|
||||||
github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834
|
github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834
|
||||||
|
golang.org/x/net v0.38.0
|
||||||
gopkg.in/yaml.v3 v3.0.1
|
gopkg.in/yaml.v3 v3.0.1
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -39,7 +40,6 @@ require (
|
|||||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
|
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
|
||||||
github.com/yuin/goldmark v1.7.13 // indirect
|
github.com/yuin/goldmark v1.7.13 // indirect
|
||||||
github.com/yuin/goldmark-emoji v1.0.6 // indirect
|
github.com/yuin/goldmark-emoji v1.0.6 // indirect
|
||||||
golang.org/x/net v0.38.0 // indirect
|
|
||||||
golang.org/x/sys v0.38.0 // indirect
|
golang.org/x/sys v0.38.0 // indirect
|
||||||
golang.org/x/term v0.36.0 // indirect
|
golang.org/x/term v0.36.0 // indirect
|
||||||
golang.org/x/text v0.30.0 // indirect
|
golang.org/x/text v0.30.0 // indirect
|
||||||
|
|||||||
+353
-53
@@ -7,26 +7,51 @@ import (
|
|||||||
"io"
|
"io"
|
||||||
"sort"
|
"sort"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"agentu/internal/llm"
|
|
||||||
"agentu/internal/tools"
|
"agentu/internal/tools"
|
||||||
|
"agentu/pkg/llm"
|
||||||
)
|
)
|
||||||
|
|
||||||
const defaultMaxToolIterations = 8
|
const (
|
||||||
|
defaultMaxToolRounds = 50
|
||||||
|
doomLoopThreshold = 3
|
||||||
|
compactRecentMessages = 6
|
||||||
|
summaryMessageLimit = 8000
|
||||||
|
contextSummaryPrefix = "[context summary]"
|
||||||
|
toolLogArgumentLimit = 400
|
||||||
|
toolLogOutputPreviewBytes = 4096
|
||||||
|
toolLogOutputMarker = "[output]"
|
||||||
|
|
||||||
|
summarizerSystemPrompt = "You are a conversation summarizer. Summarize the following conversation history into a concise but detailed summary. Preserve:\n" +
|
||||||
|
"- Key facts, decisions, and conclusions\n" +
|
||||||
|
"- File paths and code changes made\n" +
|
||||||
|
"- Tool call results (especially errors)\n" +
|
||||||
|
"- Current task state and next steps\n\n" +
|
||||||
|
"Do NOT include: tool call syntax, raw JSON, or verbatim file contents.\n" +
|
||||||
|
"Output only the summary, no preamble."
|
||||||
|
)
|
||||||
|
|
||||||
|
type toolExecutionResult struct {
|
||||||
|
call llm.ToolCall
|
||||||
|
output string
|
||||||
|
}
|
||||||
|
|
||||||
type Agent struct {
|
type Agent struct {
|
||||||
provider llm.Provider
|
provider llm.Provider
|
||||||
providerName string
|
providerName string
|
||||||
model string
|
model string
|
||||||
thinking string
|
thinking string
|
||||||
thinkingParam string
|
thinkingParam string
|
||||||
thinkingEnabled bool
|
thinkingEnabled bool
|
||||||
systemPrompt string
|
systemPrompt string
|
||||||
toolRegistry *tools.Registry
|
toolRegistry *tools.Registry
|
||||||
toolTimeout time.Duration
|
toolTimeout time.Duration
|
||||||
maxToolIterations int
|
maxToolRounds int
|
||||||
messages []llm.Message
|
maxContextTokens int
|
||||||
|
lastUsage *llm.Usage
|
||||||
|
messages []llm.Message
|
||||||
}
|
}
|
||||||
|
|
||||||
type Options struct {
|
type Options struct {
|
||||||
@@ -40,47 +65,52 @@ type Options struct {
|
|||||||
ToolRegistry *tools.Registry
|
ToolRegistry *tools.Registry
|
||||||
ToolTimeout time.Duration
|
ToolTimeout time.Duration
|
||||||
MaxToolIterations int
|
MaxToolIterations int
|
||||||
|
MaxContextTokens int
|
||||||
}
|
}
|
||||||
|
|
||||||
func New(opts Options) *Agent {
|
func New(opts Options) *Agent {
|
||||||
maxToolIterations := opts.MaxToolIterations
|
maxToolRounds := opts.MaxToolIterations
|
||||||
if maxToolIterations <= 0 {
|
if maxToolRounds <= 0 {
|
||||||
maxToolIterations = defaultMaxToolIterations
|
maxToolRounds = defaultMaxToolRounds
|
||||||
}
|
}
|
||||||
toolTimeout := opts.ToolTimeout
|
toolTimeout := opts.ToolTimeout
|
||||||
if toolTimeout <= 0 {
|
if toolTimeout <= 0 {
|
||||||
toolTimeout = 30 * time.Second
|
toolTimeout = 30 * time.Second
|
||||||
}
|
}
|
||||||
a := &Agent{
|
a := &Agent{
|
||||||
provider: opts.Provider,
|
provider: opts.Provider,
|
||||||
providerName: opts.ProviderName,
|
providerName: opts.ProviderName,
|
||||||
model: opts.Model,
|
model: opts.Model,
|
||||||
thinking: opts.Thinking,
|
thinking: opts.Thinking,
|
||||||
thinkingParam: opts.ThinkingParam,
|
thinkingParam: opts.ThinkingParam,
|
||||||
thinkingEnabled: opts.ThinkingEnabled,
|
thinkingEnabled: opts.ThinkingEnabled,
|
||||||
systemPrompt: opts.SystemPrompt,
|
systemPrompt: opts.SystemPrompt,
|
||||||
toolRegistry: opts.ToolRegistry,
|
toolRegistry: opts.ToolRegistry,
|
||||||
toolTimeout: toolTimeout,
|
toolTimeout: toolTimeout,
|
||||||
maxToolIterations: maxToolIterations,
|
maxToolRounds: maxToolRounds,
|
||||||
|
maxContextTokens: opts.MaxContextTokens,
|
||||||
}
|
}
|
||||||
a.Clear()
|
a.Clear()
|
||||||
return a
|
return a
|
||||||
}
|
}
|
||||||
|
|
||||||
type RuntimeOptions struct {
|
type RuntimeOptions struct {
|
||||||
Provider llm.Provider
|
Provider llm.Provider
|
||||||
ProviderName string
|
ProviderName string
|
||||||
Model string
|
Model string
|
||||||
Thinking string
|
Thinking string
|
||||||
ThinkingParam string
|
ThinkingParam string
|
||||||
ThinkingEnabled bool
|
ThinkingEnabled bool
|
||||||
|
MaxContextTokens int
|
||||||
}
|
}
|
||||||
|
|
||||||
type RuntimeInfo struct {
|
type RuntimeInfo struct {
|
||||||
ProviderName string
|
ProviderName string
|
||||||
Model string
|
Model string
|
||||||
Thinking string
|
Thinking string
|
||||||
ThinkingEnabled bool
|
ThinkingEnabled bool
|
||||||
|
ContextTokens int
|
||||||
|
MaxContextTokens int
|
||||||
}
|
}
|
||||||
|
|
||||||
func (a *Agent) SetRuntime(opts RuntimeOptions) {
|
func (a *Agent) SetRuntime(opts RuntimeOptions) {
|
||||||
@@ -96,18 +126,22 @@ func (a *Agent) SetRuntime(opts RuntimeOptions) {
|
|||||||
a.thinking = opts.Thinking
|
a.thinking = opts.Thinking
|
||||||
a.thinkingParam = opts.ThinkingParam
|
a.thinkingParam = opts.ThinkingParam
|
||||||
a.thinkingEnabled = opts.ThinkingEnabled
|
a.thinkingEnabled = opts.ThinkingEnabled
|
||||||
|
a.maxContextTokens = opts.MaxContextTokens
|
||||||
}
|
}
|
||||||
|
|
||||||
func (a *Agent) RuntimeInfo() RuntimeInfo {
|
func (a *Agent) RuntimeInfo() RuntimeInfo {
|
||||||
return RuntimeInfo{
|
return RuntimeInfo{
|
||||||
ProviderName: a.providerName,
|
ProviderName: a.providerName,
|
||||||
Model: a.model,
|
Model: a.model,
|
||||||
Thinking: a.thinking,
|
Thinking: a.thinking,
|
||||||
ThinkingEnabled: a.thinkingEnabled,
|
ThinkingEnabled: a.thinkingEnabled,
|
||||||
|
ContextTokens: llm.EstimateTokens(a.messages, a.toolDefinitions()),
|
||||||
|
MaxContextTokens: a.maxContextTokens,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (a *Agent) Clear() {
|
func (a *Agent) Clear() {
|
||||||
|
a.lastUsage = nil
|
||||||
a.messages = nil
|
a.messages = nil
|
||||||
if strings.TrimSpace(a.systemPrompt) != "" {
|
if strings.TrimSpace(a.systemPrompt) != "" {
|
||||||
a.messages = append(a.messages, llm.Message{
|
a.messages = append(a.messages, llm.Message{
|
||||||
@@ -117,14 +151,173 @@ func (a *Agent) Clear() {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Messages returns a copy of the current conversation messages.
|
||||||
|
func (a *Agent) Messages() []llm.Message {
|
||||||
|
out := make([]llm.Message, len(a.messages))
|
||||||
|
copy(out, a.messages)
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetMessages replaces the conversation messages (used for session restore).
|
||||||
|
func (a *Agent) SetMessages(messages []llm.Message) {
|
||||||
|
a.messages = append([]llm.Message(nil), messages...)
|
||||||
|
}
|
||||||
|
|
||||||
|
// LastUsage returns a copy of the most recent provider token usage data.
|
||||||
|
func (a *Agent) LastUsage() *llm.Usage {
|
||||||
|
if a.lastUsage == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
usage := *a.lastUsage
|
||||||
|
return &usage
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetLastUsage restores the most recent provider token usage data.
|
||||||
|
func (a *Agent) SetLastUsage(usage *llm.Usage) {
|
||||||
|
if usage == nil {
|
||||||
|
a.lastUsage = nil
|
||||||
|
return
|
||||||
|
}
|
||||||
|
copy := *usage
|
||||||
|
a.lastUsage = ©
|
||||||
|
}
|
||||||
|
|
||||||
|
// Compact summarizes older conversation messages while preserving the system
|
||||||
|
// prompt and recent messages intact.
|
||||||
|
func (a *Agent) Compact(ctx context.Context) error {
|
||||||
|
if len(a.messages) < 4 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if a.provider == nil {
|
||||||
|
return fmt.Errorf("provider is not configured")
|
||||||
|
}
|
||||||
|
|
||||||
|
prefixEnd := 0
|
||||||
|
if a.messages[0].Role == llm.RoleSystem {
|
||||||
|
prefixEnd = 1
|
||||||
|
}
|
||||||
|
if len(a.messages)-prefixEnd <= compactRecentMessages {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
recentStart := len(a.messages) - compactRecentMessages
|
||||||
|
if recentStart <= prefixEnd {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
middle := append([]llm.Message(nil), a.messages[prefixEnd:recentStart]...)
|
||||||
|
if len(middle) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
summary, usage, err := a.summarizeMessages(ctx, middle)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("compact context: %w", err)
|
||||||
|
}
|
||||||
|
summary = strings.TrimSpace(summary)
|
||||||
|
if summary == "" {
|
||||||
|
return fmt.Errorf("compact context: summary was empty")
|
||||||
|
}
|
||||||
|
|
||||||
|
compressed := make([]llm.Message, 0, prefixEnd+1+len(a.messages[recentStart:]))
|
||||||
|
compressed = append(compressed, a.messages[:prefixEnd]...)
|
||||||
|
compressed = append(compressed, llm.Message{
|
||||||
|
Role: llm.RoleAssistant,
|
||||||
|
Content: contextSummaryPrefix + "\n" + summary,
|
||||||
|
})
|
||||||
|
compressed = append(compressed, a.messages[recentStart:]...)
|
||||||
|
a.messages = compressed
|
||||||
|
if usage != nil {
|
||||||
|
a.lastUsage = usage
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Agent) shouldCompact() bool {
|
||||||
|
if a.maxContextTokens <= 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if a.lastUsage != nil && a.lastUsage.TotalTokens >= a.maxContextTokens {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return llm.EstimateTokens(a.messages, a.toolDefinitions()) >= a.maxContextTokens
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Agent) summarizeMessages(ctx context.Context, messages []llm.Message) (string, *llm.Usage, error) {
|
||||||
|
req := llm.ChatRequest{
|
||||||
|
Model: a.model,
|
||||||
|
Messages: []llm.Message{
|
||||||
|
{Role: llm.RoleSystem, Content: summarizerSystemPrompt},
|
||||||
|
{Role: llm.RoleUser, Content: formatMessagesForSummary(messages)},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
var summary strings.Builder
|
||||||
|
var usage *llm.Usage
|
||||||
|
if err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
|
||||||
|
if event.Content != "" {
|
||||||
|
summary.WriteString(event.Content)
|
||||||
|
}
|
||||||
|
if event.Usage != nil {
|
||||||
|
copy := *event.Usage
|
||||||
|
usage = ©
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}); err != nil {
|
||||||
|
return "", nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return strings.TrimSpace(summary.String()), usage, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func formatMessagesForSummary(messages []llm.Message) string {
|
||||||
|
var b strings.Builder
|
||||||
|
b.WriteString("Conversation history to summarize:\n\n")
|
||||||
|
for index, msg := range messages {
|
||||||
|
fmt.Fprintf(&b, "Message %d (%s):\n", index+1, msg.Role)
|
||||||
|
if msg.ToolCallID != "" {
|
||||||
|
fmt.Fprintf(&b, "tool_call_id: %s\n", msg.ToolCallID)
|
||||||
|
}
|
||||||
|
if msg.Content != "" {
|
||||||
|
fmt.Fprintf(&b, "%s\n", compact(msg.Content, summaryMessageLimit))
|
||||||
|
}
|
||||||
|
for _, call := range msg.ToolCalls {
|
||||||
|
fmt.Fprintf(&b, "tool_call: %s %s\n", call.Function.Name, compact(call.Function.Arguments, summaryMessageLimit))
|
||||||
|
}
|
||||||
|
b.WriteByte('\n')
|
||||||
|
}
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// roundSignature returns a deterministic signature for a set of tool calls
|
||||||
|
// in a single round, used for doom loop detection.
|
||||||
|
func roundSignature(calls []llm.ToolCall) string {
|
||||||
|
var b strings.Builder
|
||||||
|
for _, c := range calls {
|
||||||
|
b.WriteString(c.Function.Name)
|
||||||
|
b.WriteByte('\x00')
|
||||||
|
b.WriteString(strings.TrimSpace(c.Function.Arguments))
|
||||||
|
b.WriteByte('\x00')
|
||||||
|
}
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs io.Writer) error {
|
func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs io.Writer) error {
|
||||||
input = strings.TrimSpace(input)
|
input = strings.TrimSpace(input)
|
||||||
if input == "" {
|
if input == "" {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
a.messages = append(a.messages, llm.Message{Role: llm.RoleUser, Content: input})
|
a.messages = append(a.messages, llm.Message{Role: llm.RoleUser, Content: input})
|
||||||
|
if a.shouldCompact() {
|
||||||
|
if err := a.Compact(ctx); err != nil {
|
||||||
|
if logs != nil {
|
||||||
|
fmt.Fprintf(logs, "[compact] warning: %v\n", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
for iteration := 0; iteration <= a.maxToolIterations; iteration++ {
|
var recentSignatures []string
|
||||||
|
|
||||||
|
for round := 0; round < a.maxToolRounds; round++ {
|
||||||
content, calls, err := a.streamAssistant(ctx, out)
|
content, calls, err := a.streamAssistant(ctx, out)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
@@ -146,17 +339,52 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
|
|||||||
fmt.Fprintln(out)
|
fmt.Fprintln(out)
|
||||||
}
|
}
|
||||||
|
|
||||||
for _, call := range calls {
|
// Doom loop detection: check if the last N rounds all called
|
||||||
result := a.executeTool(ctx, call, logs)
|
// the exact same tools with the exact same arguments.
|
||||||
|
sig := roundSignature(calls)
|
||||||
|
recentSignatures = append(recentSignatures, sig)
|
||||||
|
if len(recentSignatures) > doomLoopThreshold {
|
||||||
|
recentSignatures = recentSignatures[len(recentSignatures)-doomLoopThreshold:]
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(recentSignatures) == doomLoopThreshold {
|
||||||
|
allSame := true
|
||||||
|
for i := 1; i < len(recentSignatures); i++ {
|
||||||
|
if recentSignatures[i] != recentSignatures[0] {
|
||||||
|
allSame = false
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if allSame {
|
||||||
|
toolNames := make([]string, 0, len(calls))
|
||||||
|
for _, c := range calls {
|
||||||
|
toolNames = append(toolNames, c.Function.Name)
|
||||||
|
}
|
||||||
|
msg := fmt.Sprintf(
|
||||||
|
"error: doom loop detected — the same tool call (%s) with identical arguments has been repeated %d times. Stopping. Try a different approach or break the task into smaller steps.",
|
||||||
|
strings.Join(toolNames, ", "),
|
||||||
|
doomLoopThreshold,
|
||||||
|
)
|
||||||
|
a.messages = append(a.messages, llm.Message{
|
||||||
|
Role: llm.RoleAssistant,
|
||||||
|
Content: msg,
|
||||||
|
})
|
||||||
|
return fmt.Errorf("%s", msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
results := a.executeTools(ctx, calls, logs)
|
||||||
|
for _, result := range results {
|
||||||
a.messages = append(a.messages, llm.Message{
|
a.messages = append(a.messages, llm.Message{
|
||||||
Role: llm.RoleTool,
|
Role: llm.RoleTool,
|
||||||
ToolCallID: call.ID,
|
ToolCallID: result.call.ID,
|
||||||
Content: result,
|
Content: result.output,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
||||||
return fmt.Errorf("tool iteration limit reached: %d", a.maxToolIterations)
|
return fmt.Errorf("tool round limit reached after %d rounds; the task may be too complex for a single turn", a.maxToolRounds)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []llm.ToolCall, error) {
|
func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []llm.ToolCall, error) {
|
||||||
@@ -180,6 +408,9 @@ func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []l
|
|||||||
}
|
}
|
||||||
|
|
||||||
err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
|
err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
|
||||||
|
if event.Usage != nil {
|
||||||
|
a.SetLastUsage(event.Usage)
|
||||||
|
}
|
||||||
if event.Content != "" {
|
if event.Content != "" {
|
||||||
content.WriteString(event.Content)
|
content.WriteString(event.Content)
|
||||||
if out != nil {
|
if out != nil {
|
||||||
@@ -212,11 +443,8 @@ func (a *Agent) toolDefinitions() []llm.Tool {
|
|||||||
return a.toolRegistry.Definitions()
|
return a.toolRegistry.Definitions()
|
||||||
}
|
}
|
||||||
|
|
||||||
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writer) string {
|
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall) string {
|
||||||
name := call.Function.Name
|
name := call.Function.Name
|
||||||
if logs != nil {
|
|
||||||
fmt.Fprintf(logs, "\n[tool] %s %s\n", name, compact(call.Function.Arguments, 400))
|
|
||||||
}
|
|
||||||
|
|
||||||
if a.toolRegistry == nil {
|
if a.toolRegistry == nil {
|
||||||
return tools.FormatError(fmt.Errorf("tool registry is disabled"))
|
return tools.FormatError(fmt.Errorf("tool registry is disabled"))
|
||||||
@@ -240,12 +468,84 @@ func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writ
|
|||||||
}
|
}
|
||||||
return tools.FormatError(err)
|
return tools.FormatError(err)
|
||||||
}
|
}
|
||||||
if logs != nil {
|
|
||||||
fmt.Fprintf(logs, "[tool] %s done\n", name)
|
|
||||||
}
|
|
||||||
return output
|
return output
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (a *Agent) hasMutatingCall(calls []llm.ToolCall) bool {
|
||||||
|
if a.toolRegistry == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
for _, call := range calls {
|
||||||
|
if tool, ok := a.toolRegistry.Get(call.Function.Name); ok {
|
||||||
|
if tools.IsMutating(tool) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Agent) executeTools(ctx context.Context, calls []llm.ToolCall, logs io.Writer) []toolExecutionResult {
|
||||||
|
results := make([]toolExecutionResult, len(calls))
|
||||||
|
|
||||||
|
// Log all tool starts sequentially so TUI gets ordered placeholders.
|
||||||
|
if logs != nil {
|
||||||
|
for _, call := range calls {
|
||||||
|
logToolStart(logs, call)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if a.hasMutatingCall(calls) {
|
||||||
|
// Sequential: avoid racing mutating tools.
|
||||||
|
for i, call := range calls {
|
||||||
|
output := a.executeTool(ctx, call)
|
||||||
|
results[i] = toolExecutionResult{call: call, output: output}
|
||||||
|
if logs != nil {
|
||||||
|
logToolOutput(logs, call, output)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return results
|
||||||
|
}
|
||||||
|
|
||||||
|
// Concurrent: independent read-only tools.
|
||||||
|
var mu sync.Mutex
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
for i, call := range calls {
|
||||||
|
wg.Add(1)
|
||||||
|
go func(idx int, c llm.ToolCall) {
|
||||||
|
defer wg.Done()
|
||||||
|
output := a.executeTool(ctx, c)
|
||||||
|
results[idx] = toolExecutionResult{call: c, output: output}
|
||||||
|
if logs != nil {
|
||||||
|
mu.Lock()
|
||||||
|
logToolOutput(logs, c, output)
|
||||||
|
mu.Unlock()
|
||||||
|
}
|
||||||
|
}(i, call)
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
return results
|
||||||
|
}
|
||||||
|
|
||||||
|
func logToolStart(logs io.Writer, call llm.ToolCall) {
|
||||||
|
fmt.Fprintf(logs, "\n[tool] %s %s\n", call.Function.Name, compact(call.Function.Arguments, toolLogArgumentLimit))
|
||||||
|
}
|
||||||
|
|
||||||
|
func logToolOutput(logs io.Writer, call llm.ToolCall, output string) {
|
||||||
|
fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s\n", call.Function.Name, compact(call.Function.Arguments, toolLogArgumentLimit), toolLogOutputMarker, toolOutputPreview(output))
|
||||||
|
}
|
||||||
|
|
||||||
|
func toolOutputPreview(output string) string {
|
||||||
|
output = strings.TrimSpace(output)
|
||||||
|
if output == "" {
|
||||||
|
return "(no output)"
|
||||||
|
}
|
||||||
|
if len(output) <= toolLogOutputPreviewBytes {
|
||||||
|
return output
|
||||||
|
}
|
||||||
|
return output[:toolLogOutputPreviewBytes] + fmt.Sprintf("\n[truncated: showing first %d bytes; omitted %d bytes]", toolLogOutputPreviewBytes, len(output)-toolLogOutputPreviewBytes)
|
||||||
|
}
|
||||||
|
|
||||||
type toolCallBuilder struct {
|
type toolCallBuilder struct {
|
||||||
index int
|
index int
|
||||||
id string
|
id string
|
||||||
|
|||||||
@@ -3,11 +3,12 @@ package agent
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"agentu/internal/llm"
|
|
||||||
"agentu/internal/tools"
|
"agentu/internal/tools"
|
||||||
|
"agentu/pkg/llm"
|
||||||
)
|
)
|
||||||
|
|
||||||
type fakeProvider struct {
|
type fakeProvider struct {
|
||||||
@@ -96,6 +97,134 @@ func (p *contentProvider) ChatStream(ctx context.Context, req llm.ChatRequest, e
|
|||||||
return emit(llm.StreamEvent{Content: p.text})
|
return emit(llm.StreamEvent{Content: p.text})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type compactProvider struct {
|
||||||
|
summaryCalls int
|
||||||
|
answerCalls int
|
||||||
|
lastSummary string
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *compactProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
if len(req.Messages) == 2 && req.Messages[0].Role == llm.RoleSystem && req.Messages[0].Content == summarizerSystemPrompt {
|
||||||
|
p.summaryCalls++
|
||||||
|
p.lastSummary = req.Messages[1].Content
|
||||||
|
if err := emit(llm.StreamEvent{Content: "old work summarized"}); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return emit(llm.StreamEvent{Usage: &llm.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12}})
|
||||||
|
}
|
||||||
|
|
||||||
|
p.answerCalls++
|
||||||
|
return emit(llm.StreamEvent{Content: "done"})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentCompact(t *testing.T) {
|
||||||
|
provider := &compactProvider{}
|
||||||
|
a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
|
||||||
|
a.SetMessages(longConversation())
|
||||||
|
|
||||||
|
if err := a.Compact(context.Background()); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
messages := a.Messages()
|
||||||
|
if len(messages) != 8 {
|
||||||
|
t.Fatalf("messages = %d, want 8", len(messages))
|
||||||
|
}
|
||||||
|
if messages[0].Role != llm.RoleSystem || messages[0].Content != "system" {
|
||||||
|
t.Fatalf("system message = %#v", messages[0])
|
||||||
|
}
|
||||||
|
if messages[1].Role != llm.RoleAssistant || !strings.HasPrefix(messages[1].Content, contextSummaryPrefix) {
|
||||||
|
t.Fatalf("summary message = %#v", messages[1])
|
||||||
|
}
|
||||||
|
if !strings.Contains(messages[1].Content, "old work summarized") {
|
||||||
|
t.Fatalf("summary content = %q", messages[1].Content)
|
||||||
|
}
|
||||||
|
if !strings.Contains(provider.lastSummary, "old-user-0") {
|
||||||
|
t.Fatalf("summary input did not include old history: %q", provider.lastSummary)
|
||||||
|
}
|
||||||
|
if messages[len(messages)-1].Content != "recent-answer-2" {
|
||||||
|
t.Fatalf("last recent message = %#v", messages[len(messages)-1])
|
||||||
|
}
|
||||||
|
if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 12 {
|
||||||
|
t.Fatalf("usage = %#v", usage)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentAutoCompact(t *testing.T) {
|
||||||
|
provider := &compactProvider{}
|
||||||
|
a := New(Options{
|
||||||
|
Provider: provider,
|
||||||
|
Model: "test",
|
||||||
|
SystemPrompt: "system",
|
||||||
|
MaxContextTokens: 1,
|
||||||
|
ToolRegistry: tools.NewRegistry(),
|
||||||
|
MaxToolIterations: 1,
|
||||||
|
})
|
||||||
|
a.SetMessages(longConversation())
|
||||||
|
|
||||||
|
var out strings.Builder
|
||||||
|
if err := a.RunTurn(context.Background(), "new question", &out, nil); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if provider.summaryCalls != 1 {
|
||||||
|
t.Fatalf("summaryCalls = %d, want 1", provider.summaryCalls)
|
||||||
|
}
|
||||||
|
if provider.answerCalls != 1 {
|
||||||
|
t.Fatalf("answerCalls = %d, want 1", provider.answerCalls)
|
||||||
|
}
|
||||||
|
if out.String() != "done" {
|
||||||
|
t.Fatalf("out = %q", out.String())
|
||||||
|
}
|
||||||
|
if !strings.Contains(joinMessageContents(a.Messages()), contextSummaryPrefix) {
|
||||||
|
t.Fatalf("messages missing summary: %#v", a.Messages())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentCompactSkipsShortHistory(t *testing.T) {
|
||||||
|
provider := &contentProvider{text: "unused"}
|
||||||
|
a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
|
||||||
|
a.SetMessages([]llm.Message{
|
||||||
|
{Role: llm.RoleSystem, Content: "system"},
|
||||||
|
{Role: llm.RoleUser, Content: "question"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "answer"},
|
||||||
|
})
|
||||||
|
|
||||||
|
if err := a.Compact(context.Background()); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if provider.called {
|
||||||
|
t.Fatal("provider should not be called for short history")
|
||||||
|
}
|
||||||
|
if len(a.Messages()) != 3 {
|
||||||
|
t.Fatalf("messages = %d, want 3", len(a.Messages()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func longConversation() []llm.Message {
|
||||||
|
messages := []llm.Message{{Role: llm.RoleSystem, Content: "system"}}
|
||||||
|
for i := 0; i < 3; i++ {
|
||||||
|
messages = append(messages,
|
||||||
|
llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("old-user-%d", i)},
|
||||||
|
llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("old-answer-%d", i)},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
for i := 0; i < 3; i++ {
|
||||||
|
messages = append(messages,
|
||||||
|
llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("recent-user-%d", i)},
|
||||||
|
llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("recent-answer-%d", i)},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
return messages
|
||||||
|
}
|
||||||
|
|
||||||
|
func joinMessageContents(messages []llm.Message) string {
|
||||||
|
var b strings.Builder
|
||||||
|
for _, msg := range messages {
|
||||||
|
b.WriteString(msg.Content)
|
||||||
|
b.WriteByte('\n')
|
||||||
|
}
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
func TestAgentLetsProviderDecideShellRequestWhenShellDisabled(t *testing.T) {
|
func TestAgentLetsProviderDecideShellRequestWhenShellDisabled(t *testing.T) {
|
||||||
provider := &contentProvider{text: "shell is disabled"}
|
provider := &contentProvider{text: "shell is disabled"}
|
||||||
a := New(Options{
|
a := New(Options{
|
||||||
@@ -137,3 +266,248 @@ func TestAgentDoesNotBlockFileRequest(t *testing.T) {
|
|||||||
t.Fatalf("out = %q", out.String())
|
t.Fatalf("out = %q", out.String())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// alwaysToolProvider always returns a tool call with the same args on every request.
|
||||||
|
type alwaysToolProvider struct {
|
||||||
|
calls int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *alwaysToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
p.calls++
|
||||||
|
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
|
||||||
|
{Index: 0, ID: fmt.Sprintf("call_%d", p.calls), Type: "function", Name: "test_echo", Arguments: `{"text":"loop"}`},
|
||||||
|
}})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentStopsOnDoomLoop(t *testing.T) {
|
||||||
|
provider := &alwaysToolProvider{}
|
||||||
|
a := New(Options{
|
||||||
|
Provider: provider,
|
||||||
|
Model: "test",
|
||||||
|
SystemPrompt: "system",
|
||||||
|
ToolRegistry: tools.NewRegistry(echoTool{}),
|
||||||
|
MaxToolIterations: 50, // High hard limit — doom loop should trigger first
|
||||||
|
})
|
||||||
|
|
||||||
|
var out strings.Builder
|
||||||
|
err := a.RunTurn(context.Background(), "loop forever", &out, nil)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error when doom loop is detected")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "doom loop detected") {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
// Should have stopped at doomLoopThreshold (3) rounds, not reached hard limit.
|
||||||
|
if provider.calls != doomLoopThreshold {
|
||||||
|
t.Fatalf("calls = %d, want %d (doom loop threshold)", provider.calls, doomLoopThreshold)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// varyingToolProvider alternates between different tool calls to avoid doom loop detection.
|
||||||
|
type varyingToolProvider struct {
|
||||||
|
calls int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *varyingToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
p.calls++
|
||||||
|
// Vary the arguments each round to avoid doom loop detection.
|
||||||
|
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
|
||||||
|
{Index: 0, ID: fmt.Sprintf("call_%d", p.calls), Type: "function", Name: "test_echo",
|
||||||
|
Arguments: fmt.Sprintf(`{"text":"round-%d"}`, p.calls)},
|
||||||
|
}})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentHardLimitStopsVaryingCalls(t *testing.T) {
|
||||||
|
provider := &varyingToolProvider{}
|
||||||
|
const limit = 5
|
||||||
|
a := New(Options{
|
||||||
|
Provider: provider,
|
||||||
|
Model: "test",
|
||||||
|
SystemPrompt: "system",
|
||||||
|
ToolRegistry: tools.NewRegistry(echoTool{}),
|
||||||
|
MaxToolIterations: limit,
|
||||||
|
})
|
||||||
|
|
||||||
|
var out strings.Builder
|
||||||
|
err := a.RunTurn(context.Background(), "vary forever", &out, nil)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error when hard limit is reached")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "tool round limit reached") {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if provider.calls != limit {
|
||||||
|
t.Fatalf("calls = %d, want %d", provider.calls, limit)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// slowTool blocks on a channel before returning.
|
||||||
|
type slowTool struct {
|
||||||
|
name string
|
||||||
|
started chan<- string
|
||||||
|
release <-chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *slowTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: t.name,
|
||||||
|
Description: "slow test tool",
|
||||||
|
Parameters: json.RawMessage(`{"type":"object","properties":{"text":{"type":"string"}}}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *slowTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args struct {
|
||||||
|
Text string `json:"text"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
t.started <- t.name
|
||||||
|
<-t.release
|
||||||
|
return args.Text, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// mutatingSlowTool is like slowTool but implements MutatingTool.
|
||||||
|
type mutatingSlowTool struct {
|
||||||
|
slowTool
|
||||||
|
}
|
||||||
|
|
||||||
|
func (*mutatingSlowTool) Mutates() bool { return true }
|
||||||
|
|
||||||
|
func TestAgentExecutesIndependentToolCallsInParallel(t *testing.T) {
|
||||||
|
started := make(chan string, 2)
|
||||||
|
release := make(chan struct{})
|
||||||
|
|
||||||
|
tool1 := &slowTool{name: "slow_one", started: started, release: release}
|
||||||
|
tool2 := &slowTool{name: "slow_two", started: started, release: release}
|
||||||
|
|
||||||
|
provider := ¶llelProvider{
|
||||||
|
responses: [][]llm.ToolCallDelta{
|
||||||
|
{
|
||||||
|
{Index: 0, ID: "call_1", Type: "function", Name: "slow_one", Arguments: `{"text":"one"}`},
|
||||||
|
{Index: 1, ID: "call_2", Type: "function", Name: "slow_two", Arguments: `{"text":"two"}`},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
finalContent: "done",
|
||||||
|
}
|
||||||
|
|
||||||
|
a := New(Options{
|
||||||
|
Provider: provider,
|
||||||
|
Model: "test",
|
||||||
|
SystemPrompt: "system",
|
||||||
|
ToolRegistry: tools.NewRegistry(tool1, tool2),
|
||||||
|
})
|
||||||
|
|
||||||
|
var out strings.Builder
|
||||||
|
done := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
done <- a.RunTurn(context.Background(), "run both", &out, nil)
|
||||||
|
}()
|
||||||
|
|
||||||
|
// Both tools must start before we release them.
|
||||||
|
names := make(map[string]bool)
|
||||||
|
for i := 0; i < 2; i++ {
|
||||||
|
select {
|
||||||
|
case name := <-started:
|
||||||
|
names[name] = true
|
||||||
|
case <-make(chan struct{}):
|
||||||
|
// Use a timer in real code, but for test simplicity we rely on deadlock detection.
|
||||||
|
t.Fatal("timed out waiting for tool to start")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !names["slow_one"] || !names["slow_two"] {
|
||||||
|
t.Fatalf("expected both tools to start, got %v", names)
|
||||||
|
}
|
||||||
|
close(release)
|
||||||
|
|
||||||
|
if err := <-done; err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if out.String() != "done" {
|
||||||
|
t.Fatalf("out = %q", out.String())
|
||||||
|
}
|
||||||
|
// Verify tool results were in order.
|
||||||
|
msgs := a.Messages()
|
||||||
|
var toolMsgs []string
|
||||||
|
for _, m := range msgs {
|
||||||
|
if m.Role == llm.RoleTool {
|
||||||
|
toolMsgs = append(toolMsgs, m.Content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(toolMsgs) != 2 || toolMsgs[0] != "one" || toolMsgs[1] != "two" {
|
||||||
|
t.Fatalf("tool messages = %v, want [one two]", toolMsgs)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentRunsMutatingToolCallsSequentially(t *testing.T) {
|
||||||
|
started := make(chan string, 2)
|
||||||
|
release := make(chan struct{})
|
||||||
|
|
||||||
|
tool1 := &mutatingSlowTool{slowTool{name: "mut_one", started: started, release: release}}
|
||||||
|
tool2 := &mutatingSlowTool{slowTool{name: "mut_two", started: started, release: release}}
|
||||||
|
|
||||||
|
provider := ¶llelProvider{
|
||||||
|
responses: [][]llm.ToolCallDelta{
|
||||||
|
{
|
||||||
|
{Index: 0, ID: "call_1", Type: "function", Name: "mut_one", Arguments: `{"text":"one"}`},
|
||||||
|
{Index: 1, ID: "call_2", Type: "function", Name: "mut_two", Arguments: `{"text":"two"}`},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
finalContent: "done",
|
||||||
|
}
|
||||||
|
|
||||||
|
a := New(Options{
|
||||||
|
Provider: provider,
|
||||||
|
Model: "test",
|
||||||
|
SystemPrompt: "system",
|
||||||
|
ToolRegistry: tools.NewRegistry(tool1, tool2),
|
||||||
|
})
|
||||||
|
|
||||||
|
var out strings.Builder
|
||||||
|
done := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
done <- a.RunTurn(context.Background(), "run both", &out, nil)
|
||||||
|
}()
|
||||||
|
|
||||||
|
// Only the first tool should start.
|
||||||
|
select {
|
||||||
|
case name := <-started:
|
||||||
|
if name != "mut_one" {
|
||||||
|
t.Fatalf("expected mut_one to start first, got %q", name)
|
||||||
|
}
|
||||||
|
case <-make(chan struct{}):
|
||||||
|
t.Fatal("timed out waiting for first tool to start")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Second tool should NOT have started yet.
|
||||||
|
select {
|
||||||
|
case name := <-started:
|
||||||
|
t.Fatalf("second tool started too early: %q", name)
|
||||||
|
default:
|
||||||
|
// Expected: second tool hasn't started.
|
||||||
|
}
|
||||||
|
|
||||||
|
close(release)
|
||||||
|
if err := <-done; err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// parallelProvider emits a batch of tool calls on the first request, then returns finalContent.
|
||||||
|
type parallelProvider struct {
|
||||||
|
responses [][]llm.ToolCallDelta
|
||||||
|
finalContent string
|
||||||
|
calls int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *parallelProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
p.calls++
|
||||||
|
if p.calls <= len(p.responses) {
|
||||||
|
return emit(llm.StreamEvent{ToolCalls: p.responses[p.calls-1]})
|
||||||
|
}
|
||||||
|
return emit(llm.StreamEvent{Content: p.finalContent})
|
||||||
|
}
|
||||||
|
|||||||
+186
-23
@@ -5,21 +5,24 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"sort"
|
"sort"
|
||||||
"strings"
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
"agentu/internal/agent"
|
"agentu/internal/agent"
|
||||||
"agentu/internal/config"
|
"agentu/pkg/config"
|
||||||
"agentu/internal/llm"
|
"agentu/pkg/llm"
|
||||||
)
|
)
|
||||||
|
|
||||||
type Manager struct {
|
type Manager struct {
|
||||||
cfg *config.Config
|
cfg *config.Config
|
||||||
agent *agent.Agent
|
agent *agent.Agent
|
||||||
httpClient *http.Client
|
httpClient *http.Client
|
||||||
providerName string
|
providerName string
|
||||||
provider config.ProviderConfig
|
provider config.ProviderConfig
|
||||||
|
store *Store
|
||||||
|
currentSession *Session
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewManager(cfg *config.Config, assistant *agent.Agent, httpClient *http.Client) (*Manager, error) {
|
func NewManager(cfg *config.Config, assistant *agent.Agent, httpClient *http.Client, store *Store) (*Manager, error) {
|
||||||
if httpClient == nil {
|
if httpClient == nil {
|
||||||
httpClient = http.DefaultClient
|
httpClient = http.DefaultClient
|
||||||
}
|
}
|
||||||
@@ -28,13 +31,162 @@ func NewManager(cfg *config.Config, assistant *agent.Agent, httpClient *http.Cli
|
|||||||
if !ok {
|
if !ok {
|
||||||
return nil, fmt.Errorf("no provider is configured")
|
return nil, fmt.Errorf("no provider is configured")
|
||||||
}
|
}
|
||||||
return &Manager{
|
m := &Manager{
|
||||||
cfg: cfg,
|
cfg: cfg,
|
||||||
agent: assistant,
|
agent: assistant,
|
||||||
httpClient: httpClient,
|
httpClient: httpClient,
|
||||||
providerName: providerName,
|
providerName: providerName,
|
||||||
provider: provider,
|
provider: provider,
|
||||||
}, nil
|
store: store,
|
||||||
|
}
|
||||||
|
return m, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ResumeLatest loads the most recent session, or creates a new one.
|
||||||
|
func (m *Manager) ResumeLatest() error {
|
||||||
|
if m.store == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
metas, err := m.store.List()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if len(metas) > 0 {
|
||||||
|
return m.Resume(metas[0].ID)
|
||||||
|
}
|
||||||
|
m.newSession()
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Resume loads a session by ID and restores it into the agent.
|
||||||
|
func (m *Manager) Resume(id string) error {
|
||||||
|
if m.store == nil {
|
||||||
|
return fmt.Errorf("session storage is not configured")
|
||||||
|
}
|
||||||
|
sess, err := m.store.Load(id)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
m.currentSession = sess
|
||||||
|
m.agent.SetMessages(sess.Messages)
|
||||||
|
m.agent.SetLastUsage(sess.LastUsage)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Save persists the current session to disk.
|
||||||
|
func (m *Manager) Save() error {
|
||||||
|
if m.store == nil || m.currentSession == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
m.currentSession.Messages = m.agent.Messages()
|
||||||
|
m.currentSession.LastUsage = m.agent.LastUsage()
|
||||||
|
m.currentSession.Provider = m.providerName
|
||||||
|
m.currentSession.Model = m.provider.Model
|
||||||
|
// Auto-name if empty
|
||||||
|
if m.currentSession.Name == "" {
|
||||||
|
m.currentSession.Name = DefaultName(m.currentSession.Messages)
|
||||||
|
}
|
||||||
|
return m.store.Save(m.currentSession)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Rename changes the current session name and saves.
|
||||||
|
func (m *Manager) Rename(name string) (string, error) {
|
||||||
|
name = strings.TrimSpace(name)
|
||||||
|
if name == "" {
|
||||||
|
return "", fmt.Errorf("name is required")
|
||||||
|
}
|
||||||
|
if m.currentSession == nil {
|
||||||
|
return "", fmt.Errorf("no active session")
|
||||||
|
}
|
||||||
|
m.currentSession.Name = name
|
||||||
|
if err := m.Save(); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("Renamed session to %q.", name), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sessions returns a formatted list of all sessions.
|
||||||
|
func (m *Manager) Sessions() (string, error) {
|
||||||
|
if m.store == nil {
|
||||||
|
return "", fmt.Errorf("session storage is not configured")
|
||||||
|
}
|
||||||
|
metas, err := m.store.List()
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if len(metas) == 0 {
|
||||||
|
return "No saved sessions.", nil
|
||||||
|
}
|
||||||
|
var b strings.Builder
|
||||||
|
b.WriteString("Sessions:\n")
|
||||||
|
for _, meta := range metas {
|
||||||
|
name := meta.Name
|
||||||
|
if name == "" {
|
||||||
|
name = "(unnamed)"
|
||||||
|
}
|
||||||
|
current := ""
|
||||||
|
if m.currentSession != nil && meta.ID == m.currentSession.ID {
|
||||||
|
current = " (current)"
|
||||||
|
}
|
||||||
|
ts := meta.UpdatedAt.Local().Format("2006-01-02 15:04")
|
||||||
|
fmt.Fprintf(&b, " %s %-30s %s %d msgs%s\n", meta.ID, fitName(name, 30), ts, meta.MessageCount, current)
|
||||||
|
}
|
||||||
|
return b.String(), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewSession starts a fresh session, saving the current one first.
|
||||||
|
func (m *Manager) NewSession() (string, error) {
|
||||||
|
// Save current if it has user messages
|
||||||
|
if m.currentSession != nil {
|
||||||
|
msgs := m.agent.Messages()
|
||||||
|
hasUser := false
|
||||||
|
for _, msg := range msgs {
|
||||||
|
if msg.Role == llm.RoleUser {
|
||||||
|
hasUser = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if hasUser {
|
||||||
|
m.currentSession.Messages = msgs
|
||||||
|
m.currentSession.Provider = m.providerName
|
||||||
|
m.currentSession.Model = m.provider.Model
|
||||||
|
if m.currentSession.Name == "" {
|
||||||
|
m.currentSession.Name = DefaultName(msgs)
|
||||||
|
}
|
||||||
|
if m.store != nil {
|
||||||
|
_ = m.store.Save(m.currentSession)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
m.newSession()
|
||||||
|
return "Session saved. Starting new session.", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) newSession() {
|
||||||
|
m.currentSession = &Session{
|
||||||
|
ID: GenerateID(),
|
||||||
|
CreatedAt: time.Now().UTC(),
|
||||||
|
Provider: m.providerName,
|
||||||
|
Model: m.provider.Model,
|
||||||
|
}
|
||||||
|
m.agent.Clear()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) CurrentSessionID() string {
|
||||||
|
if m.currentSession == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return m.currentSession.ID
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) CurrentSessionName() string {
|
||||||
|
if m.currentSession == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
if m.currentSession.Name != "" {
|
||||||
|
return m.currentSession.Name
|
||||||
|
}
|
||||||
|
return m.currentSession.ID
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *Manager) CurrentProvider() string {
|
func (m *Manager) CurrentProvider() string {
|
||||||
@@ -58,11 +210,7 @@ func (m *Manager) Info() string {
|
|||||||
fmt.Fprintf(&b, "model: %s\n", m.CurrentModel())
|
fmt.Fprintf(&b, "model: %s\n", m.CurrentModel())
|
||||||
fmt.Fprintf(&b, "thinking: %s\n", m.CurrentThinking())
|
fmt.Fprintf(&b, "thinking: %s\n", m.CurrentThinking())
|
||||||
fmt.Fprintf(&b, "providers: %s\n", strings.Join(m.providerNames(), ", "))
|
fmt.Fprintf(&b, "providers: %s\n", strings.Join(m.providerNames(), ", "))
|
||||||
if len(m.provider.Models) > 0 {
|
fmt.Fprintf(&b, "models: %s\n", strings.Join(m.provider.Models, ", "))
|
||||||
fmt.Fprintf(&b, "models: %s\n", strings.Join(m.provider.Models, ", "))
|
|
||||||
} else {
|
|
||||||
b.WriteString("models: any model name is accepted for this provider\n")
|
|
||||||
}
|
|
||||||
fmt.Fprintf(&b, "thinking levels: %s\n", strings.Join(config.ThinkingLevels(), ", "))
|
fmt.Fprintf(&b, "thinking levels: %s\n", strings.Join(config.ThinkingLevels(), ", "))
|
||||||
b.WriteString("usage: /model provider <name> | /model model <name> | /model thinking <level>")
|
b.WriteString("usage: /model provider <name> | /model model <name> | /model thinking <level>")
|
||||||
return b.String()
|
return b.String()
|
||||||
@@ -84,10 +232,10 @@ func (m *Manager) SetModel(model string) (string, error) {
|
|||||||
if model == "" {
|
if model == "" {
|
||||||
return "", fmt.Errorf("model is required")
|
return "", fmt.Errorf("model is required")
|
||||||
}
|
}
|
||||||
if len(m.provider.Models) > 0 && !contains(m.provider.Models, model) {
|
if !contains(m.provider.Models, model) {
|
||||||
return "", fmt.Errorf("model %q is not configured for provider %q; available: %s", model, m.providerName, strings.Join(m.provider.Models, ", "))
|
return "", fmt.Errorf("model %q is not configured for provider %q; available: %s", model, m.providerName, strings.Join(m.provider.Models, ", "))
|
||||||
}
|
}
|
||||||
m.provider.Model = model
|
m.provider = m.provider.WithModel(model)
|
||||||
m.cfg.Providers[m.providerName] = m.provider
|
m.cfg.Providers[m.providerName] = m.provider
|
||||||
m.syncAgent(nil)
|
m.syncAgent(nil)
|
||||||
return "Switched model.\n" + m.summary(), nil
|
return "Switched model.\n" + m.summary(), nil
|
||||||
@@ -112,14 +260,22 @@ func (m *Manager) summary() string {
|
|||||||
return fmt.Sprintf("provider: %s\nmodel: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentThinking())
|
return fmt.Sprintf("provider: %s\nmodel: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentThinking())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (m *Manager) maxContextTokens() int {
|
||||||
|
if m.cfg != nil && m.cfg.Agent.MaxContextTokens > 0 {
|
||||||
|
return m.cfg.Agent.MaxContextTokens
|
||||||
|
}
|
||||||
|
return m.provider.ContextTokens
|
||||||
|
}
|
||||||
|
|
||||||
func (m *Manager) syncAgent(provider llm.Provider) {
|
func (m *Manager) syncAgent(provider llm.Provider) {
|
||||||
m.agent.SetRuntime(agent.RuntimeOptions{
|
m.agent.SetRuntime(agent.RuntimeOptions{
|
||||||
Provider: provider,
|
Provider: provider,
|
||||||
ProviderName: m.providerName,
|
ProviderName: m.providerName,
|
||||||
Model: m.provider.Model,
|
Model: m.provider.Model,
|
||||||
Thinking: m.provider.Thinking,
|
Thinking: m.provider.Thinking,
|
||||||
ThinkingParam: m.provider.ThinkingParam,
|
ThinkingParam: m.provider.ThinkingParam,
|
||||||
ThinkingEnabled: m.provider.ThinkingConfigured,
|
ThinkingEnabled: m.provider.ThinkingConfigured,
|
||||||
|
MaxContextTokens: m.maxContextTokens(),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -140,3 +296,10 @@ func contains(values []string, value string) bool {
|
|||||||
}
|
}
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func fitName(name string, width int) string {
|
||||||
|
if len(name) <= width {
|
||||||
|
return name
|
||||||
|
}
|
||||||
|
return name[:width-3] + "..."
|
||||||
|
}
|
||||||
|
|||||||
@@ -2,12 +2,13 @@ package session
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"path/filepath"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"agentu/internal/agent"
|
"agentu/internal/agent"
|
||||||
"agentu/internal/config"
|
"agentu/pkg/config"
|
||||||
"agentu/internal/llm"
|
"agentu/pkg/llm"
|
||||||
)
|
)
|
||||||
|
|
||||||
type recordingProvider struct {
|
type recordingProvider struct {
|
||||||
@@ -19,8 +20,8 @@ func (p *recordingProvider) ChatStream(ctx context.Context, req llm.ChatRequest,
|
|||||||
return emit(llm.StreamEvent{Content: "ok"})
|
return emit(llm.StreamEvent{Content: "ok"})
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestManagerSwitchesModelAndThinking(t *testing.T) {
|
func testConfig() *config.Config {
|
||||||
cfg := &config.Config{
|
return &config.Config{
|
||||||
Providers: map[string]config.ProviderConfig{
|
Providers: map[string]config.ProviderConfig{
|
||||||
"one": {
|
"one": {
|
||||||
Name: "one",
|
Name: "one",
|
||||||
@@ -34,6 +35,11 @@ func TestManagerSwitchesModelAndThinking(t *testing.T) {
|
|||||||
},
|
},
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func testManager(t *testing.T) (*Manager, *agent.Agent, *recordingProvider) {
|
||||||
|
t.Helper()
|
||||||
|
cfg := testConfig()
|
||||||
provider := &recordingProvider{}
|
provider := &recordingProvider{}
|
||||||
assistant := agent.New(agent.Options{
|
assistant := agent.New(agent.Options{
|
||||||
Provider: provider,
|
Provider: provider,
|
||||||
@@ -43,10 +49,19 @@ func TestManagerSwitchesModelAndThinking(t *testing.T) {
|
|||||||
ThinkingParam: "thinking",
|
ThinkingParam: "thinking",
|
||||||
ThinkingEnabled: true,
|
ThinkingEnabled: true,
|
||||||
})
|
})
|
||||||
manager, err := NewManager(cfg, assistant, nil)
|
store, err := NewStore(filepath.Join(t.TempDir(), "sessions"))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
manager, err := NewManager(cfg, assistant, nil, store)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
return manager, assistant, provider
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManagerSwitchesModelAndThinking(t *testing.T) {
|
||||||
|
manager, assistant, provider := testManager(t)
|
||||||
|
|
||||||
if _, err := manager.SetModel("model-b"); err != nil {
|
if _, err := manager.SetModel("model-b"); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
@@ -80,7 +95,8 @@ func TestManagerRejectsUnknownModelAndThinking(t *testing.T) {
|
|||||||
},
|
},
|
||||||
}
|
}
|
||||||
assistant := agent.New(agent.Options{Provider: &recordingProvider{}, Model: "model-a"})
|
assistant := agent.New(agent.Options{Provider: &recordingProvider{}, Model: "model-a"})
|
||||||
manager, err := NewManager(cfg, assistant, nil)
|
store, _ := NewStore(filepath.Join(t.TempDir(), "sessions"))
|
||||||
|
manager, err := NewManager(cfg, assistant, nil, store)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
@@ -92,3 +108,226 @@ func TestManagerRejectsUnknownModelAndThinking(t *testing.T) {
|
|||||||
t.Fatal("expected invalid thinking error")
|
t.Fatal("expected invalid thinking error")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestManagerSaveAndResume(t *testing.T) {
|
||||||
|
manager, assistant, _ := testManager(t)
|
||||||
|
|
||||||
|
// Start new session and have a conversation
|
||||||
|
manager.newSession()
|
||||||
|
var out strings.Builder
|
||||||
|
assistant.RunTurn(context.Background(), "hello world", &out, nil)
|
||||||
|
assistant.SetLastUsage(&llm.Usage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15})
|
||||||
|
|
||||||
|
// Save
|
||||||
|
if err := manager.Save(); err != nil {
|
||||||
|
t.Fatalf("save: %v", err)
|
||||||
|
}
|
||||||
|
savedID := manager.CurrentSessionID()
|
||||||
|
if savedID == "" {
|
||||||
|
t.Fatal("expected session ID after save")
|
||||||
|
}
|
||||||
|
// Name should be auto-derived from first user message
|
||||||
|
if manager.CurrentSessionName() != "hello world" {
|
||||||
|
t.Fatalf("session name = %q, want %q", manager.CurrentSessionName(), "hello world")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start a new session
|
||||||
|
manager.newSession()
|
||||||
|
if manager.CurrentSessionID() == savedID {
|
||||||
|
t.Fatal("new session should have different ID")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Resume
|
||||||
|
if err := manager.Resume(savedID); err != nil {
|
||||||
|
t.Fatalf("resume: %v", err)
|
||||||
|
}
|
||||||
|
if manager.CurrentSessionID() != savedID {
|
||||||
|
t.Fatalf("after resume, ID = %q, want %q", manager.CurrentSessionID(), savedID)
|
||||||
|
}
|
||||||
|
msgs := assistant.Messages()
|
||||||
|
hasUser := false
|
||||||
|
for _, m := range msgs {
|
||||||
|
if m.Role == llm.RoleUser && m.Content == "hello world" {
|
||||||
|
hasUser = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !hasUser {
|
||||||
|
t.Fatal("resumed session should contain original user message")
|
||||||
|
}
|
||||||
|
if usage := assistant.LastUsage(); usage == nil || usage.TotalTokens != 15 {
|
||||||
|
t.Fatalf("resumed usage = %#v", usage)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManagerRename(t *testing.T) {
|
||||||
|
manager, _, _ := testManager(t)
|
||||||
|
manager.newSession()
|
||||||
|
|
||||||
|
msg, err := manager.Rename("my cool session")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("rename: %v", err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(msg, "my cool session") {
|
||||||
|
t.Fatalf("rename message = %q", msg)
|
||||||
|
}
|
||||||
|
if manager.CurrentSessionName() != "my cool session" {
|
||||||
|
t.Fatalf("name = %q", manager.CurrentSessionName())
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify persisted
|
||||||
|
_, err = manager.store.Load(manager.CurrentSessionID())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("load after rename: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManagerRenameEmpty(t *testing.T) {
|
||||||
|
manager, _, _ := testManager(t)
|
||||||
|
manager.newSession()
|
||||||
|
_, err := manager.Rename(" ")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for empty name")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManagerSessions(t *testing.T) {
|
||||||
|
manager, _, _ := testManager(t)
|
||||||
|
|
||||||
|
// No sessions yet
|
||||||
|
out, err := manager.Sessions()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "No saved sessions") {
|
||||||
|
t.Fatalf("expected 'no saved sessions', got: %s", out)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create and save a session
|
||||||
|
manager.newSession()
|
||||||
|
manager.Save()
|
||||||
|
|
||||||
|
out, err = manager.Sessions()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if strings.Contains(out, "No saved sessions") {
|
||||||
|
t.Fatal("expected session list")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManagerNewSession(t *testing.T) {
|
||||||
|
manager, assistant, _ := testManager(t)
|
||||||
|
manager.newSession()
|
||||||
|
var out strings.Builder
|
||||||
|
assistant.RunTurn(context.Background(), "first message", &out, nil)
|
||||||
|
|
||||||
|
msg, err := manager.NewSession()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("new session: %v", err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(msg, "Starting new session") {
|
||||||
|
t.Fatalf("message = %q", msg)
|
||||||
|
}
|
||||||
|
// Old session should be saved
|
||||||
|
metas, _ := manager.store.List()
|
||||||
|
if len(metas) != 1 {
|
||||||
|
t.Fatalf("expected 1 saved session, got %d", len(metas))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManagerResumeLatest(t *testing.T) {
|
||||||
|
manager, assistant, _ := testManager(t)
|
||||||
|
|
||||||
|
// Create two sessions
|
||||||
|
manager.newSession()
|
||||||
|
var out strings.Builder
|
||||||
|
assistant.RunTurn(context.Background(), "session one", &out, nil)
|
||||||
|
manager.Save()
|
||||||
|
id1 := manager.CurrentSessionID()
|
||||||
|
|
||||||
|
manager.NewSession()
|
||||||
|
assistant.RunTurn(context.Background(), "session two", &out, nil)
|
||||||
|
manager.Save()
|
||||||
|
id2 := manager.CurrentSessionID()
|
||||||
|
|
||||||
|
// Clear and resume latest
|
||||||
|
manager.newSession()
|
||||||
|
if err := manager.ResumeLatest(); err != nil {
|
||||||
|
t.Fatalf("resume latest: %v", err)
|
||||||
|
}
|
||||||
|
// Should resume the most recently updated (id2)
|
||||||
|
if manager.CurrentSessionID() != id2 {
|
||||||
|
t.Fatalf("expected to resume %s, got %s", id2, manager.CurrentSessionID())
|
||||||
|
}
|
||||||
|
|
||||||
|
// Also verify id1 is resumable
|
||||||
|
if err := manager.Resume(id1); err != nil {
|
||||||
|
t.Fatalf("resume id1: %v", err)
|
||||||
|
}
|
||||||
|
if manager.CurrentSessionID() != id1 {
|
||||||
|
t.Fatalf("expected %s, got %s", id1, manager.CurrentSessionID())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type compactingProvider struct {
|
||||||
|
summaryCalls int
|
||||||
|
answerCalls int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *compactingProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
if len(req.Messages) == 2 && req.Messages[0].Role == llm.RoleSystem {
|
||||||
|
p.summaryCalls++
|
||||||
|
return emit(llm.StreamEvent{Content: "summary"})
|
||||||
|
}
|
||||||
|
p.answerCalls++
|
||||||
|
return emit(llm.StreamEvent{Content: "answer"})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManagerUsesModelContextForAgentCompaction(t *testing.T) {
|
||||||
|
provider := &compactingProvider{}
|
||||||
|
cfg := &config.Config{Providers: map[string]config.ProviderConfig{
|
||||||
|
"one": {
|
||||||
|
Name: "one",
|
||||||
|
BaseURL: "https://one.example.com",
|
||||||
|
APIKey: "sk-test",
|
||||||
|
Model: "model-a",
|
||||||
|
Models: []string{"model-a"},
|
||||||
|
ModelConfigs: map[string]config.ModelConfig{"model-a": {ID: "model-a", Thinking: "none", ContextTokens: 1}},
|
||||||
|
ContextTokens: 1,
|
||||||
|
},
|
||||||
|
}}
|
||||||
|
assistant := agent.New(agent.Options{Provider: provider, ProviderName: "one", Model: "model-a"})
|
||||||
|
store, err := NewStore(filepath.Join(t.TempDir(), "sessions"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
manager, err := NewManager(cfg, assistant, nil, store)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := manager.SetModel("model-a"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
assistant.SetMessages([]llm.Message{
|
||||||
|
{Role: llm.RoleSystem, Content: "system"},
|
||||||
|
{Role: llm.RoleUser, Content: "old question 1"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "old answer 1"},
|
||||||
|
{Role: llm.RoleUser, Content: "old question 2"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "old answer 2"},
|
||||||
|
{Role: llm.RoleUser, Content: "recent question 1"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "recent answer 1"},
|
||||||
|
{Role: llm.RoleUser, Content: "recent question 2"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "recent answer 2"},
|
||||||
|
})
|
||||||
|
|
||||||
|
var out strings.Builder
|
||||||
|
if err := assistant.RunTurn(context.Background(), "new question", &out, nil); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if provider.summaryCalls != 1 {
|
||||||
|
t.Fatalf("summaryCalls = %d, want 1", provider.summaryCalls)
|
||||||
|
}
|
||||||
|
if provider.answerCalls != 1 {
|
||||||
|
t.Fatalf("answerCalls = %d, want 1", provider.answerCalls)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,174 @@
|
|||||||
|
package session
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/rand"
|
||||||
|
"encoding/hex"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Session represents a persisted conversation.
|
||||||
|
type Session struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
|
Provider string `json:"provider"`
|
||||||
|
Model string `json:"model"`
|
||||||
|
Messages []llm.Message `json:"messages"`
|
||||||
|
LastUsage *llm.Usage `json:"last_usage,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// SessionMeta is a lightweight summary for listing.
|
||||||
|
type SessionMeta struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
|
Provider string `json:"provider"`
|
||||||
|
Model string `json:"model"`
|
||||||
|
MessageCount int `json:"message_count"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Store manages session files under a directory.
|
||||||
|
type Store struct {
|
||||||
|
dir string
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewStore creates a store, ensuring the directory exists.
|
||||||
|
func NewStore(dir string) (*Store, error) {
|
||||||
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||||
|
return nil, fmt.Errorf("create session dir: %w", err)
|
||||||
|
}
|
||||||
|
return &Store{dir: dir}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Save writes a session to disk.
|
||||||
|
func (s *Store) Save(sess *Session) error {
|
||||||
|
if sess.ID == "" {
|
||||||
|
return fmt.Errorf("session ID is required")
|
||||||
|
}
|
||||||
|
sess.UpdatedAt = time.Now().UTC()
|
||||||
|
if sess.CreatedAt.IsZero() {
|
||||||
|
sess.CreatedAt = sess.UpdatedAt
|
||||||
|
}
|
||||||
|
data, err := json.MarshalIndent(sess, "", " ")
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("marshal session: %w", err)
|
||||||
|
}
|
||||||
|
path := s.filePath(sess.ID)
|
||||||
|
if err := os.WriteFile(path, data, 0o644); err != nil {
|
||||||
|
return fmt.Errorf("write session: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Load reads a session by ID.
|
||||||
|
func (s *Store) Load(id string) (*Session, error) {
|
||||||
|
path := s.filePath(id)
|
||||||
|
data, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
if os.IsNotExist(err) {
|
||||||
|
return nil, fmt.Errorf("session %q not found", id)
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("read session: %w", err)
|
||||||
|
}
|
||||||
|
var sess Session
|
||||||
|
if err := json.Unmarshal(data, &sess); err != nil {
|
||||||
|
return nil, fmt.Errorf("parse session: %w", err)
|
||||||
|
}
|
||||||
|
return &sess, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// List returns metadata for all sessions, sorted by updated_at descending.
|
||||||
|
func (s *Store) List() ([]SessionMeta, error) {
|
||||||
|
entries, err := os.ReadDir(s.dir)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("read session dir: %w", err)
|
||||||
|
}
|
||||||
|
var metas []SessionMeta
|
||||||
|
for _, entry := range entries {
|
||||||
|
if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
path := filepath.Join(s.dir, entry.Name())
|
||||||
|
data, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
var sess Session
|
||||||
|
if err := json.Unmarshal(data, &sess); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
metas = append(metas, SessionMeta{
|
||||||
|
ID: sess.ID,
|
||||||
|
Name: sess.Name,
|
||||||
|
CreatedAt: sess.CreatedAt,
|
||||||
|
UpdatedAt: sess.UpdatedAt,
|
||||||
|
Provider: sess.Provider,
|
||||||
|
Model: sess.Model,
|
||||||
|
MessageCount: len(sess.Messages),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
sort.Slice(metas, func(i, j int) bool {
|
||||||
|
return metas[i].UpdatedAt.After(metas[j].UpdatedAt)
|
||||||
|
})
|
||||||
|
return metas, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Delete removes a session file.
|
||||||
|
func (s *Store) Delete(id string) error {
|
||||||
|
path := s.filePath(id)
|
||||||
|
if err := os.Remove(path); err != nil {
|
||||||
|
if os.IsNotExist(err) {
|
||||||
|
return fmt.Errorf("session %q not found", id)
|
||||||
|
}
|
||||||
|
return fmt.Errorf("delete session: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Store) filePath(id string) string {
|
||||||
|
return filepath.Join(s.dir, id+".json")
|
||||||
|
}
|
||||||
|
|
||||||
|
// GenerateID creates a random 8-character hex session ID.
|
||||||
|
func GenerateID() string {
|
||||||
|
b := make([]byte, 4)
|
||||||
|
_, _ = rand.Read(b)
|
||||||
|
return hex.EncodeToString(b)
|
||||||
|
}
|
||||||
|
|
||||||
|
// DefaultName derives a session name from the first user message.
|
||||||
|
func DefaultName(messages []llm.Message) string {
|
||||||
|
for _, msg := range messages {
|
||||||
|
if msg.Role != llm.RoleUser {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
text := strings.TrimSpace(msg.Content)
|
||||||
|
if text == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
// Take first line, truncate at 40 chars at word boundary
|
||||||
|
if idx := strings.IndexByte(text, '\n'); idx >= 0 {
|
||||||
|
text = text[:idx]
|
||||||
|
}
|
||||||
|
text = strings.TrimSpace(text)
|
||||||
|
if len(text) > 40 {
|
||||||
|
text = text[:40]
|
||||||
|
if idx := strings.LastIndexByte(text, ' '); idx > 20 {
|
||||||
|
text = text[:idx]
|
||||||
|
}
|
||||||
|
text += "..."
|
||||||
|
}
|
||||||
|
return text
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
@@ -0,0 +1,174 @@
|
|||||||
|
package session
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
func tempStore(t *testing.T) *Store {
|
||||||
|
t.Helper()
|
||||||
|
dir := t.TempDir()
|
||||||
|
store, err := NewStore(dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
return store
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSaveAndLoad(t *testing.T) {
|
||||||
|
store := tempStore(t)
|
||||||
|
sess := &Session{
|
||||||
|
ID: GenerateID(),
|
||||||
|
Name: "test session",
|
||||||
|
Provider: "loveuer",
|
||||||
|
Model: "deepseek-v4-flash",
|
||||||
|
Messages: []llm.Message{
|
||||||
|
{Role: llm.RoleSystem, Content: "you are helpful"},
|
||||||
|
{Role: llm.RoleUser, Content: "hello"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "hi there"},
|
||||||
|
},
|
||||||
|
LastUsage: &llm.Usage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15},
|
||||||
|
}
|
||||||
|
if err := store.Save(sess); err != nil {
|
||||||
|
t.Fatalf("save: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
loaded, err := store.Load(sess.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("load: %v", err)
|
||||||
|
}
|
||||||
|
if loaded.ID != sess.ID {
|
||||||
|
t.Errorf("ID = %q, want %q", loaded.ID, sess.ID)
|
||||||
|
}
|
||||||
|
if loaded.Name != sess.Name {
|
||||||
|
t.Errorf("Name = %q, want %q", loaded.Name, sess.Name)
|
||||||
|
}
|
||||||
|
if len(loaded.Messages) != 3 {
|
||||||
|
t.Errorf("Messages len = %d, want 3", len(loaded.Messages))
|
||||||
|
}
|
||||||
|
if loaded.Messages[1].Content != "hello" {
|
||||||
|
t.Errorf("Messages[1].Content = %q, want %q", loaded.Messages[1].Content, "hello")
|
||||||
|
}
|
||||||
|
if loaded.LastUsage == nil || loaded.LastUsage.TotalTokens != 15 {
|
||||||
|
t.Errorf("LastUsage = %#v", loaded.LastUsage)
|
||||||
|
}
|
||||||
|
if loaded.UpdatedAt.IsZero() {
|
||||||
|
t.Error("UpdatedAt should be set")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoadNotFound(t *testing.T) {
|
||||||
|
store := tempStore(t)
|
||||||
|
_, err := store.Load("nonexistent")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for nonexistent session")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestList(t *testing.T) {
|
||||||
|
store := tempStore(t)
|
||||||
|
// Empty list
|
||||||
|
metas, err := store.List()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("list: %v", err)
|
||||||
|
}
|
||||||
|
if len(metas) != 0 {
|
||||||
|
t.Errorf("expected empty list, got %d", len(metas))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Save two sessions
|
||||||
|
store.Save(&Session{ID: "aaa", Name: "first", Messages: []llm.Message{{Role: "user", Content: "a"}}})
|
||||||
|
store.Save(&Session{ID: "bbb", Name: "second", Messages: []llm.Message{{Role: "user", Content: "b"}, {Role: "assistant", Content: "c"}}})
|
||||||
|
|
||||||
|
metas, err = store.List()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("list: %v", err)
|
||||||
|
}
|
||||||
|
if len(metas) != 2 {
|
||||||
|
t.Fatalf("expected 2 sessions, got %d", len(metas))
|
||||||
|
}
|
||||||
|
// Sorted by updated_at desc
|
||||||
|
if metas[0].ID != "bbb" && metas[0].ID != "aaa" {
|
||||||
|
t.Errorf("unexpected first session ID: %q", metas[0].ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDelete(t *testing.T) {
|
||||||
|
store := tempStore(t)
|
||||||
|
store.Save(&Session{ID: "todelete", Name: "temp"})
|
||||||
|
|
||||||
|
if err := store.Delete("todelete"); err != nil {
|
||||||
|
t.Fatalf("delete: %v", err)
|
||||||
|
}
|
||||||
|
_, err := store.Load("todelete")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error after delete")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDeleteNotFound(t *testing.T) {
|
||||||
|
store := tempStore(t)
|
||||||
|
err := store.Delete("nonexistent")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for nonexistent delete")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSaveRequiresID(t *testing.T) {
|
||||||
|
store := tempStore(t)
|
||||||
|
err := store.Save(&Session{Name: "no id"})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for empty ID")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGenerateID(t *testing.T) {
|
||||||
|
id := GenerateID()
|
||||||
|
if len(id) != 8 {
|
||||||
|
t.Errorf("ID length = %d, want 8", len(id))
|
||||||
|
}
|
||||||
|
id2 := GenerateID()
|
||||||
|
if id == id2 {
|
||||||
|
t.Error("expected unique IDs")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDefaultName(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
messages []llm.Message
|
||||||
|
want string
|
||||||
|
}{
|
||||||
|
{"empty", nil, ""},
|
||||||
|
{"no user msg", []llm.Message{{Role: "system", Content: "hi"}}, ""},
|
||||||
|
{"short", []llm.Message{{Role: "user", Content: "hello world"}}, "hello world"},
|
||||||
|
{"long", []llm.Message{{Role: "user", Content: "this is a really long message that should be truncated at some point in the middle"}}, "this is a really long message that..."},
|
||||||
|
{"multiline", []llm.Message{{Role: "user", Content: "first line\nsecond line"}}, "first line"},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
got := DefaultName(tt.messages)
|
||||||
|
if got != tt.want {
|
||||||
|
t.Errorf("DefaultName = %q, want %q", got, tt.want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewStoreCreatesDir(t *testing.T) {
|
||||||
|
dir := filepath.Join(t.TempDir(), "nested", "sessions")
|
||||||
|
store, err := NewStore(dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewStore: %v", err)
|
||||||
|
}
|
||||||
|
info, err := os.Stat(store.dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("stat: %v", err)
|
||||||
|
}
|
||||||
|
if !info.IsDir() {
|
||||||
|
t.Error("expected directory")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,219 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"go/ast"
|
||||||
|
"go/parser"
|
||||||
|
"go/printer"
|
||||||
|
"go/token"
|
||||||
|
"io/fs"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
const defaultCodeSymbolFileLimit = 50
|
||||||
|
|
||||||
|
type CodeSymbolsTool struct {
|
||||||
|
workingDir string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewCodeSymbolsTool(workingDir string) *CodeSymbolsTool {
|
||||||
|
return &CodeSymbolsTool{workingDir: workingDir}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *CodeSymbolsTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "code_symbols",
|
||||||
|
Description: "List Go code symbols (packages, imports, types, functions, methods) for a file or directory without reading full file contents.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "Go file or directory to inspect. Defaults to the configured agent working directory."},
|
||||||
|
"max_files": {"type": "integer", "description": "Maximum number of Go files to inspect when path is a directory. Defaults to 50."}
|
||||||
|
},
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type codeSymbolsArgs struct {
|
||||||
|
Path string `json:"path"`
|
||||||
|
MaxFiles int `json:"max_files"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *CodeSymbolsTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args codeSymbolsArgs
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
path, err := resolvePath(t.workingDir, args.Path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
limit := args.MaxFiles
|
||||||
|
if limit <= 0 {
|
||||||
|
limit = defaultCodeSymbolFileLimit
|
||||||
|
}
|
||||||
|
|
||||||
|
files, err := goFiles(ctx, path, limit)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if len(files) == 0 {
|
||||||
|
return "no Go files found", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
fset := token.NewFileSet()
|
||||||
|
var out strings.Builder
|
||||||
|
for i, file := range files {
|
||||||
|
if ctxErr := ctx.Err(); ctxErr != nil {
|
||||||
|
return FormatResult(out.String()), ctxErr
|
||||||
|
}
|
||||||
|
if i > 0 {
|
||||||
|
out.WriteString("\n")
|
||||||
|
}
|
||||||
|
if err := appendGoSymbols(&out, fset, file); err != nil {
|
||||||
|
fmt.Fprintf(&out, "%s\nerror: %v\n", file, err)
|
||||||
|
}
|
||||||
|
if out.Len() > MaxToolOutputBytes {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return FormatResult(out.String()), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func goFiles(ctx context.Context, path string, limit int) ([]string, error) {
|
||||||
|
info, err := os.Stat(path)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if !info.IsDir() {
|
||||||
|
if strings.HasSuffix(path, ".go") {
|
||||||
|
return []string{path}, nil
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("path is not a Go file: %s", path)
|
||||||
|
}
|
||||||
|
|
||||||
|
var files []string
|
||||||
|
stop := errors.New("file limit reached")
|
||||||
|
err = filepath.WalkDir(path, func(current string, d fs.DirEntry, err error) error {
|
||||||
|
if err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if ctxErr := ctx.Err(); ctxErr != nil {
|
||||||
|
return ctxErr
|
||||||
|
}
|
||||||
|
if d.IsDir() {
|
||||||
|
switch d.Name() {
|
||||||
|
case ".git", "vendor", "node_modules":
|
||||||
|
return filepath.SkipDir
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if strings.HasSuffix(current, ".go") {
|
||||||
|
files = append(files, current)
|
||||||
|
if len(files) >= limit {
|
||||||
|
return stop
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
if err != nil && !errors.Is(err, stop) {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
sort.Strings(files)
|
||||||
|
return files, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func appendGoSymbols(out *strings.Builder, fset *token.FileSet, path string) error {
|
||||||
|
file, err := parser.ParseFile(fset, path, nil, parser.ImportsOnly|parser.ParseComments)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
full, err := parser.ParseFile(fset, path, nil, 0)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
fmt.Fprintf(out, "%s\n", path)
|
||||||
|
fmt.Fprintf(out, " package %s\n", full.Name.Name)
|
||||||
|
if len(file.Imports) > 0 {
|
||||||
|
imports := make([]string, 0, len(file.Imports))
|
||||||
|
for _, imp := range file.Imports {
|
||||||
|
imports = append(imports, strings.Trim(imp.Path.Value, "\""))
|
||||||
|
}
|
||||||
|
fmt.Fprintf(out, " imports: %s\n", strings.Join(imports, ", "))
|
||||||
|
}
|
||||||
|
|
||||||
|
var types, funcs, methods []string
|
||||||
|
for _, decl := range full.Decls {
|
||||||
|
switch d := decl.(type) {
|
||||||
|
case *ast.GenDecl:
|
||||||
|
if d.Tok != token.TYPE {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
for _, spec := range d.Specs {
|
||||||
|
if ts, ok := spec.(*ast.TypeSpec); ok {
|
||||||
|
kind := typeKind(ts.Type)
|
||||||
|
types = append(types, fmt.Sprintf("%d: type %s %s", fset.Position(ts.Pos()).Line, ts.Name.Name, kind))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
case *ast.FuncDecl:
|
||||||
|
line := fset.Position(d.Pos()).Line
|
||||||
|
if d.Recv == nil {
|
||||||
|
funcs = append(funcs, fmt.Sprintf("%d: func %s", line, d.Name.Name))
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
methods = append(methods, fmt.Sprintf("%d: method %s.%s", line, receiverName(d.Recv), d.Name.Name))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
appendSymbolGroup(out, "types", types)
|
||||||
|
appendSymbolGroup(out, "funcs", funcs)
|
||||||
|
appendSymbolGroup(out, "methods", methods)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func appendSymbolGroup(out *strings.Builder, title string, values []string) {
|
||||||
|
if len(values) == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
fmt.Fprintf(out, " %s:\n", title)
|
||||||
|
for _, value := range values {
|
||||||
|
fmt.Fprintf(out, " - %s\n", value)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func typeKind(expr ast.Expr) string {
|
||||||
|
switch expr.(type) {
|
||||||
|
case *ast.StructType:
|
||||||
|
return "struct"
|
||||||
|
case *ast.InterfaceType:
|
||||||
|
return "interface"
|
||||||
|
case *ast.FuncType:
|
||||||
|
return "func"
|
||||||
|
default:
|
||||||
|
return "alias"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func receiverName(recv *ast.FieldList) string {
|
||||||
|
if recv == nil || len(recv.List) == 0 {
|
||||||
|
return "?"
|
||||||
|
}
|
||||||
|
var b bytes.Buffer
|
||||||
|
if err := printer.Fprint(&b, token.NewFileSet(), recv.List[0].Type); err != nil {
|
||||||
|
return "?"
|
||||||
|
}
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
+63
-11
@@ -13,7 +13,7 @@ import (
|
|||||||
"sort"
|
"sort"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"agentu/internal/llm"
|
"agentu/pkg/llm"
|
||||||
)
|
)
|
||||||
|
|
||||||
type FileReadTool struct {
|
type FileReadTool struct {
|
||||||
@@ -81,6 +81,7 @@ type FileWriteTool struct {
|
|||||||
func NewFileWriteTool(workingDir string) *FileWriteTool {
|
func NewFileWriteTool(workingDir string) *FileWriteTool {
|
||||||
return &FileWriteTool{workingDir: workingDir}
|
return &FileWriteTool{workingDir: workingDir}
|
||||||
}
|
}
|
||||||
|
func (*FileWriteTool) Mutates() bool { return true }
|
||||||
|
|
||||||
func (t *FileWriteTool) Definition() llm.Tool {
|
func (t *FileWriteTool) Definition() llm.Tool {
|
||||||
return llm.Tool{
|
return llm.Tool{
|
||||||
@@ -223,7 +224,9 @@ func (t *FileSearchTool) Definition() llm.Tool {
|
|||||||
"properties": {
|
"properties": {
|
||||||
"query": {"type": "string", "description": "Text or regular expression to search for."},
|
"query": {"type": "string", "description": "Text or regular expression to search for."},
|
||||||
"path": {"type": "string", "description": "Optional directory or file to search. Defaults to the configured agent working directory."},
|
"path": {"type": "string", "description": "Optional directory or file to search. Defaults to the configured agent working directory."},
|
||||||
"glob": {"type": "string", "description": "Optional file glob, for example '*.go'."}
|
"glob": {"type": "string", "description": "Optional file glob, for example '*.go'."},
|
||||||
|
"context_lines": {"type": "integer", "description": "Optional number of context lines around each match. Defaults to 0."},
|
||||||
|
"max_results": {"type": "integer", "description": "Optional maximum matches to return. Defaults to 100."}
|
||||||
},
|
},
|
||||||
"required": ["query"],
|
"required": ["query"],
|
||||||
"additionalProperties": false
|
"additionalProperties": false
|
||||||
@@ -233,9 +236,11 @@ func (t *FileSearchTool) Definition() llm.Tool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type fileSearchArgs struct {
|
type fileSearchArgs struct {
|
||||||
Query string `json:"query"`
|
Query string `json:"query"`
|
||||||
Path string `json:"path"`
|
Path string `json:"path"`
|
||||||
Glob string `json:"glob"`
|
Glob string `json:"glob"`
|
||||||
|
ContextLines int `json:"context_lines"`
|
||||||
|
MaxResults int `json:"max_results"`
|
||||||
}
|
}
|
||||||
|
|
||||||
func (t *FileSearchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
func (t *FileSearchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
@@ -256,8 +261,31 @@ func (t *FileSearchTool) Execute(ctx context.Context, raw json.RawMessage) (stri
|
|||||||
return t.walkSearch(ctx, args, path)
|
return t.walkSearch(ctx, args, path)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func normalizedSearchOptions(args fileSearchArgs) (contextLines int, maxResults int) {
|
||||||
|
contextLines = args.ContextLines
|
||||||
|
if contextLines < 0 {
|
||||||
|
contextLines = 0
|
||||||
|
}
|
||||||
|
if contextLines > 5 {
|
||||||
|
contextLines = 5
|
||||||
|
}
|
||||||
|
maxResults = args.MaxResults
|
||||||
|
if maxResults <= 0 {
|
||||||
|
maxResults = 100
|
||||||
|
}
|
||||||
|
if maxResults > 1000 {
|
||||||
|
maxResults = 1000
|
||||||
|
}
|
||||||
|
return contextLines, maxResults
|
||||||
|
}
|
||||||
|
|
||||||
func (t *FileSearchTool) ripgrep(ctx context.Context, args fileSearchArgs, path string) (string, error) {
|
func (t *FileSearchTool) ripgrep(ctx context.Context, args fileSearchArgs, path string) (string, error) {
|
||||||
rgArgs := []string{"--line-number", "--color", "never"}
|
rgArgs := []string{"--line-number", "--color", "never"}
|
||||||
|
contextLines, maxResults := normalizedSearchOptions(args)
|
||||||
|
if contextLines > 0 {
|
||||||
|
rgArgs = append(rgArgs, "--context", fmt.Sprint(contextLines))
|
||||||
|
}
|
||||||
|
rgArgs = append(rgArgs, "--max-count", fmt.Sprint(maxResults))
|
||||||
if args.Glob != "" {
|
if args.Glob != "" {
|
||||||
rgArgs = append(rgArgs, "--glob", args.Glob)
|
rgArgs = append(rgArgs, "--glob", args.Glob)
|
||||||
}
|
}
|
||||||
@@ -282,7 +310,10 @@ func (t *FileSearchTool) ripgrep(ctx context.Context, args fileSearchArgs, path
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (t *FileSearchTool) walkSearch(ctx context.Context, args fileSearchArgs, root string) (string, error) {
|
func (t *FileSearchTool) walkSearch(ctx context.Context, args fileSearchArgs, root string) (string, error) {
|
||||||
|
contextLines, maxResults := normalizedSearchOptions(args)
|
||||||
var out bytes.Buffer
|
var out bytes.Buffer
|
||||||
|
matchCount := 0
|
||||||
|
stop := errors.New("result limit reached")
|
||||||
err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
|
err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil
|
return nil
|
||||||
@@ -312,16 +343,22 @@ func (t *FileSearchTool) walkSearch(ctx context.Context, args fileSearchArgs, ro
|
|||||||
}
|
}
|
||||||
lines := strings.Split(string(data), "\n")
|
lines := strings.Split(string(data), "\n")
|
||||||
for i, line := range lines {
|
for i, line := range lines {
|
||||||
if strings.Contains(line, args.Query) {
|
if !strings.Contains(line, args.Query) {
|
||||||
fmt.Fprintf(&out, "%s:%d:%s\n", path, i+1, line)
|
continue
|
||||||
if out.Len() > MaxToolOutputBytes {
|
}
|
||||||
return errors.New("output limit reached")
|
matchCount++
|
||||||
}
|
appendSearchMatch(&out, path, lines, i, contextLines)
|
||||||
|
if matchCount >= maxResults {
|
||||||
|
fmt.Fprintf(&out, "[stopped after %d matches; narrow path/query or raise max_results]\n", maxResults)
|
||||||
|
return stop
|
||||||
|
}
|
||||||
|
if out.Len() > MaxToolOutputBytes {
|
||||||
|
return stop
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
})
|
})
|
||||||
if err != nil && err.Error() != "output limit reached" {
|
if err != nil && !errors.Is(err, stop) {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
if out.Len() == 0 {
|
if out.Len() == 0 {
|
||||||
@@ -329,3 +366,18 @@ func (t *FileSearchTool) walkSearch(ctx context.Context, args fileSearchArgs, ro
|
|||||||
}
|
}
|
||||||
return FormatResult(out.String()), nil
|
return FormatResult(out.String()), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func appendSearchMatch(out *bytes.Buffer, path string, lines []string, matchIndex int, contextLines int) {
|
||||||
|
start := max(0, matchIndex-contextLines)
|
||||||
|
end := min(len(lines)-1, matchIndex+contextLines)
|
||||||
|
if contextLines > 0 {
|
||||||
|
fmt.Fprintf(out, "-- %s:%d --\n", path, matchIndex+1)
|
||||||
|
}
|
||||||
|
for i := start; i <= end; i++ {
|
||||||
|
marker := ":"
|
||||||
|
if i == matchIndex {
|
||||||
|
marker = "*"
|
||||||
|
}
|
||||||
|
fmt.Fprintf(out, "%s%s%d:%s\n", path, marker, i+1, lines[i])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,178 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
type FileEditTool struct {
|
||||||
|
workingDir string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewFileEditTool(workingDir string) *FileEditTool {
|
||||||
|
return &FileEditTool{workingDir: workingDir}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (*FileEditTool) Mutates() bool { return true }
|
||||||
|
|
||||||
|
func (t *FileEditTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "file_edit",
|
||||||
|
Description: "Edit an existing UTF-8 text file by replacing an inclusive 1-based line range or an exact old text pattern.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "Existing file path relative to the configured agent working directory."},
|
||||||
|
"start_line": {"type": "integer", "description": "1-based first line for a line-range edit."},
|
||||||
|
"end_line": {"type": "integer", "description": "Optional 1-based final line; defaults to start_line when omitted."},
|
||||||
|
"old": {"type": "string", "description": "Exact text to replace for a pattern edit; not regex."},
|
||||||
|
"new": {"type": "string", "description": "Replacement text; may be empty string for deletion."},
|
||||||
|
"replace_all": {"type": "boolean", "description": "When true, replace every exact old occurrence; otherwise require exactly one occurrence."}
|
||||||
|
},
|
||||||
|
"required": ["path", "new"],
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type fileEditArgs struct {
|
||||||
|
Path string `json:"path"`
|
||||||
|
StartLine int `json:"start_line"`
|
||||||
|
EndLine int `json:"end_line"`
|
||||||
|
Old string `json:"old"`
|
||||||
|
New *string `json:"new"`
|
||||||
|
ReplaceAll bool `json:"replace_all"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *FileEditTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args fileEditArgs
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(args.Path) == "" {
|
||||||
|
return "", errors.New("path is required")
|
||||||
|
}
|
||||||
|
if args.New == nil {
|
||||||
|
return "", errors.New("new is required")
|
||||||
|
}
|
||||||
|
path, err := resolvePath(t.workingDir, args.Path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
info, err := os.Stat(path)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("cannot access %s: %w", args.Path, err)
|
||||||
|
}
|
||||||
|
if info.IsDir() {
|
||||||
|
return "", fmt.Errorf("%s is a directory, not a file", args.Path)
|
||||||
|
}
|
||||||
|
perm := info.Mode().Perm()
|
||||||
|
|
||||||
|
hasLine := args.StartLine > 0 || args.EndLine > 0
|
||||||
|
hasOld := strings.TrimSpace(args.Old) != ""
|
||||||
|
|
||||||
|
if hasLine && hasOld {
|
||||||
|
return "", errors.New("provide either start_line/end_line or old, not both")
|
||||||
|
}
|
||||||
|
if !hasLine && !hasOld {
|
||||||
|
return "", errors.New("provide either start_line/end_line or old")
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
data, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
content := string(data)
|
||||||
|
|
||||||
|
var result string
|
||||||
|
if hasOld {
|
||||||
|
result, err = editByPattern(content, args.Path, args.Old, *args.New, args.ReplaceAll)
|
||||||
|
} else {
|
||||||
|
result, err = editByLineRange(content, args.Path, args.StartLine, args.EndLine, *args.New)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(path, []byte(result), perm); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
switch {
|
||||||
|
case hasOld:
|
||||||
|
count := strings.Count(content, args.Old)
|
||||||
|
if args.ReplaceAll {
|
||||||
|
return fmt.Sprintf("edited %s: replaced %d occurrences", args.Path, count), nil
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("edited %s: replaced 1 occurrence", args.Path), nil
|
||||||
|
default:
|
||||||
|
endLine := args.EndLine
|
||||||
|
if endLine == 0 {
|
||||||
|
endLine = args.StartLine
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("edited %s: replaced lines %d-%d", args.Path, args.StartLine, endLine), nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func editByLineRange(content, displayPath string, startLine, endLine int, replacement string) (string, error) {
|
||||||
|
if startLine < 1 {
|
||||||
|
return "", errors.New("start_line must be >= 1")
|
||||||
|
}
|
||||||
|
if endLine == 0 {
|
||||||
|
endLine = startLine
|
||||||
|
}
|
||||||
|
if endLine < startLine {
|
||||||
|
return "", fmt.Errorf("end_line (%d) must be >= start_line (%d)", endLine, startLine)
|
||||||
|
}
|
||||||
|
|
||||||
|
lines := strings.Split(content, "\n")
|
||||||
|
if startLine > len(lines) {
|
||||||
|
return "", fmt.Errorf("line range %d-%d is outside file with %d lines", startLine, endLine, len(lines))
|
||||||
|
}
|
||||||
|
if endLine > len(lines) {
|
||||||
|
return "", fmt.Errorf("line range %d-%d is outside file with %d lines", startLine, endLine, len(lines))
|
||||||
|
}
|
||||||
|
|
||||||
|
var b strings.Builder
|
||||||
|
b.WriteString(strings.Join(lines[:startLine-1], "\n"))
|
||||||
|
if startLine > 1 {
|
||||||
|
b.WriteString("\n")
|
||||||
|
}
|
||||||
|
b.WriteString(replacement)
|
||||||
|
if endLine < len(lines) {
|
||||||
|
if !strings.HasSuffix(replacement, "\n") && replacement != "" {
|
||||||
|
b.WriteString("\n")
|
||||||
|
}
|
||||||
|
b.WriteString(strings.Join(lines[endLine:], "\n"))
|
||||||
|
}
|
||||||
|
|
||||||
|
return b.String(), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func editByPattern(content, displayPath, old, replacement string, replaceAll bool) (string, error) {
|
||||||
|
count := strings.Count(content, old)
|
||||||
|
if count == 0 {
|
||||||
|
return "", errors.New("old text not found")
|
||||||
|
}
|
||||||
|
if !replaceAll && count != 1 {
|
||||||
|
return "", fmt.Errorf("old text matched %d occurrences; set replace_all=true or provide a more specific old value", count)
|
||||||
|
}
|
||||||
|
return strings.ReplaceAll(content, old, replacement), nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,136 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestFileEditReplacesLineRange(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "notes.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("one\ntwo\nthree\n"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
out, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","start_line":2,"end_line":2,"new":"TWO\n"}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "replaced lines 2-2") {
|
||||||
|
t.Fatalf("output = %q", out)
|
||||||
|
}
|
||||||
|
got, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if string(got) != "one\nTWO\nthree\n" {
|
||||||
|
t.Fatalf("file content = %q", string(got))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileEditReplacesExactPattern(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "notes.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("alpha beta alpha"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
out, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"beta","new":"BETA"}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "replaced 1 occurrence") {
|
||||||
|
t.Fatalf("output = %q", out)
|
||||||
|
}
|
||||||
|
got, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if string(got) != "alpha BETA alpha" {
|
||||||
|
t.Fatalf("file content = %q", string(got))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileEditRejectsAmbiguousPattern(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "notes.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("x x"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"x","new":"y"}`))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for ambiguous pattern")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "matched 2 occurrences") {
|
||||||
|
t.Fatalf("error = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileEditReplaceAll(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "notes.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("x x"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
out, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"x","new":"y","replace_all":true}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "replaced 2 occurrences") {
|
||||||
|
t.Fatalf("output = %q", out)
|
||||||
|
}
|
||||||
|
got, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if string(got) != "y y" {
|
||||||
|
t.Fatalf("file content = %q", string(got))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileEditRejectsEscapingPath(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"../outside.txt","old":"x","new":"y"}`))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for escaping path")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "escapes") {
|
||||||
|
t.Fatalf("error = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileEditRejectsMissingNew(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "notes.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"hello"}`))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for missing new")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "new is required") {
|
||||||
|
t.Fatalf("error = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileEditRejectsNonExistentFile(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"missing.txt","old":"x","new":"y"}`))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for non-existent file")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,246 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"os/exec"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
defaultGitLogMaxCount = 20
|
||||||
|
maxGitLogMaxCount = 100
|
||||||
|
)
|
||||||
|
|
||||||
|
func runGit(ctx context.Context, workingDir string, args ...string) (string, error) {
|
||||||
|
gitArgs := append([]string{"-C", workingDir}, args...)
|
||||||
|
cmd := exec.CommandContext(ctx, "git", gitArgs...)
|
||||||
|
output, err := cmd.CombinedOutput()
|
||||||
|
text := FormatResult(string(output))
|
||||||
|
if err != nil {
|
||||||
|
if strings.Contains(err.Error(), "context deadline exceeded") || strings.Contains(err.Error(), "signal: killed") {
|
||||||
|
if text != "" {
|
||||||
|
return text, fmt.Errorf("git command timed out: %w", err)
|
||||||
|
}
|
||||||
|
return "", fmt.Errorf("git command timed out: %w", err)
|
||||||
|
}
|
||||||
|
if text != "" {
|
||||||
|
return text, fmt.Errorf("git command failed: %w", err)
|
||||||
|
}
|
||||||
|
return "", fmt.Errorf("git command failed: %w", err)
|
||||||
|
}
|
||||||
|
return text, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func gitPathspec(baseDir, path string) (string, error) {
|
||||||
|
if strings.TrimSpace(path) == "" {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
resolved, err := resolvePath(baseDir, path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
rel, err := filepath.Rel(baseDir, resolved)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return filepath.ToSlash(filepath.Clean(rel)), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeGitLogMaxCount(value int) int {
|
||||||
|
if value <= 0 {
|
||||||
|
return defaultGitLogMaxCount
|
||||||
|
}
|
||||||
|
if value > maxGitLogMaxCount {
|
||||||
|
return maxGitLogMaxCount
|
||||||
|
}
|
||||||
|
return value
|
||||||
|
}
|
||||||
|
|
||||||
|
// git_status
|
||||||
|
|
||||||
|
type GitStatusTool struct {
|
||||||
|
workingDir string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewGitStatusTool(workingDir string) *GitStatusTool {
|
||||||
|
return &GitStatusTool{workingDir: workingDir}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *GitStatusTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "git_status",
|
||||||
|
Description: "Show changed files using git status --short.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "Optional path to limit status to. Relative paths are resolved from the configured agent working directory."}
|
||||||
|
},
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type gitStatusArgs struct {
|
||||||
|
Path string `json:"path"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *GitStatusTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args gitStatusArgs
|
||||||
|
if len(raw) > 0 {
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
gitArgs := []string{"status", "--short"}
|
||||||
|
if args.Path != "" {
|
||||||
|
pathspec, err := gitPathspec(t.workingDir, args.Path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
gitArgs = append(gitArgs, "--", pathspec)
|
||||||
|
}
|
||||||
|
text, err := runGit(ctx, t.workingDir, gitArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return text, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(text) == "" {
|
||||||
|
return "Working tree clean.", nil
|
||||||
|
}
|
||||||
|
return "Changed files:\n" + text, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// git_diff
|
||||||
|
|
||||||
|
type GitDiffTool struct {
|
||||||
|
workingDir string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewGitDiffTool(workingDir string) *GitDiffTool {
|
||||||
|
return &GitDiffTool{workingDir: workingDir}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *GitDiffTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "git_diff",
|
||||||
|
Description: "Show unstaged or staged git diff output, optionally as --stat.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "Optional path to limit diff to. Relative paths are resolved from the configured agent working directory."},
|
||||||
|
"staged": {"type": "boolean", "description": "Show staged diff instead of unstaged."},
|
||||||
|
"stat": {"type": "boolean", "description": "Show diffstat summary instead of full diff."}
|
||||||
|
},
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type gitDiffArgs struct {
|
||||||
|
Path string `json:"path"`
|
||||||
|
Staged bool `json:"staged"`
|
||||||
|
Stat bool `json:"stat"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *GitDiffTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args gitDiffArgs
|
||||||
|
if len(raw) > 0 {
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
gitArgs := []string{"diff"}
|
||||||
|
if args.Staged {
|
||||||
|
gitArgs = append(gitArgs, "--cached")
|
||||||
|
}
|
||||||
|
if args.Stat {
|
||||||
|
gitArgs = append(gitArgs, "--stat")
|
||||||
|
}
|
||||||
|
if args.Path != "" {
|
||||||
|
pathspec, err := gitPathspec(t.workingDir, args.Path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
gitArgs = append(gitArgs, "--", pathspec)
|
||||||
|
}
|
||||||
|
text, err := runGit(ctx, t.workingDir, gitArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return text, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(text) == "" {
|
||||||
|
if args.Staged {
|
||||||
|
return "No staged diff.", nil
|
||||||
|
}
|
||||||
|
return "No unstaged diff.", nil
|
||||||
|
}
|
||||||
|
return text, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// git_log
|
||||||
|
|
||||||
|
type GitLogTool struct {
|
||||||
|
workingDir string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewGitLogTool(workingDir string) *GitLogTool {
|
||||||
|
return &GitLogTool{workingDir: workingDir}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *GitLogTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "git_log",
|
||||||
|
Description: "Show recent git commits with git log --oneline --decorate.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "Optional path to limit log to. Relative paths are resolved from the configured agent working directory."},
|
||||||
|
"max_count": {"type": "integer", "description": "Maximum number of commits to show. Defaults to 20, capped at 100."}
|
||||||
|
},
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type gitLogArgs struct {
|
||||||
|
Path string `json:"path"`
|
||||||
|
MaxCount int `json:"max_count"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *GitLogTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args gitLogArgs
|
||||||
|
if len(raw) > 0 {
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
maxCount := normalizeGitLogMaxCount(args.MaxCount)
|
||||||
|
gitArgs := []string{"log", "--oneline", "--decorate", fmt.Sprintf("-n%d", maxCount)}
|
||||||
|
if args.Path != "" {
|
||||||
|
pathspec, err := gitPathspec(t.workingDir, args.Path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
gitArgs = append(gitArgs, "--", pathspec)
|
||||||
|
}
|
||||||
|
text, err := runGit(ctx, t.workingDir, gitArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return text, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(text) == "" {
|
||||||
|
return "No git history.", nil
|
||||||
|
}
|
||||||
|
return text, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,142 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"os"
|
||||||
|
"os/exec"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func requireGit(t *testing.T) {
|
||||||
|
t.Helper()
|
||||||
|
if _, err := exec.LookPath("git"); err != nil {
|
||||||
|
t.Skip("git not installed")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func runGitForTest(t *testing.T, dir string, args ...string) {
|
||||||
|
t.Helper()
|
||||||
|
gitArgs := append([]string{"-C", dir}, args...)
|
||||||
|
cmd := exec.Command("git", gitArgs...)
|
||||||
|
out, err := cmd.CombinedOutput()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("git %v failed: %v\n%s", args, err, out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func initTestRepo(t *testing.T) string {
|
||||||
|
t.Helper()
|
||||||
|
dir := t.TempDir()
|
||||||
|
runGitForTest(t, dir, "init")
|
||||||
|
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "--allow-empty", "-m", "initial")
|
||||||
|
return dir
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGitStatusToolShowsChangedFiles(t *testing.T) {
|
||||||
|
requireGit(t)
|
||||||
|
dir := initTestRepo(t)
|
||||||
|
|
||||||
|
path := filepath.Join(dir, "tracked.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
runGitForTest(t, dir, "add", "tracked.txt")
|
||||||
|
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "-m", "add tracked")
|
||||||
|
|
||||||
|
if err := os.WriteFile(path, []byte("world"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
tool := NewGitStatusTool(dir)
|
||||||
|
out, err := tool.Execute(context.Background(), nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "Changed files:") {
|
||||||
|
t.Fatalf("output missing 'Changed files:': %q", out)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "M tracked.txt") {
|
||||||
|
t.Fatalf("output missing 'M tracked.txt': %q", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGitDiffToolShowsUnstagedDiff(t *testing.T) {
|
||||||
|
requireGit(t)
|
||||||
|
dir := initTestRepo(t)
|
||||||
|
|
||||||
|
path := filepath.Join(dir, "tracked.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
runGitForTest(t, dir, "add", "tracked.txt")
|
||||||
|
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "-m", "add tracked")
|
||||||
|
|
||||||
|
if err := os.WriteFile(path, []byte("world"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
tool := NewGitDiffTool(dir)
|
||||||
|
out, err := tool.Execute(context.Background(), nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "diff --git") {
|
||||||
|
t.Fatalf("output missing 'diff --git': %q", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGitDiffToolShowsStat(t *testing.T) {
|
||||||
|
requireGit(t)
|
||||||
|
dir := initTestRepo(t)
|
||||||
|
|
||||||
|
path := filepath.Join(dir, "tracked.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
runGitForTest(t, dir, "add", "tracked.txt")
|
||||||
|
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "-m", "add tracked")
|
||||||
|
|
||||||
|
if err := os.WriteFile(path, []byte("world"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
tool := NewGitDiffTool(dir)
|
||||||
|
out, err := tool.Execute(context.Background(), json.RawMessage(`{"stat":true}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "tracked.txt") {
|
||||||
|
t.Fatalf("output missing 'tracked.txt': %q", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGitLogToolShowsCommits(t *testing.T) {
|
||||||
|
requireGit(t)
|
||||||
|
dir := initTestRepo(t)
|
||||||
|
|
||||||
|
tool := NewGitLogTool(dir)
|
||||||
|
out, err := tool.Execute(context.Background(), json.RawMessage(`{"max_count":1}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "initial") {
|
||||||
|
t.Fatalf("output missing 'initial': %q", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGitToolsRejectEscapingPath(t *testing.T) {
|
||||||
|
requireGit(t)
|
||||||
|
dir := initTestRepo(t)
|
||||||
|
|
||||||
|
tool := NewGitStatusTool(dir)
|
||||||
|
_, err := tool.Execute(context.Background(), json.RawMessage(`{"path":"../outside"}`))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for escaping path")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "escapes") {
|
||||||
|
t.Fatalf("error = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
+19
-2
@@ -1,7 +1,9 @@
|
|||||||
package tools
|
package tools
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"fmt"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
)
|
)
|
||||||
|
|
||||||
func resolvePath(baseDir, path string) (string, error) {
|
func resolvePath(baseDir, path string) (string, error) {
|
||||||
@@ -9,7 +11,22 @@ func resolvePath(baseDir, path string) (string, error) {
|
|||||||
path = "."
|
path = "."
|
||||||
}
|
}
|
||||||
if filepath.IsAbs(path) {
|
if filepath.IsAbs(path) {
|
||||||
return filepath.Clean(path), nil
|
return "", fmt.Errorf("absolute paths are not allowed; use a path relative to the project directory")
|
||||||
}
|
}
|
||||||
return filepath.Abs(filepath.Join(baseDir, path))
|
resolved := filepath.Clean(filepath.Join(baseDir, path))
|
||||||
|
if !isWithinDir(resolved, baseDir) {
|
||||||
|
return "", fmt.Errorf("path escapes the project directory: %s", path)
|
||||||
|
}
|
||||||
|
return resolved, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func isWithinDir(path, dir string) bool {
|
||||||
|
rel, err := filepath.Rel(dir, path)
|
||||||
|
if err != nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return true
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,7 +9,7 @@ import (
|
|||||||
"os/exec"
|
"os/exec"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"agentu/internal/llm"
|
"agentu/pkg/llm"
|
||||||
)
|
)
|
||||||
|
|
||||||
type ShellRunTool struct {
|
type ShellRunTool struct {
|
||||||
@@ -19,6 +19,7 @@ type ShellRunTool struct {
|
|||||||
func NewShellRunTool(workingDir string) *ShellRunTool {
|
func NewShellRunTool(workingDir string) *ShellRunTool {
|
||||||
return &ShellRunTool{workingDir: workingDir}
|
return &ShellRunTool{workingDir: workingDir}
|
||||||
}
|
}
|
||||||
|
func (*ShellRunTool) Mutates() bool { return true }
|
||||||
|
|
||||||
func (t *ShellRunTool) Definition() llm.Tool {
|
func (t *ShellRunTool) Definition() llm.Tool {
|
||||||
return llm.Tool{
|
return llm.Tool{
|
||||||
|
|||||||
+76
-6
@@ -5,17 +5,32 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"sort"
|
"sort"
|
||||||
|
"strings"
|
||||||
|
|
||||||
"agentu/internal/llm"
|
"agentu/pkg/llm"
|
||||||
)
|
)
|
||||||
|
|
||||||
const MaxToolOutputBytes = 64 * 1024
|
const MaxToolOutputBytes = 64 * 1024
|
||||||
|
|
||||||
|
const truncationNoticeBudget = 160
|
||||||
|
|
||||||
type Tool interface {
|
type Tool interface {
|
||||||
Definition() llm.Tool
|
Definition() llm.Tool
|
||||||
Execute(ctx context.Context, args json.RawMessage) (string, error)
|
Execute(ctx context.Context, args json.RawMessage) (string, error)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// MutatingTool is optionally implemented by tools that modify external state.
|
||||||
|
// It is used by the agent to determine whether tool calls can run in parallel.
|
||||||
|
type MutatingTool interface {
|
||||||
|
Mutates() bool
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsMutating reports whether tool implements MutatingTool and returns true.
|
||||||
|
func IsMutating(tool Tool) bool {
|
||||||
|
mutating, ok := tool.(MutatingTool)
|
||||||
|
return ok && mutating.Mutates()
|
||||||
|
}
|
||||||
|
|
||||||
type Registry struct {
|
type Registry struct {
|
||||||
tools map[string]Tool
|
tools map[string]Tool
|
||||||
primary []string
|
primary []string
|
||||||
@@ -62,26 +77,40 @@ func ReadOnlyBuiltins(workingDir string) *Registry {
|
|||||||
fileRead := NewFileReadTool(workingDir)
|
fileRead := NewFileReadTool(workingDir)
|
||||||
fileList := NewFileListTool(workingDir)
|
fileList := NewFileListTool(workingDir)
|
||||||
fileSearch := NewFileSearchTool(workingDir)
|
fileSearch := NewFileSearchTool(workingDir)
|
||||||
|
codeSymbols := NewCodeSymbolsTool(workingDir)
|
||||||
|
gitStatus := NewGitStatusTool(workingDir)
|
||||||
|
gitDiff := NewGitDiffTool(workingDir)
|
||||||
|
gitLog := NewGitLogTool(workingDir)
|
||||||
|
|
||||||
registry := NewRegistry(fileRead, fileList, fileSearch)
|
registry := NewRegistry(fileRead, fileList, fileSearch, codeSymbols, gitStatus, gitDiff, gitLog)
|
||||||
registry.RegisterAlias("file.read", fileRead)
|
registry.RegisterAlias("file.read", fileRead)
|
||||||
registry.RegisterAlias("file.list", fileList)
|
registry.RegisterAlias("file.list", fileList)
|
||||||
registry.RegisterAlias("file.search", fileSearch)
|
registry.RegisterAlias("file.search", fileSearch)
|
||||||
|
registry.RegisterAlias("code.symbols", codeSymbols)
|
||||||
|
registry.RegisterAlias("git.status", gitStatus)
|
||||||
|
registry.RegisterAlias("git.diff", gitDiff)
|
||||||
|
registry.RegisterAlias("git.log", gitLog)
|
||||||
return registry
|
return registry
|
||||||
}
|
}
|
||||||
|
|
||||||
func Builtins(workingDir string) *Registry {
|
func Builtins(workingDir string) *Registry {
|
||||||
registry := ReadOnlyBuiltins(workingDir)
|
registry := ReadOnlyBuiltins(workingDir)
|
||||||
fileWrite := NewFileWriteTool(workingDir)
|
fileWrite := NewFileWriteTool(workingDir)
|
||||||
|
fileEdit := NewFileEditTool(workingDir)
|
||||||
shellRun := NewShellRunTool(workingDir)
|
shellRun := NewShellRunTool(workingDir)
|
||||||
registry.Register(fileWrite)
|
registry.Register(fileWrite)
|
||||||
|
registry.Register(fileEdit)
|
||||||
registry.Register(shellRun)
|
registry.Register(shellRun)
|
||||||
registry.RegisterAlias("file.write", fileWrite)
|
registry.RegisterAlias("file.write", fileWrite)
|
||||||
|
registry.RegisterAlias("file.edit", fileEdit)
|
||||||
registry.RegisterAlias("shell.run", shellRun)
|
registry.RegisterAlias("shell.run", shellRun)
|
||||||
return registry
|
return registry
|
||||||
}
|
}
|
||||||
|
|
||||||
func AgentInstructions(workingDir string, yolo bool) string {
|
func AgentInstructions(workingDir string, yolo bool) string {
|
||||||
|
searchInstructions := webSearchInstructions()
|
||||||
|
readOnlyTools := "file_list, file_read, file_search, code_symbols, git_status, git_diff, git_log, web_search, web_fetch"
|
||||||
|
yoloTools := "file_list, file_read, file_search, code_symbols, file_write, file_edit, shell_run, git_status, git_diff, git_log, web_search, web_fetch"
|
||||||
if yolo {
|
if yolo {
|
||||||
return fmt.Sprintf(`Local project context:
|
return fmt.Sprintf(`Local project context:
|
||||||
- The project working directory is %q.
|
- The project working directory is %q.
|
||||||
@@ -89,11 +118,15 @@ func AgentInstructions(workingDir string, yolo bool) string {
|
|||||||
- Do not call tools for greetings, small talk, simple Q&A, or general questions that do not require local project or command context.
|
- Do not call tools for greetings, small talk, simple Q&A, or general questions that do not require local project or command context.
|
||||||
- Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files.
|
- Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files.
|
||||||
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
|
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
|
||||||
|
- When the user asks about code structure, symbols, functions, methods, types, or definitions in Go files, use code_symbols before reading whole files.
|
||||||
- When the user asks to run, execute, test, check, debug, or inspect a shell command, use shell_run. This includes local commands written in backticks, such as go, make, git, pwd, ls, ps, and similar terminal commands.
|
- When the user asks to run, execute, test, check, debug, or inspect a shell command, use shell_run. This includes local commands written in backticks, such as go, make, git, pwd, ls, ps, and similar terminal commands.
|
||||||
|
- When editing existing files, prefer file_edit for small or localized changes; use file_write only for creating files, replacing whole files intentionally, or appending.
|
||||||
- Do not use file tools as a substitute for a command execution request.
|
- Do not use file tools as a substitute for a command execution request.
|
||||||
|
- When the user asks about changed files, diffs, or git history, use git_status, git_diff, or git_log instead of shell_run.
|
||||||
- Do not start interactive or long-running sessions. Prefer bounded, non-interactive local commands with clear completion conditions.
|
- Do not start interactive or long-running sessions. Prefer bounded, non-interactive local commands with clear completion conditions.
|
||||||
|
- %s
|
||||||
- Tool paths should normally be relative to the project working directory.
|
- Tool paths should normally be relative to the project working directory.
|
||||||
- Available tools: file_list, file_read, file_search, file_write, shell_run.`, workingDir)
|
- Available tools: %s.`, workingDir, searchInstructions, yoloTools)
|
||||||
}
|
}
|
||||||
return fmt.Sprintf(`Local project context:
|
return fmt.Sprintf(`Local project context:
|
||||||
- The project working directory is %q.
|
- The project working directory is %q.
|
||||||
@@ -101,9 +134,16 @@ func AgentInstructions(workingDir string, yolo bool) string {
|
|||||||
- Do not call tools for greetings, small talk, simple Q&A, or general questions that do not require local project context.
|
- Do not call tools for greetings, small talk, simple Q&A, or general questions that do not require local project context.
|
||||||
- Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files.
|
- Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files.
|
||||||
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
|
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
|
||||||
|
- When the user asks about code structure, symbols, functions, methods, types, or definitions in Go files, use code_symbols before reading whole files.
|
||||||
- Shell command execution is disabled in this mode. When the user asks to run, execute, test, check, debug, or inspect a shell command, do not use file tools as a substitute; explain that they need to restart agentu with --yolo to enable shell_run.
|
- Shell command execution is disabled in this mode. When the user asks to run, execute, test, check, debug, or inspect a shell command, do not use file tools as a substitute; explain that they need to restart agentu with --yolo to enable shell_run.
|
||||||
|
- When the user asks about changed files, diffs, or git history, use git_status, git_diff, or git_log instead of shell_run.
|
||||||
|
- %s
|
||||||
- Tool paths should normally be relative to the project working directory.
|
- Tool paths should normally be relative to the project working directory.
|
||||||
- Available tools: file_list, file_read, file_search.`, workingDir)
|
- Available tools: %s.`, workingDir, searchInstructions, readOnlyTools)
|
||||||
|
}
|
||||||
|
|
||||||
|
func webSearchInstructions() string {
|
||||||
|
return "When the user's request depends on current or external web information, use web_search. This includes news, prices, laws, schedules, software versions, product recommendations, or requests for the latest information. Use web_fetch to read a source URL when search snippets are not enough, when the user provides a URL, or when the user names a specific website to inspect. Base answers on returned sources and include source URLs when useful."
|
||||||
}
|
}
|
||||||
|
|
||||||
func JSONSchema(schema string) json.RawMessage {
|
func JSONSchema(schema string) json.RawMessage {
|
||||||
@@ -111,10 +151,40 @@ func JSONSchema(schema string) json.RawMessage {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func FormatResult(output string) string {
|
func FormatResult(output string) string {
|
||||||
if len(output) <= MaxToolOutputBytes {
|
return FormatResultLimit(output, MaxToolOutputBytes)
|
||||||
|
}
|
||||||
|
|
||||||
|
func FormatResultLimit(output string, limit int) string {
|
||||||
|
if limit <= 0 || len(output) <= limit {
|
||||||
return output
|
return output
|
||||||
}
|
}
|
||||||
return output[:MaxToolOutputBytes] + fmt.Sprintf("\n\n[truncated after %d bytes]", MaxToolOutputBytes)
|
if limit <= truncationNoticeBudget {
|
||||||
|
return output[:limit]
|
||||||
|
}
|
||||||
|
notice := fmt.Sprintf("\n\n[truncated: showing head and tail; omitted %d bytes]\n\n", len(output)-limit)
|
||||||
|
keep := limit - len(notice)
|
||||||
|
if keep <= 0 {
|
||||||
|
return output[:limit]
|
||||||
|
}
|
||||||
|
headLen := keep / 2
|
||||||
|
tailLen := keep - headLen
|
||||||
|
head := trimHeadAtLine(output[:headLen])
|
||||||
|
tail := trimTailAtLine(output[len(output)-tailLen:])
|
||||||
|
return head + notice + tail
|
||||||
|
}
|
||||||
|
|
||||||
|
func trimHeadAtLine(value string) string {
|
||||||
|
if idx := strings.LastIndexByte(value, '\n'); idx > 0 {
|
||||||
|
return value[:idx+1]
|
||||||
|
}
|
||||||
|
return value
|
||||||
|
}
|
||||||
|
|
||||||
|
func trimTailAtLine(value string) string {
|
||||||
|
if idx := strings.IndexByte(value, '\n'); idx >= 0 && idx+1 < len(value) {
|
||||||
|
return value[idx+1:]
|
||||||
|
}
|
||||||
|
return value
|
||||||
}
|
}
|
||||||
|
|
||||||
func FormatError(err error) string {
|
func FormatError(err error) string {
|
||||||
|
|||||||
@@ -43,6 +43,41 @@ func TestFileTools(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestFileSearchSupportsContextAndMaxResults(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "notes.txt")
|
||||||
|
content := strings.Join([]string{"before", "needle one", "middle", "needle two", "after"}, "\n")
|
||||||
|
if err := os.WriteFile(path, []byte(content), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
search := NewFileSearchTool(dir)
|
||||||
|
out, err := search.walkSearch(context.Background(), fileSearchArgs{Query: "needle", ContextLines: 1, MaxResults: 1}, dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, want := range []string{"-- ", "before", "needle one", "middle", "stopped after 1 matches"} {
|
||||||
|
if !strings.Contains(out, want) {
|
||||||
|
t.Fatalf("search output missing %q:\n%s", want, out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if strings.Contains(out, "needle two") {
|
||||||
|
t.Fatalf("search should stop after max_results:\n%s", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFormatResultKeepsHeadAndTail(t *testing.T) {
|
||||||
|
input := strings.Join([]string{"head-1", "head-2", strings.Repeat("x", 220), "tail-1", "tail-2"}, "\n")
|
||||||
|
out := FormatResultLimit(input, 220)
|
||||||
|
for _, want := range []string{"head-1", "truncated: showing head and tail", "tail-2"} {
|
||||||
|
if !strings.Contains(out, want) {
|
||||||
|
t.Fatalf("formatted output missing %q:\n%s", want, out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(out) > 220 {
|
||||||
|
t.Fatalf("formatted output len = %d, want <= 220", len(out))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestShellRunTool(t *testing.T) {
|
func TestShellRunTool(t *testing.T) {
|
||||||
dir := t.TempDir()
|
dir := t.TempDir()
|
||||||
if err := os.WriteFile(filepath.Join(dir, "x.txt"), []byte("x"), 0o644); err != nil {
|
if err := os.WriteFile(filepath.Join(dir, "x.txt"), []byte("x"), 0o644); err != nil {
|
||||||
@@ -69,12 +104,12 @@ func TestReadOnlyBuiltinsExposeOpenAICompatibleNamesAndAliases(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
for _, want := range []string{"file_list", "file_read", "file_search"} {
|
for _, want := range []string{"file_list", "file_read", "file_search", "code_symbols", "git_status", "git_diff", "git_log"} {
|
||||||
if !slices.Contains(names, want) {
|
if !slices.Contains(names, want) {
|
||||||
t.Fatalf("definitions missing %s: %#v", want, names)
|
t.Fatalf("definitions missing %s: %#v", want, names)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
for _, alias := range []string{"file.list", "file.read", "file.search"} {
|
for _, alias := range []string{"file.list", "file.read", "file.search", "code.symbols", "git.status", "git.diff", "git.log"} {
|
||||||
if _, ok := registry.Get(alias); !ok {
|
if _, ok := registry.Get(alias); !ok {
|
||||||
t.Fatalf("alias missing: %s", alias)
|
t.Fatalf("alias missing: %s", alias)
|
||||||
}
|
}
|
||||||
@@ -82,11 +117,63 @@ func TestReadOnlyBuiltinsExposeOpenAICompatibleNamesAndAliases(t *testing.T) {
|
|||||||
if _, ok := registry.Get("file_write"); ok {
|
if _, ok := registry.Get("file_write"); ok {
|
||||||
t.Fatal("read-only registry should not expose file_write")
|
t.Fatal("read-only registry should not expose file_write")
|
||||||
}
|
}
|
||||||
|
if _, ok := registry.Get("file_edit"); ok {
|
||||||
|
t.Fatal("read-only registry should not expose file_edit")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCodeSymbolsToolListsGoSymbols(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "sample.go")
|
||||||
|
if err := os.WriteFile(path, []byte(`package sample
|
||||||
|
|
||||||
|
import "context"
|
||||||
|
|
||||||
|
type Worker struct{}
|
||||||
|
type Runner interface{ Run(context.Context) error }
|
||||||
|
|
||||||
|
func NewWorker() *Worker { return &Worker{} }
|
||||||
|
func (w *Worker) Run(ctx context.Context) error { return nil }
|
||||||
|
`), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
tool := NewCodeSymbolsTool(dir)
|
||||||
|
out, err := tool.Execute(context.Background(), json.RawMessage(`{"path":"sample.go"}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, want := range []string{"package sample", "imports: context", "type Worker struct", "type Runner interface", "func NewWorker", "method *Worker.Run"} {
|
||||||
|
if !strings.Contains(out, want) {
|
||||||
|
t.Fatalf("symbols output missing %q:\n%s", want, out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRegistryCanExposeWebSearch(t *testing.T) {
|
||||||
|
registry := ReadOnlyBuiltins(t.TempDir())
|
||||||
|
registry.Register(NewBuiltinSearchTool(nil))
|
||||||
|
registry.Register(NewBuiltinFetchTool(nil))
|
||||||
|
|
||||||
|
var names []string
|
||||||
|
for _, def := range registry.Definitions() {
|
||||||
|
names = append(names, def.Function.Name)
|
||||||
|
}
|
||||||
|
for _, want := range []string{"web_search", "web_fetch"} {
|
||||||
|
if !slices.Contains(names, want) {
|
||||||
|
t.Fatalf("definitions missing %s: %#v", want, names)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, want := range []string{"web_search", "web_fetch"} {
|
||||||
|
if _, ok := registry.Get(want); !ok {
|
||||||
|
t.Fatalf("%s missing", want)
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestBuiltinsExposeYoloTools(t *testing.T) {
|
func TestBuiltinsExposeYoloTools(t *testing.T) {
|
||||||
registry := Builtins(t.TempDir())
|
registry := Builtins(t.TempDir())
|
||||||
for _, name := range []string{"file_write", "shell_run", "file.write", "shell.run"} {
|
for _, name := range []string{"file_write", "file_edit", "shell_run", "file.write", "file.edit", "shell.run"} {
|
||||||
if _, ok := registry.Get(name); !ok {
|
if _, ok := registry.Get(name); !ok {
|
||||||
t.Fatalf("tool missing: %s", name)
|
t.Fatalf("tool missing: %s", name)
|
||||||
}
|
}
|
||||||
@@ -95,16 +182,94 @@ func TestBuiltinsExposeYoloTools(t *testing.T) {
|
|||||||
|
|
||||||
func TestAgentInstructionsExplainCommandModes(t *testing.T) {
|
func TestAgentInstructionsExplainCommandModes(t *testing.T) {
|
||||||
readOnly := AgentInstructions("/tmp/project", false)
|
readOnly := AgentInstructions("/tmp/project", false)
|
||||||
for _, want := range []string{"Shell command execution is disabled", "--yolo", "Do not call tools for greetings", "Do not explore the project preemptively", "Do not call file_list just to discover context"} {
|
for _, want := range []string{"Shell command execution is disabled", "--yolo", "Do not call tools for greetings", "Do not explore the project preemptively", "Do not call file_list just to discover context", "git_status, git_diff, git_log"} {
|
||||||
if !strings.Contains(readOnly, want) {
|
if !strings.Contains(readOnly, want) {
|
||||||
t.Fatalf("read-only instructions missing %q:\n%s", want, readOnly)
|
t.Fatalf("read-only instructions missing %q:\n%s", want, readOnly)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
yolo := AgentInstructions("/tmp/project", true)
|
yolo := AgentInstructions("/tmp/project", true)
|
||||||
for _, want := range []string{"use shell_run", "commands written in backticks", "Do not start interactive", "non-interactive local commands", "Do not call tools for greetings"} {
|
for _, want := range []string{"use shell_run", "commands written in backticks", "Do not start interactive", "non-interactive local commands", "Do not call tools for greetings", "prefer file_edit"} {
|
||||||
if !strings.Contains(yolo, want) {
|
if !strings.Contains(yolo, want) {
|
||||||
t.Fatalf("yolo instructions missing %q:\n%s", want, yolo)
|
t.Fatalf("yolo instructions missing %q:\n%s", want, yolo)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestAgentInstructionsExplainWebSearch(t *testing.T) {
|
||||||
|
withSearch := AgentInstructions("/tmp/project", false)
|
||||||
|
for _, want := range []string{"use web_search", "Use web_fetch", "latest information", "source URLs", "Available tools: file_list, file_read, file_search, code_symbols, git_status, git_diff, git_log, web_search, web_fetch"} {
|
||||||
|
if !strings.Contains(withSearch, want) {
|
||||||
|
t.Fatalf("web search instructions missing %q:\n%s", want, withSearch)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolvePathRejectsAbsolutePaths(t *testing.T) {
|
||||||
|
base := t.TempDir()
|
||||||
|
_, err := resolvePath(base, "/etc/passwd")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for absolute path")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "absolute paths are not allowed") {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolvePathRejectsEscapingBaseDir(t *testing.T) {
|
||||||
|
base := filepath.Join(t.TempDir(), "project")
|
||||||
|
if err := os.MkdirAll(base, 0o755); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_, err := resolvePath(base, "../../../etc/passwd")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for path escaping base dir")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "escapes the project directory") {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolvePathAllowsSubdirTraversal(t *testing.T) {
|
||||||
|
base := t.TempDir()
|
||||||
|
if err := os.MkdirAll(filepath.Join(base, "a", "b"), 0o755); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(filepath.Join(base, "a", "target.txt"), []byte("ok"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
// "a/b/../target.txt" resolves to "a/target.txt" which is within base.
|
||||||
|
resolved, err := resolvePath(base, "a/b/../target.txt")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
want := filepath.Join(base, "a", "target.txt")
|
||||||
|
if resolved != want {
|
||||||
|
t.Fatalf("resolved = %q, want %q", resolved, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolvePathRejectsDotDotThatEscapes(t *testing.T) {
|
||||||
|
base := t.TempDir()
|
||||||
|
sub := filepath.Join(base, "a")
|
||||||
|
if err := os.MkdirAll(sub, 0o755); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_, err := resolvePath(sub, "..")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for path escaping via ..")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIsMutating(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
if IsMutating(NewFileReadTool(dir)) {
|
||||||
|
t.Fatal("FileReadTool should not be mutating")
|
||||||
|
}
|
||||||
|
if !IsMutating(NewFileWriteTool(dir)) {
|
||||||
|
t.Fatal("FileWriteTool should be mutating")
|
||||||
|
}
|
||||||
|
if !IsMutating(NewShellRunTool(dir)) {
|
||||||
|
t.Fatal("ShellRunTool should be mutating")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,597 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/url"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
|
||||||
|
"golang.org/x/net/html"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
defaultSearchBaseURL = "https://html.duckduckgo.com/html/"
|
||||||
|
defaultUserAgent = "agentu/0.0.3"
|
||||||
|
maxSearchResults = 20
|
||||||
|
maxWebReadBytes = 2 * 1024 * 1024
|
||||||
|
)
|
||||||
|
|
||||||
|
type BuiltinSearchTool struct {
|
||||||
|
baseURL string
|
||||||
|
maxResults int
|
||||||
|
client *http.Client
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewBuiltinSearchTool(client *http.Client) *BuiltinSearchTool {
|
||||||
|
return NewBuiltinSearchToolWithBaseURL(defaultSearchBaseURL, 5, client)
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewBuiltinSearchToolWithBaseURL(baseURL string, maxResults int, client *http.Client) *BuiltinSearchTool {
|
||||||
|
if strings.TrimSpace(baseURL) == "" {
|
||||||
|
baseURL = defaultSearchBaseURL
|
||||||
|
}
|
||||||
|
if maxResults <= 0 {
|
||||||
|
maxResults = 5
|
||||||
|
}
|
||||||
|
if client == nil {
|
||||||
|
client = http.DefaultClient
|
||||||
|
}
|
||||||
|
return &BuiltinSearchTool{
|
||||||
|
baseURL: baseURL,
|
||||||
|
maxResults: maxResults,
|
||||||
|
client: client,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *BuiltinSearchTool) Definition() llm.Tool {
|
||||||
|
return webSearchDefinition()
|
||||||
|
}
|
||||||
|
|
||||||
|
func webSearchDefinition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "web_search",
|
||||||
|
Description: "Search the public web for current or external information. Use for recent facts, news, prices, schedules, versions, policies, recommendations, or other questions that need online sources. Results include source URLs.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"query": {"type": "string", "description": "Search query."},
|
||||||
|
"max_results": {"type": "integer", "description": "Optional result count from 1 to 20. Defaults to 5.", "minimum": 1, "maximum": 20},
|
||||||
|
"topic": {"type": "string", "description": "Optional topic hint. Accepted for compatibility but the built-in backend may ignore it.", "enum": ["general", "news", "finance"]}
|
||||||
|
},
|
||||||
|
"required": ["query"],
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *BuiltinSearchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
args, err := parseWebSearchArgs(raw, t.maxResults)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
endpoint, err := url.Parse(t.baseURL)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("parse search endpoint: %w", err)
|
||||||
|
}
|
||||||
|
query := endpoint.Query()
|
||||||
|
query.Set("q", args.Query)
|
||||||
|
endpoint.RawQuery = query.Encode()
|
||||||
|
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint.String(), nil)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("create search request: %w", err)
|
||||||
|
}
|
||||||
|
req.Header.Set("User-Agent", defaultUserAgent)
|
||||||
|
req.Header.Set("Accept", "text/html,application/xhtml+xml")
|
||||||
|
|
||||||
|
resp, err := t.client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("search request: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||||
|
body, _ := io.ReadAll(io.LimitReader(resp.Body, 8*1024))
|
||||||
|
return "", fmt.Errorf("web search failed: status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
|
||||||
|
}
|
||||||
|
|
||||||
|
results, err := parseSearchResults(resp.Body, args.MaxResults)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return FormatResult(formatWebSearchResults(args.Query, "", results)), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type BuiltinFetchTool struct {
|
||||||
|
client *http.Client
|
||||||
|
}
|
||||||
|
|
||||||
|
type webSearchArgs struct {
|
||||||
|
Query string `json:"query"`
|
||||||
|
MaxResults int `json:"max_results"`
|
||||||
|
Topic string `json:"topic"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type webSearchResult struct {
|
||||||
|
Title string `json:"title"`
|
||||||
|
URL string `json:"url"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
Score float64 `json:"score"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type webFetchArgs struct {
|
||||||
|
URL string `json:"url"`
|
||||||
|
Query string `json:"query"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type webFetchResponse struct {
|
||||||
|
Results []webFetchResult `json:"results"`
|
||||||
|
FailedResults []webFetchFailed `json:"failed_results"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type webFetchResult struct {
|
||||||
|
URL string `json:"url"`
|
||||||
|
RawContent string `json:"raw_content"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type webFetchFailed struct {
|
||||||
|
URL string `json:"url"`
|
||||||
|
Error string `json:"error"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewBuiltinFetchTool(client *http.Client) *BuiltinFetchTool {
|
||||||
|
if client == nil {
|
||||||
|
client = http.DefaultClient
|
||||||
|
}
|
||||||
|
return &BuiltinFetchTool{client: client}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *BuiltinFetchTool) Definition() llm.Tool {
|
||||||
|
return webFetchDefinition()
|
||||||
|
}
|
||||||
|
|
||||||
|
func webFetchDefinition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "web_fetch",
|
||||||
|
Description: "Fetch and extract readable Markdown content from a public web URL. Use after web_search when snippets are not enough, or when the user gives a URL or website to inspect. Results include the source URL.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"url": {"type": "string", "description": "Public http or https URL to fetch. URLs without a scheme are treated as https URLs."},
|
||||||
|
"query": {"type": "string", "description": "Optional extraction focus. Accepted for compatibility but the built-in backend may ignore it."}
|
||||||
|
},
|
||||||
|
"required": ["url"],
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *BuiltinFetchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
args, err := parseWebFetchArgs(raw)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, args.URL, nil)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("create fetch request: %w", err)
|
||||||
|
}
|
||||||
|
req.Header.Set("User-Agent", defaultUserAgent)
|
||||||
|
req.Header.Set("Accept", "text/html,text/plain,application/xhtml+xml")
|
||||||
|
|
||||||
|
resp, err := t.client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("fetch request: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||||
|
body, _ := io.ReadAll(io.LimitReader(resp.Body, 8*1024))
|
||||||
|
return "", fmt.Errorf("web fetch failed: status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
|
||||||
|
}
|
||||||
|
|
||||||
|
data, err := io.ReadAll(io.LimitReader(resp.Body, maxWebReadBytes))
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("read response: %w", err)
|
||||||
|
}
|
||||||
|
contentType := strings.ToLower(resp.Header.Get("Content-Type"))
|
||||||
|
var content string
|
||||||
|
if strings.Contains(contentType, "html") || contentType == "" {
|
||||||
|
content, err = htmlToMarkdown(string(data), args.URL)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
} else if strings.Contains(contentType, "text/plain") || strings.Contains(contentType, "markdown") {
|
||||||
|
content = strings.TrimSpace(string(data))
|
||||||
|
} else {
|
||||||
|
return "", fmt.Errorf("unsupported content type: %s", contentType)
|
||||||
|
}
|
||||||
|
|
||||||
|
result := webFetchResponse{
|
||||||
|
Results: []webFetchResult{{
|
||||||
|
URL: args.URL,
|
||||||
|
RawContent: content,
|
||||||
|
}},
|
||||||
|
}
|
||||||
|
return FormatResult(formatWebFetchResults(args.URL, result)), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseWebSearchArgs(raw json.RawMessage, defaultMaxResults int) (webSearchArgs, error) {
|
||||||
|
var args webSearchArgs
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return args, fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
args.Query = strings.TrimSpace(args.Query)
|
||||||
|
if args.Query == "" {
|
||||||
|
return args, errors.New("query is required")
|
||||||
|
}
|
||||||
|
if args.MaxResults == 0 {
|
||||||
|
args.MaxResults = defaultMaxResults
|
||||||
|
}
|
||||||
|
if args.MaxResults < 1 || args.MaxResults > maxSearchResults {
|
||||||
|
return args, fmt.Errorf("max_results must be between 1 and %d", maxSearchResults)
|
||||||
|
}
|
||||||
|
args.Topic = strings.ToLower(strings.TrimSpace(args.Topic))
|
||||||
|
if args.Topic == "" {
|
||||||
|
args.Topic = "general"
|
||||||
|
}
|
||||||
|
switch args.Topic {
|
||||||
|
case "general", "news", "finance":
|
||||||
|
default:
|
||||||
|
return args, fmt.Errorf("topic must be one of: general, news, finance")
|
||||||
|
}
|
||||||
|
return args, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseWebFetchArgs(raw json.RawMessage) (webFetchArgs, error) {
|
||||||
|
var args webFetchArgs
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return args, fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
args.URL = strings.TrimSpace(args.URL)
|
||||||
|
if args.URL == "" {
|
||||||
|
return args, errors.New("url is required")
|
||||||
|
}
|
||||||
|
if !strings.Contains(args.URL, "://") {
|
||||||
|
args.URL = "https://" + args.URL
|
||||||
|
}
|
||||||
|
parsed, err := url.Parse(args.URL)
|
||||||
|
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Host == "" {
|
||||||
|
return args, errors.New("url must be a valid http or https URL")
|
||||||
|
}
|
||||||
|
args.Query = strings.TrimSpace(args.Query)
|
||||||
|
return args, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func formatWebSearchResults(query string, answer string, results []webSearchResult) string {
|
||||||
|
if len(results) == 0 {
|
||||||
|
return fmt.Sprintf("No web search results for %q.", query)
|
||||||
|
}
|
||||||
|
|
||||||
|
var b strings.Builder
|
||||||
|
fmt.Fprintf(&b, "Web search results for %q:\n", query)
|
||||||
|
if strings.TrimSpace(answer) != "" {
|
||||||
|
fmt.Fprintf(&b, "\nSummary: %s\n", strings.TrimSpace(answer))
|
||||||
|
}
|
||||||
|
for i, result := range results {
|
||||||
|
title := strings.TrimSpace(result.Title)
|
||||||
|
if title == "" {
|
||||||
|
title = "Untitled"
|
||||||
|
}
|
||||||
|
fmt.Fprintf(&b, "\n%d. %s\n", i+1, title)
|
||||||
|
fmt.Fprintf(&b, " Source: %s\n", strings.TrimSpace(result.URL))
|
||||||
|
if content := strings.TrimSpace(result.Content); content != "" {
|
||||||
|
fmt.Fprintf(&b, " Snippet: %s\n", content)
|
||||||
|
}
|
||||||
|
if result.Score > 0 {
|
||||||
|
fmt.Fprintf(&b, " Score: %.3f\n", result.Score)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return strings.TrimRight(b.String(), "\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
func formatWebFetchResults(inputURL string, response webFetchResponse) string {
|
||||||
|
var b strings.Builder
|
||||||
|
if len(response.Results) == 0 {
|
||||||
|
fmt.Fprintf(&b, "No readable web content extracted from %s.", inputURL)
|
||||||
|
} else {
|
||||||
|
fmt.Fprintf(&b, "Fetched web content:\n")
|
||||||
|
for i, result := range response.Results {
|
||||||
|
sourceURL := strings.TrimSpace(result.URL)
|
||||||
|
if sourceURL == "" {
|
||||||
|
sourceURL = inputURL
|
||||||
|
}
|
||||||
|
fmt.Fprintf(&b, "\n%d. Source: %s\n", i+1, sourceURL)
|
||||||
|
if content := strings.TrimSpace(result.RawContent); content != "" {
|
||||||
|
fmt.Fprintf(&b, "\n%s\n", content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, failed := range response.FailedResults {
|
||||||
|
sourceURL := strings.TrimSpace(failed.URL)
|
||||||
|
if sourceURL == "" {
|
||||||
|
sourceURL = inputURL
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(failed.Error) != "" {
|
||||||
|
fmt.Fprintf(&b, "\nFetch failed for %s: %s", sourceURL, strings.TrimSpace(failed.Error))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return strings.TrimRight(b.String(), "\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseSearchResults(r io.Reader, maxResults int) ([]webSearchResult, error) {
|
||||||
|
doc, err := html.Parse(r)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("parse search results: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
results := make([]webSearchResult, 0, maxResults)
|
||||||
|
var walk func(*html.Node)
|
||||||
|
walk = func(n *html.Node) {
|
||||||
|
if len(results) >= maxResults || n == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if n.Type == html.ElementNode && n.Data == "div" && hasClass(n, "result") {
|
||||||
|
if result, ok := parseResultNode(n); ok {
|
||||||
|
results = append(results, result)
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
for child := n.FirstChild; child != nil; child = child.NextSibling {
|
||||||
|
walk(child)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
walk(doc)
|
||||||
|
return results, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseResultNode(n *html.Node) (webSearchResult, bool) {
|
||||||
|
var result webSearchResult
|
||||||
|
var walk func(*html.Node)
|
||||||
|
walk = func(node *html.Node) {
|
||||||
|
if node == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if node.Type == html.ElementNode && node.Data == "a" {
|
||||||
|
if hasClass(node, "result__a") && result.URL == "" {
|
||||||
|
result.Title = cleanWebText(textContent(node))
|
||||||
|
result.URL = decodeSearchURL(attr(node, "href"))
|
||||||
|
}
|
||||||
|
if hasClass(node, "result__snippet") && result.Content == "" {
|
||||||
|
result.Content = cleanWebText(textContent(node))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if node.Type == html.ElementNode && (node.Data == "div" || node.Data == "span") && hasClass(node, "result__snippet") && result.Content == "" {
|
||||||
|
result.Content = cleanWebText(textContent(node))
|
||||||
|
}
|
||||||
|
for child := node.FirstChild; child != nil; child = child.NextSibling {
|
||||||
|
walk(child)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
walk(n)
|
||||||
|
return result, result.Title != "" && result.URL != ""
|
||||||
|
}
|
||||||
|
|
||||||
|
func decodeSearchURL(raw string) string {
|
||||||
|
raw = strings.TrimSpace(raw)
|
||||||
|
if raw == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
parsed, err := url.Parse(raw)
|
||||||
|
if err == nil {
|
||||||
|
if target := parsed.Query().Get("uddg"); target != "" {
|
||||||
|
if decoded, err := url.QueryUnescape(target); err == nil {
|
||||||
|
return decoded
|
||||||
|
}
|
||||||
|
return target
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return raw
|
||||||
|
}
|
||||||
|
|
||||||
|
func htmlToMarkdown(input string, sourceURL string) (string, error) {
|
||||||
|
doc, err := html.Parse(strings.NewReader(input))
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("parse html: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var b strings.Builder
|
||||||
|
var title string
|
||||||
|
var walk func(*html.Node, int)
|
||||||
|
walk = func(n *html.Node, listDepth int) {
|
||||||
|
if n == nil || skipHTMLNode(n) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if n.Type == html.ElementNode && n.Data == "title" && title == "" {
|
||||||
|
title = cleanWebText(textContent(n))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if n.Type == html.TextNode {
|
||||||
|
writeInlineText(&b, n.Data)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if n.Type != html.ElementNode {
|
||||||
|
for child := n.FirstChild; child != nil; child = child.NextSibling {
|
||||||
|
walk(child, listDepth)
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
switch n.Data {
|
||||||
|
case "h1", "h2", "h3", "h4", "h5", "h6":
|
||||||
|
writeBlockBreak(&b)
|
||||||
|
level := int(n.Data[1] - '0')
|
||||||
|
b.WriteString(strings.Repeat("#", level))
|
||||||
|
b.WriteString(" ")
|
||||||
|
writeChildren(&b, n, listDepth, walk)
|
||||||
|
writeBlockBreak(&b)
|
||||||
|
case "p", "div", "section", "article", "header", "footer", "main", "tr":
|
||||||
|
writeBlockBreak(&b)
|
||||||
|
writeChildren(&b, n, listDepth, walk)
|
||||||
|
writeBlockBreak(&b)
|
||||||
|
case "br":
|
||||||
|
b.WriteString("\n")
|
||||||
|
case "li":
|
||||||
|
writeBlockBreak(&b)
|
||||||
|
b.WriteString(strings.Repeat(" ", listDepth))
|
||||||
|
b.WriteString("- ")
|
||||||
|
writeChildren(&b, n, listDepth+1, walk)
|
||||||
|
writeBlockBreak(&b)
|
||||||
|
case "a":
|
||||||
|
label := cleanWebText(textContent(n))
|
||||||
|
href := strings.TrimSpace(attr(n, "href"))
|
||||||
|
if href != "" && label != "" {
|
||||||
|
href = resolveURL(sourceURL, href)
|
||||||
|
b.WriteString(label)
|
||||||
|
b.WriteString(" (")
|
||||||
|
b.WriteString(href)
|
||||||
|
b.WriteString(")")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeChildren(&b, n, listDepth, walk)
|
||||||
|
default:
|
||||||
|
writeChildren(&b, n, listDepth, walk)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
walk(doc, 0)
|
||||||
|
|
||||||
|
content := normalizeMarkdownText(b.String())
|
||||||
|
if title != "" && !strings.Contains(content, title) {
|
||||||
|
content = "# " + title + "\n\n" + content
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(content) == "" {
|
||||||
|
return "", errors.New("no readable content found")
|
||||||
|
}
|
||||||
|
return content, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeChildren(b *strings.Builder, n *html.Node, listDepth int, walk func(*html.Node, int)) {
|
||||||
|
for child := n.FirstChild; child != nil; child = child.NextSibling {
|
||||||
|
walk(child, listDepth)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeInlineText(b *strings.Builder, text string) {
|
||||||
|
text = cleanWebText(text)
|
||||||
|
if text == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if b.Len() > 0 {
|
||||||
|
last := b.String()[b.Len()-1]
|
||||||
|
if last != '\n' && last != ' ' {
|
||||||
|
b.WriteByte(' ')
|
||||||
|
}
|
||||||
|
}
|
||||||
|
b.WriteString(text)
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeBlockBreak(b *strings.Builder) {
|
||||||
|
if b.Len() == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
text := b.String()
|
||||||
|
if strings.HasSuffix(text, "\n\n") {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if strings.HasSuffix(text, "\n") {
|
||||||
|
b.WriteByte('\n')
|
||||||
|
return
|
||||||
|
}
|
||||||
|
b.WriteString("\n\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
func skipHTMLNode(n *html.Node) bool {
|
||||||
|
if n.Type != html.ElementNode {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
switch n.Data {
|
||||||
|
case "script", "style", "noscript", "svg", "canvas", "template":
|
||||||
|
return true
|
||||||
|
default:
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func textContent(n *html.Node) string {
|
||||||
|
var b strings.Builder
|
||||||
|
var walk func(*html.Node)
|
||||||
|
walk = func(node *html.Node) {
|
||||||
|
if node == nil || skipHTMLNode(node) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if node.Type == html.TextNode {
|
||||||
|
b.WriteString(node.Data)
|
||||||
|
b.WriteByte(' ')
|
||||||
|
return
|
||||||
|
}
|
||||||
|
for child := node.FirstChild; child != nil; child = child.NextSibling {
|
||||||
|
walk(child)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
walk(n)
|
||||||
|
return cleanWebText(b.String())
|
||||||
|
}
|
||||||
|
|
||||||
|
func hasClass(n *html.Node, class string) bool {
|
||||||
|
for _, item := range strings.Fields(attr(n, "class")) {
|
||||||
|
if item == class {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func attr(n *html.Node, name string) string {
|
||||||
|
for _, item := range n.Attr {
|
||||||
|
if item.Key == name {
|
||||||
|
return item.Val
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
func resolveURL(baseURL, raw string) string {
|
||||||
|
ref, err := url.Parse(raw)
|
||||||
|
if err != nil {
|
||||||
|
return raw
|
||||||
|
}
|
||||||
|
base, err := url.Parse(baseURL)
|
||||||
|
if err != nil {
|
||||||
|
return raw
|
||||||
|
}
|
||||||
|
return base.ResolveReference(ref).String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func cleanWebText(text string) string {
|
||||||
|
return strings.Join(strings.Fields(text), " ")
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeMarkdownText(text string) string {
|
||||||
|
lines := strings.Split(text, "\n")
|
||||||
|
out := make([]string, 0, len(lines))
|
||||||
|
blank := false
|
||||||
|
for _, line := range lines {
|
||||||
|
line = strings.TrimSpace(line)
|
||||||
|
if line == "" {
|
||||||
|
if !blank && len(out) > 0 {
|
||||||
|
out = append(out, "")
|
||||||
|
}
|
||||||
|
blank = true
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
out = append(out, line)
|
||||||
|
blank = false
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(strings.Join(out, "\n"))
|
||||||
|
}
|
||||||
@@ -0,0 +1,96 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestBuiltinSearchTool(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method != http.MethodGet {
|
||||||
|
t.Fatalf("method = %s", r.Method)
|
||||||
|
}
|
||||||
|
if got := r.URL.Query().Get("q"); got != "agentu" {
|
||||||
|
t.Fatalf("query = %q", got)
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "text/html")
|
||||||
|
_, _ = w.Write([]byte(`
|
||||||
|
<html><body>
|
||||||
|
<div class="result">
|
||||||
|
<a class="result__a" href="/l/?uddg=https%3A%2F%2Fexample.com%2Fagentu">Agentu release</a>
|
||||||
|
<a class="result__snippet">Agentu adds built-in web tools.</a>
|
||||||
|
</div>
|
||||||
|
</body></html>`))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
tool := NewBuiltinSearchToolWithBaseURL(server.URL, 5, server.Client())
|
||||||
|
out, err := tool.Execute(context.Background(), json.RawMessage(`{"query":"agentu"}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, want := range []string{"Web search results", "Agentu release", "https://example.com/agentu", "Agentu adds built-in web tools"} {
|
||||||
|
if !strings.Contains(out, want) {
|
||||||
|
t.Fatalf("output missing %q:\n%s", want, out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuiltinFetchTool(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method != http.MethodGet {
|
||||||
|
t.Fatalf("method = %s", r.Method)
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||||
|
_, _ = w.Write([]byte(`
|
||||||
|
<html>
|
||||||
|
<head><title>Agentu</title><script>hidden()</script></head>
|
||||||
|
<body>
|
||||||
|
<h1>Agentu</h1>
|
||||||
|
<p>Built-in fetch reads HTML.</p>
|
||||||
|
<ul><li>Search</li><li>Fetch</li></ul>
|
||||||
|
</body>
|
||||||
|
</html>`))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
tool := NewBuiltinFetchTool(server.Client())
|
||||||
|
out, err := tool.Execute(context.Background(), json.RawMessage(`{"url":"`+server.URL+`"}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, want := range []string{"Fetched web content", server.URL, "# Agentu", "Built-in fetch reads HTML", "- Search", "- Fetch"} {
|
||||||
|
if !strings.Contains(out, want) {
|
||||||
|
t.Fatalf("output missing %q:\n%s", want, out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if strings.Contains(out, "hidden") {
|
||||||
|
t.Fatalf("script content should be ignored:\n%s", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuiltinFetchToolAcceptsURLsWithoutScheme(t *testing.T) {
|
||||||
|
args, err := parseWebFetchArgs(json.RawMessage(`{"url":"example.com/page"}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if args.URL != "https://example.com/page" {
|
||||||
|
t.Fatalf("url = %q", args.URL)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuiltinFetchToolFormatsFailedResults(t *testing.T) {
|
||||||
|
response := webFetchResponse{
|
||||||
|
FailedResults: []webFetchFailed{{URL: "https://example.com", Error: "blocked"}},
|
||||||
|
}
|
||||||
|
out := formatWebFetchResults("https://example.com", response)
|
||||||
|
for _, want := range []string{"No readable web content", "Fetch failed", "blocked"} {
|
||||||
|
if !strings.Contains(out, want) {
|
||||||
|
t.Fatalf("output missing %q:\n%s", want, out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
+742
-74
File diff suppressed because it is too large
Load Diff
+479
-23
@@ -2,22 +2,75 @@ package tui
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"os"
|
||||||
|
"os/exec"
|
||||||
|
"path/filepath"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
|
"agentu/internal/agent"
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
|
||||||
tea "github.com/charmbracelet/bubbletea"
|
tea "github.com/charmbracelet/bubbletea"
|
||||||
|
"github.com/charmbracelet/lipgloss"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestModelRendersChatShell(t *testing.T) {
|
func TestModelRendersChatShell(t *testing.T) {
|
||||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
||||||
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
rendered := next.(model).View()
|
m = next.(model)
|
||||||
|
rendered := m.View()
|
||||||
|
|
||||||
for _, want := range []string{"agentu", "test-model", "yolo tools", "light", "ctrl+j", "alt+enter"} {
|
for _, want := range []string{"test-model", "yolo tools", "light", "ctrl+j", "alt+enter", "ctrl+p/ctrl+n history"} {
|
||||||
if !strings.Contains(rendered, want) {
|
if !strings.Contains(rendered, want) {
|
||||||
t.Fatalf("rendered view missing %q:\n%s", want, rendered)
|
t.Fatalf("rendered view missing %q:\n%s", want, rendered)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if strings.Contains(rendered, "local agent") {
|
||||||
|
t.Fatalf("top header should not render:\n%s", rendered)
|
||||||
|
}
|
||||||
|
if _, ok := m.styles.App.GetBackground().(lipgloss.NoColor); !ok {
|
||||||
|
t.Fatalf("app should not paint a global background")
|
||||||
|
}
|
||||||
|
if _, ok := m.styles.Viewport.GetBackground().(lipgloss.NoColor); !ok {
|
||||||
|
t.Fatalf("viewport should not paint an assistant background")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInputTextDoesNotPaintBackground(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
for name, color := range map[string]any{
|
||||||
|
"input box": m.styles.InputBox.GetBackground(),
|
||||||
|
"focused base": m.input.FocusedStyle.Base.GetBackground(),
|
||||||
|
"focused cursor line": m.input.FocusedStyle.CursorLine.GetBackground(),
|
||||||
|
"blurred base": m.input.BlurredStyle.Base.GetBackground(),
|
||||||
|
"blurred cursor line": m.input.BlurredStyle.CursorLine.GetBackground(),
|
||||||
|
} {
|
||||||
|
if _, ok := color.(lipgloss.NoColor); !ok {
|
||||||
|
t.Fatalf("%s should not paint input text background", name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInputBoxIsRoomier(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
|
||||||
|
if got := m.styles.InputBox.GetVerticalFrameSize(); got != 2 {
|
||||||
|
t.Fatalf("input box vertical frame should stay compact, got %d", got)
|
||||||
|
}
|
||||||
|
if got := m.styles.InputBox.GetHorizontalFrameSize(); got < 6 {
|
||||||
|
t.Fatalf("input box horizontal frame should include roomier padding, got %d", got)
|
||||||
|
}
|
||||||
|
if got, want := m.input.Width(), 100-m.styles.InputBox.GetHorizontalFrameSize(); got != want {
|
||||||
|
t.Fatalf("input content width = %d, want %d", got, want)
|
||||||
|
}
|
||||||
|
rendered := m.inputView()
|
||||||
|
line := strings.Split(rendered, "\n")[0]
|
||||||
|
if got := lipgloss.Width(line); got != 100 {
|
||||||
|
t.Fatalf("input box rendered width = %d, want 100:\n%s", got, rendered)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestModelMetaRendersBelowInput(t *testing.T) {
|
func TestModelMetaRendersBelowInput(t *testing.T) {
|
||||||
@@ -35,6 +88,20 @@ func TestModelMetaRendersBelowInput(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestModelMetaRendersContextUsage(t *testing.T) {
|
||||||
|
assistant := agent.New(agent.Options{Model: "test-model", MaxContextTokens: 1000})
|
||||||
|
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: strings.Repeat("x", 800)}})
|
||||||
|
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 120, Height: 30})
|
||||||
|
rendered := next.(model).modelMetaView()
|
||||||
|
|
||||||
|
for _, want := range []string{"ctx ~", "/1k", "%"} {
|
||||||
|
if !strings.Contains(rendered, want) {
|
||||||
|
t.Fatalf("model meta missing %q:\n%s", want, rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestInitialViewDoesNotRenderTopNote(t *testing.T) {
|
func TestInitialViewDoesNotRenderTopNote(t *testing.T) {
|
||||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
@@ -56,43 +123,97 @@ func TestRunningStateRendersActivityIndicator(t *testing.T) {
|
|||||||
m.resize()
|
m.resize()
|
||||||
|
|
||||||
rendered := m.View()
|
rendered := m.View()
|
||||||
for _, want := range []string{"Working", "thinking...", "esc cancels"} {
|
for _, want := range []string{iconWorking, "Working", "thinking...", "esc cancels"} {
|
||||||
if !strings.Contains(rendered, want) {
|
if !strings.Contains(rendered, want) {
|
||||||
t.Fatalf("running view missing %q:\n%s", want, rendered)
|
t.Fatalf("running view missing %q:\n%s", want, rendered)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if backgroundPattern.MatchString(m.activityView()) {
|
||||||
|
t.Fatalf("activity view should not paint a background:\n%q", m.activityView())
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestUserMessageRendersOnLeft(t *testing.T) {
|
func TestUserMessageRendersFullWidthBubble(t *testing.T) {
|
||||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
m = next.(model)
|
m = next.(model)
|
||||||
|
|
||||||
rendered := m.messageView(roleUser, "hello", 80)
|
rendered := m.messageView(roleUser, "hi", 80)
|
||||||
if strings.HasPrefix(rendered, " ") {
|
if !strings.Contains(rendered, "hi") {
|
||||||
t.Fatalf("user message should not be right-aligned:\n%q", rendered)
|
|
||||||
}
|
|
||||||
if !strings.Contains(rendered, "You") || !strings.Contains(rendered, "hello") {
|
|
||||||
t.Fatalf("user message missing content:\n%s", rendered)
|
t.Fatalf("user message missing content:\n%s", rendered)
|
||||||
}
|
}
|
||||||
|
if strings.Contains(rendered, "You") {
|
||||||
|
t.Fatalf("user message should not render label:\n%s", rendered)
|
||||||
|
}
|
||||||
|
if !strings.Contains(rendered, " hi") || !strings.Contains(rendered, "hi ") {
|
||||||
|
t.Fatalf("user message should have horizontal padding:\n%q", rendered)
|
||||||
|
}
|
||||||
|
lines := strings.Split(rendered, "\n")
|
||||||
|
if len(lines) < 3 {
|
||||||
|
t.Fatalf("user message background should include vertical padding:\n%q", rendered)
|
||||||
|
}
|
||||||
|
for _, line := range lines {
|
||||||
|
if got := lipgloss.Width(line); got != 80 {
|
||||||
|
t.Fatalf("user message background should span full width, line width = %d:\n%q", got, rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
m.width = 80
|
||||||
|
m.messages = []message{{role: roleUser, content: "hi"}}
|
||||||
|
rendered = m.renderMessages()
|
||||||
|
for _, line := range strings.Split(rendered, "\n") {
|
||||||
|
if got := lipgloss.Width(line); got != 80 {
|
||||||
|
t.Fatalf("rendered user transcript should span full width, line width = %d:\n%q", got, rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestAssistantMessagesRenderMarkdown(t *testing.T) {
|
func TestAssistantMessagesRenderMarkdownWithoutChrome(t *testing.T) {
|
||||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
m = next.(model)
|
m = next.(model)
|
||||||
|
|
||||||
rendered := m.messageView(roleAssistant, "# Title\n\n- item", 80)
|
rendered := m.messageView(roleAssistant, "# Title\n\n- item", 80)
|
||||||
for _, want := range []string{"Agentu", "Title", "item"} {
|
for _, want := range []string{"Title", "item"} {
|
||||||
if !strings.Contains(rendered, want) {
|
if !strings.Contains(rendered, want) {
|
||||||
t.Fatalf("assistant markdown missing %q:\n%s", want, rendered)
|
t.Fatalf("assistant markdown missing %q:\n%s", want, rendered)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
for _, unwanted := range []string{"Agentu", "│", "┃", "|"} {
|
||||||
|
if strings.Contains(rendered, unwanted) {
|
||||||
|
t.Fatalf("assistant message should not render chrome %q:\n%s", unwanted, rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
if strings.Contains(rendered, "# Title") {
|
if strings.Contains(rendered, "# Title") {
|
||||||
t.Fatalf("assistant markdown heading was not rendered:\n%s", rendered)
|
t.Fatalf("assistant markdown heading was not rendered:\n%s", rendered)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestAssistantMessageRendersPaddedBodyWithoutBackground(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
if _, ok := m.styles.AssistantMsg.GetBackground().(lipgloss.NoColor); !ok {
|
||||||
|
t.Fatalf("assistant message should not set a background color")
|
||||||
|
}
|
||||||
|
|
||||||
|
rendered := m.messageView(roleAssistant, "hello", 80)
|
||||||
|
if !strings.Contains(rendered, "hello") {
|
||||||
|
t.Fatalf("assistant message missing content:\n%s", rendered)
|
||||||
|
}
|
||||||
|
lines := strings.Split(rendered, "\n")
|
||||||
|
if len(lines) < 3 {
|
||||||
|
t.Fatalf("assistant message should include vertical padding:\n%q", rendered)
|
||||||
|
}
|
||||||
|
if got := lipgloss.Width(rendered); got <= len("hello") {
|
||||||
|
t.Fatalf("assistant message should be wider than text, width = %d:\n%q", got, rendered)
|
||||||
|
}
|
||||||
|
for _, line := range lines {
|
||||||
|
if got := lipgloss.Width(line); got >= 80 {
|
||||||
|
t.Fatalf("assistant message should not fill the row, line width = %d:\n%q", got, rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestUserMessagesDoNotRenderMarkdown(t *testing.T) {
|
func TestUserMessagesDoNotRenderMarkdown(t *testing.T) {
|
||||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
@@ -106,7 +227,7 @@ func TestUserMessagesDoNotRenderMarkdown(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestMessageRoleLabels(t *testing.T) {
|
func TestMessageRoleLabelsDoNotLabelUserOrAssistant(t *testing.T) {
|
||||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
m = next.(model)
|
m = next.(model)
|
||||||
@@ -116,9 +237,14 @@ func TestMessageRoleLabels(t *testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
rendered := m.renderMessages()
|
rendered := m.renderMessages()
|
||||||
for _, want := range []string{"You", "Agentu"} {
|
for _, want := range []string{"hello", "hi"} {
|
||||||
if !strings.Contains(rendered, want) {
|
if !strings.Contains(rendered, want) {
|
||||||
t.Fatalf("messages missing %q:\n%s", want, rendered)
|
t.Fatalf("messages missing content %q:\n%s", want, rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, unwanted := range []string{"You", "Agentu"} {
|
||||||
|
if strings.Contains(rendered, unwanted) {
|
||||||
|
t.Fatalf("message should not render label %q:\n%s", unwanted, rendered)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -153,19 +279,205 @@ func TestInputStartsAtOneLineAndGrows(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSlashCommandSuggestions(t *testing.T) {
|
func TestCommandPaletteFuzzyMatches(t *testing.T) {
|
||||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
||||||
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
m = next.(model)
|
m = next.(model)
|
||||||
m.input.SetValue("/")
|
m.input.SetValue("/mdl")
|
||||||
m.afterInputChanged()
|
m.afterInputChanged()
|
||||||
|
|
||||||
rendered := m.View()
|
plain := stripANSI(m.View())
|
||||||
for _, want := range []string{"/clear reset context", "/exit quit", "/model model/provider/thinking"} {
|
for _, want := range []string{iconCommand + " Commands", "/model model/provider/thinking"} {
|
||||||
if !strings.Contains(rendered, want) {
|
if !strings.Contains(plain, want) {
|
||||||
t.Fatalf("rendered suggestions missing %q:\n%s", want, rendered)
|
t.Fatalf("rendered palette missing %q:\n%s", want, plain)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
matches := commandMatches("m")
|
||||||
|
modelIndex := -1
|
||||||
|
renameIndex := -1
|
||||||
|
for i, match := range matches {
|
||||||
|
switch match.command.Name {
|
||||||
|
case "/model":
|
||||||
|
modelIndex = i
|
||||||
|
case "/rename":
|
||||||
|
renameIndex = i
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if modelIndex < 0 || renameIndex < 0 || modelIndex > renameIndex {
|
||||||
|
t.Fatalf("/model should sort before /rename for /m query: %#v", matches)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInputHistoryNavigation(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.rememberInput("first")
|
||||||
|
m.rememberInput("second")
|
||||||
|
|
||||||
|
next, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlP})
|
||||||
|
m = next.(model)
|
||||||
|
if got := m.input.Value(); got != "second" {
|
||||||
|
t.Fatalf("after ctrl+p input = %q", got)
|
||||||
|
}
|
||||||
|
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlP})
|
||||||
|
m = next.(model)
|
||||||
|
if got := m.input.Value(); got != "first" {
|
||||||
|
t.Fatalf("after second ctrl+p input = %q", got)
|
||||||
|
}
|
||||||
|
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlN})
|
||||||
|
m = next.(model)
|
||||||
|
if got := m.input.Value(); got != "second" {
|
||||||
|
t.Fatalf("after ctrl+n input = %q", got)
|
||||||
|
}
|
||||||
|
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlN})
|
||||||
|
m = next.(model)
|
||||||
|
if got := m.input.Value(); got != "" {
|
||||||
|
t.Fatalf("after second ctrl+n input = %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestToolMessageRendersStructuredCard(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.messages = append(m.messages, message{role: roleTool, content: `file_read {"path":"README.md"}`})
|
||||||
|
|
||||||
|
plain := stripANSI(m.renderMessages())
|
||||||
|
for _, want := range []string{iconTool + " file_read", "README.md"} {
|
||||||
|
if !strings.Contains(plain, want) {
|
||||||
|
t.Fatalf("tool card missing %q:\n%s", want, plain)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestToolMessageRendersOutputPreview(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.messages = append(m.messages, message{role: roleTool, content: "file_read {\"path\":\"README.md\"}\n[output]\nfirst line\nsecond line"})
|
||||||
|
|
||||||
|
plain := stripANSI(m.renderMessages())
|
||||||
|
for _, want := range []string{iconTool + " file_read", "README.md", "output", "first line", "second line"} {
|
||||||
|
if !strings.Contains(plain, want) {
|
||||||
|
t.Fatalf("tool card missing %q:\n%s", want, plain)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestToolLogOutputUpdatesExistingToolCard(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.messages = append(m.messages, message{role: roleTool, content: `file_read {"path":"README.md"}`})
|
||||||
|
|
||||||
|
updated := m.mergeToolLog("file_read {\"path\":\"README.md\"}\n[output]\nhello")
|
||||||
|
if !updated {
|
||||||
|
t.Fatal("expected mergeToolLog to return true")
|
||||||
|
}
|
||||||
|
if len(m.messages) != 1 {
|
||||||
|
t.Fatalf("messages count = %d, want 1", len(m.messages))
|
||||||
|
}
|
||||||
|
plain := stripANSI(m.renderMessages())
|
||||||
|
if !strings.Contains(plain, "hello") {
|
||||||
|
t.Fatalf("rendered output missing 'hello':\n%s", plain)
|
||||||
|
}
|
||||||
|
|
||||||
|
status := toolStatusText("file_read {\"path\":\"README.md\"}\n[output]\nhello world")
|
||||||
|
if strings.Contains(status, "hello") {
|
||||||
|
t.Fatalf("status should not contain output: %q", status)
|
||||||
|
}
|
||||||
|
if !strings.Contains(status, "file_read") {
|
||||||
|
t.Fatalf("status missing tool name: %q", status)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type tuiCompactProvider struct{}
|
||||||
|
|
||||||
|
func (tuiCompactProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
return emit(llm.StreamEvent{Content: "summary from tui"})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCompactCommand(t *testing.T) {
|
||||||
|
assistant := agent.New(agent.Options{Provider: tuiCompactProvider{}, Model: "test", SystemPrompt: "system"})
|
||||||
|
assistant.SetMessages([]llm.Message{
|
||||||
|
{Role: llm.RoleSystem, Content: "system"},
|
||||||
|
{Role: llm.RoleUser, Content: "old question 1"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "old answer 1"},
|
||||||
|
{Role: llm.RoleUser, Content: "old question 2"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "old answer 2"},
|
||||||
|
{Role: llm.RoleUser, Content: "recent question 1"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "recent answer 1"},
|
||||||
|
{Role: llm.RoleUser, Content: "recent question 2"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "recent answer 2"},
|
||||||
|
})
|
||||||
|
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
|
||||||
|
|
||||||
|
out, err := m.handleLocalCommand("/compact")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if out != "Context compacted." {
|
||||||
|
t.Fatalf("output = %q", out)
|
||||||
|
}
|
||||||
|
if !strings.Contains(messageContents(assistant.Messages()), "[context summary]") {
|
||||||
|
t.Fatalf("messages missing summary: %#v", assistant.Messages())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func messageContents(messages []llm.Message) string {
|
||||||
|
var b strings.Builder
|
||||||
|
for _, msg := range messages {
|
||||||
|
b.WriteString(msg.Content)
|
||||||
|
b.WriteByte('\n')
|
||||||
|
}
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWorkflowStatusDiffAndTestCommands(t *testing.T) {
|
||||||
|
if _, err := exec.LookPath("git"); err != nil {
|
||||||
|
t.Skip("git not installed")
|
||||||
|
}
|
||||||
|
dir := t.TempDir()
|
||||||
|
runInDir(t, dir, "git", "init")
|
||||||
|
runInDir(t, dir, "git", "config", "user.email", "agentu@example.test")
|
||||||
|
runInDir(t, dir, "git", "config", "user.name", "agentu")
|
||||||
|
if err := os.WriteFile(filepath.Join(dir, "go.mod"), []byte("module example.com/p1\n\ngo 1.26\n"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(filepath.Join(dir, "main_test.go"), []byte("package p1\n\nimport \"testing\"\n\nfunc TestOK(t *testing.T) {}\n"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
runInDir(t, dir, "git", "add", ".")
|
||||||
|
runInDir(t, dir, "git", "commit", "-m", "init")
|
||||||
|
if err := os.WriteFile(filepath.Join(dir, "main_test.go"), []byte("package p1\n\nimport \"testing\"\n\nfunc TestOK(t *testing.T) { t.Log(\"changed\") }\n"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", WorkingDir: dir})
|
||||||
|
status, err := m.handleStatusCommand(nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(status, "main_test.go") {
|
||||||
|
t.Fatalf("status missing changed file:\n%s", status)
|
||||||
|
}
|
||||||
|
diff, err := m.handleDiffCommand(nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(diff, "changed") {
|
||||||
|
t.Fatalf("diff missing changed content:\n%s", diff)
|
||||||
|
}
|
||||||
|
out, err := m.handleTestCommand(nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "$ go test ./...") {
|
||||||
|
t.Fatalf("test output missing default command:\n%s", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func runInDir(t *testing.T, dir string, name string, args ...string) {
|
||||||
|
t.Helper()
|
||||||
|
cmd := exec.Command(name, args...)
|
||||||
|
cmd.Dir = dir
|
||||||
|
if output, err := cmd.CombinedOutput(); err != nil {
|
||||||
|
t.Fatalf("%s %v failed: %v\n%s", name, args, err, output)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestModelCommand(t *testing.T) {
|
func TestModelCommand(t *testing.T) {
|
||||||
@@ -197,9 +509,11 @@ func TestModelThinkingCommand(t *testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type fakeModelManager struct {
|
type fakeModelManager struct {
|
||||||
provider string
|
provider string
|
||||||
model string
|
model string
|
||||||
thinking string
|
thinking string
|
||||||
|
sessionName string
|
||||||
|
saved bool
|
||||||
}
|
}
|
||||||
|
|
||||||
func (f *fakeModelManager) CurrentProvider() string {
|
func (f *fakeModelManager) CurrentProvider() string {
|
||||||
@@ -217,6 +531,17 @@ func (f *fakeModelManager) CurrentThinking() string {
|
|||||||
return f.thinking
|
return f.thinking
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (f *fakeModelManager) CurrentSessionID() string {
|
||||||
|
return "test-id"
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeModelManager) CurrentSessionName() string {
|
||||||
|
if f.sessionName == "" {
|
||||||
|
return "test-session"
|
||||||
|
}
|
||||||
|
return f.sessionName
|
||||||
|
}
|
||||||
|
|
||||||
func (f *fakeModelManager) Info() string {
|
func (f *fakeModelManager) Info() string {
|
||||||
return "provider: " + f.CurrentProvider() + "\nmodel: " + f.model + "\nthinking: " + f.thinking
|
return "provider: " + f.CurrentProvider() + "\nmodel: " + f.model + "\nthinking: " + f.thinking
|
||||||
}
|
}
|
||||||
@@ -236,6 +561,28 @@ func (f *fakeModelManager) SetThinking(level string) (string, error) {
|
|||||||
return f.Info(), nil
|
return f.Info(), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (f *fakeModelManager) Save() error {
|
||||||
|
f.saved = true
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeModelManager) Resume(id string) error {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeModelManager) Rename(name string) (string, error) {
|
||||||
|
f.sessionName = name
|
||||||
|
return "Renamed.", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeModelManager) Sessions() (string, error) {
|
||||||
|
return "No sessions.", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *fakeModelManager) NewSession() (string, error) {
|
||||||
|
return "New session.", nil
|
||||||
|
}
|
||||||
|
|
||||||
func TestParseThemeMode(t *testing.T) {
|
func TestParseThemeMode(t *testing.T) {
|
||||||
for input, want := range map[string]ThemeMode{
|
for input, want := range map[string]ThemeMode{
|
||||||
"": ThemeLight,
|
"": ThemeLight,
|
||||||
@@ -257,6 +604,115 @@ func TestParseThemeMode(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func requireStyleColor(t *testing.T, name string, got lipgloss.TerminalColor, want string) {
|
||||||
|
t.Helper()
|
||||||
|
color, ok := got.(lipgloss.Color)
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("%s = %T, want lipgloss.Color", name, got)
|
||||||
|
}
|
||||||
|
if string(color) != want {
|
||||||
|
t.Fatalf("%s = %q, want %q", name, string(color), want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func requireNoStyleColor(t *testing.T, name string, got lipgloss.TerminalColor) {
|
||||||
|
t.Helper()
|
||||||
|
if _, ok := got.(lipgloss.NoColor); !ok {
|
||||||
|
t.Fatalf("%s should not paint a color, got %T", name, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDarkThemePaintsAppBackground(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", ThemeMode: ThemeDark})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
rendered := stripANSI(m.View())
|
||||||
|
if !strings.Contains(rendered, "theme dark") {
|
||||||
|
t.Fatalf("rendered view missing 'theme dark':\n%s", rendered)
|
||||||
|
}
|
||||||
|
if m.theme.Mode != ThemeDark {
|
||||||
|
t.Fatalf("theme.Mode = %q, want %q", m.theme.Mode, ThemeDark)
|
||||||
|
}
|
||||||
|
want := darkTheme()
|
||||||
|
requireStyleColor(t, "App.Background", m.styles.App.GetBackground(), want.Background)
|
||||||
|
requireStyleColor(t, "App.Foreground", m.styles.App.GetForeground(), want.Text)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDarkThemeKeepsComponentBackgroundsLightweight(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", ThemeMode: ThemeDark})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
requireNoStyleColor(t, "Viewport.Background", m.styles.Viewport.GetBackground())
|
||||||
|
requireNoStyleColor(t, "InputBox.Background", m.styles.InputBox.GetBackground())
|
||||||
|
requireNoStyleColor(t, "AssistantMsg.Background", m.styles.AssistantMsg.GetBackground())
|
||||||
|
requireNoStyleColor(t, "Activity.Background", m.styles.Activity.GetBackground())
|
||||||
|
m.running = true
|
||||||
|
m.status = "thinking..."
|
||||||
|
m.resize()
|
||||||
|
if backgroundPattern.MatchString(m.activityView()) {
|
||||||
|
t.Fatalf("activity view should not paint a background:\n%q", m.activityView())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDarkThemeStylesUseDarkPalette(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", ThemeMode: ThemeDark})
|
||||||
|
want := darkTheme()
|
||||||
|
requireStyleColor(t, "Header.Foreground", m.styles.Header.GetForeground(), want.Title)
|
||||||
|
requireStyleColor(t, "Muted.Foreground", m.styles.Muted.GetForeground(), want.Muted)
|
||||||
|
requireStyleColor(t, "Pill.Background", m.styles.Pill.GetBackground(), want.SurfaceAlt)
|
||||||
|
requireStyleColor(t, "input.FocusedStyle.Base.Foreground", m.input.FocusedStyle.Base.GetForeground(), want.Text)
|
||||||
|
requireStyleColor(t, "input.FocusedStyle.Placeholder.Foreground", m.input.FocusedStyle.Placeholder.GetForeground(), want.Muted)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestModelCommandShowsDarkTheme(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", ThemeMode: ThemeDark})
|
||||||
|
m.input.SetValue("/model")
|
||||||
|
next, _ := m.submit()
|
||||||
|
rendered := next.(model).renderMessages()
|
||||||
|
if !strings.Contains(rendered, "theme: dark") {
|
||||||
|
t.Fatalf("model command output missing 'theme: dark':\n%s", rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestThemeSlashCommandSwitchesAtRuntime(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
if m.theme.Mode != ThemeLight {
|
||||||
|
t.Fatalf("initial theme = %q, want light", m.theme.Mode)
|
||||||
|
}
|
||||||
|
requireNoStyleColor(t, "initial App.Background", m.styles.App.GetBackground())
|
||||||
|
|
||||||
|
// Switch to dark
|
||||||
|
m.input.SetValue("/theme dark")
|
||||||
|
next, _ := m.submit()
|
||||||
|
m = next.(model)
|
||||||
|
if m.theme.Mode != ThemeDark {
|
||||||
|
t.Fatalf("after /theme dark: Mode = %q, want dark", m.theme.Mode)
|
||||||
|
}
|
||||||
|
requireStyleColor(t, "dark App.Background", m.styles.App.GetBackground(), darkTheme().Background)
|
||||||
|
last := m.messages[len(m.messages)-1]
|
||||||
|
if !strings.Contains(last.content, "theme: dark") {
|
||||||
|
t.Fatalf("expected confirmation message, got %q", last.content)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Switch back to light
|
||||||
|
m.input.SetValue("/theme light")
|
||||||
|
next, _ = m.submit()
|
||||||
|
m = next.(model)
|
||||||
|
if m.theme.Mode != ThemeLight {
|
||||||
|
t.Fatalf("after /theme light: Mode = %q, want light", m.theme.Mode)
|
||||||
|
}
|
||||||
|
requireNoStyleColor(t, "restored App.Background", m.styles.App.GetBackground())
|
||||||
|
|
||||||
|
// No arg shows usage
|
||||||
|
m.input.SetValue("/theme")
|
||||||
|
next, _ = m.submit()
|
||||||
|
m = next.(model)
|
||||||
|
last = m.messages[len(m.messages)-1]
|
||||||
|
if !strings.Contains(last.content, "usage: /theme") {
|
||||||
|
t.Fatalf("expected usage hint, got %q", last.content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestCleanToolLog(t *testing.T) {
|
func TestCleanToolLog(t *testing.T) {
|
||||||
got := cleanToolLog("\n[tool] shell_run {\"command\":\"pwd\"}\n")
|
got := cleanToolLog("\n[tool] shell_run {\"command\":\"pwd\"}\n")
|
||||||
want := `shell_run {"command":"pwd"}`
|
want := `shell_run {"command":"pwd"}`
|
||||||
|
|||||||
@@ -2,6 +2,8 @@ package tui
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"strings"
|
"strings"
|
||||||
|
"unicode"
|
||||||
|
"unicode/utf8"
|
||||||
|
|
||||||
"github.com/charmbracelet/glamour"
|
"github.com/charmbracelet/glamour"
|
||||||
"github.com/charmbracelet/glamour/ansi"
|
"github.com/charmbracelet/glamour/ansi"
|
||||||
@@ -24,7 +26,7 @@ func renderMarkdown(content string, width int, th theme) string {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return content
|
return content
|
||||||
}
|
}
|
||||||
return strings.TrimRight(rendered, "\n")
|
return trimRenderedMarkdown(rendered)
|
||||||
}
|
}
|
||||||
|
|
||||||
func markdownStyle(th theme) ansi.StyleConfig {
|
func markdownStyle(th theme) ansi.StyleConfig {
|
||||||
@@ -32,6 +34,7 @@ func markdownStyle(th theme) ansi.StyleConfig {
|
|||||||
if th.Mode == ThemeDark {
|
if th.Mode == ThemeDark {
|
||||||
cfg = glamstyles.DarkStyleConfig
|
cfg = glamstyles.DarkStyleConfig
|
||||||
}
|
}
|
||||||
|
clearMarkdownBackgrounds(&cfg)
|
||||||
|
|
||||||
zero := uint(0)
|
zero := uint(0)
|
||||||
cfg.Document.Margin = &zero
|
cfg.Document.Margin = &zero
|
||||||
@@ -51,7 +54,7 @@ func markdownStyle(th theme) ansi.StyleConfig {
|
|||||||
cfg.LinkText.Color = stringPtr(th.User)
|
cfg.LinkText.Color = stringPtr(th.User)
|
||||||
|
|
||||||
cfg.Code.Color = stringPtr(th.Tool)
|
cfg.Code.Color = stringPtr(th.Tool)
|
||||||
cfg.Code.BackgroundColor = stringPtr(th.SurfaceAlt)
|
cfg.Code.BackgroundColor = nil
|
||||||
cfg.CodeBlock.Color = stringPtr(th.Text)
|
cfg.CodeBlock.Color = stringPtr(th.Text)
|
||||||
cfg.CodeBlock.BackgroundColor = nil
|
cfg.CodeBlock.BackgroundColor = nil
|
||||||
cfg.CodeBlock.Margin = &zero
|
cfg.CodeBlock.Margin = &zero
|
||||||
@@ -59,6 +62,93 @@ func markdownStyle(th theme) ansi.StyleConfig {
|
|||||||
return cfg
|
return cfg
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func clearMarkdownBackgrounds(cfg *ansi.StyleConfig) {
|
||||||
|
cfg.Document.BackgroundColor = nil
|
||||||
|
cfg.BlockQuote.BackgroundColor = nil
|
||||||
|
cfg.Paragraph.BackgroundColor = nil
|
||||||
|
cfg.Heading.BackgroundColor = nil
|
||||||
|
cfg.H1.BackgroundColor = nil
|
||||||
|
cfg.H2.BackgroundColor = nil
|
||||||
|
cfg.H3.BackgroundColor = nil
|
||||||
|
cfg.H4.BackgroundColor = nil
|
||||||
|
cfg.H5.BackgroundColor = nil
|
||||||
|
cfg.H6.BackgroundColor = nil
|
||||||
|
cfg.Text.BackgroundColor = nil
|
||||||
|
cfg.Strikethrough.BackgroundColor = nil
|
||||||
|
cfg.Emph.BackgroundColor = nil
|
||||||
|
cfg.Strong.BackgroundColor = nil
|
||||||
|
cfg.HorizontalRule.BackgroundColor = nil
|
||||||
|
cfg.Item.BackgroundColor = nil
|
||||||
|
cfg.Enumeration.BackgroundColor = nil
|
||||||
|
cfg.Task.BackgroundColor = nil
|
||||||
|
cfg.Link.BackgroundColor = nil
|
||||||
|
cfg.LinkText.BackgroundColor = nil
|
||||||
|
cfg.Image.BackgroundColor = nil
|
||||||
|
cfg.ImageText.BackgroundColor = nil
|
||||||
|
cfg.Code.BackgroundColor = nil
|
||||||
|
cfg.CodeBlock.BackgroundColor = nil
|
||||||
|
cfg.Table.BackgroundColor = nil
|
||||||
|
cfg.DefinitionList.BackgroundColor = nil
|
||||||
|
cfg.DefinitionTerm.BackgroundColor = nil
|
||||||
|
cfg.DefinitionDescription.BackgroundColor = nil
|
||||||
|
cfg.HTMLBlock.BackgroundColor = nil
|
||||||
|
cfg.HTMLSpan.BackgroundColor = nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func trimRenderedMarkdown(value string) string {
|
||||||
|
lines := strings.Split(strings.TrimRight(value, "\n"), "\n")
|
||||||
|
for i, line := range lines {
|
||||||
|
lines[i] = trimStyledTrailingSpaces(line)
|
||||||
|
}
|
||||||
|
return strings.Join(lines, "\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
func trimStyledTrailingSpaces(line string) string {
|
||||||
|
cut := 0
|
||||||
|
seenNonSpace := false
|
||||||
|
seenWhitespaceAfterLastNonSpace := false
|
||||||
|
for i := 0; i < len(line); {
|
||||||
|
if end, ok := ansiSequenceEnd(line, i); ok {
|
||||||
|
if seenNonSpace && !seenWhitespaceAfterLastNonSpace {
|
||||||
|
cut = end
|
||||||
|
}
|
||||||
|
i = end
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
r, size := utf8.DecodeRuneInString(line[i:])
|
||||||
|
if r == utf8.RuneError && size == 0 {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
end := i + size
|
||||||
|
if unicode.IsSpace(r) {
|
||||||
|
if seenNonSpace {
|
||||||
|
seenWhitespaceAfterLastNonSpace = true
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
seenNonSpace = true
|
||||||
|
seenWhitespaceAfterLastNonSpace = false
|
||||||
|
cut = end
|
||||||
|
}
|
||||||
|
i = end
|
||||||
|
}
|
||||||
|
if cut == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return line[:cut]
|
||||||
|
}
|
||||||
|
|
||||||
|
func ansiSequenceEnd(value string, start int) (int, bool) {
|
||||||
|
if start+2 >= len(value) || value[start] != '\x1b' || value[start+1] != '[' {
|
||||||
|
return 0, false
|
||||||
|
}
|
||||||
|
for i := start + 2; i < len(value); i++ {
|
||||||
|
if value[i] >= 0x40 && value[i] <= 0x7e {
|
||||||
|
return i + 1, true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return 0, false
|
||||||
|
}
|
||||||
func stringPtr(value string) *string {
|
func stringPtr(value string) *string {
|
||||||
return &value
|
return &value
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -37,6 +37,26 @@ func TestRenderMarkdownSupportsThemesAndNarrowWidth(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestRenderMarkdownDoesNotEmitBackgroundColors(t *testing.T) {
|
||||||
|
content := "# Title\n\n`inline`\n\n```go\nfmt.Println(\"x\")\n```"
|
||||||
|
for _, mode := range []ThemeMode{ThemeLight, ThemeDark} {
|
||||||
|
rendered := renderMarkdown(content, 80, themeForMode(mode))
|
||||||
|
if backgroundPattern.MatchString(rendered) {
|
||||||
|
t.Fatalf("markdown output for %s should not emit background colors:\n%q", mode, rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRenderMarkdownTrimsGlamourTrailingSpaces(t *testing.T) {
|
||||||
|
rendered := renderMarkdown("hello", 80, themeForMode(ThemeLight))
|
||||||
|
plain := stripANSI(rendered)
|
||||||
|
if plain != "hello" {
|
||||||
|
t.Fatalf("markdown should not pad a plain assistant line:\n%q", plain)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
var backgroundPattern = regexp.MustCompile(`\x1b\[[0-9;:]*48[;:][0-9;:]*m`)
|
||||||
|
|
||||||
var ansiPattern = regexp.MustCompile(`\x1b\[[0-9;:]*[A-Za-z]`)
|
var ansiPattern = regexp.MustCompile(`\x1b\[[0-9;:]*[A-Za-z]`)
|
||||||
|
|
||||||
func stripANSI(value string) string {
|
func stripANSI(value string) string {
|
||||||
|
|||||||
+127
-62
@@ -15,6 +15,17 @@ const (
|
|||||||
ThemeDark ThemeMode = "dark"
|
ThemeDark ThemeMode = "dark"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
iconReady = "✓"
|
||||||
|
iconWorking = "◆"
|
||||||
|
iconTool = "▣"
|
||||||
|
iconError = "✕"
|
||||||
|
iconCommand = "⌘"
|
||||||
|
iconSession = "◌"
|
||||||
|
iconContext = "◔"
|
||||||
|
iconModel = "◉"
|
||||||
|
)
|
||||||
|
|
||||||
func ParseThemeMode(value string) (ThemeMode, error) {
|
func ParseThemeMode(value string) (ThemeMode, error) {
|
||||||
switch ThemeMode(strings.ToLower(strings.TrimSpace(value))) {
|
switch ThemeMode(strings.ToLower(strings.TrimSpace(value))) {
|
||||||
case "", ThemeLight:
|
case "", ThemeLight:
|
||||||
@@ -45,26 +56,35 @@ type theme struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type styles struct {
|
type styles struct {
|
||||||
App lipgloss.Style
|
App lipgloss.Style
|
||||||
Header lipgloss.Style
|
Muted lipgloss.Style
|
||||||
Title lipgloss.Style
|
Activity lipgloss.Style
|
||||||
Muted lipgloss.Style
|
Header lipgloss.Style
|
||||||
Activity lipgloss.Style
|
ModelMeta lipgloss.Style
|
||||||
ModelMeta lipgloss.Style
|
Footer lipgloss.Style
|
||||||
Footer lipgloss.Style
|
Viewport lipgloss.Style
|
||||||
Viewport lipgloss.Style
|
InputBox lipgloss.Style
|
||||||
InputBox lipgloss.Style
|
CommandHint lipgloss.Style
|
||||||
CommandHint lipgloss.Style
|
StatusReady lipgloss.Style
|
||||||
UserLabel lipgloss.Style
|
StatusWarn lipgloss.Style
|
||||||
AssistantLabel lipgloss.Style
|
StatusError lipgloss.Style
|
||||||
ToolLabel lipgloss.Style
|
StatusInfo lipgloss.Style
|
||||||
ErrorLabel lipgloss.Style
|
Pill lipgloss.Style
|
||||||
SystemLabel lipgloss.Style
|
PillAccent lipgloss.Style
|
||||||
UserMessage lipgloss.Style
|
Palette lipgloss.Style
|
||||||
AssistantMsg lipgloss.Style
|
PaletteMatch lipgloss.Style
|
||||||
ToolMessage lipgloss.Style
|
ToolName lipgloss.Style
|
||||||
ErrorMessage lipgloss.Style
|
ToolKey lipgloss.Style
|
||||||
SystemMessage lipgloss.Style
|
ToolValue lipgloss.Style
|
||||||
|
UserLabel lipgloss.Style
|
||||||
|
ToolLabel lipgloss.Style
|
||||||
|
ErrorLabel lipgloss.Style
|
||||||
|
SystemLabel lipgloss.Style
|
||||||
|
UserMessage lipgloss.Style
|
||||||
|
AssistantMsg lipgloss.Style
|
||||||
|
ToolMessage lipgloss.Style
|
||||||
|
ErrorMessage lipgloss.Style
|
||||||
|
SystemMessage lipgloss.Style
|
||||||
}
|
}
|
||||||
|
|
||||||
func themeForMode(mode ThemeMode) theme {
|
func themeForMode(mode ThemeMode) theme {
|
||||||
@@ -110,67 +130,103 @@ func darkTheme() theme {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func appStyle(t theme) lipgloss.Style {
|
||||||
|
style := lipgloss.NewStyle().Foreground(lipgloss.Color(t.Text))
|
||||||
|
if t.Mode == ThemeDark {
|
||||||
|
style = style.Background(lipgloss.Color(t.Background))
|
||||||
|
}
|
||||||
|
return style
|
||||||
|
}
|
||||||
|
|
||||||
func (t theme) styles() styles {
|
func (t theme) styles() styles {
|
||||||
return styles{
|
return styles{
|
||||||
App: lipgloss.NewStyle().
|
App: appStyle(t),
|
||||||
Foreground(lipgloss.Color(t.Text)).
|
|
||||||
Background(lipgloss.Color(t.Background)),
|
|
||||||
|
|
||||||
Header: lipgloss.NewStyle().
|
|
||||||
Foreground(lipgloss.Color(t.Text)).
|
|
||||||
Background(lipgloss.Color(t.SurfaceAlt)).
|
|
||||||
Padding(0, 1),
|
|
||||||
|
|
||||||
Title: lipgloss.NewStyle().
|
|
||||||
Bold(true).
|
|
||||||
Foreground(lipgloss.Color(t.Title)),
|
|
||||||
|
|
||||||
Muted: lipgloss.NewStyle().
|
Muted: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Muted)),
|
Foreground(lipgloss.Color(t.Muted)),
|
||||||
|
|
||||||
Activity: lipgloss.NewStyle().
|
Activity: lipgloss.NewStyle().
|
||||||
Bold(true).
|
Foreground(lipgloss.Color(t.Muted)).
|
||||||
Foreground(lipgloss.Color(t.Background)).
|
|
||||||
Background(lipgloss.Color(t.Title)).
|
|
||||||
Padding(0, 1),
|
Padding(0, 1),
|
||||||
|
|
||||||
|
Header: lipgloss.NewStyle().
|
||||||
|
Bold(true).
|
||||||
|
Foreground(lipgloss.Color(t.Title)),
|
||||||
|
|
||||||
ModelMeta: lipgloss.NewStyle().
|
ModelMeta: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Text)).
|
Foreground(lipgloss.Color(t.Muted)).
|
||||||
Background(lipgloss.Color(t.SurfaceAlt)).
|
|
||||||
Padding(0, 1),
|
Padding(0, 1),
|
||||||
|
|
||||||
Footer: lipgloss.NewStyle().
|
Footer: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Muted)).
|
Foreground(lipgloss.Color(t.Muted)).
|
||||||
Background(lipgloss.Color(t.Background)).
|
|
||||||
Padding(0, 1),
|
Padding(0, 1),
|
||||||
|
|
||||||
Viewport: lipgloss.NewStyle().
|
Viewport: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Text)).
|
Foreground(lipgloss.Color(t.Text)),
|
||||||
Background(lipgloss.Color(t.Background)),
|
|
||||||
|
|
||||||
InputBox: lipgloss.NewStyle().
|
InputBox: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Text)).
|
Foreground(lipgloss.Color(t.Text)).
|
||||||
Background(lipgloss.Color(t.Surface)).
|
Border(lipgloss.ThickBorder()).
|
||||||
Border(lipgloss.NormalBorder()).
|
|
||||||
BorderForeground(lipgloss.Color(t.Border)).
|
BorderForeground(lipgloss.Color(t.Border)).
|
||||||
Padding(0, 1),
|
Padding(0, 2),
|
||||||
|
|
||||||
CommandHint: lipgloss.NewStyle().
|
CommandHint: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Muted)).
|
Foreground(lipgloss.Color(t.Muted)).
|
||||||
|
Padding(0, 1),
|
||||||
|
|
||||||
|
StatusReady: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Assistant)),
|
||||||
|
|
||||||
|
StatusWarn: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Tool)),
|
||||||
|
|
||||||
|
StatusError: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Error)),
|
||||||
|
|
||||||
|
StatusInfo: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.User)),
|
||||||
|
|
||||||
|
Pill: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Text)).
|
||||||
Background(lipgloss.Color(t.SurfaceAlt)).
|
Background(lipgloss.Color(t.SurfaceAlt)).
|
||||||
Padding(0, 1),
|
Padding(0, 1),
|
||||||
|
|
||||||
UserLabel: labelStyle(t.User),
|
PillAccent: lipgloss.NewStyle().
|
||||||
AssistantLabel: labelStyle(t.Assistant),
|
Bold(true).
|
||||||
ToolLabel: labelStyle(t.Tool),
|
Foreground(lipgloss.Color(t.Title)).
|
||||||
ErrorLabel: labelStyle(t.Error),
|
Background(lipgloss.Color(t.SurfaceAlt)).
|
||||||
SystemLabel: labelStyle(t.System),
|
Padding(0, 1),
|
||||||
|
|
||||||
UserMessage: messageStyle(t.Text, t.User),
|
Palette: lipgloss.NewStyle().
|
||||||
AssistantMsg: messageStyle(t.Text, t.Assistant),
|
Foreground(lipgloss.Color(t.Text)).
|
||||||
ToolMessage: messageStyle(t.Text, t.Tool),
|
Border(lipgloss.ThickBorder()).
|
||||||
ErrorMessage: messageStyle(t.Error, t.Error),
|
BorderForeground(lipgloss.Color(t.Border)).
|
||||||
SystemMessage: lipgloss.NewStyle().Foreground(lipgloss.Color(t.System)).PaddingLeft(1),
|
Padding(0, 2),
|
||||||
|
|
||||||
|
PaletteMatch: lipgloss.NewStyle().
|
||||||
|
Bold(true).
|
||||||
|
Foreground(lipgloss.Color(t.User)),
|
||||||
|
|
||||||
|
ToolName: lipgloss.NewStyle().
|
||||||
|
Bold(true).
|
||||||
|
Foreground(lipgloss.Color(t.Tool)),
|
||||||
|
|
||||||
|
ToolKey: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Muted)),
|
||||||
|
|
||||||
|
ToolValue: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Text)),
|
||||||
|
|
||||||
|
UserLabel: labelStyle(t.User),
|
||||||
|
ToolLabel: labelStyle(t.Tool),
|
||||||
|
ErrorLabel: labelStyle(t.Error),
|
||||||
|
SystemLabel: labelStyle(t.System),
|
||||||
|
|
||||||
|
UserMessage: bubbleMessageStyle(t.Text, t.SurfaceAlt),
|
||||||
|
AssistantMsg: assistantMessageStyle(t.Text),
|
||||||
|
ToolMessage: messageStyle(t.Text),
|
||||||
|
ErrorMessage: messageStyle(t.Error),
|
||||||
|
SystemMessage: lipgloss.NewStyle().Foreground(lipgloss.Color(t.System)).Padding(0, 1),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -180,21 +236,30 @@ func labelStyle(color string) lipgloss.Style {
|
|||||||
Foreground(lipgloss.Color(color))
|
Foreground(lipgloss.Color(color))
|
||||||
}
|
}
|
||||||
|
|
||||||
func messageStyle(textColor, accentColor string) lipgloss.Style {
|
func messageStyle(textColor string) lipgloss.Style {
|
||||||
return lipgloss.NewStyle().
|
return lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(textColor)).
|
Foreground(lipgloss.Color(textColor)).
|
||||||
Border(lipgloss.ThickBorder(), false, false, false, true).
|
Padding(0, 1)
|
||||||
BorderForeground(lipgloss.Color(accentColor)).
|
}
|
||||||
PaddingLeft(1)
|
|
||||||
|
func assistantMessageStyle(textColor string) lipgloss.Style {
|
||||||
|
return lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(textColor)).
|
||||||
|
Padding(1, 2)
|
||||||
|
}
|
||||||
|
|
||||||
|
func bubbleMessageStyle(textColor, backgroundColor string) lipgloss.Style {
|
||||||
|
return lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(textColor)).
|
||||||
|
Background(lipgloss.Color(backgroundColor)).
|
||||||
|
Padding(1, 2)
|
||||||
}
|
}
|
||||||
|
|
||||||
func applyTextareaTheme(input *textarea.Model, t theme) {
|
func applyTextareaTheme(input *textarea.Model, t theme) {
|
||||||
focused := textarea.Style{
|
focused := textarea.Style{
|
||||||
Base: lipgloss.NewStyle().
|
Base: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Text)).
|
Foreground(lipgloss.Color(t.Text)),
|
||||||
Background(lipgloss.Color(t.Surface)),
|
CursorLine: lipgloss.NewStyle(),
|
||||||
CursorLine: lipgloss.NewStyle().
|
|
||||||
Background(lipgloss.Color(t.Surface)),
|
|
||||||
Placeholder: lipgloss.NewStyle().
|
Placeholder: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Muted)),
|
Foreground(lipgloss.Color(t.Muted)),
|
||||||
Prompt: lipgloss.NewStyle().
|
Prompt: lipgloss.NewStyle().
|
||||||
@@ -205,7 +270,7 @@ func applyTextareaTheme(input *textarea.Model, t theme) {
|
|||||||
Foreground(lipgloss.Color(t.Surface)),
|
Foreground(lipgloss.Color(t.Surface)),
|
||||||
}
|
}
|
||||||
blurred := focused
|
blurred := focused
|
||||||
blurred.CursorLine = lipgloss.NewStyle().Background(lipgloss.Color(t.Surface))
|
blurred.CursorLine = lipgloss.NewStyle()
|
||||||
|
|
||||||
input.FocusedStyle = focused
|
input.FocusedStyle = focused
|
||||||
input.BlurredStyle = blurred
|
input.BlurredStyle = blurred
|
||||||
|
|||||||
@@ -0,0 +1,143 @@
|
|||||||
|
package tui
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"os/exec"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agentu/internal/tools"
|
||||||
|
)
|
||||||
|
|
||||||
|
const workflowCommandTimeout = 2 * time.Minute
|
||||||
|
|
||||||
|
func (m model) handleStatusCommand(args []string) (string, error) {
|
||||||
|
if len(args) != 0 {
|
||||||
|
return "", fmt.Errorf("usage: /status")
|
||||||
|
}
|
||||||
|
out, err := m.runGitCommand("status", "--short")
|
||||||
|
if err != nil {
|
||||||
|
return out, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(out) == "" {
|
||||||
|
return "Working tree clean.", nil
|
||||||
|
}
|
||||||
|
return "Changed files:\n" + out, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m model) handleDiffCommand(args []string) (string, error) {
|
||||||
|
if len(args) > 0 && args[0] == "--stat" {
|
||||||
|
out, err := m.runGitCommand(append([]string{"diff", "--stat", "--"}, args[1:]...)...)
|
||||||
|
if err != nil {
|
||||||
|
return out, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(out) == "" {
|
||||||
|
return "No unstaged diff.", nil
|
||||||
|
}
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
|
out, err := m.runGitCommand(append([]string{"diff", "--"}, args...)...)
|
||||||
|
if err != nil {
|
||||||
|
return out, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(out) == "" {
|
||||||
|
return "No unstaged diff.", nil
|
||||||
|
}
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m model) handleTestCommand(args []string) (string, error) {
|
||||||
|
command := strings.TrimSpace(strings.Join(args, " "))
|
||||||
|
if command == "" {
|
||||||
|
var err error
|
||||||
|
command, err = defaultTestCommand(m.workflowDir())
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out, err := m.runShellCommand(command)
|
||||||
|
if err != nil {
|
||||||
|
if strings.TrimSpace(out) != "" {
|
||||||
|
return out, err
|
||||||
|
}
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(out) == "" {
|
||||||
|
return fmt.Sprintf("$ %s\n(no output)", command), nil
|
||||||
|
}
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m model) runGitCommand(args ...string) (string, error) {
|
||||||
|
ctx, cancel := context.WithTimeout(m.ctx, workflowCommandTimeout)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
cmd := exec.CommandContext(ctx, "git", append([]string{"-C", m.workflowDir()}, args...)...)
|
||||||
|
output, err := cmd.CombinedOutput()
|
||||||
|
text := tools.FormatResult(string(output))
|
||||||
|
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
||||||
|
return text, fmt.Errorf("git command timed out: %w", ctx.Err())
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
if strings.TrimSpace(text) != "" {
|
||||||
|
return text, fmt.Errorf("git command failed: %w", err)
|
||||||
|
}
|
||||||
|
return "", fmt.Errorf("git command failed: %w", err)
|
||||||
|
}
|
||||||
|
return text, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m model) runShellCommand(command string) (string, error) {
|
||||||
|
ctx, cancel := context.WithTimeout(m.ctx, workflowCommandTimeout)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
shell := os.Getenv("SHELL")
|
||||||
|
if shell == "" {
|
||||||
|
shell = "/bin/sh"
|
||||||
|
}
|
||||||
|
cmd := exec.CommandContext(ctx, shell, "-lc", command)
|
||||||
|
cmd.Dir = m.workflowDir()
|
||||||
|
|
||||||
|
var combined bytes.Buffer
|
||||||
|
cmd.Stdout = &combined
|
||||||
|
cmd.Stderr = &combined
|
||||||
|
err := cmd.Run()
|
||||||
|
text := strings.TrimRight(tools.FormatResult(combined.String()), "\n")
|
||||||
|
if text != "" {
|
||||||
|
text = fmt.Sprintf("$ %s\n%s", command, text)
|
||||||
|
}
|
||||||
|
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
||||||
|
return text, fmt.Errorf("command timed out: %w", ctx.Err())
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return text, fmt.Errorf("command failed: %w", err)
|
||||||
|
}
|
||||||
|
return text, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m model) workflowDir() string {
|
||||||
|
if strings.TrimSpace(m.workingDir) == "" {
|
||||||
|
return "."
|
||||||
|
}
|
||||||
|
return m.workingDir
|
||||||
|
}
|
||||||
|
|
||||||
|
func defaultTestCommand(dir string) (string, error) {
|
||||||
|
if fileExists(filepath.Join(dir, "go.mod")) {
|
||||||
|
return "go test ./...", nil
|
||||||
|
}
|
||||||
|
if fileExists(filepath.Join(dir, "package.json")) {
|
||||||
|
return "npm test", nil
|
||||||
|
}
|
||||||
|
return "", fmt.Errorf("no default test command found; use /test <command>")
|
||||||
|
}
|
||||||
|
|
||||||
|
func fileExists(path string) bool {
|
||||||
|
info, err := os.Stat(path)
|
||||||
|
return err == nil && !info.IsDir()
|
||||||
|
}
|
||||||
@@ -0,0 +1,228 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"flag"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"os"
|
||||||
|
"os/signal"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agentu/internal/agent"
|
||||||
|
"agentu/internal/session"
|
||||||
|
"agentu/internal/tools"
|
||||||
|
"agentu/internal/tui"
|
||||||
|
"agentu/pkg/config"
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
const defaultConfigContent = `# agentu configuration
|
||||||
|
# Replace the placeholder values below with your actual provider settings.
|
||||||
|
# See agentu.example.yaml for the full reference.
|
||||||
|
providers:
|
||||||
|
default:
|
||||||
|
base_url: https://api.openai.com
|
||||||
|
api_key: YOUR_API_KEY_HERE
|
||||||
|
models:
|
||||||
|
- id: gpt-4o
|
||||||
|
context: 128000
|
||||||
|
- id: gpt-4o-mini
|
||||||
|
context: 128000
|
||||||
|
`
|
||||||
|
|
||||||
|
func main() {
|
||||||
|
if err := run(); err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "agentu:", err)
|
||||||
|
os.Exit(1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func run() error {
|
||||||
|
configPath := flag.String("config", config.DefaultPath, "path to YAML config")
|
||||||
|
yolo := flag.Bool("yolo", false, "also enable automatic file writes and shell execution")
|
||||||
|
plain := flag.Bool("plain", false, "use the simple line-based REPL instead of the TUI")
|
||||||
|
themeName := flag.String("theme", "light", "TUI theme: light or dark")
|
||||||
|
resumeID := flag.String("resume", "", "resume a specific session by ID")
|
||||||
|
flag.Parse()
|
||||||
|
|
||||||
|
themeMode, err := tui.ParseThemeMode(*themeName)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
cfg, err := config.Load(*configPath)
|
||||||
|
if err != nil {
|
||||||
|
if errors.Is(err, os.ErrNotExist) {
|
||||||
|
if bootstrapErr := bootstrapConfig(*configPath); bootstrapErr != nil {
|
||||||
|
return bootstrapErr
|
||||||
|
}
|
||||||
|
return fmt.Errorf("created default config at %s; edit it to set your API key and provider, then restart agentu", *configPath)
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
providerName := cfg.DefaultProviderName()
|
||||||
|
providerConfig := cfg.Providers[providerName]
|
||||||
|
provider := llm.NewOpenAICompatibleClient(providerConfig.BaseURL, providerConfig.APIKey, http.DefaultClient)
|
||||||
|
|
||||||
|
var registry *tools.Registry
|
||||||
|
if *yolo {
|
||||||
|
registry = tools.Builtins(cfg.Agent.WorkingDir)
|
||||||
|
} else {
|
||||||
|
registry = tools.ReadOnlyBuiltins(cfg.Agent.WorkingDir)
|
||||||
|
}
|
||||||
|
registry.Register(tools.NewBuiltinSearchTool(http.DefaultClient))
|
||||||
|
registry.Register(tools.NewBuiltinFetchTool(http.DefaultClient))
|
||||||
|
systemPrompt := strings.TrimSpace(cfg.Agent.SystemPrompt + "\n\n" + tools.AgentInstructions(cfg.Agent.WorkingDir, *yolo))
|
||||||
|
maxContextTokens := cfg.Agent.MaxContextTokens
|
||||||
|
if maxContextTokens == 0 {
|
||||||
|
maxContextTokens = providerConfig.ContextTokens
|
||||||
|
}
|
||||||
|
|
||||||
|
assistant := agent.New(agent.Options{
|
||||||
|
Provider: provider,
|
||||||
|
ProviderName: providerName,
|
||||||
|
Model: providerConfig.Model,
|
||||||
|
Thinking: providerConfig.Thinking,
|
||||||
|
ThinkingParam: providerConfig.ThinkingParam,
|
||||||
|
ThinkingEnabled: providerConfig.ThinkingConfigured,
|
||||||
|
SystemPrompt: systemPrompt,
|
||||||
|
ToolRegistry: registry,
|
||||||
|
ToolTimeout: cfg.Agent.ToolTimeout,
|
||||||
|
MaxContextTokens: maxContextTokens,
|
||||||
|
})
|
||||||
|
|
||||||
|
// Initialize session store
|
||||||
|
home, err := os.UserHomeDir()
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("resolve home dir: %w", err)
|
||||||
|
}
|
||||||
|
sessionsDir := filepath.Join(home, ".agentu", "sessions")
|
||||||
|
store, err := session.NewStore(sessionsDir)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("init session store: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
modelManager, err := session.NewManager(cfg, assistant, http.DefaultClient, store)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := initializeSession(modelManager, *resumeID); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
|
||||||
|
defer stop()
|
||||||
|
|
||||||
|
if !*plain {
|
||||||
|
err := tui.Run(ctx, assistant, tui.Options{
|
||||||
|
ModelName: providerConfig.Model,
|
||||||
|
Yolo: *yolo,
|
||||||
|
ThemeMode: themeMode,
|
||||||
|
ModelManager: modelManager,
|
||||||
|
WorkingDir: cfg.Agent.WorkingDir,
|
||||||
|
})
|
||||||
|
printExitSession(modelManager)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
return repl(ctx, assistant, modelManager)
|
||||||
|
}
|
||||||
|
|
||||||
|
func initializeSession(modelManager *session.Manager, resumeID string) error {
|
||||||
|
if resumeID != "" {
|
||||||
|
if err := modelManager.Resume(resumeID); err != nil {
|
||||||
|
return fmt.Errorf("resume session: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if _, err := modelManager.NewSession(); err != nil {
|
||||||
|
return fmt.Errorf("start new session: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func bootstrapConfig(configPath string) error {
|
||||||
|
resolved, err := config.ResolvePath(configPath)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("resolve config path: %w", err)
|
||||||
|
}
|
||||||
|
dir := filepath.Dir(resolved)
|
||||||
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||||
|
return fmt.Errorf("create config directory: %w", err)
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(resolved, []byte(defaultConfigContent), 0o644); err != nil {
|
||||||
|
return fmt.Errorf("write default config: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager) error {
|
||||||
|
scanner := bufio.NewScanner(os.Stdin)
|
||||||
|
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
|
||||||
|
|
||||||
|
fmt.Println("agentu")
|
||||||
|
fmt.Println("Type /exit to quit, /clear to reset context, /compact to compress context.")
|
||||||
|
for {
|
||||||
|
fmt.Print("> ")
|
||||||
|
if !scanner.Scan() {
|
||||||
|
if err := scanner.Err(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
fmt.Println()
|
||||||
|
printExitSession(modelManager)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
input := strings.TrimSpace(scanner.Text())
|
||||||
|
switch input {
|
||||||
|
case "":
|
||||||
|
continue
|
||||||
|
case "/exit", "/quit":
|
||||||
|
printExitSession(modelManager)
|
||||||
|
return nil
|
||||||
|
case "/clear":
|
||||||
|
assistant.Clear()
|
||||||
|
fmt.Println("context cleared")
|
||||||
|
continue
|
||||||
|
case "/compact":
|
||||||
|
if err := assistant.Compact(ctx); err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "compact error:", err)
|
||||||
|
} else {
|
||||||
|
fmt.Println("context compacted")
|
||||||
|
if modelManager != nil {
|
||||||
|
_ = modelManager.Save()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := assistant.RunTurn(ctx, input, os.Stdout, os.Stderr); err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "\nerror:", err)
|
||||||
|
}
|
||||||
|
_ = modelManager.Save()
|
||||||
|
fmt.Println()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func printExitSession(m *session.Manager) {
|
||||||
|
if m == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
_ = m.Save()
|
||||||
|
id := m.CurrentSessionID()
|
||||||
|
name := m.CurrentSessionName()
|
||||||
|
if id != "" {
|
||||||
|
if name != "" && name != id {
|
||||||
|
fmt.Fprintf(os.Stderr, "session: %s (%s)\n", id, name)
|
||||||
|
} else {
|
||||||
|
fmt.Fprintf(os.Stderr, "session: %s\n", id)
|
||||||
|
}
|
||||||
|
fmt.Fprintf(os.Stderr, "resume: agentu --resume %s\n", id)
|
||||||
|
}
|
||||||
|
}
|
||||||
+101
@@ -0,0 +1,101 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agentu/internal/agent"
|
||||||
|
"agentu/internal/session"
|
||||||
|
"agentu/pkg/config"
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
type startupProvider struct{}
|
||||||
|
|
||||||
|
func (p *startupProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
return emit(llm.StreamEvent{Content: "ok"})
|
||||||
|
}
|
||||||
|
|
||||||
|
func startupTestManager(t *testing.T) (*session.Manager, *agent.Agent) {
|
||||||
|
t.Helper()
|
||||||
|
cfg := &config.Config{
|
||||||
|
Providers: map[string]config.ProviderConfig{
|
||||||
|
"one": {
|
||||||
|
Name: "one",
|
||||||
|
BaseURL: "https://one.example.com",
|
||||||
|
APIKey: "sk-test",
|
||||||
|
Model: "model-a",
|
||||||
|
Models: []string{"model-a"},
|
||||||
|
Thinking: "none",
|
||||||
|
ThinkingParam: "thinking",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
assistant := agent.New(agent.Options{
|
||||||
|
Provider: &startupProvider{},
|
||||||
|
ProviderName: "one",
|
||||||
|
Model: "model-a",
|
||||||
|
})
|
||||||
|
store, err := session.NewStore(filepath.Join(t.TempDir(), "sessions"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
manager, err := session.NewManager(cfg, assistant, nil, store)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
return manager, assistant
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInitializeSessionStartsNewByDefault(t *testing.T) {
|
||||||
|
manager, assistant := startupTestManager(t)
|
||||||
|
if _, err := manager.NewSession(); err != nil {
|
||||||
|
t.Fatalf("new session: %v", err)
|
||||||
|
}
|
||||||
|
savedID := manager.CurrentSessionID()
|
||||||
|
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: "old context"}})
|
||||||
|
if err := manager.Save(); err != nil {
|
||||||
|
t.Fatalf("save: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := initializeSession(manager, ""); err != nil {
|
||||||
|
t.Fatalf("initialize: %v", err)
|
||||||
|
}
|
||||||
|
if manager.CurrentSessionID() == savedID {
|
||||||
|
t.Fatalf("default startup resumed latest session %s", savedID)
|
||||||
|
}
|
||||||
|
for _, msg := range assistant.Messages() {
|
||||||
|
if msg.Content == "old context" {
|
||||||
|
t.Fatal("default startup carried old context into new session")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInitializeSessionResumesExplicitID(t *testing.T) {
|
||||||
|
manager, assistant := startupTestManager(t)
|
||||||
|
if _, err := manager.NewSession(); err != nil {
|
||||||
|
t.Fatalf("new session: %v", err)
|
||||||
|
}
|
||||||
|
savedID := manager.CurrentSessionID()
|
||||||
|
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: "old context"}})
|
||||||
|
if err := manager.Save(); err != nil {
|
||||||
|
t.Fatalf("save: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := initializeSession(manager, savedID); err != nil {
|
||||||
|
t.Fatalf("initialize: %v", err)
|
||||||
|
}
|
||||||
|
if manager.CurrentSessionID() != savedID {
|
||||||
|
t.Fatalf("expected resumed session %s, got %s", savedID, manager.CurrentSessionID())
|
||||||
|
}
|
||||||
|
found := false
|
||||||
|
for _, msg := range assistant.Messages() {
|
||||||
|
if msg.Content == "old context" {
|
||||||
|
found = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !found {
|
||||||
|
t.Fatal("explicit resume did not restore saved context")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -21,37 +21,50 @@ type Config struct {
|
|||||||
|
|
||||||
type ProviderConfig struct {
|
type ProviderConfig struct {
|
||||||
Name string
|
Name string
|
||||||
BaseURL string `yaml:"base_url"`
|
BaseURL string
|
||||||
Model string `yaml:"model"`
|
Model string
|
||||||
Models []string `yaml:"models"`
|
Models []string
|
||||||
APIKey string `yaml:"api_key"`
|
ModelConfigs map[string]ModelConfig
|
||||||
Thinking string `yaml:"thinking"`
|
ContextTokens int
|
||||||
ThinkingParam string `yaml:"thinking_param"`
|
APIKey string
|
||||||
|
Thinking string
|
||||||
|
ThinkingParam string
|
||||||
ThinkingConfigured bool
|
ThinkingConfigured bool
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type ModelConfig struct {
|
||||||
|
ID string `yaml:"id"`
|
||||||
|
Thinking string `yaml:"thinking"`
|
||||||
|
ContextTokens int `yaml:"context"`
|
||||||
|
}
|
||||||
|
|
||||||
type AgentConfig struct {
|
type AgentConfig struct {
|
||||||
SystemPrompt string `yaml:"system_prompt"`
|
SystemPrompt string `yaml:"system_prompt"`
|
||||||
WorkingDir string `yaml:"working_dir"`
|
WorkingDir string `yaml:"working_dir"`
|
||||||
ToolTimeout time.Duration `yaml:"tool_timeout"`
|
ToolTimeout time.Duration `yaml:"tool_timeout"`
|
||||||
|
MaxContextTokens int `yaml:"max_context_tokens"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type rawConfig struct {
|
type rawConfig struct {
|
||||||
Providers map[string]rawProviderConfig `yaml:"providers"`
|
Providers map[string]rawProviderConfig `yaml:"providers"`
|
||||||
Agent struct {
|
Agent struct {
|
||||||
SystemPrompt string `yaml:"system_prompt"`
|
SystemPrompt string `yaml:"system_prompt"`
|
||||||
WorkingDir string `yaml:"working_dir"`
|
WorkingDir string `yaml:"working_dir"`
|
||||||
ToolTimeout string `yaml:"tool_timeout"`
|
ToolTimeout string `yaml:"tool_timeout"`
|
||||||
|
MaxContextTokens int `yaml:"max_context_tokens"`
|
||||||
} `yaml:"agent"`
|
} `yaml:"agent"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type rawProviderConfig struct {
|
type rawProviderConfig struct {
|
||||||
BaseURL string `yaml:"base_url"`
|
BaseURL string `yaml:"base_url"`
|
||||||
Model string `yaml:"model"`
|
Models []rawModelConfig `yaml:"models"`
|
||||||
Models []string `yaml:"models"`
|
APIKey string `yaml:"api_key"`
|
||||||
APIKey string `yaml:"api_key"`
|
}
|
||||||
Thinking *string `yaml:"thinking"`
|
|
||||||
ThinkingParam string `yaml:"thinking_param"`
|
type rawModelConfig struct {
|
||||||
|
ID string `yaml:"id"`
|
||||||
|
Thinking string `yaml:"thinking"`
|
||||||
|
ContextTokens int `yaml:"context"`
|
||||||
}
|
}
|
||||||
|
|
||||||
func Load(path string) (*Config, error) {
|
func Load(path string) (*Config, error) {
|
||||||
@@ -76,8 +89,9 @@ func Load(path string) (*Config, error) {
|
|||||||
cfg := &Config{
|
cfg := &Config{
|
||||||
Providers: normalizeProviders(raw),
|
Providers: normalizeProviders(raw),
|
||||||
Agent: AgentConfig{
|
Agent: AgentConfig{
|
||||||
SystemPrompt: raw.Agent.SystemPrompt,
|
SystemPrompt: raw.Agent.SystemPrompt,
|
||||||
WorkingDir: raw.Agent.WorkingDir,
|
WorkingDir: raw.Agent.WorkingDir,
|
||||||
|
MaxContextTokens: raw.Agent.MaxContextTokens,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -142,8 +156,8 @@ func (c *Config) Validate() error {
|
|||||||
if strings.TrimSpace(provider.BaseURL) == "" {
|
if strings.TrimSpace(provider.BaseURL) == "" {
|
||||||
missing = append(missing, prefix+".base_url")
|
missing = append(missing, prefix+".base_url")
|
||||||
}
|
}
|
||||||
if strings.TrimSpace(provider.Model) == "" {
|
if len(provider.Models) == 0 {
|
||||||
missing = append(missing, prefix+".model")
|
missing = append(missing, prefix+".models")
|
||||||
}
|
}
|
||||||
if strings.TrimSpace(provider.APIKey) == "" {
|
if strings.TrimSpace(provider.APIKey) == "" {
|
||||||
missing = append(missing, prefix+".api_key")
|
missing = append(missing, prefix+".api_key")
|
||||||
@@ -151,6 +165,19 @@ func (c *Config) Validate() error {
|
|||||||
if len(missing) > 0 {
|
if len(missing) > 0 {
|
||||||
return fmt.Errorf("missing required config: %s", strings.Join(missing, ", "))
|
return fmt.Errorf("missing required config: %s", strings.Join(missing, ", "))
|
||||||
}
|
}
|
||||||
|
for _, model := range provider.Models {
|
||||||
|
if strings.TrimSpace(model) == "" {
|
||||||
|
return fmt.Errorf("%s.models.id is required", prefix)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, model := range provider.ModelConfigs {
|
||||||
|
if model.ContextTokens < 0 {
|
||||||
|
return fmt.Errorf("%s.models.context must be greater than or equal to zero", prefix)
|
||||||
|
}
|
||||||
|
if model.Thinking != "" && !IsThinkingLevel(model.Thinking) {
|
||||||
|
return fmt.Errorf("%s.models.thinking must be one of: %s", prefix, strings.Join(ThinkingLevels(), ", "))
|
||||||
|
}
|
||||||
|
}
|
||||||
if provider.ThinkingConfigured && !IsThinkingLevel(provider.Thinking) {
|
if provider.ThinkingConfigured && !IsThinkingLevel(provider.Thinking) {
|
||||||
return fmt.Errorf("%s.thinking must be one of: %s", prefix, strings.Join(ThinkingLevels(), ", "))
|
return fmt.Errorf("%s.thinking must be one of: %s", prefix, strings.Join(ThinkingLevels(), ", "))
|
||||||
}
|
}
|
||||||
@@ -158,6 +185,9 @@ func (c *Config) Validate() error {
|
|||||||
if c.Agent.ToolTimeout <= 0 {
|
if c.Agent.ToolTimeout <= 0 {
|
||||||
return errors.New("agent.tool_timeout must be greater than zero")
|
return errors.New("agent.tool_timeout must be greater than zero")
|
||||||
}
|
}
|
||||||
|
if c.Agent.MaxContextTokens < 0 {
|
||||||
|
return errors.New("agent.max_context_tokens must be greater than or equal to zero")
|
||||||
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -187,27 +217,59 @@ func normalizeProviders(raw rawConfig) map[string]ProviderConfig {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func normalizeProvider(name string, raw rawProviderConfig) ProviderConfig {
|
func normalizeProvider(name string, raw rawProviderConfig) ProviderConfig {
|
||||||
|
models, modelConfigs := normalizeModelConfigs(raw.Models)
|
||||||
provider := ProviderConfig{
|
provider := ProviderConfig{
|
||||||
Name: name,
|
Name: name,
|
||||||
BaseURL: raw.BaseURL,
|
BaseURL: raw.BaseURL,
|
||||||
Model: raw.Model,
|
Models: models,
|
||||||
Models: append([]string(nil), raw.Models...),
|
ModelConfigs: modelConfigs,
|
||||||
APIKey: raw.APIKey,
|
APIKey: raw.APIKey,
|
||||||
ThinkingParam: raw.ThinkingParam,
|
|
||||||
}
|
}
|
||||||
if provider.Model == "" && len(provider.Models) > 0 {
|
if len(provider.Models) > 0 {
|
||||||
provider.Model = provider.Models[0]
|
provider.Model = provider.Models[0]
|
||||||
}
|
}
|
||||||
if raw.Thinking != nil {
|
applyModelDefaults(&provider)
|
||||||
provider.Thinking = strings.ToLower(strings.TrimSpace(*raw.Thinking))
|
if provider.ThinkingConfigured && provider.ThinkingParam == "" {
|
||||||
provider.ThinkingConfigured = true
|
provider.ThinkingParam = "thinking"
|
||||||
if provider.ThinkingParam == "" {
|
|
||||||
provider.ThinkingParam = "thinking"
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
return provider
|
return provider
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func normalizeModelConfigs(rawModels []rawModelConfig) ([]string, map[string]ModelConfig) {
|
||||||
|
models := make([]string, 0, len(rawModels))
|
||||||
|
configs := make(map[string]ModelConfig, len(rawModels))
|
||||||
|
for _, raw := range rawModels {
|
||||||
|
model := ModelConfig{
|
||||||
|
ID: strings.TrimSpace(raw.ID),
|
||||||
|
Thinking: strings.ToLower(strings.TrimSpace(raw.Thinking)),
|
||||||
|
ContextTokens: raw.ContextTokens,
|
||||||
|
}
|
||||||
|
models = append(models, model.ID)
|
||||||
|
if model.ID != "" {
|
||||||
|
configs[model.ID] = model
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return models, configs
|
||||||
|
}
|
||||||
|
|
||||||
|
func applyModelDefaults(provider *ProviderConfig) {
|
||||||
|
model, ok := provider.ModelConfigs[provider.Model]
|
||||||
|
if !ok {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
provider.ContextTokens = model.ContextTokens
|
||||||
|
if model.Thinking != "" {
|
||||||
|
provider.Thinking = model.Thinking
|
||||||
|
provider.ThinkingConfigured = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p ProviderConfig) WithModel(model string) ProviderConfig {
|
||||||
|
p.Model = strings.TrimSpace(model)
|
||||||
|
applyModelDefaults(&p)
|
||||||
|
return p
|
||||||
|
}
|
||||||
|
|
||||||
func ThinkingLevels() []string {
|
func ThinkingLevels() []string {
|
||||||
return []string{"none", "middle", "high", "xhigh", "max"}
|
return []string{"none", "middle", "high", "xhigh", "max"}
|
||||||
}
|
}
|
||||||
@@ -16,8 +16,9 @@ func TestLoadExpandsEnvAndDefaults(t *testing.T) {
|
|||||||
providers:
|
providers:
|
||||||
test:
|
test:
|
||||||
base_url: https://example.com
|
base_url: https://example.com
|
||||||
model: test-model
|
|
||||||
api_key: ${AGENTU_TEST_KEY}
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
|
models:
|
||||||
|
- id: test-model
|
||||||
agent:
|
agent:
|
||||||
working_dir: .
|
working_dir: .
|
||||||
`), 0o644); err != nil {
|
`), 0o644); err != nil {
|
||||||
@@ -43,6 +44,115 @@ agent:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestLoadMinimalModelObjectConfig(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
t.Setenv("AGENTU_TEST_KEY", "sk-test")
|
||||||
|
path := filepath.Join(dir, "agentu.yaml")
|
||||||
|
if err := os.WriteFile(path, []byte(`
|
||||||
|
providers:
|
||||||
|
test:
|
||||||
|
base_url: https://example.com
|
||||||
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
|
models:
|
||||||
|
- id: model-a
|
||||||
|
thinking: none
|
||||||
|
context: 256000
|
||||||
|
`), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cfg, err := Load(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
provider := cfg.Providers["test"]
|
||||||
|
if provider.Model != "model-a" {
|
||||||
|
t.Fatalf("model = %q", provider.Model)
|
||||||
|
}
|
||||||
|
if len(provider.Models) != 1 || provider.Models[0] != "model-a" {
|
||||||
|
t.Fatalf("models = %#v", provider.Models)
|
||||||
|
}
|
||||||
|
if provider.ContextTokens != 256000 {
|
||||||
|
t.Fatalf("context tokens = %d", provider.ContextTokens)
|
||||||
|
}
|
||||||
|
if !provider.ThinkingConfigured || provider.Thinking != "none" || provider.ThinkingParam != "thinking" {
|
||||||
|
t.Fatalf("thinking config = %#v", provider)
|
||||||
|
}
|
||||||
|
if cfg.Agent.ToolTimeout != 30*time.Second {
|
||||||
|
t.Fatalf("tool timeout = %s", cfg.Agent.ToolTimeout)
|
||||||
|
}
|
||||||
|
if cfg.Agent.SystemPrompt == "" {
|
||||||
|
t.Fatal("system prompt default was not applied")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoadRejectsStringModelEntries(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "agentu.yaml")
|
||||||
|
if err := os.WriteFile(path, []byte(`
|
||||||
|
providers:
|
||||||
|
test:
|
||||||
|
base_url: https://example.com
|
||||||
|
api_key: sk-test
|
||||||
|
models:
|
||||||
|
- test-model
|
||||||
|
`), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err := Load(path)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "cannot unmarshal") {
|
||||||
|
t.Fatalf("err = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoadMaxContextTokens(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "agentu.yaml")
|
||||||
|
if err := os.WriteFile(path, []byte(`
|
||||||
|
providers:
|
||||||
|
test:
|
||||||
|
base_url: https://example.com
|
||||||
|
api_key: sk-test
|
||||||
|
models:
|
||||||
|
- id: test-model
|
||||||
|
agent:
|
||||||
|
max_context_tokens: 120000
|
||||||
|
`), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cfg, err := Load(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if cfg.Agent.MaxContextTokens != 120000 {
|
||||||
|
t.Fatalf("max context tokens = %d", cfg.Agent.MaxContextTokens)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoadRejectsNegativeMaxContextTokens(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "agentu.yaml")
|
||||||
|
if err := os.WriteFile(path, []byte(`
|
||||||
|
providers:
|
||||||
|
test:
|
||||||
|
base_url: https://example.com
|
||||||
|
api_key: sk-test
|
||||||
|
models:
|
||||||
|
- id: test-model
|
||||||
|
agent:
|
||||||
|
max_context_tokens: -1
|
||||||
|
`), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err := Load(path)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "agent.max_context_tokens") {
|
||||||
|
t.Fatalf("err = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestLoadUsesDefaultConfigPath(t *testing.T) {
|
func TestLoadUsesDefaultConfigPath(t *testing.T) {
|
||||||
dir := t.TempDir()
|
dir := t.TempDir()
|
||||||
t.Setenv("HOME", dir)
|
t.Setenv("HOME", dir)
|
||||||
@@ -56,8 +166,9 @@ func TestLoadUsesDefaultConfigPath(t *testing.T) {
|
|||||||
providers:
|
providers:
|
||||||
test:
|
test:
|
||||||
base_url: https://example.com
|
base_url: https://example.com
|
||||||
model: test-model
|
|
||||||
api_key: ${AGENTU_TEST_KEY}
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
|
models:
|
||||||
|
- id: test-model
|
||||||
`), 0o644); err != nil {
|
`), 0o644); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
@@ -93,8 +204,9 @@ func TestLoadDefaultSystemPromptDiscouragesEagerTools(t *testing.T) {
|
|||||||
providers:
|
providers:
|
||||||
test:
|
test:
|
||||||
base_url: https://example.com
|
base_url: https://example.com
|
||||||
model: test-model
|
|
||||||
api_key: ${AGENTU_TEST_KEY}
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
|
models:
|
||||||
|
- id: test-model
|
||||||
`), 0o644); err != nil {
|
`), 0o644); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
@@ -119,17 +231,15 @@ providers:
|
|||||||
loveuer:
|
loveuer:
|
||||||
base_url: https://example.com
|
base_url: https://example.com
|
||||||
api_key: ${AGENTU_TEST_KEY}
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
model: model-a
|
|
||||||
models:
|
models:
|
||||||
- model-a
|
- id: model-a
|
||||||
- model-b
|
thinking: xhigh
|
||||||
thinking: xhigh
|
- id: model-b
|
||||||
thinking_param: thinking
|
|
||||||
backup:
|
backup:
|
||||||
base_url: https://backup.example.com
|
base_url: https://backup.example.com
|
||||||
api_key: ${AGENTU_TEST_KEY}
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
models:
|
models:
|
||||||
- model-c
|
- id: model-c
|
||||||
`), 0o644); err != nil {
|
`), 0o644); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
@@ -162,11 +272,13 @@ providers:
|
|||||||
two:
|
two:
|
||||||
base_url: https://two.example.com
|
base_url: https://two.example.com
|
||||||
api_key: ${AGENTU_TEST_KEY}
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
model: model-b
|
models:
|
||||||
|
- id: model-b
|
||||||
one:
|
one:
|
||||||
base_url: https://one.example.com
|
base_url: https://one.example.com
|
||||||
api_key: ${AGENTU_TEST_KEY}
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
model: model-a
|
models:
|
||||||
|
- id: model-a
|
||||||
`), 0o644); err != nil {
|
`), 0o644); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
@@ -188,9 +300,10 @@ func TestLoadRejectsInvalidThinking(t *testing.T) {
|
|||||||
providers:
|
providers:
|
||||||
test:
|
test:
|
||||||
base_url: https://example.com
|
base_url: https://example.com
|
||||||
model: test-model
|
|
||||||
api_key: ${AGENTU_TEST_KEY}
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
thinking: huge
|
models:
|
||||||
|
- id: test-model
|
||||||
|
thinking: huge
|
||||||
`), 0o644); err != nil {
|
`), 0o644); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
@@ -224,8 +337,9 @@ func TestLoadRejectsLegacyProviderConfig(t *testing.T) {
|
|||||||
if err := os.WriteFile(path, []byte(`
|
if err := os.WriteFile(path, []byte(`
|
||||||
provider:
|
provider:
|
||||||
base_url: https://example.com
|
base_url: https://example.com
|
||||||
model: test-model
|
|
||||||
api_key: sk-test
|
api_key: sk-test
|
||||||
|
models:
|
||||||
|
- id: test-model
|
||||||
`), 0o644); err != nil {
|
`), 0o644); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
@@ -244,8 +358,9 @@ active_provider: loveuer
|
|||||||
providers:
|
providers:
|
||||||
loveuer:
|
loveuer:
|
||||||
base_url: https://example.com
|
base_url: https://example.com
|
||||||
model: test-model
|
|
||||||
api_key: sk-test
|
api_key: sk-test
|
||||||
|
models:
|
||||||
|
- id: test-model
|
||||||
`), 0o644); err != nil {
|
`), 0o644); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
@@ -10,8 +10,13 @@ import (
|
|||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
"strings"
|
"strings"
|
||||||
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
const maxChatAttempts = 3
|
||||||
|
|
||||||
|
const retryDelay = 100 * time.Millisecond
|
||||||
|
|
||||||
type OpenAICompatibleClient struct {
|
type OpenAICompatibleClient struct {
|
||||||
baseURL string
|
baseURL string
|
||||||
apiKey string
|
apiKey string
|
||||||
@@ -31,34 +36,82 @@ func NewOpenAICompatibleClient(baseURL, apiKey string, httpClient *http.Client)
|
|||||||
|
|
||||||
func (c *OpenAICompatibleClient) ChatStream(ctx context.Context, req ChatRequest, emit func(StreamEvent) error) error {
|
func (c *OpenAICompatibleClient) ChatStream(ctx context.Context, req ChatRequest, emit func(StreamEvent) error) error {
|
||||||
req.Stream = true
|
req.Stream = true
|
||||||
|
extra := make(map[string]any, len(req.Extra)+1)
|
||||||
|
for key, value := range req.Extra {
|
||||||
|
extra[key] = value
|
||||||
|
}
|
||||||
|
if _, exists := extra["stream_options"]; !exists {
|
||||||
|
extra["stream_options"] = map[string]any{"include_usage": true}
|
||||||
|
}
|
||||||
|
req.Extra = extra
|
||||||
body, err := json.Marshal(req)
|
body, err := json.Marshal(req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("marshal chat request: %w", err)
|
return fmt.Errorf("marshal chat request: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
|
var lastErr error
|
||||||
if err != nil {
|
for attempt := 1; attempt <= maxChatAttempts; attempt++ {
|
||||||
return fmt.Errorf("create chat request: %w", err)
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
|
||||||
}
|
if err != nil {
|
||||||
httpReq.Header.Set("Content-Type", "application/json")
|
return fmt.Errorf("create chat request: %w", err)
|
||||||
httpReq.Header.Set("Accept", "text/event-stream")
|
}
|
||||||
if c.apiKey != "" {
|
httpReq.Header.Set("Content-Type", "application/json")
|
||||||
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
|
httpReq.Header.Set("Accept", "text/event-stream")
|
||||||
}
|
if c.apiKey != "" {
|
||||||
|
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
|
||||||
|
}
|
||||||
|
|
||||||
resp, err := c.httpClient.Do(httpReq)
|
resp, err := c.httpClient.Do(httpReq)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("send chat request: %w", err)
|
lastErr = fmt.Errorf("send chat request: %w", err)
|
||||||
}
|
if !shouldRetryRequest(ctx, attempt, 0) {
|
||||||
defer resp.Body.Close()
|
return lastErr
|
||||||
|
}
|
||||||
|
if waitErr := waitBeforeRetry(ctx); waitErr != nil {
|
||||||
|
return waitErr
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||||
limit := io.LimitReader(resp.Body, 4096)
|
limit := io.LimitReader(resp.Body, 4096)
|
||||||
data, _ := io.ReadAll(limit)
|
data, _ := io.ReadAll(limit)
|
||||||
return fmt.Errorf("chat request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data)))
|
_ = resp.Body.Close()
|
||||||
}
|
lastErr = fmt.Errorf("chat request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data)))
|
||||||
|
if !shouldRetryRequest(ctx, attempt, resp.StatusCode) {
|
||||||
|
return lastErr
|
||||||
|
}
|
||||||
|
if waitErr := waitBeforeRetry(ctx); waitErr != nil {
|
||||||
|
return waitErr
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
return readSSE(resp.Body, emit)
|
defer resp.Body.Close()
|
||||||
|
return readSSE(resp.Body, emit)
|
||||||
|
}
|
||||||
|
return lastErr
|
||||||
|
}
|
||||||
|
|
||||||
|
func shouldRetryRequest(ctx context.Context, attempt int, status int) bool {
|
||||||
|
if ctx.Err() != nil || attempt >= maxChatAttempts {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if status == 0 {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return status == http.StatusRequestTimeout || status == http.StatusTooManyRequests || status >= 500
|
||||||
|
}
|
||||||
|
|
||||||
|
func waitBeforeRetry(ctx context.Context) error {
|
||||||
|
timer := time.NewTimer(retryDelay)
|
||||||
|
defer timer.Stop()
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return ctx.Err()
|
||||||
|
case <-timer.C:
|
||||||
|
return nil
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func readSSE(r io.Reader, emit func(StreamEvent) error) error {
|
func readSSE(r io.Reader, emit func(StreamEvent) error) error {
|
||||||
@@ -79,7 +132,7 @@ func readSSE(r io.Reader, emit func(StreamEvent) error) error {
|
|||||||
if parseErr != nil {
|
if parseErr != nil {
|
||||||
return parseErr
|
return parseErr
|
||||||
}
|
}
|
||||||
if event.Content != "" || len(event.ToolCalls) > 0 || event.FinishReason != "" {
|
if event.Content != "" || len(event.ToolCalls) > 0 || event.FinishReason != "" || event.Usage != nil {
|
||||||
if emitErr := emit(event); emitErr != nil {
|
if emitErr := emit(event); emitErr != nil {
|
||||||
return emitErr
|
return emitErr
|
||||||
}
|
}
|
||||||
@@ -100,6 +153,7 @@ type streamPayload struct {
|
|||||||
} `json:"delta"`
|
} `json:"delta"`
|
||||||
FinishReason string `json:"finish_reason"`
|
FinishReason string `json:"finish_reason"`
|
||||||
} `json:"choices"`
|
} `json:"choices"`
|
||||||
|
Usage *Usage `json:"usage,omitempty"`
|
||||||
Error *struct {
|
Error *struct {
|
||||||
Message string `json:"message"`
|
Message string `json:"message"`
|
||||||
Type string `json:"type"`
|
Type string `json:"type"`
|
||||||
@@ -127,15 +181,14 @@ func parseStreamPayload(payload string) (StreamEvent, error) {
|
|||||||
}
|
}
|
||||||
return StreamEvent{}, fmt.Errorf("provider error: %s", decoded.Error.Message)
|
return StreamEvent{}, fmt.Errorf("provider error: %s", decoded.Error.Message)
|
||||||
}
|
}
|
||||||
|
event := StreamEvent{Usage: decoded.Usage}
|
||||||
if len(decoded.Choices) == 0 {
|
if len(decoded.Choices) == 0 {
|
||||||
return StreamEvent{}, nil
|
return event, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
choice := decoded.Choices[0]
|
choice := decoded.Choices[0]
|
||||||
event := StreamEvent{
|
event.Content = choice.Delta.Content
|
||||||
Content: choice.Delta.Content,
|
event.FinishReason = choice.FinishReason
|
||||||
FinishReason: choice.FinishReason,
|
|
||||||
}
|
|
||||||
for _, tc := range choice.Delta.ToolCalls {
|
for _, tc := range choice.Delta.ToolCalls {
|
||||||
event.ToolCalls = append(event.ToolCalls, ToolCallDelta{
|
event.ToolCalls = append(event.ToolCalls, ToolCallDelta{
|
||||||
Index: tc.Index,
|
Index: tc.Index,
|
||||||
@@ -30,6 +30,10 @@ func TestOpenAICompatibleClientStreamsContentAndToolCalls(t *testing.T) {
|
|||||||
if raw["thinking"] != "high" {
|
if raw["thinking"] != "high" {
|
||||||
t.Fatalf("thinking = %#v", raw["thinking"])
|
t.Fatalf("thinking = %#v", raw["thinking"])
|
||||||
}
|
}
|
||||||
|
streamOptions, ok := raw["stream_options"].(map[string]any)
|
||||||
|
if !ok || streamOptions["include_usage"] != true {
|
||||||
|
t.Fatalf("stream_options = %#v", raw["stream_options"])
|
||||||
|
}
|
||||||
if !req.Stream {
|
if !req.Stream {
|
||||||
t.Fatal("request did not enable stream")
|
t.Fatal("request did not enable stream")
|
||||||
}
|
}
|
||||||
@@ -66,6 +70,75 @@ func TestOpenAICompatibleClientStreamsContentAndToolCalls(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestChatStreamCapturesUsage(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
body := mustReadBody(t, r)
|
||||||
|
var raw map[string]any
|
||||||
|
if err := json.NewDecoder(strings.NewReader(body)).Decode(&raw); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
streamOptions, ok := raw["stream_options"].(map[string]any)
|
||||||
|
if !ok || streamOptions["include_usage"] != true {
|
||||||
|
t.Fatalf("stream_options = %#v", raw["stream_options"])
|
||||||
|
}
|
||||||
|
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"ok"},"finish_reason":""}]}` + "\n\n"))
|
||||||
|
_, _ = w.Write([]byte(`data: {"choices":[],"usage":{"prompt_tokens":12,"completion_tokens":3,"total_tokens":15}}` + "\n\n"))
|
||||||
|
_, _ = w.Write([]byte("data: [DONE]\n\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client := NewOpenAICompatibleClient(server.URL, "sk-test", server.Client())
|
||||||
|
var usage *Usage
|
||||||
|
err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(event StreamEvent) error {
|
||||||
|
if event.Usage != nil {
|
||||||
|
usage = event.Usage
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if usage == nil {
|
||||||
|
t.Fatal("usage was not emitted")
|
||||||
|
}
|
||||||
|
if usage.PromptTokens != 12 || usage.CompletionTokens != 3 || usage.TotalTokens != 15 {
|
||||||
|
t.Fatalf("usage = %#v", usage)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAICompatibleClientRetriesTransientHTTPError(t *testing.T) {
|
||||||
|
attempts := 0
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
attempts++
|
||||||
|
if attempts == 1 {
|
||||||
|
http.Error(w, "temporary", http.StatusBadGateway)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"ok"},"finish_reason":""}]}` + "\n\n"))
|
||||||
|
_, _ = w.Write([]byte("data: [DONE]\n\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client := NewOpenAICompatibleClient(server.URL, "sk-test", server.Client())
|
||||||
|
var content strings.Builder
|
||||||
|
err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(event StreamEvent) error {
|
||||||
|
content.WriteString(event.Content)
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if attempts != 2 {
|
||||||
|
t.Fatalf("attempts = %d, want 2", attempts)
|
||||||
|
}
|
||||||
|
if content.String() != "ok" {
|
||||||
|
t.Fatalf("content = %q", content.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func mustReadBody(t *testing.T, r *http.Request) string {
|
func mustReadBody(t *testing.T, r *http.Request) string {
|
||||||
t.Helper()
|
t.Helper()
|
||||||
data, err := io.ReadAll(r.Body)
|
data, err := io.ReadAll(r.Body)
|
||||||
@@ -0,0 +1,43 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import "encoding/json"
|
||||||
|
|
||||||
|
const tokenEstimateCharRatio = 4
|
||||||
|
|
||||||
|
// EstimateTokens estimates the token cost of messages and tool definitions in a
|
||||||
|
// chat completion request. It intentionally uses a simple JSON-size heuristic:
|
||||||
|
// roughly four UTF-8 bytes per token plus a small per-message framing overhead.
|
||||||
|
func EstimateTokens(messages []Message, tools []Tool) int {
|
||||||
|
if len(messages) == 0 && len(tools) == 0 {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
payload := struct {
|
||||||
|
Messages []Message `json:"messages,omitempty"`
|
||||||
|
Tools []Tool `json:"tools,omitempty"`
|
||||||
|
}{
|
||||||
|
Messages: messages,
|
||||||
|
Tools: tools,
|
||||||
|
}
|
||||||
|
data, err := json.Marshal(payload)
|
||||||
|
if err != nil {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
return estimateBytesAsTokens(len(data)) + len(messages)*4
|
||||||
|
}
|
||||||
|
|
||||||
|
func estimateMessageTokens(msg Message) int {
|
||||||
|
data, err := json.Marshal(msg)
|
||||||
|
if err != nil {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
return estimateBytesAsTokens(len(data)) + 4
|
||||||
|
}
|
||||||
|
|
||||||
|
func estimateBytesAsTokens(n int) int {
|
||||||
|
if n <= 0 {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
return (n + tokenEstimateCharRatio - 1) / tokenEstimateCharRatio
|
||||||
|
}
|
||||||
@@ -0,0 +1,45 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import "testing"
|
||||||
|
|
||||||
|
func TestEstimateTokensEmpty(t *testing.T) {
|
||||||
|
if got := EstimateTokens(nil, nil); got != 0 {
|
||||||
|
t.Fatalf("EstimateTokens(nil, nil) = %d, want 0", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEstimateTokens(t *testing.T) {
|
||||||
|
shortMessages := []Message{{Role: RoleUser, Content: "hello"}}
|
||||||
|
short := EstimateTokens(shortMessages, nil)
|
||||||
|
if short <= 4 {
|
||||||
|
t.Fatalf("short estimate = %d, want message framing plus content", short)
|
||||||
|
}
|
||||||
|
if short > 50 {
|
||||||
|
t.Fatalf("short estimate = %d, looks too high for one short message", short)
|
||||||
|
}
|
||||||
|
|
||||||
|
longMessages := []Message{
|
||||||
|
{Role: RoleUser, Content: "hello"},
|
||||||
|
{Role: RoleAssistant, Content: "This is a longer response with enough content to exceed the short estimate."},
|
||||||
|
}
|
||||||
|
long := EstimateTokens(longMessages, nil)
|
||||||
|
if long <= short {
|
||||||
|
t.Fatalf("long estimate = %d, want > short estimate %d", long, short)
|
||||||
|
}
|
||||||
|
|
||||||
|
withTools := EstimateTokens(shortMessages, []Tool{{
|
||||||
|
Type: "function",
|
||||||
|
Function: ToolFunction{
|
||||||
|
Name: "file_read",
|
||||||
|
Description: "Read a file",
|
||||||
|
Parameters: []byte(`{"type":"object"}`),
|
||||||
|
},
|
||||||
|
}})
|
||||||
|
if withTools <= short {
|
||||||
|
t.Fatalf("withTools estimate = %d, want > short estimate %d", withTools, short)
|
||||||
|
}
|
||||||
|
|
||||||
|
if msgTokens := estimateMessageTokens(shortMessages[0]); msgTokens <= 4 {
|
||||||
|
t.Fatalf("estimateMessageTokens = %d, want message framing plus content", msgTokens)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -79,6 +79,13 @@ type StreamEvent struct {
|
|||||||
Content string
|
Content string
|
||||||
ToolCalls []ToolCallDelta
|
ToolCalls []ToolCallDelta
|
||||||
FinishReason string
|
FinishReason string
|
||||||
|
Usage *Usage
|
||||||
|
}
|
||||||
|
|
||||||
|
type Usage struct {
|
||||||
|
PromptTokens int `json:"prompt_tokens"`
|
||||||
|
CompletionTokens int `json:"completion_tokens"`
|
||||||
|
TotalTokens int `json:"total_tokens"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type ToolCallDelta struct {
|
type ToolCallDelta struct {
|
||||||
Reference in New Issue
Block a user