Merge 95393d9b68 into 412705783d
This commit is contained in:
commit
b7272b70a6
5 changed files with 316 additions and 10 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue