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
//! Non-cryptographic random identifiers, process-local run counters, and
//! caller-supplied conversation keys. Do not use generated IDs as secrets.
//!
//! ```
//! let id = rig_core::id::generate();
//! assert_eq!(id.len(), 21);
//! ```
/// URL-safe ASCII alphabet for random identifiers.
const ALPHABET: &[u8; 64] = b"_-0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
/// Number of characters returned by [`generate`].
const DEFAULT_LEN: usize = 21;
/// Generate a 21-character URL-safe random identifier. Not cryptographically
/// secure; collisions are possible.
pub fn generate() -> String {
generate_with_len(DEFAULT_LEN)
}
/// Define a nonzero process-local counter with transparent integer serde and
/// decimal parsing/display. Values increase until wraparound; raw construction
/// and deserialization do not reserve counter values.
macro_rules! counter_id {
($(#[$meta:meta])* $name:ident) => {
$(#[$meta])*
#[derive(
Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
#[serde(transparent)]
pub struct $name(std::num::NonZeroU64);
impl $name {
fn counter() -> &'static std::sync::atomic::AtomicU64 {
static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
&NEXT
}
/// Mint the next nonzero process-local counter value. Values may repeat
/// after wraparound or collide with raw or deserialized IDs.
#[allow(clippy::new_without_default)]
pub fn new() -> Self {
// Zero is reserved even after counter wraparound.
loop {
if let Some(raw) = std::num::NonZeroU64::new(
Self::counter().fetch_add(1, std::sync::atomic::Ordering::Relaxed),
) {
return Self(raw);
}
}
}
/// Build an id from its raw value. `None` for zero, which is never
/// a valid id.
pub const fn from_raw(raw: u64) -> Option<Self> {
match std::num::NonZeroU64::new(raw) {
Some(raw) => Some(Self(raw)),
None => None,
}
}
/// The id's raw value (never zero).
pub const fn to_raw(self) -> u64 {
self.0.get()
}
}
impl From<std::num::NonZeroU64> for $name {
fn from(raw: std::num::NonZeroU64) -> Self {
Self(raw)
}
}
impl From<$name> for std::num::NonZeroU64 {
fn from(id: $name) -> Self {
id.0
}
}
impl std::fmt::Display for $name {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
std::fmt::Display::fmt(&self.0, f)
}
}
impl std::fmt::Debug for $name {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}({})", stringify!($name), self.0)
}
}
impl std::str::FromStr for $name {
type Err = ParseIdError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
s.parse::<std::num::NonZeroU64>()
.map(Self)
.map_err(|_| ParseIdError(s.to_owned()))
}
}
const _: () = assert!(
std::mem::size_of::<Option<$name>>() == std::mem::size_of::<u64>()
);
};
}
/// Error from parsing a counter id: the text was not a non-zero decimal `u64`.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("invalid id: expected a non-zero integer, got {0:?}")]
pub struct ParseIdError(String);
counter_id! {
/// Process-local run identifier shared by that run's hook events.
/// Not secret; add a host or session scope for cross-process correlation.
RunId
}
/// Caller-supplied key scoping [`ConversationMemory`] and provider conversation
/// state across runs. Wraps the string without validation.
///
/// [`ConversationMemory`]: crate::memory::ConversationMemory
#[derive(
Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
#[serde(transparent)]
pub struct ConversationId(String);
impl ConversationId {
/// Wrap a caller-supplied conversation key.
pub fn new(id: impl Into<String>) -> Self {
Self(id.into())
}
/// The id as a string slice.
pub fn as_str(&self) -> &str {
&self.0
}
/// Unwrap into the underlying string.
pub fn into_string(self) -> String {
self.0
}
}
impl std::fmt::Display for ConversationId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
impl From<String> for ConversationId {
fn from(id: String) -> Self {
Self(id)
}
}
impl From<&str> for ConversationId {
fn from(id: &str) -> Self {
Self(id.to_owned())
}
}
impl AsRef<str> for ConversationId {
fn as_ref(&self) -> &str {
&self.0
}
}
/// Generate a `len`-character, URL-safe, non-cryptographic identifier.
pub(crate) fn generate_with_len(len: usize) -> String {
std::iter::repeat_with(|| {
let idx = fastrand::usize(..ALPHABET.len());
ALPHABET.get(idx).copied().unwrap_or(b'_') as char
})
.take(len)
.collect()
}
#[cfg(test)]
mod tests;