feat(agent): improve tiktoken initialization and device registration

- Configure tiktoken cache directory to isolate encoding files
- Add logging and timing for tiktoken initialization process
- Make Huawei CE and Ruijie Telnet device registration idempotent
- Add duplicate registration prevention with global flags
- Add logging for device type registration status
- Update .gitignore to exclude tiktoken cache files
This commit is contained in:
YueGuobin committed 2026-03-15 21:25:14 +08:00
1 parent 9a01e2194e
commit b71b323fd4
5 files changed
+67 -1

No files matched your search

@@ -79,6 +79,8 @@ import time
from netmiko.huawei.huawei import HuaweiBase
logger = logging.getLogger(__name__)
class GNS3HuaweiTelnetCE(HuaweiBase):
"""
@@ -373,6 +375,8 @@ class GNS3HuaweiTelnetCE(HuaweiBase):
# Register the custom device type with Netmiko
_registered = False # Flag to prevent duplicate registration
def register_custom_device_type() -> None:
"""
Register the custom GNS3HuaweiTelnetCE device type with Netmiko.
@@ -388,9 +392,18 @@ def register_custom_device_type() -> None:
or running any Netmiko tasks. Call it explicitly at the appropriate
time.
Note: This function is idempotent - multiple calls will only register once.
Returns:
None
"""
global _registered
# Prevent duplicate registration
if _registered:
logger.debug("Huawei CE device type already registered, skipping")
return
# Use importlib to avoid namespace conflicts
# Import the module using importlib to ensure we get the module,
# not a function
@@ -423,6 +436,11 @@ def register_custom_device_type() -> None:
sd.platforms_str = "\n" + "\n".join(sd.platforms_base)
sd.telnet_platforms_str = "\n" + "\n".join(sd.telnet_platforms)
# Mark as registered
_registered = True
logger.info("Successfully registered Huawei CE device type with Netmiko")
# Auto-register on import
# This ensures the device type is available when the module is imported
@@ -46,6 +46,8 @@ import time
from netmiko.ruijie.ruijie_os import RuijieOSBase
logger = logging.getLogger(__name__)
class RuijieTelnetEnhanced(RuijieOSBase):
"""
@@ -249,13 +251,24 @@ class RuijieTelnetEnhanced(RuijieOSBase):
# Register the custom device type with Netmiko
_registered = False # Flag to prevent duplicate registration
def register_custom_device_type() -> None:
"""
Register the custom RuijieTelnetEnhanced device type with Netmiko.
This function adds 'gns3_ruijie_telnet' to Netmiko's CLASS_MAPPER
and updates the platforms lists.
Note: This function is idempotent - multiple calls will only register once.
"""
global _registered
# Prevent duplicate registration
if _registered:
logger.debug("Ruijie Telnet device type already registered, skipping")
return
sd = importlib.import_module("netmiko.ssh_dispatcher")
# Register in both mappers
@@ -271,6 +284,11 @@ def register_custom_device_type() -> None:
sd.platforms_str = "\n" + "\n".join(sd.platforms_base)
sd.telnet_platforms_str = "\n" + "\n".join(sd.telnet_platforms)
# Mark as registered
_registered = True
logger.info("Successfully registered Ruijie Telnet device type with Netmiko")
# Auto-register on import
try:
@@ -368,6 +368,8 @@ class VPCSTelnet(BaseConnection):
# Register the custom device type with Netmiko
_registered = False # Flag to prevent duplicate registration
def register_custom_device_type() -> None:
"""
Register the custom VPCS Telnet device type with Netmiko.
@@ -377,7 +379,16 @@ def register_custom_device_type() -> None:
IMPORTANT: This function should be called BEFORE using the VPCS device
type with Netmiko.
Note: This function is idempotent - multiple calls will only register once.
"""
global _registered
# Prevent duplicate registration
if _registered:
logger.debug("VPCS Telnet device type already registered, skipping")
return
# Use importlib to avoid namespace conflicts
sd = importlib.import_module("netmiko.ssh_dispatcher")
@@ -408,6 +419,9 @@ def register_custom_device_type() -> None:
sd.platforms_str = "\n" + "\n".join(sd.platforms_base)
sd.telnet_platforms_str = "\n" + "\n".join(sd.telnet_platforms)
# Mark as registered
_registered = True
logger.info("Successfully registered VPCS Telnet device type with Netmiko")