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,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())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user