sqlite_graphrag/chat_api/completion.rs
1//! Response finalisation: content extraction, JSON repair and the
2//! `max_tokens` growth policy.
3//!
4//! Everything that turns a decoded [`super::wire::ChatResponse`] into the
5//! caller-facing [`ChatCompletion`], plus the truncation arithmetic that
6//! decides how much room the next attempt gets (GAP-SG-10 / GAP-SG-70/71).
7
8use super::error::ChatError;
9use super::wire::ChatResponse;
10use super::OpenRouterChatClient;
11use crate::errors::AppError;
12use crate::retry::AttemptOutcome;
13
14/// Successful [`super::OpenRouterChatClient::complete`] result (GAP-SG-72-chat).
15///
16/// `finish_reason`, `prompt_tokens` and `completion_tokens` are the raw
17/// diagnostics OpenRouter attached to the response that ultimately succeeded
18/// (after any `max_tokens` growth retries — see [`Self::value`] and the
19/// module docs). They are `None` only when the provider omitted them.
20#[derive(Debug)]
21pub struct ChatCompletion {
22 /// Model output parsed as JSON (guaranteed to be a JSON object).
23 pub value: serde_json::Value,
24 /// Cost in USD read from `usage.cost`, or `0.0` when the provider omitted it.
25 pub cost_usd: f64,
26 /// `choices[0].finish_reason` from the response that produced `value`.
27 pub finish_reason: Option<String>,
28 /// `usage.prompt_tokens` from the response that produced `value`.
29 pub prompt_tokens: Option<u32>,
30 /// `usage.completion_tokens` from the response that produced `value`.
31 pub completion_tokens: Option<u32>,
32}
33
34impl OpenRouterChatClient {
35 /// Extracts content, repairs/parses it as JSON, and enforces the
36 /// object-shape guard, attaching `finish_reason`/token diagnostics to any
37 /// failure.
38 ///
39 /// Every failure branch below (missing content, JSON-repair failure,
40 /// non-object shape) classifies as `AttemptOutcome::Transient`. This is a
41 /// deliberate, acknowledged tension with `rules_rust_retry_com_backoff.md`
42 /// (`NUNCA retentar erros de parsing ou deserialização` /
43 /// `NUNCA retentar erros de deserialização`): those rules target DETERMINISTIC parse
44 /// errors, where retrying the identical input reproduces the identical
45 /// failure. Here the "input" is `deepseek-v4-flash:nitro` sampling
46 /// variance — the SAME prompt can legitimately produce well-formed JSON
47 /// on the next generation (see GAP-SG-10). So this is a typed, bounded
48 /// hiccup, not a retry-forever loophole: it is capped by `--max-attempts`
49 /// (GAP-SG-09/GAP-SG-21) and dead-letters once attempts are exhausted.
50 pub(super) fn finish_completion(
51 &self,
52 response: ChatResponse,
53 finish_reason: Option<String>,
54 prompt_tokens: Option<u32>,
55 completion_tokens: Option<u32>,
56 ) -> Result<ChatCompletion, ChatError> {
57 let content = response
58 .choices
59 .into_iter()
60 .next()
61 .and_then(|c| c.message.content)
62 .filter(|c| !c.trim().is_empty())
63 .ok_or_else(|| {
64 AppError::Validation(crate::i18n::validation::model_no_structured_content(
65 &self.model,
66 ))
67 })
68 .map_err(|e| {
69 ChatError::with_diagnostics(
70 e,
71 finish_reason.clone(),
72 prompt_tokens,
73 completion_tokens,
74 AttemptOutcome::Transient,
75 )
76 })?;
77
78 // GAP-SG-10: deepseek-v4-flash:nitro and similar models do not honour
79 // `json_schema` strict mode reliably — they wrap output in markdown
80 // fences, add trailing commas, or omit quotes around keys. Try a strict
81 // parse first (zero cost for well-formed JSON), then fall back to the
82 // repair pass (a Rust port of `json_repair`) before giving up.
83 let value = crate::json_repair::repair_to_value(&content).map_err(|e| {
84 ChatError::with_diagnostics(
85 AppError::Validation(crate::i18n::validation::model_json_parse_failed(
86 &self.model,
87 &e,
88 )),
89 finish_reason.clone(),
90 prompt_tokens,
91 completion_tokens,
92 AttemptOutcome::Transient,
93 )
94 })?;
95
96 // GAP-SG-10: `llm_json` coerces aggressively — free text becomes a JSON
97 // string, empty input becomes `{}`, a lone delimiter becomes `null`. The
98 // enrich JUDGE contract is ALWAYS a JSON object, so a non-object result
99 // here is a malformed/refused generation, NOT a usable value. Reject it
100 // (the enrich classifier reclassifies this as a transient model hiccup,
101 // GAP-SG-09) instead of letting a coerced scalar masquerade as a
102 // valid-but-empty result downstream.
103 if !value.is_object() {
104 return Err(ChatError::with_diagnostics(
105 AppError::Validation(crate::i18n::validation::model_non_object_json(
106 &self.model,
107 json_shape_name(&value),
108 )),
109 finish_reason,
110 prompt_tokens,
111 completion_tokens,
112 AttemptOutcome::Transient,
113 ));
114 }
115
116 let cost = response.usage.and_then(|u| u.cost).unwrap_or(0.0);
117
118 Ok(ChatCompletion {
119 value,
120 cost_usd: cost,
121 finish_reason,
122 prompt_tokens,
123 completion_tokens,
124 })
125 }
126}
127
128/// Grows `current` for the next `max_tokens` retry after a truncated
129/// (`finish_reason: "length"`) response (GAP-SG-70/71). When `current` is
130/// `None` the caller left the provider default in place, so growth starts
131/// from [`crate::constants::ENRICH_INITIAL_MAX_TOKENS`] instead of an unknown
132/// base. The result is always capped at
133/// [`crate::constants::ENRICH_MAX_TOKENS_CEILING`].
134pub(super) fn grow_max_tokens(current: Option<u32>) -> u32 {
135 let base = current.unwrap_or(crate::constants::ENRICH_INITIAL_MAX_TOKENS);
136 base.saturating_mul(crate::constants::ENRICH_MAX_TOKENS_GROWTH_FACTOR)
137 .min(crate::constants::ENRICH_MAX_TOKENS_CEILING)
138}
139
140/// Names the JSON shape of `value` for diagnostics (GAP-SG-10). Used when the
141/// repaired model output is not the object the enrich JUDGE contract requires.
142fn json_shape_name(value: &serde_json::Value) -> &'static str {
143 match value {
144 serde_json::Value::Null => "null",
145 serde_json::Value::Bool(_) => "boolean",
146 serde_json::Value::Number(_) => "number",
147 serde_json::Value::String(_) => "string",
148 serde_json::Value::Array(_) => "array",
149 serde_json::Value::Object(_) => "object",
150 }
151}