Skip to content

Commit 3d4768d

Browse files
committed
feat(minichat): Persist web_search/code_interpreter completed counts in chat_turns
Signed-off-by: Jalan Kulkija <ankoo67@proton.me>
1 parent cbbbd6c commit 3d4768d

File tree

13 files changed

+377
-32
lines changed

13 files changed

+377
-32
lines changed

docs/api/api.json

Lines changed: 26 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -5127,18 +5127,6 @@
51275127
"bearerAuth": []
51285128
}
51295129
],
5130-
"x-odata-orderby": {
5131-
"allowedFields": [
5132-
"id asc",
5133-
"id desc",
5134-
"name asc",
5135-
"name desc",
5136-
"country asc",
5137-
"country desc",
5138-
"created_at asc",
5139-
"created_at desc"
5140-
]
5141-
},
51425130
"x-odata-filter": {
51435131
"allowedFields": {
51445132
"country": [
@@ -5172,6 +5160,18 @@
51725160
"in"
51735161
]
51745162
}
5163+
},
5164+
"x-odata-orderby": {
5165+
"allowedFields": [
5166+
"id asc",
5167+
"id desc",
5168+
"name asc",
5169+
"name desc",
5170+
"country asc",
5171+
"country desc",
5172+
"created_at asc",
5173+
"created_at desc"
5174+
]
51755175
}
51765176
},
51775177
"post": {
@@ -5623,6 +5623,16 @@
56235623
"bearerAuth": []
56245624
}
56255625
],
5626+
"x-odata-orderby": {
5627+
"allowedFields": [
5628+
"id asc",
5629+
"id desc",
5630+
"email asc",
5631+
"email desc",
5632+
"created_at asc",
5633+
"created_at desc"
5634+
]
5635+
},
56265636
"x-odata-filter": {
56275637
"allowedFields": {
56285638
"created_at": [
@@ -5648,16 +5658,6 @@
56485658
"in"
56495659
]
56505660
}
5651-
},
5652-
"x-odata-orderby": {
5653-
"allowedFields": [
5654-
"id asc",
5655-
"id desc",
5656-
"email asc",
5657-
"email desc",
5658-
"created_at asc",
5659-
"created_at desc"
5660-
]
56615661
}
56625662
},
56635663
"post": {
@@ -8284,15 +8284,16 @@
82848284
"type",
82858285
"title",
82868286
"status",
8287-
"detail",
8288-
"instance",
8289-
"code"
8287+
"detail"
82908288
],
82918289
"properties": {
82928290
"code": {
82938291
"type": "string",
82948292
"description": "Optional machine-readable error code defined by the application."
82958293
},
8294+
"context": {
8295+
"description": "Optional structured context (e.g. `resource_type`, `resource_name`)."
8296+
},
82968297
"detail": {
82978298
"type": "string",
82988299
"description": "A human-readable explanation specific to this occurrence of the problem."

modules/mini-chat/mini-chat/src/domain/model/finalization.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,4 +119,8 @@ pub struct OrphanFinalizationInput {
119119
pub minimal_generation_floor_applied: Option<i32>,
120120
/// `started_at` — used to derive `period_starts` for quota settlement.
121121
pub started_at: OffsetDateTime,
122+
/// Completed web search tool calls persisted from the DB (0 if pod crashed before increment).
123+
pub web_search_completed_count: u32,
124+
/// Completed code interpreter tool calls persisted from the DB (0 if pod crashed before increment).
125+
pub code_interpreter_completed_count: u32,
122126
}

modules/mini-chat/mini-chat/src/domain/repos/mod.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,8 @@ pub(crate) use quota_usage_repo::{IncrementReserveParams, QuotaUsageRepository,
3232
pub(crate) use reaction_repo::{ReactionRepository, UpsertReactionParams};
3333
pub(crate) use thread_summary_repo::{ThreadSummaryModel, ThreadSummaryRepository};
3434
pub(crate) use turn_repo::{
35-
CasCompleteParams, CasTerminalParams, CreateTurnParams, TurnRepository, UpdatePreflightParams,
35+
CasCompleteParams, CasTerminalParams, CreateTurnParams, ToolCallType, TurnRepository,
36+
UpdatePreflightParams,
3637
};
3738
pub(crate) use user_limits_provider::UserLimitsProvider;
3839
pub(crate) use vector_store_repo::{InsertVectorStoreParams, VectorStoreRepository};

modules/mini-chat/mini-chat/src/domain/repos/turn_repo.rs

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,14 @@ pub struct UpdatePreflightParams {
6969
pub minimal_generation_floor_applied: i32,
7070
}
7171

72+
/// Identifies which completed tool call counter to increment.
73+
#[domain_model]
74+
#[derive(Debug, Clone, Copy)]
75+
pub enum ToolCallType {
76+
WebSearch,
77+
CodeInterpreter,
78+
}
79+
7280
/// Repository trait for turn persistence operations.
7381
#[async_trait]
7482
#[allow(dead_code)]
@@ -212,4 +220,16 @@ pub trait TurnRepository: Send + Sync {
212220
scope: &AccessScope,
213221
params: UpdatePreflightParams,
214222
) -> Result<u64, DomainError>;
223+
224+
/// Atomically increment the completed tool call counter for a turn.
225+
///
226+
/// Called from the streaming task on `ToolPhase::Done`. Best-effort:
227+
/// callers are expected to log and swallow errors.
228+
async fn increment_tool_calls<C: DBRunner>(
229+
&self,
230+
runner: &C,
231+
scope: &AccessScope,
232+
turn_id: Uuid,
233+
tool: ToolCallType,
234+
) -> Result<(), DomainError>;
215235
}

modules/mini-chat/mini-chat/src/domain/service/finalization_service.rs

Lines changed: 60 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -518,8 +518,8 @@ impl<TR: TurnRepository + 'static, MR: MessageRepository + 'static> Finalization
518518
.unwrap_or(0),
519519
settlement_path: SettlementPath::Estimated,
520520
period_starts,
521-
web_search_calls: 0,
522-
code_interpreter_calls: 0,
521+
web_search_calls: input.web_search_completed_count,
522+
code_interpreter_calls: input.code_interpreter_completed_count,
523523
};
524524

525525
let scope = modkit_security::AccessScope::allow_all();
@@ -571,8 +571,8 @@ impl<TR: TurnRepository + 'static, MR: MessageRepository + 'static> Finalization
571571
actual_credits_micro,
572572
settlement_method: settlement_method_str.to_owned(),
573573
policy_version_applied: input.policy_version_applied.unwrap_or(0),
574-
web_search_calls: 0,
575-
code_interpreter_calls: 0,
574+
web_search_calls: input.web_search_completed_count,
575+
code_interpreter_calls: input.code_interpreter_completed_count,
576576
timestamp: now,
577577
};
578578
outbox_enqueuer
@@ -1758,6 +1758,8 @@ mod tests {
17581758
policy_version_applied: Some(1),
17591759
minimal_generation_floor_applied: Some(10),
17601760
started_at: time::OffsetDateTime::now_utc(),
1761+
web_search_completed_count: 0,
1762+
code_interpreter_completed_count: 0,
17611763
};
17621764

