use std::fmt;
use std::str::FromStr;
use tea_protocol::{ModelId, ModelRef, ProfileId, ProtocolTimestamp, ProviderId, SessionId};
use thiserror::Error;
use crate::{SessionStoreError, SessionStoreFuture};
pub const MAX_SESSION_NAME_BYTES: usize = 256;
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct SessionName(String);
impl SessionName {
pub fn new(value: impl Into<String>) -> Result<Self, SessionNameError> {
let value = value.into();
let value = value.trim();
if value.is_empty() {
return Err(SessionNameError::Empty);
}
if value.len() > MAX_SESSION_NAME_BYTES {
return Err(SessionNameError::TooLong);
}
if value.chars().any(char::is_control) {
return Err(SessionNameError::ControlCharacter);
}
Ok(Self(value.to_owned()))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for SessionName {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(self.as_str())
}
}
impl FromStr for SessionName {
type Err = SessionNameError;
fn from_str(value: &str) -> Result<Self, Self::Err> {
Self::new(value)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum SessionNameError {
#[error("session name cannot be empty")]
Empty,
#[error("session name exceeds {MAX_SESSION_NAME_BYTES} bytes")]
TooLong,
#[error("session name cannot contain control characters")]
ControlCharacter,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SessionCatalogEntry {
session_id: SessionId,
name: Option<SessionName>,
updated_at: ProtocolTimestamp,
profile_id: ProfileId,
model: Option<ModelRef>,
message_count: usize,
pending_approval_count: usize,
}
impl SessionCatalogEntry {
#[allow(clippy::too_many_arguments)]
const fn new(
session_id: SessionId,
name: Option<SessionName>,
updated_at: ProtocolTimestamp,
profile_id: ProfileId,
model: Option<ModelRef>,
message_count: usize,
pending_approval_count: usize,
) -> Self {
Self {
session_id,
name,
updated_at,
profile_id,
model,
message_count,
pending_approval_count,
}
}
pub fn from_snapshot(
snapshot: &crate::SessionSnapshot,
name: Option<SessionName>,
) -> Result<Self, SessionStoreError> {
catalog_entry(snapshot, name)
}
#[must_use]
pub const fn session_id(&self) -> SessionId {
self.session_id
}
#[must_use]
pub const fn name(&self) -> Option<&SessionName> {
self.name.as_ref()
}
#[must_use]
pub const fn updated_at(&self) -> ProtocolTimestamp {
self.updated_at
}
#[must_use]
pub const fn profile_id(&self) -> &ProfileId {
&self.profile_id
}
#[must_use]
pub const fn model_id(&self) -> Option<&ModelId> {
match &self.model {
Some(model) => Some(model.model_id()),
None => None,
}
}
#[must_use]
pub const fn provider_id(&self) -> Option<&ProviderId> {
match &self.model {
Some(model) => Some(model.provider_id()),
None => None,
}
}
#[must_use]
pub const fn model_ref(&self) -> Option<&ModelRef> {
self.model.as_ref()
}
#[must_use]
pub const fn message_count(&self) -> usize {
self.message_count
}
#[must_use]
pub const fn pending_approval_count(&self) -> usize {
self.pending_approval_count
}
}
pub trait SessionCatalog: fmt::Debug + Send + Sync {
fn list_sessions(&self) -> SessionStoreFuture<'_, Vec<SessionCatalogEntry>>;
fn set_session_name(
&self,
session_id: SessionId,
name: Option<SessionName>,
) -> SessionStoreFuture<'_, ()>;
fn session_name(&self, session_id: SessionId) -> SessionStoreFuture<'_, Option<SessionName>>;
}
pub(crate) fn catalog_entry(
snapshot: &crate::SessionSnapshot,
name: Option<SessionName>,
) -> Result<SessionCatalogEntry, SessionStoreError> {
let state = snapshot.state();
let updated_at = snapshot
.records()
.last()
.map(tea_protocol::RecordEnvelope::timestamp)
.ok_or_else(|| {
SessionStoreError::new(
crate::SessionStoreErrorCode::InvalidRecord,
"stored session has no durable records",
)
})?;
Ok(SessionCatalogEntry::new(
state.session_id(),
name,
updated_at,
state.configuration().profile_id().clone(),
state.configuration().model_ref().cloned(),
state.messages().len(),
state.pending_approvals().len(),
))
}
pub(crate) fn sort_catalog(entries: &mut [SessionCatalogEntry]) {
entries.sort_by(|left, right| {
right
.updated_at()
.cmp(&left.updated_at())
.then_with(|| left.session_id().cmp(&right.session_id()))
});
}