| 1 | from __future__ import annotations |
| 2 | |
| 3 | import json |
| 4 | import os |
| 5 | import re |
| 6 | import time |
| 7 | import uuid |
| 8 | from pathlib import Path, PurePosixPath |
| 9 | |
| 10 | from agent import AgentContext |
| 11 | from helpers import files, persist_chat, projects as a0_projects |
| 12 | from helpers.api import ApiHandler, Input, Output, Request, Response |
| 13 | from plugins._migrate_agents.api.migration_preview import uploaded_files |
| 14 | from plugins._migrate_agents.helpers.migration import Asset, Project, build_a0_chat, parse_bundle |
| 15 | |
| 16 | |
| 17 | def _slug(value: str, fallback: str) -> str: |
| 18 | result = re.sub(r"[^a-z0-9_-]+", "_", value.lower()).strip("_") |
| 19 | return (result or fallback)[:80] |
| 20 | |
| 21 | |
| 22 | def _unique_dir(parent: Path, name: str) -> Path: |
| 23 | candidate = parent / name |
| 24 | index = 2 |
| 25 | while candidate.exists(): |
| 26 | candidate = parent / f"{name}_{index}" |
| 27 | index += 1 |
| 28 | return candidate |
| 29 | |
| 30 | |
| 31 | def _atomic_write(path: Path, data: bytes) -> None: |
| 32 | path.parent.mkdir(parents=True, exist_ok=True) |
| 33 | temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") |
| 34 | try: |
| 35 | with temporary.open("wb") as handle: |
| 36 | handle.write(data) |
| 37 | handle.flush() |
| 38 | os.fsync(handle.fileno()) |
| 39 | temporary.replace(path) |
| 40 | finally: |
| 41 | temporary.unlink(missing_ok=True) |
| 42 | |
| 43 | |
| 44 | def _knowledge_path(run_root: Path, asset: Asset) -> Path: |
| 45 | source = PurePosixPath(asset.path) |
| 46 | name = _slug("_".join(source.parts[-3:]), "knowledge") |
| 47 | if not name.endswith(".md"): |
| 48 | name += ".md" |
| 49 | return run_root / name |
| 50 | |
| 51 | |
| 52 | def import_knowledge(source: str, category: str, assets: list[Asset]) -> list[str]: |
| 53 | if not assets: |
| 54 | return [] |
| 55 | base = Path(files.get_abs_path("usr", "knowledge", "_migrate_agents", source, category)) |
| 56 | run = _unique_dir(base, time.strftime("%Y%m%d_%H%M%S")) |
| 57 | written: list[str] = [] |
| 58 | for asset in assets: |
| 59 | destination = _knowledge_path(run, asset) |
| 60 | header = ( |
| 61 | f"# Imported from {source}\n\n" |
| 62 | f"> Original path: `{asset.path}` \n" |
| 63 | "> Imported by Migrate Agents. Review before sharing.\n\n" |
| 64 | ).encode() |
| 65 | _atomic_write(destination, header + asset.data) |
| 66 | written.append(files.deabsolute_path(str(destination))) |
| 67 | return written |
| 68 | |
| 69 | |
| 70 | def import_skills(source: str, skills: dict[str, list[Asset]]) -> list[str]: |
| 71 | base = Path(files.get_abs_path("usr", "skills", "_migrate_agents", source)) |
| 72 | written: list[str] = [] |
| 73 | for name, assets in skills.items(): |
| 74 | destination = _unique_dir(base, _slug(name, "skill")) |
| 75 | for asset in assets: |
| 76 | relative = PurePosixPath(asset.path) |
| 77 | if relative.is_absolute() or ".." in relative.parts: |
| 78 | raise ValueError(f"Unsafe skill path: {asset.path}") |
| 79 | _atomic_write(destination.joinpath(*relative.parts), asset.data) |
| 80 | written.append(files.deabsolute_path(str(destination))) |
| 81 | return written |
| 82 | |
| 83 | |
| 84 | def import_projects(source: str, items: list[Project]) -> tuple[list[str], dict[str, str]]: |
| 85 | parent = Path(a0_projects.get_projects_parent_folder()) |
| 86 | written: list[str] = [] |
| 87 | chat_projects: dict[str, str] = {} |
| 88 | for item in items: |
| 89 | base = _slug(f"{source}_{item.title}", f"{source}_project") |
| 90 | name = base |
| 91 | index = 2 |
| 92 | while (parent / name).exists(): |
| 93 | name = f"{base}_{index}" |
| 94 | index += 1 |
| 95 | a0_projects.create_project( |
| 96 | name, |
| 97 | { |
| 98 | "title": item.title, |
| 99 | "description": f"Imported from {source}. Original workspace: {item.path}", |
| 100 | "instructions": "", |
| 101 | "include_agents_md": True, |
| 102 | "color": "#6366f1", |
| 103 | "git_url": "", |
| 104 | }, |
| 105 | ) |
| 106 | written.append(name) |
| 107 | chat_projects.update({chat_id: name for chat_id in item.conversation_ids}) |
| 108 | return written, chat_projects |
| 109 | |
| 110 | |
| 111 | def import_chats(source: str, conversations, chat_projects: dict[str, str] | None = None) -> list[str]: |
| 112 | payloads = [json.dumps(build_a0_chat(item, source), ensure_ascii=False) for item in conversations] |
| 113 | if not payloads: |
| 114 | return [] |
| 115 | ctxids = persist_chat.load_json_chats(payloads) |
| 116 | if len(ctxids) != len(conversations): |
| 117 | raise RuntimeError("Imported chat count does not match the migration preview") |
| 118 | chat_projects = chat_projects or {} |
| 119 | for ctxid, conversation in zip(ctxids, conversations): |
| 120 | context = AgentContext.get(ctxid) |
| 121 | if context is None: |
| 122 | raise RuntimeError(f"Imported chat was not loaded: {ctxid}") |
| 123 | project_name = chat_projects.get(conversation.source_id) |
| 124 | if project_name: |
| 125 | a0_projects.activate_project(ctxid, project_name) |
| 126 | else: |
| 127 | persist_chat.save_tmp_chat(context) |
| 128 | return ctxids |
| 129 | |
| 130 | |
| 131 | class MigrationImport(ApiHandler): |
| 132 | async def process(self, input: Input, request: Request) -> Output: |
| 133 | try: |
| 134 | source = str(request.form.get("source") or "").strip().lower() |
| 135 | include_chats = request.form.get("include_chats", "true").lower() == "true" |
| 136 | include_projects = request.form.get("include_projects", "true").lower() == "true" |
| 137 | legacy_knowledge = request.form.get("include_knowledge", "true") |
| 138 | include_memories = request.form.get("include_memories", legacy_knowledge).lower() == "true" |
| 139 | include_instructions = request.form.get("include_instructions", legacy_knowledge).lower() == "true" |
| 140 | include_skills = request.form.get("include_skills", "true").lower() == "true" |
| 141 | bundle = parse_bundle(source, uploaded_files(request)) |
| 142 | project_names, chat_projects = import_projects(source, bundle.projects) if include_projects else ([], {}) |
| 143 | ctxids = import_chats(source, bundle.conversations, chat_projects) if include_chats else [] |
| 144 | memories = import_knowledge(source, "memories", bundle.memories) if include_memories else [] |
| 145 | instructions = import_knowledge(source, "instructions", bundle.instructions) if include_instructions else [] |
| 146 | skills = import_skills(source, bundle.skills) if include_skills else [] |
| 147 | return { |
| 148 | "ok": True, |
| 149 | "ctxids": ctxids, |
| 150 | "projects": project_names, |
| 151 | "memories": memories, |
| 152 | "instructions": instructions, |
| 153 | "knowledge": [*memories, *instructions], |
| 154 | "skills": skills, |
| 155 | "summary": { |
| 156 | "chats": len(ctxids), |
| 157 | "projects": len(project_names), |
| 158 | "memories": len(memories), |
| 159 | "instructions": len(instructions), |
| 160 | "knowledge": len(memories) + len(instructions), |
| 161 | "skills": len(skills), |
| 162 | "redactions": bundle.redactions, |
| 163 | }, |
| 164 | "warnings": bundle.warnings, |
| 165 | } |
| 166 | except ValueError as exc: |
| 167 | return Response(str(exc), 400) |