conduit_core/
storage.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::num::NonZeroU32;
4use std::pin::Pin;
5use std::sync::Arc;
6
7use serde_json::Value;
8use uuid::Uuid;
9
10use crate::Provider;
11
12#[derive(Clone)]
13pub struct Storages {
14    pub transit: Arc<dyn TransitStorage>,
15}
16
17pub trait TransitStorage: Send + Sync {
18    fn store_identities(
19        &self,
20        declarations: Vec<IdentityDeclaration>,
21    ) -> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send + '_>>;
22
23    fn store_usages(
24        &self,
25        declarations: Vec<UsageDeclaration>,
26    ) -> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send + '_>>;
27
28    fn list_transits(
29        &self,
30        query: TransitQuery,
31    ) -> Pin<Box<dyn Future<Output = anyhow::Result<TransitPage>> + Send + '_>>;
32
33    fn get_transits(
34        &self,
35        transit_ids: Vec<Uuid>,
36    ) -> Pin<Box<dyn Future<Output = anyhow::Result<Vec<TransitRecord>>> + Send + '_>>;
37
38    fn list_unsent(
39        &self,
40        limit: u32,
41    ) -> Pin<Box<dyn Future<Output = anyhow::Result<Vec<TransitRecord>>> + Send + '_>>;
42
43    fn mark_sent(
44        &self,
45        transit_ids: Vec<Uuid>,
46    ) -> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send + '_>>;
47}
48
49pub struct IdentityDeclaration {
50    pub transit_id: Uuid,
51    pub provider: Provider,
52    pub header_id: Option<String>,
53    pub body_id: Option<String>,
54    pub vh_headers: Option<HashMap<String, String>>,
55}
56
57pub struct UsageDeclaration {
58    pub transit_id: Uuid,
59    pub provider: Provider,
60    pub model: Option<String>,
61    pub usage: Value,
62}
63
64pub struct TransitRecord {
65    pub transit_id: Uuid,
66    pub stored_at: chrono::DateTime<chrono::Utc>,
67    pub provider: Provider,
68    pub header_id: Option<String>, // LLM provider's identifier from the response headers i.e. "debugging id"
69    pub body_id: Option<String>, // LLM provider's identifier from the response body i.e. "correlation id"
70    pub vh_headers: Option<HashMap<String, String>>, // "X-VH-" headers from the request
71    pub model: Option<String>,
72    pub usage: Option<Value>,
73    pub estimated_cost_usd: Option<f64>,
74}
75
76impl TransitRecord {
77    pub fn estimate_cost(&self) -> Option<f64> {
78        if let Some(cost) = self.estimated_cost_usd {
79            return Some(cost);
80        }
81        let model = self.model.as_deref()?;
82        let usage = self.usage.as_ref()?;
83        crate::cost::estimate_cost(self.provider, model, usage)
84    }
85}
86
87pub struct TransitQuery {
88    pub cursor: Option<chrono::DateTime<chrono::Utc>>,
89    pub direction: Direction,
90    pub limit: NonZeroU32,
91}
92
93pub enum Direction {
94    After,  // records after the cursor, or from first existing if no cursor
95    Before, // records before the cursor, or from last existing if no cursor
96}
97
98pub struct TransitPage {
99    pub records: Vec<TransitRecord>,
100    pub has_more: bool,
101}
conduit_core/storage.rs

conduit_core/
storage.rs
