Skip to content

Commit 918862e

Browse files
ksylvanksylvan
authored
Merge pull request #1568 from ksylvan/0703-enhanced-anthropic-search-tool
Runtime Web Search Control via Command-Line Flag
2 parents 095890a + d9b8bc3 commit 918862e

File tree

6 files changed

+251
-25
lines changed

6 files changed

+251
-25
lines changed

‎cli/flags.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,8 @@ type Flags struct {
7474
ListStrategies bool `long:"liststrategies" description:"List all strategies"`
7575
ListVendors bool `long:"listvendors" description:"List all vendors"`
7676
ShellCompleteOutput bool `long:"shell-complete-list" description:"Output raw list without headers/formatting (for shell completion)"`
77+
Search bool `long:"search" description:"Enable web search tool for supported models (Anthropic)"`
78+
SearchLocation string `long:"search-location" description:"Set location for web search results (e.g., 'America/Los_Angeles')"`
7779
}
7880

7981
var debug = false
@@ -263,6 +265,8 @@ func (o *Flags) BuildChatOptions() (ret *common.ChatOptions) {
263265
Raw: o.Raw,
264266
Seed: o.Seed,
265267
ModelContextLength: o.ModelContextLength,
268+
Search: o.Search,
269+
SearchLocation: o.SearchLocation,
266270
}
267271
return
268272
}

‎common/domain.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ type ChatOptions struct {
2626
Seed int
2727
ModelContextLength int
2828
MaxTokens int
29+
Search bool
30+
SearchLocation string
2931
}
3032

3133
// NormalizeMessages remove empty messages and ensure messages order user-assist-user

‎plugins/ai/anthropic/anthropic.go

