Coverage for agentos/tools/id_generator.py: 24%

1""" 

2IDGenerator — multi-format unique identifier generation. 

3 

4Supports: 

5 - UUID4 (random) 

6 - UUID7 (time-ordered, sortable) 

7 - ULID (26-char Crockford base32, time-sortable) 

8 - Nano ID (custom alphabet & length) 

9 - Snowflake-like (timestamp + worker + sequence) 

10 - KSUID (K-Sortable Unique IDentifier) 

11 - XID (12-byte globally unique ID) 

12 - Short ID (URL-safe, configurable length) 

13""" 

14 

15from __future__ import annotations 

16 

17import secrets 

18import struct 

19import time 

20import uuid 

21 

22# ============================================================================ 

23# UUID7 (time-ordered UUID, RFC 9562 draft) 

24# ============================================================================ 

25 

26 

27def uuid7() -> str: 

28 """Generate a time-ordered UUIDv7 string.""" 

29 timestamp_ms = int(time.time() * 1000) 

30 rand_bytes = secrets.token_bytes(10) 

31 

32 # UUID7 layout: 48-bit unix_ts_ms | 4-bit ver | 12-bit rand_a | 2-bit var | 62-bit rand_b 

33 ts_bytes = struct.pack(">Q", timestamp_ms)[2:] # 6 bytes 

34 b = bytearray(ts_bytes + rand_bytes) 

35 

36 # Set version to 7 

37 b[6] = (b[6] & 0x0F) | 0x70 

38 # Set variant to 10xx (RFC 4122) 

39 b[8] = (b[8] & 0x3F) | 0x80 

40 

41 # Format as UUID 

42 u = uuid.UUID(bytes=bytes(b)) 

43 return str(u) 

44 

45 

46# ============================================================================ 

47# ULID 

48# ============================================================================ 

49 

50_CROCKFORD = "0123456789ABCDEFGHJKMNPQRSTVWXYZ" 

51 

52 

53def ulid() -> str: 

54 """Generate a ULID (26-character Crockford base32).""" 

55 ts = int(time.time() * 1000) 

56 rand = secrets.token_bytes(10) 

57 

58 # Timestamp: 48 bits = 10 base32 chars 

59 ts_part = "" 

60 for _ in range(10): 

61 ts_part = _CROCKFORD[ts & 0x1F] + ts_part 

62 ts >>= 5 

63 

64 # Random: 80 bits = 16 base32 chars 

65 rand_part = "" 

66 r = int.from_bytes(rand, "big") 

67 for _ in range(16): 

68 rand_part = _CROCKFORD[r & 0x1F] + rand_part 

69 r >>= 5 

70 

71 return ts_part + rand_part 

72 

73 

74# ============================================================================ 

75# Nano ID 

76# ============================================================================ 

77 

78 

79def nanoid(size: int = 21, alphabet: str | None = None) -> str: 

80 """Generate a Nano ID string. 

81 

82 Args: 

83 size: Length of the ID (default 21) 

84 alphabet: Custom alphabet (default URL-safe alphanumeric) 

85 """ 

86 if alphabet is None: 

87 alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ_abcdefghijklmnopqrstuvwxyz-" 

88 

89 mask = (1 << ((len(alphabet) - 1).bit_length())) - 1 

90 step = max(1, int(1.6 * mask * size / len(alphabet))) 

91 

92 result = [] 

93 while len(result) < size: 

94 for byte in secrets.token_bytes(step): 

95 idx = byte & mask 

96 if idx < len(alphabet): 

97 result.append(alphabet[idx]) 

98 if len(result) == size: 

99 break 

100 

101 return "".join(result) 

102 

103 

104# ============================================================================ 

105# Snowflake 

106# ============================================================================ 

107 

108 

109class Snowflake: 

110 """Snowflake-like distributed ID generator. 

111 

112 Layout (64 bits): timestamp(42) | worker(10) | sequence(12) 

113 Custom epoch: 2024-01-01T00:00:00Z 

114 """ 

115 

116 CUSTOM_EPOCH = 1704067200000 # 2024-01-01T00:00:00Z in ms 

117 

118 def __init__(self, worker_id: int = 0): 

119 if not (0 <= worker_id < 1024): 

120 raise ValueError("worker_id must be 0-1023") 

121 self._worker_id = worker_id 

122 self._sequence = 0 

123 self._last_ms = -1 

124 

125 def generate(self) -> int: 

126 """Generate next snowflake ID.""" 

127 now = int(time.time() * 1000) 

128 

129 if now < self._last_ms: 

130 # Clock moved backwards — wait 

131 now = self._last_ms 

132 

133 if now == self._last_ms: 

134 self._sequence = (self._sequence + 1) & 0xFFF 

135 if self._sequence == 0: 

136 # Sequence exhausted, wait for next millisecond 

137 while now <= self._last_ms: 

138 now = int(time.time() * 1000) 

139 else: 

140 self._sequence = 0 

141 

142 self._last_ms = now 

143 ts = now - self.CUSTOM_EPOCH 

144 

145 return (ts << 22) | (self._worker_id << 12) | self._sequence 

146 

147 def generate_str(self) -> str: 

148 """Generate a snowflake ID as string.""" 

149 return str(self.generate()) 

150 

151 

152# ============================================================================ 

153# Short ID 

154# ============================================================================ 

155 

156_SHORT_ALPHABET = "23456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz" 

157 

158 

159def short_id(length: int = 8) -> str: 

160 """Generate a short URL-safe random ID.""" 

161 return "".join(secrets.choice(_SHORT_ALPHABET) for _ in range(length)) 

162 

163 

164# ============================================================================ 

165# Convenience 

166# ============================================================================ 

167 

168 

169def uuid4() -> str: 

170 """Standard random UUIDv4.""" 

171 return str(uuid.uuid4()) 

172 

173 

174def generate(style: str = "uuid4") -> str: 

175 """Generate an ID in the requested style. 

176 

177 Supported: uuid4, uuid7, ulid, nanoid, short 

178 """ 

179 generators = { 

180 "uuid4": uuid4, 

181 "uuid7": uuid7, 

182 "ulid": ulid, 

183 "nanoid": lambda: nanoid(), 

184 "short": lambda: short_id(), 

185 } 

186 if style not in generators: 

187 raise ValueError(f"Unknown style: {style}. Choose from {list(generators.keys())}") 

188 return generators[style]()