UN-2946 [FEAT] Prompt Studio lookups bridge, executor hook, and IDE wiring (OSS side)#1929
Open
chandrasekharan-zipstack wants to merge 45 commits intomainfrom
Open
UN-2946 [FEAT] Prompt Studio lookups bridge, executor hook, and IDE wiring (OSS side)#1929chandrasekharan-zipstack wants to merge 45 commits intomainfrom
chandrasekharan-zipstack wants to merge 45 commits intomainfrom
Conversation
…prompt list endpoint
- Add CustomToolListSerializer for the list action to avoid N+1 queries
(profile lookups, prompt fetching, coverage calculation per tool)
- Add ToolStudioPromptListSerializer with only prompt_id, prompt_key,
enforce_type, sequence_number
- Add GET /prompt-studio/prompt/?tool_id={uuid} list endpoint
- List action uses select_related and Subquery annotation for prompt_count
- Detail endpoint unchanged (still uses full CustomToolSerializer)
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Add lookup-studio plugin detection with dynamic import - Add PromptStudioPopoverContent for hover submenu (Projects / Look-Ups) following the same Popover pattern as HITL and Platform Settings - Register lookups/* route in useMainAppRoutes.js Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…ofile_manager Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…on enrichment - Add lookup_config export in prompt_studio_helper and registry_helper via cloud plugin guard (try/except ImportError) - Store raw output in PromptStudioOutputManager, enriched in cloud LookupOutputResult — preserving both for UI tab display - Add LookupEnrichmentProtocol and plugin call in post-extraction pipeline using ExecutorPluginLoader (no-op in OSS) - Track lookup LLM usage via standard metrics pipeline (usage_kwargs with run_id/execution_id, capture_metrics) - Move webhook postprocessing from answer_prompt to pipeline - Frontend: dynamic plugin imports for LookupMenuItem, LookupIndicator, LookupOutputTabs in prompt cards; fetch lookup outputs on page load - Add scroll-to-prompt support via query param in DocumentParser Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…age reason - Extract get_lookup_config() to prompt_studio/lookup_utils.py, replacing 4 identical try/except ImportError blocks across prompt_studio_helper and prompt_studio_registry_helper - Add LOOKUP to LLMUsageReason choices (was missing, causing invalid choice on usage records from lookup enrichment LLM calls) - Migration: usage_v2/0004_add_lookup_usage_reason Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Store prompt/completion/total token counts from the most recent complete() call on the LLM object itself, making usage data queryable without relying on the Audit pipeline roundtrip. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…hook methods Move lookup result-application logic to the cloud plugin, matching the challenge plugin pattern where the plugin owns metadata mutation. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…port fallback Add reusable extraction_complete/extraction_error callback tasks to the ide_callback worker, replacing the need for Django-based celery workers for text extraction. Add ExtractionAPIClient for internal API calls. Add polling fallback to WebSocket transport for local dev reliability. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…ntrusive UX" This reverts commit d6e136d.
Replace inline DRAFT lookup check with pluggable cloud-only hook. Uses try/except ImportError pattern — zero lookup code in OSS. Collects all DRAFT lookups in one pass with markdown-linked error messages. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Consolidate cloud imports into lookup_utils.py with persist_lookup_output() and validate_lookups_for_export() wrappers - Fix LookupEnrichmentProtocol.run() return type to None matching challenge/evaluation pattern - Revert logger.info to logger.debug in websocket_views.py - Eliminate duplicated LookupOutputTabs ternary with renderWithLookupWrapper helper - Move lookups menu constants from SideNavBar.jsx to cloud plugin - Harden DocumentParser.jsx scrollTo with UUID validation and fix useEffect dependency - Revert SocketContext transport to ["websocket"] Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Move Prompt Studio / Look-Ups navigation from a hover popover on the sidebar into a Segmented control within the ToolNavBar. CustomTools dynamically imports LookupList from the plugin and renders tabs when available, falling back to projects-only view in OSS mode. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Switch from eager per-call Audit HTTP push to a deferred batch write pattern for adapter usage. LLM/embedding calls stash records in-memory; the executor flushes them into ExecutionResult metadata; the Celery task batch-writes via a new internal endpoint. Adds 5 nullable columns to Usage (reference_id, reference_type, execution_time_ms, status, error_message) and a composite index for lookup dashboard queries. Extensible choice lists allow cloud plugins to register additional usage reasons and reference types. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…izer Bridge function in lookup_utils.py lets cloud plugins enrich PromptStudioOutputSerializer with lookup data (enriched output, lookup name). Enables real-time lookup results via WebSocket without page refresh. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
… metadata passthrough - Wire usage_kwargs_extra from lookup config into LLM usage_kwargs for execution observability - Add error handling around enricher.run() with explicit ERROR usage records - Generic passthrough of _usage_kwargs into usage records for arbitrary metadata (e.g. reference_id) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…hooks Add dynamic import of getEnrichedCopyText so the copy button copies enriched lookup output when the Enriched tab is active. Applied to both single-pass and multi-profile output paths. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…modified_at fix - Add /prompt-studio/<pk>/lookup-validation/ endpoint backing the FE Export/Deploy gate; multi-var block check accepts prompt_ids so a single prompt run isn't blocked by an unrelated multi-var lookup. - Add /prompt-output/latest-by-keys/ endpoint that returns the most recent raw output per prompt_key for the test panel's "Use Latest Outputs" helper. - Fix prompt output modified_at not refreshing on re-runs (QuerySet.update bypasses auto_now); set timezone.now() explicitly in the update args. - lookup_utils: bridge get_lookup_validation_for_tool and get_multi_var_lookups_for_tool with prompt_ids scoping. - Header wires useLookupExportGate via try-import (no-op stub in OSS). - TokenUsage treats all-null Usage rows as empty. - CombinedOutput / JsonView build enriched dict from metadata.lookup_outputs to back the Raw|Enriched output toggle. - .gitignore: widen docker/compose.*.yaml. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ab on back - ProfileInfoBar: swap Row/Col for plain flex-wrap div — kills Ant Row negative-margin quirk that overlapped wrapped pills in combined output. - CustomTools: honor location.state.activeTab so back navigation from lookup detail lands on the Look-Ups tab instead of defaulting to Projects. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Introduces nullable last_exported_at on CustomTool (populated on first successful export) so staleness checks can compare against downstream mutations without a data backfill. NULL is treated as "unknown" and suppresses the lookup-dirty flag to avoid false alarms on pre-feature projects. Adds the get_latest_lookup_mutation_for_tool bridge in lookup_utils so OSS stays decoupled from the cloud plugin. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
When enricher.run() raises, surface a user-visible ERROR log line in the workflow execution log alongside the existing usage record. Keeps lookup failures observable next to the other pre/post lookup lines we already emit. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Loads useLookupDirtySeed (server-side is_lookup_dirty) and useLookupExportGate from the cloud plugin via dynamic imports so the reminder banner reflects lookup changes across page reloads and the banner's Export flow goes through the same validation modal as the main buttons. Also adds a titleAdornment slot on ToolNavBar for rendering the onboarding tooltip and relaxes EmptyState.text to accept nodes for the tagline + link composition. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… state - Delegate production lookup enrichment to LookupEnrichment.run_with_metrics so the executor and the IDE test path share LLM construction, error handling, and usage-record emission. - Let ExecutionLogs callers pass an arbitrary backRouteState via location state so nested UI restore (e.g. a sub-tab) no longer needs special casing in this component. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Bucket of hardening fixes driven by a staff-level PR re-review: - Org-scope latest_outputs_by_keys (was cross-tenant readable via raw .objects.filter() that bypassed OrganizationFilterBackend). - Hide lookup payload shape from OSS: three new opaque bridge helpers (get_original_value_if_enriched, attach_combined_output_enrichment, extract_prompt_output_enrichment) replace direct reads of metadata["lookup_outputs"] / _lookup_outputs / lookup_outputs in output_manager_helper, CombinedOutput.jsx, and usePromptOutput.js. - Split usage_v2 index into a new 0005 migration that uses AddIndexConcurrently + atomic=False so prod doesn't lock the billing table during build. - Delete stale workers/tests/test_usage.py that imported the removed UsageHelper module. - SDK1 LLM gains public get_last_usage_record() so downstream code stops reaching into _pending_usage across plugin boundaries. - legacy_executor stamps metadata["lookup_errors"][prompt] on a failed lookup outcome for dashboards that surface partial-failure runs. - extraction_client docstring notes the cloud-only endpoint contract. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- Module-level probe in prompt_studio/lookup_utils.py — swap per-function try/except ImportError for a single LOOKUPS_AVAILABLE flag. Add attach_lookup_config / attach_lookup_configs_to_tool_settings helpers so the direct metadata["lookup_errors"] write and the lookup_config key stamping both route through the bridge. - Reject org=None in UsageBatchCreateView (usage_v2/internal_views.py). - Lift useLookupExportGate to a single mount in ToolIde.jsx; thread checkLookups down into custom-tools/header/Header.jsx (eliminates the double modal-portal risk). - Delete the direct metadata["lookup_errors"] write from workers/executor/executors/legacy_executor.py — flat summary is now stamped by LookupEnrichment.write_lookup_error in the cloud plugin. Replace hardcoded "lookup_llm" metrics key with lookup_cls.METRICS_KEY. - Trim boilerplate comments across CombinedOutput, PromptCardItems, PromptOutput, usePromptOutput, prompt-card/Header, CustomToolsHelper, SideNavBar — keep the why-comments, drop the what-comments.
OSS counterpart to the cloud-side data-model change. Wires the prompt studio runtime to the new wire shape, surfaces lookup runnability state in the prompt card, and adds the usage_v2 enum entries the cloud side records against (lookup as an LLM usage reason, lookup_version as a reference type). Partially working — known follow-up: TODO: rework lookup input UX. The current variable-mapping flow is awkward (separate rows, manual prompt selection per variable); needs a redesign that mirrors how users actually compose a lookup template. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Bring in: - modified_at auto-bump fix (BaseModelQuerySet/BaseModelManager + save() override) - legacy_executor: tool-context warning, backtick log formatting, organization IDs threaded into ExecutorToolShim - prompt-service log streaming + markdown rendering - minio bucket-listing scope fix and CustomMarkdown URL-safety helper Conflict resolutions: - workers/executor/executors/legacy_executor.py: combine HEAD's _usage_records field and typed annotations with main's _execution_id/_organization_id context fields and tool-context warning. Keep both the usage flush after challenge and main's backticked stream_log message. Drop redundant explicit modified_at workarounds now that BaseModelQuerySet auto-bumps modified_at on QuerySet.update(): - prompt_studio_output_manager_v2/output_manager_helper.py: remove the "modified_at": timezone.now() entry passed to PromptStudioOutputManager .objects.filter(...).update(**args), and the now-unused timezone import.
When a configured lookup runs but extraction returned None for the source prompt, _run_lookup_enrichment used to fall through silently — leaving users wondering why enrichment didn't appear. Stream a one-line workflow log via shim.stream_log so the skip is visible alongside other tool-run events.
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds lookup/enrichment plumbing with OSS-safe fallbacks, lookup validation and export staleness tracking, buffered/batched usage reporting with schema migrations and internal bulk API, extraction completion/error callbacks and client, worker-side usage buffering/flush, and multiple frontend integrations for lookup UI and enriched output. Changes
Sequence Diagram(s)sequenceDiagram
participant Executor as Worker Executor
participant Orchestrator as Orchestrator/Task
participant WorkerTask as Celery Task
participant UsageClient as UsageAPIClient
participant Backend as Backend API
participant DB as Database
Executor->>Orchestrator: run execution, collect _usage_records
Orchestrator->>WorkerTask: return ExecutionResult (metadata includes usage_records)
WorkerTask->>UsageClient: bulk_create_usage(usage_records, organization_id)
UsageClient->>Backend: POST /v1/usage/batch/ (records)
Backend->>DB: INSERT usage rows
DB-->>Backend: OK
Backend-->>UsageClient: 200 { created: N } / success
UsageClient-->>WorkerTask: success/failure
WorkerTask-->>Orchestrator: log result / continue
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
# Conflicts: # backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py # frontend/src/components/custom-tools/prompt-card/PromptCardItems.jsx
athul-rs
reviewed
Apr 28, 2026
Contributor
athul-rs
left a comment
athul-rs
left a comment
There was a problem hiding this comment.
Code review
Found 4 issues:
APIException(detail=..., code=400)returns HTTP 500, not 400 (DRF'sAPIException.status_codedefaults to 500; thecode=kwarg sets the JSONdetail.codestring, not the HTTP status). Use aValidationErroror a custom exception class withstatus_code = 400._run_webhook_postprocessingis called unconditionally after_apply_type_conversion, so when JSON parsing fails insidehandle_json(which setsstructured_output[prompt] = {}and returns), the webhook now fires withparsed_data={}. Previously the webhook only ran on successful parses — this is a behavioral regression for any deployment usingpostprocessing_webhook_url.LookupEnrichmentProtocoldeclares onlydef run(self) -> None, but the call site (legacy_executor.py:2009) invokeslookup_cls.run_with_metrics(...)with 8 keyword args.@runtime_checkablewould accept any class withrun(), masking a missingrun_with_metrics.- Migration 0004 hardcodes
("lookup", "Lookup")into the OSS migration state forUsage.llm_usage_reason, butmodels.pyonly adds"lookup"at runtime when the cloud plugin is importable. In OSS-only environments the model and migration choices diverge, somanage.py makemigrations --checkwill exit non-zero.
Line-anchored notes inline.
🤖 Generated with Claude Code
- If this code review was useful, please react with 👍. Otherwise, react with 👎.
When a prompt has a lookup configured, switching its enforce_type to a complex type (table / line-item / agentic_table) is blocked: the cloud plugin exposes useEnforceTypeSwitchGate, which OSS dynamically imports with a no-op fallback. Blocked switches surface an alert and keep the previous enforce_type.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (2)
backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py (2)
216-225:⚠️ Potential issue | 🟠 MajorNarrow lookup persistence error handling; generic catch masks enrichment write failures.
Line 220 catches all exceptions and only logs, so lookup persistence failures are hidden and the request path still looks successful. Catch expected DB/validation failures explicitly and re-raise unknown errors.
#!/bin/bash # Verify the broad catch and whether any failure signal is propagated. rg -n -C4 'persist_lookup_output|except Exception|logger.warning' backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py rg -n -C3 'lookup_persist_warning|persist_warning|enrichment_warning' backend/prompt_studio🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py` around lines 216 - 225, The current broad except Exception around persist_lookup_output hides unexpected failures; update the block in output_manager_helper.py that checks prompt_lookup and calls persist_lookup_output to catch only expected errors (e.g., database/validation exceptions your codebase uses such as ValueError, ValidationError, or your DB exception class) and log them with logger.warning including exc_info, but re-raise any other unknown exceptions so the request path surfaces failures; specifically modify the try/except that references persist_lookup_output, prompt_lookup, prompt_output and prompt.prompt_key to only handle known exception types and use raise for all others after logging.
281-299:⚠️ Potential issue | 🟠 MajorDefault output collector is nondeterministic and does extra DB/plugin work.
Line 287 (
exists()) + Line 291 (full iteration) adds extra query/work, and Line 292 overwritesvalueper row without ordering, so returned output is unstable. Also, theObjectDoesNotExistbranch at Line 298 is not a realistic path for this queryset flow.Proposed refactor
def _collect_default_output_for_prompt( tool_prompt: ToolStudioPrompt, profile_manager_id: str, document_manager_id: str, enrichment_by_key: dict[str, Any], ) -> Any: from prompt_studio.lookup_utils import enrich_prompt_output - try: - queryset = PromptStudioOutputManager.objects.filter( - prompt_id=str(tool_prompt.prompt_id), - profile_manager=profile_manager_id, - is_single_pass_extract=False, - document_manager_id=document_manager_id, - ) - if not queryset.exists(): - return "" - - value: Any = "" - for output in queryset: - value = output.output - enriched = enrich_prompt_output(output, {}) - bundle = extract_prompt_output_enrichment(enriched) - if bundle is not None: - enrichment_by_key[tool_prompt.prompt_key] = bundle - return value - except ObjectDoesNotExist: - return "" + output = ( + PromptStudioOutputManager.objects.filter( + prompt_id=str(tool_prompt.prompt_id), + profile_manager=profile_manager_id, + is_single_pass_extract=False, + document_manager_id=document_manager_id, + ) + .order_by("-modified_at") + .first() + ) + if output is None: + return "" + + enriched = enrich_prompt_output(output, {}) + bundle = extract_prompt_output_enrichment(enriched) + if bundle is not None: + enrichment_by_key[tool_prompt.prompt_key] = bundle + return output.output#!/bin/bash # Verify query shape + unreachable-except pattern in the collector. rg -n -C5 '_collect_default_output_for_prompt|exists\(|for output in queryset|except ObjectDoesNotExist|order_by' backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py` around lines 281 - 299, The default collector (_collect_default_output_for_prompt) currently does an exists() then iterates the entire queryset and overwrites value per row, making results nondeterministic and doing extra DB/plugin work; change this to fetch a single deterministic row (e.g., use queryset.order_by("-created_at") or .order_by("id").first() or .values("output").first()) and only call enrich_prompt_output/extract_prompt_output_enrichment for that one record to populate enrichment_by_key[tool_prompt.prompt_key]; remove the try/except ObjectDoesNotExist branch since .filter().first() will return None rather than raising and handle the None case by returning "".
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@backend/prompt_studio/prompt_studio_core_v2/views.py`:
- Around line 514-517: The lookup gating (_multi_var_lookup_block_response) is
being invoked before validating the required input prompt_id, causing callers
without prompt_id to get a lookup-related 400 instead of a clear validation
error; move the prompt_id presence/validation check so it runs before calling
_multi_var_lookup_block_response (update both occurrences that reference
prompt_id gating), and ensure any early return returns the appropriate
validation error when prompt_id is missing rather than executing the lookup
gate.
In `@frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx`:
- Around line 247-259: The copy action currently calls getEnrichedCopyText
directly so if getEnrichedCopyText throws the enriched text path aborts and the
fallback (displayPromptResult) never runs; create a small safe resolver function
(e.g., resolveSafeCopyText) that takes promptOutputId, promptOutput,
promptDetails and internally calls getEnrichedCopyText inside a try/catch,
returning the enriched text on success or the fallback displayPromptResult(...)
on error/undefined, and then replace both copyToClipboard call sites (the inline
handler using getEnrichedCopyText and the other similar usage) to call
resolveSafeCopyText and pass its result into copyOutputToClipboard to ensure
copy never fails.
- Around line 54-68: The empty catch blocks around the dynamic imports for
LookupOutputTabs and getEnrichedCopyText swallow real errors; update the two
try/catch blocks that import
"../../../plugins/lookup-studio/prompt-card/LookupOutputTabs" and
"../../../plugins/lookup-studio/prompt-card/getEnrichedCopyText" so they
preserve the OSS no-op behavior but log unexpected failures (e.g., catch (err) {
console.error("Failed to load LookupOutputTabs:", err) } or use your app logger)
to surface chunk/load errors while still leaving LookupOutputTabs and
getEnrichedCopyText undefined when the plugin is absent.
---
Duplicate comments:
In
`@backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py`:
- Around line 216-225: The current broad except Exception around
persist_lookup_output hides unexpected failures; update the block in
output_manager_helper.py that checks prompt_lookup and calls
persist_lookup_output to catch only expected errors (e.g., database/validation
exceptions your codebase uses such as ValueError, ValidationError, or your DB
exception class) and log them with logger.warning including exc_info, but
re-raise any other unknown exceptions so the request path surfaces failures;
specifically modify the try/except that references persist_lookup_output,
prompt_lookup, prompt_output and prompt.prompt_key to only handle known
exception types and use raise for all others after logging.
- Around line 281-299: The default collector
(_collect_default_output_for_prompt) currently does an exists() then iterates
the entire queryset and overwrites value per row, making results
nondeterministic and doing extra DB/plugin work; change this to fetch a single
deterministic row (e.g., use queryset.order_by("-created_at") or
.order_by("id").first() or .values("output").first()) and only call
enrich_prompt_output/extract_prompt_output_enrichment for that one record to
populate enrichment_by_key[tool_prompt.prompt_key]; remove the try/except
ObjectDoesNotExist branch since .filter().first() will return None rather than
raising and handle the None case by returning "".
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: d7d4bdf8-3be2-42c9-afb3-d9ecfd54448d
📒 Files selected for processing (11)
backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.pybackend/prompt_studio/prompt_studio_core_v2/views.pybackend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.pybackend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.pyfrontend/src/components/custom-tools/prompt-card/Header.jsxfrontend/src/components/custom-tools/prompt-card/PromptCard.jsxfrontend/src/components/custom-tools/prompt-card/PromptCardItems.jsxfrontend/src/components/custom-tools/prompt-card/PromptOutput.jsxunstract/sdk1/src/unstract/sdk1/llm.pyworkers/executor/executors/legacy_executor.pyworkers/ide_callback/tasks.py
✅ Files skipped from review due to trivial changes (1)
- backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py
🚧 Files skipped from review as they are similar to previous changes (3)
- frontend/src/components/custom-tools/prompt-card/PromptCardItems.jsx
- backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py
- frontend/src/components/custom-tools/prompt-card/Header.jsx
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
- Guard embedding flush on callback_manager is None so public-adapter
embedding usage stops being silently swallowed by the broad except.
- Initialise LegacyExecutor._usage_records in __init__ so an
early-execute crash can't leave the orchestrator's getattr fallback
reading None.
- Wrap challenge and summarisation completion calls in try/finally so
flush_pending_usage() runs even on exception — transient errors
no longer drop those LLMs' billing rows.
- Replace bare except: pass on the FE error stream with
logger.debug; the secondary failure is now recoverable.
- Tighten UsageAPIClient.bulk_create_usage success heuristic against
future partial-body contracts; capture and log the bool return at the
Celery task with run_id + organization_id.
- Drop redundant organization_id kwarg passed alongside
set_organization_context in tasks.py.
- SDK records: spread _usage_kwargs first so explicit billing fields
win; write run_id/llm_usage_reason as None instead of "" so
UUIDField/choice columns don't reject the row.
- Use rsplit('/', 1)[-1] for display_model so multi-segment IDs
(e.g. bedrock/anthropic/claude) collapse to the trailing segment.
- Log the litellm cost_per_token failure for embeddings (matches the
LLM path).
- Tighten LookupEnrichmentProtocol to declare run_with_metrics and METRICS_KEY so it actually matches the plugin call site — run() -> None was structurally identical to the no-op protocols. - Wrap the cloud run_with_metrics call in _run_lookup_enrichment in a defensive try/except that logs + streams a WARN to the IDE log, so plugin contract drift degrades the lookup gracefully instead of aborting the answer-prompt batch. - Hoist llm_cls out of the per-prompt hot loop — it was being unpacked from a 7-tuple inside _run_lookup_enrichment on every prompt; the caller already has it. - Extract the is_empty ladder into a module-level _is_blank helper so the predicate (and the falsy-but-valid 0/False rationale comment) lives next to the predicate, not the executor. - _run_webhook_postprocessing: when webhook_enabled and output is non-JSON, log + shim.stream_log so the user sees the skip in the IDE panel instead of silently never receiving the webhook call. - Narrow persist_lookup_output catch to (IntegrityError, ValidationError) and promote to logger.error — broad catches were hiding plugin schema drift while reporting success. - Wrap enrich_prompt_output in the prompt-output serializer's to_representation so an enrichment exception no longer 500s the list endpoint; matches the surrounding log-and-continue policy.
- UsageStatus(TextChoices) so the status field has an enforced domain instead of free-text — producers (llm.py, usage_handler.py) already write the canonical "SUCCESS" string so no producer changes needed. - Add usage_reference_pair_consistent CheckConstraint so a row with reference_id set but reference_type NULL (or vice versa) is rejected at the DB. Cheap to validate on apply because both fields landed together in lookups V2 — legacy rows have both NULL. The sibling (usage_type, llm_usage_reason) constraint is deferred to a follow-up issue: legacy embedding rows have llm_usage_reason='' from the old SDK default, and a full-table backfill or default ADD CONSTRAINT scan would lock the billing table for too long in prod. - internal_views.py: write llm_usage_reason as None instead of "" when missing so the choice column doesn't store an out-of-domain value. - lookup_utils.py: narrow the ImportError catch to the four cloud lookup modules so a failing transitive import inside the cloud plugin re-raises instead of silently degrading the whole feature to a no-op. Annotate get_original_value_if_enriched return as tuple[Any, dict] | None and rephrase the docstring to match the actual shape callers tuple-unpack. - Drop the attach_lookup_config / attach_lookup_configs_to_tool_settings / get_lookup_config_from_output wrappers — they were trivial dict ops over a hardcoded key already used directly by the executor and single-pass plugin, so the "key owned by the bridge" framing was misleading. Inline at all five callsites. - Drop the "future prompt_studio" forward-looking comment in ide_callback/tasks.py — only source='lookup' is wired up, and the cloud lookups plugin is the only registrant of the underlying endpoints.
- latest_outputs_by_keys: switch the per-prompt latest pick to
order_by('prompt_id', '-modified_at').distinct('prompt_id') so
Postgres returns at most one row per prompt instead of materialising
every historical run + relying on a Python break.
- fetch_default_output_response: collapse the previous N+1
(exists() + for output in queryset per prompt) into a single
DISTINCT ON (prompt_id, profile_manager_id) query, plus a
per-tool cache for the default-profile lookup. Combined Output is a
hot path — the old shape made every panel switch O(prompts × runs)
with a plugin invocation per matching row.
- Drop the unused _resolve_profile_for_prompt and
_collect_default_output_for_prompt helpers, and the dead
except ObjectDoesNotExist: return '' (neither .exists() nor
the queryset iteration ever raised that exception).
- PromptOutput.jsx: replace empty catch {} on the lookup
dynamic-import sites with console.warn so unexpected chunk-load
failures don't masquerade as OSS-mode behaviour. Add a
resolveCopyText helper that wraps getEnrichedCopyText in
try/catch + fallback so a plugin-side throw can't break the Copy
button at either of the two call sites.
- usePromptOutput.js: same catch (error) -> console.warn for
the two existing dynamic-import sites; wrap the per-item
handleLookupOutput call in try/catch so a single bad enrichment
payload no longer aborts the surrounding forEach and skips the
prompt-output state update.
- prompt_studio_core_v2/views.py: validate prompt_id before
the multi-var lookup gate so a missing field returns a clean 400
instead of a lookup-related error.
- ide_callback/tasks.py: (result_dict.get('data') or {}) so an
explicit data=None from the executor doesn't AttributeError into
a generic ERROR callback. Replace the inner except: pass swallow
with logger.debug so a secondary WS-emit failure during the
outer extraction_complete fallback is recoverable from logs.
….debug
- _run_webhook_postprocessing: skip when structured_output[prompt_name]
is empty or non-iterable. Pre-refactor, the webhook lived inside
handle_json after its parse-failure early-return (which sets the
output to {}), so a malformed JSON answer never dispatched the
webhook — the new explicit gate restores that behaviour. Subscribers
no longer receive empty-payload calls they didn't see before.
- .gitignore: re-include docker/compose.debug.yaml after the
broader docker/compose.*.yaml rule so a delete + recreate doesn't
make the tracked file look untracked, and so teammate-added compose
files aren't silently masked.
`APIException(code=400)` only sets `detail.code` in the JSON body; `status_code` is hardcoded to HTTP_500_INTERNAL_SERVER_ERROR. Switch to `ValidationError` so the missing-tool_id and tool-not-found branches in `latest_outputs_by_keys` and `get_output_for_tool_default` actually respond with 400 as the comment / docstring imply.
Contributor
Author
Review comment dispositionsThanks for the thorough reviews. The latest round of commits on this branch ( Fixed in this round
Already addressed in earlier rounds (this branch)
Stale
Deferred to follow-ups
Skipped — out of scope
🤖 Generated with Claude Code |
Contributor
|
✅ All
The engineering rationale for the deferred items is sound — especially the I'll resolve my open threads and run a fresh review pass against the latest commits. ✅ Actions performedReview triggered.
|
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (5)
workers/ide_callback/tasks.py (1)
592-598:⚠️ Potential issue | 🟠 MajorTreat missing extraction callback endpoints as terminal OSS no-ops.
These endpoints are documented as cloud-only in
workers/shared/clients/extraction_client.py, but both calls currently treat a 404 like a real failure. In OSS builds that meansextraction_complete()can re-raise into retries, andextraction_error()just logs noise instead of exiting cleanly. Catch the client error type here and swallow HTTP 404 explicitly.Also applies to: 666-671
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@workers/ide_callback/tasks.py` around lines 592 - 598, The calls to api.mark_extraction_complete and the corresponding extraction_error handler should treat a 404 from the extraction client as a terminal no-op for OSS; update extraction_complete() and extraction_error() to catch the extraction client exception (the client error type thrown by the extraction client) around the api.mark_extraction_complete / api.mark_extraction_error calls, check if the exception represents HTTP 404 and if so swallow it and return/exit cleanly, otherwise re-raise or propagate the error as before so real failures still surface.workers/executor/executors/legacy_executor.py (1)
2038-2067:⚠️ Potential issue | 🟠 MajorKeep lookup enrichment resilient after the plugin call returns.
The new
try/exceptonly protectsrun_with_metrics()itself. If the plugin returns an object missingusage_records,llm_metrics, orMETRICS_KEY, lines 2066-2067 still raise and fail the prompt after extraction already succeeded. Guard those attributes and skip enrichment when the outcome shape is invalid.Suggested fix
- self._usage_records.extend(outcome.usage_records) - metrics.setdefault(prompt_name, {})[lookup_cls.METRICS_KEY] = outcome.llm_metrics + usage_records = getattr(outcome, "usage_records", None) + llm_metrics = getattr(outcome, "llm_metrics", None) + metrics_key = getattr(lookup_cls, "METRICS_KEY", None) + + if isinstance(usage_records, list): + self._usage_records.extend(usage_records) + if metrics_key is not None and llm_metrics is not None: + metrics.setdefault(prompt_name, {})[metrics_key] = llm_metrics🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@workers/executor/executors/legacy_executor.py` around lines 2038 - 2067, The current try/except only wraps lookup_cls.run_with_metrics and then assumes outcome has the expected shape; validate the returned outcome before mutating state: after calling lookup_cls.run_with_metrics (or inside the same try block), check that outcome has an iterable usage_records and a llm_metrics attribute and that lookup_cls.METRICS_KEY is defined; if any of these are missing or invalid, log the error with logger.exception, call shim.stream_log with a WARN message similar to the existing one, and return/skip enrichment instead of calling self._usage_records.extend or writing into metrics; keep these checks adjacent to lookup_cls.run_with_metrics so malformed plugin returns are handled the same way as raised exceptions.workers/executor/tasks.py (1)
112-123:⚠️ Potential issue | 🟠 MajorEscalate a failed usage flush instead of only logging it.
bulk_create_usage()already has a non-exception failure signal. This branch logs and then returns the original extraction result, so a transient usage API outage still drops the whole batch permanently and Celery never retries. Re-raise here, or push the records to a retryable fallback/DLQ whenokis false.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@workers/executor/tasks.py` around lines 112 - 123, The failure branch after calling usage_client.bulk_create_usage(usage_records) only logs and returns, which silently drops records; change it so a non-exception failure becomes retryable: when ok is False (inside the try surrounding usage_client.bulk_create_usage in tasks.py) either raise a specific exception (e.g., UsageFlushError with run_id/organization_id/record_count in the message) to let Celery retry, or push the usage_records into your retryable fallback/DLQ before returning; update the logger.error to include the same contextual details and then raise or enqueue so transient outages don't permanently drop the batch.backend/prompt_studio/prompt_studio_output_manager_v2/views.py (1)
74-95:⚠️ Potential issue | 🟡 MinorReject malformed
tool_idbefore the ORM filter.This only handles the missing case. A non-UUID string still flows into
ToolStudioPrompt.objects.filter(tool_id=tool_id, ...); withtool_idbacked by a UUID FK, Django raises before the query executes, so this action still returns a 500 instead of the intended 400. The same guard should be reused inget_output_for_tool_default()below.🔧 Suggested guard
+ from uuid import UUID + tool_id = request.GET.get("tool_id") keys_param = request.GET.get("prompt_keys", "") if not tool_id: # ``APIException(code=400)`` only sets ``detail.code`` in the body; # ``status_code`` is hardcoded to 500. Use ``ValidationError`` so # callers get the intended 400. raise ValidationError(detail=PromptOutputManagerErrorMessage.TOOL_VALIDATION) + try: + UUID(str(tool_id)) + except (TypeError, ValueError): + raise ValidationError( + detail=PromptOutputManagerErrorMessage.TOOL_VALIDATION + )#!/bin/bash # Verify that ToolStudioPrompt.tool_id points to a UUID-backed CustomTool key. set -e fd -i '^models\.py$' backend/prompt_studio -x rg -n -C2 'class ToolStudioPrompt|class CustomTool|tool_id\s*=|prompt_id\s*=' {}Expected result:
ToolStudioPrompt.tool_idresolves toCustomTool, andCustomTooluses a UUID-backed identifier. If so, malformed UUIDs should be rejected before the ORM filter.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/prompt_studio/prompt_studio_output_manager_v2/views.py` around lines 74 - 95, Reject malformed tool_id before it reaches the ORM by validating its UUID form right after obtaining tool_id from request.GET; replace the current "if not tool_id" check with a guard that also attempts to parse tool_id as a UUID (or use Django's UUIDField validator) and raises ValidationError(detail=PromptOutputManagerErrorMessage.TOOL_VALIDATION) on parse failure, and apply the same validation in get_output_for_tool_default() so ToolStudioPrompt.objects.filter(tool_id=tool_id, ...) never receives an invalid UUID string.backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py (1)
350-355:⚠️ Potential issue | 🟠 MajorKeep combined-output enrichment fail-open as well.
resultalready has the raw values here, butenrich_prompt_output()andattach_combined_output_enrichment()still cross the cloud lookup bridge without a guard. One plugin-side exception will 500 the whole combined-output endpoint again instead of returning the raw payload.🛡️ Suggested guard
- enriched = enrich_prompt_output(output, {}) - bundle = extract_prompt_output_enrichment(enriched) - if bundle is not None: - enrichment_by_key[tool_prompt.prompt_key] = bundle + try: + enriched = enrich_prompt_output(output, {}) + bundle = extract_prompt_output_enrichment(enriched) + if bundle is not None: + enrichment_by_key[tool_prompt.prompt_key] = bundle + except Exception: + logger.exception( + "Failed to enrich combined output for prompt %s", + tool_prompt.prompt_key, + ) - attach_combined_output_enrichment(result, enrichment_by_key) + try: + attach_combined_output_enrichment(result, enrichment_by_key) + except Exception: + logger.exception("Failed to attach combined output enrichment")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py` around lines 350 - 355, Wrap the enrichment/attachment steps in a try/except so plugin/cloud lookup failures don't crash the combined-output endpoint: when iterating over tool prompts, call enrich_prompt_output(output, {}) and extract_prompt_output_enrichment(enriched) inside a try block and only add to enrichment_by_key on success (use tool_prompt.prompt_key), and then call attach_combined_output_enrichment(result, enrichment_by_key) inside its own try/except (or skip if enrichment_by_key is empty); on exception log the error and proceed to return the existing raw result unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py`:
- Around line 1180-1182: The legacy single-pass flow is missing lookup_configs:
when building the legacy single-pass payload in _fetch_single_pass_response
(which constructs single_pass_extraction) you must include the same
lookup_configs produced by get_lookup_configs_for_tool(tool, prompts=prompts) so
lookups continue to work; either call get_lookup_configs_for_tool there and add
the result into the single_pass_extraction/tool_settings, or refactor to reuse
build_single_pass_payload to produce a canonical payload that both
prompt_responder() and _fetch_single_pass_response() consume, ensuring
lookup_configs is present in both paths.
In `@backend/prompt_studio/prompt_studio_core_v2/views.py`:
- Around line 109-123: The guard currently checks the tool's setting via
custom_tool.single_pass_extraction_mode which allows callers like
fetch_response() and bulk_fetch_response() to bypass the multi-variable lookup
protection; change the helper to accept an explicit is_single_pass boolean
(e.g., add parameter is_single_pass=False) and use that to decide whether to
skip the block instead of reading custom_tool.single_pass_extraction_mode,
update all callers: call with is_single_pass=False from fetch_response() and
bulk_fetch_response(), and either call with is_single_pass=True (or skip the
helper) from single_pass_extraction(); keep the existing Response behavior and
use get_multi_var_lookups_for_tool and names logic unchanged.
In
`@backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py`:
- Around line 301-304: The lookup validation is being run against the full
prompts list (lookup_configs, lookup_error =
validate_lookups_for_export(prompts)) before the export loop filters out NOTES
and inactive prompts, causing spurious failures; update the code to perform
validate_lookups_for_export on the filtered set that will actually be exported
(i.e., the prompts after the logic that drops NOTES and inactive entries) or
pass the already-filtered collection into validate_lookups_for_export so only
exportable prompts are validated (refer to validate_lookups_for_export and the
export filtering loop that removes NOTES/inactive prompts).
In `@unstract/sdk1/src/unstract/sdk1/usage_handler.py`:
- Around line 108-154: The code assumes self.token_counter exists when recording
embeddings; guard accesses to avoid AttributeError by checking if
self.token_counter is not None before reading total_embedding_token_count and
before calling reset_counts in the embedding branch (where embed_model is
checked). If token_counter is None, use 0 for embedding_tokens (and any token
counts) and skip calling reset_counts; keep the rest of the _pending_usage
payload the same. Reference symbols: self.token_counter,
total_embedding_token_count, reset_counts, embed_model, and the block that
appends to self._pending_usage.
---
Duplicate comments:
In
`@backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py`:
- Around line 350-355: Wrap the enrichment/attachment steps in a try/except so
plugin/cloud lookup failures don't crash the combined-output endpoint: when
iterating over tool prompts, call enrich_prompt_output(output, {}) and
extract_prompt_output_enrichment(enriched) inside a try block and only add to
enrichment_by_key on success (use tool_prompt.prompt_key), and then call
attach_combined_output_enrichment(result, enrichment_by_key) inside its own
try/except (or skip if enrichment_by_key is empty); on exception log the error
and proceed to return the existing raw result unchanged.
In `@backend/prompt_studio/prompt_studio_output_manager_v2/views.py`:
- Around line 74-95: Reject malformed tool_id before it reaches the ORM by
validating its UUID form right after obtaining tool_id from request.GET; replace
the current "if not tool_id" check with a guard that also attempts to parse
tool_id as a UUID (or use Django's UUIDField validator) and raises
ValidationError(detail=PromptOutputManagerErrorMessage.TOOL_VALIDATION) on parse
failure, and apply the same validation in get_output_for_tool_default() so
ToolStudioPrompt.objects.filter(tool_id=tool_id, ...) never receives an invalid
UUID string.
In `@workers/executor/executors/legacy_executor.py`:
- Around line 2038-2067: The current try/except only wraps
lookup_cls.run_with_metrics and then assumes outcome has the expected shape;
validate the returned outcome before mutating state: after calling
lookup_cls.run_with_metrics (or inside the same try block), check that outcome
has an iterable usage_records and a llm_metrics attribute and that
lookup_cls.METRICS_KEY is defined; if any of these are missing or invalid, log
the error with logger.exception, call shim.stream_log with a WARN message
similar to the existing one, and return/skip enrichment instead of calling
self._usage_records.extend or writing into metrics; keep these checks adjacent
to lookup_cls.run_with_metrics so malformed plugin returns are handled the same
way as raised exceptions.
In `@workers/executor/tasks.py`:
- Around line 112-123: The failure branch after calling
usage_client.bulk_create_usage(usage_records) only logs and returns, which
silently drops records; change it so a non-exception failure becomes retryable:
when ok is False (inside the try surrounding usage_client.bulk_create_usage in
tasks.py) either raise a specific exception (e.g., UsageFlushError with
run_id/organization_id/record_count in the message) to let Celery retry, or push
the usage_records into your retryable fallback/DLQ before returning; update the
logger.error to include the same contextual details and then raise or enqueue so
transient outages don't permanently drop the batch.
In `@workers/ide_callback/tasks.py`:
- Around line 592-598: The calls to api.mark_extraction_complete and the
corresponding extraction_error handler should treat a 404 from the extraction
client as a terminal no-op for OSS; update extraction_complete() and
extraction_error() to catch the extraction client exception (the client error
type thrown by the extraction client) around the api.mark_extraction_complete /
api.mark_extraction_error calls, check if the exception represents HTTP 404 and
if so swallow it and return/exit cleanly, otherwise re-raise or propagate the
error as before so real failures still surface.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0e13657d-89a4-4202-ab9f-951b5ca9868a
📒 Files selected for processing (20)
.gitignorebackend/prompt_studio/lookup_utils.pybackend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.pybackend/prompt_studio/prompt_studio_core_v2/views.pybackend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.pybackend/prompt_studio/prompt_studio_output_manager_v2/serializers.pybackend/prompt_studio/prompt_studio_output_manager_v2/views.pybackend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.pybackend/usage_v2/internal_views.pybackend/usage_v2/migrations/0006_alter_usage_status_and_more.pybackend/usage_v2/models.pyfrontend/src/components/custom-tools/prompt-card/PromptOutput.jsxfrontend/src/hooks/usePromptOutput.jsunstract/sdk1/src/unstract/sdk1/llm.pyunstract/sdk1/src/unstract/sdk1/usage_handler.pyworkers/executor/executors/legacy_executor.pyworkers/executor/executors/plugins/protocols.pyworkers/executor/tasks.pyworkers/ide_callback/tasks.pyworkers/shared/clients/usage_client.py
✅ Files skipped from review due to trivial changes (1)
- workers/executor/executors/plugins/protocols.py
🚧 Files skipped from review as they are similar to previous changes (5)
- .gitignore
- frontend/src/hooks/usePromptOutput.js
- frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx
- backend/usage_v2/internal_views.py
- backend/prompt_studio/lookup_utils.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (10)
workers/executor/tasks.py (1)
113-130:⚠️ Potential issue | 🟠 MajorUsage flush failure is still non-recoverable (records can be dropped).
When
bulk_create_usagefails, this path only logs and continues. That still allows silent billing-row loss on transient outages.Proposed fix
if not ok: logger.error( "bulk_create_usage returned failure for %d records " "(run_id=%s organization_id=%s)", len(usage_records), context.run_id, context.organization_id, ) + raise ConnectionError("bulk_create_usage returned failure") except Exception: logger.error( "Failed to flush %d usage records for run_id=%s organization_id=%s", len(usage_records), context.run_id, context.organization_id, exc_info=True, ) + raise🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@workers/executor/tasks.py` around lines 113 - 130, The current flush path in tasks.py logs failures from bulk_create_usage but swallows them, risking dropped billing rows; modify the failure handling inside the block that calls bulk_create_usage so that on a non-ok result or on exception you propagate an error to the caller (either by re-raising the caught exception or raising a RecoverableError) instead of only logging, so the orchestration can retry; keep the existing logger.error messages (including run_id and organization_id) but follow them with a raise to fail the task and preserve usage_records/context for retry handling.backend/usage_v2/internal_views.py (1)
164-193:⚠️ Potential issue | 🟠 MajorValidate record shape before building
Usageobjects.Current permissive defaults (
usage_type,adapter_instance_id,model_name, etc.) can mask producer bugs, and one malformed row can fail the entire batch write.Proposed fix sketch
usage_objects = [] - for r in records: + for i, r in enumerate(records): + if not isinstance(r, dict): + return JsonResponse( + {"success": False, "error": f"Invalid record at index {i}"}, + status=status.HTTP_400_BAD_REQUEST, + ) + missing = [k for k in ("usage_type", "adapter_instance_id", "model_name") if k not in r] + if missing: + return JsonResponse( + { + "success": False, + "error": f"Missing required keys at index {i}: {', '.join(missing)}", + }, + status=status.HTTP_400_BAD_REQUEST, + ) usage_objects.append( Usage(🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/usage_v2/internal_views.py` around lines 164 - 193, The loop that constructs Usage objects from records (the for r in records -> Usage(...) block) is too permissive and can mask producer bugs; add explicit validation before building each Usage: verify required fields (e.g., workflow_id, execution_id, usage_type, adapter_instance_id or run_id as applicable) are present and of the expected type, coerce/validate numeric fields (embedding_tokens, prompt_tokens, completion_tokens, total_tokens as ints; cost_in_dollars as float), ensure model_name is non-empty when usage_type == "llm", and enforce the llm_usage_reason None/coercion rule; for records that fail validation, log a clear error including the offending record and either skip the record or collect/report them so the batch write isn't aborted by a single malformed row, then only append validated Usage instances to usage_objects.unstract/sdk1/src/unstract/sdk1/usage_handler.py (1)
108-112:⚠️ Potential issue | 🟠 MajorGuard
token_counterbefore dereference in embedding callback.
token_counteris nullable, but Line 111 and Line 154 assume it is always present. This can raiseAttributeErrorin production callback flows and skip usage persistence.Proposed fix
): if self.embed_model is None: return + if self.token_counter is None: + logger.warning( + "Embedding usage callback invoked without token_counter; skipping usage record" + ) + return model_name = self.embed_model.model_name embedding_tokens = self.token_counter.total_embedding_token_count @@ - self.token_counter.reset_counts() + self.token_counter.reset_counts()Also applies to: 154-154
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@unstract/sdk1/src/unstract/sdk1/usage_handler.py` around lines 108 - 112, The code dereferences self.token_counter (accessing total_embedding_token_count) without checking for None, which can raise AttributeError and skip usage persistence; fix by guarding token_counter in the embedding callback(s): check if self.token_counter is not None before reading self.token_counter.total_embedding_token_count (use 0 or an appropriate default when None) and apply the same guard wherever total_embedding_token_count is accessed (e.g., the lines referencing self.token_counter at the embedding callback and at the later usage around line 154); ensure stream_log and any usage persistence code use the guarded/ default value instead of assuming token_counter exists.workers/ide_callback/tasks.py (2)
589-590:⚠️ Potential issue | 🟡 MinorNormalize
extracted_textwhen present-but-None.At Line 590,
len(extracted_text)can still fail ifdata["extracted_text"]is explicitlyNone.Proposed fix
- extracted_text = (result_dict.get("data") or {}).get("extracted_text", "") + data = result_dict.get("data") or {} + extracted_text = data.get("extracted_text") or "" token_count = len(extracted_text) // 4🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@workers/ide_callback/tasks.py` around lines 589 - 590, The code assigns extracted_text = (result_dict.get("data") or {}).get("extracted_text", "") but if data["extracted_text"] exists and is None len(extracted_text) will raise; normalize extracted_text to a string before computing token_count by coercing None to "" (e.g., replace None with "" for the extracted_text variable used in token_count calculation in tasks.py), then compute token_count = len(extracted_text) // 4 so len() is safe.
567-572:⚠️ Potential issue | 🟠 MajorTreat HTTP 404 from extraction persistence endpoints as terminal OSS no-op.
If cloud-only extraction endpoints are absent, persistence calls can 404 and currently flow into the outer exception path, then re-raise on Line 644. That breaks OSS-safe degradation for this callback path.
Proposed fix
@@ def extraction_complete( @@ - try: + def _is_http_404(exc: Exception) -> bool: + status_code = getattr(exc, "status_code", None) + if status_code is None: + status_code = getattr(getattr(exc, "response", None), "status_code", None) + return status_code == 404 + + try: @@ - api.mark_extraction_error( - source=source, - file_id=file_id, - error=error_msg, - organization_id=org_id, - ) + try: + api.mark_extraction_error( + source=source, + file_id=file_id, + error=error_msg, + organization_id=org_id, + ) + except Exception as e: + if _is_http_404(e): + logger.info( + "Skipping extraction_error persistence (404 endpoint missing): source=%s file=%s", + source, + file_id, + ) + return {"status": "ok", "note": "extraction callback endpoint not present (OSS)"} + raise @@ - api.mark_extraction_complete( - source=source, - file_id=file_id, - token_count=token_count, - extracted_text_path=extracted_text_path, - organization_id=org_id, - ) + try: + api.mark_extraction_complete( + source=source, + file_id=file_id, + token_count=token_count, + extracted_text_path=extracted_text_path, + organization_id=org_id, + ) + except Exception as e: + if _is_http_404(e): + logger.info( + "Skipping extraction_complete persistence (404 endpoint missing): source=%s file=%s", + source, + file_id, + ) + return {"status": "ok", "note": "extraction callback endpoint not present (OSS)"} + raiseAlso applies to: 592-598, 620-644
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@workers/ide_callback/tasks.py` around lines 567 - 572, The persistence calls to the API (e.g., api.mark_extraction_error called with source, file_id, organization_id/org_id and similar calls at the other locations) must treat an HTTP 404 from the extraction persistence endpoints as a terminal OSS no-op: wrap each api.* persistence call (api.mark_extraction_error, the matching api.mark_extraction_success/persist call(s) around lines ~592-598 and ~620-644) in a try/except that catches the HTTP error type your client raises (e.g., HTTPError/ApiError) and if the caught exception indicates response.status_code == 404, swallow it and return/exit this callback path instead of propagating; for any other exception re-raise so existing error handling still applies. Ensure you reference the same parameters (source, file_id, org_id/organization_id, error_msg, extracted_files, checksum) when locating and wrapping the calls.backend/prompt_studio/prompt_studio_core_v2/views.py (2)
923-926:⚠️ Potential issue | 🔴 Critical
prompt_studio_toolis still effectively caller-controlled on this detail route.Lines 923-926 only change which tool is used for the profile-count check. On Line 934,
self.perform_create(serializer)still persists anyprompt_studio_toolthe request supplied, so a caller with access to one tool can create a profile manager under another tool ID. Force the save againstself.get_object()here, or make the serializer field permission-scoped.Suggested fix
- prompt_studio_tool = ( - serializer.validated_data.get(ProfileManagerKeys.PROMPT_STUDIO_TOOL) - or self.get_object() - ) + prompt_studio_tool = self.get_object() @@ - self.perform_create(serializer) + serializer.save(prompt_studio_tool=prompt_studio_tool)Also applies to: 934-934
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/prompt_studio/prompt_studio_core_v2/views.py` around lines 923 - 926, The code currently allows a client-supplied prompt_studio_tool to be persisted because perform_create(serializer) uses serializer data; before calling self.perform_create(serializer) (in the view method surrounding get_object and serializer.validated_data), force the tool to the detail object by setting serializer.validated_data[ProfileManagerKeys.PROMPT_STUDIO_TOOL] = self.get_object() or call serializer.save(prompt_studio_tool=self.get_object()) so the persisted ProfileManager always uses self.get_object() instead of any caller-supplied value; alternatively, enforce the field is read-only/permission-scoped on the serializer to prevent client override.
99-110:⚠️ Potential issue | 🟠 MajorUse the requested execution path to decide whether to skip this gate.
Line 109 still keys the bypass off
custom_tool.single_pass_extraction_mode, sofetch_response()/bulk_fetch_response()can skip the non-single-pass lookup block just by calling those endpoints on a tool configured for single-pass mode. Pass an explicitis_single_passflag from each caller instead of inferring it from tool settings.Suggested direction
-def _multi_var_lookup_block_response(custom_tool, prompt_ids=None): +def _multi_var_lookup_block_response( + custom_tool, *, is_single_pass: bool, prompt_ids=None +): @@ - if getattr(custom_tool, "single_pass_extraction_mode", False): + if is_single_pass: return None- if err := _multi_var_lookup_block_response(custom_tool, prompt_ids=[prompt_id]): + if err := _multi_var_lookup_block_response( + custom_tool, is_single_pass=False, prompt_ids=[prompt_id] + ): return err- if err := _multi_var_lookup_block_response(custom_tool, prompt_ids=prompt_ids): + if err := _multi_var_lookup_block_response( + custom_tool, is_single_pass=False, prompt_ids=prompt_ids + ): return err🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/prompt_studio/prompt_studio_core_v2/views.py` around lines 99 - 110, The gate function _multi_var_lookup_block_response currently checks custom_tool.single_pass_extraction_mode directly; change its signature to accept an explicit is_single_pass boolean (e.g., def _multi_var_lookup_block_response(custom_tool, prompt_ids=None, is_single_pass=False)) and use that flag to decide skipping instead of getattr(custom_tool, "single_pass_extraction_mode", False). Update each caller (e.g., fetch_response and bulk_fetch_response) to pass the correct execution path by computing and passing is_single_pass from the request context/parameters rather than relying on tool settings. Ensure all call sites of _multi_var_lookup_block_response are updated to include the new parameter to avoid defaulting to the wrong behavior.backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py (1)
301-304:⚠️ Potential issue | 🟠 MajorValidate lookup assignments only for prompts that are actually exportable.
Validation is still executed on the full prompt list, but NOTES/inactive prompts are skipped during export. This can fail export for prompts that will never be emitted.
🔧 Suggested fix
- lookup_configs, lookup_error = validate_lookups_for_export(prompts) + lookup_validation_prompts = [ + prompt + for prompt in prompts + if prompt.prompt_type != JsonSchemaKey.NOTES and prompt.active + ] + lookup_configs, lookup_error = validate_lookups_for_export( + lookup_validation_prompts + ) if lookup_error: raise InValidCustomToolError(lookup_error)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py` around lines 301 - 304, The code calls validate_lookups_for_export(prompts) over the full prompts list which includes NOTES/inactive prompts and can block export for items that are never emitted; change the call to run only on the subset of exportable prompts (filter prompts to exclude NOTES/inactive entries before calling validate_lookups_for_export) so that validate_lookups_for_export only inspects prompts that will be exported and only raise InValidCustomToolError when those exportable prompts fail validation.backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py (1)
1180-1182:⚠️ Potential issue | 🟠 MajorSingle-pass lookup wiring is still missing in the legacy execution path.
This only updates
build_single_pass_payload().prompt_responder()still routes no-idflow through_execute_prompts_in_single_pass()→_fetch_single_pass_response(), wheretool_settings["lookup_configs"]is still not attached.🔧 Suggested follow-up
def _fetch_single_pass_response(...): @@ tool_settings[TSPKeys.SIMILARITY_TOP_K] = default_profile.similarity_top_k + + lookup_configs = get_lookup_configs_for_tool(tool, prompts=prompts) + if lookup_configs: + tool_settings["lookup_configs"] = lookup_configs🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py` around lines 1180 - 1182, The no-id legacy execution path in prompt_responder still fails to pass lookup configs into the single-pass flow, so ensure tool_settings["lookup_configs"] is populated for that path by calling get_lookup_configs_for_tool(tool, prompts=prompts) and assigning the result to tool_settings["lookup_configs"] before invoking _execute_prompts_in_single_pass() (and consequently _fetch_single_pass_response()); update prompt_responder to mirror the build_single_pass_payload() change by computing lookup_configs from prompts and attaching them to the tool_settings used in the no-id flow.backend/prompt_studio/prompt_studio_output_manager_v2/views.py (1)
74-96:⚠️ Potential issue | 🟠 MajorValidate
tool_idformat before ORM filters.A malformed
tool_idcan still bubble out as a server error instead of a clean 400. Add a UUID shape guard before both query paths usetool_id.🔧 Suggested fix
+from uuid import UUID @@ tool_id = request.GET.get("tool_id") @@ if not tool_id: raise ValidationError(detail=PromptOutputManagerErrorMessage.TOOL_VALIDATION) + try: + UUID(str(tool_id)) + except (TypeError, ValueError): + raise ValidationError( + detail=PromptOutputManagerErrorMessage.TOOL_VALIDATION + ) @@ tool_id = request.GET.get("tool_id") @@ if not tool_id: raise ValidationError(detail=tool_validation_message) + try: + UUID(str(tool_id)) + except (TypeError, ValueError): + raise ValidationError(detail=tool_validation_message)Please verify with an API test that
tool_id=not-a-uuidreturns HTTP 400 on both endpoints.Also applies to: 127-138
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/prompt_studio/prompt_studio_output_manager_v2/views.py` around lines 74 - 96, Add a UUID-shape guard that validates tool_id before any ORM usage: attempt to parse request.GET.get("tool_id") with uuid.UUID(...) in a try/except and if parsing fails raise ValidationError(detail=PromptOutputManagerErrorMessage.TOOL_VALIDATION). Apply this check in the current handler before building prompt_keys and before the other query path referenced in the comment (the block using ToolStudioPrompt.objects.filter and the similar block at lines 127-138), so no malformed tool_id reaches the ORM; keep using the same ValidationError and message to ensure the endpoint returns HTTP 400.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@backend/prompt_studio/lookup_utils.py`:
- Around line 18-38: The ImportError guard misses the parent package name, so
when pluggable_apps is absent the exception is re-raised; update
_CLOUD_LOOKUP_MODULES to include "pluggable_apps" (in addition to the existing
entries) so the except block will treat missing the parent package as expected
and set LOOKUPS_AVAILABLE = False; ensure this change covers the import block
that imports execution/output_enrichment/staleness/validation and
LookupOutputResult (symbols: _CLOUD_LOOKUP_MODULES, _execution,
_output_enrichment, _staleness, _validation, _LookupOutputResult,
LOOKUPS_AVAILABLE, and the ImportError handling using e.name).
---
Duplicate comments:
In `@backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py`:
- Around line 1180-1182: The no-id legacy execution path in prompt_responder
still fails to pass lookup configs into the single-pass flow, so ensure
tool_settings["lookup_configs"] is populated for that path by calling
get_lookup_configs_for_tool(tool, prompts=prompts) and assigning the result to
tool_settings["lookup_configs"] before invoking
_execute_prompts_in_single_pass() (and consequently
_fetch_single_pass_response()); update prompt_responder to mirror the
build_single_pass_payload() change by computing lookup_configs from prompts and
attaching them to the tool_settings used in the no-id flow.
In `@backend/prompt_studio/prompt_studio_core_v2/views.py`:
- Around line 923-926: The code currently allows a client-supplied
prompt_studio_tool to be persisted because perform_create(serializer) uses
serializer data; before calling self.perform_create(serializer) (in the view
method surrounding get_object and serializer.validated_data), force the tool to
the detail object by setting
serializer.validated_data[ProfileManagerKeys.PROMPT_STUDIO_TOOL] =
self.get_object() or call serializer.save(prompt_studio_tool=self.get_object())
so the persisted ProfileManager always uses self.get_object() instead of any
caller-supplied value; alternatively, enforce the field is
read-only/permission-scoped on the serializer to prevent client override.
- Around line 99-110: The gate function _multi_var_lookup_block_response
currently checks custom_tool.single_pass_extraction_mode directly; change its
signature to accept an explicit is_single_pass boolean (e.g., def
_multi_var_lookup_block_response(custom_tool, prompt_ids=None,
is_single_pass=False)) and use that flag to decide skipping instead of
getattr(custom_tool, "single_pass_extraction_mode", False). Update each caller
(e.g., fetch_response and bulk_fetch_response) to pass the correct execution
path by computing and passing is_single_pass from the request context/parameters
rather than relying on tool settings. Ensure all call sites of
_multi_var_lookup_block_response are updated to include the new parameter to
avoid defaulting to the wrong behavior.
In `@backend/prompt_studio/prompt_studio_output_manager_v2/views.py`:
- Around line 74-96: Add a UUID-shape guard that validates tool_id before any
ORM usage: attempt to parse request.GET.get("tool_id") with uuid.UUID(...) in a
try/except and if parsing fails raise
ValidationError(detail=PromptOutputManagerErrorMessage.TOOL_VALIDATION). Apply
this check in the current handler before building prompt_keys and before the
other query path referenced in the comment (the block using
ToolStudioPrompt.objects.filter and the similar block at lines 127-138), so no
malformed tool_id reaches the ORM; keep using the same ValidationError and
message to ensure the endpoint returns HTTP 400.
In
`@backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py`:
- Around line 301-304: The code calls validate_lookups_for_export(prompts) over
the full prompts list which includes NOTES/inactive prompts and can block export
for items that are never emitted; change the call to run only on the subset of
exportable prompts (filter prompts to exclude NOTES/inactive entries before
calling validate_lookups_for_export) so that validate_lookups_for_export only
inspects prompts that will be exported and only raise InValidCustomToolError
when those exportable prompts fail validation.
In `@backend/usage_v2/internal_views.py`:
- Around line 164-193: The loop that constructs Usage objects from records (the
for r in records -> Usage(...) block) is too permissive and can mask producer
bugs; add explicit validation before building each Usage: verify required fields
(e.g., workflow_id, execution_id, usage_type, adapter_instance_id or run_id as
applicable) are present and of the expected type, coerce/validate numeric fields
(embedding_tokens, prompt_tokens, completion_tokens, total_tokens as ints;
cost_in_dollars as float), ensure model_name is non-empty when usage_type ==
"llm", and enforce the llm_usage_reason None/coercion rule; for records that
fail validation, log a clear error including the offending record and either
skip the record or collect/report them so the batch write isn't aborted by a
single malformed row, then only append validated Usage instances to
usage_objects.
In `@unstract/sdk1/src/unstract/sdk1/usage_handler.py`:
- Around line 108-112: The code dereferences self.token_counter (accessing
total_embedding_token_count) without checking for None, which can raise
AttributeError and skip usage persistence; fix by guarding token_counter in the
embedding callback(s): check if self.token_counter is not None before reading
self.token_counter.total_embedding_token_count (use 0 or an appropriate default
when None) and apply the same guard wherever total_embedding_token_count is
accessed (e.g., the lines referencing self.token_counter at the embedding
callback and at the later usage around line 154); ensure stream_log and any
usage persistence code use the guarded/ default value instead of assuming
token_counter exists.
In `@workers/executor/tasks.py`:
- Around line 113-130: The current flush path in tasks.py logs failures from
bulk_create_usage but swallows them, risking dropped billing rows; modify the
failure handling inside the block that calls bulk_create_usage so that on a
non-ok result or on exception you propagate an error to the caller (either by
re-raising the caught exception or raising a RecoverableError) instead of only
logging, so the orchestration can retry; keep the existing logger.error messages
(including run_id and organization_id) but follow them with a raise to fail the
task and preserve usage_records/context for retry handling.
In `@workers/ide_callback/tasks.py`:
- Around line 589-590: The code assigns extracted_text =
(result_dict.get("data") or {}).get("extracted_text", "") but if
data["extracted_text"] exists and is None len(extracted_text) will raise;
normalize extracted_text to a string before computing token_count by coercing
None to "" (e.g., replace None with "" for the extracted_text variable used in
token_count calculation in tasks.py), then compute token_count =
len(extracted_text) // 4 so len() is safe.
- Around line 567-572: The persistence calls to the API (e.g.,
api.mark_extraction_error called with source, file_id, organization_id/org_id
and similar calls at the other locations) must treat an HTTP 404 from the
extraction persistence endpoints as a terminal OSS no-op: wrap each api.*
persistence call (api.mark_extraction_error, the matching
api.mark_extraction_success/persist call(s) around lines ~592-598 and ~620-644)
in a try/except that catches the HTTP error type your client raises (e.g.,
HTTPError/ApiError) and if the caught exception indicates response.status_code
== 404, swallow it and return/exit this callback path instead of propagating;
for any other exception re-raise so existing error handling still applies.
Ensure you reference the same parameters (source, file_id,
org_id/organization_id, error_msg, extracted_files, checksum) when locating and
wrapping the calls.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: ce8433f8-3196-4387-824b-9de175c39622
📒 Files selected for processing (20)
.gitignorebackend/prompt_studio/lookup_utils.pybackend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.pybackend/prompt_studio/prompt_studio_core_v2/views.pybackend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.pybackend/prompt_studio/prompt_studio_output_manager_v2/serializers.pybackend/prompt_studio/prompt_studio_output_manager_v2/views.pybackend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.pybackend/usage_v2/internal_views.pybackend/usage_v2/migrations/0006_alter_usage_status_and_more.pybackend/usage_v2/models.pyfrontend/src/components/custom-tools/prompt-card/PromptOutput.jsxfrontend/src/hooks/usePromptOutput.jsunstract/sdk1/src/unstract/sdk1/llm.pyunstract/sdk1/src/unstract/sdk1/usage_handler.pyworkers/executor/executors/legacy_executor.pyworkers/executor/executors/plugins/protocols.pyworkers/executor/tasks.pyworkers/ide_callback/tasks.pyworkers/shared/clients/usage_client.py
✅ Files skipped from review due to trivial changes (2)
- backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py
- workers/executor/executors/legacy_executor.py
🚧 Files skipped from review as they are similar to previous changes (1)
- .gitignore
Five threads on PR #1929 raised against the latest disposition push: * prompt_studio_core_v2/views.py: drop the ``single_pass_extraction_mode`` bypass from ``_multi_var_lookup_block_response`` — fetch_response / bulk_fetch_response are always non-SP, so the tool-setting check just let multi-var lookups slip past the gate when the tool happened to be configured for SP. * prompt_studio_registry_helper.py: filter out NOTES + inactive prompts *before* calling ``validate_lookups_for_export`` so an incomplete lookup on a non-exportable prompt no longer fails the whole export. * unstract/sdk1/usage_handler.py: guard ``self.token_counter is None`` in the embedding-end branch — degrade with a warning instead of an AttributeError on early callbacks. * workers/executor/legacy_executor.py: move the ``outcome.usage_records`` / ``outcome.llm_metrics`` access inside the ``try`` so plugin contract drift hits the same graceful-degrade branch as a thrown ``run_with_metrics``. * backend/prompt_studio/lookup_utils.py: include ``pluggable_apps`` itself in ``_CLOUD_LOOKUP_MODULES``. Pure OSS images don't have the parent package, so ``ImportError.name`` surfaces as ``"pluggable_apps"`` and the previous filter re-raised instead of setting LOOKUPS_AVAILABLE=False.
Contributor
Frontend Lint Report (Biome)✅ All checks passed! No linting or formatting issues found. |
Contributor
Author
Round-2 dispositionOSS commit
🤖 Generated with Claude Code |
Contributor
|
Tip For best results, initiate chat on the files or code changes.
The I'll resolve all my open threads now. [resolve] |
Contributor
Test ResultsSummary
Runner Tests - Full Report
SDK1 Tests - Full Report
|
|
Contributor
Code reviewNo issues found. Checked for bugs and CLAUDE.md compliance. 🤖 Generated with Claude Code - If this code review was useful, please react with 👍. Otherwise, react with 👎. |
| ) | ||
|
|
||
|
|
||
| class UsageBatchCreateView(APIView): |
Contributor
There was a problem hiding this comment.
The UsageBatchCreateView has no explicit authentication or permission class. While the existing internal views in this file may rely on some middleware/URL-level gating, this is a
billing-critical write endpoint that bulk-creates usage records. If the internal API is accessible without auth (e.g., via service mesh), any caller can inject arbitrary usage rows.
Suggestion: Explicitly add permission_classes or document which middleware gates this endpoint. At minimum add a comment clarifying the auth contract.
| ) | ||
|
|
||
| usage_objects = [] | ||
| for r in records: |
Contributor
There was a problem hiding this comment.
The records list is iterated and each dict's values are passed directly to Usage() model construction via .get() with no schema validation. A malformed record (wrong types, missing keys,
extra keys) will only surface as a DB-level error caught by the broad except Exception.
Suggestion: Add a serializer or at minimum validate required fields before constructing model instances. This is a billing-critical table — silent failures or corrupt data would be costly.
| prompt_tokens=usage_data.get("prompt_tokens", 0), | ||
| completion_tokens=usage_data.get("completion_tokens", 0), | ||
| total_tokens=usage_data.get("total_tokens", 0), | ||
| prompt_tokens = usage_data.get("prompt_tokens", 0) |
Contributor
There was a problem hiding this comment.
The old _record_usage called token_counter(model=model, messages=messages) from litellm to independently count prompt tokens. The new version blindly trusts usage_data.get("prompt_tokens",
0) from the API response. Some providers return 0 or omit token counts. This silently records 0 tokens where previously litellm would compute an accurate count.
Suggestion: Add a fallback: if usage_data returns 0 for prompt_tokens but messages are non-empty, use litellm.token_counter() as a fallback.
| ProfileManagerKeys.PROMPT_STUDIO_TOOL | ||
| ] | ||
| prompt_studio_tool = ( | ||
| serializer.validated_data.get(ProfileManagerKeys.PROMPT_STUDIO_TOOL) |
Contributor
There was a problem hiding this comment.
This changes a KeyError (immediate failure) to a silent fallback to self.get_object(). If the serializer is supposed to always include the tool, this masks bugs. If the field is genuinely
optional now, the serializer's field definition should reflect that.
| # callers get the intended 400. | ||
| raise ValidationError(detail=PromptOutputManagerErrorMessage.TOOL_VALIDATION) | ||
|
|
||
| prompt_keys = [k.strip() for k in keys_param.split(",") if k.strip()] |
Contributor
There was a problem hiding this comment.
The new latest_outputs_by_keys endpoint accepts a comma-separated prompt_keys param with no limit on how many keys can be requested. A request with hundreds of keys would generate a large
SQL query and potentially large response.
prompt_keys = [k.strip() for k in keys_param.split(",") if k.strip()]
Suggestion: Cap prompt_keys to a reasonable maximum (e.g., 100) and return 400 if exceeded.
| }, [scrollToBottom]); | ||
|
|
||
| // Handle scrollTo query param for cross-linking from Lookup Studio | ||
| const UUID_RE = |
Contributor
There was a problem hiding this comment.
The regex is declared inside the component body, so it's recreated on every render:
Suggestion: Move UUID_RE outside the component as a module-level constant.
| level=LogLevel.ERROR, | ||
| ) | ||
|
|
||
| def _run_lookup_enrichment( |
Contributor
There was a problem hiding this comment.
REFACTOR : Can move look-up related helpers to a separate directory for maintainability and logical separation. By this way we can prevent maintaining large files.
| prompt_name, | ||
| ) | ||
| return | ||
| if output_type != PSKeys.JSON: |
Contributor
There was a problem hiding this comment.
Previously, webhook postprocessing lived inside answer_prompt.handle_json and ran for any output type (the code was inside the JSON handler but after the parse). Now it's explicitly gated
to output_type == PSKeys.JSON. This is a behavioral change for any user who has webhooks configured on non-JSON prompts — their webhooks will silently stop firing.
Suggestion: This should be called out prominently in the PR description and migration notes. Consider a deprecation warning period rather than silently breaking existing webhook
configurations.
| names = get_multi_var_lookups_for_tool(custom_tool, prompt_ids=prompt_ids) | ||
| if not names: | ||
| return None | ||
| return Response( |
Contributor
There was a problem hiding this comment.
_multi_var_lookup_block_response returns {"error": "..."} but DRF's standard pattern uses {"detail": "..."}. Other endpoints in this file use both. Lets standardize on one.
| return { | ||
| output, | ||
| enriched: buildEnrichedFromBundle(output, bundle, displayPromptResult), | ||
| hasEnriched: true, |
Contributor
There was a problem hiding this comment.
hasEnriched Is Always true in Default Profile Path. This means the enriched toggle may appear even when there's no enrichment data. It should be hasEnriched: bundle != null.
| # paths; guard against AttributeError so a benign empty-data response | ||
| # doesn't escalate to a generic "ERROR" callback. | ||
| extracted_text = (result_dict.get("data") or {}).get("extracted_text", "") | ||
| token_count = len(extracted_text) // 4 |
Contributor
There was a problem hiding this comment.
This character-count-based estimation (÷4) is a very rough approximation and can be significantly off for non-English text or code. Consider using litellm.token_counter since litellm is
already a dependency.
| from django.db import migrations, models | ||
|
|
||
|
|
||
| class Migration(migrations.Migration): |
Contributor
There was a problem hiding this comment.
Migration 0004 adds the status field without choices, and 0006 immediately alters it to add choices. Since this is a new feature branch, these could be squashed into one migration to avoid
unnecessary ALTER TABLE operations on deployment.
| // OSS: plugin may not exist; cloud: surface unexpected chunk-load | ||
| // failures so they don't degrade silently to OSS-mode behaviour. | ||
| // eslint-disable-next-line no-console | ||
| console.warn("[PromptOutput] LookupOutputTabs unavailable:", error); |
Contributor
There was a problem hiding this comment.
In OSS builds, this will emit console warnings on every page load for every user. These cloud-only plugin imports should use a silent catch (like other similar patterns in this PR do with
empty catch {}).
| if vector_db: | ||
| vector_db.close() | ||
|
|
||
| def _init_llm_and_retrieval( |
Contributor
There was a problem hiding this comment.
If vector_db_cls() throws, llm and embedding are created but the except re-raises as LegacyExecutorError. The caller's try/finally with _flush_per_prompt_metrics (which calls
vector_db.close() and flushes usage) is never entered because the destructuring assignment llm, embedding, vector_db = self._init_llm_and_retrieval(...) fails. The llm._pending_usage
records from construction-time calls (if any) are silently dropped.
Suggestion: Flush llm.flush_pending_usage() in the except block before re-raising, and extend self._usage_records with whatever was collected.
| ws_room = cb.get("ws_room", "") | ||
| ws_event = cb.get("ws_event", "") | ||
|
|
||
| api = _get_extraction_client() |
Contributor
There was a problem hiding this comment.
If BaseAPIClient manages HTTP sessions, ExtractionAPIClient leaks connections on every Celery task invocation.
Compare with UsageAPIClient in workers/executor/tasks.py which correctly uses with:
with UsageAPIClient(config) as usage_client:
| ws_room = cb.get("ws_room", "") | ||
| ws_event = cb.get("ws_event", "") | ||
|
|
||
| api = _get_extraction_client() |
Contributor
There was a problem hiding this comment.
If BaseAPIClient manages HTTP sessions, ExtractionAPIClient leaks connections on every Celery task invocation.
Compare with UsageAPIClient in workers/executor/tasks.py which correctly uses with:
with UsageAPIClient(config) as usage_client:
| if embedding.callback_manager is not None: | ||
| for handler in embedding.callback_manager.handlers: | ||
| if hasattr(handler, "flush_pending_usage"): | ||
| self._usage_records.extend(handler.flush_pending_usage()) |
Contributor
There was a problem hiding this comment.
If one handler's flush_pending_usage() throws, all subsequent handlers' records are lost. This loop should be guarded per-handler:
for handler in embedding.callback_manager.handlers:
if hasattr(handler, "flush_pending_usage"):
try:
self._usage_records.extend(handler.flush_pending_usage())
except Exception:
logger.warning("Failed to flush usage from handler", exc_info=True)
| ) | ||
| # Inside the try so a missing/renamed attribute on the outcome | ||
| # (plugin contract drift) hits the same graceful-degrade branch. | ||
| self._usage_records.extend(outcome.usage_records) |
Contributor
There was a problem hiding this comment.
If run_with_metrics makes an LLM call internally and then crashes (e.g., on post-processing the result), the lookup LLM's token consumption happened but the usage records are trapped
inside the plugin and never flushed. This is billing data loss.
Suggestion: The plugin protocol should expose a way to recover partial usage even on failure, or the executor should pass its _usage_records list to the plugin so records are appended in
real-time.
| self._metrics = new_metrics | ||
|
|
||
| # Stamp timing onto the most recent pending usage record | ||
| pending = getattr(self, "_pending_usage", []) |
Contributor
There was a problem hiding this comment.
This always stamps the last record in _pending_usage. If a decorated method makes multiple LLM calls (e.g., retry logic inside a single complete() call), only the last call gets timing.
Earlier calls get execution_time_ms = None in the DB, which is misleading — it looks like the data wasn't captured rather than being an intermediate call.
| # unused choice entries on OSS. | ||
|
|
||
| LLM_USAGE_REASON_CHOICES: list[tuple[str, str]] = [ | ||
| ("extraction", "Extraction"), |
Contributor
There was a problem hiding this comment.
Any code that references LLMUsageReason.EXTRACTION (the enum member) will now fail with ImportError or AttributeError. If any serializer, filter, or admin code references the old enum,
it's broken.
Suggestion: Grep the codebase for LLMUsageReason usages and verify none remain. Consider keeping the enum class and deriving the choices list from it.
| } | ||
| const useDefaultProfile = activeKey === "0" && !isSimplePromptStudio; | ||
| const { output, enriched } = useDefaultProfile | ||
| ? buildDefaultProfileOutputs(data) |
Contributor
There was a problem hiding this comment.
buildDefaultProfileOutputs, buildPerPromptOutput, and buildSelectedProfileOutputs are defined inside the component body and recreated on every render. They capture selectedProfile,
displayPromptResult, and other values via closure. Since they're only used inside fetchCombinedOutput (an async callback), they should be wrapped in useCallback or moved inside the
fetchCombinedOutput body to make the lifecycle clearer.
|
|
||
| // Extend Prompt Studio active state to include /lookups paths | ||
| if (lookupStudioEnabled && isUnstract) { | ||
| const psItem = data[0]?.subMenu?.find((el) => el.id === 1.1); |
Contributor
There was a problem hiding this comment.
This directly mutates the data array's child object. If data comes from state or a shared reference, this mutation can cause stale renders or bugs in other components reading the same
data. Use spread/clone instead.
athul-rs
reviewed
Apr 30, 2026
Contributor
athul-rs
left a comment
athul-rs
left a comment
There was a problem hiding this comment.
A few nits from a follow-up pass — non-blocking.
Comment on lines
+925
to
+926
| serializer.validated_data.get(ProfileManagerKeys.PROMPT_STUDIO_TOOL) | ||
| or self.get_object() |
Contributor
There was a problem hiding this comment.
nit: validated_data.get(...) or self.get_object() only patches the local variable. serializer.validated_data still lacks prompt_studio_tool, and ProfileManager.prompt_studio_tool is null=True, so if the field is ever omitted from the request perform_create() could persist a row with prompt_studio_tool=NULL — orphaned from filter(prompt_studio_tool=...) queries. Either make the serializer field required, or write the resolved tool back into validated_data before perform_create.
| ) | ||
| # Inside the try so a missing/renamed attribute on the outcome | ||
| # (plugin contract drift) hits the same graceful-degrade branch. | ||
| self._usage_records.extend(outcome.usage_records) |
Contributor
There was a problem hiding this comment.
nit: self._usage_records.extend(outcome.usage_records) lives inside the try, so any LLM usage already accumulated by the lookup plugin is dropped if run_with_metrics raises after partial work — billing loss. _run_challenge_if_enabled in this same file already establishes the right pattern (flush in finally). Consider moving the extend into a finally block guarded by if outcome is not None.
| ) | ||
| return | ||
| if output_type != PSKeys.JSON: | ||
| # The pre-refactor behaviour fired the webhook regardless of |
Contributor
There was a problem hiding this comment.
nit: this comment is inaccurate — pre-refactor, the webhook was invoked from inside AnswerPromptService.handle_json, which the executor only called from the elif output_type == PSKeys.JSON: branch. So webhooks were never fired for non-JSON outputs. The new restriction matches old behaviour rather than tightening it; worth rewording so future readers don't assume a regression surface.
Contributor
|
nit (couldn't anchor inline — line is unchanged context): the docstring at |
athul-rs
approved these changes
Apr 30, 2026
Contributor
athul-rs
left a comment
athul-rs
left a comment
There was a problem hiding this comment.
@chandrasekharan-zipstack please chekc the NIt comments added
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
OSS-side glue for Lookups V2 — a cloud-only Prompt Studio feature that runs an LLM-powered post-extraction enrichment over a prompt's structured output. This PR is the OSS half: a try-import bridge module, executor hook for non-SP runs, IDE wiring for the Raw/Enriched tabs and gating modals, and a few small schema extensions.
The cloud-side PR (where the feature actually lives) is on the
unstract-cloudrepo, same branch name, same ticket: https://github.com/Zipstack/unstract-cloud/pull/1463.Why
The cloud plugin needs to integrate with extraction (executor), with the per-prompt output flow (
prompt_studio_helper,prompt_studio_output_manager_v2), with export (prompt_studio_registry_helper), and with the IDE UI (Raw/Enriched tabs, gate modal, dirty-seed banner). All of that has to live in OSS — but OSS must never import cloud schema directly. This PR is that boundary.How
backend/prompt_studio/lookup_utils.pytry-imports the cloud plugin and exposes ~10 stable, opaque functions. Every function is a no-op in OSS. Callers don't know cloud key names or model shapes.legacy_executor._run_lookup_enrichmentresolves the worker plugin viaExecutorPluginLoader.get("lookup-enrichment")and runs after each prompt's extraction. Skips with a workflow-log entry when the source prompt extracted nothing.fetch_response/bulk_fetch_response, scoped to the prompt set being run.lookup-validationaction onCustomToolViewreturns the validation buckets that the FE modal renders.CustomTool.last_exported_at(new field) plusis_lookup_dirtyincheck_deployment_usagepowers the "needs re-export" banner.Usagemodel adds a few cloud-extensible enum choices (reference_type,llm_usage_reason="lookup",status,error_message,execution_time_ms).Reviewer aid
A condensed one-page primer for human reviewers — system shape, data model, execution diagrams, gates, smoke checklist — is at
lookups-v2-reviewer-refresher.html(attached, same file on both PRs). Read it before opening the diff.The full KB lives in the reviewer's vault at
~/Documents/Obsidian Vault/zipstuff/lookups/.Main files to review (sorted by significance)
The bridge — central to the OSS↔cloud contract
backend/prompt_studio/lookup_utils.py— every cloud call funnels through here. Tiny, opaque, no-op in OSS.Executor — non-SP entry
2.
workers/executor/executors/legacy_executor.py—_run_lookup_enrichmentdispatches the worker plugin; new "skip-with-log" branch when extraction yieldsNone.3.
workers/executor/executor_tool_shim.py—stream_logrouting (confirms the user-facing path for the new skip log).Prompt Studio integration
4.
backend/prompt_studio/prompt_studio_core_v2/views.py— multi-var pre-flight,lookup-validationaction,is_lookup_dirtyincheck_deployment_usage.5.
backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py— three call sites that attachlookup_configto per-prompt output / tool settings.6.
backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py— export-time validation + per-prompt config attach.7.
backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py— persists lookup output via the bridge after enrichment.8.
backend/prompt_studio/prompt_studio_core_v2/models.py— addsCustomTool.last_exported_at.Usage / observability
9.
backend/usage_v2/models.py— extensibleLLM_USAGE_REASON_CHOICES/REFERENCE_TYPE_CHOICESvia try-import; new fields for lookup attribution.Frontend wiring
10.
frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx— Raw vs Enriched tabs.11.
frontend/src/components/custom-tools/header/Header.jsx— wiresuseLookupExportGate(try-imported, no-op in OSS).12.
frontend/src/components/custom-tools/tool-ide/ToolIde.jsx— wiresuseLookupDirtySeed.13.
frontend/src/components/custom-tools/combined-output/{CombinedOutput,JsonView}.jsx— combined-output enrichment passthrough.Can this PR break any existing features?
No production breakage expected. All cloud calls flow through
lookup_utilsand degrade toNone/[]/{ok: True}in OSS, so existing flows behave identically when no lookup is assigned or when the cloud plugin isn't installed.Watch areas:
_run_lookup_enrichmentruns after every prompt in the non-SP path. Whenlookup_configisNone(no assignment) it returns immediately — should be invisible to non-lookup users.metadata["lookup_outputs"]is a new key on prompt output. FE consumers reflecting the fullmetadatashape need to ignore unknown keys.Usagetable now sees rows withreference_type="lookup_version"; dashboards should treat these as additional rows.legacy_executor.pyskip log usesshim.stream_logwhich routes throughLogPublisherto the workflow log UI — verified path inexecutor_tool_shim.py.Database Migrations
usage_v2/0005_usage_reason_ref_created_idxusage_v2/0006_alter_usage_llm_usage_reason_and_moreprompt_studio_core_v2/0007_customtool_last_exported_atEnv Config
None new.
Relevant Docs
~/Documents/Obsidian Vault/zipstuff/lookups/.lookups-v2-reviewer-refresher.html.Related Issues or PRs
Dependencies Versions
No new dependencies.
Notes on Testing
Smoke checklist in the reviewer refresher. The OSS-side hooks and bridge can be exercised against a cloud build; in pure OSS, every lookup-related code path is exercised but no-ops out via
LOOKUPS_AVAILABLE = False.Screenshots
UI screenshots live on the cloud PR (the OSS-side FE changes are wiring only).
Checklist
I have read and understood the Contribution Guidelines.