Skip to content

Commit fe66b05

Browse files
chandrasekharan-zipstackclaude
chandrasekharan-zipstack
and
claude
committed
fix(migration): files phase end-to-end + report quieting + tool_instance NOT FOUND
Multiple fixes uncovered by the first local-stack run: client.py - list_prompt_documents: mount under prompt-studio/ prefix (BE include in urls_v2.py). - download_prompt_file: use ?document_id=, matching fetch_contents_ide serializer (was ?file_name=, BE ignored it and returned 400 ValidationError). - upload_prompt_file: drop trailing slash, BE pattern is prompt-studio/file/<uuid:pk> with no slash so POST 404'd. - add get_custom_tool / update_custom_tool for default-doc PATCH. phases/files.py - After upload loop per tool, mirror source's CustomTool.output by filename so FE auto-selects on load. Fall back to first target doc. Preserve any existing target output (operator may have already picked manually on a re-run). phases/tool_instance.py - Detect source serializer sentinels ([X NOT FOUND], [DELETED ADAPTER ...], [NEEDS UPDATE]) in stored metadata and skip the PATCH instead of round-tripping a broken adapter reference. ToolInstance row exists with backend defaults; operator re-binds in UI. report.py - Drop the full source->target UUID map from rendered output (noisy on large migrations). Print per-entity counts only; full map still in as_dict() and at DEBUG log level. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1 parent 05055ce commit fe66b05

5 files changed

Lines changed: 301 additions & 37 deletions

File tree

‎src/unstract/migration/client.py‎

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,22 @@ def list_custom_tools(self) -> list[dict[str, Any]]:
169169
result = self._request("GET", "prompt-studio/")
170170
return result if isinstance(result, list) else result.get("results", [])
171171

172+
def get_custom_tool(self, tool_id: str) -> dict[str, Any]:
173+
"""Fetch a single prompt-studio project (full serializer).
174+
175+
Returns ``fields = "__all__"`` per ``CustomToolSerializer`` —
176+
notably includes ``output`` (the default DocumentManager id the
177+
FE binds to ``selectedDoc`` on load).
178+
"""
179+
return self._request("GET", f"prompt-studio/{tool_id}/")
180+
181+
def update_custom_tool(
182+
self, tool_id: str, body: dict[str, Any]
183+
) -> dict[str, Any]:
184+
"""PATCH a prompt-studio project. Used to set ``output`` (the
185+
default doc id) after the files phase populates DM rows."""
186+
return self._request("PATCH", f"prompt-studio/{tool_id}/", json=body)
187+
172188
def list_profiles(self, tool_id: str) -> list[dict[str, Any]]:
173189
"""List ProfileManager rows for a tool.
174190
@@ -261,24 +277,25 @@ def list_prompt_documents(self, tool_id: str) -> list[dict[str, Any]]:
261277
``to_representation`` filter).
262278
"""
263279
result = self._request(
264-
"GET", "prompt-document/", params={"tool_id": tool_id}
280+
"GET", "prompt-studio/prompt-document/", params={"tool_id": tool_id}
265281
)
266282
return result if isinstance(result, list) else result.get("results", [])
267283

268284
def download_prompt_file(
269-
self, tool_id: str, file_name: str
285+
self, tool_id: str, document_id: str
270286
) -> dict[str, Any]:
271-
"""GET a Prompt Studio document by tool + filename.
287+
"""GET a Prompt Studio document by tool + DM row id.
272288
273-
Returns the backend's ``{"data": ..., "mime_type": ...}`` envelope
274-
verbatim. PDFs come back as base64; text/csv as decoded utf-8;
275-
Excel returns a placeholder string (not real bytes) — callers must
276-
treat unsupported mime types as needing manual re-upload.
289+
``fetch_contents_ide`` resolves the filename internally from the
290+
DocumentManager row, so the SDK passes the ``document_id`` it
291+
already has from ``list_prompt_documents`` rather than reposting
292+
the filename. Returns ``{"data": ..., "mime_type": ...}`` —
293+
PDFs base64, text/csv utf-8, Excel placeholder.
277294
"""
278295
return self._request(
279296
"GET",
280297
f"prompt-studio/file/{tool_id}",
281-
params={"file_name": file_name},
298+
params={"document_id": document_id},
282299
)
283300

284301
def upload_prompt_file(
@@ -297,7 +314,7 @@ def upload_prompt_file(
297314
"""
298315
files = {"file": (file_name, data, mime_type)}
299316
return self._request(
300-
"POST", f"prompt-studio/file/{tool_id}/", files=files
317+
"POST", f"prompt-studio/file/{tool_id}", files=files
301318
)
302319