Lines changed: 46 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,6 @@ import (
55
"fmt"
66
"strings"
77

8-
"github.com/samber/lo"
9-
108
"github.com/anthropics/anthropic-sdk-go"
119
"github.com/anthropics/anthropic-sdk-go/option"
1210
"github.com/danielmiessler/fabric/chat"
@@ -16,6 +14,10 @@ import (
1614

1715
const defaultBaseUrl = "https://api.anthropic.com/"
1816

17+
const webSearchToolName = "web_search"
18+
const webSearchToolType = "web_search_20250305"
19+
const sourcesHeader = "## Sources"
20+
1921
func NewClient() (ret *Client) {
2022
vendorName := "Anthropic"
2123
ret = &Client{}
@@ -29,9 +31,6 @@ func NewClient() (ret *Client) {
2931
ret.ApiBaseURL = ret.AddSetupQuestion("API Base URL", false)
3032
ret.ApiBaseURL.Value = defaultBaseUrl
3133
ret.ApiKey = ret.PluginBase.AddSetupQuestion("API key", true)
32-
ret.UseWebTool = ret.AddSetupQuestionBool("Web Search Tool Enabled", false)
33-
ret.WebToolLocation = ret.AddSetupQuestionCustom("Web Search Tool Location", false,
34-
"Enter your approximate timezone location for web search (e.g., 'America/Los_Angeles', see https://en.wikipedia.org/wiki/List_of_tz_database_time_zones).")
3534

3635
ret.maxTokens = 4096
3736
ret.defaultRequiredUserMessage = "Hi"
@@ -49,10 +48,8 @@ func NewClient() (ret *Client) {
4948

5049
type Client struct {
5150
*plugins.PluginBase
52-
ApiBaseURL *plugins.SetupQuestion
53-
ApiKey *plugins.SetupQuestion
54-
UseWebTool *plugins.SetupQuestion
55-
WebToolLocation *plugins.SetupQuestion
51+
ApiBaseURL *plugins.SetupQuestion
52+
ApiKey *plugins.SetupQuestion
5653

5754
maxTokens int
5855
defaultRequiredUserMessage string
@@ -127,20 +124,17 @@ func (an *Client) buildMessageParams(msgs []anthropic.MessageParam, opts *common
127124
Messages: msgs,
128125
}
129126

130-
if plugins.ParseBoolElseFalse(an.UseWebTool.Value) {
127+
if opts.Search {
131128
// Build the web-search tool definition:
132129
webTool := anthropic.WebSearchTool20250305Param{
133-
Name: "web_search", // string literal instead of constant
134-
Type: "web_search_20250305", // string literal instead of constant
130+
Name: webSearchToolName,
131+
Type: webSearchToolType,
135132
CacheControl: anthropic.NewCacheControlEphemeralParam(),
136-
// Optional: restrict domains or max uses
137-
// AllowedDomains: []string{"wikipedia.org", "openai.com"},
138-
// MaxUses: anthropic.Opt[int64](5),
139133
}
140134

141-
if an.WebToolLocation.Value != "" {
135+
if opts.SearchLocation != "" {
142136
webTool.UserLocation.Type = "approximate"
143-
webTool.UserLocation.Timezone = anthropic.Opt(an.WebToolLocation.Value)
137+
webTool.UserLocation.Timezone = anthropic.Opt(opts.SearchLocation)
144138
}
145139

146140
// Wrap it in the union:
@@ -165,13 +159,42 @@ func (an *Client) Send(ctx context.Context, msgs []*chat.ChatCompletionMessage,
165159
return
166160
}
167161

168-
texts := lo.FilterMap(message.Content, func(block anthropic.ContentBlockUnion, _ int) (ret string, ok bool) {
169-
if ok = block.Type == "text" && block.Text != ""; ok {
170-
ret = block.Text
162+
var textParts []string
163+
var citations []string
164+
citationMap := make(map[string]bool) // To avoid duplicate citations
165+
166+
for _, block := range message.Content {
167+
if block.Type == "text" && block.Text != "" {
168+
textParts = append(textParts, block.Text)
169+
170+
// Extract citations from this text block
171+
for _, citation := range block.Citations {
172+
if citation.Type == "web_search_result_location" {
173+
citationKey := citation.URL + "|" + citation.Title
174+
if !citationMap[citationKey] {
175+
citationMap[citationKey] = true
176+
citationText := fmt.Sprintf("- [%s](%s)", citation.Title, citation.URL)
177+
if citation.CitedText != "" {
178+
citationText += fmt.Sprintf(" - \"%s\"", citation.CitedText)
179+
}
180+
citations = append(citations, citationText)
181+
}
182+
}
183+
}
171184
}
172-
return
173-
})
174-
ret = strings.Join(texts, "")
185+
}
186+
187+
var resultBuilder strings.Builder
188+
resultBuilder.WriteString(strings.Join(textParts, ""))
189+
190+
// Append citations if any were found
191+
if len(citations) > 0 {
192+
resultBuilder.WriteString("\n\n")
193+
resultBuilder.WriteString(sourcesHeader)
194+
resultBuilder.WriteString("\n\n")
195+
resultBuilder.WriteString(strings.Join(citations, "\n"))
196+
}
197+
ret = resultBuilder.String()
175198

176199
return
177200
}

‎plugins/ai/anthropic/anthropic_test.go

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
11
package anthropic
22

33
import (
4+
"strings"
45
"testing"
6+
7+
"github.com/anthropics/anthropic-sdk-go"
8+
"github.com/danielmiessler/fabric/common"
59
)
610

711
// Test generated using Keploy
@@ -63,3 +67,192 @@ func TestClient_ListModels_ReturnsCorrectModels(t *testing.T) {
6367
}
6468
}
6569
}
70+
71+
func TestBuildMessageParams_WithoutSearch(t *testing.T) {
72+
client := NewClient()
73+
opts := &common.ChatOptions{
74+
Model: "claude-3-5-sonnet-latest",
75+
Temperature: 0.7,
76+
Search: false,
77+
}
78+
79+
messages := []anthropic.MessageParam{
80+
anthropic.NewUserMessage(anthropic.NewTextBlock("Hello")),
81+
}
82+
83+
params := client.buildMessageParams(messages, opts)
84+
85+
if params.Tools != nil {
86+
t.Error("Expected no tools when search is disabled, got tools")
87+
}
88+
89+
if params.Model != anthropic.Model(opts.Model) {
90+
t.Errorf("Expected model %s, got %s", opts.Model, params.Model)
91+
}
92+
93+
if params.Temperature.Value != opts.Temperature {
94+
t.Errorf("Expected temperature %f, got %f", opts.Temperature, params.Temperature.Value)
95+
}
96+
}
97+
98+
func TestBuildMessageParams_WithSearch(t *testing.T) {
99+
client := NewClient()
100+
opts := &common.ChatOptions{
101+
Model: "claude-3-5-sonnet-latest",
102+
Temperature: 0.7,
103+
Search: true,
104+
}
105+
106+
messages := []anthropic.MessageParam{
107+
anthropic.NewUserMessage(anthropic.NewTextBlock("What's the weather today?")),
108+
}
109+
110+
params := client.buildMessageParams(messages, opts)
111+
112+
if params.Tools == nil {
113+
t.Fatal("Expected tools when search is enabled, got nil")
114+
}
115+
116+
if len(params.Tools) != 1 {
117+
t.Errorf("Expected 1 tool, got %d", len(params.Tools))
118+
}
119+
120+
webTool := params.Tools[0].OfWebSearchTool20250305
121+
if webTool == nil {
122+
t.Fatal("Expected web search tool, got nil")
123+
}
124+
125+
if webTool.Name != "web_search" {
126+
t.Errorf("Expected tool name 'web_search', got %s", webTool.Name)
127+
}
128+
129+
if webTool.Type != "web_search_20250305" {
130+
t.Errorf("Expected tool type 'web_search_20250305', got %s", webTool.Type)
131+
}
132+
}
133+
134+
func TestBuildMessageParams_WithSearchAndLocation(t *testing.T) {
135+
client := NewClient()
136+
opts := &common.ChatOptions{
137+
Model: "claude-3-5-sonnet-latest",
138+
Temperature: 0.7,
139+
Search: true,
140+
SearchLocation: "America/Los_Angeles",
141+
}
142+
143+
messages := []anthropic.MessageParam{
144+
anthropic.NewUserMessage(anthropic.NewTextBlock("What's the weather in San Francisco?")),
145+
}
146+
147+
params := client.buildMessageParams(messages, opts)
148+
149+
if params.Tools == nil {
150+
t.Fatal("Expected tools when search is enabled, got nil")
151+
}
152+
153+
webTool := params.Tools[0].OfWebSearchTool20250305
154+
if webTool == nil {
155+
t.Fatal("Expected web search tool, got nil")
156+
}
157+
158+
if webTool.UserLocation.Type != "approximate" {
159+
t.Errorf("Expected location type 'approximate', got %s", webTool.UserLocation.Type)
160+
}
161+
162+
if webTool.UserLocation.Timezone.Value != opts.SearchLocation {
163+
t.Errorf("Expected timezone %s, got %s", opts.SearchLocation, webTool.UserLocation.Timezone.Value)
164+
}
165+
}
166+
167+
func TestCitationFormatting(t *testing.T) {
168+
// Test the citation formatting logic by creating a mock message with citations
169+
message := &anthropic.Message{
170+
Content: []anthropic.ContentBlockUnion{
171+
{
172+
Type: "text",
173+
Text: "Based on recent research, artificial intelligence is advancing rapidly.",
174+
Citations: []anthropic.TextCitationUnion{
175+
{
176+
Type: "web_search_result_location",
177+
URL: "https://example.com/ai-research",
178+
Title: "AI Research Advances 2025",
179+
CitedText: "artificial intelligence is advancing rapidly",
180+
},
181+
{
182+
Type: "web_search_result_location",
183+
URL: "https://another-source.com/tech-news",
184+
Title: "Technology News Today",
185+
CitedText: "recent developments in AI",
186+
},
187+
},
188+
},
189+
{
190+
Type: "text",
191+
Text: " Machine learning models are becoming more sophisticated.",
192+
Citations: []anthropic.TextCitationUnion{
193+
{
194+
Type: "web_search_result_location",
195+
URL: "https://example.com/ai-research", // Duplicate URL should be deduplicated
196+
Title: "AI Research Advances 2025",
197+
CitedText: "machine learning models",
198+
},
199+
},
200+
},
201+
},
202+
}
203+
204+
// Extract text and citations using the same logic as the Send method
205+
var textParts []string
206+
var citations []string
207+
citationMap := make(map[string]bool)
208+
209+
for _, block := range message.Content {
210+
if block.Type == "text" && block.Text != "" {
211+
textParts = append(textParts, block.Text)
212+
213+
for _, citation := range block.Citations {
214+
if citation.Type == "web_search_result_location" {
215+
citationKey := citation.URL + "|" + citation.Title
216+
if !citationMap[citationKey] {
217+
citationMap[citationKey] = true
218+
citationText := "- [" + citation.Title + "](" + citation.URL + ")"
219+
if citation.CitedText != "" {
220+
citationText += " - \"" + citation.CitedText + "\""
221+
}
222+
citations = append(citations, citationText)
223+
}
224+
}
225+
}
226+
}
227+
}
228+
229+
result := strings.Join(textParts, "")
230+
if len(citations) > 0 {
231+
result += "\n\n## Sources\n\n" + strings.Join(citations, "\n")
232+
}
233+
234+
// Verify the result contains the expected text
235+
expectedText := "Based on recent research, artificial intelligence is advancing rapidly. Machine learning models are becoming more sophisticated."
236+
if !strings.Contains(result, expectedText) {
237+
t.Errorf("Expected result to contain text: %s", expectedText)
238+
}
239+
240+
// Verify citations are included
241+
if !strings.Contains(result, "## Sources") {
242+
t.Error("Expected result to contain Sources section")
243+
}
244+
245+
if !strings.Contains(result, "[AI Research Advances 2025](https://example.com/ai-research)") {
246+
t.Error("Expected result to contain first citation")
247+
}
248+
249+
if !strings.Contains(result, "[Technology News Today](https://another-source.com/tech-news)") {
250+
t.Error("Expected result to contain second citation")
251+
}
252+
253+
// Verify deduplication - should only have 2 unique citations, not 3
254+
citationCount := strings.Count(result, "- [")
255+
if citationCount != 2 {
256+
t.Errorf("Expected 2 unique citations, got %d", citationCount)
257+
}
258+
}

‎plugins/ai/dryrun/dryrun.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,12 @@ func (c *Client) formatOptions(opts *common.ChatOptions) string {
7676
if opts.ModelContextLength != 0 {
7777
builder.WriteString(fmt.Sprintf("ModelContextLength: %d\n", opts.ModelContextLength))
7878
}
79+
if opts.Search {
80+
builder.WriteString("Search: enabled\n")
81+
if opts.SearchLocation != "" {
82+
builder.WriteString(fmt.Sprintf("SearchLocation: %s\n", opts.SearchLocation))
83+
}
84+
}
7985

8086
return builder.String()
8187
}

‎plugins/plugin.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,6 @@ func (o *Setting) FillEnvFileContent(buffer *bytes.Buffer) {
152152
}
153153
buffer.WriteString("\n")
154154
}
155-
return
156155
}
157156

158157
func ParseBoolElseFalse(val string) (ret bool) {
@@ -279,7 +278,6 @@ func (o Settings) FillEnvFileContent(buffer *bytes.Buffer) {
279278
for _, setting := range o {
280279
setting.FillEnvFileContent(buffer)
281280
}
282-
return
283281
}
284282

285283
type SetupQuestions []*SetupQuestion

0 commit comments

Comments
 (0)