Skip to content

Commit b7c9000

Browse files
maximilienmaximilien
authored
Merge pull request #33 from ajbozarth/search
feat: add search endpoint for direct doc retrieval
2 parents e32d51c + 67f8806 commit b7c9000

File tree

10 files changed

+396
-11
lines changed

10 files changed

+396
-11
lines changed

‎cli/src/main.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,12 @@ func addContextualHelp() {
8080
ShowContextualHelp("query", "")
8181
}
8282
}
83+
84+
searchCmd.PostRun = func(cmd *cobra.Command, args []string) {
85+
if !silent {
86+
ShowContextualHelp("search", "")
87+
}
88+
}
8389
}
8490

8591
func main() {
@@ -115,6 +121,7 @@ A command-line interface for working with Maestro Knowledge configurations.`,
115121
rootCmd.AddCommand(documentCmd)
116122
rootCmd.AddCommand(embeddingCmd)
117123
rootCmd.AddCommand(queryCmd)
124+
rootCmd.AddCommand(searchCmd)
118125
rootCmd.AddCommand(validateCmd)
119126
rootCmd.AddCommand(statusCmd)
120127

‎cli/src/mcp_client.go‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -669,6 +669,44 @@ func (c *MCPClient) Query(dbName, query string, limit int, collectionName string
669669
return resultStr, nil
670670
}
671671

672+
// Search calls the search tool on the MCP server
673+
func (c *MCPClient) Search(dbName, query string, limit int, collectionName string) (string, error) {
674+
params := map[string]interface{}{
675+
"input": map[string]interface{}{
676+
"db_name": dbName,
677+
"query": query,
678+
"limit": limit,
679+
"collection_name": collectionName,
680+
},
681+
}
682+
683+
response, err := c.callMCPServer("search", params)
684+
if err != nil {
685+
return "", err
686+
}
687+
688+
// Check for error in response
689+
if response.Error != nil {
690+
return "", fmt.Errorf("MCP server error: %s", response.Error.Message)
691+
}
692+
693+
// The response should be a string with the search result
694+
if response.Result == nil {
695+
return "", fmt.Errorf("no response from MCP server")
696+
}
697+
698+
prettyJSON, err := json.MarshalIndent(response.Result, "", " ")
699+
if err != nil {
700+
resultStr, ok := response.Result.(string)
701+
if !ok {
702+
return resultStr, nil
703+
}
704+
return "", fmt.Errorf("unexpected response type from MCP server")
705+
}
706+
707+
return string(prettyJSON), nil
708+
}
709+
672710
// Close closes the MCP client
673711
func (c *MCPClient) Close() error {
674712
// Cancel the context to prevent context leaks

‎cli/src/query.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,11 +45,11 @@ The query agent will search through the documents and provide relevant answers.`
4545
collectionName = "Test_collection_lowercase" // Default collection
4646
}
4747

48-
return queryVectorDatabase(vdbName, query)
48+
return queryVectorDatabase(vdbName, query, collectionName)
4949
},
5050
}
5151

