Ben/category enrichment - #32
Conversation
…2sql-onboarding into ben/category-enrichment
…ding into ben/category-enrichment
📝 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 💡 1📝 Generate docstrings 💡
🧪 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 |
|
@coderabbitai review |
✅ Action performedReview finished.
|
| if not esca_write_enabled or not raw_data_ref: | ||
| if inline_result_rows is not None: | ||
| limit = 5 | ||
| limit = 10 |
There was a problem hiding this comment.
I know this is not your code but we need to add an env variable RESULT_ROW_COUNT_LIMIT that we check if it exists it will get only the first RESULT_ROW_COUNT_LIMIT and if not it will use everything (all inline_result_columns).
This is because sometimes we don't want to trim the whole rows we get back from the db (when not using esca)
![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
| MAX_REFINER_ITERATIONS: int = Field(default=10, gt=0) | ||
| REFINER_SCHEMA_CONTEXT_TABLES: int = Field(default=8, gt=0) |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
Assess the cost of 10 refiner iterations.
MAX_REFINER_ITERATIONS moves from 3 to 10. Each iteration in the new graph runs one LLM call in agent_node plus one Trino execution in trino_exec_node. Worst-case latency and token cost per request increase by more than three times, and a failing query now holds a Trino connection for up to 10 attempts. Confirm that a request timeout or a cost guard bounds this loop.
🤖 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/config.py` around lines 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.
| # Suppress unverified HTTPS warnings for dev internal endpoints | ||
| urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) | ||
| warnings.filterwarnings("ignore", category=urllib3.exceptions.InsecureRequestWarning) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Global warning suppression hides TLS problems for the whole process.
urllib3.disable_warnings and the warnings.filterwarnings call are process-wide. Importing this module silences InsecureRequestWarning for every HTTP client in the agent, including the Jeen MCP client and the ESCA client, not only for Langfuse. Certificate verification failures then become invisible in production.
Restrict the suppression to the development case, for example by gating it on a setting, and keep verification enabled elsewhere. urllib3 is also imported directly here; confirm it is a declared dependency of the agent package rather than a transitive one.
🔒 Proposed fix
-# Suppress unverified HTTPS warnings for dev internal endpoints
-urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
-warnings.filterwarnings("ignore", category=urllib3.exceptions.InsecureRequestWarning)
+# Suppress unverified HTTPS warnings for dev internal endpoints only.
+if getattr(settings, "ALLOW_INSECURE_INTERNAL_TLS", False):
+ urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+ warnings.filterwarnings(
+ "ignore", category=urllib3.exceptions.InsecureRequestWarning
+ )🤖 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/langfuse_client.py` around lines 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.
| for short, full in table_mappings.items(): | ||
| match = re.match(r'\"([^\"]+)\"\.\"([^\"]+)\"\.\"([^\"]+)\"', full) | ||
| if match: | ||
| cat, sch, tbl = match.groups() | ||
| sql = re.sub(rf'(?<![\.\w\"])\"{re.escape(sch)}\"\.\"{re.escape(tbl)}\"', full, sql) | ||
| sql = re.sub(rf'(?<![\.\w\"]){re.escape(sch)}\.{re.escape(tbl)}(?![\.\w\"])', full, sql) | ||
| sql = re.sub(rf'(?<![\.\w])\"{re.escape(short)}\"(?![\.\w])', full, sql) | ||
| sql = re.sub(rf'(?<![\.\w\"]){re.escape(short)}(?![\.\w\"])', full, sql) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Table-name rewriting also edits string literals and aliases.
The substitutions at lines 240-241 replace the short table name anywhere in the query text. The lookarounds block only adjacent word characters, dots, and quotes. They do not exclude string literals, column names, or aliases. Two concrete failures:
SELECT * FROM users WHERE role = 'users'becomes... WHERE role = '"hive"."prod"."users_table"', which silently changes the predicate value.- A column named
usersor an aliasAS usersis rewritten into a three-part name and the query fails to parse.
The repository already depends on sqlglot (see agent/src/agent/services/sql_transformer.py). Parse the SQL once and qualify table references on the AST. That removes the literal and alias hazards and replaces six regex passes with one traversal.
Add a test for the literal case in agent/tests/refiner/test_refiner_node_trino.py, next to test_trino_exec_table_alias_word_boundary.
🧰 Tools
🪛 ast-grep (0.45.0)
[warning] 237-237: Regex pattern passed to re is built from a non-literal (variable, call, concatenation, or f-string) value. If that value is attacker-controlled it can introduce a malicious pattern with catastrophic backtracking (ReDoS). Use a hardcoded literal pattern, or validate/escape untrusted input with re.escape() and bound the regex complexity before compiling.
Context: re.sub(rf'(?<![.\w"])"{re.escape(sch)}"."{re.escape(tbl)}"', full, sql)
Note: [CWE-1333] Inefficient Regular Expression Complexity.
(redos-non-literal-regex-python)
[warning] 238-238: Regex pattern passed to re is built from a non-literal (variable, call, concatenation, or f-string) value. If that value is attacker-controlled it can introduce a malicious pattern with catastrophic backtracking (ReDoS). Use a hardcoded literal pattern, or validate/escape untrusted input with re.escape() and bound the regex complexity before compiling.
Context: re.sub(rf'(?<![.\w"]){re.escape(sch)}.{re.escape(tbl)}(?![.\w"])', full, sql)
Note: [CWE-1333] Inefficient Regular Expression Complexity.
(redos-non-literal-regex-python)
[warning] 239-239: Regex pattern passed to re is built from a non-literal (variable, call, concatenation, or f-string) value. If that value is attacker-controlled it can introduce a malicious pattern with catastrophic backtracking (ReDoS). Use a hardcoded literal pattern, or validate/escape untrusted input with re.escape() and bound the regex complexity before compiling.
Context: re.sub(rf'(?<![.\w])"{re.escape(short)}"(?![.\w])', full, sql)
Note: [CWE-1333] Inefficient Regular Expression Complexity.
(redos-non-literal-regex-python)
[warning] 240-240: Regex pattern passed to re is built from a non-literal (variable, call, concatenation, or f-string) value. If that value is attacker-controlled it can introduce a malicious pattern with catastrophic backtracking (ReDoS). Use a hardcoded literal pattern, or validate/escape untrusted input with re.escape() and bound the regex complexity before compiling.
Context: re.sub(rf'(?<![.\w"]){re.escape(short)}(?![.\w"])', full, sql)
Note: [CWE-1333] Inefficient Regular Expression Complexity.
(redos-non-literal-regex-python)
🤖 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 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.
| @staticmethod | ||
| def extract(sql: str, schema: Dict[str, Dict[str, str]]) -> List[SQLFilterParams]: |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Decompose extract into module-level helpers.
extract spans roughly 300 lines and defines six nested closures (lowercase_identifiers, nest_schema, get_unaliased_table_name, resolve_col_ref, extract_literal_val, get_leaf_comparisons). The closures capture table_alias_map and unnest_map, so none of them can be unit tested on its own. agent/tests/test_filter_extractor.py can only exercise the whole pipeline through SQL strings.
Extract extract_literal_val and get_leaf_comparisons to module level, since they capture nothing. Move the scope-resolution state into a small helper class so resolve_col_ref becomes a testable method. extract_literal_val is also duplicated verbatim in agent/src/agent/services/sql_transformer.py lines 24-183; a shared module-level function removes that duplication.
🤖 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 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.
| # 1. Trino catalog workaround: replace '@' in table references with '$' | ||
| sql_processed: str = sql.replace("@", "$") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Global @→$ substitution corrupts literal filter values.
Line 42 replaces every @ in the whole SQL string, not only in table references as the comment states. The extractor reverses the substitution for source_table (line 302) and for original_expression (line 333), but never for value.
A predicate such as WHERE user_email = 'jane@corp.com' therefore yields value == "jane$corp.com". That corrupted value flows into HybridSearcher.search and into the f"{col}#@#{val}" candidate keys and LLM prompt built by EnrichmentOrchestrator.enrich_query (agent/src/agent/services/enrichment_orchestrator.py lines 90-226), so lookups miss and the LLM sees a wrong original value. No test in agent/tests/test_filter_extractor.py covers a literal that contains @.
Reverse the substitution on extracted values.
🐛 Proposed fix
def extract_literal_val(node: Optional[exp.Expression]) -> Any:
"""Translates sqlglot AST literal/boolean node values to Python primitives."""
if node is None:
return None
if isinstance(node, exp.Literal):
if node.is_string:
- return node.this
+ return node.this.replace("$", "@") if "$" in node.this else node.this
try:A safer alternative restricts the substitution to table references only, instead of rewriting the raw SQL string. Then no reversal is needed for values.
Also applies to: 287-296
🤖 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 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'.
| embedding: Any | None = Field(default=None, sa_column=Column(Vector(768))) | ||
| embedder_model: str = Field(default="nomic-embed-text") |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial
Plan an ANN index for the embedding column.
large_category_values holds one row per distinct categorical value per column, so the table grows with cardinality. Semantic search over embedding without an HNSW or IVFFlat index performs a sequential scan of every row. Add a vector index in the migration, and scope lookups by table_id and column_name so the scan stays bounded.
🤖 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 `@core/src/core/models/models.py` around lines 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.
| created_at: datetime = Field(default_factory=datetime.now) | ||
| updated_at: datetime = Field(default_factory=datetime.now) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win
updated_at never changes after insert.
updated_at uses default_factory only. SQLAlchemy does not refresh it on UPDATE. Re-ingestion of a category value therefore keeps the original timestamp, and LargeCategoryValueRead.updated_at reports stale freshness data. Add an onupdate hook.
♻️ Proposed fix
created_at: datetime = Field(default_factory=datetime.now)
- updated_at: datetime = Field(default_factory=datetime.now)
+ updated_at: datetime = Field(
+ default_factory=datetime.now,
+ sa_column_kwargs={"onupdate": datetime.now},
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| created_at: datetime = Field(default_factory=datetime.now) | |
| updated_at: datetime = Field(default_factory=datetime.now) | |
| created_at: datetime = Field(default_factory=datetime.now) | |
| updated_at: datetime = Field( | |
| default_factory=datetime.now, | |
| sa_column_kwargs={"onupdate": datetime.now}, | |
| ) |
🤖 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 `@core/src/core/models/models.py` around lines 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.
| import urllib3 | ||
|
|
||
| # Suppress unverified HTTPS warnings for dev internal endpoints | ||
| urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm urllib3 is declared in the core manifest and check import ordering in core/trino.py.
fd -t f 'pyproject.toml|requirements.*txt' | xargs -r rg -n 'urllib3'
sed -n '1,30p' core/src/core/trino.pyRepository: StavPonte11/text2sql-onboarding
Length of output: 787
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== manifests containing urllib3/trino/requests ==="
python3 - <<'PY'
from pathlib import Path
for p in Path('.').rglob('*'):
if p.is_file() and p.name in {
'pyproject.toml','requirements.txt','requirements-dev.txt',
'requirements-core.txt','Pipfile','poetry.lock','uv.lock','Pipfile.lock'
}:
try:
text = p.read_text(errors='ignore').lower()
except Exception:
continue
if 'urllib3' in text or 'trino' in text or 'requests' in p.name.lower():
print(f'--- {p} ---')
for i,line in enumerate(p.read_text(errors='ignore').splitlines(),1):
if any(q.lower() in line.lower() for q in ['urllib3','trino','requests']):
print(f'{i}:{line}')
PY
echo "=== trino.py top section with line numbers ==="
cat -n core/src/core/trino.py | sed -n '1,120p'
echo "=== all urllib3.disable_warnings occurrences ==="
rg -n "disable_warnings|InsecureRequestWarning|TRINO_VERIFY|verify=" core/src || trueRepository: StavPonte11/text2sql-onboarding
Length of output: 14686
Scope the urllib3 warning suppression and declare the dependency.
urllib3.disable_warnings(...) runs at module import, before settings is available, and disables InsecureRequestWarning for the whole process. Gate this only when settings.TRINO_VERIFY is false, and list urllib3 in core/pyproject.toml instead of relying on transitive imports.
🤖 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 `@core/src/core/trino.py` around lines 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.
| - LLM_API_KEY=sk-pj7x2rp8o2tli9 | ||
| - LLM_BASE_URL=https://pj7x2rp8o2tli9-8000.proxy.runpod.net/v1 | ||
| - LLM_MODEL=openai/gpt-oss-120b | ||
| - JEEN_LLM_CORE_URL=http://schema-modeler.dev161.internal/api/mcp | ||
| - JEEN_API_KEY=mcp_ecd023ab04f849b36aef5d797525365c4c095052e6e07577d065bfe82507ae67 |
There was a problem hiding this comment.
🔒 Security & Privacy | 🔴 Critical | ⚡ Quick win
Remove and rotate the exposed credentials.
Lines 568 and 572 commit usable API keys into the repository. Docker Compose also exposes them through the resolved container configuration.
Move both values to an ignored secret source or Docker Compose secrets. Rotate both credentials immediately. Also remove the identical JEEN_API_KEY default in agent/src/agent/config.py:8-35, or the rotated Jeen key can remain exposed there.
🧰 Tools
🪛 Betterleaks (1.7.3)
[high] 568-568: Detected a Generic API Key, potentially exposing access to various services and sensitive operations.
(generic-api-key)
[high] 572-572: Detected a Generic API Key, potentially exposing access to various services and sensitive operations.
(generic-api-key)
🤖 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 `@docker-compose.yml` around lines 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.
Source: Linters/SAST tools
| refiner_started = False | ||
| last_trino_error = None | ||
| last_trino_row_count = None | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Report failed executions as failures.
When the refiner emits end_fail, Line 306 still prints a success banner. When agent_graph.astream() raises, Lines 308-311 print the error and return normally. Both paths produce exit status 0.
Track the terminal outcome. Return failure from run_flow for end_fail and exceptions. Make main exit non-zero for a failed non-interactive run.
Proposed fix
refiner_started = False
+ refiner_failed = False
last_trino_error = None
elif node_name == "end_fail":
+ refiner_failed = True
reason = updates.get("escalation_reason", "Refinement limit reached")
print_refiner_step_header("END FAIL", f"{RED}Refinement Exited ({reason}){RESET}")
- print_banner("Execution Completed Successfully!", GREEN)
+ if refiner_failed:
+ print_banner("Execution Failed", RED)
+ return False
+ print_banner("Execution Completed Successfully!", GREEN)
+ return True
except Exception as exc:
print_banner(f"Execution Encountered Error: {exc}", RED)
import traceback
traceback.print_exc()
+ return False
if args.query:
- asyncio.run(run_flow(args.query, auto_approve=not args.require_approval))
+ succeeded = asyncio.run(run_flow(args.query, auto_approve=not args.require_approval))
+ if not succeeded:
+ raise SystemExit(1)Also applies to: 225-227, 306-311, 321-322
🤖 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 `@scripts/inspect_flow.py` around lines 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.
| JEEN_LLM_CORE_URL: str = "http://schema-modeler.dev161.internal/api/mcp" # If empty, agent gracefully skips fetching | ||
| JEEN_API_KEY: str = "mcp_ecd023ab04f849b36aef5d797525365c4c095052e6e07577d065bfe82507ae67" # If empty, agent gracefully skips fetching |
There was a problem hiding this comment.
🔒 Security & Privacy | 🔴 Critical | ⚡ Quick win
Remove the hard-coded API key from source.
Line 29 commits a live-looking mcp_... bearer token as the default value of JEEN_API_KEY. The value is now in Git history and is readable by anyone with repository access. The same pattern exists at line 42 for JEEN_METADATA_MCP_KEY. Restore empty defaults and supply the keys through the environment. Rotate both keys, because they must be treated as compromised.
The inline comment "If empty, agent gracefully skips fetching" also no longer matches a non-empty default.
🔒 Proposed fix
# ── Jeen Integration ──────────────────────────────────────────────────────
- JEEN_LLM_CORE_URL: str = "http://schema-modeler.dev161.internal/api/mcp" # If empty, agent gracefully skips fetching
- JEEN_API_KEY: str = "mcp_ecd023ab04f849b36aef5d797525365c4c095052e6e07577d065bfe82507ae67" # If empty, agent gracefully skips fetching
+ JEEN_LLM_CORE_URL: str = "" # If empty, agent gracefully skips fetching
+ JEEN_API_KEY: str = "" # If empty, agent gracefully skips fetching📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| JEEN_LLM_CORE_URL: str = "http://schema-modeler.dev161.internal/api/mcp" # If empty, agent gracefully skips fetching | |
| JEEN_API_KEY: str = "mcp_ecd023ab04f849b36aef5d797525365c4c095052e6e07577d065bfe82507ae67" # If empty, agent gracefully skips fetching | |
| JEEN_LLM_CORE_URL: str = "" # If empty, agent gracefully skips fetching | |
| JEEN_API_KEY: str = "" # If empty, agent gracefully skips fetching |
🤖 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/config.py` around lines 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.
| LANGFUSE_PROMPT_FINALIZER: str = "text2sql/finalizer" | ||
| LANGFUSE_PROMPT_REJECTION_ROUTER: str = "text2sql/rejection_router" | ||
| LANGFUSE_PROMPT_CATEGORY_ENRICHMENT: str = "text2sql/category_enrichment" | ||
| LANGFUSE_PROMPT_LOC_EXTRACTOR: str = "text2sql/extractor" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find all uses of the two extractor prompt settings and their prompt templates.
rg -n -C3 'LANGFUSE_PROMPT_LOC_EXTRACTOR\b|LANGFUSE_PROMPT_EXTRACTOR\b'
fd -i 'location' --extension txt --extension md --exec cat -n {}Repository: StavPonte11/text2sql-onboarding
Length of output: 169
🏁 Script executed:
#!/bin/bash
set -u
echo "== repo files sample =="
git ls-files | sed -n '1,120p'
echo "== config file =="
fd -i 'config.py' -x sh -c 'echo "--- $0"; wc -l "$0"; sed -n "1,120p" "$0"' {}
echo "== prompt extractor occurrences case-insensitive full text =="
rg -n -C3 'LANGFUSE_PROMPT(extractor|LOC_EXTRACTOR)|text2sql/extractor|text2sql/location|extractor' . || true
echo "== locate location extractor files =="
fd -i 'location' --extension py --extension ts --extension js --extension txt --extension md || trueRepository: StavPonte11/text2sql-onboarding
Length of output: 50389
🏁 Script executed:
#!/bin/bash
set -u
echo "== agent config relevant lines =="
sed -n '45,75p' agent/src/agent/config.py
echo "== location extractor implementation =="
fd -i 'location_extractor.py' -x sh -c 'echo "--- $0"; wc -l "$0"; sed -n "$1,\$p" "$0"' {} \;
echo "== upload script relevant prompt/setting entries =="
sed -n '16,45p;70,115p;180,235p' agent/scripts/upload_all_prompts.py
echo "== config references to location prompt identifiers =="
rg -n -C2 'LANGFUSE_PROMPT_LOC_EXTRACTOR|LANGFUSE_PROMPT_LOC_EXTRACTOR_INSTRUCTION|text2sql/extractor|text2sql/location' agent/src/agent agent/scripts/worker || true
echo "== all config.py references exact identifiers =="
rg -n 'LANGFUSE_PROMPT_(LOC_EXTRACTOR|EXTRACTOR|LOC_EXTRACTOR_INSTRUCTION)' . || trueRepository: StavPonte11/text2sql-onboarding
Length of output: 18556
🏁 Script executed:
#!/bin/bash
set -u
echo "== location_extractor_implementation sections =="
sed -n '1,260p' agent/src/agent/services/location_extractor.py | cat -n
echo "== tests relevant expectations/calls =="
sed -n '135,230p' agent/tests/test_location_extractor.py | cat -n
echo "== upload script location/instruction references =="
rg -n -C4 'location|extractor|LOCATION|instruction|wkt' agent/scripts/upload_all_prompts.py || trueRepository: StavPonte11/text2sql-onboarding
Length of output: 20683
Point LANGFUSE_PROMPT_LOC_EXTRACTOR at the location prompt.
LocationExtractorAgent._build_prompt() uses LANGFUSE_PROMPT_LOC_EXTRACTOR, but this setting is still "text2sql/extractor", the regular query-enrichment prompt. If the shared prompt does not return the Hebrew-name to standard-location JSON map, _parse_llm_json() returns {} and location WKT injection is skipped. Use the intended location extractor prompt identifier or add the missing prompt.
🤖 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/config.py` at 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().
| def _safe_update_current_span(*args, **kwargs): | ||
| current_span = otel_trace_api.get_current_span() | ||
| if current_span is otel_trace_api.INVALID_SPAN: | ||
| return | ||
| return _orig_update_current_span(*args, **kwargs) | ||
|
|
||
|
|
||
| def _safe_get_current_trace_id(*args, **kwargs): | ||
| current_span = otel_trace_api.get_current_span() | ||
| if current_span is otel_trace_api.INVALID_SPAN: | ||
| return None | ||
| return _orig_get_current_trace_id(*args, **kwargs) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
opentelemetry python get_current_span INVALID_SPAN NonRecordingSpan span context is_valid
💡 Result:
In OpenTelemetry Python, get_current_span is used to retrieve the span currently stored in the execution context [1][2]. If no span is present in the context, it returns an INVALID_SPAN object [3][1]. INVALID_SPAN is a specific instance of a NonRecordingSpan [3][4]. A NonRecordingSpan is a "no-op" span implementation used when no full tracing SDK or instrumentation is active [3][5][6]. Its primary purpose is to allow applications to continue using the OpenTelemetry API (e.g., propagating trace contexts) without crashing or requiring a fully functional tracing backend [3][5]. Key points regarding these components: 1. INVALID_SPAN: This is the default returned by get_current_span when no active span is found [1][2]. It contains an "invalid" SpanContext [4]. 2. NonRecordingSpan: This class implements the Span interface but performs no actual recording of events, attributes, or status updates [3][6]. Its method is_recording consistently returns False [5][6]. 3. SpanContext and is_valid: The SpanContext associated with a span has a property called is_valid [3][6]. This property indicates whether the context contains valid trace and span identifiers [3][7]. In the case of INVALID_SPAN, the associated SpanContext has is_valid set to False (or its equivalent state indicating an invalid/empty context) [5]. You can check the validity of a span's context by accessing span.get_span_context.is_valid [3][6]. If you find that get_current_span is returning an INVALID_SPAN unexpectedly, it often means that no span has been activated in the current execution context [8]. This typically occurs if you have not started a span using a tracer (e.g., using tracer.start_as_current_span or the context manager) or if the scope where the span was activated has already closed [8].
Citations:
- 1: https://github.com/open-telemetry/opentelemetry-python/blob/main/opentelemetry-api/src/opentelemetry/trace/propagation/__init__.py
- 2: https://opentelemetry-python.readthedocs.io/en/stable/_modules/opentelemetry/trace/propagation.html
- 3: https://opentelemetry-python.readthedocs.io/en/stable/api/trace.html
- 4: https://github.com/open-telemetry/opentelemetry-python/blob/main/opentelemetry-api/src/opentelemetry/trace/__init__.py
- 5: https://opentelemetry.io/docs/specs/otel/trace/api/
- 6: https://opentelemetry-python.readthedocs.io/en/latest/%5Fmodules/opentelemetry/trace/span.html
- 7: Refactor is_valid to be an instance attribute open-telemetry/opentelemetry-python#1005
- 8: trace.get_current_span doesn't provide active span under test open-telemetry/opentelemetry-python#3999
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Repository files matching langfuse_client.py:"
fd -a 'langfuse_client.py' . || true
echo
echo "Relevant file excerpt:"
if [ -f agent/src/agent/langfuse_client.py ]; then
cat -n agent/src/agent/langfuse_client.py | sed -n '1,120p'
else
echo "file not found"
fi
echo
echo "OpenTelemetry API/imports and span context usage:"
rg -n "INVALID_SPAN|get_current_span|NonRecordingSpan|is_valid|get_span_context" -S . --glob '!node_modules' --glob '!dist' --glob '!build' || true
echo
echo "Check installed opentelemetry API availability and behavior without importing repo code:"
python3 - <<'PY'
import importlib.util
for mod in ['opentelemetry.trace','opentelemetry.trace.span']:
print(f'{mod}: {importlib.util.find_spec(mod) is not None}')
if importlib.util.find_spec('opentelemetry.trace') is not None:
import opentelemetry.trace as trace_api
print('INVALID_SPAN id:', id(trace_api.INVALID_SPAN))
print('INVALID_SPAN __class__:', trace_api.INVALID_SPAN.__class__)
print('INVALID_SPAN is_valid:', trace_api.INVALID_SPAN.get_span_context().is_valid if hasattr(trace_api.INVALID_SPAN,'get_span_context') else 'no get_span_context')
print('get_current_span():', trace_api.get_current_span())
print('current is same INSTANCE:', trace_api.get_current_span() is trace_api.INVALID_SPAN)
print('current is_valid:', trace_api.get_current_span().get_span_context().is_valid if hasattr(trace_api.get_current_span(),'get_span_context') else 'no get_span_context')
PYRepository: StavPonte11/text2sql-onboarding
Length of output: 2762
🌐 Web query:
OpenTelemetry Python api span propagate active span NON_RECORDING_span InvalidSpan get_current_span behavior
💡 Result:
In OpenTelemetry Python, the behavior of span-related objects and the retrieval of the current span are designed to ensure applications do not crash when tracing is not active or when a span is explicitly configured not to record data [1][2]. get_current_span and InvalidSpan The function opentelemetry.trace.get_current_span is used to retrieve the currently active span from the context [3][4]. If no span is active in the current context, it returns an INVALID_SPAN object [3][4]. INVALID_SPAN is a singleton that implements the Span interface as a no-op; calling methods on it (such as set_attribute or end) does nothing, ensuring that your application code remains safe even if the tracing infrastructure is not initialized [5][2]. NonRecordingSpan A NonRecordingSpan is a specific type of span that exists to facilitate context propagation even when a span is not intended to be recorded (e.g., due to sampling decisions) [5][1]. - Functionality: Like INVALID_SPAN, a NonRecordingSpan is essentially a no-op span [5]. Its primary purpose is to hold SpanContext information (such as TraceId and SpanId) so that trace context can be propagated to downstream services without incurring the overhead of recording span data [5][6]. - Propagation: When trace context is extracted (e.g., from incoming HTTP headers), a NonRecordingSpan may be created if the incoming context indicates that the trace should not be sampled (e.g., the sampled flag is set to 0 in the W3C traceparent header) [7][6][8]. - Confusion in Behavior: Users often encounter confusion when they see spans are not recorded when a traceparent header is present [7][8]. This is typically due to the default ParentBased sampler, which respects the sampling decision propagated from the upstream service [7]. If the upstream service decided not to sample the request, the Python application will create a NonRecordingSpan for the child, causing it to appear as if tracing is broken [7][8]. Summary of Key Differences - INVALID_SPAN: Returned by get_current_span when absolutely no span is associated with the current context [3][4]. It is the "default" fallback [2]. - NonRecordingSpan: Returned when a context exists and needs to be propagated, but the sampling configuration (or inherited sampling decision) dictates that no span data should be recorded [5][7][6]. If you are experiencing unexpected non-recording behavior, verify your sampling configuration (e.g., OTEL_TRACES_SAMPLER environment variable) and ensure you are not explicitly passing or inheriting a non-sampled context [7][8].
Citations:
- 1: https://github.com/open-telemetry/opentelemetry-python/blob/main/opentelemetry-api/src/opentelemetry/trace/__init__.py
- 2: TracerProvider now required? open-telemetry/opentelemetry-python#881
- 3: https://github.com/open-telemetry/opentelemetry-python/blob/main/opentelemetry-api/src/opentelemetry/trace/propagation/__init__.py
- 4: https://opentelemetry-python.readthedocs.io/en/stable/_modules/opentelemetry/trace/propagation.html
- 5: https://opentelemetry-python.readthedocs.io/en/stable/api/trace.html
- 6: https://github.com/open-telemetry/opentelemetry-python/blob/main/opentelemetry-api/src/opentelemetry/trace/propagation/tracecontext.py
- 7: Spans not recorded by default if context is propagated open-telemetry/opentelemetry-python#4430
- 8: FastAPI OTEL not recording span when propagation is in use? open-telemetry/opentelemetry-python-contrib#3267
🌐 Web query:
Python opentelemetry NonRecordingSpan is_recording get span context is_valid
💡 Result:
In the OpenTelemetry Python API, NonRecordingSpan is a concrete implementation of the Span interface used when no functional tracing SDK is configured or active [1][2]. Its primary purpose is to allow applications to use the OpenTelemetry API without side effects, while still supporting context propagation [1][3]. NonRecordingSpan and its methods behave as follows: is_recording: This method always returns False for a NonRecordingSpan [2][4]. This signals that the span is not capturing any data such as attributes, events, or status updates [3]. All operations that would typically record data (e.g., set_attribute, add_event) are no-op in this class [2][4]. get_span_context: This method returns the SpanContext associated with the NonRecordingSpan [2][4]. Even when a span is non-recording, it may hold a valid SpanContext to ensure that trace and span identifiers can be propagated to downstream services or child spans [3][5]. SpanContext.is_valid: The is_valid property on a SpanContext indicates whether the context contains a properly formatted, non-zero trace ID and span ID [6][7]. A SpanContext is considered valid if the trace ID is greater than the INVALID_TRACE_ID and the span ID is greater than the INVALID_SPAN_ID [7]. In summary, a NonRecordingSpan is a lightweight, no-op span that maintains a SpanContext to facilitate trace propagation, but does not perform any data recording [3][4]. Its is_recording method consistently returns False [2][4], while the validity of its context is determined by the internal state of the attached SpanContext [7].
Citations:
- 1: https://opentelemetry-python.readthedocs.io/en/stable/api/trace.html
- 2: https://opentelemetry-python.readthedocs.io/en/latest/%5Fmodules/opentelemetry/trace/span.html
- 3: https://opentelemetry.io/docs/specs/otel/trace/api/
- 4: https://github.com/open-telemetry/opentelemetry-python/blob/eed100c3/opentelemetry-api/src/opentelemetry/trace/span.py
- 5: https://github.com/open-telemetry/opentelemetry-python/blob/main/opentelemetry-api/src/opentelemetry/trace/__init__.py
- 6: https://github.com/open-telemetry/opentelemetry-specification/blob/master/specification/trace/api.md
- 7: https://opentelemetry-python.readthedocs.io/en/stable/_modules/opentelemetry/trace/span.html
Test span recording and context validity, not object identity.
INVALID_SPAN is returned when the empty context has no span. Other OpenTelemetry no-op spans, such as propagated non-recording spans and sampled child spans with no SDK, are distinct objects and still cannot record Langfuse data. Check span.get_span_context().is_valid and span.get_span_context().is_remote, or the span’s recording-capability, before calling update_current_span and get_current_trace_id.
🤖 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/langfuse_client.py` around lines 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.
| except Exception as e: | ||
| return f"Error retrieving data preview from Esca: {e}" |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟡 Minor | ⚡ Quick win
Log the exception and return a neutral preview message.
get_esca_preview returns the raw exception text. finalizer_node passes that text to the LLM as sql_results. The model can echo internal storage details into the user-facing summary. Log the exception and return a message that contains no internal detail.
🛡️ Proposed fix
- except Exception as e:
- return f"Error retrieving data preview from Esca: {e}"
+ except Exception:
+ logger.exception("Failed to retrieve data preview from Esca for id=%s", esca_id)
+ return "Data preview is unavailable."Add the logger import if it is absent:
import logging
logger = logging.getLogger(__name__)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| except Exception as e: | |
| return f"Error retrieving data preview from Esca: {e}" | |
| except Exception: | |
| logger.exception("Failed to retrieve data preview from Esca for id=%s", esca_id) | |
| return "Data preview 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/nodes/finalizer.py` around lines 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.
| prompt_name = getattr( | ||
| settings, "LANGFUSE_PROMPT_FINALIZER", "text2sql/finalizer" | ||
| ) | ||
| prompt_summary = ChatPromptTemplate.from_messages( | ||
| langfuse_prompt_summary.get_langchain_prompt() | ||
| langfuse_prompt = langfuse_client.get_prompt(prompt_name) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check the declared Langfuse prompt settings fields.
fd -t f 'config.py' agent/src | xargs rg -n 'LANGFUSE_PROMPT'Repository: StavPonte11/text2sql-onboarding
Length of output: 982
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== settings class files =="
fd -t f 'config.py' agent/src
echo
echo "== finalizer node context =="
fd -t f 'finalizer.py' agent/src/agent/nodes | xargs -r sed -n '1,130p'
echo
echo "== query_builder context =="
fd -t f 'query_builder.py' agent/src/agent/nodes | xargs -r sed -n '1,130p'Repository: StavPonte11/text2sql-onboarding
Length of output: 8643
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== all Langfuse prompt usages in nodes =="
fd -t f -i '*.py' agent/src/agent/nodes | xargs rg -n 'LANGFUSE_PROMPT|get_prompt|getattr' || true
echo
echo "== direct attribute access pattern for prompt names =="
fd -t f -i '*.py' agent/src/agent/nodes | xargs rg -n 'settings\.LANGFUSE_PROMPT|getattr\([^,]+,\s*"LANGFUSE' || trueRepository: StavPonte11/text2sql-onboarding
Length of output: 11398
Replace the fallback with direct attribute access for LANGFUSE_PROMPT_FINALIZER.
settings declares LANGFUSE_PROMPT_FINALIZER, so this getattr fallback is unnecessary. Use settings.LANGFUSE_PROMPT_FINALIZER to match the query_builder.py prompt-name access pattern.
🤖 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/finalizer.py` around lines 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.
| def make_literal(val_str: str) -> exp.Expression: | ||
| try: | ||
| if "." in val_str: | ||
| return exp.Literal.number(float(val_str)) | ||
| return exp.Literal.number(int(val_str)) | ||
| except ValueError: | ||
| return exp.Literal.string(val_str) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
make_literal turns numeric-looking category values into number literals.
refined_values holds candidate strings taken from a categorical value store. A candidate such as "444" or "52" becomes exp.Literal.number, so the generated predicate is place = 444 instead of place = '444'. Trino then compares a varchar column with a bigint and the query fails, or an implicit cast changes the match. A candidate such as "1.0" is also rewritten to a float and loses its original text.
The transformation plan replaces string category values only. Always emit string literals.
🐛 Proposed fix
def make_literal(val_str: str) -> exp.Expression:
- try:
- if "." in val_str:
- return exp.Literal.number(float(val_str))
- return exp.Literal.number(int(val_str))
- except ValueError:
- return exp.Literal.string(val_str)
+ # Candidate values come from a categorical string store, so the
+ # replacement literal must keep the string type of the column.
+ return exp.Literal.string(val_str)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def make_literal(val_str: str) -> exp.Expression: | |
| try: | |
| if "." in val_str: | |
| return exp.Literal.number(float(val_str)) | |
| return exp.Literal.number(int(val_str)) | |
| except ValueError: | |
| return exp.Literal.string(val_str) | |
| def make_literal(val_str: str) -> exp.Expression: | |
| # Candidate values come from a categorical string store, so the | |
| # replacement literal must keep the string type of the column. | |
| return exp.Literal.string(val_str) |
🤖 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 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".
| # Load the project's .env file automatically so users don't have to source it | ||
| env_path = os.path.join(os.path.dirname(__file__), "..", ".env") | ||
| load_dotenv(env_path) | ||
|
|
||
| if not os.environ.get("LANGFUSE_PUBLIC_KEY"): | ||
| os.environ["LANGFUSE_PUBLIC_KEY"] = "pk-lf-123" | ||
| os.environ["LANGFUSE_SECRET_KEY"] = "sk-lf-123" | ||
| os.environ["LANGFUSE_BASE_URL"] = "http://localhost:3000" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Set each Langfuse variable independently.
The guard checks only LANGFUSE_PUBLIC_KEY. If the developer .env defines LANGFUSE_PUBLIC_KEY but omits LANGFUSE_SECRET_KEY or LANGFUSE_BASE_URL, the fallbacks are skipped and the missing variables stay unset. Tests that read those variables then fail or target an unintended host.
🛠️ Proposed fix
-if not os.environ.get("LANGFUSE_PUBLIC_KEY"):
- os.environ["LANGFUSE_PUBLIC_KEY"] = "pk-lf-123"
- os.environ["LANGFUSE_SECRET_KEY"] = "sk-lf-123"
- os.environ["LANGFUSE_BASE_URL"] = "http://localhost:3000"
+os.environ.setdefault("LANGFUSE_PUBLIC_KEY", "pk-lf-123")
+os.environ.setdefault("LANGFUSE_SECRET_KEY", "sk-lf-123")
+os.environ.setdefault("LANGFUSE_BASE_URL", "http://localhost:3000")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Load the project's .env file automatically so users don't have to source it | |
| env_path = os.path.join(os.path.dirname(__file__), "..", ".env") | |
| load_dotenv(env_path) | |
| if not os.environ.get("LANGFUSE_PUBLIC_KEY"): | |
| os.environ["LANGFUSE_PUBLIC_KEY"] = "pk-lf-123" | |
| os.environ["LANGFUSE_SECRET_KEY"] = "sk-lf-123" | |
| os.environ["LANGFUSE_BASE_URL"] = "http://localhost:3000" | |
| # Load the project's .env file automatically so users don't have to source it | |
| env_path = os.path.join(os.path.dirname(__file__), "..", ".env") | |
| load_dotenv(env_path) | |
| os.environ.setdefault("LANGFUSE_PUBLIC_KEY", "pk-lf-123") | |
| os.environ.setdefault("LANGFUSE_SECRET_KEY", "sk-lf-123") | |
| os.environ.setdefault("LANGFUSE_BASE_URL", "http://localhost:3000") |
🤖 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/conftest.py` around lines 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.
| def patch_graph_infrastructure(): | ||
| """ | ||
| Patches all external I/O (Redis, Langfuse, ESCA) across the entire subgraph | ||
| to prevent network crashes during E2E testing. | ||
| """ | ||
| return ( | ||
| patch("agent.nodes.refiner.publish_node_event", new_callable=AsyncMock), | ||
| patch( | ||
| "agent.nodes.enrichment_orchestrator.publish_node_event", | ||
| new_callable=AsyncMock, | ||
| ), | ||
| patch("agent.nodes.refiner.langfuse_client"), | ||
| patch("agent.services.enrichment_orchestrator.langfuse_client"), | ||
| patch("agent.nodes.refiner.get_esca_client", MagicMock()), | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm the orchestrator module path and helper usage.
fd -t f 'enrichment_orchestrator.py' agent
rg -n 'patch_graph_infrastructure' agentRepository: StavPonte11/text2sql-onboarding
Length of output: 348
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== test helper =="
sed -n '1,80p' agent/tests/refiner/test_refiner_e2e_mocked.py
echo
echo "== orchestrator publish_node_event / imports =="
rg -n 'publish_node_event|from agent\.nodes\.enrichment_orchestrator import|import agent\.nodes\.enrichment_orchestrator|import importlib\.util|load_module_from_path|load_module' agent/src/agent/services/enrichment_orchestrator.py agent/src || true
echo
echo "== all patch_graph infrastructure references =="
rg -n 'patch_graph_infrastructure|agent\.nodes\.enrichment_orchestrator|agent\.services\.enrichment_orchestrator' .Repository: StavPonte11/text2sql-onboarding
Length of output: 10032
Remove the unused helper or fix its patch target.
patch_graph_infrastructure is not called by any test. The helper also returns unstarted patcher objects, so a caller gets no patching without with or .start(). The patch for agent.nodes.enrichment_orchestrator.publish_node_event targets a module that does not exist; use agent.services.enrichment_orchestrator.publish_node_event if this should remain.
🤖 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 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.
| assert final_state.get("is_satisfied") is True, ( | ||
| f"Failed to self-correct: {final_state.get('last_error') or final_state.get('escalation_reason')}" | ||
| ) | ||
| assert final_state.get("trino_error") is None | ||
|
|
||
| query = final_state["sql_query"].lower() | ||
|
|
||
| # Verify all issues were fixed | ||
| assert "limit 3" in query | ||
| assert "top 3" not in query | ||
| assert "coalesce" in query or "isnull" not in query | ||
| assert "123" in query |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Relax the strict token assertions in this live test.
The test asserts exact tokens in LLM-generated SQL, for example "limit 3" in query. A correct model output can use FETCH FIRST 3 ROWS ONLY or a different ordering and still satisfy the user request. This test then fails without a product defect. Assert the observable outcome instead: execution success, is_satisfied, absence of the SQL Server constructs, and a bounded row count.
🤖 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_real.py` around lines 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.
| import json | ||
| import pytest | ||
| from unittest.mock import AsyncMock, patch, MagicMock | ||
| from agent.nodes.finalizer import finalizer_node, get_esca_preview |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Add coverage for get_esca_preview and the Esca branch.
The file imports get_esca_preview, but no test exercises it. The Esca branch of finalizer_node and the "No data reference found." branch also stay uncovered. Add a test that enables ESCA_WRITE_ENABLED, sets raw_data_ref, mocks get_esca_client, and asserts the serialized preview. Add a second test for the missing-reference branch.
I can generate these tests if you want.
🤖 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/test_finalizer.py` at 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.
| if catalog: | ||
| return catalog | ||
|
|
||
| table_profiles = state.get("table_profiles") or [] |
There was a problem hiding this comment.
Can remove table_profiles since we don't use it anymore, we get all the data on catalog from the tool get_catalog_prompt in schema explorer
| execution_path = state.get("execution_path") or [] | ||
|
|
||
| sql = state.get("sql_query") | ||
| table_profiles = state.get("table_profiles") |
There was a problem hiding this comment.
Can remove the table_profiles here as well
| langfuse_prompt = langfuse_client.get_prompt(prompt_key) | ||
| prompt = ChatPromptTemplate.from_messages(langfuse_prompt.get_langchain_prompt()) | ||
| except Exception as e: | ||
| logger.warning(f"Could not fetch prompt '{prompt_key}' from Langfuse: {e}. Trying base refiner prompt.") | ||
| fallback_key = settings.LANGFUSE_PROMPT_REFINER | ||
| try: | ||
| result = await asyncio.to_thread(execute_query_sync, sql) | ||
| success = result.success | ||
| trino_error = result.error_message or "Unknown Trino error" | ||
| if not success: | ||
| error_history.append(trino_error) | ||
| except Exception as e: | ||
| success = False | ||
| trino_error = str(e) | ||
| error_history.append(trino_error) | ||
| result = None | ||
| langfuse_prompt = langfuse_client.get_prompt(fallback_key) | ||
| prompt = ChatPromptTemplate.from_messages(langfuse_prompt.get_langchain_prompt()) | ||
| except Exception as e2: | ||
| logger.error(f"Failed to load any refiner prompt from Langfuse: {e2}") | ||
| raise RuntimeError(f"Could not load refiner prompts from Langfuse: {e2}") from e2 |
There was a problem hiding this comment.
This is weird, why we fallback to other prompt if there is an error there?
I think if we can't get the prompt we need ,we will just throw and error so the user will see
|
|
||
| # 2. Short Table Names -> Fully Qualified Names | ||
| table_mappings: dict[str, str] = {} | ||
| table_profiles = state.get("table_profiles") or [] |
There was a problem hiding this comment.
remove table_profiles
| full_name = f'"{cat}"."{sch}"."{tbl}"' | ||
| table_mappings[tbl] = full_name | ||
|
|
||
| for short, full in table_mappings.items(): |
| return END | ||
|
|
||
| return "refiner" | ||
| def end_success_node(state: AgentState): |
There was a problem hiding this comment.
If it does nothing but just route to END, maybe just route to END instead of this
| @@ -0,0 +1 @@ | |||
| # agent.services package initialization | |||
There was a problem hiding this comment.
What is that comment
| volumes: | ||
| - ./core:/app/core | ||
| - ./backend:/app/backend | ||
| - /app/backend/.venv |
There was a problem hiding this comment.
Why mount them into the image? This is already in the image from the Dockerfile
| volumes: | ||
| - ./core:/app/core | ||
| - ./agent:/app/agent | ||
| - /app/agent/.venv |
There was a problem hiding this comment.
Why mount them into the image? This is already in the image from the Dockerfile
| return emb | ||
|
|
||
|
|
||
| def ingest_large_category_values(db_session: Session, profile_result: TableProfilingResult, batch_size: int | None = None): |
There was a problem hiding this comment.
This needs to be changed to use the large category of jeen_metadata
2 participants
Summary by CodeRabbit
New Features
Bug Fixes
Tests