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

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 

21from typing import Optional 

22 

23 

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

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

26# ============================================================================ 

27 

28def uuid7() -> str: 

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

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

31 rand_bytes = secrets.token_bytes(10) 

32 

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

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

35 b = bytearray(ts_bytes + rand_bytes) 

36 

37 # Set version to 7 

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

39 # Set variant to 10xx (RFC 4122) 

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

41 

42 # Format as UUID 

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

44 return str(u) 

45 

46 

47# ============================================================================ 

48# ULID 

49# ============================================================================ 

50 

51_CROCKFORD = "0123456789ABCDEFGHJKMNPQRSTVWXYZ" 

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 

78def nanoid(size: int = 21, alphabet: Optional[str] = None) -> str: 

79 """Generate a Nano ID string. 

80 

81 Args: 

82 size: Length of the ID (default 21) 

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

84 """ 

85 if alphabet is None: 

86 alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ_abcdefghijklmnopqrstuvwxyz-" 

87 

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

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

90 

91 result = [] 

92 while len(result) < size: 

93 for byte in secrets.token_bytes(step): 

94 idx = byte & mask 

95 if idx < len(alphabet): 

96 result.append(alphabet[idx]) 

97 if len(result) == size: 

98 break 

99 

100 return "".join(result) 

101 

102 

103# ============================================================================ 

104# Snowflake 

105# ============================================================================ 

106 

107class Snowflake: 

108 """Snowflake-like distributed ID generator. 

109 

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

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

112 """ 

113 

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

115 

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

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

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

119 self._worker_id = worker_id 

120 self._sequence = 0 

121 self._last_ms = -1 

122 

123 def generate(self) -> int: 

124 """Generate next snowflake ID.""" 

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

126 

127 if now < self._last_ms: 

128 # Clock moved backwards — wait 

129 now = self._last_ms 

130 

131 if now == self._last_ms: 

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

133 if self._sequence == 0: 

134 # Sequence exhausted, wait for next millisecond 

135 while now <= self._last_ms: 

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

137 else: 

138 self._sequence = 0 

139 

140 self._last_ms = now 

141 ts = now - self.CUSTOM_EPOCH 

142 

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

144 

145 def generate_str(self) -> str: 

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

147 return str(self.generate()) 

148 

149 

150# ============================================================================ 

151# Short ID 

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

153 

154_SHORT_ALPHABET = "23456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz" 

155 

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

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

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

159 

160 

161# ============================================================================ 

162# Convenience 

163# ============================================================================ 

164 

165def uuid4() -> str: 

166 """Standard random UUIDv4.""" 

167 return str(uuid.uuid4()) 

168 

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

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

171 

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

173 """ 

174 generators = { 

175 "uuid4": uuid4, 

176 "uuid7": uuid7, 

177 "ulid": ulid, 

178 "nanoid": lambda: nanoid(), 

179 "short": lambda: short_id(), 

180 } 

181 if style not in generators: 

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

183 return generators[style]()