feat(wren-ui): adjustment API (ai-env-changed)#1501
feat(wren-ui): adjustment API (ai-env-changed)#1501wwwy3y3 merged 22 commits intoepic/adjust-answerfrom
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the WalkthroughThis pull request introduces a series of changes that add adjustment functionality to thread responses. A new migration creates an Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Resolver
participant Service
participant Repository
participant BackgroundTracker
participant DB
Client->>Resolver: Mutation adjustThreadResponse(responseId, data)
Resolver->>Service: adjustThreadResponseAnswer()/with SQL method
Service->>Repository: Retrieve thread response
Repository->>DB: Query ThreadResponse table (with adjustment column)
DB-->>Repository: Return thread response data
Service->>BackgroundTracker: Initiate adjustment task
BackgroundTracker->>DB: Poll for task updates
DB-->>BackgroundTracker: Return adjustment status/result
BackgroundTracker-->>Service: Adjustment task result
Service-->>Resolver: Return updated ThreadResponse
Resolver-->>Client: Response with adjustment details
Possibly related PRs
Poem
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
wwwy3y3
requested review from
andreashimin,
Copilot,
fredalai and
onlyjackfrost
and removed request for
Copilot and
onlyjackfrost
There was a problem hiding this comment.
Pull Request Overview
This PR introduces a new adjustment API feature for thread responses, enabling users to either provide feedback on SQL generation reasoning or directly apply SQL changes. The changes include the addition of new GraphQL types and resolvers, updates to service layers and repositories to support adjustment data, and corresponding migrations and telemetry events.
Reviewed Changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| wren-ui/src/common.ts | Injects the new askingTaskRepository dependency. |
| wren-ui/src/apollo/server/telemetry/telemetry.ts | Adds a new telemetry event for thread response adjustments. |
| wren-ui/src/apollo/server/services/askingTaskTracker.ts | Updates the return object to explicitly cast taskRecord.detail as AskResult. |
| wren-ui/src/apollo/server/services/askingService.ts | Adds new methods for adjustment operations (apply SQL, provide reasoning, cancel, rerun, and retrieve tasks). |
| wren-ui/src/apollo/server/schema.ts | Introduces new GraphQL types, inputs, enums, and mutations/queries for adjustments. |
| wren-ui/src/apollo/server/resolvers/askingResolver.ts | Adds resolvers for adjustment operations and integrates them into the existing API. |
| wren-ui/src/apollo/server/repositories/threadResponseRepository.ts | Updates thread response repository to store and parse adjustment data. |
| wren-ui/src/apollo/server/repositories/askingTaskRepository.ts | Extends the asking task detail type to include adjustment information. |
| wren-ui/src/apollo/server/models/adaptor.ts | Adds AskFeedback types and interfaces to support the new adjustment functionality. |
| Other files (backgrounds, wrenAIAdaptor, migrations) | Support the new adjustment feature by introducing relevant background tracking, API calls, and schema migrations. |
Comments suppressed due to low confidence (2)
wren-ui/src/apollo/server/services/askingTaskTracker.ts:407
- [nitpick] Consider ensuring that the explicit cast to AskResult accurately reflects the structure of taskRecord.detail to avoid potential type mismatches.
...(taskRecord?.detail as AskResult),
wren-ui/src/apollo/server/resolvers/askingResolver.ts:725
- [nitpick] Verify that using parent.askingTaskId to retrieve the adjustment task is appropriate and clearly distinguishes between adjustment and standard asking tasks.
parent.askingTaskId,
|
@coderabbitai summary |
✅ Actions performedSummary regeneration triggered. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (10)
wren-ui/src/apollo/server/backgrounds/adjustmentBackgroundTracker.ts (6)
17-18: Consider using environment-based logging levels.
By settinglogger.level = 'debug'unconditionally, there is a risk of exposing detailed logs in production, which may include sensitive or verbose output. It is often preferable to load log levels from environment variables or a configuration file to ensure that debug logs are only enabled in development.
61-93: Validate polling strategy in constructor.
Starting the polling in the constructor may be desired behavior, but it can lead to potential race conditions if other dependencies are not fully initialized. Consider offering a separateinit()method or a factory to ensure controlled startup, especially if future expansions require dependency readiness checks.
229-245: Avoid partial memory-database discrepancy.
getAdjustmentResultfirst checks the in-memory map, then falls back to the database. This approach is fine, but if a task is canceled or removed from memory, you may want to ensure the database record also reflects that status to avoid stale entries if the user queries the task again later.🧰 Tools
🪛 Biome (1.9.4)
[error] 235-235: Change to an optional chain.
Unsafe fix: Change to an optional chain.
(lint/complexity/useOptionalChain)
258-260: Memory cleanup on cancellation.
cancelAdjustmentTaskdoes not remove the canceled task from the in-memory tracking maps. If a user repeatedly creates and cancels tasks, the in-memory store could accumulate stale tasks unless the tasks are removed in a follow-up operation.
274-368: Parallel polling might overwhelm external services.
Promise.allSettledfor all tasks is efficient, but if the number of tasks scales high, it could stress the AI service or database. Consider throttling or batching the jobs in smaller chunks to avoid performance bottlenecks or rate limits.
440-450: Potential for extended change detection.
isResultChangedcurrently checks only status diffs. If the response contents (like text or partial results) can evolve without changing status, consider expanding the check to detect those changes too.wren-ui/src/apollo/server/schema.ts (2)
659-663: Clear input for adjustments.
AdjustThreadResponseInputeffectively captures the key data required for either explaining reasoning (tables,sqlGenerationReasoning) or applying direct SQL changes (sql). Ensure all fields are optional or validated if certain keys only apply to specific adjustment types.
721-724: Payload typed as JSON.
Storing the payload inJSONallows maximum flexibility. If you need stricter contracts, consider using custom scalar types or inline object fields. However,JSONis often the most practical choice for rapid iteration.wren-ui/src/apollo/server/resolvers/askingResolver.ts (1)
468-501: Check for consistent naming and error handling.
The logic flow is correct: ifdata.sqlis present, apply it directly; otherwise, create a feedback-based adjustment. Consider adding defensive checks to ensuredata.sqlordata.tablesis valid.wren-ui/src/apollo/server/services/askingService.ts (1)
1062-1085:adjustThreadResponseWithSQLdirectly clones the context.
While the creation of a newThreadResponseis correct, be sure it's clear to users that any subsequent modifications are not reflected in the original.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (14)
wren-ui/migrations/20250510000000_add_adjustment_to_thread_response.js(1 hunks)wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts(4 hunks)wren-ui/src/apollo/server/backgrounds/adjustmentBackgroundTracker.ts(1 hunks)wren-ui/src/apollo/server/backgrounds/index.ts(1 hunks)wren-ui/src/apollo/server/models/adaptor.ts(1 hunks)wren-ui/src/apollo/server/repositories/askingTaskRepository.ts(1 hunks)wren-ui/src/apollo/server/repositories/threadResponseRepository.ts(6 hunks)wren-ui/src/apollo/server/resolvers.ts(2 hunks)wren-ui/src/apollo/server/resolvers/askingResolver.ts(6 hunks)wren-ui/src/apollo/server/schema.ts(5 hunks)wren-ui/src/apollo/server/services/askingService.ts(9 hunks)wren-ui/src/apollo/server/services/askingTaskTracker.ts(1 hunks)wren-ui/src/apollo/server/telemetry/telemetry.ts(1 hunks)wren-ui/src/common.ts(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
wren-ui/src/apollo/server/repositories/threadResponseRepository.ts (1)
Learnt from: andreashimin
PR: Canner/WrenAI#1117
File: wren-ui/src/apollo/server/repositories/threadResponseRepository.ts:37-37
Timestamp: 2025-03-27T09:01:59.977Z
Learning: In the ThreadResponseRepository, the `chartType` property is intentionally typed as string in the database layer, while the ChartType enum is used specifically for AI adaptor input/output in the application layer.
🧬 Code Definitions (4)
wren-ui/src/apollo/server/services/askingTaskTracker.ts (1)
wren-ui/src/apollo/server/models/adaptor.ts (1)
AskResult(119-134)
wren-ui/src/apollo/server/repositories/askingTaskRepository.ts (2)
wren-ui/src/apollo/server/models/adaptor.ts (2)
AskResult(119-134)AskFeedbackResult(313-321)wren-ui/src/apollo/server/resolvers/askingResolver.ts (1)
AskingTask(51-65)
wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts (1)
wren-ui/src/apollo/server/models/adaptor.ts (3)
AskFeedbackInput(296-302)AsyncQueryResponse(68-70)AskFeedbackResult(313-321)
wren-ui/src/apollo/server/resolvers/askingResolver.ts (3)
wren-ui/src/apollo/server/models/adaptor.ts (1)
WrenAIError(5-8)wren-ui/src/apollo/server/repositories/threadResponseRepository.ts (1)
ThreadResponse(65-76)wren-ui/src/apollo/client/graphql/__types__.ts (1)
ThreadResponse(1201-1212)
🪛 Biome (1.9.4)
wren-ui/src/apollo/server/backgrounds/adjustmentBackgroundTracker.ts
[error] 235-235: Change to an optional chain.
Unsafe fix: Change to an optional chain.
(lint/complexity/useOptionalChain)
⏰ Context from checks skipped due to timeout of 90000ms (2)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (go)
🔇 Additional comments (50)
wren-ui/src/apollo/server/backgrounds/index.ts (1)
3-3: LGTM: Properly exposing the new background trackerThe export follows the established pattern, making the adjustment background tracker available to the rest of the application.
wren-ui/src/common.ts (1)
131-131: LGTM: Appropriate dependency injection for AskingServiceAdding the
askingTaskRepositoryto the AskingService constructor enables the service to directly interact with adjustment tasks, consistent with the PR objective of supporting thread response adjustments.wren-ui/src/apollo/server/telemetry/telemetry.ts (1)
59-60: LGTM: Added telemetry event for tracking adjustment actionsThe new telemetry event follows naming conventions and will properly track user interactions with the adjustment feature. Good organizational structure with the dedicated comment section.
wren-ui/src/apollo/server/services/askingTaskTracker.ts (1)
407-407: Added type safety for task detail propertyThe explicit type cast to
AskResultensures that the properties being spread fromtaskRecord?.detailconform to the expected type structure, improving type safety in the returned object.wren-ui/migrations/20250510000000_add_adjustment_to_thread_response.js (2)
1-14: Database migration looks well-structured and appropriately documented.The
upfunction correctly adds a nullable JSONB column calledadjustmentwith a clear comment describing its purpose. Using JSONB is a good choice for storing structured adjustment data that may vary in format.
16-24: Rollback function properly handles column removal.The
downfunction correctly implements the rollback logic to drop theadjustmentcolumn if needed.wren-ui/src/apollo/server/models/adaptor.ts (3)
295-302: The AskFeedbackInput interface captures all necessary fields for the adjustment API.The interface includes all the required parameters for providing feedback on SQL generation:
- tables selection
- SQL generation reasoning
- the SQL itself
- project identifier
- optional configuration settings
This aligns well with the PR objective of supporting adjustments to thread responses.
304-311: The AskFeedbackStatus enum follows established patterns in the codebase.The status enum maintains consistency with other status enums in the file (particularly AskResultStatus). The progression from UNDERSTANDING through GENERATING, CORRECTING, to terminal states like FINISHED, FAILED, or STOPPED follows a logical flow for background task management.
313-321: AskFeedbackResult interface properly structures the adjustment response data.The interface includes all necessary fields for tracking and displaying the result of an adjustment operation. The response array is correctly typed to only allow LLM-generated SQL adjustments, which matches the PR objective of allowing users to adjust SQL generation.
wren-ui/src/apollo/server/resolvers.ts (2)
37-39: Query resolver for adjustment task retrieval is properly added.The resolver for fetching adjustment task status is appropriately placed in the Query section with a clear header comment indicating its purpose.
103-107:Details
❓ Verification inconclusive
Mutation resolvers for adjustment operations follow consistent naming patterns.
The three mutation resolvers for adjustment operations (adjust, cancel, rerun) are well-organized under a dedicated "Adjustment" section header. This follows the established pattern in the codebase for grouping related resolvers.
However, there appears to be some inconsistency in the method naming. While
adjustThreadResponsefollows the standard pattern, the cancel and rerun methods usecancelAdjustThreadResponseAnswerandrerunAdjustThreadResponseAnswerwhich include "Answer" in their names, suggesting they might be specifically for answer adjustments rather than general thread response adjustments.
🏁 Script executed:
#!/bin/bash # Check if there are other adjustment-related resolver methods in the codebase # to verify the naming convention rg -A 2 -B 2 "adjustThreadResponse|cancelAdjust|rerunAdjust" --type tsLength of output: 17513
Action: Clarify Naming Consistency for Adjustment Resolvers
The adjustment-related GraphQL mutations are neatly grouped, but there’s a potential naming inconsistency worth verifying:
- The resolver for adjusting a thread response is named
adjustThreadResponse, whereas the cancellation and rerun operations are implemented ascancelAdjustThreadResponseAnswerandrerunAdjustThreadResponseAnswer(note the extra “Answer” suffix).- The GraphQL schema, however, exposes these as
cancelAdjustmentTaskandrerunAdjustmentTask, which do not mention “Answer.”- Please confirm whether the “Answer” suffix is intentional to indicate that these methods specifically handle answer adjustments or if they should be renamed for consistency with
adjustThreadResponse.wren-ui/src/apollo/server/repositories/askingTaskRepository.ts (3)
10-11: Import statement correctly includes the new AskFeedbackResult type.The import has been updated to include AskFeedbackResult alongside the existing AskResult, correctly supporting the union type defined below.
12-16: New AskingTaskDetail union type properly extends the data model.The union type approach is a good design choice that maintains backward compatibility with existing code while adding support for the new adjustment functionality. The inclusion of an optional
adjustmentboolean flag provides a clear indicator of whether the task detail represents an adjustment task.
22-22: The AskingTask interface is properly updated to use the new union type.Updating the
detailproperty to use the new AskingTaskDetail union type ensures that the repository can handle both regular AskResult objects and the new AskFeedbackResult objects with adjustment data.wren-ui/src/apollo/server/backgrounds/adjustmentBackgroundTracker.ts (4)
49-59: Interface clarity is good.
TheIAdjustmentBackgroundTaskTrackerinterface clearly defines the methods needed to manage adjustment tasks, promoting clean architecture and making it easier to mock in tests.
95-153: Robust error handling enhances reliability.
ThecreateAdjustmentTaskmethod properly wraps AI requests and database operations in a try/catch block. Consider including more contextual information (e.g., input parameters) when logging errors to aid debugging.
159-227: Potential concurrency considerations for rerun.
When rerunning an adjustment task, multiple users or processes could trigger concurrent reruns on the samethreadResponseId. Ensure the database logic (e.g., re-updating the same asking task) has sufficient concurrency controls, or confirm that your usage patterns allow for safe updates without row locking or version checks.
370-380: Finalizing thread response logic is concise.
updateThreadResponseWhenTaskFinalizedelegantly updates only the SQL field when a final result is available. This approach maintains minimal changes. Good job ensuring no extra fields are overwritten inadvertently.wren-ui/src/apollo/server/repositories/threadResponseRepository.ts (4)
42-45: Enum usage promotes clarity.
DefiningThreadResponseAdjustmentTypeas an enum strongly communicates the possible adjustment types and reduces the chance of typos or mismatches.
94-95: JSON column registration is consistent.
Adding'adjustment'tojsonbColumnskeeps data parsing consistent across repository methods. Maintaining a single array for JSON columns helps ensure new fields won't be missed when reading/writing from the DB.
130-133: Safe JSON parsing on loaded data.
The logic to parseres.adjustmentif it's a string preserves schema integrity. Good job ensuring the data is not double-parsed. Usingadjustment: adjustment || nullalso helps avoid undefined references in the application layer.Also applies to: 139-139
153-153: Data alignment between domain and storage.
Marking theadjustmentproperty inThreadResponseand converting it to a JSON string onupdateOneensures database consistency. This approach follows the repository pattern neatly, keeping the domain model and DB representation distinct.Also applies to: 170-170
wren-ui/src/apollo/server/schema.ts (5)
716-719: Exposing adjustment types to the client.
GraphQL enumThreadResponseAdjustmentTypematches your backend enum, providing a straightforward contract for clients. This consistency is beneficial for strongly-typed or code-generated clients.
726-732: Adjustment task exposure.
DefiningAdjustmentTaskin the schema clarifies real-time or eventual result retrieval for building polling UIs. Be sure to handle potential partial data (e.g.,sqlmight be absent if reasoning is not complete).
744-745: Aggregating all response data is user-friendly.
By including bothadjustmentandadjustmentTaskwithinThreadResponse, clients can retrieve the relevant details in a single query, simplifying client-side state management.
984-986: Separate query foradjustmentTask.
Allowing direct retrieval ofadjustmentTaskbytaskIdis helpful. Double check thattaskIdis unique across tasks, or else ensure additional fields for disambiguation.
1098-1105: Comprehensive mutation set.
Allowing adjusting, canceling, and rerunning tasks covers all major user flows. This separation of concerns also allows fine-grained control of user permissions, should you need to restrict certain actions.wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts (3)
33-35: Imports align well with new features.
No issues found here; these imports are consistent with the rest of the file changes.
125-131: Doc comments for new ask feedback methods look good.
This aligns with the existing style of the interface; consistent docstrings help maintain clarity.
877-878: Status union extension is appropriate.
EnsuringAskFeedbackStatusis recognized intransformStatusAndErrorhelps unify status handling across all feedback tasks.wren-ui/src/apollo/server/resolvers/askingResolver.ts (8)
8-8: The new import is consistent.
AskFeedbackStatusis required for adjustment tasks; no issues here.
43-49: NewAdjustmentTaskinterface enhances clarity.
This interface straightforwardly captures the essential fields for feedback tasks.
121-126: Constructor bindings are well-structured.
Binding new methods ensures they can be correctly referenced across the resolver. No issues spotted.
503-512: Cancellation logic appears straightforward.
This method properly forwards cancellation to the service; no further concerns.
514-530: Rerun flow is consistent with existing patterns.
Calling back intoaskingService.rerunAdjustThreadResponseAnswerkeeps logic centralized.
532-547:getAdjustmentTaskshapes a user-friendly response.
By mapping the service response intoAdjustmentTask, client-side usage is simplified. No issues found.
705-707: Conditional return helps avoid conflicts.
SkippingaskingTaskifadjustmentexists clarifies the code path and prevents race conditions between tasks.
715-735:adjustmentTasknested resolver is consistent.
Properly returns the newAdjustmentTaskdata; no major concerns.wren-ui/src/apollo/server/services/askingService.ts (12)
12-12: AddingProjectConfigurationsis appropriate.
This import is necessary for typed project-level configs.
20-20:ThreadResponseAdjustmentTypeimport clarifies constants usage.
Provides clear semantic meaning for the new adjustment type.
30-30: Importing repositories andProjectis aligned with service expansions.
Interfaces fromrepositoriesandProjectunify cross-service data usage.
41-41: New background trackers are well-integrated.
IncludingAdjustmentBackgroundTaskTrackerandTrackedAdjustmentResultfits the pattern of other background trackers.
418-418:private askingTaskRepositoryfield is a good addition.
Encourages internal referencing of the repository, consistent with other service fields.
419-419:adjustmentBackgroundTrackerensures new tasks remain decoupled.
This reference keeps the adjustment logic encapsulated inside a dedicated tracker.
485-488: Properly instantiatingadjustmentBackgroundTracker.
Dependencies likewrenAIAdaptor,askingTaskRepository, andthreadResponseRepositoryare passed, ensuring centralized management for feedback tasks.
1087-1112:adjustThreadResponseAnswerdelegates tasks to the background tracker.
This is well-structured for asynchronous feedback processing. Make sure all exceptions from the tracker are properly propagated or handled.
1114-1117:cancelAdjustThreadResponseAnsweris straightforward.
This properly calls the background tracker’s cancellation method. No issues.
1119-1142: Rerunning adjustments inherits from existing design.
This approach aligns with the patterns used for rerunning other tasks. Good consistency.
1144-1148:getAdjustmentTaskretrieval logic is consistent.
No concerns. The tracker call is encapsulated well.
1150-1154:getAdjustmentTaskByIdcomplementsgetAdjustmentTask.
Both ID-based and task-based retrieval are useful. This finalizes the adjustment workflow.
wwwy3y3
force-pushed
the
feature/adjustment-api
branch
from
6242c9f to
d20c370
Compare
github-actions
bot
changed the title
andreashimin
force-pushed
the
feature/adjustment-api
branch
from
d0e8a51 to
b7556bc
Compare
Description
Thread Response Adjustment
This PR introduces a new feature that allows users to adjust thread responses by providing feedback on SQL generation reasoning or directly applying SQL modifications.
Features
Added support for two types of adjustments:
Background task tracking for adjustment processes with status updates
Support for canceling and rerunning adjustment tasks
Persistence of adjustment data in the database
API Examples
1. Adjust thread response by providing feedback on SQL generation reasoning
payload
Bring in thread response id,
sqlGenerationReasoning&tables2. Tracking status
payload
response
3. Thread response query
adjustmentTaskfield.adjustmentrecorded what changes did users makeadjustment.payload.typeshows what type of adjustment users made. If thru reasoning adjustments, type would beREASONINGadjustment.payload.originalThreadResponseIdstands for the thread response we make adjustment based on.Response
4. cancel
payload
5. rerun
payload
6. Apply SQL directly to do adjustment
adjustment.payload.typeshows what type of adjustment users made. If thru reasoning adjustments, type would beAPPLY_SQLadjustment.payload.originalThreadResponseIdstands for the thread response we make adjustment based on.payload
thread response query will be:
Summary by CodeRabbit