import click from pathlib import Path from datetime import datetime from jinja2 import Template from typing import Optional from dataflow_agent.logger import get_logger log = get_logger(__name__) TEMPLATE_DIR = Path(__file__).parent / "templates" # ---------- util ---------- def to_snake(s: str) -> str: import re s = re.sub(r'[\- ]+', '_', s).strip('_') parts = re.split(r'[_]', s) return '_'.join(p.lower() for p in parts if p) def to_camel(s: str) -> str: return ''.join(p.capitalize() for p in to_snake(s).split('_')) # ---------- CLI ---------- @click.group() def cli(): """DataFlow-Agent command line.""" pass @cli.command("create") @click.option("--wf_name", help="要创建的 workflow 名称") @click.option("--agent_name", help="要创建的 agent 名称") @click.option("--gradio_name", help="要创建的 gradio page 名称") @click.option("--prompt_name", help="要创建的 prompt template 名称") @click.option("--agent_as_tool_name", help="要创建的 agent-as-tool 名称") @click.option("--state_name", help="要创建的 state 名称") def create_artifact(wf_name: Optional[str] = None, agent_name: Optional[str] = None, gradio_name: Optional[str] = None, prompt_name: Optional[str] = None, agent_as_tool_name: Optional[str] = None, state_name: Optional[str] = None): """ dfa create --wf_name xxx dfa create --agent_name yyy dfa create --gradio_name zzz dfa create --prompt_name zzz dfa create --agent_as_tool_name aaa dfa create --state_name bbb """ opts = [bool(wf_name), bool(agent_name), bool(gradio_name), bool(prompt_name), bool(agent_as_tool_name), bool(state_name)] if sum(opts) != 1: click.echo(" --wf_name / --agent_name / --gradio_name / --prompt_name / --agent_as_tool_name / --state_name 必须且只能选一个", err=True) raise SystemExit(1) timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") # ------------------------------------------------------------------ # 1. Workflow # ------------------------------------------------------------------ if wf_name: wf_name_snake = to_snake(wf_name) # 1.1 workflow 源码 wf_dest = Path(__file__).parent / "workflow" / f"wf_{wf_name_snake}.py" wf_tpl_path = TEMPLATE_DIR / "workflow.py.jinja" wf_context = dict( wf_name=wf_name, wf_name_snake=wf_name_snake, entry="step1", timestamp=timestamp, ) _generate_file(wf_dest, wf_tpl_path, wf_context, "workflow") # 1.2 对应测试 project_root = Path(__file__).parent.parent test_dest = project_root / "tests" / f"test_{wf_name_snake}.py" test_tpl = TEMPLATE_DIR / "test_workflow.py.jinja" test_ctx = dict( wf_name=wf_name, wf_name_snake=wf_name_snake, timestamp=timestamp, ) _generate_file(test_dest, test_tpl, test_ctx, "test") # ------------------------------------------------------------------ # 2. Agent # ------------------------------------------------------------------ elif agent_name: agent_snake = to_snake(agent_name) dest = Path(__file__).parent / "agentroles" / "common_agents" / f"{agent_snake}_agent.py" tpl_path = TEMPLATE_DIR / "agent.py.jinja" ctx = dict( agent_name=agent_name, agent_name_snake=agent_snake, agent_name_camel=to_camel(agent_name), timestamp=timestamp, ) _generate_file(dest, tpl_path, ctx, "agent") # ------------------------------------------------------------------ # 3. Gradio Page # ------------------------------------------------------------------ elif gradio_name: page_snake = to_snake(gradio_name) project_root = Path(__file__).parent.parent dest = project_root / "gradio_app" / "pages" / f"page_{page_snake}.py" tpl_path = TEMPLATE_DIR / "gradio_page.py.jinja" ctx = dict( page_name=gradio_name, page_name_snake=page_snake, timestamp=timestamp, ) _generate_file(dest, tpl_path, ctx, "gradio page") # ------------------------------------------------------------------ # 4. Prompt Template # ------------------------------------------------------------------ elif prompt_name: prompt_snake = to_snake(prompt_name) dest = Path(__file__).parent / "promptstemplates" / "resources" / f"pt_{prompt_snake}_repo.py" tpl_path = TEMPLATE_DIR / "prompt_repo.py.jinja" ctx = dict( prompt_name=prompt_name, prompt_name_snake=prompt_snake, prompt_name_camel=to_camel(prompt_name), timestamp=timestamp, ) _generate_file(dest, tpl_path, ctx, "prompt template") # ------------------------------------------------------------------ # 5. Agent-as-Tool # ------------------------------------------------------------------ elif agent_as_tool_name: agent_snake = to_snake(agent_as_tool_name) dest = Path(__file__).parent / "agentroles" / "common_agents" / f"{agent_snake}_agent.py" tpl_path = TEMPLATE_DIR / "agent_as_tool_name.py.jinja" ctx = dict( agent_name=agent_as_tool_name, agent_name_snake=agent_snake, agent_name_camel=to_camel(agent_as_tool_name), timestamp=timestamp, ) _generate_file(dest, tpl_path, ctx, "agent-as-tool") # ------------------------------------------------------------------ # 6. State # ------------------------------------------------------------------ else: state_snake = to_snake(state_name) dest = Path(__file__).parent / "states" / f"{state_snake}_state.py" tpl_path = TEMPLATE_DIR / "state_name.py.jinja" ctx = dict( state_name=state_name, state_name_snake=state_snake, state_name_camel=to_camel(state_name), timestamp=timestamp, ) _generate_file(dest, tpl_path, ctx, "state") # ---------- helper ---------- def _generate_file(dest: Path, tpl_path: Path, context: dict, file_type: str): """ 通用文件生成函数 """ dest.parent.mkdir(parents=True, exist_ok=True) if dest.exists(): # click.echo(f" {dest} 已存在，跳过生成") log.error(f" {dest} 已存在，跳过生成") return if not tpl_path.exists(): log.error(f" 模板不存在: {tpl_path}", err=True) raise SystemExit(1) rendered = Template(tpl_path.read_text(encoding="utf-8")).render(**context) dest.write_text(rendered, encoding="utf-8") try: rel_path = dest.relative_to(Path.cwd()) except ValueError: rel_path = dest log.critical(f" 已生成 {file_type}: {rel_path}") if __name__ == "__main__": cli()