132 lines
3.5 KiB
Go
132 lines
3.5 KiB
Go
package seahorse
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"github.com/sipeed/picoclaw/pkg/tools"
|
|
)
|
|
|
|
// ExpandTool recovers full message content by ID.
|
|
type ExpandTool struct {
|
|
engine *RetrievalEngine
|
|
}
|
|
|
|
func NewExpandTool(engine *RetrievalEngine) *ExpandTool {
|
|
return &ExpandTool{engine: engine}
|
|
}
|
|
|
|
func (t *ExpandTool) Name() string {
|
|
return "short_expand"
|
|
}
|
|
|
|
func (t *ExpandTool) Description() string {
|
|
return `Get full message content by ID.
|
|
|
|
Use when short_grep returns messages and you need complete content (not just snippet).
|
|
|
|
Parameters:
|
|
- message_ids (required): Array of message ID strings (from short_grep results)
|
|
|
|
Returns message with:
|
|
- content: Full text content
|
|
- parts: Structured content
|
|
- text: Full text
|
|
- tool_use: name, arguments, toolCallId
|
|
- tool_result: toolCallId only (content omitted - re-run tool if needed)
|
|
- media: mediaUri (file path), mimeType
|
|
|
|
Notes:
|
|
- tool_result content is not returned (can be large). Re-run the tool if you need the result.
|
|
- Media files are stored on disk at mediaUri path, use bash to access.
|
|
|
|
Example:
|
|
{"message_ids": ["10", "25"]}`
|
|
}
|
|
|
|
func (t *ExpandTool) Parameters() map[string]any {
|
|
return map[string]any{
|
|
"type": "object",
|
|
"properties": map[string]any{
|
|
"message_ids": map[string]any{
|
|
"type": "array",
|
|
"items": map[string]any{"type": "string"},
|
|
"description": "Message IDs to expand (from short_grep results, e.g., [\"10\", \"25\"])",
|
|
},
|
|
},
|
|
"required": []string{"message_ids"},
|
|
}
|
|
}
|
|
|
|
func (t *ExpandTool) Execute(ctx context.Context, args map[string]any) *tools.ToolResult {
|
|
idsRaw, ok := args["message_ids"].([]any)
|
|
if !ok || len(idsRaw) == 0 {
|
|
return tools.ErrorResult(
|
|
"Missing required 'message_ids' argument. " +
|
|
"Example: {\"message_ids\": [\"10\", \"25\"]}")
|
|
}
|
|
|
|
// Parse message IDs
|
|
messageIDs := make([]int64, 0, len(idsRaw))
|
|
for _, id := range idsRaw {
|
|
switch v := id.(type) {
|
|
case string:
|
|
var n int64
|
|
if _, err := fmt.Sscanf(v, "%d", &n); err != nil {
|
|
return tools.ErrorResult(fmt.Sprintf("Invalid message_id %q: %v", v, err))
|
|
}
|
|
messageIDs = append(messageIDs, n)
|
|
case float64:
|
|
messageIDs = append(messageIDs, int64(v))
|
|
}
|
|
}
|
|
|
|
result, err := t.engine.ExpandMessages(ctx, messageIDs)
|
|
if err != nil {
|
|
return tools.ErrorResult("Expand failed: " + err.Error())
|
|
}
|
|
|
|
// Build response with filtered parts
|
|
messages := make([]map[string]any, 0, len(result.Messages))
|
|
for _, msg := range result.Messages {
|
|
parts := make([]map[string]any, 0, len(msg.Parts))
|
|
for _, p := range msg.Parts {
|
|
part := map[string]any{"type": p.Type}
|
|
switch p.Type {
|
|
case "text":
|
|
part["text"] = p.Text
|
|
case "tool_use":
|
|
part["name"] = p.Name
|
|
part["arguments"] = p.Arguments
|
|
part["toolCallId"] = p.ToolCallID
|
|
case "tool_result":
|
|
// Omit content - can be large, re-run tool if needed
|
|
part["toolCallId"] = p.ToolCallID
|
|
case "media":
|
|
part["mediaUri"] = p.MediaURI
|
|
part["mimeType"] = p.MimeType
|
|
}
|
|
parts = append(parts, part)
|
|
}
|
|
|
|
messages = append(messages, map[string]any{
|
|
"id": fmt.Sprintf("%d", msg.ID),
|
|
"role": msg.Role,
|
|
"content": msg.Content,
|
|
"parts": parts,
|
|
"conversationId": msg.ConversationID,
|
|
})
|
|
}
|
|
|
|
output := map[string]any{
|
|
"success": true,
|
|
"tokenCount": result.TokenCount,
|
|
"messages": messages,
|
|
}
|
|
data, err := json.Marshal(output)
|
|
if err != nil {
|
|
return tools.ErrorResult(fmt.Sprintf("failed to marshal expand result: %v", err))
|
|
}
|
|
return tools.NewToolResult(string(data))
|
|
}
|