diff --git a/pkg/seahorse/short_retrieval.go b/pkg/seahorse/short_retrieval.go index 3e94eec14..c431efa0b 100644 --- a/pkg/seahorse/short_retrieval.go +++ b/pkg/seahorse/short_retrieval.go @@ -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) diff --git a/pkg/seahorse/tool_expand.go b/pkg/seahorse/tool_expand.go index 749c9cd6c..fbc4305f1 100644 --- a/pkg/seahorse/tool_expand.go +++ b/pkg/seahorse/tool_expand.go @@ -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)) diff --git a/pkg/seahorse/tool_expand_test.go b/pkg/seahorse/tool_expand_test.go index fc726a7a0..9e18b5f10 100644 --- a/pkg/seahorse/tool_expand_test.go +++ b/pkg/seahorse/tool_expand_test.go @@ -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") + } +} diff --git a/pkg/seahorse/tool_grep.go b/pkg/seahorse/tool_grep.go index 9671d2a7f..31cfdc39f 100644 --- a/pkg/seahorse/tool_grep.go +++ b/pkg/seahorse/tool_grep.go @@ -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, diff --git a/pkg/seahorse/tool_grep_test.go b/pkg/seahorse/tool_grep_test.go index 050d9deeb..ecb5f61f8 100644 --- a/pkg/seahorse/tool_grep_test.go +++ b/pkg/seahorse/tool_grep_test.go @@ -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) + } +}