opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

scope seahorse retrieval tools to current session

Browse source
This commit is contained in:
Anton Bogdanovich 2026-05-03 21:10:01 -07:00
parent 6e6293e596
commit ff9fdf774f
5 changed files with 149 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
pkg/seahorse/short_retrieval.go
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"`
@ -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
}

7
pkg/seahorse/tool_expand.go
View file

@ -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())
}

38
pkg/seahorse/tool_expand_test.go
View file

@ -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)
}
}

7
pkg/seahorse/tool_grep.go
View file

@ -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)
}

58
pkg/seahorse/tool_grep_test.go
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,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)
}
}