| 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" |
|
|
| |
| 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('_')) |
|
|
| |
| @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") |
|
|
| |
| |
| |
| if wf_name: |
| wf_name_snake = to_snake(wf_name) |
|
|
| |
| 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") |
|
|
| |
| 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") |
|
|
| |
| |
| |
| 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") |
|
|
| |
| |
| |
| 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") |
|
|
| |
| |
| |
| 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") |
|
|
| |
| |
| |
| 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") |
|
|
| |
| |
| |
| 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") |
|
|
|
|
| |
| def _generate_file(dest: Path, tpl_path: Path, context: dict, file_type: str): |
| """ |
| 通用文件生成函数 |
| """ |
| dest.parent.mkdir(parents=True, exist_ok=True) |
|
|
| if dest.exists(): |
| |
| 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() |
|
|
