硬编码改为配置文件

This commit is contained in:
2026-08-25 20:35:52 +08:00
parent bd5926aeed
commit b3895a905e
7 changed files with 113 additions and 18 deletions
+1 -1
View File
@@ -1,3 +1,3 @@
__pycache__/ __pycache__/
*.pyc *.pyc
output/ output/
+11 -3
View File
@@ -6,17 +6,25 @@ from agent.prompt_builder import PromptBuilder
from agent.api_client import APIClient from agent.api_client import APIClient
from agent.output_writer import OutputWriter from agent.output_writer import OutputWriter
from agent.models import ProgramMeta, FileInfo, CopyField, KeyInfo, TableColumn, TableInfo from agent.models import ProgramMeta, FileInfo, CopyField, KeyInfo, TableColumn, TableInfo
from agent.config import load_config
def generate(design_md: str, source_cbl: str, file_db_md: str, def generate(design_md: str, source_cbl: str, file_db_md: str,
cpy_dir: str, db_md: str, output_dir: str = "output", cpy_dir: str, db_md: str, output_dir: str = "output",
api_key: str = "sk-6156cccdc9c14d949cf5bfc5afc67a03", api_key: str = None, api_model: str = None,
api_model: str = "deepseek-v4-flash",
rules_dir: str = "rules", rules_dir: str = "rules",
max_tokens: int = 32768) -> dict: max_tokens: int = None) -> dict:
"""生成测试数据的主入口函数。""" """生成测试数据的主入口函数。"""
import os import os
cfg = load_config()['api']
if api_key is None:
api_key = cfg['api_key']
if api_model is None:
api_model = cfg['model']
if max_tokens is None:
max_tokens = cfg['max_tokens']
print(f"== 解析入力: {design_md}") print(f"== 解析入力: {design_md}")
parser = InputParser(design_md, source_cbl, file_db_md, cpy_dir, db_md) parser = InputParser(design_md, source_cbl, file_db_md, cpy_dir, db_md)
+12 -10
View File
@@ -4,20 +4,22 @@ from typing import Dict, Any, Optional
import requests import requests
from agent.config import load_config
class APIClient: class APIClient:
"""DeepSeek API 客户端,含重试逻辑。""" """DeepSeek API 客户端,含重试逻辑。"""
def __init__(self, api_key: str, model: str = 'deepseek-v4-flash', def __init__(self, api_key: str = None, model: str = None,
base_url: str = 'https://api.deepseek.com/chat/completions', base_url: str = None, max_retries: int = None,
max_retries: int = 3, timeout: int = 120, timeout: int = None, max_tokens: int = None):
max_tokens: int = 32768): cfg = load_config()['api']
self.api_key = api_key self.api_key = api_key if api_key else cfg['api_key']
self.model = model self.model = model if model else cfg['model']
self.base_url = base_url self.base_url = base_url if base_url else cfg['base_url']
self.max_retries = max_retries self.max_retries = max_retries if max_retries is not None else cfg['max_retries']
self.timeout = timeout self.timeout = timeout if timeout is not None else cfg['timeout']
self.max_tokens = max_tokens self.max_tokens = max_tokens if max_tokens is not None else cfg['max_tokens']
def generate(self, prompt: str) -> Dict[str, Any]: def generate(self, prompt: str) -> Dict[str, Any]:
"""发送 prompt 并返回 AI 生成的结果。""" """发送 prompt 并返回 AI 生成的结果。"""
+62
View File
@@ -0,0 +1,62 @@
import copy
import json
import os
DEFAULT_CONFIG = {
"api": {
"api_key": None,
"model": "deepseek-v4-flash",
"base_url": "https://api.deepseek.com/chat/completions",
"timeout": 120,
"max_retries": 3,
"max_tokens": 32768,
}
}
ENV_OVERRIDES = {
"api_key": "DEEPSEEK_API_KEY",
"model": "DEEPSEEK_MODEL",
"base_url": "DEEPSEEK_BASE_URL",
"timeout": "DEEPSEEK_TIMEOUT",
"max_retries": "DEEPSEEK_MAX_RETRIES",
"max_tokens": "DEEPSEEK_MAX_TOKENS",
}
_INT_KEYS = {"timeout", "max_retries", "max_tokens"}
def _project_root() -> str:
return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
def _default_config_path() -> str:
return os.path.join(_project_root(), 'config.json')
def load_config(path: str = None) -> dict:
config = copy.deepcopy(DEFAULT_CONFIG)
if path is None:
path = _default_config_path()
if os.path.exists(path):
with open(path, 'r', encoding='utf-8') as f:
user_config = json.load(f)
_deep_merge(config, user_config)
for key, env_name in ENV_OVERRIDES.items():
value = os.environ.get(env_name)
if value is not None:
if key in _INT_KEYS:
value = int(value)
config['api'][key] = value
return config
def _deep_merge(base: dict, override: dict) -> None:
for key, value in override.items():
if isinstance(value, dict) and isinstance(base.get(key), dict):
_deep_merge(base[key], value)
else:
base[key] = value
+10
View File
@@ -0,0 +1,10 @@
{
"api": {
"api_key": "替换为你的 DeepSeek API Key,或用环境变量 DEEPSEEK_API_KEY",
"model": "deepseek-v4-flash",
"base_url": "https://api.deepseek.com/chat/completions",
"timeout": 120,
"max_retries": 3,
"max_tokens": 32768
}
}
+10
View File
@@ -0,0 +1,10 @@
{
"api": {
"api_key": "sk-6156cccdc9c14d949cf5bfc5afc67a03",
"model": "deepseek-v4-flash",
"base_url": "https://api.deepseek.com/chat/completions",
"timeout": 120,
"max_retries": 3,
"max_tokens": 32768
}
}
+7 -4
View File
@@ -12,11 +12,14 @@ import os
import sys import sys
from agent import generate from agent import generate
from agent.config import load_config
DEFAULT_RULES_DIR = os.path.join(os.path.dirname(__file__), 'rules') DEFAULT_RULES_DIR = os.path.join(os.path.dirname(__file__), 'rules')
def main(): def main():
cfg = load_config()['api']
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description='COBOLテストデータ生成Agent' description='COBOLテストデータ生成Agent'
) )
@@ -27,11 +30,11 @@ def main():
parser.add_argument('--cpy', required=True, help='COPYBOOK 格納ディレクトリ') parser.add_argument('--cpy', required=True, help='COPYBOOK 格納ディレクトリ')
parser.add_argument('--db-md', required=True, help='DB 定義書 .md のパス') parser.add_argument('--db-md', required=True, help='DB 定義書 .md のパス')
parser.add_argument('--output', default='output', help='出力ディレクトリ') parser.add_argument('--output', default='output', help='出力ディレクトリ')
parser.add_argument('--api-key', default='sk-6156cccdc9c14d949cf5bfc5afc67a03', parser.add_argument('--api-key', default=cfg['api_key'],
help='DeepSeek API Key') help='DeepSeek API Key(默认读取 config.json 或环境变量 DEEPSEEK_API_KEY')
parser.add_argument('--model', default='deepseek-v4-flash', help='API モデル名') parser.add_argument('--model', default=cfg['model'], help='API モデル名')
parser.add_argument('--rules', default=DEFAULT_RULES_DIR, help='ルール格納ディレクトリ') parser.add_argument('--rules', default=DEFAULT_RULES_DIR, help='ルール格納ディレクトリ')
parser.add_argument('--max-tokens', type=int, default=32768, parser.add_argument('--max-tokens', type=int, default=cfg['max_tokens'],
help='API 生成トークン上限') help='API 生成トークン上限')
args = parser.parse_args() args = parser.parse_args()