Skip to content

Commit 7d27d4a

Browse files
feat(clone): cloud entity phases (lookups, manual review, agentic studio)
Extend the org clone with cloud-only (enterprise) entity support, gated so OSS runs are unchanged. Foundation: - Capability probe (PlatformClient.probe + ctx.feature_present cache) and a Phase.probe_path gate. The orchestrator probes source/target before a cloud phase: absent on source -> skip silently (OSS looks like today); present on source but absent on target -> warn + skip + continue. report.warnings added. - custom_tool now records a src->tgt prompt-id remap (matched by prompt_key, real + dry-run planned) so prompt-scoped cloud config can rewrite its FKs. Phases (each create-or-adopt-by-name, FK remap, dry-run aware, probe-gated): - lookups: definition + draft template/adapters + reference-file blobs + draft-pinned assignments (prompt + lookup + variable_mappings remap). v1 defers published-version replay (logged, non-fatal). - manual_review: workflow-scoped RuleEngine (+ confidence filters) and HITLSettings rebased onto the cloned workflow, org-level AutoApprovalSettings, ReviewApiKey recreated (secret re-minted -> operator re-wire warning). - agentic_studio: project (4 adapter FKs remapped) + prompt-versions (parent-before-child) + schemas + settings, registry republished via export. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CsGrHbs5SWmQkKqiimg6CF
1 parent cd9bcf5 commit 7d27d4a

15 files changed

Lines changed: 2905 additions & 1 deletion

src/unstract/clone/client.py

Lines changed: 327 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,21 @@ def get_post_schema(self, entity_path: str) -> frozenset[str]:
113113
self._post_schema_cache[entity_path] = writable
114114
return writable
115115

116+
def probe(self, path: str) -> bool:
117+
"""Capability probe: is this feature's route installed on this deployment?
118+
119+
GET ``path`` and return True on 200, False on 404 (route absent =
120+
feature not built into this deployment). Any other status / transport
121+
error re-raises — a real failure must not look like "feature missing".
122+
"""
123+
try:
124+
self._request("GET", path)
125+
except PlatformAPIError as e:
126+
if e.status_code == 404:
127+
return False
128+
raise
129+
return True
130+
116131
# ----- org users & groups -----
117132

118133
def list_users(self) -> list[dict[str, Any]]:
@@ -238,6 +253,18 @@ def list_profiles(self, tool_id: str) -> list[dict[str, Any]]:
238253
result = self._request("GET", f"prompt-studio/prompt-studio-profile/{tool_id}/")
239254
return result if isinstance(result, list) else result.get("results", [])
240255

