Implement InferenceQueries for Postgres #5915
Open
Conversation
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
f098854 to
f98ca8b
Compare
7c0807d to
a662ea1
Compare
f98ca8b to
73bf4fa
Compare
a662ea1 to
3d54063
Compare
73bf4fa to
b30f8b4
Compare
3d54063 to
cdfd5a9
Compare
fd9a7b3 to
a1c96c6
Compare
b30f8b4 to
e5d90b3
Compare
64cacd9 to
6e9cc21
Compare
a1c96c6 to
9e2cdae
Compare
5bcdb93 to
cb64a5e
Compare
4c85997 to
b8838e8
Compare
cb64a5e to
ac7926c
Compare
ac7926c to
43cfaef
Compare
Member
Author
|
@codex review |
|
Codex review is not enabled for this repo. Please contact the admins of this repo to enable Codex. |
Member
Author
|
@claude review |
937a420 to
033b112
Compare
Member
Author
|
@BugBot review |
21 tasks
033b112 to
5420511
Compare
b1f86f3 to
aabceaa
Compare
Uncomment tests Implement delegating trait Push order and limit down to postgres
aabceaa to
4d4abc2
Compare
virajmehta
reviewed
Jan 29, 2026
|
|
||
| fn apply_demonstration_feedback_filter( | ||
| query_builder: &mut QueryBuilder<sqlx::Postgres>, | ||
| df: &crate::endpoints::stored_inferences::v1::types::DemonstrationFeedbackFilter, |
virajmehta
reviewed
Jan 30, 2026
Comment on lines
+122
to
+130
| // Use EXISTS subquery to filter by metric value | ||
| query_builder | ||
| .push("EXISTS (SELECT 1 FROM tensorzero.boolean_metric_feedback f WHERE f.target_id = "); | ||
| query_builder.push(join_column); | ||
| query_builder.push(" AND f.metric_name = "); | ||
| query_builder.push_bind(bm.metric_name.clone()); | ||
| query_builder.push(" AND f.value = "); | ||
| query_builder.push_bind(bm.value); | ||
| query_builder.push(")"); |
Member
There was a problem hiding this comment.
assuming this type of predicate is indexed is this fast?
virajmehta
reviewed
Jan 30, 2026
| let operator = tag.comparison_operator.to_postgres_operator(); | ||
|
|
||
| // For Postgres JSONB, we use the ->> operator to extract text and compare | ||
| // We also check that the key exists to handle != correctly |
Member
There was a problem hiding this comment.
let's update the comment to say what the correct handling is (the tag must exist with a different value)
virajmehta
reviewed
Jan 30, 2026
Comment on lines
+283
to
+288
| SELECT DISTINCT ON (target_id) | ||
| target_id, | ||
| value | ||
| FROM {table_name} | ||
| WHERE metric_name = '{metric_name}' | ||
| ORDER BY target_id, created_at DESC |
Member
There was a problem hiding this comment.
will PG know to push down the join into this predicate?
virajmehta
reviewed
Jan 30, 2026
| // Convert JSONB to string | ||
| Ok(result.map(|r| r.output.to_string())) | ||
| } | ||
| } |
Member
There was a problem hiding this comment.
aren't we just gonna re-parse on the outside of this method?
virajmehta
reviewed
Jan 30, 2026
Comment on lines
+537
to
+550
| let input_json = serde_json::to_value(&row.input).unwrap_or_default(); | ||
| let output_json = serde_json::to_value(&row.output).unwrap_or_default(); | ||
| let inference_params_json = | ||
| serde_json::to_value(&row.inference_params).unwrap_or_default(); | ||
| let tags_json = serde_json::to_value(&row.tags).unwrap_or_default(); | ||
| let extra_body_json = serde_json::to_value(&row.extra_body).unwrap_or_default(); | ||
|
|
||
| let SerializedToolParams { | ||
| dynamic_tools, | ||
| dynamic_provider_tools, | ||
| allowed_tools, | ||
| tool_choice, | ||
| parallel_tool_calls, | ||
| } = serialize_tool_params(row.tool_params.as_ref()).unwrap_or_default(); |
Member
There was a problem hiding this comment.
all these unwrap_or_defaults() are bad. Would like to propagate errors if possible vs failing silently.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Sorry about the large PR!
This implements the main set of inference queries: list/get inferences, insert inferences, count inferences, and so on.
There's some duplication across InferenceQueries and InferenceCountQueries - will remove the InferenceCountQueries later. (Those were directly ported from the frontend.)
There's a small risk of SQL injection with metric names in joins, because of sqlx QueryBuilder interface limitations. We should consider revisiting it but because metric names are controlled by the clients, this is not a huge risk in practice (malicious users can't provide arbitrary metric names).
Step towards #5691.
Note
Medium Risk
Medium risk due to adding a new Postgres query/insert path for inferences and wiring it into live inference and batch write flows (with new pagination/order/filter logic and a schema migration). Potential correctness/performance issues or edge cases in query building (notably metric ORDER BY joins) could affect inference listing/counting and background writes.
Overview
Implements
InferenceQueriesfor Postgres, including listing/counting inferences acrosschat_inferencesandjson_inferences, fetching function info/output/tool params/output schema, and batch inserting inference rows.Refactors pagination validation into
ListInferencesParams::validate_pagination()and updates ClickHouse queries to use it, while adding a Postgres-specific filter/ORDER BY builder (including metric-based ordering via JOINs) and a migration to convertchat_inferences.tool_choiceto JSONB.Wires inference writes in
/inferenceand batch inference completion throughDelegatingDatabaseConnectionto always write to ClickHouse and optionally dual-write to Postgres behind feature flags, and adds/updates e2e tests plus newsqlxquery metadata.Written by Cursor Bugbot for commit 033b112. This will update automatically on new commits. Configure here.