diff --git a/gns3server/agent/gns3_copilot/agent/gns3_copilot.py b/gns3server/agent/gns3_copilot/agent/gns3_copilot.py index 16b78733d..6d6ae0be6 100644 --- a/gns3server/agent/gns3_copilot/agent/gns3_copilot.py +++ b/gns3server/agent/gns3_copilot/agent/gns3_copilot.py @@ -66,6 +66,9 @@ from gns3server.agent.gns3_copilot.agent.model_factory import ( create_title_model, ) from gns3server.agent.gns3_copilot.gns3_client import GNS3TopologyTool +from gns3server.agent.gns3_copilot.gns3_client.context_helpers import ( + get_current_llm_config, +) from gns3server.agent.gns3_copilot.prompts import TITLE_PROMPT, load_system_prompt from gns3server.agent.gns3_copilot.tools_v2 import ( ExecuteMultipleDeviceConfigCommands, @@ -148,7 +151,6 @@ def llm_call(state: dict, config: RunnableConfig | None = None): logger.info("LLM call node invoked") # Get llm_config from request-scoped context variable - from gns3server.agent.gns3_copilot.gns3_client import get_current_llm_config llm_config = get_current_llm_config() if not llm_config: @@ -246,7 +248,6 @@ def generate_title(state: MessagesState, config: RunnableConfig | None = None) - """ # Get llm_config from request-scoped context variable - from gns3server.agent.gns3_copilot.gns3_client import get_current_llm_config llm_config = get_current_llm_config() if not llm_config: diff --git a/gns3server/agent/gns3_copilot/agent_service.py b/gns3server/agent/gns3_copilot/agent_service.py index a46d79bca..911b2c537 100644 --- a/gns3server/agent/gns3_copilot/agent_service.py +++ b/gns3server/agent/gns3_copilot/agent_service.py @@ -45,6 +45,10 @@ from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver from gns3server.agent.gns3_copilot.agent.gns3_copilot import agent_builder from gns3server.agent.gns3_copilot.chat_sessions_repository import ChatSessionsRepository +from gns3server.agent.gns3_copilot.gns3_client.context_helpers import ( + set_current_jwt_token, + set_current_llm_config, +) from gns3server.agent.gns3_copilot.utils.message_converters import convert_langchain_to_openai log = logging.getLogger(__name__) @@ -239,11 +243,9 @@ class AgentService: # Set request-scoped context variables (memory-only, not persisted) if jwt_token: - from gns3server.agent.gns3_copilot.gns3_client import set_current_jwt_token set_current_jwt_token(jwt_token) log.debug("JWT token set in context") if llm_config: - from gns3server.agent.gns3_copilot.gns3_client import set_current_llm_config set_current_llm_config(llm_config) log.debug("LLM config set in context: provider=%s, model=%s", llm_config.get("provider"), llm_config.get("model")) diff --git a/gns3server/agent/gns3_copilot/gns3_client/__init__.py b/gns3server/agent/gns3_copilot/gns3_client/__init__.py index b515dcc81..aff51a8ba 100644 --- a/gns3server/agent/gns3_copilot/gns3_client/__init__.py +++ b/gns3server/agent/gns3_copilot/gns3_client/__init__.py @@ -52,10 +52,12 @@ from .connector_factory import ( get_gns3_connector_with_llm_config, get_gns3_server_host, get_llm_config, - set_current_jwt_token, +) +from .context_helpers import ( get_current_jwt_token, - set_current_llm_config, get_current_llm_config, + set_current_jwt_token, + set_current_llm_config, ) from .custom_gns3fy import ( CONSOLE_TYPES, diff --git a/gns3server/agent/gns3_copilot/gns3_client/connector_factory.py b/gns3server/agent/gns3_copilot/gns3_client/connector_factory.py index 920121f87..e18fb89d8 100644 --- a/gns3server/agent/gns3_copilot/gns3_client/connector_factory.py +++ b/gns3server/agent/gns3_copilot/gns3_client/connector_factory.py @@ -42,50 +42,20 @@ import asyncio import concurrent.futures import logging import threading -from contextvars import ContextVar from typing import Optional from uuid import UUID # Local imports from gns3server.agent.gns3_copilot.gns3_client.custom_gns3fy import Gns3Connector +from gns3server.agent.gns3_copilot.gns3_client.context_helpers import ( + get_current_jwt_token, + get_current_llm_config, + set_current_jwt_token, + set_current_llm_config, +) logger = logging.getLogger(__name__) -# Context variables for request-scoped data -# Automatically cleaned up when request context ends -_jwt_token_context: ContextVar[Optional[str]] = ContextVar("_jwt_token_context", default=None) -_llm_config_context: ContextVar[Optional[dict]] = ContextVar("_llm_config_context", default=None) - -def set_current_jwt_token(token: str) -> None: - """Set the JWT token for the current request context.""" - _jwt_token_context.set(token) - logger.debug("JWT token set in context") - -def get_current_jwt_token() -> Optional[str]: - """Get the JWT token for the current request context.""" - token = _jwt_token_context.get() - if token: - logger.debug("JWT token retrieved from context") - else: - logger.warning("JWT token not found in context") - return token - -def set_current_llm_config(config: dict) -> None: - """Set the LLM config for the current request context.""" - _llm_config_context.set(config) - logger.debug("LLM config set in context: provider=%s, model=%s", - config.get("provider"), config.get("model")) - -def get_current_llm_config() -> Optional[dict]: - """Get the LLM config for the current request context.""" - config = _llm_config_context.get() - if config: - logger.debug("LLM config retrieved from context: provider=%s, model=%s", - config.get("provider"), config.get("model")) - else: - logger.warning("LLM config not found in context") - return config - # Fallback default URL DEFAULT_GNS3_URL = "http://127.0.0.1:3080" diff --git a/gns3server/agent/gns3_copilot/gns3_client/context_helpers.py b/gns3server/agent/gns3_copilot/gns3_client/context_helpers.py new file mode 100644 index 000000000..acae30aaa --- /dev/null +++ b/gns3server/agent/gns3_copilot/gns3_client/context_helpers.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: GPL-3.0-or-later +# +# GNS3-Copilot - AI-powered Network Lab Assistant for GNS3 +# +# This file is part of GNS3-Copilot project. +# +# GNS3-Copilot is free software: you can redistribute it and/or modify it +# under the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or (at your +# option) any later version. +# +# GNS3-Copilot is distributed in the hope that it will be useful, but +# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY +# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License +# for more details. +# +# You should have received a copy of the GNU General Public License +# along with GNS3-Copilot. If not, see . +# +# Copyright (C) 2025 Guobin Yue +# Author: Guobin Yue +# +# Project Home: https://github.com/yueguobin/gns3-copilot +# + +""" +Request-Scoped Context Helpers for GNS3-Copilot + +This module provides context variable management for request-scoped data +(JWT tokens and LLM configurations) using Python's contextvars module. + +These functions are separated into their own module to avoid circular imports +between agent_service.py, gns3_copilot.py, and gns3_client modules. + +Key Features: +- Thread-safe and async-safe request-scoped context management +- Automatic cleanup when request context ends +- No manual cleanup required + +Usage: + from gns3server.agent.gns3_copilot.gns3_client.context_helpers import ( + set_current_jwt_token, + get_current_jwt_token, + set_current_llm_config, + get_current_llm_config, + ) + + # In request handler + set_current_jwt_token(token) + set_current_llm_config(config) + + # In downstream code + token = get_current_jwt_token() + config = get_current_llm_config() +""" + +import logging +from contextvars import ContextVar +from typing import Optional + +logger = logging.getLogger(__name__) + +# Context variables for request-scoped data +# Automatically cleaned up when request context ends +_jwt_token_context: ContextVar[Optional[str]] = ContextVar("_jwt_token_context", default=None) +_llm_config_context: ContextVar[Optional[dict]] = ContextVar("_llm_config_context", default=None) + + +def set_current_jwt_token(token: str) -> None: + """Set the JWT token for the current request context. + + Args: + token: JWT token string + """ + _jwt_token_context.set(token) + logger.debug("JWT token set in context") + + +def get_current_jwt_token() -> Optional[str]: + """Get the JWT token for the current request context. + + Returns: + JWT token string if available, None otherwise + """ + token = _jwt_token_context.get() + if token: + logger.debug("JWT token retrieved from context") + else: + logger.warning("JWT token not found in context") + return token + + +def set_current_llm_config(config: dict) -> None: + """Set the LLM config for the current request context. + + Args: + config: LLM configuration dictionary with provider, model, api_key, etc. + """ + _llm_config_context.set(config) + logger.debug("LLM config set in context: provider=%s, model=%s", + config.get("provider"), config.get("model")) + + +def get_current_llm_config() -> Optional[dict]: + """Get the LLM config for the current request context. + + Returns: + LLM configuration dict if available, None otherwise + """ + config = _llm_config_context.get() + if config: + logger.debug("LLM config retrieved from context: provider=%s, model=%s", + config.get("provider"), config.get("model")) + else: + logger.warning("LLM config not found in context") + return config + + +__all__ = [ + "set_current_jwt_token", + "get_current_jwt_token", + "set_current_llm_config", + "get_current_llm_config", +]