refactor(agent): Agent modular refactoring (#1487)

dbgpt/util/id_generator.py

@@ -0,0 +1,73 @@
+from typing import Optional
+
+from snowflake import Snowflake, SnowflakeGenerator
+
+_GLOBAL_GENERATOR = SnowflakeGenerator(42)
+
+
+def initialize_id_generator(
+    instance: int, *, seq: int = 0, epoch: int = 0, timestamp: Optional[int] = None
+):
+    """Initialize the global ID generator.
+
+    Args:
+        instance (int): The identifier combining both data center and machine ID in
+            traditional Snowflake algorithm. This single value serves to uniquely
+            identify the source of the ID generation request within distributed
+            environments. In standard Snowflake, this would be split into datacenter_id
+            and worker_id, but here it is combined into one for simplicity.
+
+        seq (int, optional): The initial sequence number for the generator. Default is
+            0. The sequence number increments within the same millisecond to allow
+            multiple IDs to be generated in quick succession. It resets when the
+                timestamp advances.
+
+        epoch (int, optional): The epoch time in milliseconds that acts as an offset
+            for the generator. This value helps to reduce the length of the generated
+            number by setting a custom "start time" for the timestamp component.
+            Default is 0.
+
+        timestamp (int, optional): The initial timestamp for the generator in
+            milliseconds since epoch. If not provided, the generator will use the
+            current system time. This can be used for testing or in scenarios where a
+            fixed start time is required.
+    """
+    global _GLOBAL_GENERATOR
+    _GLOBAL_GENERATOR = SnowflakeGenerator(
+        instance, seq=seq, epoch=epoch, timestamp=timestamp
+    )
+
+
+def new_id() -> int:
+    """Generate a new Snowflake ID.
+
+    Returns:
+        int: A new Snowflake ID.
+    """
+    return next(_GLOBAL_GENERATOR)
+
+
+def parse(snowflake_id: int, epoch: int = 0) -> Snowflake:
+    """Parse a Snowflake ID into its components.
+
+    Example:
+        .. code-block:: python
+
+            from dbgpt.util.id_generator import parse, new_id
+
+            snowflake_id = new_id()
+            snowflake = parse(snowflake_id)
+            print(snowflake.timestamp)
+            print(snowflake.instance)
+            print(snowflake.seq)
+            print(snowflake.datetime)
+
+    Args:
+        snowflake_id (int): The Snowflake ID to parse.
+        epoch (int, optional): The epoch time in milliseconds that acts as an offset
+            for the generator.
+
+    Returns:
+        Snowflake: The parsed Snowflake object.
+    """
+    return Snowflake.parse(snowflake_id, epoch=epoch)