use super::Result;
use bitflags::parser::to_writer;
use genai::chat::{ChatStream, ChatStreamEvent, StreamEnd};
use tokio_stream::StreamExt;
#[macro_export]
macro_rules! get_option_value {
($struct:ident.$field:ident) => {
$struct.$field.ok_or(concat!("Should have ", stringify!($field)))?
};
}
bitflags::bitflags! {
#[derive(Clone)]
pub struct Check: u8 {
const REASONING = 0b00000001;
const REASONING_USAGE = 0b00000010;
const USAGE = 0b00000100;
}
}
impl std::fmt::Debug for Check {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut buffer = String::new();
to_writer(self, &mut buffer).unwrap();
write!(f, "{}", buffer)
}
}
pub fn contains_checks(checks: Option<Check>, matching_check: Check) -> bool {
let Some(checks) = checks else { return false };
checks.contains(matching_check)
}
pub fn validate_checks(checks: Option<Check>, valid_flags: Check) -> Result<()> {
let Some(checks) = checks else { return Ok(()) };
let unsupported = checks - valid_flags;
if !unsupported.is_empty() {
return Err(format!("Unsupported flags passed for this test: {:?}", unsupported).into());
}
Ok(())
}
pub struct StreamExtract {
pub stream_end: StreamEnd,
pub content: Option<String>,
pub reasoning_content: Option<String>,
}
pub async fn extract_stream_end(mut chat_stream: ChatStream) -> Result<StreamExtract> {
let mut stream_end: Option<StreamEnd> = None;
let mut content: Vec<String> = Vec::new();
let mut reasoning_content: Vec<String> = Vec::new();
while let Some(Ok(stream_event)) = chat_stream.next().await {
match stream_event {
ChatStreamEvent::Start => (), ChatStreamEvent::Chunk(s_chunk) => content.push(s_chunk.content),
ChatStreamEvent::ReasoningChunk(s_chunk) => reasoning_content.push(s_chunk.content),
ChatStreamEvent::End(s_end) => {
stream_end = Some(s_end);
break;
}
}
}
let stream_end = stream_end.ok_or("Should have a StreamEnd event")?;
let content = (!content.is_empty()).then(|| content.join(""));
let reasoning_content = (!reasoning_content.is_empty()).then(|| reasoning_content.join(""));
Ok(StreamExtract {
stream_end,
content,
reasoning_content,
})
}