This commit is contained in:
Anton Bogdanovich 2026-05-15 11:29:28 +03:00 committed by GitHub
commit b7272b70a6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 316 additions and 10 deletions

View file

@ -44,6 +44,7 @@ type GrepInput struct {
Pattern string `json:"pattern"`
Scope string `json:"scope,omitempty"` // "both" (default), "summary", or "message"
Role string `json:"role,omitempty"` // "user", "assistant", or "" (all)
ConversationID int64 `json:"conversationId,omitempty"`
AllConversations bool `json:"allConversations,omitempty"`
Since *time.Time `json:"since,omitempty"`
Before *time.Time `json:"before,omitempty"`
@ -84,8 +85,9 @@ type GrepMessageResult struct {
// ExpandMessagesResult contains expanded messages.
type ExpandMessagesResult struct {
Messages []Message `json:"messages"`
TokenCount int `json:"tokenCount"`
Messages []Message `json:"messages"`
RejectedMessageIDs []int64 `json:"rejectedMessageIds,omitempty"`
TokenCount int `json:"tokenCount"`
}
// Grep searches summaries and messages for matching content.
@ -120,6 +122,7 @@ func (r *RetrievalEngine) Grep(ctx context.Context, input GrepInput) (*GrepResul
Pattern: input.Pattern,
Mode: mode,
Role: input.Role,
ConversationID: input.ConversationID,
AllConversations: input.AllConversations,
Since: since,
Before: input.Before,
@ -193,8 +196,43 @@ func (r *RetrievalEngine) Grep(ctx context.Context, input GrepInput) (*GrepResul
return result, nil
}
// ConversationIDForSession returns the storage conversation ID for a session key.
func (r *RetrievalEngine) ConversationIDForSession(ctx context.Context, sessionKey string) (int64, bool, error) {
if strings.TrimSpace(sessionKey) == "" {
return 0, false, nil
}
conv, err := r.store.GetConversationBySessionKey(ctx, sessionKey)
if err != nil {
return 0, false, err
}
if conv == nil {
return 0, false, nil
}
return conv.ConversationID, true, nil
}
// ExpandMessages retrieves full message content by IDs.
func (r *RetrievalEngine) ExpandMessages(ctx context.Context, messageIDs []int64) (*ExpandMessagesResult, error) {
return r.expandMessages(ctx, messageIDs, 0, true)
}
// ExpandMessagesScoped retrieves full message content by IDs, restricted to a conversation
// unless allConversations is true.
func (r *RetrievalEngine) ExpandMessagesScoped(
ctx context.Context,
messageIDs []int64,
conversationID int64,
allConversations bool,
) (*ExpandMessagesResult, error) {
return r.expandMessages(ctx, messageIDs, conversationID, allConversations)
}
func (r *RetrievalEngine) expandMessages(
ctx context.Context,
messageIDs []int64,
conversationID int64,
allConversations bool,
) (*ExpandMessagesResult, error) {
result := &ExpandMessagesResult{
Messages: make([]Message, 0, len(messageIDs)),
}
@ -202,6 +240,11 @@ func (r *RetrievalEngine) ExpandMessages(ctx context.Context, messageIDs []int64
for _, msgID := range messageIDs {
msg, err := r.store.GetMessageByID(ctx, msgID)
if err != nil {
result.RejectedMessageIDs = append(result.RejectedMessageIDs, msgID)
continue
}
if !allConversations && (conversationID <= 0 || msg.ConversationID != conversationID) {
result.RejectedMessageIDs = append(result.RejectedMessageIDs, msgID)
continue
}
result.Messages = append(result.Messages, *msg)

View file

@ -28,6 +28,7 @@ Use when short_grep returns messages and you need complete content (not just sni
Parameters:
- message_ids (required): Array of message ID strings (from short_grep results)
- all_conversations: Expand IDs from any conversation (default: current conversation only)
Returns message with:
- content: Full text content
@ -40,9 +41,12 @@ Returns message with:
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.
- By default, IDs outside the current conversation are rejected and reported in rejectedMessageIds.
- If short_grep used all_conversations: true, pass all_conversations: true to expand those IDs.
Example:
{"message_ids": ["10", "25"]}`
{"message_ids": ["10", "25"]}
{"message_ids": ["10", "25"], "all_conversations": true}`
}
func (t *ExpandTool) Parameters() map[string]any {
@ -54,6 +58,10 @@ func (t *ExpandTool) Parameters() map[string]any {
"items": map[string]any{"type": "string"},
"description": "Message IDs to expand (from short_grep results, e.g., [\"10\", \"25\"])",
},
"all_conversations": map[string]any{
"type": "boolean",
"description": "Expand IDs across all conversations (default: current conversation only)",
},
},
"required": []string{"message_ids"},
}
@ -82,7 +90,23 @@ func (t *ExpandTool) Execute(ctx context.Context, args map[string]any) *tools.To
}
}
result, err := t.engine.ExpandMessages(ctx, messageIDs)
allConversations, _ := args["all_conversations"].(bool)
var conversationID int64
if !allConversations {
var found bool
var err error
conversationID, found, err = t.engine.ConversationIDForSession(ctx, tools.ToolSessionKey(ctx))
if err != nil {
return tools.ErrorResult("Expand failed: resolve current conversation: " + err.Error())
}
if !found {
return tools.ErrorResult(
"Expand failed: no current conversation found for this session. Use all_conversations: true to expand across conversations.",
)
}
}
result, err := t.engine.ExpandMessagesScoped(ctx, messageIDs, conversationID, allConversations)
if err != nil {
return tools.ErrorResult("Expand failed: " + err.Error())
}
@ -120,9 +144,10 @@ func (t *ExpandTool) Execute(ctx context.Context, args map[string]any) *tools.To
}
output := map[string]any{
"success": true,
"tokenCount": result.TokenCount,
"messages": messages,
"success": true,
"tokenCount": result.TokenCount,
"messages": messages,
"rejectedMessageIds": result.RejectedMessageIDs,
}
data, _ := json.Marshal(output)
return tools.NewToolResult(string(data))

View file

@ -5,6 +5,8 @@ import (
"encoding/json"
"fmt"
"testing"
"github.com/sipeed/picoclaw/pkg/tools"
)
func TestExpandToolByMessageIDs(t *testing.T) {
@ -19,7 +21,8 @@ func TestExpandToolByMessageIDs(t *testing.T) {
tool := NewExpandTool(re)
result := tool.Execute(ctx, map[string]any{
"message_ids": []any{fmt.Sprintf("%d", msg1.ID), fmt.Sprintf("%d", msg2.ID)},
"message_ids": []any{fmt.Sprintf("%d", msg1.ID), fmt.Sprintf("%d", msg2.ID)},
"all_conversations": true,
})
if result.IsError {
@ -76,7 +79,8 @@ func TestExpandToolWithParts(t *testing.T) {
tool := NewExpandTool(re)
result := tool.Execute(ctx, map[string]any{
"message_ids": []any{fmt.Sprintf("%d", msg.ID)},
"message_ids": []any{fmt.Sprintf("%d", msg.ID)},
"all_conversations": true,
})
if result.IsError {
@ -134,3 +138,108 @@ func TestExpandToolWithParts(t *testing.T) {
t.Error("missing tool_result part")
}
}
func TestExpandToolScopesToCurrentSession(t *testing.T) {
s := openTestStore(t)
ctx := context.Background()
current, _ := s.GetOrCreateConversation(ctx, "session:current")
other, _ := s.GetOrCreateConversation(ctx, "session:other")
currentMsg, _ := s.AddMessage(ctx, current.ConversationID, "user", "current message", 5)
otherMsg, _ := s.AddMessage(ctx, other.ConversationID, "user", "other message", 5)
tool := NewExpandTool(&RetrievalEngine{store: s})
toolCtx := tools.WithToolSessionContext(ctx, "agent", "session:current", nil)
result := tool.Execute(toolCtx, map[string]any{
"message_ids": []any{
float64(currentMsg.ID),
float64(otherMsg.ID),
},
})
if result.IsError {
t.Fatalf("Execute returned error: %s", result.ContentForLLM())
}
var output struct {
Messages []struct {
Content string `json:"content"`
} `json:"messages"`
RejectedMessageIDs []int64 `json:"rejectedMessageIds"`
}
if err := json.Unmarshal([]byte(result.ContentForLLM()), &output); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if len(output.Messages) != 1 {
t.Fatalf("messages = %d, want 1: %#v", len(output.Messages), output.Messages)
}
if output.Messages[0].Content != "current message" {
t.Fatalf("content = %q, want current message", output.Messages[0].Content)
}
if len(output.RejectedMessageIDs) != 1 || output.RejectedMessageIDs[0] != otherMsg.ID {
t.Fatalf("rejectedMessageIds = %#v, want [%d]", output.RejectedMessageIDs, otherMsg.ID)
}
}
func TestExpandToolCanExpandAllConversations(t *testing.T) {
s := openTestStore(t)
ctx := context.Background()
current, _ := s.GetOrCreateConversation(ctx, "session:current")
other, _ := s.GetOrCreateConversation(ctx, "session:other")
currentMsg, _ := s.AddMessage(ctx, current.ConversationID, "user", "current message", 5)
otherMsg, _ := s.AddMessage(ctx, other.ConversationID, "user", "other message", 5)
tool := NewExpandTool(&RetrievalEngine{store: s})
toolCtx := tools.WithToolSessionContext(ctx, "agent", "session:current", nil)
result := tool.Execute(toolCtx, map[string]any{
"message_ids": []any{
float64(currentMsg.ID),
float64(otherMsg.ID),
},
"all_conversations": true,
})
if result.IsError {
t.Fatalf("Execute returned error: %s", result.ContentForLLM())
}
var output struct {
Messages []struct {
Content string `json:"content"`
} `json:"messages"`
RejectedMessageIDs []int64 `json:"rejectedMessageIds"`
}
if err := json.Unmarshal([]byte(result.ContentForLLM()), &output); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if len(output.Messages) != 2 {
t.Fatalf("messages = %d, want 2: %#v", len(output.Messages), output.Messages)
}
if len(output.RejectedMessageIDs) != 0 {
t.Fatalf("rejectedMessageIds = %#v, want none", output.RejectedMessageIDs)
}
}
func TestExpandToolUnknownSessionErrors(t *testing.T) {
s := openTestStore(t)
ctx := context.Background()
conv, _ := s.GetOrCreateConversation(ctx, "session:current")
msg, _ := s.AddMessage(ctx, conv.ConversationID, "user", "current message", 5)
tool := NewExpandTool(&RetrievalEngine{store: s})
toolCtx := tools.WithToolSessionContext(ctx, "agent", "session:missing", nil)
result := tool.Execute(toolCtx, map[string]any{
"message_ids": []any{float64(msg.ID)},
})
if !result.IsError {
t.Fatal("expected error for unknown current session")
}
}
func TestExpandToolSupportsAllConversationsParameter(t *testing.T) {
s := openTestStore(t)
tool := NewExpandTool(&RetrievalEngine{store: s})
params := tool.Parameters()
props := params["properties"].(map[string]any)
if _, ok := props["all_conversations"]; !ok {
t.Error("Parameters missing 'all_conversations' field")
}
}

View file

@ -131,6 +131,21 @@ func (t *GrepTool) Execute(ctx context.Context, args map[string]any) *tools.Tool
if allConv, ok := args["all_conversations"].(bool); ok {
input.AllConversations = allConv
}
if !input.AllConversations {
conversationID, found, err := t.engine.ConversationIDForSession(ctx, tools.ToolSessionKey(ctx))
if err != nil {
return tools.ErrorResult("Grep failed: resolve current conversation: " + err.Error())
}
if !found {
return grepJSONResult(&GrepResult{
Success: true,
Summaries: make([]GrepSummaryResult, 0),
Messages: make([]GrepMessageResult, 0),
Hint: "No current conversation found for this session. Use all_conversations: true to search across conversations.",
})
}
input.ConversationID = conversationID
}
if limit, ok := args["limit"].(float64); ok {
input.Limit = int(limit)
}
@ -155,7 +170,10 @@ func (t *GrepTool) Execute(ctx context.Context, args map[string]any) *tools.Tool
return tools.ErrorResult("Grep failed: " + err.Error())
}
// Build response
return grepJSONResult(result)
}
func grepJSONResult(result *GrepResult) *tools.ToolResult {
output := map[string]any{
"success": result.Success,
"summaries": result.Summaries,

View file

@ -2,7 +2,10 @@ package seahorse
import (
"context"
"encoding/json"
"testing"
"github.com/sipeed/picoclaw/pkg/tools"
)
func TestGrepSearchSummaries(t *testing.T) {
@ -70,3 +73,111 @@ func TestGrepToolSupportsAllConversations(t *testing.T) {
t.Error("Parameters missing 'all_conversations' field")
}
}
func TestGrepToolScopesToCurrentSessionByDefault(t *testing.T) {
s := openTestStore(t)
ctx := context.Background()
current, _ := s.GetOrCreateConversation(ctx, "session:current")
other, _ := s.GetOrCreateConversation(ctx, "session:other")
s.AddMessage(ctx, current.ConversationID, "user", "shared needle from current topic", 5)
s.AddMessage(ctx, other.ConversationID, "user", "shared needle from other topic", 5)
tool := NewGrepTool(&RetrievalEngine{store: s})
toolCtx := tools.WithToolSessionContext(ctx, "agent", "session:current", nil)
result := tool.Execute(toolCtx, map[string]any{"pattern": "needle"})
if result.IsError {
t.Fatalf("Execute returned error: %s", result.ContentForLLM())
}
var output struct {
Messages []GrepMessageResult `json:"messages"`
}
if err := json.Unmarshal([]byte(result.ContentForLLM()), &output); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if len(output.Messages) != 1 {
t.Fatalf("messages = %d, want 1: %#v", len(output.Messages), output.Messages)
}
if output.Messages[0].ConversationID != current.ConversationID {
t.Fatalf("conversation id = %d, want %d", output.Messages[0].ConversationID, current.ConversationID)
}
}
func TestGrepToolCanSearchAllConversations(t *testing.T) {
s := openTestStore(t)
ctx := context.Background()
current, _ := s.GetOrCreateConversation(ctx, "session:current")
other, _ := s.GetOrCreateConversation(ctx, "session:other")
s.AddMessage(ctx, current.ConversationID, "user", "shared needle from current topic", 5)
s.AddMessage(ctx, other.ConversationID, "user", "shared needle from other topic", 5)
tool := NewGrepTool(&RetrievalEngine{store: s})
toolCtx := tools.WithToolSessionContext(ctx, "agent", "session:current", nil)
result := tool.Execute(toolCtx, map[string]any{"pattern": "needle", "all_conversations": true})
if result.IsError {
t.Fatalf("Execute returned error: %s", result.ContentForLLM())
}
var output struct {
Messages []GrepMessageResult `json:"messages"`
}
if err := json.Unmarshal([]byte(result.ContentForLLM()), &output); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if len(output.Messages) != 2 {
t.Fatalf("messages = %d, want 2: %#v", len(output.Messages), output.Messages)
}
}
func TestGrepToolUnknownSessionDoesNotSearchAllConversations(t *testing.T) {
s := openTestStore(t)
ctx := context.Background()
current, _ := s.GetOrCreateConversation(ctx, "session:current")
other, _ := s.GetOrCreateConversation(ctx, "session:other")
s.AddMessage(ctx, current.ConversationID, "user", "shared needle from current topic", 5)
s.AddMessage(ctx, other.ConversationID, "user", "shared needle from other topic", 5)
tool := NewGrepTool(&RetrievalEngine{store: s})
toolCtx := tools.WithToolSessionContext(ctx, "agent", "session:missing", nil)
result := tool.Execute(toolCtx, map[string]any{"pattern": "needle"})
if result.IsError {
t.Fatalf("Execute returned error: %s", result.ContentForLLM())
}
var output struct {
Messages []GrepMessageResult `json:"messages"`
Hint string `json:"hint"`
}
if err := json.Unmarshal([]byte(result.ContentForLLM()), &output); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if len(output.Messages) != 0 {
t.Fatalf("messages = %d, want 0: %#v", len(output.Messages), output.Messages)
}
if output.Hint == "" {
t.Fatal("expected hint for missing current conversation")
}
}
func TestGrepToolEmptySessionDoesNotSearchAllConversations(t *testing.T) {
s := openTestStore(t)
ctx := context.Background()
conv, _ := s.GetOrCreateConversation(ctx, "session:current")
s.AddMessage(ctx, conv.ConversationID, "user", "shared needle from current topic", 5)
tool := NewGrepTool(&RetrievalEngine{store: s})
result := tool.Execute(ctx, map[string]any{"pattern": "needle"})
if result.IsError {
t.Fatalf("Execute returned error: %s", result.ContentForLLM())
}
var output struct {
Messages []GrepMessageResult `json:"messages"`
}
if err := json.Unmarshal([]byte(result.ContentForLLM()), &output); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if len(output.Messages) != 0 {
t.Fatalf("messages = %d, want 0: %#v", len(output.Messages), output.Messages)
}
}