recall_echo/graph/
llm.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5//! Minimal LLM provider trait for knowledge graph operations.
6//!
7//! recall-graph defines its own trait to stay independent of pulse-system-types.
8//! Callers implement this to bridge their actual LLM backend.
9
10use serde::{Deserialize, Serialize};
11
12use super::error::GraphError;
13
14/// Tokens a provider *reported* for one call.
15///
16/// Measured, never inferred: this type is only ever built from numbers a
17/// provider printed. Where a provider says nothing, the caller estimates and
18/// says which it is doing — see [`crate::graph::types::IngestionReport`].
19#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
20pub struct TokenUsage {
21    /// Prompt tokens, as the provider counted them.
22    pub input_tokens: u64,
23    /// Completion tokens, as the provider counted them.
24    pub output_tokens: u64,
25}
26
27impl TokenUsage {
28    /// Usage from a pair of counts, or `None` when neither was reported.
29    ///
30    /// A provider that reports one side and not the other still measured
31    /// something, and half a real number beats a whole invented one.
32    #[must_use]
33    pub fn from_counts(input_tokens: Option<u64>, output_tokens: Option<u64>) -> Option<Self> {
34        match (input_tokens, output_tokens) {
35            (None, None) => None,
36            (input, output) => Some(Self {
37                input_tokens: input.unwrap_or(0),
38                output_tokens: output.unwrap_or(0),
39            }),
40        }
41    }
42
43    /// Total tokens billed for the call.
44    #[must_use]
45    pub fn total(self) -> u64 {
46        self.input_tokens + self.output_tokens
47    }
48}
49
50/// One completion, and what it cost if the provider was willing to say.
51#[derive(Debug, Clone, PartialEq, Eq)]
52pub struct Completion {
53    /// The answer text — exactly what [`LlmProvider::complete`] returns.
54    pub text: String,
55    /// The provider's own token counts, when it reported any.
56    pub usage: Option<TokenUsage>,
57}
58
59impl Completion {
60    /// An answer from a provider that reported no usage.
61    #[must_use]
62    pub fn unmeasured(text: impl Into<String>) -> Self {
63        Self {
64            text: text.into(),
65            usage: None,
66        }
67    }
68
69    /// An answer with the provider's token counts attached.
70    #[must_use]
71    pub fn measured(text: impl Into<String>, usage: Option<TokenUsage>) -> Self {
72        Self {
73            text: text.into(),
74            usage,
75        }
76    }
77}
78
79/// Minimal LLM provider for extraction and deduplication.
80///
81/// Implementors bridge this to their actual LLM backend:
82/// - recall-echo bridges to `echo_system_types::LmProvider`
83/// - Standalone users can implement with any HTTP client
84#[async_trait::async_trait]
85pub trait LlmProvider: Send + Sync {
86    /// Send a system prompt and user message, get back a text response.
87    async fn complete(
88        &self,
89        system_prompt: &str,
90        user_message: &str,
91        max_tokens: u32,
92    ) -> Result<String, GraphError>;
93
94    /// The same call, plus whatever the provider reported about its cost.
95    ///
96    /// Defaulted to [`LlmProvider::complete`] with no usage, so an existing
97    /// implementor keeps working and simply goes on being estimated. Override
98    /// it wherever the backend prints real numbers — codex's
99    /// `turn.completed.usage`, grok's `usage`, the Anthropic and OpenAI
100    /// envelopes.
101    async fn complete_measured(
102        &self,
103        system_prompt: &str,
104        user_message: &str,
105        max_tokens: u32,
106    ) -> Result<Completion, GraphError> {
107        let text = self
108            .complete(system_prompt, user_message, max_tokens)
109            .await?;
110        Ok(Completion::unmeasured(text))
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn a_provider_that_reports_nothing_measures_nothing() {
120        assert_eq!(TokenUsage::from_counts(None, None), None);
121    }
122
123    #[test]
124    fn one_reported_side_is_still_a_measurement() {
125        assert_eq!(
126            TokenUsage::from_counts(None, Some(5)),
127            Some(TokenUsage {
128                input_tokens: 0,
129                output_tokens: 5,
130            })
131        );
132    }
133
134    #[test]
135    fn total_sums_both_sides() {
136        let usage = TokenUsage::from_counts(Some(13_658), Some(5)).unwrap();
137        assert_eq!(usage.total(), 13_663);
138    }
139}
recall_echo/graph/llm.rs

recall_echo/graph/
llm.rs
