Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
59 commits
Select commit Hold shift + click to select a range
36463bb
UN-2946 [FEAT] Add lightweight list serializer for Prompt Studio and …
chandrasekharan-zipstack Apr 1, 2026
9300a4d
UN-2946 [FEAT] Add Look-Ups plugin integration in sidebar nav and routes
chandrasekharan-zipstack Apr 1, 2026
e36562d
UN-2946 [FIX] Use .get() fallback for prompt_studio_tool in create_pr…
chandrasekharan-zipstack Apr 2, 2026
4765e99
UN-2946 [FEAT] Add Lookups V2 OSS integration hooks for post-extracti…
chandrasekharan-zipstack Apr 3, 2026
68e394b
UN-2946 Removed unnecessary gitignore
chandrasekharan-zipstack Apr 3, 2026
059ceed
UN-2946 [REFACTOR] Deduplicate lookup config helper and add lookup us…
chandrasekharan-zipstack Apr 3, 2026
7699441
UN-2946 [FEAT] Add get_last_usage() to SDK1 LLM for token tracking
chandrasekharan-zipstack Apr 5, 2026
6e36896
UN-2946 [REFACTOR] Split post-extraction pipeline into lookup and web…
chandrasekharan-zipstack Apr 5, 2026
5b8c06d
UN-2946 [FEAT] Generic async extraction callbacks and WebSocket trans…
chandrasekharan-zipstack Apr 5, 2026
df49569
Reduce success notification duration from 2s to 1s for less intrusive UX
chandrasekharan-zipstack Apr 5, 2026
5e39f70
Revert "Reduce success notification duration from 2s to 1s for less i…
chandrasekharan-zipstack Apr 5, 2026
e4c023e
UN-2946 [REFACTOR] Pluggable lookup export validation
chandrasekharan-zipstack Apr 6, 2026
9bf072e
UN-2946 [REFACTOR] Lookups V2 review cleanup
chandrasekharan-zipstack Apr 7, 2026
ae4ba0a
UN-2946 [UI] Replace sidebar popover with in-page tabs for Lookups
chandrasekharan-zipstack Apr 7, 2026
8aa6758
UN-2946 [FEAT] Deferred batch usage tracking with operation metrics
chandrasekharan-zipstack Apr 8, 2026
0d29c9d
UN-2946 [FEAT] Add plugin hook for lookup output enrichment in serial…
chandrasekharan-zipstack Apr 10, 2026
17ffe38
UN-2946 [FEAT] Add lookup usage observability with error handling and…
chandrasekharan-zipstack Apr 12, 2026
6aad216
UN-2946 [FEAT] Support enriched output copy and lookup drawer plugin …
chandrasekharan-zipstack Apr 15, 2026
dafefa9
UN-2946 [FEAT] Lookup export validation gate, raw-latest helper, and …
chandrasekharan-zipstack Apr 21, 2026
3cf18c9
UN-2946 [UI] Fix combined output pill overlap and preserve Look-Ups t…
chandrasekharan-zipstack Apr 21, 2026
d63767b
UN-2946 [FEAT] Add last_exported_at and wire lookup staleness bridge
chandrasekharan-zipstack Apr 21, 2026
eced338
UN-2946 [FEAT] Stream lookup enrichment failures to workflow logs
chandrasekharan-zipstack Apr 21, 2026
58f2dbd
UN-2946 [UI] Wire lookup dirty-seed and export gate into ToolIde
chandrasekharan-zipstack Apr 21, 2026
3bddd7c
UN-2946 [FEAT] Share lookup test wrapper + generic ExecutionLogs back…
chandrasekharan-zipstack Apr 23, 2026
1bf0b03
UN-2946 [REFACTOR] Round-2 review fixes — OSS side
chandrasekharan-zipstack Apr 23, 2026
baaf203
UN-2946 [REFACTOR] Round-3 review fixes — OSS side
chandrasekharan-zipstack Apr 24, 2026
1ce0d6e
UN-2946 [FEAT] Reference prompts by UUID + missing-file gate — OSS side
chandrasekharan-zipstack Apr 24, 2026
57b8734
Merge branch 'main' into feat/lookups-v2
chandrasekharan-zipstack Apr 26, 2026
07dd880
UN-2946 [FIX] Surface skipped lookups when source prompt has no value
chandrasekharan-zipstack Apr 26, 2026
288df58
UN-2946 [REFACTOR] Address Sonar findings on lookups V2 PR
chandrasekharan-zipstack Apr 26, 2026
8e5491c
UN-2946 [FIX] Preserve usage records on executor failure paths
chandrasekharan-zipstack Apr 27, 2026
a23c7a6
UN-2946 [REFACTOR] Squash usage_v2 migrations 3 → 2 for lookups V2
chandrasekharan-zipstack Apr 27, 2026
0be7a0f
UN-2946 [REFACTOR] Update OSS↔cloud lookup bridge for app rename
chandrasekharan-zipstack Apr 27, 2026
a4679b1
UN-2946 [FIX] Address greptile review on lookups V2 OSS
chandrasekharan-zipstack Apr 27, 2026
5e8a223
UN-2946 [FIX] Static usage choices to fix migration drift in OSS CI
chandrasekharan-zipstack Apr 28, 2026
b93abd0
Merge remote-tracking branch 'origin/main' into feat/lookups-v2
chandrasekharan-zipstack Apr 28, 2026
f9d3844
UN-2946 [FEAT] Block enforce_type switch via lookup plugin gate
chandrasekharan-zipstack Apr 28, 2026
d898697
UN-2946 [FIX] Harden billing/usage paths against silent drops
chandrasekharan-zipstack Apr 28, 2026
9ef8c80
UN-2946 [FIX] Cross-cutting hygiene around lookup enrichment & webhook
chandrasekharan-zipstack Apr 28, 2026
29a7050
UN-2946 [FIX] Tighten Usage choices & lookup_utils contracts
chandrasekharan-zipstack Apr 28, 2026
993ae95
UN-2946 [PERF] Push Combined Output queries into SQL
chandrasekharan-zipstack Apr 28, 2026
cc544e4
UN-2946 [FIX] Frontend & callback hygiene around lookup hooks
chandrasekharan-zipstack Apr 28, 2026
301dc9b
UN-2946 [FIX] Skip webhook on JSON parse failure & re-include compose…
chandrasekharan-zipstack Apr 28, 2026
2fd9cf5
UN-2946 [FIX] Return 400 for missing tool_id (was 500)
chandrasekharan-zipstack Apr 28, 2026
ba6c32a
UN-2946 [FIX] Address remaining post-disposition review comments (OSS)
chandrasekharan-zipstack Apr 28, 2026
4aea7e4
UN-2946 [DOCS] Tighten comments across lookups V2 OSS surface
chandrasekharan-zipstack May 5, 2026
21caca2
UN-2946 [REFACTOR] Drop unused token_count param from ExtractionAPICl…
chandrasekharan-zipstack May 5, 2026
7b4438a
UN-2946 [REFACTOR] DRF-ify Usage internal batch endpoint + squash mig…
chandrasekharan-zipstack May 5, 2026
1104cf4
UN-2946 [REFACTOR] Harden SDK / worker billing path + extract lookup …
chandrasekharan-zipstack May 5, 2026
dab3ffa
UN-2946 [FIX] Lookup-related FE polish + DRF error envelope
chandrasekharan-zipstack May 5, 2026
4b9e0fd
UN-2946 [FIX] Type run_id / reference_id as UUIDField in batch serial…
chandrasekharan-zipstack May 5, 2026
4bbf2de
UN-2946 [DOCS] Tighten extraction_complete docstring
chandrasekharan-zipstack May 5, 2026
3734332
UN-2946 [FIX] Pre-bind validated_file_execution_id in usage client
chandrasekharan-zipstack May 6, 2026
23aa52c
Merge remote-tracking branch 'origin/main' into feat/lookups-v2
chandrasekharan-zipstack May 6, 2026
e67be15
[CHORE] Ignore .pi/ tooling directory
chandrasekharan-zipstack May 7, 2026
a4de618
UN-3494 [REFACTOR] Replace polymorphic Usage attribution with typed c…
chandrasekharan-zipstack May 8, 2026
c935d58
UN-3494 [REVIEW] Idempotent hook registration + lookup-usage partial …
chandrasekharan-zipstack May 8, 2026
a4969ff
UN-3494 [FIX] Forward answer-step metadata in structure pipeline
chandrasekharan-zipstack May 8, 2026
c53c51d
UN-3494 [REFACTOR] Unify usage-records carrier and propagate executio…
chandrasekharan-zipstack May 8, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -653,7 +653,11 @@ docker/*.env
!docker/sample*.env
docker/public_tools.json
docker/proxy_overrides.yaml
docker/compose.override.yaml
docker/compose.*.yaml
Comment thread
chandrasekharan-zipstack marked this conversation as resolved.
# ``docker/compose.debug.yaml`` is checked-in tooling — keep it out of the
# broader ``compose.*.yaml`` ignore so a delete + recreate doesn't make it
# look untracked, and so teammates can spot it.
!docker/compose.debug.yaml
docker/workflow_data/

# Tool development
Expand Down Expand Up @@ -696,6 +700,12 @@ CLAUDE.md
CONTRIBUTION_GUIDE.md
.mcp.json

# Codex
AGENTS.md

# Pi
.pi/

# Windsurf
.qodo
.windsurfrules
Expand Down
171 changes: 171 additions & 0 deletions backend/prompt_studio/lookup_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,171 @@
"""Shared utility for lookup operations. No-ops in OSS.

Only the absence of ``pluggable_apps.lookups`` itself is treated as
"cloud not installed"; an ImportError from a transitive dependency
re-raises so we don't silently degrade to a no-op on a real bug.
"""

import logging
from typing import Any

logger = logging.getLogger(__name__)

_CLOUD_LOOKUP_MODULES = {
# OSS images lack the parent ``pluggable_apps`` package, so include it.
"pluggable_apps",
"pluggable_apps.lookups",
"pluggable_apps.lookups.execution",
"pluggable_apps.lookups.output_enrichment",
"pluggable_apps.lookups.staleness",
"pluggable_apps.lookups.validation",
"pluggable_apps.lookups.models",
}

try:
Comment thread
chandrasekharan-zipstack marked this conversation as resolved.
from pluggable_apps.lookups import execution as _execution
from pluggable_apps.lookups import output_enrichment as _output_enrichment
from pluggable_apps.lookups import staleness as _staleness
from pluggable_apps.lookups import validation as _validation
from pluggable_apps.lookups.models import LookupOutputResult as _LookupOutputResult

LOOKUPS_AVAILABLE = True
except ImportError as e:
if e.name not in _CLOUD_LOOKUP_MODULES:
raise
LOOKUPS_AVAILABLE = False
Comment thread
coderabbitai[bot] marked this conversation as resolved.


def get_lookup_config(prompt) -> dict | None:
"""Return lookup config for a prompt, or None if lookups are unavailable."""
if not LOOKUPS_AVAILABLE:
return None
return _execution.build_lookup_config_for_prompt(prompt)


def get_lookup_configs_for_tool(tool, prompts=None) -> list[dict] | None:
"""Return lookup configs for a tool (single pass), or None in OSS.

``prompts`` scopes validation to the run's prompts so unrelated
incomplete assignments on the tool don't block it.
"""
if not LOOKUPS_AVAILABLE:
return None
return _execution.build_lookup_configs_for_tool(tool, prompts=prompts)


def get_multi_var_lookups_for_tool(tool, prompt_ids=None) -> list[str]:
"""Return names of multi-variable lookups linked to the tool, [] in OSS.

``prompt_ids`` scopes the check so a run is only blocked when the
multi-var lookup is actually used by it.
"""
if not LOOKUPS_AVAILABLE:
return []
_, names = _execution.has_multi_var_lookups(tool, prompt_ids=prompt_ids)
return names


def persist_lookup_output(prompt_output, prompt_lookup: dict) -> None:
"""Persist lookup enrichment result. No-op in OSS."""
if not LOOKUPS_AVAILABLE:
return
lookup_meta = prompt_lookup.get("meta", {})
lookup_id = lookup_meta.get("lookup_id")
if not lookup_id:
return
defaults = {
"lookup_definition_id": lookup_id,
"output": prompt_lookup.get("enriched", ""),
}
version_id = lookup_meta.get("version_id")
if version_id:
defaults["version_id"] = version_id
_LookupOutputResult.objects.update_or_create(
prompt_output=prompt_output,
defaults=defaults,
)


def enrich_prompt_output(prompt_output, data: dict) -> dict:
"""Let cloud plugins enrich serialized prompt output with lookup data.

No-op in OSS.
"""
if not LOOKUPS_AVAILABLE:
return data
return _output_enrichment.enrich_with_lookup_output(prompt_output, data)


def validate_lookups_for_export(prompts) -> tuple[dict, str | None]:
"""Validate lookup assignments before export. Returns ({}, None) in OSS."""
if not LOOKUPS_AVAILABLE:
return {}, None
return _validation.validate_lookups_for_export(prompts)


def get_latest_lookup_mutation_for_tool(tool):
"""Max ``modified_at`` across lookup-related records linked to the tool
(version, reference file, assignment) — feeds the staleness banner.
None if unavailable or nothing linked.
"""
if not LOOKUPS_AVAILABLE:
return None
return _staleness.get_latest_lookup_mutation_for_tool(tool)


def get_original_value_if_enriched(
metadata: dict, prompt_key: str
) -> tuple[Any, dict] | None:
"""Return ``(original_value, prompt_lookup_dict)`` if ``prompt_key`` was
enriched, or ``None`` otherwise.

Pure metadata-shape check — safe to call even when LOOKUPS_AVAILABLE
is False (returns None because the shape won't match).
"""
if not isinstance(metadata, dict):
return None
lookup_outputs = metadata.get("lookup_outputs") or {}
prompt_lookup = lookup_outputs.get(prompt_key)
if isinstance(prompt_lookup, dict) and "original" in prompt_lookup:
return prompt_lookup.get("original"), prompt_lookup
return None


def attach_combined_output_enrichment(result: dict, enriched_by_key: dict) -> None:
"""Stamp the combined-output payload with enriched-output metadata.

Comment thread
chandrasekharan-zipstack marked this conversation as resolved.
Key name stays cloud-side so the FE-plugin shape can evolve without
coordinating with OSS.
"""
if not LOOKUPS_AVAILABLE:
return
_output_enrichment.attach_combined_output_enrichment(result, enriched_by_key)


def extract_prompt_output_enrichment(item) -> dict | None:
"""Pick enriched-output data off a serialized prompt-output row.

Returns a plugin-opaque dict (FE-only) or None when no enrichment
is present / plugin missing.
"""
if not LOOKUPS_AVAILABLE:
return None
return _output_enrichment.extract_prompt_output_enrichment(item)


def get_lookup_validation_for_tool(tool) -> dict:
"""Pre-emptive lookup validation for FE Export / Deploy gating.

Returns an "always ok" payload in OSS so the FE gate is a no-op.
"""
if not LOOKUPS_AVAILABLE:
return {
"ok": True,
"draft_lookups": [],
"multi_var_lookups": [],
"incomplete_lookups": [],
"single_pass_enabled": bool(
getattr(tool, "single_pass_extraction_mode", False)
),
}
return _validation.get_lookup_validation_for_tool(tool)
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
# Generated by Django 4.2.1 on 2026-04-21 20:20

from django.db import migrations, models


class Migration(migrations.Migration):
dependencies = [
("prompt_studio_core_v2", "0006_add_custom_data_to_customtool"),
]

operations = [
migrations.AddField(
model_name="customtool",
name="last_exported_at",
field=models.DateTimeField(
blank=True,
db_comment="Timestamp of the last successful export; NULL if never exported since the field was introduced.",
null=True,
),
),
]
9 changes: 9 additions & 0 deletions backend/prompt_studio/prompt_studio_core_v2/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,15 @@ class CustomTool(DefaultOrganizationMixin, BaseModel):
db_comment="Flag to share this custom tool with all users in the organization",
)

# NULL on pre-feature tools; populated on first successful export.
# Drives staleness checks (e.g. lookup-change banner) without requiring
# a data backfill.
last_exported_at = models.DateTimeField(
null=True,
blank=True,
db_comment="Timestamp of the last successful export; NULL if never exported since the field was introduced.",
)

objects = CustomToolModelManager()

def delete(self, organization_id=None, *args, **kwargs):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,17 @@
from django.db import transaction
from django.db.models.manager import BaseManager
from plugins import get_plugin
from rest_framework.exceptions import APIException
from rest_framework.request import Request
from utils.file_storage.constants import FileStorageKeys
from utils.file_storage.helpers.prompt_studio_file_helper import PromptStudioFileHelper
from utils.local_context import StateStore

from backend.celery_service import app as celery_app
from prompt_studio.lookup_utils import (
get_lookup_config,
get_lookup_configs_for_tool,
)
from prompt_studio.prompt_profile_manager_v2.models import ProfileManager
from prompt_studio.prompt_profile_manager_v2.profile_manager_helper import (
ProfileManagerHelper,
Expand Down Expand Up @@ -387,6 +392,9 @@ def _build_prompt_output(
if webhook_enabled:
output[TSPKeys.POSTPROCESSING_WEBHOOK_URL] = webhook_url

if lookup_config := get_lookup_config(prompt):
output["lookup_config"] = lookup_config

output[TSPKeys.EVAL_SETTINGS] = {}
output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_EVALUATE] = prompt.evaluate
output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_MONITOR_LLM] = [monitor_llm]
Expand Down Expand Up @@ -798,6 +806,9 @@ def build_fetch_response_payload(
if webhook_enabled:
output[TSPKeys.POSTPROCESSING_WEBHOOK_URL] = webhook_url

if lookup_config := get_lookup_config(prompt):
output["lookup_config"] = lookup_config

output[TSPKeys.EVAL_SETTINGS] = {}
output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_EVALUATE] = prompt.evaluate
output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_MONITOR_LLM] = [monitor_llm]
Expand Down Expand Up @@ -1166,6 +1177,10 @@ def build_single_pass_payload(
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
Comment thread
coderabbitai[bot] marked this conversation as resolved.

for p in prompts:
if not p.prompt:
raise EmptyPromptError()
Expand Down Expand Up @@ -1607,6 +1622,9 @@ def _execute_single_prompt(
is_single_pass=False,
profile_manager_id=profile_manager_id,
)
except APIException:
# Validation responses are user-facing; DRF renders them as-is.
raise
except Exception as e:
logger.error(
f"[{tool.tool_id}] Error while fetching response for "
Expand Down Expand Up @@ -1672,6 +1690,9 @@ def _execute_prompts_in_single_pass(
document_id=document_id,
is_single_pass=True,
)
except APIException:
# Validation responses are user-facing; DRF renders them as-is.
raise
except Exception as e:
logger.error(
f"[{tool.tool_id}] Error while fetching single pass response: {e}"
Expand Down Expand Up @@ -1911,6 +1932,8 @@ def _fetch_response(
output[TSPKeys.ENABLE_POSTPROCESSING_WEBHOOK] = webhook_enabled
if webhook_enabled:
output[TSPKeys.POSTPROCESSING_WEBHOOK_URL] = webhook_url
if lookup_config := get_lookup_config(prompt):
output["lookup_config"] = lookup_config
# Eval settings for the prompt
output[TSPKeys.EVAL_SETTINGS] = {}
output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_EVALUATE] = prompt.evaluate
Expand Down
9 changes: 9 additions & 0 deletions backend/prompt_studio/prompt_studio_core_v2/urls.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,10 @@

prompt_studio_task_status = PromptStudioCoreView.as_view({"get": "task_status"})

prompt_studio_lookup_validation = PromptStudioCoreView.as_view(
{"get": "lookup_validation"}
)


urlpatterns = format_suffix_patterns(
[
Expand Down Expand Up @@ -165,5 +169,10 @@
prompt_studio_task_status,
name="prompt-studio-task-status",
),
path(
"prompt-studio/<uuid:pk>/lookup-validation/",
prompt_studio_lookup_validation,
name="prompt-studio-lookup-validation",
),
]
)
Loading