256+
def list_prompts(self, tool_id: str) -> list[dict[str, Any]]:
257+
"""List a tool's prompts (``prompt_id`` + ``prompt_key`` per row).
258+
259+
Used to map source prompt ids to the target prompts created by
260+
``import_project`` / ``sync_prompts`` (matched by ``prompt_key``),
261+
so prompt-scoped cloud config can remap its FKs.
262+
"""
263+
result = self._request(
264+
"GET", "prompt-studio/prompt/", params={"tool_id": tool_id}
265+
)
266+
return result if isinstance(result, list) else result.get("results", [])
267+
241268
def export_project(self, tool_id: str) -> dict[str, Any]:
242269
"""Export a prompt-studio project as a portable JSON blob.
243270
@@ -533,3 +560,303 @@ def create_api_key(self, payload: dict[str, Any]) -> dict[str, Any]:
533560
and cannot be carried over from source.
534561
"""
535562
return self._request("POST", "api/keys/api/", json=payload)
563+
564+
# ----- lookups (cloud-only) -----
565+
566+
def list_lookup_definitions(self) -> list[dict[str, Any]]:
567+
"""List lookup definitions in this org. Also the capability-probe path."""
568+
result = self._request("GET", "lookups/definitions/")
569+
return result if isinstance(result, list) else (result or {}).get("results", [])
570+
571+
def get_lookup_definition(self, lookup_id: str) -> dict[str, Any]:
572+
"""Fetch a lookup definition's detail.
573+
574+
Detail inlines the draft content: ``prompt_template``,
575+
``draft_version_id``, ``input_vars``, and ``adapters`` (a dict with
576+
``llm`` / ``x2text`` adapter UUIDs, either possibly ``None``).
577+
"""
578+
return self._request("GET", f"lookups/definitions/{lookup_id}/")
579+
580+
def create_lookup_definition(self, payload: dict[str, Any]) -> dict[str, Any]:
581+
"""Create a lookup definition. Backend auto-creates an empty DRAFT
582+
version with default adapters; populate it via the draft/adapters/file
583+
endpoints below.
584+
"""
585+
return self._request("POST", "lookups/definitions/", json=payload)
586+
587+
def update_lookup_draft_template(
588+
self, lookup_id: str, prompt_template: str
589+
) -> dict[str, Any]:
590+
"""Set the draft version's prompt template."""
591+
return self._request(
592+
"PATCH",
593+
f"lookups/definitions/{lookup_id}/draft/",
594+
json={"prompt_template": prompt_template},
595+
)
596+
597+
def update_lookup_draft_adapters(
598+
self, lookup_id: str, adapters: dict[str, str]
599+
) -> dict[str, Any]:
600+
"""Set the draft version's LLM and/or X2Text adapters by target UUID.
601+
602+
``adapters`` may carry either or both of ``llm`` / ``x2text``; absent
603+
keys leave the existing draft adapter untouched.
604+
"""
605+
return self._request(
606+
"PATCH",
607+
f"lookups/definitions/{lookup_id}/adapters/",
608+
json=adapters,
609+
)
610+
611+
def list_lookup_files(self, lookup_id: str) -> list[dict[str, Any]]:
612+
"""List a lookup's draft reference files (rows carry ``file_id``,
613+
``file_name``, ``file_size``).
614+
"""
615+
result = self._request("GET", f"lookups/definitions/{lookup_id}/files/")
616+
return result if isinstance(result, list) else (result or {}).get("results", [])
617+
618+
def download_lookup_file(self, lookup_id: str, file_id: str) -> bytes:
619+
"""Download a reference file's original bytes.
620+
621+
Returns raw bytes — the content route serves an ``HttpResponse`` body
622+
(not a JSON envelope), so this bypasses the JSON-decoding request path.
623+
"""
624+
url = self._url(f"lookups/definitions/{lookup_id}/files/{file_id}/content/")
625+
logger.debug("GET %s", url)
626+
resp = self._session.get(url, timeout=self.timeout, verify=self.verify)
627+
if not 200 <= resp.status_code < 300:
628+
raise PlatformAPIError(
629+
f"GET lookups/definitions/{lookup_id}/files/{file_id}/content/ "
630+
f"returned {resp.status_code}",
631+
status_code=resp.status_code,
632+
body=resp.text[:2000],
633+
)
634+
return resp.content
635+
636+
def upload_lookup_file(
637+
self, lookup_id: str, file_name: str, data: bytes, mime_type: str
638+
) -> dict[str, Any]:
639+
"""Upload a reference file into a lookup's draft version.
640+
641+
Backend writes bytes to storage, creates the row, and dispatches
642+
re-extraction server-side. The draft enforces a unique filename per
643+
version, so callers pre-check via ``list_lookup_files`` to avoid a 409.
644+
"""
645+
files = {"file": (file_name, data, mime_type)}
646+
return self._request(
647+
"POST", f"lookups/definitions/{lookup_id}/files/", files=files
648+
)
649+
650+
def list_lookup_assignments(self) -> list[dict[str, Any]]:
651+
"""List PromptLookupAssignment rows in this org.
652+
653+
Each row carries ``assignment_id``, ``prompt`` (src ToolStudioPrompt
654+
uuid), ``version`` (src LookupVersion uuid), ``lookup_definition``
655+
(src lookup_id), ``is_draft_version``, and ``variable_mappings``.
656+
"""
657+
result = self._request("GET", "lookups/assignments/")
658+
return result if isinstance(result, list) else (result or {}).get("results", [])
659+
660+
def create_lookup_assignment(self, payload: dict[str, Any]) -> dict[str, Any]:
661+
"""Create a prompt-lookup assignment.
662+
663+
Writable: ``prompt``, ``lookup_definition`` (required), ``version``,
664+
``variable_mappings``. Backend enforces one assignment per prompt
665+
(``one_lookup_per_prompt``), so callers pre-check target assignments.
666+
"""
667+
return self._request("POST", "lookups/assignments/", json=payload)
668+
669+
# ----- manual review / HITL (cloud-only) -----
670+
#
671+
# Each workflow can hold one RuleEngine row per ``rule_type`` (DB / API)
672+
# and one HITLSettings row. The "using_workflow" GET routes take the
673+
# workflow id in the URL path and wrap the row in ``{"data": ...}``;
674+
# they 404 (rules) / 500 (settings) when none exists — callers treat a
675+
# missing row as "nothing to clone", not an error.
676+
677+
MR_RULE_TYPES: tuple[str, ...] = ("DB", "API")
678+
679+
def get_review_rule(
680+
self, workflow_id: str, rule_type: str
681+
) -> dict[str, Any] | None:
682+
"""Fetch a workflow's RuleEngine row for one ``rule_type``.
683+
684+
Returns the rule dict (with nested ``confidence_filters``) or ``None``
685+
when no rule of that type exists (backend answers 404).
686+
"""
687+
try:
688+
body = self._request(
689+
"GET",
690+
f"manual_review/rule_engine/workflow/{workflow_id}/",
691+
params={"rule_type": rule_type},
692+
)
693+
except PlatformAPIError as e:
694+
if e.status_code == 404:
695+
return None
696+
raise
697+
return (body or {}).get("data")
698+
699+
def create_review_rule(self, payload: dict[str, Any]) -> dict[str, Any]:
700+
"""Create a RuleEngine row (+ nested ``confidence_filters``).
701+
702+
Writable: ``workflow`` (required), ``rule_type``, ``percentage``,
703+
``rule_string``, ``rule_json``, ``rule_logic``, ``confidence_filters``.
704+
Unique per (workflow, rule_type, organization).
705+
"""
706+
return self._request("POST", "manual_review/rule_engine/", json=payload)
707+
708+
def get_review_settings(self, workflow_id: str) -> dict[str, Any] | None:
709+
"""Fetch a workflow's HITLSettings row, or ``None`` if absent.
710+
711+
The backend's ``settings_using_workflow`` raises ``DoesNotExist``
712+
(→ 500) when no row exists rather than 404, so any error here is
713+
treated as "no settings to clone".
714+
"""
715+
try:
716+
body = self._request(
717+
"GET", f"manual_review/settings/workflow/{workflow_id}/"
718+
)
719+
except PlatformAPIError:
720+
return None
721+
return (body or {}).get("data")
722+
723+
def create_review_settings(self, payload: dict[str, Any]) -> dict[str, Any]:
724+
"""Create a HITLSettings row.
725+
726+
Writable: ``workflow`` (OneToOne, required), ``sync_with``,
727+
``ttl_hours``.
728+
"""
729+
return self._request("POST", "manual_review/settings/", json=payload)
730+
731+
def list_auto_approval_settings(self) -> list[dict[str, Any]]:
732+
"""List org-level AutoApprovalSettings (0 or 1 row per org).
733+
734+
Plain ModelViewSet ``list`` — 200s bare with no query params, so it
735+
doubles as the manual-review capability probe path.
736+
"""
737+
result = self._request("GET", "manual_review/auto_approval_settings/")
738+
return result if isinstance(result, list) else (result or {}).get("results", [])
739+
740+
def create_auto_approval_settings(
741+
self, payload: dict[str, Any]
742+
) -> dict[str, Any]:
743+
"""Create org-level AutoApprovalSettings.
744+
745+
Writable: ``auto_approved_document_classes``, ``auto_approved_users``.
746+
``organization`` is server-set. Unique per organization.
747+
"""
748+
return self._request(
749+
"POST", "manual_review/auto_approval_settings/", json=payload
750+
)
751+
752+
def list_review_api_keys(self) -> list[dict[str, Any]]:
753+
"""List ReviewApiKey rows in this org."""
754+
result = self._request("GET", "manual_review/api/keys/")
755+
return result if isinstance(result, list) else (result or {}).get("results", [])
756+
757+
def create_review_api_key(self, payload: dict[str, Any]) -> dict[str, Any]:
758+
"""Create a ReviewApiKey. The ``api_key`` secret is server-minted
759+
(``uuid4``, non-editable) and cannot be carried over from source.
760+
Writable: ``class_name``, ``description``, ``is_active``.
761+
"""
762+
return self._request("POST", "manual_review/api/key/", json=payload)
763+
764+
# ----- agentic studio (cloud-only) -----
765+
766+
def list_agentic_projects(self) -> list[dict[str, Any]]:
767+
"""List agentic projects in this org. Also the capability-probe path.
768+
769+
Source platform key is a service account, which sees every project.
770+
Rows carry ``id``, ``name``, ``description``, the four adapter FK ids
771+
(``llm_connector_id`` / ``agent_llm_connector_id`` /
772+
``lightweight_llm_connector_id`` / ``text_extractor_connector_id``),
773+
and ``canary_fields``.
774+
"""
775+
result = self._request("GET", "agentic/projects/")
776+
return result if isinstance(result, list) else (result or {}).get("results", [])
777+
778+
def create_agentic_project(self, payload: dict[str, Any]) -> dict[str, Any]:
779+
"""Create an agentic project. Returns the created row (carries ``id``)."""
780+
return self._request("POST", "agentic/projects/", json=payload)
781+
782+
def list_agentic_prompt_versions(
783+
self, *, project_id: str | None = None
784+
) -> list[dict[str, Any]]:
785+
"""List agentic prompt versions, optionally scoped to a project.
786+
787+
Rows carry ``id``, ``project``, ``version``, ``prompt_text``,
788+
``accuracy``, ``is_active``, and the self-FK ``parent_version``.
789+
"""
790+
params: dict[str, Any] = {}
791+
if project_id is not None:
792+
params["project_id"] = project_id
793+
result = self._request("GET", "agentic/prompt-versions/", params=params)
794+
return result if isinstance(result, list) else (result or {}).get("results", [])
795+
796+
def create_agentic_prompt_version(self, payload: dict[str, Any]) -> dict[str, Any]:
797+
"""Create an agentic prompt version (flat endpoint, ``project`` in body)."""
798+
return self._request("POST", "agentic/prompt-versions/", json=payload)
799+
800+
def list_agentic_schemas(
801+
self, *, project_id: str | None = None
802+
) -> list[dict[str, Any]]:
803+
"""List agentic schemas, optionally scoped to a project.
804+
805+
Rows carry ``id``, ``project``, ``json_schema``, ``version``,
806+
``is_active``.
807+
"""
808+
params: dict[str, Any] = {}
809+
if project_id is not None:
810+
params["project_id"] = project_id
811+
result = self._request("GET", "agentic/schemas/", params=params)
812+
return result if isinstance(result, list) else (result or {}).get("results", [])
813+
814+
def create_agentic_schema(self, payload: dict[str, Any]) -> dict[str, Any]:
815+
"""Create an agentic schema (flat endpoint, ``project`` in body)."""
816+
return self._request("POST", "agentic/schemas/", json=payload)
817+
818+
def list_agentic_settings(self) -> list[dict[str, Any]]:
819+
"""List agentic settings. Org-wide key/value rows (no project FK)."""
820+
result = self._request("GET", "agentic/settings/")
821+
return result if isinstance(result, list) else (result or {}).get("results", [])
822+
823+
def create_agentic_setting(self, payload: dict[str, Any]) -> dict[str, Any]:
824+
"""Create an org-wide agentic setting (``key`` is globally unique)."""
825+
return self._request("POST", "agentic/settings/", json=payload)
826+
827+
def update_agentic_setting(
828+
self, setting_id: str, payload: dict[str, Any]
829+
) -> dict[str, Any]:
830+
"""PATCH an existing agentic setting by id."""
831+
return self._request("PATCH", f"agentic/settings/{setting_id}/", json=payload)
832+
833+
def export_agentic_project(self, project_id: str, *, force: bool = True) -> Any:
834+
"""Republish ``AgenticStudioRegistry`` from the project's active
835+
schema + prompt. Mirror of ``export_custom_tool``.
836+
837+
Backend requires an active schema and active prompt; ``force_export``
838+
bypasses the wizard-completion guard. Records nothing — caller re-reads
839+
the registry to learn the new id.
840+
"""
841+
return self._request(
842+
"POST",
843+
f"agentic/projects/{project_id}/export/",
844+
json={
845+
"is_shared_with_org": False,
846+
"user_ids": [],
847+
"force_export": force,
848+
},
849+
)
850+
851+
def list_agentic_registries(
852+
self, *, agentic_project: str | None = None
853+
) -> list[dict[str, Any]]:
854+
"""List AgenticStudioRegistry rows. The list endpoint returns nothing
855+
unless a filter is supplied; pass ``agentic_project`` to look up the
856+
registry id for a given project.
857+
"""
858+
params: dict[str, Any] = {}
859+
if agentic_project is not None:
860+
params["agentic_project"] = agentic_project
861+
result = self._request("GET", "agentic-studio-registry/", params=params)
862+
return result if isinstance(result, list) else (result or {}).get("results", [])

src/unstract/clone/context.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,3 +145,20 @@ class CloneContext:
145145
# touches them once per endpoint, never per resource).
146146
share_cache: dict[str, Any] = field(default_factory=dict)
147147
share_cache_lock: threading.Lock = field(default_factory=threading.Lock)
148+
# Capability-probe memo: (id(client), feature_path) -> present?. Probed
149+
# once per (deployment, feature) so cloud-phase gating costs one GET total.
150+
probe_cache: dict[tuple[int, str], bool] = field(default_factory=dict)
151+
152+
def feature_present(self, client: "PlatformClient", path: str) -> bool:
153+
"""Is ``path`` (a feature's list endpoint) installed on ``client``?
154+
155+
Memoised per run. # ponytail: plain dict, no lock — probing runs in
156+
the single-threaded orchestrator loop, before any parallel_map fan-out.
157+
"""
158+
key = (id(client), path)
159+
cached = self.probe_cache.get(key)
160+
if cached is not None:
161+
return cached
162+
present = client.probe(path)
163+
self.probe_cache[key] = present
164+
return present

0 commit comments

Comments
 (0)