use serde::de::DeserializeOwned;
use schemars::JsonSchema;
use super::super::{StreamItem, SSEProvider};
use super::core::processor::SSEStreamProcessor;
use super::extractors::{JsonPathExtractor, ExtractorConfig, openai_config, claude_config};
use futures_core::stream::Stream;
use bytes::Bytes;
use std::pin::Pin;
#[derive(Debug, Clone)]
pub struct ConfigurableSSEProvider {
config: ExtractorConfig,
}
impl ConfigurableSSEProvider {
pub fn new(config: ExtractorConfig) -> Self {
Self { config }
}
pub fn openai() -> Self {
Self::new(openai_config())
}
pub fn claude() -> Self {
Self::new(claude_config())
}
}
impl<T> SSEProvider<T> for ConfigurableSSEProvider
where
T: DeserializeOwned + JsonSchema + Send + 'static,
{
fn parse_sse_stream(
&self,
byte_stream: Pin<Box<dyn Stream<Item = Result<Bytes, super::super::error::AIError>> + Send>>
) -> Pin<Box<dyn Stream<Item = Result<StreamItem<T>, super::super::error::QueryResolverError>> + Send>> {
let extractor = JsonPathExtractor::new(self.config.clone());
SSEStreamProcessor::process_sse_stream(byte_stream, extractor)
}
}
pub type OpenAISSEProvider = ConfigurableSSEProvider;
pub type ClaudeSSEProvider = ConfigurableSSEProvider;
#[derive(Debug, Clone, Default)]
pub struct OpenAISSEProviderDefault;
impl OpenAISSEProviderDefault {
pub fn provider() -> ConfigurableSSEProvider {
ConfigurableSSEProvider::openai()
}
}
#[derive(Debug, Clone, Default)]
pub struct ClaudeSSEProviderDefault;
impl ClaudeSSEProviderDefault {
pub fn provider() -> ConfigurableSSEProvider {
ConfigurableSSEProvider::claude()
}
}
pub fn openai_provider() -> ConfigurableSSEProvider {
ConfigurableSSEProvider::openai()
}
pub fn claude_provider() -> ConfigurableSSEProvider {
ConfigurableSSEProvider::claude()
}