1package extensions
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "time"
8
9 "github.com/patriceckhart/zot/packages/core"
10 "github.com/patriceckhart/zot/packages/provider"
11)
12
13// extensionTool wraps a single extension-registered tool as a
14// core.Tool. The agent's tool registry contains one of these per
15// extension tool; Execute round-trips through the manager to the
16// owning subprocess.
17//
18// One concrete type instead of a closure-driven anonymous tool
19// keeps the schema, name, and ownership inspectable for logs and
20// dialogs.
21type extensionTool struct {
22 name string
23 description string
24 schema json.RawMessage
25 extension string
26 manager *Manager
27 timeout time.Duration
28 deferred bool
29}
30
31// NewTool returns a core.Tool that round-trips invocations through
32// mgr to the extension that registered (name, schema). The default
33// per-call timeout is 60 seconds; callers can override.
34func NewTool(mgr *Manager, info ToolInfo) core.Tool {
35 return &extensionTool{
36 name: info.Name,
37 description: info.Description,
38 schema: info.Schema,
39 extension: info.Extension,
40 manager: mgr,
41 timeout: 60 * time.Second,
42 deferred: info.Deferred,
43 }
44}
45
46func (t *extensionTool) Name() string { return t.name }
47func (t *extensionTool) Description() string { return t.description }
48func (t *extensionTool) Schema() json.RawMessage { return t.schema }
49func (t *extensionTool) Extension() string { return t.extension }
50func (t *extensionTool) Deferred() bool { return t.deferred }
51
52// Execute is what the agent calls when the LLM invokes the tool. It
53// hands args to the owning extension, waits up to t.timeout for the
54// reply, and converts the response into a core.ToolResult.
55func (t *extensionTool) Execute(ctx context.Context, args json.RawMessage, _ func(string)) (core.ToolResult, error) {
56 if len(args) == 0 {
57 args = json.RawMessage(`{}`)
58 }
59 resp, err := t.manager.InvokeTool(ctx, t.name, args, t.timeout)
60 if err != nil {
61 return core.ToolResult{
62 IsError: true,
63 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("extension %s/%s failed: %v", t.extension, t.name, err)}},
64 }, nil
65 }
66 out := core.ToolResult{IsError: resp.IsError, ActivateTools: resp.ActivateTools}
67 for _, b := range resp.Content {
68 switch b.Type {
69 case "text":
70 if b.Text != "" {
71 out.Content = append(out.Content, provider.TextBlock{Text: b.Text})
72 }
73 case "image":
74 data, dErr := decodeBase64(b.Data)
75 if dErr != nil {
76 out.IsError = true
77 out.Content = append(out.Content, provider.TextBlock{Text: fmt.Sprintf("extension %s/%s returned undecodable image: %v", t.extension, t.name, dErr)})
78 continue
79 }
80 out.Content = append(out.Content, provider.ImageBlock{
81 MimeType: b.MimeType,
82 Data: data,
83 })
84 default:
85 // Unknown block type: stringify for debug visibility.
86 out.Content = append(out.Content, provider.TextBlock{Text: fmt.Sprintf("[unknown block type %q from extension %s]", b.Type, t.extension)})
87 }
88 }
89 if len(out.Content) == 0 {
90 // Defensive: an empty content slice would confuse the model.
91 out.Content = []provider.Content{provider.TextBlock{Text: "(extension returned no content)"}}
92 }
93 out.Details = map[string]any{"extension": t.extension, "tool": t.name}
94 return out, nil
95}
96
97// decodeBase64 is a tiny wrapper around encoding/base64 so we can
98// validate the extension's image data in one place.
99func decodeBase64(s string) ([]byte, error) {
100 if s == "" {
101 return nil, nil
102 }
103 return base64DecodeStd(s)
104}