diff --git a/black-box-data-create/.gitignore b/black-box-data-create/.gitignore index f849cca..5a7222a 100644 --- a/black-box-data-create/.gitignore +++ b/black-box-data-create/.gitignore @@ -1,3 +1,3 @@ __pycache__/ *.pyc -output/ +output/ \ No newline at end of file diff --git a/black-box-data-create/agent/__init__.py b/black-box-data-create/agent/__init__.py index 25c2104..02dab42 100644 --- a/black-box-data-create/agent/__init__.py +++ b/black-box-data-create/agent/__init__.py @@ -6,17 +6,25 @@ from agent.prompt_builder import PromptBuilder from agent.api_client import APIClient from agent.output_writer import OutputWriter 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, cpy_dir: str, db_md: str, output_dir: str = "output", - api_key: str = "sk-6156cccdc9c14d949cf5bfc5afc67a03", - api_model: str = "deepseek-v4-flash", + api_key: str = None, api_model: str = None, rules_dir: str = "rules", - max_tokens: int = 32768) -> dict: + max_tokens: int = None) -> dict: """生成测试数据的主入口函数。""" 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}") parser = InputParser(design_md, source_cbl, file_db_md, cpy_dir, db_md) diff --git a/black-box-data-create/agent/api_client.py b/black-box-data-create/agent/api_client.py index 7b9697b..86b15a9 100644 --- a/black-box-data-create/agent/api_client.py +++ b/black-box-data-create/agent/api_client.py @@ -4,20 +4,22 @@ from typing import Dict, Any, Optional import requests +from agent.config import load_config + class APIClient: """DeepSeek API 客户端，含重试逻辑。""" - def __init__(self, api_key: str, model: str = 'deepseek-v4-flash', - base_url: str = 'https://api.deepseek.com/chat/completions', - max_retries: int = 3, timeout: int = 120, - max_tokens: int = 32768): - self.api_key = api_key - self.model = model - self.base_url = base_url - self.max_retries = max_retries - self.timeout = timeout - self.max_tokens = max_tokens + def __init__(self, api_key: str = None, model: str = None, + base_url: str = None, max_retries: int = None, + timeout: int = None, max_tokens: int = None): + cfg = load_config()['api'] + self.api_key = api_key if api_key else cfg['api_key'] + self.model = model if model else cfg['model'] + self.base_url = base_url if base_url else cfg['base_url'] + self.max_retries = max_retries if max_retries is not None else cfg['max_retries'] + self.timeout = timeout if timeout is not None else cfg['timeout'] + self.max_tokens = max_tokens if max_tokens is not None else cfg['max_tokens'] def generate(self, prompt: str) -> Dict[str, Any]: """发送 prompt 并返回 AI 生成的结果。""" diff --git a/black-box-data-create/agent/config.py b/black-box-data-create/agent/config.py new file mode 100644 index 0000000..b1d26fe --- /dev/null +++ b/black-box-data-create/agent/config.py @@ -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 diff --git a/black-box-data-create/config.example.json b/black-box-data-create/config.example.json new file mode 100644 index 0000000..926f405 --- /dev/null +++ b/black-box-data-create/config.example.json @@ -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 + } +} diff --git a/black-box-data-create/config.json b/black-box-data-create/config.json new file mode 100644 index 0000000..71ac6c2 --- /dev/null +++ b/black-box-data-create/config.json @@ -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 + } +} diff --git a/black-box-data-create/main.py b/black-box-data-create/main.py index 17bacab..8c4d444 100644 --- a/black-box-data-create/main.py +++ b/black-box-data-create/main.py @@ -12,11 +12,14 @@ import os import sys from agent import generate +from agent.config import load_config DEFAULT_RULES_DIR = os.path.join(os.path.dirname(__file__), 'rules') def main(): + cfg = load_config()['api'] + parser = argparse.ArgumentParser( description='COBOLテストデータ生成Agent' ) @@ -27,11 +30,11 @@ def main(): parser.add_argument('--cpy', required=True, help='COPYBOOK 格納ディレクトリ') parser.add_argument('--db-md', required=True, help='DB 定義書 .md のパス') parser.add_argument('--output', default='output', help='出力ディレクトリ') - parser.add_argument('--api-key', default='sk-6156cccdc9c14d949cf5bfc5afc67a03', - help='DeepSeek API Key') - parser.add_argument('--model', default='deepseek-v4-flash', help='API モデル名') + parser.add_argument('--api-key', default=cfg['api_key'], + help='DeepSeek API Key（默认读取 config.json 或环境变量 DEEPSEEK_API_KEY）') + parser.add_argument('--model', default=cfg['model'], help='API モデル名') 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 生成トークン上限') args = parser.parse_args()