Skip to content

Commit a8e016e

Browse files
feat: add flexible configuration file support for agent and simunet services
Add command-line arguments and environment variable support for specifying custom configuration files when starting agent and simunet services, enabling flexible deployment across different environments. Changes: - Add CLI arguments: -c/--config, --host, -p/--port, --reload/--no-reload - Support environment variables: NETDRIVER_AGENT_CONFIG and NETDRIVER_SIMUNET_CONFIG - Maintain backward compatibility with default config paths - Dynamic config reloading when custom config is specified - Fix incorrect import in conn_req_handler.py (Session from pytest -> netdriver.client.session) Usage examples: poetry run agent -c config/production.yml --no-reload poetry run simunet -c config/test-simunet.yml -p 9001 NETDRIVER_AGENT_CONFIG=config/custom.yml poetry run agent 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
1 parent 672cd77 commit a8e016e

File tree

5 files changed

+132
-5
lines changed

5 files changed

+132
-5
lines changed
Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#!/usr/bin/env python3
22
# -*- coding: utf-8 -*-
33

4+
import os
45
from dependency_injector.containers import DeclarativeContainer
56
from dependency_injector.providers import Factory, Configuration
67
from netdriver.agent.handlers.cmd_req_handler import CommandRequestHandler
@@ -9,8 +10,15 @@
910

1011
class Container(DeclarativeContainer):
1112
""" IoC container of netdriver agent. """
12-
config = Configuration(yaml_files=["config/agent/agent.yml"])
13+
config = Configuration()
1314
cmd_req_handler = Factory(CommandRequestHandler)
1415
conn_req_handler = Factory(ConnectRequestHandler)
1516

17+
18+
def get_config_file() -> str:
19+
"""Get config file path from environment variable or use default."""
20+
return os.getenv("NETDRIVER_AGENT_CONFIG", "config/agent/agent.yml")
21+
22+
1623
container = Container()
24+
container.config.from_yaml(get_config_file())

bases/netdriver/agent/handlers/conn_req_handler.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
#!/usr/bin/env python3
22
# -*- coding: utf-8 -*-
3-
from pytest import Session
43
from netdriver.agent.models.common import CommonResponse
54
from netdriver.agent.models.conn import ConnectRequest
65
from netdriver.client.pool import SessionPool
6+
from netdriver.client.session import Session
77
from netdriver.exception.errors import ConnectFailed
88

99

bases/netdriver/agent/main.py

