use std::collections::BTreeSet;
use std::convert::Infallible;
use std::fmt;
use std::future::Future;
use std::sync::Arc;
use reqwest::multipart::Part;
use serde::{Deserialize, Serialize};
use crate::KitError;
use crate::core::filter_tools_for_choice;
use crate::error::Operation;
use crate::error::ValidationViolation;
use crate::generation::{
FunctionCall, FunctionResult, ProviderFunctionTool, ProviderToolSpec, ToolChoice, Usage,
};
use crate::models::ModelId;
use crate::provider::ProviderRuntime;
use crate::providers::openai::OpenAiSession;
#[derive(Clone, Debug, PartialEq)]
pub struct SessionConfig {
instructions: Option<String>,
input_audio: Option<(String, bool)>,
output_audio: Option<(String, bool)>,
turn_detection: SessionTurnDetection,
functions: Vec<ProviderFunctionTool>,
tool_choice: Option<ToolChoice>,
}
impl Default for SessionConfig {
fn default() -> Self {
Self {
instructions: None,
input_audio: None,
output_audio: None,
turn_detection: SessionTurnDetection::Automatic,
functions: Vec::new(),
tool_choice: None,
}
}
}
impl SessionConfig {
pub fn with_instructions(mut self, instructions: impl Into<String>) -> Result<Self, KitError> {
let instructions = instructions.into();
if instructions.trim().is_empty() {
return Err(KitError::invalid_input(
ValidationViolation::Empty {
field: "session.instructions",
},
"interactive session instructions must not be blank",
));
}
self.instructions = Some(instructions);
Ok(self)
}
pub fn with_input_audio(
mut self,
media_type: impl Into<String>,
transcription: bool,
) -> Result<Self, KitError> {
self.input_audio = Some((
Self::validate_audio_format(media_type.into())?,
transcription,
));
Ok(self)
}
pub fn with_output_audio(
mut self,
media_type: impl Into<String>,
transcription: bool,
) -> Result<Self, KitError> {
self.output_audio = Some((
Self::validate_audio_format(media_type.into())?,
transcription,
));
Ok(self)
}
#[must_use]
pub const fn with_turn_detection(mut self, turn_detection: SessionTurnDetection) -> Self {
self.turn_detection = turn_detection;
self
}
pub fn with_functions(
mut self,
functions: Vec<ProviderFunctionTool>,
tool_choice: Option<ToolChoice>,
) -> Result<Self, KitError> {
let mut names = BTreeSet::new();
for function in &functions {
if !names.insert(function.name()) {
return Err(KitError::invalid_input(
ValidationViolation::Duplicate {
field: "session.functions.name",
},
format!(
"interactive function name `{}` is duplicated",
function.name()
),
));
}
}
if functions.is_empty()
&& matches!(
tool_choice.as_ref(),
Some(ToolChoice::Required | ToolChoice::Any)
)
{
return Err(KitError::invalid_input(
ValidationViolation::MissingField {
field: "session.functions",
},
"interactive required tool choice needs at least one function",
));
}
let specs = functions
.iter()
.cloned()
.map(ProviderToolSpec::Function)
.collect::<Vec<_>>();
filter_tools_for_choice(&specs, tool_choice.as_ref())?;
self.functions = functions;
self.tool_choice = tool_choice;
Ok(self)
}
#[must_use]
pub fn instructions(&self) -> Option<&str> {
self.instructions.as_deref()
}
#[must_use]
pub fn input_audio_format(&self) -> Option<&str> {
self.input_audio
.as_ref()
.map(|(media_type, _)| media_type.as_str())
}
#[must_use]
pub fn output_audio_format(&self) -> Option<&str> {
self.output_audio
.as_ref()
.map(|(media_type, _)| media_type.as_str())
}
#[must_use]
pub fn input_transcription(&self) -> bool {
self.input_audio
.as_ref()
.is_some_and(|(_, transcription)| *transcription)
}
#[must_use]
pub fn output_transcription(&self) -> bool {
self.output_audio
.as_ref()
.is_some_and(|(_, transcription)| *transcription)
}
#[must_use]
pub const fn turn_detection(&self) -> SessionTurnDetection {
self.turn_detection
}
#[must_use]
pub fn functions(&self) -> &[ProviderFunctionTool] {
&self.functions
}
#[must_use]
pub const fn tool_choice(&self) -> Option<&ToolChoice> {
self.tool_choice.as_ref()
}
fn validate_audio_format(media_type: String) -> Result<String, KitError> {
if !media_type.starts_with("audio/")
|| media_type.trim() != media_type
|| media_type.chars().any(char::is_control)
|| Part::bytes(Vec::new()).mime_str(&media_type).is_err()
{
return Err(KitError::invalid_input(
ValidationViolation::InvalidFormat {
field: "session.audio_format",
rule: "unpadded IANA audio media type",
},
"interactive audio format must be an unpadded IANA audio media type",
));
}
Ok(media_type)
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SessionTurnDetection {
#[default]
Automatic,
Manual,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
pub enum SessionInput {
Text {
text: String,
},
Audio {
media_type: String,
#[serde(with = "crate::core::content::base64_bytes")]
data: Vec<u8>,
},
FunctionResult(FunctionResult),
TurnStart,
TurnEnd,
CancelTurn,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "event", rename_all = "snake_case", deny_unknown_fields)]
pub enum SessionEvent {
TextDelta {
delta: String,
},
Audio {
media_type: String,
#[serde(with = "crate::core::content::base64_bytes")]
data: Vec<u8>,
},
InputTranscriptionDelta {
delta: String,
},
OutputTranscriptionDelta {
delta: String,
},
FunctionCall(FunctionCall),
FunctionCallsCancelled {
call_ids: Vec<String>,
},
Usage(Usage),
TurnCancelled,
Completed,
}
pub trait SessionBackend: Send + 'static {
fn id(&self) -> &str;
fn model_id(&self) -> &ModelId;
fn send(
&mut self,
input: SessionInput,
) -> impl Future<Output = Result<(), KitError>> + Send + '_;
fn recv(&mut self) -> impl Future<Output = Result<SessionEvent, KitError>> + Send + '_;
fn close(self) -> impl Future<Output = Result<(), KitError>> + Send;
}
impl SessionBackend for Infallible {
fn id(&self) -> &str {
match *self {}
}
fn model_id(&self) -> &ModelId {
match *self {}
}
async fn send(&mut self, _input: SessionInput) -> Result<(), KitError> {
match *self {}
}
async fn recv(&mut self) -> Result<SessionEvent, KitError> {
match *self {}
}
async fn close(self) -> Result<(), KitError> {
match self {}
}
}
pub(crate) enum SessionInner {
OpenAi(OpenAiSession),
}
pub struct Session {
inner: SessionInner,
provider: Arc<ProviderRuntime>,
}
impl Session {
pub(crate) fn new(inner: SessionInner, provider: Arc<ProviderRuntime>) -> Self {
Self { inner, provider }
}
}
impl SessionBackend for SessionInner {
fn id(&self) -> &str {
match self {
Self::OpenAi(session) => session.id(),
}
}
fn model_id(&self) -> &ModelId {
match self {
Self::OpenAi(session) => session.model_id(),
}
}
async fn send(&mut self, input: SessionInput) -> Result<(), KitError> {
match self {
Self::OpenAi(session) => session.send(input).await,
}
}
async fn recv(&mut self) -> Result<SessionEvent, KitError> {
match self {
Self::OpenAi(session) => session.recv().await,
}
}
async fn close(self) -> Result<(), KitError> {
match self {
Self::OpenAi(session) => session.close().await,
}
}
}
impl SessionBackend for Session {
fn id(&self) -> &str {
self.inner.id()
}
fn model_id(&self) -> &ModelId {
self.inner.model_id()
}
async fn send(&mut self, input: SessionInput) -> Result<(), KitError> {
let _admission = self.provider.admit(Operation::SessionSend)?;
self.inner
.send(input)
.await
.map_err(|error| error.at_operation(Operation::SessionSend))
}
async fn recv(&mut self) -> Result<SessionEvent, KitError> {
let _admission = self.provider.admit(Operation::SessionReceive)?;
self.inner
.recv()
.await
.map_err(|error| error.at_operation(Operation::SessionReceive))
}
async fn close(self) -> Result<(), KitError> {
self.inner
.close()
.await
.map_err(|error| error.at_operation(Operation::SessionClose))
}
}
impl fmt::Debug for Session {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("Session")
.field("id", &self.id())
.field("model_id", &self.model_id())
.finish()
}
}