17631765
let result = svc.finalize_orphan_turn(input, 60).await.unwrap();
@@ -1803,6 +1805,56 @@ mod tests {
18031805
);
18041806
}
18051807

1808+
#[tokio::test]
1809+
async fn finalize_orphan_tool_counts_propagate_to_usage_event() {
1810+
let db = mock_db_provider(inmem_db().await);
1811+
let (svc, outbox) = build_finalization_service(Arc::clone(&db));
1812+
1813+
let tenant_id = Uuid::new_v4();
1814+
let chat_id = Uuid::new_v4();
1815+
let turn_id = Uuid::new_v4();
1816+
let request_id = Uuid::new_v4();
1817+
let user_id = Uuid::new_v4();
1818+
1819+
insert_test_chat(&db, tenant_id, chat_id, user_id).await;
1820+
insert_running_turn(&db, tenant_id, chat_id, turn_id, request_id).await;
1821+
1822+
let conn = db.conn().unwrap();
1823+
backdate_turn_progress(&conn, turn_id).await;
1824+
1825+
let input = crate::domain::model::finalization::OrphanFinalizationInput {
1826+
turn_id,
1827+
tenant_id,
1828+
chat_id,
1829+
request_id,
1830+
user_id: Some(user_id),
1831+
requester_type: mini_chat_sdk::RequesterType::User,
1832+
effective_model: Some("gpt-5.2".to_owned()),
1833+
reserve_tokens: Some(100),
1834+
max_output_tokens_applied: Some(4096),
1835+
reserved_credits_micro: Some(1000),
1836+
policy_version_applied: Some(1),
1837+
minimal_generation_floor_applied: Some(10),
1838+
started_at: time::OffsetDateTime::now_utc(),
1839+
web_search_completed_count: 3,
1840+
code_interpreter_completed_count: 2,
1841+
};
1842+
1843+
let result = svc.finalize_orphan_turn(input, 60).await.unwrap();
1844+
assert!(result, "should be CAS winner");
1845+
1846+
let usage_events = outbox.usage_events.lock().unwrap();
1847+
assert_eq!(usage_events.len(), 1);
1848+
assert_eq!(
1849+
usage_events[0].web_search_calls, 3,
1850+
"web_search_calls must reflect persisted count"
1851+
);
1852+
assert_eq!(
1853+
usage_events[0].code_interpreter_calls, 2,
1854+
"code_interpreter_calls must reflect persisted count"
1855+
);
1856+
}
1857+
18061858
#[tokio::test]
18071859
async fn finalize_orphan_cas_loser() {
18081860
let db = mock_db_provider(inmem_db().await);
@@ -1852,6 +1904,8 @@ mod tests {
18521904
policy_version_applied: Some(1),
18531905
minimal_generation_floor_applied: Some(10),
18541906
started_at: time::OffsetDateTime::now_utc(),
1907+
web_search_completed_count: 0,
1908+
code_interpreter_completed_count: 0,
18551909
};
18561910

18571911
let result = svc.finalize_orphan_turn(input, 60).await.unwrap();
@@ -1901,6 +1955,8 @@ mod tests {
19011955
policy_version_applied: Some(1),
19021956
minimal_generation_floor_applied: Some(10),
19031957
started_at: time::OffsetDateTime::now_utc(),
1958+
web_search_completed_count: 0,
1959+
code_interpreter_completed_count: 0,
19041960
};
19051961

19061962
let result = svc.finalize_orphan_turn(input, 60).await.unwrap();

modules/mini-chat/mini-chat/src/domain/service/replay.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -289,6 +289,8 @@ mod tests {
289289
effective_model,
290290
minimal_generation_floor_applied: Some(10),
291291
web_search_enabled: false,
292+
web_search_completed_count: 0,
293+
code_interpreter_completed_count: 0,
292294
deleted_at: None,
293295
replaced_by_request_id: None,
294296
started_at: OffsetDateTime::now_utc(),

modules/mini-chat/mini-chat/src/domain/service/stream_service/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4829,6 +4829,8 @@ mod tests {
48294829
effective_model: Set(None),
48304830
minimal_generation_floor_applied: Set(None),
48314831
web_search_enabled: Set(false),
4832+
web_search_completed_count: Set(0),
4833+
code_interpreter_completed_count: Set(0),
48324834
deleted_at: Set(None),
48334835
replaced_by_request_id: Set(None),
48344836
started_at: Set(now),

modules/mini-chat/mini-chat/src/domain/service/stream_service/provider_task.rs

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use tracing::{Instrument, debug, info, warn};
88

99
use crate::domain::llm::ToolPhase;
1010
use crate::domain::ports::metric_labels::{stage, trigger};
11-
use crate::domain::repos::{MessageRepository, TurnRepository};
11+
use crate::domain::repos::{MessageRepository, ToolCallType, TurnRepository};
1212
use crate::domain::stream_events::{DoneData, ErrorData, StreamEvent};
1313
use crate::infra::db::entity::chat_turn::TurnState;
1414
use crate::infra::llm::{
@@ -394,6 +394,18 @@ pub(super) fn spawn_provider_task<TR: TurnRepository + 'static, MR: MessageRepos
394394
}
395395
ToolPhase::Done => {
396396
web_search_completed_count += 1;
397+
if let Some(ref fctx) = fin_ctx {
398+
match fctx.db.conn() {
399+
Ok(conn) => {
400+
if let Err(e) = fctx.turn_repo.increment_tool_calls(&conn, &fctx.scope, fctx.turn_id, ToolCallType::WebSearch).await {
401+
warn!(turn_id = %fctx.turn_id, error = %e, "failed to persist web_search_completed_count");
402+
}
403+
}
404+
Err(e) => {
405+
warn!(turn_id = %fctx.turn_id, error = %e, "failed to acquire DB connection for web_search_completed_count");
406+
}
407+
}
408+
}
397409
}
398410
}
399411
}
@@ -480,6 +492,18 @@ pub(super) fn spawn_provider_task<TR: TurnRepository + 'static, MR: MessageRepos
480492
}
481493
ToolPhase::Done => {
482494
code_interpreter_completed_count += 1;
495+
if let Some(ref fctx) = fin_ctx {
496+
match fctx.db.conn() {
497+
Ok(conn) => {
498+
if let Err(e) = fctx.turn_repo.increment_tool_calls(&conn, &fctx.scope, fctx.turn_id, ToolCallType::CodeInterpreter).await {
499+
warn!(turn_id = %fctx.turn_id, error = %e, "failed to persist code_interpreter_completed_count");
500+
}
501+
}
502+
Err(e) => {
503+
warn!(turn_id = %fctx.turn_id, error = %e, "failed to acquire DB connection for code_interpreter_completed_count");
504+
}
505+
}
506+
}
483507
}
484508
}
485509
}

