Ben/category enrichment - #32
Merged
Merged
Conversation
…2sql-onboarding into ben/category-enrichment
…ding into ben/category-enrichment
|
Important Review available on request
Reviews should be triggered manually for repositories with fewer than 10 stars. Select Trigger review above or comment ⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📝 WalkthroughWalkthroughThe refiner workflow now separates enrichment, agent reasoning, and Trino execution. The change adds category-value extraction, hybrid search, SQL transformation, vector-backed storage, expanded tests, updated runtime configuration, and an execution-inspection CLI. ChangesAgent workflow
Category enrichment
Category storage
Validation and tooling
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 2📝 Generate docstrings 💡
⚔️ Resolve merge conflicts 💡
🧪 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 |
Collaborator
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
yuvalkh
commented
Aug 6, 2026
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 60
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@agent/src/agent/config.py`:
- Around line 71-72: Validate the refiner loop governed by
MAX_REFINER_ITERATIONS and add or reuse a request timeout or cost guard that
bounds cumulative LLM calls, Trino executions, and connection occupancy across
all attempts. Ensure the guard applies to the graph’s agent_node and
trino_exec_node flow while preserving the configured iteration limit.
- Around line 28-29: Replace the hard-coded defaults for JEEN_API_KEY and
JEEN_METADATA_MCP_KEY in the configuration definitions with empty values,
preserving the documented skip-fetch behavior; load both keys from their
environment variables instead, and rotate the exposed credentials outside the
source change.
- Line 63: Update the LANGFUSE_PROMPT_LOC_EXTRACTOR configuration used by
LocationExtractorAgent._build_prompt() to reference the intended
location-extractor prompt identifier rather than the regular text2sql/extractor
prompt, ensuring it returns the Hebrew-name-to-standard-location JSON map
required by _parse_llm_json().
In `@agent/src/agent/langfuse_client.py`:
- Around line 7-9: Restrict the InsecureRequestWarning suppression near the
Langfuse client initialization to an explicit development-only setting, leaving
TLS verification and warnings enabled in production and other clients. Remove
any unconditional process-wide suppression, and ensure the directly imported
urllib3 dependency is declared by the agent package.
- Around line 22-33: Update _safe_update_current_span and
_safe_get_current_trace_id to validate the current span’s context and recording
capability instead of comparing it only with INVALID_SPAN. Skip Langfuse
operations for invalid, remote, or non-recording spans, while preserving the
existing delegation for valid recording spans.
In `@agent/src/agent/nodes/finalizer.py`:
- Around line 81-84: Update the prompt-name assignment in the finalizer flow to
use direct access via settings.LANGFUSE_PROMPT_FINALIZER instead of getattr with
a fallback, matching the existing query_builder.py access pattern. Leave the
subsequent langfuse_client.get_prompt call unchanged.
- Around line 35-36: Update get_esca_preview to log the caught exception through
the module logger, then return a neutral preview message that excludes the
exception text and any internal storage details; keep finalizer_node’s existing
sql_results flow unchanged.
In `@agent/src/agent/nodes/refiner.py`:
- Around line 234-241: Replace the regex-based table rewriting in the refiner’s
table-mapping loop with a single sqlglot parse and AST traversal that updates
only table-reference nodes, preserving string literals, column names, and
aliases. Reuse the repository’s existing sqlglot conventions from
sql_transformer.py, regenerate the SQL from the transformed AST, and add the
requested string-literal regression test beside
test_trino_exec_table_alias_word_boundary.
- Around line 205-215: Guard the sql_query preprocessing in the refiner node so
missing or null state["sql_query"] does not reach the regex substitutions.
Update the flow around sql, locations_dict, and the subsequent table-name
replacements to handle the absent query before calling re.sub, while preserving
normal substitution behavior for valid SQL and ensuring the failure is recorded
through the existing Trino error-handling path.
- Around line 169-172: Update the prompt-state construction in the refiner flow
so last_result_success is true only when error_msg is empty and a query
execution has occurred. Distinguish the initial no-execution state from
successful execution using the existing execution/result state symbols, and
ensure satisfaction failures represented by error_msg produce a false success
flag while preserving last_result_error.
- Around line 27-31: Update build_refiner_schema_context to fall back to
settings.REFINER_SCHEMA_CONTEXT_TABLES when runtime_flags lacks
REFINER_SCHEMA_CONTEXT_TABLES, and normalize string runtime values such as "8"
before applying the limit. Preserve the existing table_profiles slicing and JSON
serialization behavior, while ignoring invalid or non-positive values
consistently.
- Around line 176-196: Update the refiner return logic around clean_sql and the
sql_query field so QUERY_SATISFIED responses containing only a TRANSLATION block
preserve state.get("sql_query") instead of storing prose. Use the newly cleaned
SQL only when the response includes a SQL code block or a valid unfenced query,
while leaving the existing satisfaction and explanation handling unchanged.
- Around line 154-158: Update the schema-context tagging in the refiner node to
use Langfuse’s public attribute-propagation API, or wrap the existing private
call in try/except so SDK changes or Langfuse failures cannot escape the node.
Preserve the current trace-id guard and both schema-context and step tags.
In `@agent/src/agent/services/enrichment_orchestrator.py`:
- Line 169: Update the local annotation for llm in the orchestrator flow to use
the already imported BaseChatModel instead of the undefined ChatOpenAI type,
while leaving the get_orchestrator_llm() call and return behavior unchanged.
- Around line 129-142: Update the matching logic in the search-results
formatting loop to recognize filters whose value is a collection by matching the
individual result value against its members, while retaining scalar matching for
non-collection filters. Ensure matching_filter resolves the original filter so
its operator remains IN (or the appropriate multi-value operator) and
SQLTransformer receives the correct operator instead of the "=" fallback.
- Around line 194-214: Normalize the column name consistently in the ghost-value
validation around the enrichment-details loop: use the normalized form for the
direct search_results key and normalize k_col before comparing it with
tf.column. Preserve the existing candidate validation and warning behavior once
the pool is found.
In `@agent/src/agent/services/filter_extractor.py`:
- Line 77: Remove the unused cte_select_map declaration and its population pass,
or update resolve_col_ref to consult that map and eliminate its inline
projection walk. Ensure each scope’s projection expressions are traversed only
once while preserving existing column-reference resolution behavior.
- Around line 41-42: Update the extraction logic around sql_processed and value
construction so literal filter values containing “@” are restored to their
original form before being returned. Ensure the fix covers all extracted value
paths, including the logic around lines 287–296, while preserving the existing
source_table and original_expression handling; add coverage for a predicate such
as user_email = 'jane@corp.com'.
- Around line 27-28: The extract method is overly large and contains untestable
nested helpers, including duplicated literal parsing. Move extract_literal_val
and get_leaf_comparisons to module-level helpers, introduce a small helper class
to own table_alias_map and unnest_map so resolve_col_ref becomes a method, and
update extract to use these symbols while preserving existing behavior. Reuse
the shared module-level literal extraction helper from sql_transformer instead
of retaining duplicate logic.
- Around line 340-342: The broad exception handler in the filter extraction
method masks defects by converting all failures into an empty filter list.
Narrow the try/except to only the SQL parse and qualification steps, preserving
the existing parse-failure handling, and allow errors from the subsequent
extraction loop—including scope resolution and match-type mapping—to propagate
to callers and tests.
- Around line 340-342: The blanket exception handler in FilterExtractor.extract
currently converts internal failures into an empty filter list. In
agent/src/agent/services/filter_extractor.py lines 340-342, narrow handling to
the parse and qualify operations or return a distinct failure signal that
EnrichmentOrchestrator.enrich_query can distinguish from genuinely having no
filters; in agent/tests/test_filter_extractor.py lines 164-197, update
test_extract_no_filters and test_extract_ignore_column_to_column with caplog
assertions so swallowed exceptions cannot satisfy the empty-filter assertions.
- Around line 228-232: Update get_leaf_comparisons to handle exp.Not-wrapped
predicates instead of dropping them, and ensure SQLTransformer.transform_node
does not rewrite predicates under negation because it matches only by column and
operator. Propagate negation metadata or explicitly mark these predicates so
transformed plans skip negated filters, preserving correct semantics when
positive and negated comparisons coexist.
In `@agent/src/agent/services/hybrid_searcher.py`:
- Around line 210-220: Rename rerank_candidates to cap_candidates and update its
callers at the relevant search paths so the function’s name reflects that it
only truncates results. Replace the hard-coded 5 with a module-level constant or
existing setting for a tunable candidate limit, and update the docstring
accordingly; do not implement cross-encoder reranking.
- Around line 382-392: Extract the duplicated param.value normalization from the
task-queueing loop and output-mapping loop into a shared normalize_values(param)
helper. Update both the search-target construction and the s_val/cache-key
construction to use this helper, preserving list handling, None filtering,
string conversion, and LIKE percent stripping so both paths produce identical
values.
- Around line 406-415: Normalize the result-key column casing consistently in
HybridSearcher.search: lowercase param.source_column when constructing key,
cache_key, and both results assignments around the task/result handling flow.
Add a mixed-case source_column case such as "Order_Status" to
agent/tests/test_hybrid_searcher.py lines 156-197 and assert the returned key is
"order_status#@#active".
- Around line 288-292: Update unit_id_workflow’s fused candidate list before it
is returned so combined results are capped to the same five-item limit used by
search_workflow and consumed by EnrichmentOrchestrator.enrich_query. Preserve
the existing ranking/order while truncating only the final list, and revise the
nearby fetch comment to state the actual query_db_digits_match limit of 20
instead of 50.
- Around line 55-65: Update the table-resolution logic around the session query
and `.first()` so ambiguous one-part or two-part source_table matches are
detected instead of selecting an arbitrary row. Require a unique matching Table
record, returning no table ID when multiple rows match, while preserving the
existing three-part catalog/schema/name filtering and single-match behavior.
In `@agent/src/agent/services/location_extractor.py`:
- Around line 111-126: Update the successful-location processing around
_make_var_name so duplicate normalized identifiers are logged and disambiguated
with unique identifiers, ensuring each location retains its own polygon in
coords_dict. Store the final identifier for each location in var_names, and
build locations_dict_str from those stored values instead of calling
_make_var_name again.
- Around line 131-143: Update the prompt handling in the location extraction
flow to distinguish chat prompts from text prompts instead of checking only for
compile support. Route ChatPromptClient values through
ChatPromptTemplate.from_messages(...).format(...) so instruction_text is a
formatted string, while preserving compile(...) for TextPromptClient values and
avoiding Python repr storage of message dictionaries.
In `@agent/src/agent/services/sql_transformer.py`:
- Around line 36-40: The SQL transformation must avoid globally replacing @ and
$ characters, which corrupts existing dollar content and string literals. Update
the preprocessing near sql_processed and the corresponding reverse logic to
replace only extractor placeholder tokens matching the `@name_wkt`@ form, record
each exact substituted token, and restore only those recorded tokens after
transformation.
- Around line 177-183: Update make_literal so every candidate value from
refined_values is emitted with exp.Literal.string, without parsing integers or
floating-point values. Preserve the original text exactly, including values such
as "444", "52", and "1.0".
- Around line 159-169: Update the transformation logic around the
unmatched-value branch and refined-values handling so values without a matching
plan remain in the rebuilt predicate alongside transformed values. Preserve the
original operator semantics, especially for BETWEEN: do not force
target_operator to IN merely because two refined values exist; retain BETWEEN
when the source predicate is a range.
In `@agent/tests/conftest.py`:
- Around line 7-14: Update the environment setup around load_dotenv so
LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY, and LANGFUSE_BASE_URL each receive
their fallback independently when unset, rather than gating all defaults on
LANGFUSE_PUBLIC_KEY. Preserve values loaded from the developer’s .env file.
In `@agent/tests/refiner/test_refiner_e2e_mocked.py`:
- Around line 9-23: Remove the unused patch_graph_infrastructure helper, or make
it functional by correcting the publish_node_event target to
agent.services.enrichment_orchestrator.publish_node_event and ensuring all
returned patchers are started or used through a context manager. Keep the
helper’s external-I/O mocking behavior intact if it remains.
In `@agent/tests/refiner/test_refiner_e2e_real.py`:
- Around line 519-530: Relax the SQL token assertions in the live test around
the final_state checks: retain execution success, is_satisfied, and absence of
SQL Server constructs, but replace exact “limit 3”/“top 3”/literal-token checks
with validation that the executed result is bounded to three rows or fewer.
In `@agent/tests/test_enrichment_orchestrator.py`:
- Around line 506-533: The fixture schema and AgentSQLTable in the
car-registration test use dataverse.registered_cars while the SQL references
registered_cars. Update the schema key and AgentSQLTable.name to
registered_cars, matching test_extract_real_world_car_registrations and the FROM
clause so FilterExtractor resolves the table and columns correctly.
- Around line 67-94: The tests currently pass through
EnrichmentOrchestrator.enrich_query’s outer exception fallback instead of
exercising their intended paths. In
agent/tests/test_enrichment_orchestrator.py:67-94, mock HybridSearcher.search to
return {} and assert it was awaited, proving the fast path follows an empty
candidate set. In agent/tests/test_enrichment_orchestrator.py:145-161, configure
mock_llm.ainvoke as an AsyncMock so the inner handler reaches
parse_transformation_plan without a TypeError.
- Around line 117-118: Update each test setup for get_orchestrator_llm to
configure the structured-output mock through
mock_llm.with_structured_output.return_value, rather than calling
with_structured_output() during setup. Apply this at all repeated locations,
including the setups near lines 146, 198, 269, 381, and 502, so the production
arguments remain observable and verifiable.
In `@agent/tests/test_filter_extractor.py`:
- Around line 230-235: Update the unresolved-table test around
FilterExtractor.extract to assert f.source_table as part of the contract, using
"dataverse.unknown_table" or the extractor’s intended unresolved-table value
while preserving the existing assertions.
- Around line 164-197: Update test_extract_no_filters and
test_extract_ignore_column_to_column to accept the caplog fixture and assert
that FilterExtractor.extract emits no error log during each call, while
retaining the existing empty-list assertions. Use caplog to distinguish a
legitimate no-filter result from an exception swallowed by
FilterExtractor.extract.
In `@agent/tests/test_finalizer.py`:
- Line 4: Add tests in the finalizer test module covering get_esca_preview and
both relevant finalizer_node branches: enable ESCA_WRITE_ENABLED, provide
raw_data_ref, mock get_esca_client, and assert the serialized preview; then add
a separate test with no data reference asserting the “No data reference found.”
behavior.
In `@agent/tests/test_hybrid_searcher.py`:
- Around line 21-40: Strengthen test_find_table_id_qualified by capturing each
statement passed to mock_session.exec and asserting its compiled filter matches
the expected catalog.schema.table, schema.table, and table branches
respectively. Keep the existing ID assertions, but ensure each input verifies
the corresponding lookup constraint rather than relying on the shared mock row.
- Around line 156-197: Add a case-sensitivity test alongside
test_hybrid_searcher_routing using a mixed-case source_column such as
Order_Status, while keeping the table metadata lowercase. Mock the category
workflow and assert HybridSearcher.search returns the lowercase key
order_status#@#active expected by EnrichmentOrchestrator.enrich_query.
In `@agent/tests/test_sql_transformer.py`:
- Around line 372-446: The transformation plan currently matches filters without
table qualification, allowing identical predicates on joined tables to be
rewritten together. Propagate FilterExtractor’s source_table into
FilterTransformation, update SQLTransformer.transform_node matching to require
the table qualifier when available, and add a joined-table test with the same
column, operator, and literal verifying only the intended table’s predicate
changes.
- Around line 470-478: Add coverage in the SQLTransformer filter tests for an
all-digit refined categorical value such as "444", including preservation of
leading zeros, and update make_literal to retain string-literal rendering when
the source predicate compared against a string literal. Ensure numeric literals
remain numeric for genuinely numeric source predicates while string-column
comparisons emit quoted, unchanged values.
- Around line 78-79: Update the assertion following SQLTransformer.apply in the
no-change test to compare the complete expected predicate, not merely the
presence of “active”. Verify that changed_filter=False preserves the original
predicate exactly, including its operator and values.
In `@backend/alembic/versions/ed40dd0a57ad_add_large_category_values_table.py`:
- Around line 41-44: Update the migration creating the large_category_values
indexes to add an ANN index on embedding using the cosine operator class and an
accompanying trigram GIN index on value_text for similarity searches. Place the
ANN index creation after the initial bulk-ingestion step, while retaining the
existing lookup indexes.
- Around line 22-24: Update the migration’s upgrade() to create the pg_trgm
extension alongside vector before trigram queries run, and update downgrade() to
drop pg_trgm only when no other objects depend on it, preserving the existing
extension lifecycle.
- Line 9: Update the migration imports to explicitly load the
pgvector.sqlalchemy submodule before the vector type is accessed. Ensure the
VECTOR(dim=768) usages in the migration resolve through the imported SQLAlchemy
integration, including the usage near line 32.
In `@backend/app/infra_init.py`:
- Around line 1230-1239: Add ingest_large_category_values to the local
dependency import block inside _ensure_airlines_registered so the successful
result path can resolve it. Update the Vector ingestion failed logger.error call
to pass exc_info=True, preserving the existing context and exception message
while retaining traceback details for genuine failures.
In `@backend/app/services/category_ingestion.py`:
- Around line 50-58: Update the distinct-value query in the category ingestion
flow to escape embedded double quotes in col_name and safely quote table_fqn
using the project’s existing identifier-quoting approach. Add an explicit LIMIT
to bound rows collected into trino_values, preserving the existing null
filtering and ingestion behavior.
- Around line 80-100: Update the embedding loop in the category ingestion flow
to process values concurrently with a bounded thread pool or supported batch
embedding API, while preserving batch sizing and record construction. Track each
value whose get_query_embedding call fails or returns None, and log the total
failed count after processing so complete or partial embedding failures are
visible.
- Around line 108-114: Update the summary logging in the embedding pipeline
around the batch_size condition so the total_saved count is logged
unconditionally for both batched and unbatched ingestion. Preserve the existing
per-chunk logger.info message inside the batch path and ensure the final summary
uses total_saved and col_name.
- Around line 102-106: Update the batch persistence logic in the category
ingestion function around db_session.commit() to handle unique-constraint
conflicts safely, preferably by using PostgreSQL INSERT ... ON CONFLICT DO
NOTHING for uq_large_category_val. If commits remain, catch IntegrityError, roll
back the session for that batch, and continue processing subsequent batches and
columns without losing already committed records.
In `@backend/pyproject.toml`:
- Around line 28-31: Update the dependency declarations in pyproject.toml by
removing unused langchain, langchain-core, and langchain-openai packages; if any
are required, constrain them below version 1.0.0. Also remove sqlglot from the
backend dependencies unless backend enrichment code imports it, relocating it to
the appropriate dependency group when needed.
In `@core/src/core/models/models.py`:
- Around line 603-604: The embedding column needs an ANN index and bounded
semantic-search lookups. Add an HNSW or IVFFlat index for
large_category_values.embedding in the relevant migration, then update the
embedding query to filter by table_id and column_name before ranking results.
- Around line 606-607: The updated_at field in the model definition needs an
update-time hook in addition to its default_factory. Add SQLAlchemy onupdate
behavior to updated_at so every UPDATE, including category value re-ingestion,
refreshes the timestamp while preserving the existing creation-time default and
read-model behavior.
In `@core/src/core/trino.py`:
- Around line 11-14: Update the module-level urllib3 warning suppression in
trino.py to run only after settings is available and only when
settings.TRINO_VERIFY is false, avoiding global suppression for verified
connections. Add urllib3 as an explicit dependency in pyproject.toml rather than
relying on a transitive import.
In `@docker-compose.yml`:
- Around line 568-572: Remove the hardcoded LLM_API_KEY and JEEN_API_KEY values
from the Docker Compose environment, source them from an ignored secret
mechanism or Docker Compose secrets, and rotate both credentials immediately. In
agent configuration around the JEEN_API_KEY default, remove the identical
fallback so the rotated key cannot remain exposed; update consumers to require
the injected secret instead.
In `@scripts/inspect_flow.py`:
- Around line 139-142: Track the terminal outcome in run_flow, including
end_fail events and exceptions from agent_graph.astream(), and return a failure
status for both paths instead of printing a success result. Update main to
propagate that status as a non-zero process exit for failed non-interactive runs
while preserving successful and interactive behavior.
🪄 Autofix
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: ASSERTIVE
Plan: Pro Plus
Run ID: ef9970e0-996a-440e-a9c2-02c21d3323d3
⛔ Files ignored due to path filters (1)
backend/uv.lockis excluded by!**/*.lock
📒 Files selected for processing (45)
agent/.coverageagent/pyproject.tomlagent/src/agent/config.pyagent/src/agent/graph.pyagent/src/agent/langfuse_client.pyagent/src/agent/llm.pyagent/src/agent/nodes/finalizer.pyagent/src/agent/nodes/query_builder.pyagent/src/agent/nodes/refiner.pyagent/src/agent/nodes/refiner_graph.pyagent/src/agent/nodes/satisfaction_check.pyagent/src/agent/nodes/schema_explorer.pyagent/src/agent/services/__init__.pyagent/src/agent/services/enrichment_models.pyagent/src/agent/services/enrichment_orchestrator.pyagent/src/agent/services/filter_extractor.pyagent/src/agent/services/hybrid_searcher.pyagent/src/agent/services/location_extractor.pyagent/src/agent/services/sql_transformer.pyagent/src/agent/state.pyagent/src/agent/utils/flag_bridge.pyagent/src/agent/utils/jeen_metadata_client.pyagent/tests/conftest.pyagent/tests/refiner/test_refiner_e2e_mocked.pyagent/tests/refiner/test_refiner_e2e_real.pyagent/tests/refiner/test_refiner_node_agent.pyagent/tests/refiner/test_refiner_node_enrichment.pyagent/tests/refiner/test_refiner_node_trino.pyagent/tests/test_cache_and_gates.pyagent/tests/test_enrichment_orchestrator.pyagent/tests/test_filter_extractor.pyagent/tests/test_finalizer.pyagent/tests/test_hybrid_searcher.pyagent/tests/test_query_builder.pyagent/tests/test_routing.pyagent/tests/test_sql_transformer.pybackend/alembic/versions/ed40dd0a57ad_add_large_category_values_table.pybackend/alembic/versions/merge_heads_d3d006362f40_ed40dd0a57ad.pybackend/app/infra_init.pybackend/app/services/category_ingestion.pybackend/pyproject.tomlcore/src/core/models/models.pycore/src/core/trino.pydocker-compose.ymlscripts/inspect_flow.py
💤 Files with no reviewable changes (1)
- agent/src/agent/nodes/satisfaction_check.py
yuvalkh
commented
Aug 6, 2026
yuvalkh
commented
Aug 6, 2026
yuvalkh
commented
Aug 6, 2026
yuvalkh
commented
Aug 6, 2026
yuvalkh
commented
Aug 6, 2026
yuvalkh
commented
Aug 6, 2026
yuvalkh
commented
Aug 6, 2026
yuvalkh
commented
Aug 6, 2026
yuvalkh
commented
Aug 6, 2026
yuvalkh
commented
Aug 6, 2026
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (8)
agent/src/agent/nodes/refiner.py (1)
139-151: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDo not fall back to the generic refiner prompt.
If
LANGFUSE_PROMPT_REFINER_STEP2is unavailable, stop the run with a clear error. The generic prompt can use a different input and output contract, so this fallback can produce invalid SQL or incorrect satisfaction state.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@agent/src/agent/nodes/refiner.py` around lines 139 - 151, Update the prompt-loading logic in the refiner flow to use only settings.LANGFUSE_PROMPT_REFINER_STEP2. Remove the fallback_key and generic prompt retrieval, and when the step-two prompt cannot be fetched or converted, log a clear error and raise RuntimeError so execution stops without producing a result from the generic prompt.agent/tests/refiner/test_refiner_e2e_mocked.py (1)
226-229: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winSeed and verify the first SQL execution.
Line 226 does not provide an initial
sql_query. Line 224 makesmock_execsucceed for every input. The test can pass when the firsttrino_execcall receives empty SQL. Add an initial query and assert that the first execution input is non-empty.Proposed test update
state = { "user_query": "get data", + "sql_query": "SELECT * FROM A", "runtime_flags": {"SATISFACTION_CHECK_ENABLED": False}, } @@ path = final_state["execution_path"] + assert mock_exec.call_args_list[0].args[0].strip()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@agent/tests/refiner/test_refiner_e2e_mocked.py` around lines 226 - 229, Update the test state in the end-to-end refiner setup to include a non-empty initial sql_query, then verify the first trino_exec/mock_exec invocation receives non-empty SQL before allowing subsequent mocked executions. Preserve the existing runtime_flags and test flow.agent/src/agent/services/filter_extractor.py (4)
121-138: 📐 Maintainability & Code Quality | 🔵 TrivialUse or remove
cte_select_map.The second pass populates
cte_select_map, butresolve_col_refnever reads it. The resolver walks each projection again. Remove the dead pass or use the map.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@agent/src/agent/services/filter_extractor.py` around lines 121 - 138, Remove the unused second-pass loop that populates cte_select_map, since resolve_col_ref does not read it and already walks projections directly. Delete the associated cte_select_map initialization if it has no other consumers, while preserving the existing scope and column-resolution behavior.
347-349: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMake enrichment failures explicit at both entry points.
Both methods catch every exception and return a value that callers treat like normal output:
FilterExtractor.extract()returns[], soEnrichmentOrchestrator.enrich_query()logs “No query filters extracted” and skips enrichment;SQLTransformer.apply()returns the original SQL and the caller logs it as enriched. Narrow these handlers to expected recovery errors, or return a typed success/failure result so transformation failures are not reported as successful enrichment.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@agent/src/agent/services/filter_extractor.py` around lines 347 - 349, Update both FilterExtractor.extract() in agent/src/agent/services/filter_extractor.py:347-349 and SQLTransformer.apply() in agent/src/agent/services/sql_transformer.py:229-231 so enrichment failures cannot appear as successful or empty normal output. Narrow the exception handlers to expected recoverable errors, or introduce a typed success/failure result that EnrichmentOrchestrator.enrich_query() and the SQL transformation caller can detect and report explicitly; preserve fallback values only for genuinely recoverable cases.
253-269: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winSkip subquery-backed
INpredicates explicitly.
Inusesexpressionsfor literal members, butIN (SELECT ...)is stored inquery. Lines 253-269 do not inspectquery, soid IN (SELECT id FROM other_table)yieldsvalue=[]at line 296. Checkleaf.queryand skip or represent non-literal RHS predicates before emitting the filter.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@agent/src/agent/services/filter_extractor.py` around lines 253 - 269, Update the RHS validation in the filter extraction logic to inspect the In predicate’s leaf.query in addition to rhs_keys. When query contains a subquery, skip that predicate or otherwise prevent it from being emitted as a literal filter with an empty value; preserve existing handling for literal expressions and other RHS fields.Source: MCP tools
271-275: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftPreserve
ILIKEthrough the filter extraction contract.SQLGlot represents case-insensitive predicates with
exp.ILike, butfilter_extractor.pymaps bothLikeandILiketoLIKE. This lossy step can make enrichment lose the original case-insensitivity and the regenerated SQL may emit a case-sensitiveLIKE. KeepILIKEdistinct in the extractedSQLFilterParams.operatorand update the field descriptions/docs to cover it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@agent/src/agent/services/filter_extractor.py` around lines 271 - 275, Preserve case-insensitive predicates in the filter extraction contract by mapping exp.ILike to the distinct "ILIKE" operator while retaining exp.Like as "LIKE" in the operator extraction logic of filter_extractor.py. Update the SQLFilterParams field descriptions and related documentation in agent/src/agent/services/sql_transformer.py (lines 85-88) to document both operators and ensure downstream SQL regeneration preserves ILIKE semantics.Source: MCP tools
agent/src/agent/services/sql_transformer.py (2)
121-126: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftMatch transformation plans by source identity.
The matcher compares
tf.columnand the operator, but it does not compare the source table ornode.this.table. In a join containingorders.statusandcustomers.status, a plan for one predicate can transform both predicates when their values and operators match.Carry source identity through
FilterTransformationand require it here. Skip ambiguous matches if source identity is unavailable.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@agent/src/agent/services/sql_transformer.py` around lines 121 - 126, Update the FilterTransformation data flow and the matching loop around matching_tfs to carry and compare source-table identity from node.this.table alongside column and operator. Require both identities to match, and skip transformation-plan matches when either source identity is unavailable so predicates from joined tables cannot be conflated.
128-132: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winCompare
LIKEpatterns without stripping%.Removing
%makesactive%,%active%, and%activeidentical.FilterExtractorassigns different match types for these patterns, but this matcher discards that distinction. A plan for one pattern can rewrite another.Match the raw pattern or include
match_typein the plan key.Also applies to: 145-147
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@agent/src/agent/services/sql_transformer.py` around lines 128 - 132, The matching logic around the current_vals loop must preserve LIKE wildcard distinctions instead of stripping '%' from values. Update the comparison used to append matching_tfs so raw patterns are matched, or include each filter transformation’s match_type in the plan key; ensure active%, %active%, and %active remain distinct.
♻️ Duplicate comments (4)
agent/src/agent/services/sql_transformer.py (3)
178-184: 🎯 Functional Correctness | 🔴 CriticalEmit refined category values as strings.
refined_valuesare categorical strings.make_literalconverts values such as"444"and"1.0"into numeric literals. The generated SQL can then compare a string column with a number. Always useexp.Literal.stringfor category replacements.Verify the pinned literal constructors
#!/bin/bash set -euo pipefail curl -fsSL https://raw.githubusercontent.com/tobymao/sqlglot/v25.0.0/sqlglot/expressions.py | grep -nE 'def number|def string|class Literal' || true🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@agent/src/agent/services/sql_transformer.py` around lines 178 - 184, Update make_literal in the refined_values transformation to always return exp.Literal.string(val_str), removing numeric parsing so categorical replacements such as "444" and "1.0" remain SQL strings.Source: MCP tools
36-38: 🗄️ Data Integrity & Integration | 🟠 MajorUse a reversible token map for catalog markers.
The helper now protects single-quoted literals, but the reverse pass still rewrites every unquoted
$. A pre-existing$in a double-quoted identifier can become@in the returned SQL.Replace only known catalog tokens and restore those exact tokens.
Also applies to: 225-226
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@agent/src/agent/services/sql_transformer.py` around lines 36 - 38, Update the SQL transformation flow around replace_unquoted_char to use a reversible token map for catalog markers: replace each catalog @ with a unique temporary token, then restore only those exact tokens during the reverse pass. Do not globally rewrite unquoted $, so pre-existing $ characters in double-quoted identifiers remain unchanged.
150-170: 🎯 Functional Correctness | 🟠 MajorKeep unmatched values and retain
BETWEEN.When one value changes, the unmatched values are dropped.
IN ('a', 'b')can become only the transformeda, which widens the result set.The length check also converts any two refined values to
IN. ABETWEENpredicate can therefore become membership testing. Keep unmatched literals and rebuildBETWEENexplicitly.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@agent/src/agent/services/sql_transformer.py` around lines 150 - 170, Update the transformation logic around matched_tf and target_operator so unmatched original values are preserved in refined_values instead of being dropped when another value changes. Rebuild a BETWEEN predicate explicitly when the original operator is BETWEEN, rather than converting it to IN based only on refined_values length; retain IN conversion only for applicable membership predicates.agent/src/agent/services/filter_extractor.py (1)
226-239: 🎯 Functional Correctness | 🟠 MajorPreserve negation when extracting predicates.
get_leaf_comparisonsreturns no leaf forexp.Not. This omits predicates such asNOT IN,NOT LIKE, andIS NOT NULL. The pinned sqlglot parser representsNOTwithexp.Not. (github.com)Carry negation through
SQLFilterParamsand the transformation plan, or exclude negated nodes from both extraction and transformation.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@agent/src/agent/services/filter_extractor.py` around lines 226 - 239, Update get_leaf_comparisons to handle exp.Not so negated predicates such as NOT IN, NOT LIKE, and IS NOT NULL are not silently dropped. Preserve the negation through SQLFilterParams and the transformation plan, or consistently exclude those exp.Not nodes from both predicate extraction and transformation.Source: MCP tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@agent/src/agent/nodes/refiner.py`:
- Around line 263-280: Update the table-mapping logic around the table-mapping
loop to match qualified references using both table.db/schema and table.name, so
audit.users resolves only to the corresponding catalog mapping. Use short_name
only when it maps uniquely; do not let duplicate unqualified names select
whichever mapping was processed last, while preserving the existing
full_name_sql AST replacement behavior.
In `@agent/tests/refiner/test_refiner_e2e_real.py`:
- Around line 282-284: Strengthen the satisfied-state assertions in the
final_state validation: retain the existing customer check, assert that
final_state["sql_query"] no longer references the invalid email column, and
verify Trino completed without an error using the existing execution-result
field or assertion pattern.
In `@backend/alembic/versions/f9a3d1c8e205_add_config_schema_flags.py`:
- Line 96: Add a new Alembic revision after f9a3d1c8e205 that updates existing
installations’ persisted MAX_REFINER_ITERATIONS flag and the high_quality and
benchmark overrides to the new values. Implement the changes in the revision’s
upgrade path and provide a downgrade restoring the previous values, using the
existing migration patterns and configuration symbols.
In `@scripts/inspect_flow.py`:
- Around line 169-171: Update the error-display logic around last_trino_error so
whitespace-only values do not cause splitlines()[0] to raise IndexError. Extract
the first non-empty trimmed line and use a fallback message when none exists,
while preserving the existing Trigger output for errors with text.
---
Outside diff comments:
In `@agent/src/agent/nodes/refiner.py`:
- Around line 139-151: Update the prompt-loading logic in the refiner flow to
use only settings.LANGFUSE_PROMPT_REFINER_STEP2. Remove the fallback_key and
generic prompt retrieval, and when the step-two prompt cannot be fetched or
converted, log a clear error and raise RuntimeError so execution stops without
producing a result from the generic prompt.
In `@agent/src/agent/services/filter_extractor.py`:
- Around line 121-138: Remove the unused second-pass loop that populates
cte_select_map, since resolve_col_ref does not read it and already walks
projections directly. Delete the associated cte_select_map initialization if it
has no other consumers, while preserving the existing scope and
column-resolution behavior.
- Around line 347-349: Update both FilterExtractor.extract() in
agent/src/agent/services/filter_extractor.py:347-349 and SQLTransformer.apply()
in agent/src/agent/services/sql_transformer.py:229-231 so enrichment failures
cannot appear as successful or empty normal output. Narrow the exception
handlers to expected recoverable errors, or introduce a typed success/failure
result that EnrichmentOrchestrator.enrich_query() and the SQL transformation
caller can detect and report explicitly; preserve fallback values only for
genuinely recoverable cases.
- Around line 253-269: Update the RHS validation in the filter extraction logic
to inspect the In predicate’s leaf.query in addition to rhs_keys. When query
contains a subquery, skip that predicate or otherwise prevent it from being
emitted as a literal filter with an empty value; preserve existing handling for
literal expressions and other RHS fields.
- Around line 271-275: Preserve case-insensitive predicates in the filter
extraction contract by mapping exp.ILike to the distinct "ILIKE" operator while
retaining exp.Like as "LIKE" in the operator extraction logic of
filter_extractor.py. Update the SQLFilterParams field descriptions and related
documentation in agent/src/agent/services/sql_transformer.py (lines 85-88) to
document both operators and ensure downstream SQL regeneration preserves ILIKE
semantics.
In `@agent/src/agent/services/sql_transformer.py`:
- Around line 121-126: Update the FilterTransformation data flow and the
matching loop around matching_tfs to carry and compare source-table identity
from node.this.table alongside column and operator. Require both identities to
match, and skip transformation-plan matches when either source identity is
unavailable so predicates from joined tables cannot be conflated.
- Around line 128-132: The matching logic around the current_vals loop must
preserve LIKE wildcard distinctions instead of stripping '%' from values. Update
the comparison used to append matching_tfs so raw patterns are matched, or
include each filter transformation’s match_type in the plan key; ensure active%,
%active%, and %active remain distinct.
In `@agent/tests/refiner/test_refiner_e2e_mocked.py`:
- Around line 226-229: Update the test state in the end-to-end refiner setup to
include a non-empty initial sql_query, then verify the first
trino_exec/mock_exec invocation receives non-empty SQL before allowing
subsequent mocked executions. Preserve the existing runtime_flags and test flow.
---
Duplicate comments:
In `@agent/src/agent/services/filter_extractor.py`:
- Around line 226-239: Update get_leaf_comparisons to handle exp.Not so negated
predicates such as NOT IN, NOT LIKE, and IS NOT NULL are not silently dropped.
Preserve the negation through SQLFilterParams and the transformation plan, or
consistently exclude those exp.Not nodes from both predicate extraction and
transformation.
In `@agent/src/agent/services/sql_transformer.py`:
- Around line 178-184: Update make_literal in the refined_values transformation
to always return exp.Literal.string(val_str), removing numeric parsing so
categorical replacements such as "444" and "1.0" remain SQL strings.
- Around line 36-38: Update the SQL transformation flow around
replace_unquoted_char to use a reversible token map for catalog markers: replace
each catalog @ with a unique temporary token, then restore only those exact
tokens during the reverse pass. Do not globally rewrite unquoted $, so
pre-existing $ characters in double-quoted identifiers remain unchanged.
- Around line 150-170: Update the transformation logic around matched_tf and
target_operator so unmatched original values are preserved in refined_values
instead of being dropped when another value changes. Rebuild a BETWEEN predicate
explicitly when the original operator is BETWEEN, rather than converting it to
IN based only on refined_values length; retain IN conversion only for applicable
membership predicates.
🪄 Autofix
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: ASSERTIVE
Plan: Pro Plus
Run ID: db028865-94f0-4b92-a588-dcf91fdc9914
📒 Files selected for processing (22)
agent/src/agent/config.pyagent/src/agent/graph.pyagent/src/agent/nodes/detect_ambiguity.pyagent/src/agent/nodes/extractor.pyagent/src/agent/nodes/refiner.pyagent/src/agent/nodes/refiner_graph.pyagent/src/agent/services/filter_extractor.pyagent/src/agent/services/sql_transformer.pyagent/src/agent/state.pyagent/src/agent/utils/flag_bridge.pyagent/src/agent/utils/schema_enrichment.pyagent/src/agent/utils/sql.pyagent/tests/refiner/test_refiner_e2e_mocked.pyagent/tests/refiner/test_refiner_e2e_real.pyagent/tests/refiner/test_refiner_node_agent.pyagent/tests/refiner/test_refiner_node_enrichment.pyagent/tests/refiner/test_refiner_node_trino.pyagent/tests/test_extractor.pyagent/tests/test_resilience.pyagent/tests/test_routing.pybackend/alembic/versions/f9a3d1c8e205_add_config_schema_flags.pyscripts/inspect_flow.py
💤 Files with no reviewable changes (3)
- agent/src/agent/utils/flag_bridge.py
- agent/src/agent/nodes/extractor.py
- agent/src/agent/utils/schema_enrichment.py
yuvalkh
commented
Aug 13, 2026
| return "No schema context available." | ||
|
|
||
|
|
||
| def parse_jeen_catalog_tables(jeen_catalog: str) -> tuple[dict, list[AgentSQLTable]]: |
Collaborator
Author
There was a problem hiding this comment.
Instead of parsing the catalog_tables through the string search, maybe it's better to call the mcp functions get_tables, and get_columns?
It will look better in code and also will be easier to take the values there such as type, examples, description
Collaborator
Author
There was a problem hiding this comment.
yuvalkh
commented
Aug 13, 2026
yuvalkh
commented
Aug 13, 2026
| "last_result_row_count": len(inline_result_rows) | ||
| if inline_result_rows | ||
| else 0, | ||
| "last_result_data": str([inline_result_columns] + inline_result_rows[:15]) |
Collaborator
Author
There was a problem hiding this comment.
Currently this will be [:15] but we will change it after the merge to the new env variable that tells us how many preview rows to show
yuvalkh
commented
Aug 15, 2026
yuvalkh
commented
Aug 15, 2026
yuvalkh
commented
Aug 16, 2026
yuvalkh
commented
Aug 16, 2026
yuvalkh
commented
Aug 16, 2026
yuvalkh
commented
Aug 16, 2026
yuvalkh
commented
Aug 16, 2026
yuvalkh
commented
Aug 16, 2026
yuvalkh
commented
Aug 16, 2026
yuvalkh
commented
Aug 16, 2026
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
2 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.
Summary by CodeRabbit
New Features
Bug Fixes
Tests