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 2026-03-15 21:25:14 +08:00
parent 9a01e2194e
commit b71b323fd4
5 changed files with 67 additions and 1 deletions

3
.gitignore vendored
View File

@ -72,3 +72,6 @@ venv
.claude/ .claude/
!.claude/development.md # Exception: allow development docs !.claude/development.md # Exception: allow development docs
# PROJECT_CONTEXT.md # Commented out: allow tracking project context # PROJECT_CONTEXT.md # Commented out: allow tracking project context
# Tiktoken cache files
gns3server/agent/gns3_copilot/cache/tiktoken/

View File

@ -43,10 +43,17 @@ Requirements:
import json import json
import logging import logging
import os
import warnings import warnings
from pathlib import Path
from typing import Any from typing import Any
from typing import Callable from typing import Callable
# Configure tiktoken cache directory (must be set before importing tiktoken)
_cache_dir = Path(__file__).parent.parent / "cache" / "tiktoken"
_cache_dir.mkdir(parents=True, exist_ok=True)
os.environ["TIKTOKEN_CACHE_DIR"] = str(_cache_dir)
import tiktoken import tiktoken
from langchain_core.messages import BaseMessage from langchain_core.messages import BaseMessage
from langchain_core.messages import SystemMessage from langchain_core.messages import SystemMessage
@ -59,8 +66,14 @@ logger = logging.getLogger(__name__)
# ============================================================================ # ============================================================================
# Initialize tiktoken encoding (required dependency) # Initialize tiktoken encoding (required dependency)
import time
logger.info("Initializing tiktoken encoding (cl100k_base)...")
logger.info(f"Cache directory: {_cache_dir}")
logger.info("This may take a moment on first run (downloading ~1.6MB encoding file from openaipublic.blob.core.windows.net)")
start_time = time.time()
_tiktoken_encoding = tiktoken.get_encoding("cl100k_base") _tiktoken_encoding = tiktoken.get_encoding("cl100k_base")
logger.info("Using tiktoken (cl100k_base) for accurate token counting") elapsed = time.time() - start_time
logger.info(f"✓ tiktoken encoding loaded successfully (took {elapsed:.2f}s)")
# ============================================================================ # ============================================================================
# Constants # Constants

View File

@ -79,6 +79,8 @@ import time
from netmiko.huawei.huawei import HuaweiBase from netmiko.huawei.huawei import HuaweiBase
logger = logging.getLogger(__name__)
class GNS3HuaweiTelnetCE(HuaweiBase): class GNS3HuaweiTelnetCE(HuaweiBase):
""" """
@ -373,6 +375,8 @@ class GNS3HuaweiTelnetCE(HuaweiBase):
# Register the custom device type with Netmiko # Register the custom device type with Netmiko
_registered = False # Flag to prevent duplicate registration
def register_custom_device_type() -> None: def register_custom_device_type() -> None:
""" """
Register the custom GNS3HuaweiTelnetCE device type with Netmiko. 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 or running any Netmiko tasks. Call it explicitly at the appropriate
time. time.
Note: This function is idempotent - multiple calls will only register once.
Returns: Returns:
None None
""" """
global _registered
# Prevent duplicate registration
if _registered:
logger.debug("Huawei CE device type already registered, skipping")
return
# Use importlib to avoid namespace conflicts # Use importlib to avoid namespace conflicts
# Import the module using importlib to ensure we get the module, # Import the module using importlib to ensure we get the module,
# not a function # not a function
@ -423,6 +436,11 @@ def register_custom_device_type() -> None:
sd.platforms_str = "\n" + "\n".join(sd.platforms_base) sd.platforms_str = "\n" + "\n".join(sd.platforms_base)
sd.telnet_platforms_str = "\n" + "\n".join(sd.telnet_platforms) 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 # Auto-register on import
# This ensures the device type is available when the module is imported # This ensures the device type is available when the module is imported

View File

@ -46,6 +46,8 @@ import time
from netmiko.ruijie.ruijie_os import RuijieOSBase from netmiko.ruijie.ruijie_os import RuijieOSBase
logger = logging.getLogger(__name__)
class RuijieTelnetEnhanced(RuijieOSBase): class RuijieTelnetEnhanced(RuijieOSBase):
""" """
@ -249,13 +251,24 @@ class RuijieTelnetEnhanced(RuijieOSBase):
# Register the custom device type with Netmiko # Register the custom device type with Netmiko
_registered = False # Flag to prevent duplicate registration
def register_custom_device_type() -> None: def register_custom_device_type() -> None:
""" """
Register the custom RuijieTelnetEnhanced device type with Netmiko. Register the custom RuijieTelnetEnhanced device type with Netmiko.
This function adds 'gns3_ruijie_telnet' to Netmiko's CLASS_MAPPER This function adds 'gns3_ruijie_telnet' to Netmiko's CLASS_MAPPER
and updates the platforms lists. 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") sd = importlib.import_module("netmiko.ssh_dispatcher")
# Register in both mappers # Register in both mappers
@ -271,6 +284,11 @@ def register_custom_device_type() -> None:
sd.platforms_str = "\n" + "\n".join(sd.platforms_base) sd.platforms_str = "\n" + "\n".join(sd.platforms_base)
sd.telnet_platforms_str = "\n" + "\n".join(sd.telnet_platforms) 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 # Auto-register on import
try: try:

View File

@ -368,6 +368,8 @@ class VPCSTelnet(BaseConnection):
# Register the custom device type with Netmiko # Register the custom device type with Netmiko
_registered = False # Flag to prevent duplicate registration
def register_custom_device_type() -> None: def register_custom_device_type() -> None:
""" """
Register the custom VPCS Telnet device type with Netmiko. 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 IMPORTANT: This function should be called BEFORE using the VPCS device
type with Netmiko. 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 # Use importlib to avoid namespace conflicts
sd = importlib.import_module("netmiko.ssh_dispatcher") 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.platforms_str = "\n" + "\n".join(sd.platforms_base)
sd.telnet_platforms_str = "\n" + "\n".join(sd.telnet_platforms) 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") logger.info("Successfully registered VPCS Telnet device type with Netmiko")