diff --git a/pkg/seahorse/short_retrieval.go b/pkg/seahorse/short_retrieval.go index 3e94eec14..8f5b7ecb7 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"` @@ -120,6 +121,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 +195,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, error) { + if strings.TrimSpace(sessionKey) == "" { + return 0, nil + } + conv, err := r.store.GetConversationBySessionKey(ctx, sessionKey) + if err != nil { + return 0, err + } + if conv == nil { + return 0, nil + } + return conv.ConversationID, 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, false) +} + +// 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)), } @@ -204,6 +241,9 @@ func (r *RetrievalEngine) ExpandMessages(ctx context.Context, messageIDs []int64 if err != nil { continue } + if !allConversations && conversationID > 0 && msg.ConversationID != conversationID { + continue + } result.Messages = append(result.Messages, *msg) result.TokenCount += msg.TokenCount } diff --git a/pkg/seahorse/tool_expand.go b/pkg/seahorse/tool_expand.go index 749c9cd6c..b37691071 100644 --- a/pkg/seahorse/tool_expand.go +++ b/pkg/seahorse/tool_expand.go @@ -82,7 +82,12 @@ func (t *ExpandTool) Execute(ctx context.Context, args map[string]any) *tools.To } } - result, err := t.engine.ExpandMessages(ctx, messageIDs) + conversationID, err := t.engine.ConversationIDForSession(ctx, tools.ToolSessionKey(ctx)) + if err != nil { + return tools.ErrorResult("Expand failed: resolve current conversation: " + err.Error()) + } + + result, err := t.engine.ExpandMessagesScoped(ctx, messageIDs, conversationID, false) if err != nil { return tools.ErrorResult("Expand failed: " + err.Error()) } diff --git a/pkg/seahorse/tool_expand_test.go b/pkg/seahorse/tool_expand_test.go index fc726a7a0..bd303df69 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) { @@ -134,3 +136,39 @@ 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"` + } + 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) + } +} diff --git a/pkg/seahorse/tool_grep.go b/pkg/seahorse/tool_grep.go index 9671d2a7f..88fff137a 100644 --- a/pkg/seahorse/tool_grep.go +++ b/pkg/seahorse/tool_grep.go @@ -131,6 +131,13 @@ 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, err := t.engine.ConversationIDForSession(ctx, tools.ToolSessionKey(ctx)) + if err != nil { + return tools.ErrorResult("Grep failed: resolve current conversation: " + err.Error()) + } + input.ConversationID = conversationID + } if limit, ok := args["limit"].(float64); ok { input.Limit = int(limit) } diff --git a/pkg/seahorse/tool_grep_test.go b/pkg/seahorse/tool_grep_test.go index 050d9deeb..999a2f3d5 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,58 @@ 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) + } +}