mirror of
https://github.com/GNS3/gns3-server.git
synced 2026-09-23 12:21:11 +03:00
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:
parent
9a01e2194e
commit
b71b323fd4
3
.gitignore
vendored
3
.gitignore
vendored
@ -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/
|
||||||
|
|||||||
@ -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
|
||||||
|
|||||||
@ -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
|
||||||
|
|||||||
@ -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:
|
||||||
|
|||||||
@ -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")
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user