149 lines
4.3 KiB
Go
149 lines
4.3 KiB
Go
package vision
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
|
)
|
|
|
|
// toolParameters returns the shared input JSON Schema for the vision and OCR
|
|
// tools. It builds a fresh map each call rather than sharing a package-level
|
|
// variable.
|
|
func toolParameters() map[string]any {
|
|
return map[string]any{
|
|
"type": "object",
|
|
"properties": map[string]any{
|
|
"image": map[string]any{
|
|
"type": "string",
|
|
"description": "The image: a local path, an http(s) URL, or a base64/data: URI.",
|
|
},
|
|
"prompt": map[string]any{
|
|
"type": "string",
|
|
"description": "Optional instruction for the vision model (vision tool only).",
|
|
},
|
|
},
|
|
"required": []any{"image"},
|
|
}
|
|
}
|
|
|
|
// VisionTool wraps a VisionAdapter as an sdk.Tool so the M2 tool loop can invoke
|
|
// it and return its description as a tool result, keeping the image out of the
|
|
// main-model context. It is safe for concurrent use.
|
|
type VisionTool struct {
|
|
name string
|
|
adapter VisionAdapter
|
|
}
|
|
|
|
// NewVisionTool builds a VisionTool with the given name and adapter.
|
|
func NewVisionTool(name string, adapter VisionAdapter) *VisionTool {
|
|
return &VisionTool{name: name, adapter: adapter}
|
|
}
|
|
|
|
// Name returns the tool identifier the agent uses to invoke it.
|
|
func (t *VisionTool) Name() string { return t.name }
|
|
|
|
// Description explains what the tool does.
|
|
func (t *VisionTool) Description() string {
|
|
return "Understand an image and return a description of its contents."
|
|
}
|
|
|
|
// Parameters returns the input JSON Schema.
|
|
func (t *VisionTool) Parameters() (map[string]any, error) {
|
|
return toolParameters(), nil
|
|
}
|
|
|
|
// Execute runs the adapter on the image from args and returns the description.
|
|
func (t *VisionTool) Execute(ctx context.Context, args any) (any, error) {
|
|
image, prompt, err := toolArgs(args)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
ref, err := ParseImageRef(image)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
desc, err := t.adapter.Describe(ctx, ref, prompt)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return map[string]any{"description": desc}, nil
|
|
}
|
|
|
|
// OCRTool wraps an OCRAdapter as an sdk.Tool. It returns the extracted text as a
|
|
// tool result. It is safe for concurrent use.
|
|
type OCRTool struct {
|
|
name string
|
|
adapter OCRAdapter
|
|
}
|
|
|
|
// NewOCRTool builds an OCRTool with the given name and adapter.
|
|
func NewOCRTool(name string, adapter OCRAdapter) *OCRTool {
|
|
return &OCRTool{name: name, adapter: adapter}
|
|
}
|
|
|
|
// Name returns the tool identifier.
|
|
func (t *OCRTool) Name() string { return t.name }
|
|
|
|
// Description explains what the tool does.
|
|
func (t *OCRTool) Description() string {
|
|
return "Extract text from an image (OCR) and return the recognized text."
|
|
}
|
|
|
|
// Parameters returns the input JSON Schema.
|
|
func (t *OCRTool) Parameters() (map[string]any, error) {
|
|
return toolParameters(), nil
|
|
}
|
|
|
|
// Execute runs the OCR adapter on the image from args and returns the text.
|
|
func (t *OCRTool) Execute(ctx context.Context, args any) (any, error) {
|
|
image, _, err := toolArgs(args)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
ref, err := ParseImageRef(image)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
text, err := t.adapter.Extract(ctx, ref)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return map[string]any{"text": text}, nil
|
|
}
|
|
|
|
// RegisterTools registers the vision and OCR tools with r via the SDK helper.
|
|
// Each tool is only registered when its pointer is non-nil. It returns the
|
|
// first registration error, or nil when every tool registered.
|
|
func RegisterTools(r sdk.Registrar, tool *VisionTool, ocr *OCRTool) error {
|
|
if r == nil {
|
|
return sdk.ErrNilRegistrar
|
|
}
|
|
if tool != nil {
|
|
if err := sdk.RegisterTool(r, tool); err != nil {
|
|
return fmt.Errorf("vision: register tool %q: %w", tool.Name(), err)
|
|
}
|
|
}
|
|
if ocr != nil {
|
|
if err := sdk.RegisterTool(r, ocr); err != nil {
|
|
return fmt.Errorf("vision: register tool %q: %w", ocr.Name(), err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// toolArgs decodes sdk.Tool.Execute args into the image string and an optional
|
|
// prompt. Args must be a map[string]any with a non-empty "image" string.
|
|
func toolArgs(args any) (image, prompt string, err error) {
|
|
m, ok := args.(map[string]any)
|
|
if !ok {
|
|
return "", "", fmt.Errorf("vision: %w: args must be an object", ErrBadArgs)
|
|
}
|
|
image, _ = m["image"].(string)
|
|
if image == "" {
|
|
return "", "", fmt.Errorf("vision: %w: missing image", ErrBadArgs)
|
|
}
|
|
prompt, _ = m["prompt"].(string)
|
|
return image, prompt, nil
|
|
}
|