303320
def export_custom_tool(self, tool_id: str, *, force: bool = True) -> Any:

‎src/unstract/migration/phases/files.py‎

Lines changed: 113 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,8 @@ def _migrate_tool(
114114

115115
for doc in src_docs:
116116
file_name = doc.get("document_name")
117-
if not file_name:
117+
src_document_id = doc.get("document_id")
118+
if not file_name or not src_document_id:
118119
continue
119120
if file_name in target_names:
120121
result.skipped += 1
@@ -131,7 +132,18 @@ def _migrate_tool(
131132
)
132133
continue
133134
self._migrate_one_file(
134-
src_tool_id, tgt_tool_id, tool_name, file_name, report, result
135+
src_tool_id,
136+
tgt_tool_id,
137+
tool_name,
138+
file_name,
139+
src_document_id,
140+
report,
141+
result,
142+
)
143+
144+
if not self.ctx.options.dry_run:
145+
self._ensure_default_doc(
146+
src_tool_id, tgt_tool_id, tool_name, src_docs
135147
)
136148

137149
def _migrate_one_file(
@@ -140,12 +152,15 @@ def _migrate_one_file(
140152
tgt_tool_id: str,
141153
tool_name: str,
142154
file_name: str,
155+
src_document_id: str,
143156
report: MigrationReport,
144157
result: PhaseResult,
145158
) -> None:
146159
try:
147160
payload = self._with_retry(
148-
lambda: self.ctx.source.download_prompt_file(src_tool_id, file_name),
161+
lambda: self.ctx.source.download_prompt_file(
162+
src_tool_id, src_document_id
163+
),
149164
op=f"download {tool_name}/{file_name}",
150165
)
151166
except Exception as e:
@@ -284,6 +299,101 @@ def _decode_payload(
284299
# not real bytes. Round-trip would corrupt the file.
285300
return None
286301

302+
def _ensure_default_doc(
303+
self,
304+
src_tool_id: str,
305+
tgt_tool_id: str,
306+
tool_name: str,
307+
src_docs: list[dict[str, Any]],
308+
) -> None:
309+
"""Set target ``CustomTool.output`` so the FE auto-selects a doc.
310+
311+
Mirror source's chosen doc by filename when possible; fall back
312+
to the first available target doc. Skip if target already has
313+
``output`` set — never override an operator's later choice on
314+
re-runs.
315+
"""
316+
try:
317+
tgt_tool = self.ctx.target.get_custom_tool(tgt_tool_id)
318+
except Exception as e:
319+
logger.warning(
320+
"files: skipping default-doc set for tool=%s — fetch tgt failed: %s",
321+
tool_name, e,
322+
)
323+
return
324+
325+
if tgt_tool.get("output"):
326+
logger.debug(
327+
"files: target tool=%s already has default doc; leaving as-is",
328+
tool_name,
329+
)
330+
return
331+
332+
try:
333+
tgt_docs = self.ctx.target.list_prompt_documents(tgt_tool_id)
334+
except Exception as e:
335+
logger.warning(
336+
"files: skipping default-doc set for tool=%s — list tgt docs failed: %s",
337+
tool_name, e,
338+
)
339+
return
340+
if not tgt_docs:
341+
return
342+
343+
chosen_id = self._pick_default_doc_id(
344+
src_tool_id, src_docs, tgt_docs, tool_name
345+
)
346+
if not chosen_id:
347+
return
348+
349+
try:
350+
self.ctx.target.update_custom_tool(tgt_tool_id, {"output": chosen_id})
351+
logger.info(
352+
"files: set default doc tool=%s doc_id=%s", tool_name, chosen_id
353+
)
354+
except Exception as e:
355+
logger.warning(
356+
"files: PATCH default doc failed tool=%s: %s", tool_name, e
357+
)
358+
359+
def _pick_default_doc_id(
360+
self,
361+
src_tool_id: str,
362+
src_docs: list[dict[str, Any]],
363+
tgt_docs: list[dict[str, Any]],
364+
tool_name: str,
365+
) -> str | None:
366+
# Try mirroring the source's selection by filename. If source
367+
# GET fails or source has no chosen doc, fall back to the first
368+
# target doc so the FE doesn't render an empty selector.
369+
try:
370+
src_tool = self.ctx.source.get_custom_tool(src_tool_id)
371+
src_output = src_tool.get("output")
372+
except Exception as e:
373+
logger.debug(
374+
"files: source CustomTool fetch failed for tool=%s (%s); "
375+
"falling back to first target doc",
376+
tool_name, e,
377+
)
378+
src_output = None
379+
380+
if src_output:
381+
src_name = next(
382+
(d.get("document_name") for d in src_docs
383+
if d.get("document_id") == src_output),
384+
None,
385+
)
386+
if src_name:
387+
matched = next(
388+
(d.get("document_id") for d in tgt_docs
389+
if d.get("document_name") == src_name),
390+
None,
391+
)
392+
if matched:
393+
return matched
394+
395+
return tgt_docs[0].get("document_id")
396+
287397
def _lookup_tool_name(self, tgt_tool_id: str) -> str | None:
288398
# CustomToolPhase doesn't record names; fetch lazily for log clarity.
289399
# One call per tool is cheap relative to the per-file traffic.

‎src/unstract/migration/phases/tool_instance.py‎

Lines changed: 44 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,28 @@
2727

2828
logger = logging.getLogger(__name__)
2929

30+
# Source backend's ToolInstanceSerializer.to_representation emits these
31+
# sentinel strings when an adapter UUID/name in the stored metadata can
32+
# no longer be resolved (deleted or renamed on source). Round-tripping
33+
# them to target produces an AdapterNotFound on PATCH, so we detect and
34+
# skip the metadata PATCH instead — the ToolInstance row exists with the
35+
# backend's safe defaults and the operator can re-bind in the UI.
36+
_BROKEN_ADAPTER_SENTINELS: tuple[str, ...] = (
37+
"NOT FOUND",
38+
"[DELETED ADAPTER",
39+
"[NEEDS UPDATE]",
40+
)
41+
42+
43+
def _broken_adapter_keys(metadata: dict[str, Any]) -> list[str]:
44+
broken: list[str] = []
45+
for key, value in metadata.items():
46+
if isinstance(value, str) and any(
47+
s in value for s in _BROKEN_ADAPTER_SENTINELS
48+
):
49+
broken.append(f"{key}={value!r}")
50+
return broken
51+
3052

3153
class ToolInstancePhase(Phase):
3254
name = "tool_instance"
@@ -120,14 +142,28 @@ def _migrate_workflow_tools(
120142
# PATCH the metadata regardless of created/adopted — keeps tool config
121143
# aligned with source on every run.
122144
src_metadata = src_ti.get("metadata") or {}
123-
try:
124-
self.ctx.target.update_tool_instance_metadata(tgt_ti["id"], src_metadata)
125-
except Exception as e:
126-
logger.exception(
127-
"Failed to PATCH tool_instance %s metadata: %s", tgt_ti["id"], e
145+
broken = _broken_adapter_keys(src_metadata)
146+
if broken:
147+
logger.warning(
148+
"skipping metadata PATCH for tool_instance src=%s tgt=%s — "
149+
"source metadata carries broken adapter refs %s; "
150+
"row exists with backend defaults, re-bind in UI",
151+
src_ti_id, tgt_ti["id"], broken,
128152
)
129-
result.failed += 1
130-
result.errors.append(f"patch metadata {tgt_ti['id']}: {e}")
131-
return
153+
result.errors.append(
154+
f"stale adapter refs on src tool_instance {src_ti_id}: {broken}"
155+
)
156+
else:
157+
try:
158+
self.ctx.target.update_tool_instance_metadata(
159+
tgt_ti["id"], src_metadata
160+
)
161+
except Exception as e:
162+
logger.exception(
163+
"Failed to PATCH tool_instance %s metadata: %s", tgt_ti["id"], e
164+
)
165+
result.failed += 1
166+
result.errors.append(f"patch metadata {tgt_ti['id']}: {e}")
167+
return
132168

133169
self.ctx.remap.record("tool_instance", src_ti_id, tgt_ti["id"])

‎src/unstract/migration/report.py‎

Lines changed: 24 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,12 @@
77

88
from __future__ import annotations
99

10+
import logging
1011
from dataclasses import dataclass, field
1112
from typing import Any
1213

14+
logger = logging.getLogger(__name__)
15+
1316

1417
@dataclass
1518
class PhaseResult:
@@ -65,15 +68,7 @@ def render(self) -> str:
6568
if self.skipped_phases:
6669
console.print(f"[dim]Skipped phases:[/dim] {', '.join(self.skipped_phases)}")
6770
self._render_files_sections(console)
68-
if self.remap_snapshot:
69-
remap = Table(title="Source -> Target UUID Map")
70-
remap.add_column("Entity")
71-
remap.add_column("Source UUID")
72-
remap.add_column("Target UUID")
73-
for entity, mapping in self.remap_snapshot.items():
74-
for src, tgt in mapping.items():
75-
remap.add_row(entity, src, tgt)
76-
console.print(remap)
71+
self._render_remap_summary(console_print=console.print)
7772
if self.aborted:
7873
console.print(f"[red]ABORTED:[/red] {self.abort_reason}")
7974
return buf.getvalue()
@@ -89,13 +84,7 @@ def _render_plain(self) -> str:
8984
if self.skipped_phases:
9085
lines.append(f"Skipped phases: {', '.join(self.skipped_phases)}")
9186
lines.extend(self._files_sections_plain())
92-
if self.remap_snapshot:
93-
lines.append("")
94-
lines.append("Source -> Target UUID Map")
95-
lines.append("-" * 60)
96-
for entity, mapping in self.remap_snapshot.items():
97-
for src, tgt in mapping.items():
98-
lines.append(f" {entity:<12} {src} -> {tgt}")
87+
self._render_remap_summary(console_print=lines.append)
9988
if self.aborted:
10089
lines.append(f"ABORTED: {self.abort_reason}")
10190
return "\n".join(lines)
@@ -124,6 +113,25 @@ def as_dict(self) -> dict[str, Any]:
124113
"failed_files": list(self.failed_files),
125114
}
126115

116+
def _render_remap_summary(self, console_print: Any) -> None:
117+
"""Summarise the remap snapshot. Full map is large and noisy, so
118+
we only print per-entity counts here; the full mapping is emitted
119+
at DEBUG and remains in ``as_dict()`` for programmatic consumers.
120+
"""
121+
if not self.remap_snapshot:
122+
return
123+
counts = ", ".join(
124+
f"{entity}={len(mapping)}"
125+
for entity, mapping in self.remap_snapshot.items()
126+
if mapping
127+
)
128+
if counts:
129+
console_print(f"Remap entries: {counts}")
130+
if logger.isEnabledFor(logging.DEBUG):
131+
for entity, mapping in self.remap_snapshot.items():
132+
for src, tgt in mapping.items():
133+
logger.debug("remap %s %s -> %s", entity, src, tgt)
134+
127135
def _render_files_sections(self, console: Any) -> None:
128136
if self.uploaded_files:
129137
console.print(

0 commit comments

Comments
 (0)