硬编码改为配置文件
This commit is contained in:
@@ -1,3 +1,3 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
output/
|
||||
output/
|
||||
@@ -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)
|
||||
|
||||
@@ -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 生成的结果。"""
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user