modules/mini-chat/mini-chat/src/infra/db/entity/chat_turn.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ pub struct Model {
3030
pub effective_model: Option<String>,
3131
pub minimal_generation_floor_applied: Option<i32>,
3232
pub web_search_enabled: bool,
33+
pub web_search_completed_count: i32,
34+
pub code_interpreter_completed_count: i32,
3335
pub deleted_at: Option<OffsetDateTime>,
3436
pub replaced_by_request_id: Option<Uuid>,
3537
pub started_at: OffsetDateTime,
Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
use sea_orm_migration::prelude::*;
2+
use sea_orm_migration::sea_orm::ConnectionTrait;
3+
4+
#[derive(DeriveMigrationName)]
5+
pub struct Migration;
6+
7+
#[async_trait::async_trait]
8+
impl MigrationTrait for Migration {
9+
async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
10+
let conn = manager.get_connection();
11+
conn.execute_unprepared(
12+
"ALTER TABLE chat_turns ADD COLUMN web_search_completed_count INT NOT NULL DEFAULT 0",
13+
)
14+
.await?;
15+
conn.execute_unprepared(
16+
"ALTER TABLE chat_turns ADD COLUMN code_interpreter_completed_count INT NOT NULL DEFAULT 0",
17+
)
18+
.await?;
19+
Ok(())
20+
}
21+
22+
async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
23+
let conn = manager.get_connection();
24+
conn.execute_unprepared("ALTER TABLE chat_turns DROP COLUMN web_search_completed_count")
25+
.await?;
26+
conn.execute_unprepared(
27+
"ALTER TABLE chat_turns DROP COLUMN code_interpreter_completed_count",
28+
)
29+
.await?;
30+
Ok(())
31+
}
32+
}

0 commit comments

Comments
 (0)