Repository navigation
Added Brave Search Cost - #327
Conversation
📝 WalkthroughWalkthroughThis PR adds configurable per-request Brave web-search cost, exposes it via backend settings and frontend UI, counts successful web_search tool calls in agentic chat to compute total searchCost, and persists searchCost with message cost through DAO implementations. ChangesWeb Search Cost Configuration and Persistence
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Code Review
This pull request introduces agentic chat capabilities, enabling models with tool-calling support to perform web searches using a new Brave Search tool. The implementation includes an agentic chat flow, search cost tracking in the DAO layer, and UI updates for search configuration. Feedback highlights that the chat history extraction logic is fragile and fails on non-text content, the use of http.DefaultClient lacks necessary timeouts, and the API URL construction should be more robust.
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
backend/chatservice/dao/dao_sqlite.go (1)
185-231:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd the same
searchCostvalidation here to keep DAO behavior safe and consistent.Line 231 adds
searchCostdirectly into persisted totals. Without bounds/finite checks, invalid values can poison cost accounting.Suggested fix
import ( "database/sql" "encoding/json" "fmt" "log" "log/slog" + "math" proto "sortedstartup/chatservice/proto" "strings" @@ func (s *SQLiteDAO) AddChatMessageWithTokens( @@ ) (MessageSummary, error) { + if searchCost < 0 || math.IsNaN(searchCost) || math.IsInf(searchCost, 0) { + return MessageSummary{}, fmt.Errorf("invalid searchCost") + } // Handle contentImage - use NULL if empty🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@backend/chatservice/dao/dao_sqlite.go` around lines 185 - 231, Add the same validation for searchCost in AddChatMessageWithTokens: before you do cost += searchCost, import math and check if math.IsNaN(searchCost) || math.IsInf(searchCost, 0) || searchCost < 0 and handle it the same way used elsewhere (e.g. clamp to 0 or return an error consistent with DAO behavior); reference the searchCost variable and the AddChatMessageWithTokens function so you can place the check immediately after computing cost and before persisting totals.backend/chatservice/dao/dao_postgres.go (1)
217-283:⚠️ Potential issue | 🟠 Major | ⚡ Quick winValidate
searchCostbefore persisting cost totals.Line 257 and Line 282 apply raw
searchCostto stored cost. Negative or non-finite values can corrupt chat/message accounting.Suggested fix
import ( "database/sql" "encoding/json" "errors" "fmt" "log/slog" + "math" "regexp" proto "sortedstartup/chatservice/proto" "strconv" "strings" @@ func (p *PostgresDAO) AddChatMessageWithTokens( @@ ) (MessageSummary, error) { slog.Info("dao_postgres:AddChatMessageWithTokens", "userID", userID, "chatId", chatId, "role", role, "model", model) + if searchCost < 0 || math.IsNaN(searchCost) || math.IsInf(searchCost, 0) { + return MessageSummary{}, fmt.Errorf("invalid searchCost") + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@backend/chatservice/dao/dao_postgres.go` around lines 217 - 283, Validate and sanitize the searchCost parameter before using it in the DB query: inside AddChatMessageWithTokens (or the function containing variables searchCost, sqlQuery, and the p.db.QueryRow call) check that searchCost is finite and non-negative (e.g., !math.IsNaN(searchCost) && !math.IsInf(searchCost,0) && searchCost >= 0); if it fails, either set searchCost = 0 (preferred for robustness) or return an error, then proceed to run the existing sqlQuery so the CTE calc and upd_chat always get a safe numeric value instead of a negative or non-finite one.
🧹 Nitpick comments (2)
backend/chatservice/dao/dao_sqlite_test.go (1)
220-220: ⚡ Quick winAdd at least one non-zero
searchCostcase to validate the new cost path.Line 220 always passes
0, so this test suite does not verify thatsearchCostis included insummary.Costandchat_list.cost.Suggested test update
tokenTestCases := []struct { name string role string message string inputTokens int outputTokens int cachedTokens int + searchCost float64 references string ragEnabled bool expectError bool errorContains string }{ @@ { name: "Valid assistant message", @@ cachedTokens: 0, + searchCost: 0, @@ }, + { + name: "Includes non-zero search cost", + role: "assistant", + message: "Used web search", + inputTokens: 10, + outputTokens: 20, + cachedTokens: 0, + searchCost: 0.02, + references: "", + ragEnabled: false, + expectError: false, + }, @@ - 0, + tc.searchCost, @@ - expectedCost := (float64(tc.inputTokens)*model.inputTokenCost + + expectedCost := (float64(tc.inputTokens)*model.inputTokenCost + float64(tc.outputTokens)*model.outputTokenCost + - float64(tc.cachedTokens)*model.cachedTokenCost) / 1000000.0 + float64(tc.cachedTokens)*model.cachedTokenCost) / 1000000.0 + expectedCost += tc.searchCost🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@backend/chatservice/dao/dao_sqlite_test.go` at line 220, Update the test that currently passes a literal 0 for searchCost to include at least one non-zero case (e.g., 50) so the new cost aggregation path is exercised; locate the test that constructs the call with the 0 argument (the parameter named searchCost) and add another subtest or table row that supplies a non-zero searchCost, then adjust the expected values for summary.Cost and chat_list.cost in the assertions to include the searchCost contribution and ensure an assertion explicitly verifies searchCost is reflected in summary.Cost and the chat_list.cost field.backend/chatservice/sortedagents/runner_test.go (1)
26-57: ⚡ Quick winAdd a stream-usage event test path.
MockLLM.CallStreamonly emits text deltas, so the newUsageEventpath inRunStreamremains unverified. Please add a test that emits a chunk withUsageand assertsInputTokens,OutputTokens, andCachedTokenspropagation.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@backend/chatservice/sortedagents/runner_test.go` around lines 26 - 57, Add a test path that verifies the UsageEvent branch in RunStream by changing the test harness to emit a stream chunk carrying a Usage payload and asserting that RunStream propagates InputTokens, OutputTokens, and CachedTokens; specifically, update or add a test that uses MockLLM.CallStream to send one chunk with a Usage structure (instead of only text delta) and then invokes RunStream (or the test wrapper that calls RunStream) and asserts the resulting UsageEvent fields (InputTokens, OutputTokens, CachedTokens) are received and equal to the values you emitted.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@backend/chatservice/service/agentic_chat.go`:
- Around line 221-224: In the switch case handling *sortedagents.UsageEvent
(where inputTokens, outputTokens, cachedTokens are currently assigned), change
the logic to accumulate values instead of overwriting: add e.InputTokens to
inputTokens, e.OutputTokens to outputTokens, and e.CachedTokens to cachedTokens
so totals reflect multiple turns (e.g., tool calls across MAX_TURNS) rather than
only the last UsageEvent.
In `@backend/chatservice/service/brave_search_tool.go`:
- Around line 68-71: The logs in BraveSearchTool.Execute (and the other log
sites around the same method at lines noted) currently emit raw query args and
full upstream error/response bodies; update the logging to redact or truncate
sensitive content instead of logging the raw values. Locate the Execute method
on BraveSearchTool and replace direct logging of args and upstream bodies with a
short, redacted/truncated representation (e.g., log a fixed placeholder or first
N characters with "...") by introducing or calling a small helper like
redactOrTruncate(value, maxLen) and use that helper in the slog.Info/slog.Error
calls (including the error body/logging at the other referenced spots). Ensure
you do not change the returned errors or behavior—only the logged strings.
- Around line 36-52: NewBraveSearchTool and NewBraveSearchToolWithConfig
currently return http.DefaultClient (which has no timeout), so Execute's
t.httpClient.Do(req) can hang; modify these constructors to create and assign a
bounded *http.Client with a sensible Timeout (e.g., 5–15s) instead of
http.DefaultClient, or ensure the request uses a context with a deadline before
calling Execute; update NewBraveSearchTool, NewBraveSearchToolWithConfig to set
t.httpClient = &http.Client{Timeout: <duration>} and ensure Execute respects the
passed ctx when building the request.
In `@backend/chatservice/service/service.go`:
- Around line 460-483: The current code falls back to the direct LLM path on any
error from s.runAgenticChat, but runAgenticChat can already have started
streaming/persisting side effects; change the contract so the caller can
distinguish "agent never started" from "agent failed after side effects": update
s.runAgenticChat to return an error sentinel or a boolean (e.g.,
ErrAgentNotStarted or started bool) indicating whether the agent began emitting
side effects, then in the caller only perform the direct-LLM fallback when the
return indicates the agent did NOT start; keep all other errors (agent started)
as non-fallback failures to avoid double-streams/duplicate persistence.
In `@backend/chatservice/service/settings_service.go`:
- Around line 429-452: getWebSearchSettings may return settings with missing,
non-numeric, or negative Cost values which can corrupt cost accounting; update
getWebSearchSettings and the early-return branches (when sql.ErrNoRows and when
value is empty) to always populate a validated non-negative cost by parsing
settings.Cost (or the default struct) into a numeric form and if parsing fails
or is negative set settings.Cost = defaultBraveSearchCost; ensure the final
returned webSearchSettings instance always has settings.Cost set to a valid
non-negative value (use WEBSEARCH_SETTINGS_KEY, webSearchSettings,
defaultBraveSearchCost and the getWebSearchSettings function names to locate the
changes).
In `@docs/design` decisions/WebSearch_tool.md:
- Line 49: Update the limitation line that says "Brave Search API usage cost is
not yet tracked" to reflect the implemented change: either remove the line
entirely or reword it to a narrower, accurate limitation (for example "Brave
Search per-tool-call granularity is not tracked" or similar), and mention that
searchCost is now persisted (reference: searchCost and the Added Brave Search
Cost change) so the doc matches the current implementation and PR scope.
- Line 11: Fix wording and typos in the routing conditions and the
future-improvement bullet: change "websearch api key" to "WebSearch API key" and
ensure consistent casing for "WebSearch", correct "teh" to "the", and replace
the sentence at the future-improvement bullet (the line currently hard to parse)
with "If the model does not support tool calling, start the agent without tools
instead of making a direct LLM call." Update both occurrences referenced in the
doc to keep phrasing and capitalization consistent.
---
Outside diff comments:
In `@backend/chatservice/dao/dao_postgres.go`:
- Around line 217-283: Validate and sanitize the searchCost parameter before
using it in the DB query: inside AddChatMessageWithTokens (or the function
containing variables searchCost, sqlQuery, and the p.db.QueryRow call) check
that searchCost is finite and non-negative (e.g., !math.IsNaN(searchCost) &&
!math.IsInf(searchCost,0) && searchCost >= 0); if it fails, either set
searchCost = 0 (preferred for robustness) or return an error, then proceed to
run the existing sqlQuery so the CTE calc and upd_chat always get a safe numeric
value instead of a negative or non-finite one.
In `@backend/chatservice/dao/dao_sqlite.go`:
- Around line 185-231: Add the same validation for searchCost in
AddChatMessageWithTokens: before you do cost += searchCost, import math and
check if math.IsNaN(searchCost) || math.IsInf(searchCost, 0) || searchCost < 0
and handle it the same way used elsewhere (e.g. clamp to 0 or return an error
consistent with DAO behavior); reference the searchCost variable and the
AddChatMessageWithTokens function so you can place the check immediately after
computing cost and before persisting totals.
---
Nitpick comments:
In `@backend/chatservice/dao/dao_sqlite_test.go`:
- Line 220: Update the test that currently passes a literal 0 for searchCost to
include at least one non-zero case (e.g., 50) so the new cost aggregation path
is exercised; locate the test that constructs the call with the 0 argument (the
parameter named searchCost) and add another subtest or table row that supplies a
non-zero searchCost, then adjust the expected values for summary.Cost and
chat_list.cost in the assertions to include the searchCost contribution and
ensure an assertion explicitly verifies searchCost is reflected in summary.Cost
and the chat_list.cost field.
In `@backend/chatservice/sortedagents/runner_test.go`:
- Around line 26-57: Add a test path that verifies the UsageEvent branch in
RunStream by changing the test harness to emit a stream chunk carrying a Usage
payload and asserting that RunStream propagates InputTokens, OutputTokens, and
CachedTokens; specifically, update or add a test that uses MockLLM.CallStream to
send one chunk with a Usage structure (instead of only text delta) and then
invokes RunStream (or the test wrapper that calls RunStream) and asserts the
resulting UsageEvent fields (InputTokens, OutputTokens, CachedTokens) are
received and equal to the values you emitted.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b4736968-640c-4546-a76c-32bea2f479ad
⛔ Files ignored due to path filters (3)
backend/chatservice/proto/chatservice.pb.gois excluded by!**/*.pb.gobackend/chatservice/proto/chatservice_grpc.pb.gois excluded by!**/*.pb.gobackend/go.work.sumis excluded by!**/*.sum
📒 Files selected for processing (19)
backend/chatservice/dao/dao.gobackend/chatservice/dao/dao_postgres.gobackend/chatservice/dao/dao_sqlite.gobackend/chatservice/dao/dao_sqlite_test.gobackend/chatservice/dao/db/postgres/scripts/seed/2_add_support_tool_calling.up.sqlbackend/chatservice/dao/db/sqlite/scripts/seed/2_add_support_tool_calling.up.sqlbackend/chatservice/dao/models.gobackend/chatservice/service/agentic_chat.gobackend/chatservice/service/brave_search_tool.gobackend/chatservice/service/service.gobackend/chatservice/service/settings_service.gobackend/chatservice/sortedagents/llm.gobackend/chatservice/sortedagents/runner.gobackend/chatservice/sortedagents/runner_test.godocs/design decisions/WebSearch_tool.mdfrontend/proto/chatservice.tsfrontend/src/pages/setting.tsxfrontend/src/store/setting.tsproto/chatservice.proto
There was a problem hiding this comment.
🧹 Nitpick comments (1)
backend/chatservice/dao/dao_sqlite_test.go (1)
210-223: ⚡ Quick winAdd a non-zero
searchCosttest case to cover the new billing path.Current assertions only validate token-based cost because the new arg is always
0. Add at least one case withsearchCost > 0and assert bothsummary.Costandchat_list.costinclude it.Proposed minimal test update
- summary, err := daoInstance.AddChatMessageWithTokenCount( + searchCost := 0.0 + if tc.name == "Valid assistant message" { + searchCost = 0.25 + } + summary, err := daoInstance.AddChatMessageWithTokenCount( userID, chatID, tc.role, tc.message, "", model.model, tc.inputTokens, tc.outputTokens, tc.cachedTokens, - 0, + searchCost, tc.references, tc.ragEnabled, ) @@ - expectedCost := (float64(tc.inputTokens)*model.inputTokenCost + + expectedCost := (float64(tc.inputTokens)*model.inputTokenCost + float64(tc.outputTokens)*model.outputTokenCost + - float64(tc.cachedTokens)*model.cachedTokenCost) / 1000000.0 + float64(tc.cachedTokens)*model.cachedTokenCost) / 1000000.0 + expectedCost += searchCostAlso applies to: 257-265, 291-293
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@backend/chatservice/dao/dao_sqlite_test.go` around lines 210 - 223, Add a test case that passes a non-zero searchCost into AddChatMessageWithTokenCount so the billing path that sums token cost + search cost is exercised; specifically, in the tests that call AddChatMessageWithTokenCount (the call using parameters userID, chatID, tc.role, tc.message, "", model.model, tc.inputTokens, tc.outputTokens, tc.cachedTokens, 0, tc.references, tc.ragEnabled) replace the zero searchCost with a positive value in at least one table-driven case, then update assertions to check both summary.Cost and the persisted chat_list cost include that searchCost (i.e., expect summary.Cost == tokenCost+searchCost and the chat list entry’s cost reflects the same total); make the same change for the other two call sites mentioned so all paths validate non-zero searchCost handling.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@backend/chatservice/dao/dao_sqlite_test.go`:
- Around line 210-223: Add a test case that passes a non-zero searchCost into
AddChatMessageWithTokenCount so the billing path that sums token cost + search
cost is exercised; specifically, in the tests that call
AddChatMessageWithTokenCount (the call using parameters userID, chatID, tc.role,
tc.message, "", model.model, tc.inputTokens, tc.outputTokens, tc.cachedTokens,
0, tc.references, tc.ragEnabled) replace the zero searchCost with a positive
value in at least one table-driven case, then update assertions to check both
summary.Cost and the persisted chat_list cost include that searchCost (i.e.,
expect summary.Cost == tokenCost+searchCost and the chat list entry’s cost
reflects the same total); make the same change for the other two call sites
mentioned so all paths validate non-zero searchCost handling.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: f26f7201-d79b-41b4-8998-f4ef49b7990d
📒 Files selected for processing (7)
backend/chatservice/dao/dao.gobackend/chatservice/dao/dao_postgres.gobackend/chatservice/dao/dao_sqlite.gobackend/chatservice/dao/dao_sqlite_test.gobackend/chatservice/service/agentic_chat.gobackend/chatservice/service/settings_service.gobackend/chatservice/settings/settings.go
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@backend/chatservice/service/brave_search_tool.go`:
- Line 84: In BraveSearchTool:Execute the slog.Error currently logs t.apiURL
verbatim which may leak secrets; instead parse t.apiURL (use url.Parse), clear
sensitive parts (u.User = nil; u.RawQuery = "") and build a redacted string such
as u.Scheme + "://" + u.Host + u.Path (or omit the URL entirely) and log that
redacted value in the slog.Error call along with the error (err) — update the
slog.Error invocation to reference the redacted URL variable rather than
t.apiURL.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: f1fe5cad-0fb7-417a-9704-6a17778b3ee6
📒 Files selected for processing (2)
backend/chatservice/service/brave_search_tool.gobackend/chatservice/service/service.go
Summary by CodeRabbit