52-
func queryVectorDatabase(dbName, query string) error {
52+
func queryVectorDatabase(dbName, query, collectionName string) error {
5353
// Initialize progress indicator
5454
var progress *ProgressIndicator
5555
if ShouldShowProgress() {
@@ -101,7 +101,7 @@ func queryVectorDatabase(dbName, query string) error {
101101
}
102102

103103
// Call the query method
104-
result, err := client.Query(dbName, query, docLimit, "Test_collection_lowercase")
104+
result, err := client.Query(dbName, query, docLimit, collectionName)
105105
if err != nil {
106106
if progress != nil {
107107
progress.StopWithError("Query failed")

‎cli/src/search.go‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
package main
2+
3+
import (
4+
"fmt"
5+
"strings"
6+
7+
"github.com/spf13/cobra"
8+
)
9+
10+
var searchCmd = &cobra.Command{
11+
Use: "search",
12+
Short: "Search documents using natural language",
13+
Long: `Search documents in a vector database using natural language.
14+
15+
This command allows you to ask questions about documents stored in a vector database.
16+
The query agent will search through the documents and return relevant documents.`,
17+
Example: ` maestro-k search "What is the main topic of the documents?" --vdb=my-vdb
18+
maestro-k search "Find information about API endpoints" --vdb=my-vdb --doc-limit 10`,
19+
Args: cobra.ExactArgs(1),
20+
RunE: func(cmd *cobra.Command, args []string) error {
21+
query := args[0]
22+
vdbName, _ := cmd.Flags().GetString("vdb")
23+
collectionName, _ := cmd.Flags().GetString("collection")
24+
25+
// Validate inputs
26+
if strings.TrimSpace(query) == "" {
27+
return fmt.Errorf("search cannot be empty")
28+
}
29+
30+
// Use interactive selection if vdb name is not provided
31+
if vdbName == "" {
32+
var err error
33+
vdbName, err = PromptForVectorDatabase(vdbName)
34+
if err != nil {
35+
return fmt.Errorf("failed to select vector database: %w", err)
36+
}
37+
}
38+
39+
// Use default collection if not specified
40+
if collectionName == "" {
41+
collectionName = "Test_collection_lowercase" // Default collection
42+
}
43+
44+
return searchVectorDatabase(vdbName, query, collectionName)
45+
},
46+
}
47+
48+
func searchVectorDatabase(dbName, query, collectionName string) error {
49+
// Initialize progress indicator
50+
var progress *ProgressIndicator
51+
if ShouldShowProgress() {
52+
progress = NewProgressIndicator("Processing search...")
53+
progress.Start()
54+
}
55+
56+
if verbose {
57+
fmt.Println("Searching vector database...")
58+
}
59+
60+
if dryRun {
61+
if progress != nil {
62+
progress.Stop("Dry run completed")
63+
}
64+
fmt.Println("[DRY RUN] Would search vector database")
65+
return nil
66+
}
67+
68+
if progress != nil {
69+
progress.Update("Connecting to MCP server...")
70+
}
71+
72+
// Get MCP server URI
73+
serverURI, err := getMCPServerURI(mcpServerURI)
74+
if err != nil {
75+
if progress != nil {
76+
progress.StopWithError("Failed to get MCP server URI")
77+
}
78+
return fmt.Errorf("failed to get MCP server URI: %w", err)
79+
}
80+
81+
if verbose {
82+
fmt.Printf("Connecting to MCP server at: %s\n", serverURI)
83+
}
84+
85+
// Create MCP client
86+
client, err := NewMCPClient(serverURI)
87+
if err != nil {
88+
if progress != nil {
89+
progress.StopWithError("Failed to create MCP client")
90+
}
91+
return fmt.Errorf("failed to create MCP client: %w", err)
92+
}
93+
defer client.Close()
94+
95+
if progress != nil {
96+
progress.Update("Executing search...")
97+
}
98+
99+
// Call the search method
100+
result, err := client.Search(dbName, query, docLimit, collectionName)
101+
if err != nil {
102+
if progress != nil {
103+
progress.StopWithError("Search failed")
104+
}
105+
return fmt.Errorf("failed to search vector database: %w", err)
106+
}
107+
108+
if progress != nil {
109+
progress.Stop("Search completed successfully")
110+
}
111+
112+
fmt.Println(result)
113+
return nil
114+
}
115+
116+
func init() {
117+
// Add flags to search command
118+
searchCmd.Flags().String("vdb", "", "Vector database name")
119+
searchCmd.Flags().String("collection", "", "Collection name to search in")
120+
searchCmd.Flags().IntVarP(&docLimit, "doc-limit", "d", 5, "Maximum number of documents to consider")
121+
}

‎cli/tests/search_test.go‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
package main
2+
3+
import (
4+
"os/exec"
5+
"testing"
6+
)
7+
8+
// TestSearchHelp tests the search command help
9+
func TestSearchHelp(t *testing.T) {
10+
cmd := exec.Command("../maestro-k", "search", "--help")
11+
output, err := cmd.Output()
12+
13+
if err != nil {
14+
t.Fatalf("Failed to run query help command: %v", err)
15+
}
16+
17+
helpOutput := string(output)
18+
19+
// Check for expected search help content
20+
expectedContent := []string{
21+
"search",
22+
"Search documents in a vector database using natural language",
23+
"doc-limit",
24+
"Maximum number of documents to consider",
25+
"Examples:",
26+
}
27+
28+
for _, expected := range expectedContent {
29+
if !contains(helpOutput, expected) {
30+
t.Errorf("Search help output should contain '%s'", expected)
31+
}
32+
}
33+
}
34+
35+
// TestSearchNoArgs tests the search command with no arguments
36+
func TestSearchNoArgs(t *testing.T) {
37+
cmd := exec.Command("../maestro-k", "search")
38+
output, err := cmd.CombinedOutput()
39+
40+
if err == nil {
41+
t.Error("Search command should fail with no arguments")
42+
}
43+
44+
outputStr := string(output)
45+
if !contains(outputStr, "Error:") {
46+
t.Error("Search command should show error message with no arguments")
47+
}
48+
}
49+
50+
// TestSearchEmptyDatabaseName tests the search command with missing vdb flag
51+
func TestSearchEmptyDatabaseName(t *testing.T) {
52+
cmd := exec.Command("../maestro-k", "search", "test query")
53+
output, err := cmd.CombinedOutput()
54+
55+
if err == nil {
56+
t.Error("Search command should fail with missing --vdb flag")
57+
}
58+
59+
outputStr := string(output)
60+
if !contains(outputStr, "Error:") {
61+
t.Error("Search command should show error message with missing --vdb flag")
62+
}
63+
}
64+
65+
// TestSearchEmptySearch tests the search command with empty query
66+
func TestSearchEmptySearch(t *testing.T) {
67+
cmd := exec.Command("../maestro-k", "search", "", "--vdb=test-db")
68+
output, err := cmd.CombinedOutput()
69+
70+
if err == nil {
71+
t.Error("Search command should fail with empty query")
72+
}
73+
74+
outputStr := string(output)
75+
if !contains(outputStr, "Error:") {
76+
t.Error("Search command should show error message with empty query")
77+
}
78+
}
79+
80+
// TestSearchWithDryRunFlag tests the search command with dry-run flag
81+
func TestSearchWithDryRunFlag(t *testing.T) {
82+
cmd := exec.Command("../maestro-k", "search", "test query", "--vdb=test-db", "--dry-run")
83+
output, err := cmd.Output()
84+
85+
if err != nil {
86+
t.Fatalf("Search command with dry-run should not fail: %v", err)
87+
}
88+
89+
outputStr := string(output)
90+
if !contains(outputStr, "[DRY RUN]") {
91+
t.Error("Search command with dry-run should show dry run message")
92+
}
93+
}
94+
95+
// TestSearchWithSpecialCharacters tests the search command with special characters
96+
func TestSearchWithSpecialCharacters(t *testing.T) {
97+
cmd := exec.Command("../maestro-k", "search", "What's the deal with API endpoints? (v2.0)", "--vdb=test-db", "--dry-run")
98+
output, err := cmd.Output()
99+
100+
if err != nil {
101+
t.Fatalf("Search command with special characters should not fail: %v", err)
102+
}
103+
104+
outputStr := string(output)
105+
if !contains(outputStr, "[DRY RUN]") {
106+
t.Error("Search command with special characters should show dry run message")
107+
}
108+
}
109+
110+
// TestSearchWithLongSearch tests the search command with a long query
111+
func TestSearchWithLongSearch(t *testing.T) {
112+
longSearch := "This is a very long query that contains many words and should test the ability of the command to handle long input strings without any issues or problems"
113+
cmd := exec.Command("../maestro-k", "search", longSearch, "--vdb=test-db", "--dry-run")
114+
output, err := cmd.Output()
115+
116+
if err != nil {
117+
t.Fatalf("Search command with long query should not fail: %v", err)
118+
}
119+
120+
outputStr := string(output)
121+
if !contains(outputStr, "[DRY RUN]") {
122+
t.Error("Search command with long query should show dry run message")
123+
}
124+
}
125+
126+
// TestSearchInvalidDocLimit tests the search command with invalid doc-limit values
127+
func TestSearchInvalidDocLimit(t *testing.T) {
128+
testCases := []string{"abc", "1.5"}
129+
130+
for _, invalidLimit := range testCases {
131+
t.Run("InvalidLimit_"+invalidLimit, func(t *testing.T) {
132+
cmd := exec.Command("../maestro-k", "search", "test-db", "test query", "--doc-limit", invalidLimit)
133+
output, err := cmd.CombinedOutput()
134+
135+
// The command should handle invalid limits gracefully
136+
outputStr := string(output)
137+
138+
// It might fail due to argument parsing or MCP server, but shouldn't crash
139+
if err == nil && !contains(outputStr, "[DRY RUN]") {
140+
t.Error("Search command should handle invalid doc-limit gracefully")
141+
}
142+
})
143+
}
144+
}
145+
146+
// TestSearchCommandExists tests that the search command exists in the CLI
147+
func TestSearchCommandExists(t *testing.T) {
148+
cmd := exec.Command("../maestro-k", "--help")
149+
output, err := cmd.Output()
150+
151+
if err != nil {
152+
t.Fatalf("Failed to run help command: %v", err)
153+
}
154+
155+
helpOutput := string(output)
156+
if !contains(helpOutput, "search") {
157+
t.Error("Help output should contain 'search' command")
158+
}
159+
}

0 commit comments

Comments
 (0)