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