Lines changed: 58 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@
44
This is the main module for the agent.
55
It is responsible for starting the FastAPI server.
66
"""
7+
import os
8+
import sys
9+
import argparse
710
from contextlib import asynccontextmanager
811

912
import uvicorn
@@ -68,4 +71,58 @@ async def root() -> dict:
6871

6972

7073
def start():
71-
uvicorn.run("netdriver.agent.main:app", host="0.0.0.0", port=8000, reload=True)
74+
"""Start the agent server with optional configuration file parameter."""
75+
parser = argparse.ArgumentParser(description="NetDriver Agent Server")
76+
parser.add_argument(
77+
"-c", "--config",
78+
type=str,
79+
default=None,
80+
help="Path to configuration file (default: config/agent/agent.yml or NETDRIVER_AGENT_CONFIG env var)"
81+
)
82+
parser.add_argument(
83+
"--host",
84+
type=str,
85+
default="0.0.0.0",
86+
help="Host to bind (default: 0.0.0.0)"
87+
)
88+
parser.add_argument(
89+
"-p", "--port",
90+
type=int,
91+
default=8000,
92+
help="Port to bind (default: 8000)"
93+
)
94+
parser.add_argument(
95+
"--reload",
96+
action="store_true",
97+
default=True,
98+
help="Enable auto-reload (default: True)"
99+
)
100+
parser.add_argument(
101+
"--no-reload",
102+
action="store_true",
103+
help="Disable auto-reload"
104+
)
105+
106+
args = parser.parse_args()
107+
108+
# Set config file path via environment variable if specified
109+
if args.config:
110+
os.environ["NETDRIVER_AGENT_CONFIG"] = args.config
111+
# Reload container configuration with new config file
112+
container.config.from_yaml(args.config)
113+
# Reconfigure logging with new config
114+
logman.configure_logman(
115+
level=container.config.logging.level(),
116+
intercept_loggers=container.config.logging.intercept_loggers(),
117+
log_file=container.config.logging.log_file()
118+
)
119+
120+
# Handle reload flag
121+
reload = args.reload and not args.no_reload
122+
123+
uvicorn.run(
124+
"netdriver.agent.main:app",
125+
host=args.host,
126+
port=args.port,
127+
reload=reload
128+
)

bases/netdriver/simunet/containers.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#!/usr/bin/env python3
22
# -*- coding: utf-8 -*-
3+
import os
34
from dependency_injector.containers import DeclarativeContainer
45
from dependency_injector.providers import Configuration
56

@@ -9,5 +10,10 @@ class Container(DeclarativeContainer):
910
config = Configuration()
1011

1112

13+
def get_config_file() -> str:
14+
"""Get config file path from environment variable or use default."""
15+
return os.getenv("NETDRIVER_SIMUNET_CONFIG", "config/simunet/simunet.yml")
16+
17+
1218
container = Container()
13-
container.config.from_yaml("config/simunet/simunet.yml", required=True)
19+
container.config.from_yaml(get_config_file(), required=True)

bases/netdriver/simunet/main.py

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#!/usr/bin/env python3.10.6
22
# -*- coding: utf-8 -*-
3+
import os
4+
import argparse
35
import asyncio
46
from typing import AsyncGenerator
57

@@ -54,4 +56,58 @@ async def on_shutdown() -> None:
5456

5557

5658
def start():
57-
uvicorn.run("netdriver.simunet.main:app", host="0.0.0.0", port=8001, reload=True)
59+
"""Start the simunet server with optional configuration file parameter."""
60+
parser = argparse.ArgumentParser(description="NetDriver SimuNet Server")
61+
parser.add_argument(
62+
"-c", "--config",
63+
type=str,
64+
default=None,
65+
help="Path to configuration file (default: config/simunet/simunet.yml or NETDRIVER_SIMUNET_CONFIG env var)"
66+
)
67+
parser.add_argument(
68+
"--host",
69+
type=str,
70+
default="0.0.0.0",
71+
help="Host to bind (default: 0.0.0.0)"
72+
)
73+
parser.add_argument(
74+
"-p", "--port",
75+
type=int,
76+
default=8001,
77+
help="Port to bind (default: 8001)"
78+
)
79+
parser.add_argument(
80+
"--reload",
81+
action="store_true",
82+
default=True,
83+
help="Enable auto-reload (default: True)"
84+
)
85+
parser.add_argument(
86+
"--no-reload",
87+
action="store_true",
88+
help="Disable auto-reload"
89+
)
90+
91+
args = parser.parse_args()
92+
93+
# Set config file path via environment variable if specified
94+
if args.config:
95+
os.environ["NETDRIVER_SIMUNET_CONFIG"] = args.config
96+
# Reload container configuration with new config file
97+
container.config.from_yaml(args.config, required=True)
98+
# Reconfigure logging with new config
99+
logman.configure_logman(
100+
level=container.config.logging.level(),
101+
intercept_loggers=container.config.logging.intercept_loggers(),
102+
log_file=container.config.logging.log_file()
103+
)
104+
105+
# Handle reload flag
106+
reload = args.reload and not args.no_reload
107+
108+
uvicorn.run(
109+
"netdriver.simunet.main:app",
110+
host=args.host,
111+
port=args.port,
112+
reload=reload
113+
)

0 commit comments

Comments
 (0)