硬编码改为配置文件
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user