v0.0.1: agentu mvp
Add OpenAI-compatible providers, TUI, slash commands, and model switching. Add local file tools plus optional yolo shell execution. Document config, usage, and development workflow.
This commit is contained in:
loveuer
@@ -0,0 +1,4 @@
|
||||
agentu*.yaml
|
||||
/agentu
|
||||
/coverage.out
|
||||
*.swp
|
||||
@@ -0,0 +1,144 @@
|
||||
# agentu
|
||||
|
||||
agentu is a small local agent written in Go. The current MVP is a terminal chat
|
||||
UI with OpenAI-compatible providers, streaming responses, session-level
|
||||
model/provider switching, and local tools.
|
||||
|
||||
## Current MVP
|
||||
|
||||
- TUI-first chat experience with a light theme by default.
|
||||
- OpenAI-compatible `/v1/chat/completions` provider support.
|
||||
- Multiple providers from `~/.agentu/config.yaml`.
|
||||
- Session-only provider, model, and thinking-level switching.
|
||||
- Read-only project tools enabled by default.
|
||||
- `--yolo` mode for file writes and shell execution.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```sh
|
||||
mkdir -p ~/.agentu
|
||||
cp agentu.example.yaml ~/.agentu/config.yaml
|
||||
export AGENTU_API_KEY='your-api-key'
|
||||
go run ./cmd/agentu
|
||||
```
|
||||
|
||||
Use `--yolo` only when you want agentu to write files or run shell commands:
|
||||
|
||||
```sh
|
||||
go run ./cmd/agentu --yolo
|
||||
```
|
||||
|
||||
## Config
|
||||
|
||||
By default, agentu reads:
|
||||
|
||||
```text
|
||||
~/.agentu/config.yaml
|
||||
```
|
||||
|
||||
Use `--config <path>` to load another file.
|
||||
|
||||
```yaml
|
||||
providers:
|
||||
loveuer:
|
||||
base_url: https://ai.loveuer.com
|
||||
api_key: ${AGENTU_API_KEY}
|
||||
model: deepseek-v4-flash
|
||||
models:
|
||||
- deepseek-v4-flash
|
||||
thinking: none
|
||||
thinking_param: thinking
|
||||
|
||||
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
|
||||
```
|
||||
|
||||
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
|
||||
current session.
|
||||
|
||||
`models` is optional. If present, `/model model <name>` is restricted to that
|
||||
list. If omitted, any model name is accepted for that provider.
|
||||
|
||||
`thinking` is optional. Supported values are:
|
||||
|
||||
```text
|
||||
none, middle, high, xhigh, max
|
||||
```
|
||||
|
||||
When configured or changed with `/model thinking ...`, the value is sent as a
|
||||
top-level OpenAI-compatible request field. `thinking_param` controls that field
|
||||
name and defaults to `thinking`.
|
||||
|
||||
## CLI Flags
|
||||
|
||||
```sh
|
||||
go run ./cmd/agentu [flags]
|
||||
```
|
||||
|
||||
- `--config <path>` loads a YAML config file. Default: `~/.agentu/config.yaml`.
|
||||
- `--theme light|dark` selects the TUI theme. Default: `light`.
|
||||
- `--plain` uses the simple line-based REPL instead of the TUI.
|
||||
- `--yolo` enables automatic file writes and shell execution.
|
||||
|
||||
## TUI Controls
|
||||
|
||||
- `enter` sends the message.
|
||||
- `ctrl+j` or `alt+enter` inserts a newline.
|
||||
- The input starts at one line and grows up to six lines.
|
||||
- `page up` / `page down` scrolls the transcript.
|
||||
- `esc` cancels a running response.
|
||||
- Type `/` to show available commands.
|
||||
|
||||
Slash commands:
|
||||
|
||||
- `/clear` resets conversation history.
|
||||
- `/model` shows current provider, model, thinking, and switch usage.
|
||||
- `/model provider <name>` switches provider 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.
|
||||
- `/exit` or `/quit` exits.
|
||||
|
||||
## Local Tools
|
||||
|
||||
Read-only tools are available by default:
|
||||
|
||||
- `file_read`
|
||||
- `file_list`
|
||||
- `file_search`
|
||||
|
||||
`--yolo` additionally enables:
|
||||
|
||||
- `file_write`
|
||||
- `shell_run`
|
||||
|
||||
Use `--yolo` for prompts that ask agentu to run or inspect local commands, for
|
||||
example:
|
||||
|
||||
```text
|
||||
请帮我运行 `go test ./...` 并总结结果
|
||||
```
|
||||
|
||||
## Development
|
||||
|
||||
Run the verification suite:
|
||||
|
||||
```sh
|
||||
go test ./...
|
||||
go vet ./...
|
||||
go build ./cmd/agentu
|
||||
```
|
||||
|
||||
Local config and build outputs are intentionally ignored by git.
|
||||
|
||||
## Notes
|
||||
|
||||
- Provider/model/thinking changes are session-only; agentu does not write them
|
||||
back to `~/.agentu/config.yaml`.
|
||||
- The config schema is intentionally strict during MVP development. Deprecated
|
||||
fields such as `provider` or `active_provider` are rejected.
|
||||
@@ -0,0 +1,17 @@
|
||||
providers:
|
||||
loveuer:
|
||||
base_url: https://ai.loveuer.com
|
||||
api_key: ${AGENTU_API_KEY}
|
||||
model: deepseek-v4-flash
|
||||
models:
|
||||
- deepseek-v4-flash
|
||||
thinking: none
|
||||
thinking_param: thinking
|
||||
|
||||
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
|
||||
@@ -0,0 +1,129 @@
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
module agentu
|
||||
|
||||
go 1.26
|
||||
|
||||
require (
|
||||
github.com/charmbracelet/bubbles v1.0.0
|
||||
github.com/charmbracelet/bubbletea v1.3.10
|
||||
github.com/charmbracelet/lipgloss v1.1.0
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/atotto/clipboard v0.1.4 // indirect
|
||||
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
|
||||
github.com/charmbracelet/colorprofile v0.4.1 // indirect
|
||||
github.com/charmbracelet/x/ansi v0.11.6 // indirect
|
||||
github.com/charmbracelet/x/cellbuf v0.0.15 // indirect
|
||||
github.com/charmbracelet/x/term v0.2.2 // indirect
|
||||
github.com/clipperhouse/displaywidth v0.9.0 // indirect
|
||||
github.com/clipperhouse/stringish v0.1.1 // indirect
|
||||
github.com/clipperhouse/uax29/v2 v2.5.0 // indirect
|
||||
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/mattn/go-localereader v0.0.1 // indirect
|
||||
github.com/mattn/go-runewidth v0.0.19 // indirect
|
||||
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect
|
||||
github.com/muesli/cancelreader v0.2.2 // indirect
|
||||
github.com/muesli/termenv v0.16.0 // indirect
|
||||
github.com/rivo/uniseg v0.4.7 // indirect
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
|
||||
golang.org/x/sys v0.38.0 // indirect
|
||||
golang.org/x/text v0.3.8 // indirect
|
||||
)
|
||||
@@ -0,0 +1,60 @@
|
||||
github.com/MakeNowJust/heredoc v1.0.0 h1:cXCdzVdstXyiTqTvfqk9SDHpKNjxuom+DOlyEeQ4pzQ=
|
||||
github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE=
|
||||
github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4=
|
||||
github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
|
||||
github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
|
||||
github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
|
||||
github.com/aymanbagabas/go-udiff v0.3.1 h1:LV+qyBQ2pqe0u42ZsUEtPiCaUoqgA9gYRDs3vj1nolY=
|
||||
github.com/aymanbagabas/go-udiff v0.3.1/go.mod h1:G0fsKmG+P6ylD0r6N/KgQD/nWzgfnl8ZBcNLgcbrw8E=
|
||||
github.com/charmbracelet/bubbles v1.0.0 h1:12J8/ak/uCZEMQ6KU7pcfwceyjLlWsDLAxB5fXonfvc=
|
||||
github.com/charmbracelet/bubbles v1.0.0/go.mod h1:9d/Zd5GdnauMI5ivUIVisuEm3ave1XwXtD1ckyV6r3E=
|
||||
github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw=
|
||||
github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4=
|
||||
github.com/charmbracelet/colorprofile v0.4.1 h1:a1lO03qTrSIRaK8c3JRxJDZOvhvIeSco3ej+ngLk1kk=
|
||||
github.com/charmbracelet/colorprofile v0.4.1/go.mod h1:U1d9Dljmdf9DLegaJ0nGZNJvoXAhayhmidOdcBwAvKk=
|
||||
github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY=
|
||||
github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30=
|
||||
github.com/charmbracelet/x/ansi v0.11.6 h1:GhV21SiDz/45W9AnV2R61xZMRri5NlLnl6CVF7ihZW8=
|
||||
github.com/charmbracelet/x/ansi v0.11.6/go.mod h1:2JNYLgQUsyqaiLovhU2Rv/pb8r6ydXKS3NIttu3VGZQ=
|
||||
github.com/charmbracelet/x/cellbuf v0.0.15 h1:ur3pZy0o6z/R7EylET877CBxaiE1Sp1GMxoFPAIztPI=
|
||||
github.com/charmbracelet/x/cellbuf v0.0.15/go.mod h1:J1YVbR7MUuEGIFPCaaZ96KDl5NoS0DAWkskup+mOY+Q=
|
||||
github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk=
|
||||
github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI=
|
||||
github.com/clipperhouse/displaywidth v0.9.0 h1:Qb4KOhYwRiN3viMv1v/3cTBlz3AcAZX3+y9OLhMtAtA=
|
||||
github.com/clipperhouse/displaywidth v0.9.0/go.mod h1:aCAAqTlh4GIVkhQnJpbL0T/WfcrJXHcj8C0yjYcjOZA=
|
||||
github.com/clipperhouse/stringish v0.1.1 h1:+NSqMOr3GR6k1FdRhhnXrLfztGzuG+VuFDfatpWHKCs=
|
||||
github.com/clipperhouse/stringish v0.1.1/go.mod h1:v/WhFtE1q0ovMta2+m+UbpZ+2/HEXNWYXQgCt4hdOzA=
|
||||
github.com/clipperhouse/uax29/v2 v2.5.0 h1:x7T0T4eTHDONxFJsL94uKNKPHrclyFI0lm7+w94cO8U=
|
||||
github.com/clipperhouse/uax29/v2 v2.5.0/go.mod h1:Wn1g7MK6OoeDT0vL+Q0SQLDz/KpfsVRgg6W7ihQeh4g=
|
||||
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4=
|
||||
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM=
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag=
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
|
||||
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
||||
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4=
|
||||
github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88=
|
||||
github.com/mattn/go-runewidth v0.0.19 h1:v++JhqYnZuu5jSKrk9RbgF5v4CGUjqRfBm05byFGLdw=
|
||||
github.com/mattn/go-runewidth v0.0.19/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs=
|
||||
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI=
|
||||
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo=
|
||||
github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA=
|
||||
github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo=
|
||||
github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc=
|
||||
github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk=
|
||||
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
|
||||
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
|
||||
golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI=
|
||||
golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo=
|
||||
golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc=
|
||||
golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
||||
golang.org/x/text v0.3.8 h1:nAL+RVCQ9uMn3vJZbV+MRnydTJFPf8qqY42YiA6MrqY=
|
||||
golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
@@ -0,0 +1,311 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"agentu/internal/llm"
|
||||
"agentu/internal/tools"
|
||||
)
|
||||
|
||||
const defaultMaxToolIterations = 8
|
||||
|
||||
type Agent struct {
|
||||
provider llm.Provider
|
||||
providerName string
|
||||
model string
|
||||
thinking string
|
||||
thinkingParam string
|
||||
thinkingEnabled bool
|
||||
systemPrompt string
|
||||
toolRegistry *tools.Registry
|
||||
toolTimeout time.Duration
|
||||
maxToolIterations int
|
||||
messages []llm.Message
|
||||
}
|
||||
|
||||
type Options struct {
|
||||
Provider llm.Provider
|
||||
ProviderName string
|
||||
Model string
|
||||
Thinking string
|
||||
ThinkingParam string
|
||||
ThinkingEnabled bool
|
||||
SystemPrompt string
|
||||
ToolRegistry *tools.Registry
|
||||
ToolTimeout time.Duration
|
||||
MaxToolIterations int
|
||||
}
|
||||
|
||||
func New(opts Options) *Agent {
|
||||
maxToolIterations := opts.MaxToolIterations
|
||||
if maxToolIterations <= 0 {
|
||||
maxToolIterations = defaultMaxToolIterations
|
||||
}
|
||||
toolTimeout := opts.ToolTimeout
|
||||
if toolTimeout <= 0 {
|
||||
toolTimeout = 30 * time.Second
|
||||
}
|
||||
a := &Agent{
|
||||
provider: opts.Provider,
|
||||
providerName: opts.ProviderName,
|
||||
model: opts.Model,
|
||||
thinking: opts.Thinking,
|
||||
thinkingParam: opts.ThinkingParam,
|
||||
thinkingEnabled: opts.ThinkingEnabled,
|
||||
systemPrompt: opts.SystemPrompt,
|
||||
toolRegistry: opts.ToolRegistry,
|
||||
toolTimeout: toolTimeout,
|
||||
maxToolIterations: maxToolIterations,
|
||||
}
|
||||
a.Clear()
|
||||
return a
|
||||
}
|
||||
|
||||
type RuntimeOptions struct {
|
||||
Provider llm.Provider
|
||||
ProviderName string
|
||||
Model string
|
||||
Thinking string
|
||||
ThinkingParam string
|
||||
ThinkingEnabled bool
|
||||
}
|
||||
|
||||
type RuntimeInfo struct {
|
||||
ProviderName string
|
||||
Model string
|
||||
Thinking string
|
||||
ThinkingEnabled bool
|
||||
}
|
||||
|
||||
func (a *Agent) SetRuntime(opts RuntimeOptions) {
|
||||
if opts.Provider != nil {
|
||||
a.provider = opts.Provider
|
||||
}
|
||||
if opts.ProviderName != "" {
|
||||
a.providerName = opts.ProviderName
|
||||
}
|
||||
if opts.Model != "" {
|
||||
a.model = opts.Model
|
||||
}
|
||||
a.thinking = opts.Thinking
|
||||
a.thinkingParam = opts.ThinkingParam
|
||||
a.thinkingEnabled = opts.ThinkingEnabled
|
||||
}
|
||||
|
||||
func (a *Agent) RuntimeInfo() RuntimeInfo {
|
||||
return RuntimeInfo{
|
||||
ProviderName: a.providerName,
|
||||
Model: a.model,
|
||||
Thinking: a.thinking,
|
||||
ThinkingEnabled: a.thinkingEnabled,
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) Clear() {
|
||||
a.messages = nil
|
||||
if strings.TrimSpace(a.systemPrompt) != "" {
|
||||
a.messages = append(a.messages, llm.Message{
|
||||
Role: llm.RoleSystem,
|
||||
Content: a.systemPrompt,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs io.Writer) error {
|
||||
input = strings.TrimSpace(input)
|
||||
if input == "" {
|
||||
return nil
|
||||
}
|
||||
a.messages = append(a.messages, llm.Message{Role: llm.RoleUser, Content: input})
|
||||
|
||||
for iteration := 0; iteration <= a.maxToolIterations; iteration++ {
|
||||
content, calls, err := a.streamAssistant(ctx, out)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(calls) == 0 {
|
||||
a.messages = append(a.messages, llm.Message{
|
||||
Role: llm.RoleAssistant,
|
||||
Content: content,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
a.messages = append(a.messages, llm.Message{
|
||||
Role: llm.RoleAssistant,
|
||||
Content: content,
|
||||
ToolCalls: calls,
|
||||
})
|
||||
if content != "" {
|
||||
fmt.Fprintln(out)
|
||||
}
|
||||
|
||||
for _, call := range calls {
|
||||
result := a.executeTool(ctx, call, logs)
|
||||
a.messages = append(a.messages, llm.Message{
|
||||
Role: llm.RoleTool,
|
||||
ToolCallID: call.ID,
|
||||
Content: result,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return fmt.Errorf("tool iteration limit reached: %d", a.maxToolIterations)
|
||||
}
|
||||
|
||||
func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []llm.ToolCall, error) {
|
||||
var content strings.Builder
|
||||
calls := make(map[int]*toolCallBuilder)
|
||||
|
||||
req := llm.ChatRequest{
|
||||
Model: a.model,
|
||||
Messages: append([]llm.Message(nil), a.messages...),
|
||||
Tools: a.toolDefinitions(),
|
||||
}
|
||||
if a.thinkingEnabled {
|
||||
thinkingParam := a.thinkingParam
|
||||
if thinkingParam == "" {
|
||||
thinkingParam = "thinking"
|
||||
}
|
||||
req.Extra = map[string]any{thinkingParam: a.thinking}
|
||||
}
|
||||
if len(req.Tools) > 0 {
|
||||
req.ToolChoice = "auto"
|
||||
}
|
||||
|
||||
err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
|
||||
if event.Content != "" {
|
||||
content.WriteString(event.Content)
|
||||
if out != nil {
|
||||
if _, err := io.WriteString(out, event.Content); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, delta := range event.ToolCalls {
|
||||
builder := calls[delta.Index]
|
||||
if builder == nil {
|
||||
builder = &toolCallBuilder{index: delta.Index}
|
||||
calls[delta.Index] = builder
|
||||
}
|
||||
builder.apply(delta)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
|
||||
return content.String(), finalizeToolCalls(calls), nil
|
||||
}
|
||||
|
||||
func (a *Agent) toolDefinitions() []llm.Tool {
|
||||
if a.toolRegistry == nil {
|
||||
return nil
|
||||
}
|
||||
return a.toolRegistry.Definitions()
|
||||
}
|
||||
|
||||
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writer) string {
|
||||
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 {
|
||||
return tools.FormatError(fmt.Errorf("tool registry is disabled"))
|
||||
}
|
||||
tool, ok := a.toolRegistry.Get(name)
|
||||
if !ok {
|
||||
return tools.FormatError(fmt.Errorf("unknown tool: %s", name))
|
||||
}
|
||||
|
||||
args := strings.TrimSpace(call.Function.Arguments)
|
||||
if args == "" {
|
||||
args = "{}"
|
||||
}
|
||||
|
||||
toolCtx, cancel := context.WithTimeout(ctx, a.toolTimeout)
|
||||
defer cancel()
|
||||
output, err := tool.Execute(toolCtx, json.RawMessage(args))
|
||||
if err != nil {
|
||||
if output != "" {
|
||||
return output + "\n" + tools.FormatError(err)
|
||||
}
|
||||
return tools.FormatError(err)
|
||||
}
|
||||
if logs != nil {
|
||||
fmt.Fprintf(logs, "[tool] %s done\n", name)
|
||||
}
|
||||
return output
|
||||
}
|
||||
|
||||
type toolCallBuilder struct {
|
||||
index int
|
||||
id string
|
||||
callType string
|
||||
name string
|
||||
arguments strings.Builder
|
||||
}
|
||||
|
||||
func (b *toolCallBuilder) apply(delta llm.ToolCallDelta) {
|
||||
if delta.ID != "" {
|
||||
b.id = delta.ID
|
||||
}
|
||||
if delta.Type != "" {
|
||||
b.callType = delta.Type
|
||||
}
|
||||
if delta.Name != "" {
|
||||
b.name = delta.Name
|
||||
}
|
||||
if delta.Arguments != "" {
|
||||
b.arguments.WriteString(delta.Arguments)
|
||||
}
|
||||
}
|
||||
|
||||
func finalizeToolCalls(builders map[int]*toolCallBuilder) []llm.ToolCall {
|
||||
if len(builders) == 0 {
|
||||
return nil
|
||||
}
|
||||
indexes := make([]int, 0, len(builders))
|
||||
for index := range builders {
|
||||
indexes = append(indexes, index)
|
||||
}
|
||||
sort.Ints(indexes)
|
||||
|
||||
calls := make([]llm.ToolCall, 0, len(indexes))
|
||||
for _, index := range indexes {
|
||||
builder := builders[index]
|
||||
id := builder.id
|
||||
if id == "" {
|
||||
id = fmt.Sprintf("call_%d", index)
|
||||
}
|
||||
callType := builder.callType
|
||||
if callType == "" {
|
||||
callType = "function"
|
||||
}
|
||||
calls = append(calls, llm.ToolCall{
|
||||
ID: id,
|
||||
Type: callType,
|
||||
Function: llm.FunctionCall{
|
||||
Name: builder.name,
|
||||
Arguments: builder.arguments.String(),
|
||||
},
|
||||
})
|
||||
}
|
||||
return calls
|
||||
}
|
||||
|
||||
func compact(value string, limit int) string {
|
||||
value = strings.TrimSpace(value)
|
||||
if len(value) <= limit {
|
||||
return value
|
||||
}
|
||||
return value[:limit] + "...[truncated]"
|
||||
}
|
||||
@@ -0,0 +1,139 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"agentu/internal/llm"
|
||||
"agentu/internal/tools"
|
||||
)
|
||||
|
||||
type fakeProvider struct {
|
||||
calls int
|
||||
t *testing.T
|
||||
}
|
||||
|
||||
func (p *fakeProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||
p.calls++
|
||||
switch p.calls {
|
||||
case 1:
|
||||
if len(req.Tools) == 0 {
|
||||
p.t.Fatal("expected tools")
|
||||
}
|
||||
if req.ToolChoice != "auto" {
|
||||
p.t.Fatalf("tool_choice = %q", req.ToolChoice)
|
||||
}
|
||||
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
|
||||
{Index: 0, ID: "call_1", Type: "function", Name: "test_echo", Arguments: `{"text":"he`},
|
||||
{Index: 0, Arguments: `llo"}`},
|
||||
}})
|
||||
case 2:
|
||||
last := req.Messages[len(req.Messages)-1]
|
||||
if last.Role != llm.RoleTool || last.Content != "hello" {
|
||||
p.t.Fatalf("last message = %#v", last)
|
||||
}
|
||||
return emit(llm.StreamEvent{Content: "done"})
|
||||
default:
|
||||
p.t.Fatalf("unexpected call count %d", p.calls)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
type echoTool struct{}
|
||||
|
||||
func (echoTool) Definition() llm.Tool {
|
||||
return llm.Tool{
|
||||
Type: "function",
|
||||
Function: llm.ToolFunction{
|
||||
Name: "test_echo",
|
||||
Description: "echo test",
|
||||
Parameters: json.RawMessage(`{"type":"object"}`),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (echoTool) 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
|
||||
}
|
||||
return args.Text, nil
|
||||
}
|
||||
|
||||
func TestAgentRunsToolLoop(t *testing.T) {
|
||||
provider := &fakeProvider{t: t}
|
||||
a := New(Options{
|
||||
Provider: provider,
|
||||
Model: "test",
|
||||
SystemPrompt: "system",
|
||||
ToolRegistry: tools.NewRegistry(echoTool{}),
|
||||
})
|
||||
|
||||
var out strings.Builder
|
||||
var logs strings.Builder
|
||||
if err := a.RunTurn(context.Background(), "use a tool", &out, &logs); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if out.String() != "done" {
|
||||
t.Fatalf("out = %q", out.String())
|
||||
}
|
||||
if !strings.Contains(logs.String(), "[tool] test_echo") {
|
||||
t.Fatalf("logs = %q", logs.String())
|
||||
}
|
||||
}
|
||||
|
||||
type contentProvider struct {
|
||||
called bool
|
||||
text string
|
||||
}
|
||||
|
||||
func (p *contentProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||
p.called = true
|
||||
return emit(llm.StreamEvent{Content: p.text})
|
||||
}
|
||||
|
||||
func TestAgentLetsProviderDecideShellRequestWhenShellDisabled(t *testing.T) {
|
||||
provider := &contentProvider{text: "shell is disabled"}
|
||||
a := New(Options{
|
||||
Provider: provider,
|
||||
Model: "test",
|
||||
SystemPrompt: "system",
|
||||
ToolRegistry: tools.ReadOnlyBuiltins(t.TempDir()),
|
||||
})
|
||||
|
||||
var out strings.Builder
|
||||
if err := a.RunTurn(context.Background(), "能帮我看看 `ssh home-dev` 的运行情况么", &out, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !provider.called {
|
||||
t.Fatal("provider should be called so the model can decide whether to call a tool or answer")
|
||||
}
|
||||
if out.String() != "shell is disabled" {
|
||||
t.Fatalf("out = %q", out.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentDoesNotBlockFileRequest(t *testing.T) {
|
||||
provider := &contentProvider{text: "ok"}
|
||||
a := New(Options{
|
||||
Provider: provider,
|
||||
Model: "test",
|
||||
SystemPrompt: "system",
|
||||
ToolRegistry: tools.ReadOnlyBuiltins(t.TempDir()),
|
||||
})
|
||||
|
||||
var out strings.Builder
|
||||
if err := a.RunTurn(context.Background(), "请读取 `README.md`", &out, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !provider.called {
|
||||
t.Fatal("provider should be called for file requests")
|
||||
}
|
||||
if out.String() != "ok" {
|
||||
t.Fatalf("out = %q", out.String())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
const DefaultPath = "~/.agentu/config.yaml"
|
||||
|
||||
type Config struct {
|
||||
Providers map[string]ProviderConfig `yaml:"providers"`
|
||||
Agent AgentConfig `yaml:"agent"`
|
||||
}
|
||||
|
||||
type ProviderConfig struct {
|
||||
Name string
|
||||
BaseURL string `yaml:"base_url"`
|
||||
Model string `yaml:"model"`
|
||||
Models []string `yaml:"models"`
|
||||
APIKey string `yaml:"api_key"`
|
||||
Thinking string `yaml:"thinking"`
|
||||
ThinkingParam string `yaml:"thinking_param"`
|
||||
ThinkingConfigured bool
|
||||
}
|
||||
|
||||
type AgentConfig struct {
|
||||
SystemPrompt string `yaml:"system_prompt"`
|
||||
WorkingDir string `yaml:"working_dir"`
|
||||
ToolTimeout time.Duration `yaml:"tool_timeout"`
|
||||
}
|
||||
|
||||
type rawConfig struct {
|
||||
Providers map[string]rawProviderConfig `yaml:"providers"`
|
||||
Agent struct {
|
||||
SystemPrompt string `yaml:"system_prompt"`
|
||||
WorkingDir string `yaml:"working_dir"`
|
||||
ToolTimeout string `yaml:"tool_timeout"`
|
||||
} `yaml:"agent"`
|
||||
}
|
||||
|
||||
type rawProviderConfig struct {
|
||||
BaseURL string `yaml:"base_url"`
|
||||
Model string `yaml:"model"`
|
||||
Models []string `yaml:"models"`
|
||||
APIKey string `yaml:"api_key"`
|
||||
Thinking *string `yaml:"thinking"`
|
||||
ThinkingParam string `yaml:"thinking_param"`
|
||||
}
|
||||
|
||||
func Load(path string) (*Config, error) {
|
||||
path, err := ResolvePath(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read config: %w", err)
|
||||
}
|
||||
|
||||
expanded := os.ExpandEnv(string(data))
|
||||
var raw rawConfig
|
||||
decoder := yaml.NewDecoder(strings.NewReader(expanded))
|
||||
decoder.KnownFields(true)
|
||||
if err := decoder.Decode(&raw); err != nil {
|
||||
return nil, fmt.Errorf("parse config: %w", err)
|
||||
}
|
||||
|
||||
cfg := &Config{
|
||||
Providers: normalizeProviders(raw),
|
||||
Agent: AgentConfig{
|
||||
SystemPrompt: raw.Agent.SystemPrompt,
|
||||
WorkingDir: raw.Agent.WorkingDir,
|
||||
},
|
||||
}
|
||||
|
||||
if cfg.Agent.WorkingDir == "" {
|
||||
cfg.Agent.WorkingDir = "."
|
||||
}
|
||||
absWorkingDir, err := filepath.Abs(cfg.Agent.WorkingDir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("resolve working_dir: %w", err)
|
||||
}
|
||||
cfg.Agent.WorkingDir = absWorkingDir
|
||||
|
||||
if raw.Agent.ToolTimeout == "" {
|
||||
cfg.Agent.ToolTimeout = 30 * time.Second
|
||||
} else {
|
||||
timeout, err := time.ParseDuration(raw.Agent.ToolTimeout)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse agent.tool_timeout: %w", err)
|
||||
}
|
||||
cfg.Agent.ToolTimeout = timeout
|
||||
}
|
||||
|
||||
if cfg.Agent.SystemPrompt == "" {
|
||||
cfg.Agent.SystemPrompt = "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."
|
||||
}
|
||||
|
||||
if err := cfg.Validate(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
func ResolvePath(path string) (string, error) {
|
||||
if strings.TrimSpace(path) == "" {
|
||||
path = DefaultPath
|
||||
}
|
||||
if path == "~" {
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("resolve home dir: %w", err)
|
||||
}
|
||||
return home, nil
|
||||
}
|
||||
if strings.HasPrefix(path, "~/") {
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("resolve home dir: %w", err)
|
||||
}
|
||||
return filepath.Join(home, strings.TrimPrefix(path, "~/")), nil
|
||||
}
|
||||
return path, nil
|
||||
}
|
||||
|
||||
func (c *Config) Validate() error {
|
||||
if len(c.Providers) == 0 {
|
||||
return errors.New("missing required config: providers")
|
||||
}
|
||||
for name, provider := range c.Providers {
|
||||
var missing []string
|
||||
prefix := "providers." + name
|
||||
if strings.TrimSpace(provider.BaseURL) == "" {
|
||||
missing = append(missing, prefix+".base_url")
|
||||
}
|
||||
if strings.TrimSpace(provider.Model) == "" {
|
||||
missing = append(missing, prefix+".model")
|
||||
}
|
||||
if strings.TrimSpace(provider.APIKey) == "" {
|
||||
missing = append(missing, prefix+".api_key")
|
||||
}
|
||||
if len(missing) > 0 {
|
||||
return fmt.Errorf("missing required config: %s", strings.Join(missing, ", "))
|
||||
}
|
||||
if provider.ThinkingConfigured && !IsThinkingLevel(provider.Thinking) {
|
||||
return fmt.Errorf("%s.thinking must be one of: %s", prefix, strings.Join(ThinkingLevels(), ", "))
|
||||
}
|
||||
}
|
||||
if c.Agent.ToolTimeout <= 0 {
|
||||
return errors.New("agent.tool_timeout must be greater than zero")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Config) DefaultProviderName() string {
|
||||
names := c.ProviderNames()
|
||||
if len(names) == 0 {
|
||||
return ""
|
||||
}
|
||||
return names[0]
|
||||
}
|
||||
|
||||
func (c *Config) ProviderNames() []string {
|
||||
names := make([]string, 0, len(c.Providers))
|
||||
for name := range c.Providers {
|
||||
names = append(names, name)
|
||||
}
|
||||
sort.Strings(names)
|
||||
return names
|
||||
}
|
||||
|
||||
func normalizeProviders(raw rawConfig) map[string]ProviderConfig {
|
||||
providers := make(map[string]ProviderConfig)
|
||||
for name, provider := range raw.Providers {
|
||||
providers[name] = normalizeProvider(name, provider)
|
||||
}
|
||||
return providers
|
||||
}
|
||||
|
||||
func normalizeProvider(name string, raw rawProviderConfig) ProviderConfig {
|
||||
provider := ProviderConfig{
|
||||
Name: name,
|
||||
BaseURL: raw.BaseURL,
|
||||
Model: raw.Model,
|
||||
Models: append([]string(nil), raw.Models...),
|
||||
APIKey: raw.APIKey,
|
||||
ThinkingParam: raw.ThinkingParam,
|
||||
}
|
||||
if provider.Model == "" && len(provider.Models) > 0 {
|
||||
provider.Model = provider.Models[0]
|
||||
}
|
||||
if raw.Thinking != nil {
|
||||
provider.Thinking = strings.ToLower(strings.TrimSpace(*raw.Thinking))
|
||||
provider.ThinkingConfigured = true
|
||||
if provider.ThinkingParam == "" {
|
||||
provider.ThinkingParam = "thinking"
|
||||
}
|
||||
}
|
||||
return provider
|
||||
}
|
||||
|
||||
func ThinkingLevels() []string {
|
||||
return []string{"none", "middle", "high", "xhigh", "max"}
|
||||
}
|
||||
|
||||
func IsThinkingLevel(value string) bool {
|
||||
value = strings.ToLower(strings.TrimSpace(value))
|
||||
for _, level := range ThinkingLevels() {
|
||||
if value == level {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,257 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestLoadExpandsEnvAndDefaults(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
|
||||
model: test-model
|
||||
api_key: ${AGENTU_TEST_KEY}
|
||||
agent:
|
||||
working_dir: .
|
||||
`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
cfg, err := Load(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
provider := cfg.Providers[cfg.DefaultProviderName()]
|
||||
if provider.APIKey != "sk-test" {
|
||||
t.Fatalf("api key = %q", provider.APIKey)
|
||||
}
|
||||
if providerName := cfg.DefaultProviderName(); providerName != "test" {
|
||||
t.Fatalf("default provider = %q", providerName)
|
||||
}
|
||||
if cfg.Agent.ToolTimeout != 30*time.Second {
|
||||
t.Fatalf("tool timeout = %s", cfg.Agent.ToolTimeout)
|
||||
}
|
||||
if !filepath.IsAbs(cfg.Agent.WorkingDir) {
|
||||
t.Fatalf("working dir is not absolute: %s", cfg.Agent.WorkingDir)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadUsesDefaultConfigPath(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
t.Setenv("HOME", dir)
|
||||
t.Setenv("AGENTU_TEST_KEY", "sk-test")
|
||||
configDir := filepath.Join(dir, ".agentu")
|
||||
if err := os.MkdirAll(configDir, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
path := filepath.Join(configDir, "config.yaml")
|
||||
if err := os.WriteFile(path, []byte(`
|
||||
providers:
|
||||
test:
|
||||
base_url: https://example.com
|
||||
model: test-model
|
||||
api_key: ${AGENTU_TEST_KEY}
|
||||
`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
cfg, err := Load("")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := cfg.DefaultProviderName(); got != "test" {
|
||||
t.Fatalf("default provider = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolvePathExpandsHome(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
t.Setenv("HOME", dir)
|
||||
|
||||
got, err := ResolvePath("~/.agentu/config.yaml")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := filepath.Join(dir, ".agentu", "config.yaml")
|
||||
if got != want {
|
||||
t.Fatalf("path = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadDefaultSystemPromptDiscouragesEagerTools(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
|
||||
model: test-model
|
||||
api_key: ${AGENTU_TEST_KEY}
|
||||
`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
cfg, err := Load(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, want := range []string{"Answer simple conversation directly", "Use tools only when"} {
|
||||
if !strings.Contains(cfg.Agent.SystemPrompt, want) {
|
||||
t.Fatalf("system prompt missing %q: %s", want, cfg.Agent.SystemPrompt)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadMultiProviderConfig(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:
|
||||
loveuer:
|
||||
base_url: https://example.com
|
||||
api_key: ${AGENTU_TEST_KEY}
|
||||
model: model-a
|
||||
models:
|
||||
- model-a
|
||||
- model-b
|
||||
thinking: xhigh
|
||||
thinking_param: thinking
|
||||
backup:
|
||||
base_url: https://backup.example.com
|
||||
api_key: ${AGENTU_TEST_KEY}
|
||||
models:
|
||||
- model-c
|
||||
`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
cfg, err := Load(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := cfg.DefaultProviderName(); got != "backup" {
|
||||
t.Fatalf("default provider = %q", got)
|
||||
}
|
||||
provider := cfg.Providers["loveuer"]
|
||||
if got := provider.Model; got != "model-a" {
|
||||
t.Fatalf("active model = %q", got)
|
||||
}
|
||||
if got := cfg.Providers["backup"].Model; got != "model-c" {
|
||||
t.Fatalf("backup default model = %q", got)
|
||||
}
|
||||
if !provider.ThinkingConfigured || provider.Thinking != "xhigh" || provider.ThinkingParam != "thinking" {
|
||||
t.Fatalf("thinking config = %#v", provider)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadDefaultProviderNameIsDeterministic(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:
|
||||
two:
|
||||
base_url: https://two.example.com
|
||||
api_key: ${AGENTU_TEST_KEY}
|
||||
model: model-b
|
||||
one:
|
||||
base_url: https://one.example.com
|
||||
api_key: ${AGENTU_TEST_KEY}
|
||||
model: model-a
|
||||
`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
cfg, err := Load(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := cfg.DefaultProviderName(); got != "one" {
|
||||
t.Fatalf("default provider = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadRejectsInvalidThinking(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
|
||||
model: test-model
|
||||
api_key: ${AGENTU_TEST_KEY}
|
||||
thinking: huge
|
||||
`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err := Load(path)
|
||||
if err == nil || !strings.Contains(err.Error(), "thinking") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadRequiresProviderFields(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
|
||||
`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err := Load(path)
|
||||
if err == nil {
|
||||
t.Fatal("expected error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadRejectsLegacyProviderConfig(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "agentu.yaml")
|
||||
if err := os.WriteFile(path, []byte(`
|
||||
provider:
|
||||
base_url: https://example.com
|
||||
model: test-model
|
||||
api_key: sk-test
|
||||
`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err := Load(path)
|
||||
if err == nil || !strings.Contains(err.Error(), "field provider not found") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadRejectsActiveProviderConfig(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "agentu.yaml")
|
||||
if err := os.WriteFile(path, []byte(`
|
||||
active_provider: loveuer
|
||||
providers:
|
||||
loveuer:
|
||||
base_url: https://example.com
|
||||
model: test-model
|
||||
api_key: sk-test
|
||||
`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err := Load(path)
|
||||
if err == nil || !strings.Contains(err.Error(), "field active_provider not found") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type OpenAICompatibleClient struct {
|
||||
baseURL string
|
||||
apiKey string
|
||||
httpClient *http.Client
|
||||
}
|
||||
|
||||
func NewOpenAICompatibleClient(baseURL, apiKey string, httpClient *http.Client) *OpenAICompatibleClient {
|
||||
if httpClient == nil {
|
||||
httpClient = http.DefaultClient
|
||||
}
|
||||
return &OpenAICompatibleClient{
|
||||
baseURL: strings.TrimRight(baseURL, "/"),
|
||||
apiKey: apiKey,
|
||||
httpClient: httpClient,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *OpenAICompatibleClient) ChatStream(ctx context.Context, req ChatRequest, emit func(StreamEvent) error) error {
|
||||
req.Stream = true
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal chat request: %w", err)
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return fmt.Errorf("create chat request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Accept", "text/event-stream")
|
||||
if c.apiKey != "" {
|
||||
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
|
||||
}
|
||||
|
||||
resp, err := c.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
return fmt.Errorf("send chat request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
limit := io.LimitReader(resp.Body, 4096)
|
||||
data, _ := io.ReadAll(limit)
|
||||
return fmt.Errorf("chat request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data)))
|
||||
}
|
||||
|
||||
return readSSE(resp.Body, emit)
|
||||
}
|
||||
|
||||
func readSSE(r io.Reader, emit func(StreamEvent) error) error {
|
||||
reader := bufio.NewReader(r)
|
||||
for {
|
||||
line, err := reader.ReadString('\n')
|
||||
if err != nil && !errors.Is(err, io.EOF) {
|
||||
return fmt.Errorf("read stream: %w", err)
|
||||
}
|
||||
|
||||
line = strings.TrimSpace(line)
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if payload == "[DONE]" {
|
||||
return nil
|
||||
}
|
||||
event, parseErr := parseStreamPayload(payload)
|
||||
if parseErr != nil {
|
||||
return parseErr
|
||||
}
|
||||
if event.Content != "" || len(event.ToolCalls) > 0 || event.FinishReason != "" {
|
||||
if emitErr := emit(event); emitErr != nil {
|
||||
return emitErr
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if errors.Is(err, io.EOF) {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type streamPayload struct {
|
||||
Choices []struct {
|
||||
Delta struct {
|
||||
Content string `json:"content"`
|
||||
ToolCalls []toolCallDeltaJSON `json:"tool_calls"`
|
||||
} `json:"delta"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
Type string `json:"type"`
|
||||
} `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type toolCallDeltaJSON struct {
|
||||
Index int `json:"index"`
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
}
|
||||
|
||||
func parseStreamPayload(payload string) (StreamEvent, error) {
|
||||
var decoded streamPayload
|
||||
if err := json.Unmarshal([]byte(payload), &decoded); err != nil {
|
||||
return StreamEvent{}, fmt.Errorf("parse stream payload: %w", err)
|
||||
}
|
||||
if decoded.Error != nil {
|
||||
if decoded.Error.Type != "" {
|
||||
return StreamEvent{}, fmt.Errorf("provider error: %s: %s", decoded.Error.Type, decoded.Error.Message)
|
||||
}
|
||||
return StreamEvent{}, fmt.Errorf("provider error: %s", decoded.Error.Message)
|
||||
}
|
||||
if len(decoded.Choices) == 0 {
|
||||
return StreamEvent{}, nil
|
||||
}
|
||||
|
||||
choice := decoded.Choices[0]
|
||||
event := StreamEvent{
|
||||
Content: choice.Delta.Content,
|
||||
FinishReason: choice.FinishReason,
|
||||
}
|
||||
for _, tc := range choice.Delta.ToolCalls {
|
||||
event.ToolCalls = append(event.ToolCalls, ToolCallDelta{
|
||||
Index: tc.Index,
|
||||
ID: tc.ID,
|
||||
Type: tc.Type,
|
||||
Name: tc.Function.Name,
|
||||
Arguments: tc.Function.Arguments,
|
||||
})
|
||||
}
|
||||
return event, nil
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestOpenAICompatibleClientStreamsContentAndToolCalls(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/chat/completions" {
|
||||
t.Fatalf("path = %s", r.URL.Path)
|
||||
}
|
||||
if got := r.Header.Get("Authorization"); got != "Bearer sk-test" {
|
||||
t.Fatalf("authorization = %q", got)
|
||||
}
|
||||
body := mustReadBody(t, r)
|
||||
var raw map[string]any
|
||||
if err := json.NewDecoder(strings.NewReader(body)).Decode(&raw); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var req ChatRequest
|
||||
if err := json.NewDecoder(strings.NewReader(body)).Decode(&req); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if raw["thinking"] != "high" {
|
||||
t.Fatalf("thinking = %#v", raw["thinking"])
|
||||
}
|
||||
if !req.Stream {
|
||||
t.Fatal("request did not enable stream")
|
||||
}
|
||||
if req.ToolChoice != "auto" && len(req.Tools) > 0 {
|
||||
t.Fatalf("tool_choice = %q", req.ToolChoice)
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"hi "},"finish_reason":""}]}` + "\n\n"))
|
||||
_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"file_read","arguments":"{\"path\""}}]},"finish_reason":""}]}` + "\n\n"))
|
||||
_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":":\"README.md\"}"}}]},"finish_reason":"tool_calls"}]}` + "\n\n"))
|
||||
_, _ = w.Write([]byte("data: [DONE]\n\n"))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewOpenAICompatibleClient(server.URL, "sk-test", server.Client())
|
||||
var content strings.Builder
|
||||
var calls []ToolCallDelta
|
||||
err := client.ChatStream(context.Background(), ChatRequest{Model: "test", Extra: map[string]any{"thinking": "high"}}, func(event StreamEvent) error {
|
||||
content.WriteString(event.Content)
|
||||
calls = append(calls, event.ToolCalls...)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if content.String() != "hi " {
|
||||
t.Fatalf("content = %q", content.String())
|
||||
}
|
||||
if len(calls) != 2 {
|
||||
t.Fatalf("calls = %#v", calls)
|
||||
}
|
||||
if calls[0].Name != "file_read" || calls[1].Arguments != ":\"README.md\"}" {
|
||||
t.Fatalf("unexpected calls: %#v", calls)
|
||||
}
|
||||
}
|
||||
|
||||
func mustReadBody(t *testing.T, r *http.Request) string {
|
||||
t.Helper()
|
||||
data, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return string(data)
|
||||
}
|
||||
|
||||
func TestOpenAICompatibleClientHTTPError(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Error(w, "bad key", http.StatusUnauthorized)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewOpenAICompatibleClient(server.URL, "sk-test", server.Client())
|
||||
err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(StreamEvent) error {
|
||||
return nil
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "status=401") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
)
|
||||
|
||||
const (
|
||||
RoleSystem = "system"
|
||||
RoleUser = "user"
|
||||
RoleAssistant = "assistant"
|
||||
RoleTool = "tool"
|
||||
)
|
||||
|
||||
type Provider interface {
|
||||
ChatStream(ctx context.Context, req ChatRequest, emit func(StreamEvent) error) error
|
||||
}
|
||||
|
||||
type ChatRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []Message `json:"messages"`
|
||||
Tools []Tool `json:"tools,omitempty"`
|
||||
ToolChoice string `json:"tool_choice,omitempty"`
|
||||
Stream bool `json:"stream"`
|
||||
Extra map[string]any `json:"-"`
|
||||
}
|
||||
|
||||
func (r ChatRequest) MarshalJSON() ([]byte, error) {
|
||||
payload := map[string]any{
|
||||
"model": r.Model,
|
||||
"messages": r.Messages,
|
||||
"stream": r.Stream,
|
||||
}
|
||||
if len(r.Tools) > 0 {
|
||||
payload["tools"] = r.Tools
|
||||
}
|
||||
if r.ToolChoice != "" {
|
||||
payload["tool_choice"] = r.ToolChoice
|
||||
}
|
||||
for key, value := range r.Extra {
|
||||
if _, exists := payload[key]; exists {
|
||||
continue
|
||||
}
|
||||
payload[key] = value
|
||||
}
|
||||
return json.Marshal(payload)
|
||||
}
|
||||
|
||||
type Message struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content,omitempty"`
|
||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||
}
|
||||
|
||||
type ToolCall struct {
|
||||
ID string `json:"id,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Function FunctionCall `json:"function,omitempty"`
|
||||
}
|
||||
|
||||
type FunctionCall struct {
|
||||
Name string `json:"name,omitempty"`
|
||||
Arguments string `json:"arguments,omitempty"`
|
||||
}
|
||||
|
||||
type Tool struct {
|
||||
Type string `json:"type"`
|
||||
Function ToolFunction `json:"function"`
|
||||
}
|
||||
|
||||
type ToolFunction struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Parameters json.RawMessage `json:"parameters"`
|
||||
}
|
||||
|
||||
type StreamEvent struct {
|
||||
Content string
|
||||
ToolCalls []ToolCallDelta
|
||||
FinishReason string
|
||||
}
|
||||
|
||||
type ToolCallDelta struct {
|
||||
Index int
|
||||
ID string
|
||||
Type string
|
||||
Name string
|
||||
Arguments string
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"agentu/internal/agent"
|
||||
"agentu/internal/config"
|
||||
"agentu/internal/llm"
|
||||
)
|
||||
|
||||
type Manager struct {
|
||||
cfg *config.Config
|
||||
agent *agent.Agent
|
||||
httpClient *http.Client
|
||||
providerName string
|
||||
provider config.ProviderConfig
|
||||
}
|
||||
|
||||
func NewManager(cfg *config.Config, assistant *agent.Agent, httpClient *http.Client) (*Manager, error) {
|
||||
if httpClient == nil {
|
||||
httpClient = http.DefaultClient
|
||||
}
|
||||
providerName := cfg.DefaultProviderName()
|
||||
provider, ok := cfg.Providers[providerName]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no provider is configured")
|
||||
}
|
||||
return &Manager{
|
||||
cfg: cfg,
|
||||
agent: assistant,
|
||||
httpClient: httpClient,
|
||||
providerName: providerName,
|
||||
provider: provider,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (m *Manager) CurrentProvider() string {
|
||||
return m.providerName
|
||||
}
|
||||
|
||||
func (m *Manager) CurrentModel() string {
|
||||
return m.provider.Model
|
||||
}
|
||||
|
||||
func (m *Manager) CurrentThinking() string {
|
||||
if !m.provider.ThinkingConfigured {
|
||||
return "unset"
|
||||
}
|
||||
return m.provider.Thinking
|
||||
}
|
||||
|
||||
func (m *Manager) Info() string {
|
||||
var b strings.Builder
|
||||
fmt.Fprintf(&b, "provider: %s\n", m.CurrentProvider())
|
||||
fmt.Fprintf(&b, "model: %s\n", m.CurrentModel())
|
||||
fmt.Fprintf(&b, "thinking: %s\n", m.CurrentThinking())
|
||||
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, ", "))
|
||||
} else {
|
||||
b.WriteString("models: any model name is accepted for this provider\n")
|
||||
}
|
||||
fmt.Fprintf(&b, "thinking levels: %s\n", strings.Join(config.ThinkingLevels(), ", "))
|
||||
b.WriteString("usage: /model provider <name> | /model model <name> | /model thinking <level>")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func (m *Manager) SetProvider(name string) (string, error) {
|
||||
provider, ok := m.cfg.Providers[name]
|
||||
if !ok {
|
||||
return "", fmt.Errorf("unknown provider %q; available: %s", name, strings.Join(m.providerNames(), ", "))
|
||||
}
|
||||
m.providerName = name
|
||||
m.provider = provider
|
||||
m.syncAgent(llm.NewOpenAICompatibleClient(provider.BaseURL, provider.APIKey, m.httpClient))
|
||||
return "Switched provider.\n" + m.summary(), nil
|
||||
}
|
||||
|
||||
func (m *Manager) SetModel(model string) (string, error) {
|
||||
model = strings.TrimSpace(model)
|
||||
if model == "" {
|
||||
return "", fmt.Errorf("model is required")
|
||||
}
|
||||
if len(m.provider.Models) > 0 && !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, ", "))
|
||||
}
|
||||
m.provider.Model = model
|
||||
m.cfg.Providers[m.providerName] = m.provider
|
||||
m.syncAgent(nil)
|
||||
return "Switched model.\n" + m.summary(), nil
|
||||
}
|
||||
|
||||
func (m *Manager) SetThinking(level string) (string, error) {
|
||||
level = strings.ToLower(strings.TrimSpace(level))
|
||||
if !config.IsThinkingLevel(level) {
|
||||
return "", fmt.Errorf("thinking must be one of: %s", strings.Join(config.ThinkingLevels(), ", "))
|
||||
}
|
||||
m.provider.Thinking = level
|
||||
m.provider.ThinkingConfigured = true
|
||||
if m.provider.ThinkingParam == "" {
|
||||
m.provider.ThinkingParam = "thinking"
|
||||
}
|
||||
m.cfg.Providers[m.providerName] = m.provider
|
||||
m.syncAgent(nil)
|
||||
return "Updated thinking.\n" + m.summary(), nil
|
||||
}
|
||||
|
||||
func (m *Manager) summary() string {
|
||||
return fmt.Sprintf("provider: %s\nmodel: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentThinking())
|
||||
}
|
||||
|
||||
func (m *Manager) syncAgent(provider llm.Provider) {
|
||||
m.agent.SetRuntime(agent.RuntimeOptions{
|
||||
Provider: provider,
|
||||
ProviderName: m.providerName,
|
||||
Model: m.provider.Model,
|
||||
Thinking: m.provider.Thinking,
|
||||
ThinkingParam: m.provider.ThinkingParam,
|
||||
ThinkingEnabled: m.provider.ThinkingConfigured,
|
||||
})
|
||||
}
|
||||
|
||||
func (m *Manager) providerNames() []string {
|
||||
names := make([]string, 0, len(m.cfg.Providers))
|
||||
for name := range m.cfg.Providers {
|
||||
names = append(names, name)
|
||||
}
|
||||
sort.Strings(names)
|
||||
return names
|
||||
}
|
||||
|
||||
func contains(values []string, value string) bool {
|
||||
for _, item := range values {
|
||||
if item == value {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"agentu/internal/agent"
|
||||
"agentu/internal/config"
|
||||
"agentu/internal/llm"
|
||||
)
|
||||
|
||||
type recordingProvider struct {
|
||||
req llm.ChatRequest
|
||||
}
|
||||
|
||||
func (p *recordingProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||
p.req = req
|
||||
return emit(llm.StreamEvent{Content: "ok"})
|
||||
}
|
||||
|
||||
func TestManagerSwitchesModelAndThinking(t *testing.T) {
|
||||
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", "model-b"},
|
||||
Thinking: "none",
|
||||
ThinkingParam: "thinking",
|
||||
ThinkingConfigured: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
provider := &recordingProvider{}
|
||||
assistant := agent.New(agent.Options{
|
||||
Provider: provider,
|
||||
ProviderName: "one",
|
||||
Model: "model-a",
|
||||
Thinking: "none",
|
||||
ThinkingParam: "thinking",
|
||||
ThinkingEnabled: true,
|
||||
})
|
||||
manager, err := NewManager(cfg, assistant, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if _, err := manager.SetModel("model-b"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := manager.SetThinking("high"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var out strings.Builder
|
||||
if err := assistant.RunTurn(context.Background(), "hi", &out, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if provider.req.Model != "model-b" {
|
||||
t.Fatalf("model = %q", provider.req.Model)
|
||||
}
|
||||
if provider.req.Extra["thinking"] != "high" {
|
||||
t.Fatalf("thinking = %#v", provider.req.Extra["thinking"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerRejectsUnknownModelAndThinking(t *testing.T) {
|
||||
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"},
|
||||
},
|
||||
},
|
||||
}
|
||||
assistant := agent.New(agent.Options{Provider: &recordingProvider{}, Model: "model-a"})
|
||||
manager, err := NewManager(cfg, assistant, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if _, err := manager.SetModel("model-b"); err == nil {
|
||||
t.Fatal("expected unknown model error")
|
||||
}
|
||||
if _, err := manager.SetThinking("huge"); err == nil {
|
||||
t.Fatal("expected invalid thinking error")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,331 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io/fs"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"agentu/internal/llm"
|
||||
)
|
||||
|
||||
type FileReadTool struct {
|
||||
workingDir string
|
||||
}
|
||||
|
||||
func NewFileReadTool(workingDir string) *FileReadTool {
|
||||
return &FileReadTool{workingDir: workingDir}
|
||||
}
|
||||
|
||||
func (t *FileReadTool) Definition() llm.Tool {
|
||||
return llm.Tool{
|
||||
Type: "function",
|
||||
Function: llm.ToolFunction{
|
||||
Name: "file_read",
|
||||
Description: "Read a local file as UTF-8 text.",
|
||||
Parameters: JSONSchema(`{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "File path to read. Relative paths are resolved from the configured agent working directory."},
|
||||
"max_bytes": {"type": "integer", "description": "Optional maximum bytes to return."}
|
||||
},
|
||||
"required": ["path"],
|
||||
"additionalProperties": false
|
||||
}`),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
type fileReadArgs struct {
|
||||
Path string `json:"path"`
|
||||
MaxBytes int `json:"max_bytes"`
|
||||
}
|
||||
|
||||
func (t *FileReadTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||
var args fileReadArgs
|
||||
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")
|
||||
}
|
||||
path, err := resolvePath(t.workingDir, args.Path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
limit := MaxToolOutputBytes
|
||||
if args.MaxBytes > 0 && args.MaxBytes < limit {
|
||||
limit = args.MaxBytes
|
||||
}
|
||||
if len(data) > limit {
|
||||
data = append(data[:limit], []byte(fmt.Sprintf("\n\n[truncated after %d bytes]", limit))...)
|
||||
}
|
||||
return string(data), ctx.Err()
|
||||
}
|
||||
|
||||
type FileWriteTool struct {
|
||||
workingDir string
|
||||
}
|
||||
|
||||
func NewFileWriteTool(workingDir string) *FileWriteTool {
|
||||
return &FileWriteTool{workingDir: workingDir}
|
||||
}
|
||||
|
||||
func (t *FileWriteTool) Definition() llm.Tool {
|
||||
return llm.Tool{
|
||||
Type: "function",
|
||||
Function: llm.ToolFunction{
|
||||
Name: "file_write",
|
||||
Description: "Write UTF-8 text to a local file. Creates parent directories when needed.",
|
||||
Parameters: JSONSchema(`{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "File path to write. Relative paths are resolved from the configured agent working directory."},
|
||||
"content": {"type": "string", "description": "Text content to write."},
|
||||
"append": {"type": "boolean", "description": "Append instead of overwriting."}
|
||||
},
|
||||
"required": ["path", "content"],
|
||||
"additionalProperties": false
|
||||
}`),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
type fileWriteArgs struct {
|
||||
Path string `json:"path"`
|
||||
Content string `json:"content"`
|
||||
Append bool `json:"append"`
|
||||
}
|
||||
|
||||
func (t *FileWriteTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||
var args fileWriteArgs
|
||||
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")
|
||||
}
|
||||
path, err := resolvePath(t.workingDir, args.Path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
return "", err
|
||||
}
|
||||
flag := os.O_CREATE | os.O_WRONLY | os.O_TRUNC
|
||||
if args.Append {
|
||||
flag = os.O_CREATE | os.O_WRONLY | os.O_APPEND
|
||||
}
|
||||
file, err := os.OpenFile(path, flag, 0o644)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer file.Close()
|
||||
if _, err := file.WriteString(args.Content); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
action := "wrote"
|
||||
if args.Append {
|
||||
action = "appended"
|
||||
}
|
||||
return fmt.Sprintf("%s %d bytes to %s", action, len(args.Content), path), nil
|
||||
}
|
||||
|
||||
type FileListTool struct {
|
||||
workingDir string
|
||||
}
|
||||
|
||||
func NewFileListTool(workingDir string) *FileListTool {
|
||||
return &FileListTool{workingDir: workingDir}
|
||||
}
|
||||
|
||||
func (t *FileListTool) Definition() llm.Tool {
|
||||
return llm.Tool{
|
||||
Type: "function",
|
||||
Function: llm.ToolFunction{
|
||||
Name: "file_list",
|
||||
Description: "List files and directories at a local path.",
|
||||
Parameters: JSONSchema(`{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "Directory path to list. Defaults to the configured agent working directory."}
|
||||
},
|
||||
"additionalProperties": false
|
||||
}`),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
type fileListArgs struct {
|
||||
Path string `json:"path"`
|
||||
}
|
||||
|
||||
func (t *FileListTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||
var args fileListArgs
|
||||
if len(raw) > 0 {
|
||||
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
|
||||
}
|
||||
entries, err := os.ReadDir(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
lines := make([]string, 0, len(entries))
|
||||
for _, entry := range entries {
|
||||
name := entry.Name()
|
||||
if entry.IsDir() {
|
||||
name += "/"
|
||||
}
|
||||
lines = append(lines, name)
|
||||
}
|
||||
sort.Strings(lines)
|
||||
if err := ctx.Err(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return FormatResult(strings.Join(lines, "\n")), nil
|
||||
}
|
||||
|
||||
type FileSearchTool struct {
|
||||
workingDir string
|
||||
}
|
||||
|
||||
func NewFileSearchTool(workingDir string) *FileSearchTool {
|
||||
return &FileSearchTool{workingDir: workingDir}
|
||||
}
|
||||
|
||||
func (t *FileSearchTool) Definition() llm.Tool {
|
||||
return llm.Tool{
|
||||
Type: "function",
|
||||
Function: llm.ToolFunction{
|
||||
Name: "file_search",
|
||||
Description: "Search local text files for a pattern. Uses ripgrep when available.",
|
||||
Parameters: JSONSchema(`{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"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."},
|
||||
"glob": {"type": "string", "description": "Optional file glob, for example '*.go'."}
|
||||
},
|
||||
"required": ["query"],
|
||||
"additionalProperties": false
|
||||
}`),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
type fileSearchArgs struct {
|
||||
Query string `json:"query"`
|
||||
Path string `json:"path"`
|
||||
Glob string `json:"glob"`
|
||||
}
|
||||
|
||||
func (t *FileSearchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||
var args fileSearchArgs
|
||||
if err := json.Unmarshal(raw, &args); err != nil {
|
||||
return "", fmt.Errorf("parse arguments: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(args.Query) == "" {
|
||||
return "", errors.New("query is required")
|
||||
}
|
||||
path, err := resolvePath(t.workingDir, args.Path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err := exec.LookPath("rg"); err == nil {
|
||||
return t.ripgrep(ctx, args, path)
|
||||
}
|
||||
return t.walkSearch(ctx, args, path)
|
||||
}
|
||||
|
||||
func (t *FileSearchTool) ripgrep(ctx context.Context, args fileSearchArgs, path string) (string, error) {
|
||||
rgArgs := []string{"--line-number", "--color", "never"}
|
||||
if args.Glob != "" {
|
||||
rgArgs = append(rgArgs, "--glob", args.Glob)
|
||||
}
|
||||
rgArgs = append(rgArgs, args.Query, path)
|
||||
cmd := exec.CommandContext(ctx, "rg", rgArgs...)
|
||||
output, err := cmd.CombinedOutput()
|
||||
text := FormatResult(string(output))
|
||||
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
||||
return text, fmt.Errorf("search timed out: %w", ctx.Err())
|
||||
}
|
||||
if err != nil {
|
||||
var exitErr *exec.ExitError
|
||||
if errors.As(err, &exitErr) && exitErr.ExitCode() == 1 {
|
||||
return "no matches", nil
|
||||
}
|
||||
if text != "" {
|
||||
return text, fmt.Errorf("search failed: %w", err)
|
||||
}
|
||||
return "", fmt.Errorf("search failed: %w", err)
|
||||
}
|
||||
return text, nil
|
||||
}
|
||||
|
||||
func (t *FileSearchTool) walkSearch(ctx context.Context, args fileSearchArgs, root string) (string, error) {
|
||||
var out bytes.Buffer
|
||||
err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
if ctxErr := ctx.Err(); ctxErr != nil {
|
||||
return ctxErr
|
||||
}
|
||||
if d.IsDir() {
|
||||
name := d.Name()
|
||||
if name == ".git" || name == "vendor" || name == "node_modules" {
|
||||
return filepath.SkipDir
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if args.Glob != "" {
|
||||
match, err := filepath.Match(args.Glob, filepath.Base(path))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !match {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil || bytes.IndexByte(data, 0) >= 0 {
|
||||
return nil
|
||||
}
|
||||
lines := strings.Split(string(data), "\n")
|
||||
for i, line := range lines {
|
||||
if strings.Contains(line, args.Query) {
|
||||
fmt.Fprintf(&out, "%s:%d:%s\n", path, i+1, line)
|
||||
if out.Len() > MaxToolOutputBytes {
|
||||
return errors.New("output limit reached")
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil && err.Error() != "output limit reached" {
|
||||
return "", err
|
||||
}
|
||||
if out.Len() == 0 {
|
||||
return "no matches", nil
|
||||
}
|
||||
return FormatResult(out.String()), nil
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
)
|
||||
|
||||
func resolvePath(baseDir, path string) (string, error) {
|
||||
if path == "" {
|
||||
path = "."
|
||||
}
|
||||
if filepath.IsAbs(path) {
|
||||
return filepath.Clean(path), nil
|
||||
}
|
||||
return filepath.Abs(filepath.Join(baseDir, path))
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"strings"
|
||||
|
||||
"agentu/internal/llm"
|
||||
)
|
||||
|
||||
type ShellRunTool struct {
|
||||
workingDir string
|
||||
}
|
||||
|
||||
func NewShellRunTool(workingDir string) *ShellRunTool {
|
||||
return &ShellRunTool{workingDir: workingDir}
|
||||
}
|
||||
|
||||
func (t *ShellRunTool) Definition() llm.Tool {
|
||||
return llm.Tool{
|
||||
Type: "function",
|
||||
Function: llm.ToolFunction{
|
||||
Name: "shell_run",
|
||||
Description: "Run a non-interactive shell command on the local machine. Use for builds, tests, local status inspection, and local automation. Avoid commands that require an interactive TTY, credentials, or open long-running sessions.",
|
||||
Parameters: JSONSchema(`{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {"type": "string", "description": "Command to execute."},
|
||||
"workdir": {"type": "string", "description": "Optional working directory. Relative paths are resolved from the configured agent working directory."}
|
||||
},
|
||||
"required": ["command"],
|
||||
"additionalProperties": false
|
||||
}`),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
type shellRunArgs struct {
|
||||
Command string `json:"command"`
|
||||
Workdir string `json:"workdir"`
|
||||
}
|
||||
|
||||
func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||
var args shellRunArgs
|
||||
if err := json.Unmarshal(raw, &args); err != nil {
|
||||
return "", fmt.Errorf("parse arguments: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(args.Command) == "" {
|
||||
return "", errors.New("command is required")
|
||||
}
|
||||
|
||||
workdir, err := resolvePath(t.workingDir, args.Workdir)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
shell := os.Getenv("SHELL")
|
||||
if shell == "" {
|
||||
shell = "/bin/sh"
|
||||
}
|
||||
cmd := exec.CommandContext(ctx, shell, "-lc", args.Command)
|
||||
cmd.Dir = workdir
|
||||
|
||||
output, err := cmd.CombinedOutput()
|
||||
text := FormatResult(string(output))
|
||||
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
||||
return text, fmt.Errorf("command timed out: %w", ctx.Err())
|
||||
}
|
||||
if err != nil {
|
||||
if text != "" {
|
||||
return text, fmt.Errorf("command failed: %w", err)
|
||||
}
|
||||
return "", fmt.Errorf("command failed: %w", err)
|
||||
}
|
||||
return text, nil
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
|
||||
"agentu/internal/llm"
|
||||
)
|
||||
|
||||
const MaxToolOutputBytes = 64 * 1024
|
||||
|
||||
type Tool interface {
|
||||
Definition() llm.Tool
|
||||
Execute(ctx context.Context, args json.RawMessage) (string, error)
|
||||
}
|
||||
|
||||
type Registry struct {
|
||||
tools map[string]Tool
|
||||
primary []string
|
||||
}
|
||||
|
||||
func NewRegistry(items ...Tool) *Registry {
|
||||
r := &Registry{tools: make(map[string]Tool, len(items))}
|
||||
for _, item := range items {
|
||||
r.Register(item)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
func (r *Registry) Register(tool Tool) {
|
||||
name := tool.Definition().Function.Name
|
||||
if _, exists := r.tools[name]; !exists {
|
||||
r.primary = append(r.primary, name)
|
||||
}
|
||||
r.tools[name] = tool
|
||||
}
|
||||
|
||||
func (r *Registry) RegisterAlias(alias string, tool Tool) {
|
||||
r.tools[alias] = tool
|
||||
}
|
||||
|
||||
func (r *Registry) Get(name string) (Tool, bool) {
|
||||
tool, ok := r.tools[name]
|
||||
return tool, ok
|
||||
}
|
||||
|
||||
func (r *Registry) Definitions() []llm.Tool {
|
||||
defs := make([]llm.Tool, 0, len(r.primary))
|
||||
for _, name := range r.primary {
|
||||
tool := r.tools[name]
|
||||
defs = append(defs, tool.Definition())
|
||||
}
|
||||
sort.Slice(defs, func(i, j int) bool {
|
||||
return defs[i].Function.Name < defs[j].Function.Name
|
||||
})
|
||||
return defs
|
||||
}
|
||||
|
||||
func ReadOnlyBuiltins(workingDir string) *Registry {
|
||||
fileRead := NewFileReadTool(workingDir)
|
||||
fileList := NewFileListTool(workingDir)
|
||||
fileSearch := NewFileSearchTool(workingDir)
|
||||
|
||||
registry := NewRegistry(fileRead, fileList, fileSearch)
|
||||
registry.RegisterAlias("file.read", fileRead)
|
||||
registry.RegisterAlias("file.list", fileList)
|
||||
registry.RegisterAlias("file.search", fileSearch)
|
||||
return registry
|
||||
}
|
||||
|
||||
func Builtins(workingDir string) *Registry {
|
||||
registry := ReadOnlyBuiltins(workingDir)
|
||||
fileWrite := NewFileWriteTool(workingDir)
|
||||
shellRun := NewShellRunTool(workingDir)
|
||||
registry.Register(fileWrite)
|
||||
registry.Register(shellRun)
|
||||
registry.RegisterAlias("file.write", fileWrite)
|
||||
registry.RegisterAlias("shell.run", shellRun)
|
||||
return registry
|
||||
}
|
||||
|
||||
func AgentInstructions(workingDir string, yolo bool) string {
|
||||
if yolo {
|
||||
return fmt.Sprintf(`Local project context:
|
||||
- The project working directory is %q.
|
||||
- Tools are available, but use them only when needed to satisfy the user's request.
|
||||
- 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.
|
||||
- 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 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.
|
||||
- Do not use file tools as a substitute for a command execution request.
|
||||
- Do not start interactive or long-running sessions. Prefer bounded, non-interactive local commands with clear completion conditions.
|
||||
- Tool paths should normally be relative to the project working directory.
|
||||
- Available tools: file_list, file_read, file_search, file_write, shell_run.`, workingDir)
|
||||
}
|
||||
return fmt.Sprintf(`Local project context:
|
||||
- The project working directory is %q.
|
||||
- Read-only project file tools are available, but use them only when needed to satisfy the user's request.
|
||||
- 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.
|
||||
- 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.
|
||||
- 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.
|
||||
- Tool paths should normally be relative to the project working directory.
|
||||
- Available tools: file_list, file_read, file_search.`, workingDir)
|
||||
}
|
||||
|
||||
func JSONSchema(schema string) json.RawMessage {
|
||||
return json.RawMessage(schema)
|
||||
}
|
||||
|
||||
func FormatResult(output string) string {
|
||||
if len(output) <= MaxToolOutputBytes {
|
||||
return output
|
||||
}
|
||||
return output[:MaxToolOutputBytes] + fmt.Sprintf("\n\n[truncated after %d bytes]", MaxToolOutputBytes)
|
||||
}
|
||||
|
||||
func FormatError(err error) string {
|
||||
return "error: " + err.Error()
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestFileTools(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
ctx := context.Background()
|
||||
|
||||
write := NewFileWriteTool(dir)
|
||||
if out, err := write.Execute(ctx, json.RawMessage(`{"path":"notes/a.txt","content":"hello world"}`)); err != nil {
|
||||
t.Fatal(err)
|
||||
} else if !strings.Contains(out, "wrote 11 bytes") {
|
||||
t.Fatalf("write output = %q", out)
|
||||
}
|
||||
|
||||
read := NewFileReadTool(dir)
|
||||
if out, err := read.Execute(ctx, json.RawMessage(`{"path":"notes/a.txt"}`)); err != nil {
|
||||
t.Fatal(err)
|
||||
} else if out != "hello world" {
|
||||
t.Fatalf("read output = %q", out)
|
||||
}
|
||||
|
||||
list := NewFileListTool(dir)
|
||||
if out, err := list.Execute(ctx, json.RawMessage(`{"path":"notes"}`)); err != nil {
|
||||
t.Fatal(err)
|
||||
} else if out != "a.txt" {
|
||||
t.Fatalf("list output = %q", out)
|
||||
}
|
||||
|
||||
search := NewFileSearchTool(dir)
|
||||
if out, err := search.Execute(ctx, json.RawMessage(`{"query":"hello","path":"notes"}`)); err != nil {
|
||||
t.Fatal(err)
|
||||
} else if !strings.Contains(out, "hello world") {
|
||||
t.Fatalf("search output = %q", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShellRunTool(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(dir, "x.txt"), []byte("x"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
tool := NewShellRunTool(dir)
|
||||
out, err := tool.Execute(context.Background(), json.RawMessage(`{"command":"pwd && ls"}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(out, dir) || !strings.Contains(out, "x.txt") {
|
||||
t.Fatalf("output = %q", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadOnlyBuiltinsExposeOpenAICompatibleNamesAndAliases(t *testing.T) {
|
||||
registry := ReadOnlyBuiltins(t.TempDir())
|
||||
var names []string
|
||||
for _, def := range registry.Definitions() {
|
||||
names = append(names, def.Function.Name)
|
||||
if strings.Contains(def.Function.Name, ".") {
|
||||
t.Fatalf("tool name contains dot: %s", def.Function.Name)
|
||||
}
|
||||
}
|
||||
|
||||
for _, want := range []string{"file_list", "file_read", "file_search"} {
|
||||
if !slices.Contains(names, want) {
|
||||
t.Fatalf("definitions missing %s: %#v", want, names)
|
||||
}
|
||||
}
|
||||
for _, alias := range []string{"file.list", "file.read", "file.search"} {
|
||||
if _, ok := registry.Get(alias); !ok {
|
||||
t.Fatalf("alias missing: %s", alias)
|
||||
}
|
||||
}
|
||||
if _, ok := registry.Get("file_write"); ok {
|
||||
t.Fatal("read-only registry should not expose file_write")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuiltinsExposeYoloTools(t *testing.T) {
|
||||
registry := Builtins(t.TempDir())
|
||||
for _, name := range []string{"file_write", "shell_run", "file.write", "shell.run"} {
|
||||
if _, ok := registry.Get(name); !ok {
|
||||
t.Fatalf("tool missing: %s", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentInstructionsExplainCommandModes(t *testing.T) {
|
||||
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"} {
|
||||
if !strings.Contains(readOnly, want) {
|
||||
t.Fatalf("read-only instructions missing %q:\n%s", want, readOnly)
|
||||
}
|
||||
}
|
||||
|
||||
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"} {
|
||||
if !strings.Contains(yolo, want) {
|
||||
t.Fatalf("yolo instructions missing %q:\n%s", want, yolo)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,622 @@
|
||||
package tui
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"strings"
|
||||
|
||||
"agentu/internal/agent"
|
||||
|
||||
"github.com/charmbracelet/bubbles/textarea"
|
||||
"github.com/charmbracelet/bubbles/viewport"
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
"github.com/charmbracelet/lipgloss"
|
||||
)
|
||||
|
||||
const (
|
||||
minInputLines = 1
|
||||
maxInputLines = 6
|
||||
)
|
||||
|
||||
type Options struct {
|
||||
ModelName string
|
||||
Yolo bool
|
||||
ThemeMode ThemeMode
|
||||
ModelManager ModelManager
|
||||
}
|
||||
|
||||
type ModelManager interface {
|
||||
CurrentProvider() string
|
||||
CurrentModel() string
|
||||
CurrentThinking() string
|
||||
Info() string
|
||||
SetProvider(name string) (string, error)
|
||||
SetModel(model string) (string, error)
|
||||
SetThinking(level string) (string, error)
|
||||
}
|
||||
|
||||
func Run(ctx context.Context, assistant *agent.Agent, opts Options) error {
|
||||
model := newModel(ctx, assistant, opts)
|
||||
_, err := tea.NewProgram(model, tea.WithAltScreen()).Run()
|
||||
return err
|
||||
}
|
||||
|
||||
type role string
|
||||
|
||||
const (
|
||||
roleUser role = "user"
|
||||
roleAssistant role = "assistant"
|
||||
roleTool role = "tool"
|
||||
roleError role = "error"
|
||||
roleSystem role = "system"
|
||||
)
|
||||
|
||||
type message struct {
|
||||
role role
|
||||
content string
|
||||
}
|
||||
|
||||
type model struct {
|
||||
ctx context.Context
|
||||
agent *agent.Agent
|
||||
modelName string
|
||||
yolo bool
|
||||
models ModelManager
|
||||
theme theme
|
||||
styles styles
|
||||
|
||||
viewport viewport.Model
|
||||
input textarea.Model
|
||||
|
||||
width int
|
||||
height int
|
||||
|
||||
messages []message
|
||||
status string
|
||||
running bool
|
||||
cancel context.CancelFunc
|
||||
events <-chan tea.Msg
|
||||
}
|
||||
|
||||
type assistantChunkMsg string
|
||||
type toolLogMsg string
|
||||
type turnStartedMsg struct{}
|
||||
type turnDoneMsg struct {
|
||||
err error
|
||||
}
|
||||
|
||||
type slashCommand struct {
|
||||
Name string
|
||||
Description string
|
||||
}
|
||||
|
||||
var slashCommands = []slashCommand{
|
||||
{Name: "/clear", Description: "reset context"},
|
||||
{Name: "/exit", Description: "quit"},
|
||||
{Name: "/quit", Description: "quit"},
|
||||
{Name: "/model", Description: "model/provider/thinking"},
|
||||
}
|
||||
|
||||
func newModel(ctx context.Context, assistant *agent.Agent, opts Options) model {
|
||||
mode := opts.ThemeMode
|
||||
if mode == "" {
|
||||
mode = ThemeLight
|
||||
}
|
||||
th := themeForMode(mode)
|
||||
st := th.styles()
|
||||
|
||||
input := textarea.New()
|
||||
input.Placeholder = "Message agentu..."
|
||||
input.Prompt = ""
|
||||
input.ShowLineNumbers = false
|
||||
input.MaxHeight = maxInputLines
|
||||
input.SetHeight(minInputLines)
|
||||
input.SetWidth(80)
|
||||
applyTextareaTheme(&input, th)
|
||||
input.Focus()
|
||||
|
||||
vp := viewport.New(80, 20)
|
||||
vp.Style = st.Viewport
|
||||
vp.SetContent("")
|
||||
|
||||
return model{
|
||||
ctx: ctx,
|
||||
agent: assistant,
|
||||
modelName: opts.ModelName,
|
||||
yolo: opts.Yolo,
|
||||
models: opts.ModelManager,
|
||||
theme: th,
|
||||
styles: st,
|
||||
viewport: vp,
|
||||
input: input,
|
||||
status: "ready",
|
||||
messages: []message{
|
||||
{role: roleSystem, content: "Ask anything. Enter sends, ctrl+j or alt+enter inserts a newline, esc cancels a running turn."},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (m model) Init() tea.Cmd {
|
||||
return textarea.Blink
|
||||
}
|
||||
|
||||
func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
var cmds []tea.Cmd
|
||||
|
||||
switch msg := msg.(type) {
|
||||
case tea.WindowSizeMsg:
|
||||
m.width = msg.Width
|
||||
m.height = msg.Height
|
||||
m.resize()
|
||||
m.refreshViewport(true)
|
||||
return m, nil
|
||||
|
||||
case tea.KeyMsg:
|
||||
switch msg.String() {
|
||||
case "ctrl+c":
|
||||
if m.cancel != nil {
|
||||
m.cancel()
|
||||
}
|
||||
return m, tea.Quit
|
||||
case "esc":
|
||||
if m.running && m.cancel != nil {
|
||||
m.cancel()
|
||||
m.status = "canceling..."
|
||||
return m, nil
|
||||
}
|
||||
case "enter":
|
||||
if m.running {
|
||||
return m, nil
|
||||
}
|
||||
return m.submit()
|
||||
case "alt+enter", "ctrl+j":
|
||||
return m.insertInputNewline()
|
||||
}
|
||||
|
||||
case assistantChunkMsg:
|
||||
m.appendAssistantChunk(string(msg))
|
||||
m.status = "streaming..."
|
||||
m.refreshViewport(true)
|
||||
return m, m.waitForEvent()
|
||||
|
||||
case toolLogMsg:
|
||||
text := cleanToolLog(string(msg))
|
||||
if text != "" {
|
||||
m.messages = append(m.messages, message{role: roleTool, content: text})
|
||||
m.status = text
|
||||
m.refreshViewport(true)
|
||||
}
|
||||
return m, m.waitForEvent()
|
||||
|
||||
case turnDoneMsg:
|
||||
m.running = false
|
||||
m.events = nil
|
||||
m.cancel = nil
|
||||
if msg.err != nil {
|
||||
m.messages = append(m.messages, message{role: roleError, content: msg.err.Error()})
|
||||
m.status = "error"
|
||||
} else {
|
||||
m.status = "ready"
|
||||
}
|
||||
m.refreshViewport(true)
|
||||
cmds = append(cmds, textarea.Blink)
|
||||
}
|
||||
|
||||
if keyMsg, ok := msg.(tea.KeyMsg); ok {
|
||||
if isViewportKey(keyMsg.String()) {
|
||||
nextViewport, cmd := m.viewport.Update(msg)
|
||||
m.viewport = nextViewport
|
||||
return m, cmd
|
||||
}
|
||||
if !m.running {
|
||||
nextInput, cmd := m.input.Update(msg)
|
||||
m.input = nextInput
|
||||
m.afterInputChanged()
|
||||
return m, cmd
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
if !m.running {
|
||||
nextInput, cmd := m.input.Update(msg)
|
||||
m.input = nextInput
|
||||
m.afterInputChanged()
|
||||
cmds = append(cmds, cmd)
|
||||
}
|
||||
|
||||
nextViewport, cmd := m.viewport.Update(msg)
|
||||
m.viewport = nextViewport
|
||||
cmds = append(cmds, cmd)
|
||||
|
||||
return m, tea.Batch(cmds...)
|
||||
}
|
||||
|
||||
func (m model) View() string {
|
||||
if m.width <= 0 || m.height <= 0 {
|
||||
return "agentu"
|
||||
}
|
||||
|
||||
header := m.headerView()
|
||||
status := m.statusView()
|
||||
input := m.inputView()
|
||||
body := lipgloss.JoinVertical(lipgloss.Left, header, m.viewport.View(), input, status)
|
||||
return m.styles.App.Width(m.width).Height(m.height).Render(body)
|
||||
}
|
||||
|
||||
func (m *model) submit() (tea.Model, tea.Cmd) {
|
||||
input := strings.TrimSpace(m.input.Value())
|
||||
if input == "" {
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
switch input {
|
||||
case "/exit", "/quit":
|
||||
return *m, tea.Quit
|
||||
case "/clear":
|
||||
m.agent.Clear()
|
||||
m.messages = []message{{role: roleSystem, content: "Context cleared."}}
|
||||
m.input.SetValue("")
|
||||
m.afterInputChanged()
|
||||
m.status = "ready"
|
||||
m.refreshViewport(true)
|
||||
return *m, nil
|
||||
case "/model":
|
||||
return m.runLocalCommand(input)
|
||||
}
|
||||
|
||||
if strings.HasPrefix(input, "/") {
|
||||
return m.runLocalCommand(input)
|
||||
}
|
||||
|
||||
m.messages = append(m.messages, message{role: roleUser, content: input})
|
||||
m.input.SetValue("")
|
||||
m.afterInputChanged()
|
||||
m.running = true
|
||||
m.status = "thinking..."
|
||||
m.refreshViewport(true)
|
||||
|
||||
turnCtx, cancel := context.WithCancel(m.ctx)
|
||||
m.cancel = cancel
|
||||
|
||||
events := make(chan tea.Msg, 256)
|
||||
m.events = events
|
||||
|
||||
start := func() tea.Msg {
|
||||
go func() {
|
||||
defer close(events)
|
||||
out := eventWriter{ctx: turnCtx, events: events, assistant: true}
|
||||
logs := eventWriter{ctx: turnCtx, events: events}
|
||||
err := m.agent.RunTurn(turnCtx, input, out, logs)
|
||||
if errors.Is(err, context.Canceled) {
|
||||
err = nil
|
||||
}
|
||||
select {
|
||||
case events <- turnDoneMsg{err: err}:
|
||||
case <-turnCtx.Done():
|
||||
}
|
||||
}()
|
||||
return turnStartedMsg{}
|
||||
}
|
||||
|
||||
return *m, tea.Batch(start, m.waitForEvent())
|
||||
}
|
||||
|
||||
func (m *model) insertInputNewline() (tea.Model, tea.Cmd) {
|
||||
enter := tea.KeyMsg{Type: tea.KeyEnter}
|
||||
nextInput, cmd := m.input.Update(enter)
|
||||
m.input = nextInput
|
||||
m.afterInputChanged()
|
||||
return *m, cmd
|
||||
}
|
||||
|
||||
func (m model) waitForEvent() tea.Cmd {
|
||||
events := m.events
|
||||
if events == nil {
|
||||
return nil
|
||||
}
|
||||
return func() tea.Msg {
|
||||
msg, ok := <-events
|
||||
if !ok {
|
||||
return turnDoneMsg{}
|
||||
}
|
||||
return msg
|
||||
}
|
||||
}
|
||||
|
||||
func (m *model) appendAssistantChunk(chunk string) {
|
||||
if chunk == "" {
|
||||
return
|
||||
}
|
||||
if len(m.messages) == 0 || m.messages[len(m.messages)-1].role != roleAssistant {
|
||||
m.messages = append(m.messages, message{role: roleAssistant})
|
||||
}
|
||||
m.messages[len(m.messages)-1].content += chunk
|
||||
}
|
||||
|
||||
func (m *model) resize() {
|
||||
width := max(40, m.width)
|
||||
inputWidth := max(20, width-4)
|
||||
m.input.SetWidth(inputWidth)
|
||||
m.syncInputHeight()
|
||||
|
||||
headerHeight := 1
|
||||
statusHeight := 1
|
||||
inputHeight := m.inputBlockHeight()
|
||||
m.viewport.Width = width
|
||||
m.viewport.Height = max(4, m.height-headerHeight-statusHeight-inputHeight)
|
||||
m.viewport.Style = m.styles.Viewport
|
||||
}
|
||||
|
||||
func (m *model) refreshViewport(bottom bool) {
|
||||
m.viewport.SetContent(m.renderMessages())
|
||||
if bottom {
|
||||
m.viewport.GotoBottom()
|
||||
}
|
||||
}
|
||||
|
||||
func (m model) renderMessages() string {
|
||||
width := max(40, m.width)
|
||||
contentWidth := max(20, width-6)
|
||||
parts := make([]string, 0, len(m.messages))
|
||||
for _, msg := range m.messages {
|
||||
content := strings.TrimRight(msg.content, "\n")
|
||||
if content == "" {
|
||||
continue
|
||||
}
|
||||
parts = append(parts, m.messageView(msg.role, content, contentWidth))
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
return m.styles.Muted.Width(contentWidth).Render("Start a conversation.")
|
||||
}
|
||||
return strings.Join(parts, "\n\n")
|
||||
}
|
||||
|
||||
func (m model) messageView(r role, content string, width int) string {
|
||||
label := string(r)
|
||||
style := m.styles.AssistantMsg
|
||||
labelStyle := m.styles.AssistantLabel
|
||||
|
||||
switch r {
|
||||
case roleUser:
|
||||
label = "you"
|
||||
style = m.styles.UserMessage
|
||||
labelStyle = m.styles.UserLabel
|
||||
case roleTool:
|
||||
label = "tool"
|
||||
style = m.styles.ToolMessage
|
||||
labelStyle = m.styles.ToolLabel
|
||||
case roleError:
|
||||
label = "error"
|
||||
style = m.styles.ErrorMessage
|
||||
labelStyle = m.styles.ErrorLabel
|
||||
case roleSystem:
|
||||
label = "note"
|
||||
style = m.styles.SystemMessage
|
||||
labelStyle = m.styles.SystemLabel
|
||||
default:
|
||||
label = "agentu"
|
||||
}
|
||||
|
||||
labelLine := labelStyle.Render(label)
|
||||
body := style.Width(width).Render(content)
|
||||
return lipgloss.JoinVertical(lipgloss.Left, labelLine, body)
|
||||
}
|
||||
|
||||
func (m model) headerView() string {
|
||||
mode := "read-only tools"
|
||||
if m.yolo {
|
||||
mode = "yolo tools"
|
||||
}
|
||||
modelName := m.modelName
|
||||
thinking := ""
|
||||
if m.models != nil {
|
||||
modelName = m.models.CurrentModel()
|
||||
if currentThinking := m.models.CurrentThinking(); currentThinking != "" && currentThinking != "unset" {
|
||||
thinking = " · thinking " + currentThinking
|
||||
}
|
||||
}
|
||||
title := m.styles.Title.Render("agentu")
|
||||
meta := m.styles.Muted.Render(fmt.Sprintf("%s · %s%s · %s", modelName, mode, thinking, m.theme.Mode))
|
||||
line := lipgloss.JoinHorizontal(lipgloss.Center, title, " ", meta)
|
||||
return m.styles.Header.Width(max(1, m.width)).Render(line)
|
||||
}
|
||||
|
||||
func (m model) inputView() string {
|
||||
title := m.styles.Muted.Render("Enter sends · ctrl+j newline · alt+enter newline · /clear · /exit")
|
||||
if m.running {
|
||||
title = m.styles.Muted.Render("Working · esc cancel · ctrl+c quit")
|
||||
}
|
||||
suggestions := m.slashSuggestionsView()
|
||||
box := m.styles.InputBox.Width(max(20, m.width-2)).Render(m.input.View())
|
||||
if suggestions != "" {
|
||||
return lipgloss.JoinVertical(lipgloss.Left, suggestions, title, box)
|
||||
}
|
||||
return lipgloss.JoinVertical(lipgloss.Left, title, box)
|
||||
}
|
||||
|
||||
func (m model) statusView() string {
|
||||
left := m.status
|
||||
if left == "" {
|
||||
left = "ready"
|
||||
}
|
||||
return m.styles.Status.Width(max(1, m.width)).Render(left)
|
||||
}
|
||||
|
||||
type eventWriter struct {
|
||||
ctx context.Context
|
||||
events chan<- tea.Msg
|
||||
assistant bool
|
||||
}
|
||||
|
||||
func (w eventWriter) Write(p []byte) (int, error) {
|
||||
if len(p) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
text := string(append([]byte(nil), p...))
|
||||
var msg tea.Msg
|
||||
if w.assistant {
|
||||
msg = assistantChunkMsg(text)
|
||||
} else {
|
||||
msg = toolLogMsg(text)
|
||||
}
|
||||
select {
|
||||
case w.events <- msg:
|
||||
return len(p), nil
|
||||
case <-w.ctx.Done():
|
||||
return 0, w.ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func cleanToolLog(text string) string {
|
||||
text = strings.TrimSpace(text)
|
||||
text = strings.TrimPrefix(text, "[tool]")
|
||||
return strings.TrimSpace(text)
|
||||
}
|
||||
|
||||
func isViewportKey(key string) bool {
|
||||
switch key {
|
||||
case "pgup", "pgdown", "ctrl+up", "ctrl+down":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (m *model) afterInputChanged() {
|
||||
m.syncInputHeight()
|
||||
if m.width > 0 && m.height > 0 {
|
||||
m.resize()
|
||||
m.refreshViewport(true)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *model) syncInputHeight() {
|
||||
m.input.SetHeight(m.desiredInputHeight())
|
||||
}
|
||||
|
||||
func (m model) desiredInputHeight() int {
|
||||
value := m.input.Value()
|
||||
if value == "" {
|
||||
return minInputLines
|
||||
}
|
||||
|
||||
textWidth := max(1, m.input.Width())
|
||||
lines := 0
|
||||
for _, line := range strings.Split(value, "\n") {
|
||||
displayWidth := lipgloss.Width(line)
|
||||
wrapped := int(math.Ceil(float64(max(1, displayWidth)) / float64(textWidth)))
|
||||
lines += max(1, wrapped)
|
||||
}
|
||||
return min(max(lines, minInputLines), maxInputLines)
|
||||
}
|
||||
|
||||
func (m model) inputBlockHeight() int {
|
||||
height := 1 + m.input.Height() + m.styles.InputBox.GetVerticalFrameSize()
|
||||
if m.showSlashSuggestions() {
|
||||
height++
|
||||
}
|
||||
return height
|
||||
}
|
||||
|
||||
func (m model) showSlashSuggestions() bool {
|
||||
value := strings.TrimSpace(m.input.Value())
|
||||
return strings.HasPrefix(value, "/") && !strings.Contains(value, " ")
|
||||
}
|
||||
|
||||
func (m model) slashSuggestionsView() string {
|
||||
if !m.showSlashSuggestions() {
|
||||
return ""
|
||||
}
|
||||
prefix := strings.TrimSpace(m.input.Value())
|
||||
var parts []string
|
||||
for _, command := range slashCommands {
|
||||
if !strings.HasPrefix(command.Name, prefix) {
|
||||
continue
|
||||
}
|
||||
parts = append(parts, fmt.Sprintf("%s %s", command.Name, command.Description))
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
parts = append(parts, "No matching commands")
|
||||
}
|
||||
return m.styles.CommandHint.Width(max(20, m.width-2)).Render(strings.Join(parts, " "))
|
||||
}
|
||||
|
||||
func (m model) modelInfo() string {
|
||||
if m.models != nil {
|
||||
return m.models.Info()
|
||||
}
|
||||
mode := "read-only tools"
|
||||
if m.yolo {
|
||||
mode = "yolo tools"
|
||||
}
|
||||
return fmt.Sprintf("model: %s\nmode: %s\ntheme: %s", m.modelName, mode, m.theme.Mode)
|
||||
}
|
||||
|
||||
func (m *model) runLocalCommand(input string) (tea.Model, tea.Cmd) {
|
||||
output, err := m.handleLocalCommand(input)
|
||||
role := roleSystem
|
||||
status := "command"
|
||||
if err != nil {
|
||||
output = err.Error()
|
||||
role = roleError
|
||||
status = "command error"
|
||||
}
|
||||
m.messages = append(m.messages, message{role: role, content: output})
|
||||
m.input.SetValue("")
|
||||
m.afterInputChanged()
|
||||
m.status = status
|
||||
m.refreshViewport(true)
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
func (m model) handleLocalCommand(input string) (string, error) {
|
||||
fields := strings.Fields(input)
|
||||
if len(fields) == 0 {
|
||||
return "", fmt.Errorf("empty command")
|
||||
}
|
||||
switch fields[0] {
|
||||
case "/model":
|
||||
return m.handleModelCommand(fields[1:])
|
||||
default:
|
||||
return "", fmt.Errorf("unknown command: %s", fields[0])
|
||||
}
|
||||
}
|
||||
|
||||
func (m model) handleModelCommand(args []string) (string, error) {
|
||||
if m.models == nil {
|
||||
if len(args) > 0 {
|
||||
return "", fmt.Errorf("model switching is not configured")
|
||||
}
|
||||
return m.modelInfo(), nil
|
||||
}
|
||||
if len(args) == 0 || args[0] == "help" {
|
||||
return m.models.Info(), nil
|
||||
}
|
||||
if len(args) == 1 {
|
||||
return m.models.SetModel(args[0])
|
||||
}
|
||||
switch args[0] {
|
||||
case "provider":
|
||||
if len(args) != 2 {
|
||||
return "", fmt.Errorf("usage: /model provider <name>")
|
||||
}
|
||||
return m.models.SetProvider(args[1])
|
||||
case "model":
|
||||
if len(args) != 2 {
|
||||
return "", fmt.Errorf("usage: /model model <name>")
|
||||
}
|
||||
return m.models.SetModel(args[1])
|
||||
case "thinking":
|
||||
if len(args) != 2 {
|
||||
return "", fmt.Errorf("usage: /model thinking <none|middle|high|xhigh|max>")
|
||||
}
|
||||
return m.models.SetThinking(args[1])
|
||||
default:
|
||||
return "", fmt.Errorf("unknown /model command %q", args[0])
|
||||
}
|
||||
}
|
||||
|
||||
var _ io.Writer = eventWriter{}
|
||||
@@ -0,0 +1,163 @@
|
||||
package tui
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
)
|
||||
|
||||
func TestModelRendersChatShell(t *testing.T) {
|
||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
||||
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||
rendered := next.(model).View()
|
||||
|
||||
for _, want := range []string{"agentu", "test-model", "yolo tools", "light", "ctrl+j", "alt+enter"} {
|
||||
if !strings.Contains(rendered, want) {
|
||||
t.Fatalf("rendered view missing %q:\n%s", want, rendered)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCtrlJInsertsInputNewline(t *testing.T) {
|
||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||
m.input.SetValue("hello")
|
||||
|
||||
next, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlJ})
|
||||
got := next.(model).input.Value()
|
||||
if got != "hello\n" {
|
||||
t.Fatalf("input value = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInputStartsAtOneLineAndGrows(t *testing.T) {
|
||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||
if got := m.input.Height(); got != 1 {
|
||||
t.Fatalf("initial input height = %d", got)
|
||||
}
|
||||
|
||||
m.input.SetValue("one\ntwo\nthree")
|
||||
m.afterInputChanged()
|
||||
if got := m.input.Height(); got != 3 {
|
||||
t.Fatalf("multiline input height = %d", got)
|
||||
}
|
||||
|
||||
m.input.SetValue(strings.Repeat("line\n", 20))
|
||||
m.afterInputChanged()
|
||||
if got := m.input.Height(); got != maxInputLines {
|
||||
t.Fatalf("clamped input height = %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlashCommandSuggestions(t *testing.T) {
|
||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
||||
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||
m = next.(model)
|
||||
m.input.SetValue("/")
|
||||
m.afterInputChanged()
|
||||
|
||||
rendered := m.View()
|
||||
for _, want := range []string{"/clear reset context", "/exit quit", "/model model/provider/thinking"} {
|
||||
if !strings.Contains(rendered, want) {
|
||||
t.Fatalf("rendered suggestions missing %q:\n%s", want, rendered)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelCommand(t *testing.T) {
|
||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
||||
m.input.SetValue("/model")
|
||||
|
||||
next, _ := m.submit()
|
||||
rendered := next.(model).renderMessages()
|
||||
for _, want := range []string{"model: test-model", "mode: yolo tools", "theme: light"} {
|
||||
if !strings.Contains(rendered, want) {
|
||||
t.Fatalf("model command output missing %q:\n%s", want, rendered)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelThinkingCommand(t *testing.T) {
|
||||
manager := &fakeModelManager{model: "test-model", thinking: "none"}
|
||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model", ModelManager: manager})
|
||||
m.input.SetValue("/model thinking high")
|
||||
|
||||
next, _ := m.submit()
|
||||
rendered := next.(model).renderMessages()
|
||||
if manager.thinking != "high" {
|
||||
t.Fatalf("thinking = %q", manager.thinking)
|
||||
}
|
||||
if !strings.Contains(rendered, "thinking: high") {
|
||||
t.Fatalf("model command output missing thinking:\n%s", rendered)
|
||||
}
|
||||
}
|
||||
|
||||
type fakeModelManager struct {
|
||||
provider string
|
||||
model string
|
||||
thinking string
|
||||
}
|
||||
|
||||
func (f *fakeModelManager) CurrentProvider() string {
|
||||
if f.provider == "" {
|
||||
return "default"
|
||||
}
|
||||
return f.provider
|
||||
}
|
||||
|
||||
func (f *fakeModelManager) CurrentModel() string {
|
||||
return f.model
|
||||
}
|
||||
|
||||
func (f *fakeModelManager) CurrentThinking() string {
|
||||
return f.thinking
|
||||
}
|
||||
|
||||
func (f *fakeModelManager) Info() string {
|
||||
return "provider: " + f.CurrentProvider() + "\nmodel: " + f.model + "\nthinking: " + f.thinking
|
||||
}
|
||||
|
||||
func (f *fakeModelManager) SetProvider(name string) (string, error) {
|
||||
f.provider = name
|
||||
return f.Info(), nil
|
||||
}
|
||||
|
||||
func (f *fakeModelManager) SetModel(model string) (string, error) {
|
||||
f.model = model
|
||||
return f.Info(), nil
|
||||
}
|
||||
|
||||
func (f *fakeModelManager) SetThinking(level string) (string, error) {
|
||||
f.thinking = level
|
||||
return f.Info(), nil
|
||||
}
|
||||
|
||||
func TestParseThemeMode(t *testing.T) {
|
||||
for input, want := range map[string]ThemeMode{
|
||||
"": ThemeLight,
|
||||
"light": ThemeLight,
|
||||
"dark": ThemeDark,
|
||||
"DARK": ThemeDark,
|
||||
} {
|
||||
got, err := ParseThemeMode(input)
|
||||
if err != nil {
|
||||
t.Fatalf("ParseThemeMode(%q): %v", input, err)
|
||||
}
|
||||
if got != want {
|
||||
t.Fatalf("ParseThemeMode(%q) = %q, want %q", input, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := ParseThemeMode("sepia"); err == nil {
|
||||
t.Fatal("expected invalid theme error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCleanToolLog(t *testing.T) {
|
||||
got := cleanToolLog("\n[tool] shell_run {\"command\":\"pwd\"}\n")
|
||||
want := `shell_run {"command":"pwd"}`
|
||||
if got != want {
|
||||
t.Fatalf("got %q want %q", got, want)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
package tui
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/charmbracelet/bubbles/textarea"
|
||||
"github.com/charmbracelet/lipgloss"
|
||||
)
|
||||
|
||||
type ThemeMode string
|
||||
|
||||
const (
|
||||
ThemeLight ThemeMode = "light"
|
||||
ThemeDark ThemeMode = "dark"
|
||||
)
|
||||
|
||||
func ParseThemeMode(value string) (ThemeMode, error) {
|
||||
switch ThemeMode(strings.ToLower(strings.TrimSpace(value))) {
|
||||
case "", ThemeLight:
|
||||
return ThemeLight, nil
|
||||
case ThemeDark:
|
||||
return ThemeDark, nil
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported theme %q; expected light or dark", value)
|
||||
}
|
||||
}
|
||||
|
||||
type theme struct {
|
||||
Mode ThemeMode
|
||||
|
||||
Background string
|
||||
Surface string
|
||||
SurfaceAlt string
|
||||
Border string
|
||||
Text string
|
||||
Muted string
|
||||
|
||||
Title string
|
||||
User string
|
||||
Assistant string
|
||||
Tool string
|
||||
Error string
|
||||
System string
|
||||
}
|
||||
|
||||
type styles struct {
|
||||
App lipgloss.Style
|
||||
Header lipgloss.Style
|
||||
Title lipgloss.Style
|
||||
Muted lipgloss.Style
|
||||
Status lipgloss.Style
|
||||
Viewport lipgloss.Style
|
||||
InputBox lipgloss.Style
|
||||
CommandHint lipgloss.Style
|
||||
UserLabel lipgloss.Style
|
||||
AssistantLabel 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 {
|
||||
if mode == ThemeDark {
|
||||
return darkTheme()
|
||||
}
|
||||
return lightTheme()
|
||||
}
|
||||
|
||||
func lightTheme() theme {
|
||||
return theme{
|
||||
Mode: ThemeLight,
|
||||
Background: "#F8FAFC",
|
||||
Surface: "#FFFFFF",
|
||||
SurfaceAlt: "#EEF2F7",
|
||||
Border: "#CBD5E1",
|
||||
Text: "#111827",
|
||||
Muted: "#64748B",
|
||||
Title: "#0F766E",
|
||||
User: "#2563EB",
|
||||
Assistant: "#0F766E",
|
||||
Tool: "#B45309",
|
||||
Error: "#DC2626",
|
||||
System: "#64748B",
|
||||
}
|
||||
}
|
||||
|
||||
func darkTheme() theme {
|
||||
return theme{
|
||||
Mode: ThemeDark,
|
||||
Background: "#111827",
|
||||
Surface: "#1F2937",
|
||||
SurfaceAlt: "#0F172A",
|
||||
Border: "#475569",
|
||||
Text: "#E5E7EB",
|
||||
Muted: "#94A3B8",
|
||||
Title: "#5EEAD4",
|
||||
User: "#93C5FD",
|
||||
Assistant: "#5EEAD4",
|
||||
Tool: "#FBBF24",
|
||||
Error: "#FCA5A5",
|
||||
System: "#94A3B8",
|
||||
}
|
||||
}
|
||||
|
||||
func (t theme) styles() styles {
|
||||
return styles{
|
||||
App: lipgloss.NewStyle().
|
||||
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().
|
||||
Foreground(lipgloss.Color(t.Muted)),
|
||||
|
||||
Status: lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(t.Muted)).
|
||||
Background(lipgloss.Color(t.SurfaceAlt)).
|
||||
Padding(0, 1),
|
||||
|
||||
Viewport: lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(t.Text)).
|
||||
Background(lipgloss.Color(t.Background)),
|
||||
|
||||
InputBox: lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(t.Text)).
|
||||
Background(lipgloss.Color(t.Surface)).
|
||||
Border(lipgloss.NormalBorder()).
|
||||
BorderForeground(lipgloss.Color(t.Border)).
|
||||
Padding(0, 1),
|
||||
|
||||
CommandHint: lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(t.Muted)).
|
||||
Background(lipgloss.Color(t.SurfaceAlt)).
|
||||
Padding(0, 1),
|
||||
|
||||
UserLabel: labelStyle(t.User),
|
||||
AssistantLabel: labelStyle(t.Assistant),
|
||||
ToolLabel: labelStyle(t.Tool),
|
||||
ErrorLabel: labelStyle(t.Error),
|
||||
SystemLabel: labelStyle(t.System),
|
||||
|
||||
UserMessage: messageStyle(t.Text, t.User),
|
||||
AssistantMsg: messageStyle(t.Text, t.Assistant),
|
||||
ToolMessage: messageStyle(t.Text, t.Tool),
|
||||
ErrorMessage: messageStyle(t.Error, t.Error),
|
||||
SystemMessage: lipgloss.NewStyle().Foreground(lipgloss.Color(t.System)).PaddingLeft(1),
|
||||
}
|
||||
}
|
||||
|
||||
func labelStyle(color string) lipgloss.Style {
|
||||
return lipgloss.NewStyle().
|
||||
Bold(true).
|
||||
Foreground(lipgloss.Color(color))
|
||||
}
|
||||
|
||||
func messageStyle(textColor, accentColor string) lipgloss.Style {
|
||||
return lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(textColor)).
|
||||
Border(lipgloss.ThickBorder(), false, false, false, true).
|
||||
BorderForeground(lipgloss.Color(accentColor)).
|
||||
PaddingLeft(1)
|
||||
}
|
||||
|
||||
func applyTextareaTheme(input *textarea.Model, t theme) {
|
||||
focused := textarea.Style{
|
||||
Base: lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(t.Text)).
|
||||
Background(lipgloss.Color(t.Surface)),
|
||||
CursorLine: lipgloss.NewStyle().
|
||||
Background(lipgloss.Color(t.Surface)),
|
||||
Placeholder: lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(t.Muted)),
|
||||
Prompt: lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(t.Muted)),
|
||||
Text: lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(t.Text)),
|
||||
EndOfBuffer: lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(t.Surface)),
|
||||
}
|
||||
blurred := focused
|
||||
blurred.CursorLine = lipgloss.NewStyle().Background(lipgloss.Color(t.Surface))
|
||||
|
||||
input.FocusedStyle = focused
|
||||
input.BlurredStyle = blurred
|
||||
}
|
||||
Reference in New Issue
Block a user