union_square 0.2.0

A proxy/wire-tap service for making LLM calls and recording everything that happens in a session for later analysis and test-case extraction
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186

//! Provider abstraction and routing for LLM API providers
//!
//! This module implements the provider abstraction layer as defined in ADR-0011,
//! supporting multiple LLM providers with URL-based routing and preserving API compatibility.

pub mod bedrock;
pub mod constants;
pub mod response_processor;

use crate::proxy::types::ProxyError;
use async_trait::async_trait;
use axum::body::Body;
use hyper::{Request, Response};
use nutype::nutype;
use std::sync::Arc;

/// Provider identifier newtype for type safety
#[nutype(
    sanitize(trim, lowercase),
    validate(not_empty, regex = r"^[a-z][a-z0-9-]*$"),
    derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Display)
)]
pub struct ProviderId(String);

impl ProviderId {
    pub const BEDROCK: &'static str = "bedrock";
    pub const OPENAI: &'static str = "openai";
    pub const ANTHROPIC: &'static str = "anthropic";

    /// Create a new ProviderId for well-known providers without validation
    pub fn bedrock() -> Self {
        Self::try_new(Self::BEDROCK.to_string()).unwrap()
    }

    pub fn openai() -> Self {
        Self::try_new(Self::OPENAI.to_string()).unwrap()
    }

    pub fn anthropic() -> Self {
        Self::try_new(Self::ANTHROPIC.to_string()).unwrap()
    }
}

/// Request ID from provider for tracking
#[nutype(
    sanitize(trim),
    validate(not_empty, regex = r"^[a-zA-Z0-9-]+$"),
    derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Display)
)]
pub struct RequestId(String);

/// Registry of all available providers
#[derive(Default)]
pub struct ProviderRegistry {
    providers: Vec<Arc<dyn Provider>>,
}

impl ProviderRegistry {
    /// Create a new provider registry
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a provider
    pub fn register(&mut self, provider: Arc<dyn Provider>) {
        self.providers.push(provider);
    }

    /// Route a request to the appropriate provider based on path
    pub fn route(&self, path: &str) -> Option<Arc<dyn Provider>> {
        self.providers
            .iter()
            .find(|p| p.matches_path(path))
            .cloned()
    }
}

/// Core provider trait for LLM API providers
#[async_trait]
pub trait Provider: Send + Sync {
    /// Provider identifier (openai, anthropic, bedrock, etc.)
    fn id(&self) -> ProviderId;

    /// Check if this provider handles the given path
    fn matches_path(&self, path: &str) -> bool;

    /// Transform Union Square URL to provider URL
    fn transform_url(&self, url: &hyper::Uri) -> Result<hyper::Uri, ProviderError>;

    /// Validate and forward the request
    async fn forward_request(
        &self,
        request: Request<Body>,
        client: &hyper_util::client::legacy::Client<
            hyper_util::client::legacy::connect::HttpConnector,
            Body,
        >,
    ) -> Result<Response<Body>, ProviderError>;

    /// Extract metadata for audit logging
    fn extract_metadata(
        &self,
        request: &Request<Body>,
        response: &Response<Body>,
    ) -> ProviderMetadata;

    /// Provider-specific health check
    async fn health_check(
        &self,
        client: &hyper_util::client::legacy::Client<
            hyper_util::client::legacy::connect::HttpConnector,
            Body,
        >,
    ) -> HealthStatus;
}

/// Provider-specific error type
#[derive(Debug, thiserror::Error)]
pub enum ProviderError {
    #[error("Invalid path: {0}")]
    InvalidPath(String),

    #[error("Invalid URL: {0}")]
    InvalidUrl(String),

    #[error("Authentication error: {0}")]
    AuthenticationError(String),

    #[error("Provider unavailable: {0}")]
    Unavailable(String),

    #[error("Request failed: {0}")]
    RequestFailed(String),

    #[error("Internal error: {0}")]
    Internal(String),
}

impl From<ProviderError> for ProxyError {
    fn from(err: ProviderError) -> Self {
        match err {
            ProviderError::InvalidPath(msg) | ProviderError::InvalidUrl(msg) => {
                ProxyError::InvalidTargetUrl(msg)
            }
            ProviderError::AuthenticationError(msg) => {
                ProxyError::Internal(format!("Authentication error: {msg}"))
            }
            ProviderError::Unavailable(msg) | ProviderError::RequestFailed(msg) => {
                ProxyError::Internal(format!("Provider error: {msg}"))
            }
            ProviderError::Internal(msg) => ProxyError::Internal(msg),
        }
    }
}

/// Provider metadata for audit logging using type-safe domain types
#[derive(Debug, Clone)]
pub struct ProviderMetadata {
    pub provider_id: ProviderId,
    pub model_id: Option<crate::providers::bedrock::types::ModelId>,
    pub request_tokens: Option<crate::providers::bedrock::types::InputTokens>,
    pub response_tokens: Option<crate::providers::bedrock::types::OutputTokens>,
    pub total_tokens: Option<crate::providers::bedrock::types::TotalTokens>,
    pub provider_request_id: Option<RequestId>,
}

impl Default for ProviderMetadata {
    fn default() -> Self {
        Self {
            provider_id: ProviderId::bedrock(), // Default to bedrock for now
            model_id: None,
            request_tokens: None,
            response_tokens: None,
            total_tokens: None,
            provider_request_id: None,
        }
    }
}

/// Health status for a provider
#[derive(Debug, Clone)]
pub enum HealthStatus {
    Healthy,
    Degraded(String),
    Unhealthy(String),
}