feat(wren-ai-service): Add Instructions for SQL Generation#1376
Conversation
WalkthroughThis pull request introduces new pipeline entries for handling instructions in the Wren AI service. It adds two new entries— Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant C as Client
participant R as FastAPI Router (/instructions)
participant S as InstructionsService
participant P as Instructions Pipeline
participant DS as Document Store (qdrant)
C->>R: POST /instructions with instructions payload
R->>S: Trigger indexing via background task
S->>P: Invoke instructions_indexing pipeline
P->>P: Convert instructions to documents
P->>P: Generate embeddings with litellm_embedder.default
P->>DS: Write documents to qdrant
P-->>S: Return indexing result
S-->>R: Update event status in TTL cache
R-->>C: Respond with event ID
sequenceDiagram
autonumber
participant C as Client
participant A as AskService
participant IR as Instructions Retrieval Pipeline
participant DS as Document Store (qdrant)
participant SQL as SQL Generation Pipeline
C->>A: Send query with project ID
A->>IR: Call instructions_retrieval pipeline
IR->>DS: Retrieve documents using qdrant
IR-->>A: Return retrieved instructions
A->>SQL: Pass instructions along with SQL samples
SQL-->>A: Generate SQL query response
A-->>C: Return final response
Possibly related PRs
Tip ⚡🧪 Multi-step agentic review comment chat (experimental)
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms (3)
🪧 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 (
|
paopa
force-pushed
the
feat/instruction-pipeline-and-endpoints
branch
from
9127e1a to
da50bc0
Compare
paopa
added
module/ai-service
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (16)
wren-ai-service/src/pipelines/generation/utils/sql.py (1)
552-558: Consider addressing the TODO commentThere's a TODO comment about refactoring the format of the instructions. Consider specifying what refactoring is needed or creating a follow-up task for this refactoring.
- # todo: refactor the format of the instructions + # TODO: Refactor the format of the instructions to support structured metadata and additional propertieswren-ai-service/src/web/v1/services/ask.py (1)
367-375: Implementation of instructions retrieval looks good with minor typo.The code successfully retrieves instructions using the instructions_retrieval pipeline, but there's a typo in the first TODO comment ("retireve" should be "retrieve").
-# todo: consider to retireve at the same time with sql_samples +# todo: consider to retrieve at the same time with sql_sampleswren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py (1)
21-21: Use a unique pipeline instance for each test to prevent cross-test interference.While the tests appear to be functioning correctly, consider creating a fresh pipeline instance inside each test function (or as a fixture) to avoid potential cross-test interference when tests are run concurrently or in parallel.
wren-ai-service/src/web/v1/services/instructions.py (2)
34-41: Be mindful of large TTL cache sizes.
TTLCacheis configured withmaxsize=1_000_000, which may lead to high memory usage in long-running environments. Ensure your system capacity can handle this or consider a smaller cache limit if usage patterns do not require storing that many events.
78-86: Skip indexing instructions with empty questions if desired.Currently, each question spawns a separate instruction for indexing. If empty or trivial questions are present, it might create redundant documents. Consider filtering them out if your use case requires it.
wren-ai-service/tests/pytest/services/test_instructions.py (1)
13-25: Recreate index in a fixture teardown to ensure isolation.While
recreate_index=Trueis used, consider a teardown step that consistently wipes the index after each test to avoid accidental state leaks if future tests mutate the store without re-creating.wren-ai-service/src/web/v1/routers/instructions.py (2)
88-108: Ensure thread safety when storing events
Storingevent_idreferences in theservicedictionary could expose race conditions in high-concurrency scenarios. If multiple indexing requests come in simultaneously, consider using an asynchronous or thread-safe data structure/persistence to avoid data overwrites.
116-139: Return more specific HTTP status codes
Currently, the endpoint sets a 500 status code whenevent.status == "failed". In certain failures (e.g., invalid IDs), a 400 or 404 might be more appropriate. Consider refining error-handling to reflect the nature of failures accurately.wren-ai-service/src/pipelines/indexing/instructions.py (4)
20-25: Clarify the role ofquestionvs.instruction
TheInstructionmodel includes bothinstructionandquestionfields. Explain or rename them to avoid ambiguity. It’s not entirely clear how each field is used downstream.
28-51: Validateinstructionsbefore document conversion
TheInstructionsConverterreturnsDocumentobjects with minimal validation. If instructions are malformed, the indexing pipeline may fail downstream. Consider adding guards or skipping invalid instructions.
80-87: Allow for empty instructions
Wheninstructionslist is empty, the pipeline returns an empty dict quietly. Confirm if this is the correct behavior or if a warning should be logged for visibility.
123-177: Add coverage for error-handling paths
TheInstructionspipeline has a well-defined flow for indexing but lacks explicit error-handling logic if components fail. Consider adding try/except blocks or centralized error handling to prevent partial writes or inconsistent states.wren-ai-service/src/pipelines/retrieval/instructions.py (4)
18-35: Handle empty or inconsistent metadata
When constructing outputs, confirm that the documents’metafields (e.g.,"instruction","instruction_id") have valid contents. An unexpectedNonecould propagate. Optionally, add fallback values or warning logs.
88-103: Expose reason for filtered documents
filtered_documentscan remove many relevant results ifsimilarity_thresholdortop_kis strict. Consider logging how many documents were filtered out, for transparency in debugging.
127-139: Combine results more transparently
formatted_outputmerges default instructions with retrieved documents. Log or document the final ordering logic in case the user expects default instructions to appear first or last.
144-171: Standardize pipeline usage
TheInstructionsretrieval pipeline is straightforward, but ensure that the naming of pipeline steps matches the indexing pipeline. Differences in naming or usage across indexing vs. retrieval might confuse maintainers.Would you like help aligning naming conventions between indexing and retrieval?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (27)
deployment/kustomizations/base/cm.yaml(2 hunks)docker/config.example.yaml(2 hunks)wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/src/config.py(1 hunks)wren-ai-service/src/globals.py(4 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(1 hunks)wren-ai-service/src/pipelines/indexing/__init__.py(2 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(1 hunks)wren-ai-service/src/pipelines/indexing/sql_pairs.py(2 hunks)wren-ai-service/src/pipelines/retrieval/__init__.py(2 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)wren-ai-service/src/web/v1/routers/__init__.py(2 hunks)wren-ai-service/src/web/v1/routers/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/__init__.py(2 hunks)wren-ai-service/src/web/v1/services/ask.py(3 hunks)wren-ai-service/src/web/v1/services/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/semantics_preparation.py(2 hunks)wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py(1 hunks)wren-ai-service/tests/pytest/services/test_instructions.py(1 hunks)wren-ai-service/tools/config/config.example.yaml(2 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (2)
- GitHub Check: pytest
- GitHub Check: pytest
🔇 Additional comments (49)
wren-ai-service/docs/config_examples/config.groq.yaml (1)
127-132: New Pipeline Entries for Instructions Indexing and RetrievalThe new entries for
instructions_indexingandinstructions_retrievalare correctly defined with an embedder set tolitellm_embedder.defaultand a document store ofqdrant. Their structure is consistent with the other pipeline definitions in this configuration file. Please ensure that any related runtime logic or documentation references these entries appropriately.wren-ai-service/docs/config_examples/config.deepseek.yaml (1)
145-150: Consistent Addition of Instructions PipelinesThe additions of
instructions_indexingandinstructions_retrievalusinglitellm_embedder.defaultanddocument_store: qdrantare properly integrated into the pipelines section. This change aligns well with the configuration style observed in other parts of the file. Consider verifying that any changes in pipeline behavior due to these new entries are covered by your integration tests.wren-ai-service/tools/config/config.full.yaml (1)
152-157: Integration of Instructions Pipeline in Full ConfigurationThe configuration now includes two new pipeline entries,
instructions_indexingandinstructions_retrieval, which are both correctly configured with the expected embedder and document store. Their placement within the full configuration file maintains consistency with other similar pipeline entries. Ensure that the documentation and any related configuration references are updated accordingly.wren-ai-service/docs/config_examples/config.google_ai_studio.yaml (1)
134-139: Addition of Instructions Pipelines for Google AI Studio ConfigThe new pipeline entries for
instructions_indexingandinstructions_retrievalare implemented with the same consistent parameters (litellm_embedder.defaultandqdrant). The additions are clear and follow the expected YAML syntax. Please double-check that any service consuming these pipelines is aware of the new configuration options.wren-ai-service/tools/config/config.example.yaml (1)
154-159: Incorporating Instructions Pipelines into the Example ConfigurationThe new entries for
instructions_indexingandinstructions_retrievalhave been added in a manner consistent with the rest of the pipeline entries. They utilize the same embedder and document store settings as other similar pipelines. This consistency helps ensure that the new SQL generation instruction capabilities integrate smoothly across different deployments. Verify that the example configuration documentation reflects these changes.docker/config.example.yaml (3)
5-30: Well-structured model configurationThe restructuring of the
modelssection under thellmtype to use a hyphen-prefixed list format improves readability and standardizes the YAML structure. This format is more consistent with YAML best practices.
37-41: Consistent formatting for embedder modelsThe embedder model configuration follows the same improved structure as the LLM models, maintaining consistency throughout the configuration file.
136-141: New instructions pipeline entries align with PR objectivesThe addition of
instructions_indexingandinstructions_retrievalpipelines successfully implements the core functionality described in the PR objectives. These entries correctly use thelitellm_embedder.defaultembedder and theqdrantdocument store, maintaining consistency with other similar pipeline configurations.wren-ai-service/src/web/v1/services/__init__.py (2)
66-66: Import added in alphabetical orderThe
InstructionsServiceimport is correctly added in alphabetical order within the import section, adhering to good code organization practices.
89-89: Service properly exposed through allAdding
InstructionsServiceto the__all__list ensures that it's correctly exposed as part of the module's public API, allowing it to be imported by other components.wren-ai-service/src/web/v1/routers/__init__.py (2)
8-8: Import added in alphabetical orderThe
instructionsmodule import is correctly added in alphabetical order within the import list, which is good practice for code organization.
32-32: Router integrationThe
instructions.routeris properly included in the main router, which will expose the new instructions endpoints in the API. This completes the integration of the new feature into the API routing system.wren-ai-service/src/config.py (1)
35-36: Configuration parameters for instructions featureThe addition of
instructions_similarity_thresholdandinstructions_top_kparameters with sensible defaults (0.7 and 10 respectively) aligns with the PR objectives. These parameters are appropriately placed in the indexing and retrieval config section, following a similar pattern to the existingsql_pairsparameters.Would you like to add documentation comments for these new parameters to match the style of other documented parameters in this file?
wren-ai-service/src/pipelines/indexing/__init__.py (1)
93-93: Proper addition of the Instructions module to the indexing pipelineThe addition of the Instructions import and its inclusion in the
__all__list correctly exposes the new instructions functionality to the rest of the codebase. This aligns well with the PR objective of implementing a new instructions indexing and retrieval pipeline.Also applies to: 102-102
wren-ai-service/docs/config_examples/config.ollama.yaml (1)
124-129: Correctly configured instructions pipeline entriesThe addition of
instructions_indexingandinstructions_retrievalpipeline entries with appropriate configurations for embedder and document store matches the PR objectives. This implementation ensures that the Ollama LLM provider can properly utilize the new instructions functionality.wren-ai-service/docs/config_examples/config.azure.yaml (1)
135-140: Well-implemented instructions pipeline for AzureThe new pipeline entries for instructions indexing and retrieval are properly configured with the Azure embedder and document store. This implementation maintains consistency with the other LLM providers and ensures that custom instructions can be incorporated during SQL generation when using Azure.
wren-ai-service/src/pipelines/retrieval/__init__.py (1)
2-2: Instructions retrieval pipeline properly integratedThe addition of the Instructions import and its inclusion in the
__all__list correctly exposes the instructions retrieval functionality. This complements the indexing pipeline and completes the implementation of the instructions feature.Also applies to: 14-14
deployment/kustomizations/base/cm.yaml (1)
184-189: New pipeline entries for instructions handling look goodThe addition of
instructions_indexingandinstructions_retrievalpipelines is well-structured and consistent with existing pipeline definitions. These entries use the same embedder and document store as other similar pipelines, which maintains consistency in the configuration.wren-ai-service/src/pipelines/generation/followup_sql_generation.py (4)
78-78: Parameter addition for instructions looks goodThe optional
instructionsparameter with proper type hints matches the PR's goal of incorporating custom instructions into the SQL generation process.
95-96: Correctly passing instructions to construct_instructionsThe parameter is correctly passed to the
construct_instructionsfunction.
160-161: Parameter addition to run method looks goodThe optional
instructionsparameter is consistently added to the class's run method with the same type definition as in the prompt function.
176-177: Correctly including instructions in pipeline inputsThe instructions parameter is properly included in the inputs dictionary passed to the pipeline execution.
wren-ai-service/src/pipelines/generation/sql_generation.py (4)
68-69: Parameter addition for instructions looks goodThe optional
instructionsparameter with proper type hints matches the PR's goal of incorporating custom instructions into the SQL generation process.
81-82: Correctly passing instructions to construct_instructionsThe parameter is correctly passed to the
construct_instructionsfunction.
148-149: Parameter addition to run method looks goodThe optional
instructionsparameter is consistently added to the class's run method with the same type definition as in the prompt function.
161-162: Correctly including instructions in pipeline inputsThe instructions parameter is properly included in the inputs dictionary passed to the pipeline execution.
wren-ai-service/src/pipelines/generation/utils/sql.py (2)
537-538: Parameter addition to construct_instructions function looks goodThe optional
instructionsparameter is added with appropriate type hints, maintaining consistency with changes in other files.
542-543: Good variable rename to avoid parameter name clashRenaming the variable from
instructionsto_instructionsavoids a name clash with the parameter of the same name.wren-ai-service/src/web/v1/services/semantics_preparation.py (4)
92-96: Code reformatting looks good.The reformatting of the assignment to
self._prepare_semantics_statusesimproves readability by properly structuring the multi-line assignment with parentheses.
100-108: Consistency in formatting applied correctly.Similar to the previous block, this reformatting maintains consistent style throughout the file while improving readability.
138-138: Parameter addition supports extensibility.The addition of
**kwargsto the method signature allows for passing arbitrary keyword arguments, making the method more flexible for future extensions.
149-149: Instructions successfully integrated into deletion process.Adding "instructions" to the list of names ensures that instruction data is properly cleaned up when semantics are deleted, which aligns with the PR objective of adding instruction handling capabilities.
wren-ai-service/src/web/v1/services/ask.py (2)
426-426: Instructions successfully integrated into followup SQL generation.The retrieved instructions are correctly passed to the followup_sql_generation pipeline, aligning with the PR objective of incorporating custom instructions during the query generation process.
440-440: Instructions successfully integrated into SQL generation.The retrieved instructions are correctly passed to the sql_generation pipeline, ensuring consistent instruction handling across both initial queries and follow-up queries.
wren-ai-service/src/pipelines/indexing/sql_pairs.py (2)
128-128: Default parameter increases flexibility.Making the
embeddingparameter optional with a default empty dictionary provides more flexibility in how the function can be called, which is particularly useful when this pipeline is integrated with the new instructions functionality.
218-218: Default parameter consistency maintained.Similarly, making the
sql_pairsparameter in thecleanmethod optional with a default empty list maintains consistency with the above change and increases the method's flexibility.wren-ai-service/src/globals.py (4)
29-29: New service properly added to container class.The addition of the
instructions_servicefield to theServiceContainerclass properly extends the service architecture to support the new instructions functionality.
71-73: Instructions indexing pipeline correctly configured.The instructions indexing pipeline is properly integrated into the semantics preparation service, which allows for managing instruction data alongside other semantics components.
99-103: Instructions retrieval pipeline well-configured with settings.The code correctly initializes the instructions retrieval pipeline with configurable similarity threshold and top-k parameters, which aligns with the PR objective of adding new settings for managing instructions similarity thresholds and retrieval options.
247-254: Instructions service properly instantiated.The
instructions_serviceis correctly instantiated with the appropriate pipeline configuration, completing the integration of the new instructions functionality into the service container.wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py (1)
44-49: Consider checking the exact content ofdocument.meta.You currently verify that specific keys exist in document metadata. For deeper coverage, consider asserting the exact values of
instruction_id,instruction, oris_defaultto ensure the pipeline’s indexing logic is fully correct.wren-ai-service/src/web/v1/services/instructions.py (1)
125-127: Validate instruction IDs.Ensure that invalid or empty instruction IDs are effectively handled. Currently, the code does not appear to guard against empty or malformed IDs. If needed, add validation to prevent indexing or deleting instructions with invalid IDs.
wren-ai-service/tests/pytest/services/test_instructions.py (1)
229-246: Validate coverage for instructions referencing multiple questions.Your tests demonstrate indexing a single question per instruction. Consider adding a case with multiple questions in a single instruction to confirm that each question is properly expanded into separate documents.
wren-ai-service/src/web/v1/routers/instructions.py (2)
79-82: Consider adding validation for instruction fields
Currently, thePostRequestmodel referencesInstructionsService.Instructionbut does not impose additional constraints (e.g., max length). For robustness and security, ensure that the instruction data is validated (e.g., length, special characters) before indexing.
148-154: Handle missing event IDs more gracefully
Accessingcontainer.instructions_service[event_id]may raise aKeyErrorifevent_idis not found, leading to a 500 response. Consider adding explicit error handling to return a 404 or custom error response if the ID is invalid or no longer in memory.wren-ai-service/src/pipelines/indexing/instructions.py (2)
53-75: Confirm index existence before deletion
When deleting documents, ensure the targeted collection or index (instructionsdataset) exists in the document store. Consider logging or handling the case where no matching documents are found for the specified IDs.
97-110: Consider concurrency checks for partial deletes
Thecleanfunction could experience conflicts if multiple concurrent deletions target different sets of instruction IDs. Consider concurrency handling or locking if partial overlap in instruction IDs might cause issues.wren-ai-service/src/pipelines/retrieval/instructions.py (2)
38-54: Ensurecount_documentsaligns with indexing filter logic
count_documentsapplies a project ID filter. Confirm that the filters used here match the ones used during indexing so that the counted dataset is consistent with what was stored.
105-125: Use consistent approach for retrieving default instructions
default_instructionsuses synchronousstore.filter_documents, unlike the async calls in other pipeline steps. For consistency, consider bridging to an async method or describing why a sync call is acceptable here.
paopa
force-pushed
the
feat/instruction-pipeline-and-endpoints
branch
from
8719a3d to
c21cadf
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🔭 Outside diff range comments (2)
wren-ai-service/src/web/v1/services/instructions.py (2)
141-143:⚠️ Potential issueUndefined variable “request” causes a runtime error.
The line
return self._cache[request.event_id].with_metadata()referencesrequest, which isn’t defined here. This will fail if called. You likely intended to returnself._cache[id].with_metadata().def __setitem__(self, id: str, value: Event): self._cache[id] = value - return self._cache[request.event_id].with_metadata() + return self._cache[id].with_metadata()
144-161:⚠️ Potential issueDuplicate code segment likely due to merge conflict.
Lines 144-161 repeat the
__getitem__and__setitem__definitions with slightly different parameter names. This is confusing and can cause unpredictable behavior. Remove or reconcile one version.
🧹 Nitpick comments (15)
wren-ai-service/src/globals.py (2)
96-100: Add validation for similarity threshold and top-k.While passing
settings.instructions_similarity_thresholdandsettings.instructions_top_kinto the pipeline is good, it’s safer to validate or constrain these values to avoid unexpected behavior (e.g., negative integers, threshold greater than 1, etc.). Consider bounding them or raise exceptions if they’re out of range.def create_service_container(...): ... "instructions_retrieval": retrieval.Instructions( **pipe_components["instructions_retrieval"], + similarity_threshold=min(max(settings.instructions_similarity_threshold, 0.0), 1.0), + top_k=max(settings.instructions_top_k, 1), ), ...
244-251: Consider clarifying or documenting usage ofinstructions_service.The
instructions_servicehandles indexing logic for instructions. Adding a concise docstring or inline comment clarifying its scope and how it integrates with retrieval flows can help new contributors understand its purpose.wren-ai-service/src/web/v1/services/instructions.py (4)
17-20: Allow for multiple error codes.Currently, the
Errormodel only accepts"OTHERS"ascode. If later you need to handle different error categories (e.g., "INVALID_INPUT", "NOT_FOUND"), usingLiteralwith multiple values or anEnummight be more extensible.
27-35: Externalize cache configuration.Large TTL caches can raise memory concerns. Consider moving
maxsizeandttlto the application’s config or environment variables for easier tuning without code changes.
57-89: Consider partial or async-friendly indexing improvements.For very large instruction sets, you might want to process indexing in chunks or in background tasks to avoid timeouts. If you expect large data, consider streaming or a queue-based approach.
95-101: Consider handling partial deletions vs. total failures.If some instruction IDs don’t exist, you might want to differentiate partial from full failures by returning which IDs were successfully removed and which were not. Right now, any exception triggers a full failure state.
wren-ai-service/src/pipelines/retrieval/instructions.py (3)
24-25: Avoid overshadowing built-in names.You are using
listas a variable name insideOutputFormatter.run, overshadowing the built-inlisttype. This can lead to confusion and potential bugs if the built-in type is needed later in the method or file. Consider renaming the variable toformatted_documentsor something more descriptive.- list = [] + formatted_documents = []
65-67: Confirm retrieval behavior for empty embeddings.If
embeddingis empty, the retrieval step immediately returns{}. This design choice is valid, but ensure business logic expects no results instead of a partial fallback. You might want to log a debug message or handle an edge case where an embedding error yields an empty dictionary.
118-140: Add concurrency considerations or docstrings for pipeline initialization.The
Instructionsclass sets up asynchronous components, but there’s no mention of concurrency strategies (e.g., concurrency-limited usage ofdocument_store). If multiple queries or pipeline runs occur simultaneously, consider verifying thread-safe or process-safe interactions with the Qdrant store and embedder.wren-ai-service/src/pipelines/indexing/instructions.py (1)
132-138: ClarifyDuplicatePolicy.OVERWRITEusage.The pipeline uses
OVERWRITEfor duplicates. This can be desirable when frequently updating existing instructions, but might lead to overwriting previously indexed instructions if an ID reoccurs. Verify with stakeholders whether this aligns with the intended design.wren-ai-service/src/web/v1/routers/instructions.py (1)
96-109: Consider logging success/failure for background tasks explicitly.While a background task is added to index instructions, the code only sets the initial
event_idstatus, then returns immediately. Logging the outcome upon completion (success or failure) in the background process can help with operational visibility and debugging.wren-ai-service/tests/pytest/services/test_instructions.py (4)
103-118: Missing document count assertionUnlike other test functions, this test doesn't verify the document count after indexing empty instructions. Consider adding an assertion to confirm that no documents are created, similar to what's done in the
test_with_empty_questionsfunction.assert response.status == "finished" + + pipe_components = generate_components(settings.components) + document_store_provider: DocumentStoreProvider = pipe_components[ + "instructions_indexing" + ]["document_store_provider"] + store = document_store_provider.get_store( + dataset_name="instructions", + ) + assert await store.count_documents() == 0
28-66: Consider refactoring common document store access patternThe code to retrieve the document store and check document count is repeated in multiple test functions. Consider extracting this into a helper function to reduce duplication and improve maintainability.
async def get_document_count(): pipe_components = generate_components(settings.components) document_store_provider: DocumentStoreProvider = pipe_components[ "instructions_indexing" ]["document_store_provider"] store = document_store_provider.get_store( dataset_name="instructions", ) return await store.count_documents()Then you could simplify assertions to:
assert await get_document_count() == 2Also applies to: 69-100, 121-171, 174-222, 225-267
28-66: Consider verifying document contentThe test verifies the document count after indexing but doesn't check if the indexed documents contain the correct content. Consider adding assertions to verify the documents were indexed with the expected data.
You could extend the test by retrieving and verifying documents:
# After asserting count documents = await store.get_documents() assert len(documents) == 2 # Verify document content matches the instructions assert any(doc.text == "This is a test instruction" for doc in documents) assert any(doc.text == "This is another test instruction" for doc in documents)
1-268: Consider adding error case testsThe test suite effectively covers the happy paths for indexing and deletion but doesn't test error handling. Consider adding tests for error cases like malformed instructions, duplicate IDs, or service failures.
Examples of error cases to test:
- Duplicate instruction IDs
- Invalid instruction format
- Service failure during indexing/deletion
- Deleting non-existent instructions
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (56)
wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/tools/config/config.example.yaml(2 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(1 hunks)wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/tools/config/config.example.yaml(1 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)deployment/kustomizations/base/cm.yaml(2 hunks)docker/config.example.yaml(1 hunks)wren-ai-service/src/web/v1/services/instructions.py(1 hunks)wren-ai-service/src/config.py(1 hunks)wren-ai-service/src/globals.py(3 hunks)wren-ai-service/src/pipelines/indexing/__init__.py(2 hunks)wren-ai-service/src/pipelines/retrieval/__init__.py(2 hunks)wren-ai-service/src/web/v1/services/__init__.py(2 hunks)wren-ai-service/src/web/v1/routers/__init__.py(2 hunks)wren-ai-service/src/web/v1/routers/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/instructions.py(6 hunks)wren-ai-service/src/globals.py(1 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/semantics_preparation.py(2 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(1 hunks)wren-ai-service/src/pipelines/indexing/sql_pairs.py(2 hunks)wren-ai-service/src/web/v1/services/semantics_preparation.py(2 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(2 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(3 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(2 hunks)wren-ai-service/src/web/v1/routers/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/instructions.py(3 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(0 hunks)wren-ai-service/src/web/v1/services/ask.py(3 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/tools/config/config.example.yaml(1 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py(1 hunks)wren-ai-service/tests/pytest/services/test_instructions.py(1 hunks)docker/config.example.yaml(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- wren-ai-service/src/pipelines/retrieval/instructions.py
🚧 Files skipped from review as they are similar to previous changes (46)
- wren-ai-service/src/pipelines/indexing/instructions.py
- wren-ai-service/tools/config/config.example.yaml
- wren-ai-service/docs/config_examples/config.groq.yaml
- docker/config.example.yaml
- wren-ai-service/docs/config_examples/config.google_ai_studio.yaml
- wren-ai-service/docs/config_examples/config.azure.yaml
- wren-ai-service/docs/config_examples/config.deepseek.yaml
- wren-ai-service/src/config.py
- wren-ai-service/src/web/v1/routers/init.py
- wren-ai-service/docs/config_examples/config.azure.yaml
- wren-ai-service/docs/config_examples/config.deepseek.yaml
- wren-ai-service/src/pipelines/indexing/instructions.py
- wren-ai-service/tools/config/config.full.yaml
- wren-ai-service/docs/config_examples/config.ollama.yaml
- wren-ai-service/tools/config/config.example.yaml
- wren-ai-service/docs/config_examples/config.deepseek.yaml
- wren-ai-service/docs/config_examples/config.ollama.yaml
- wren-ai-service/docs/config_examples/config.groq.yaml
- wren-ai-service/tools/config/config.full.yaml
- wren-ai-service/tools/config/config.example.yaml
- wren-ai-service/src/web/v1/services/semantics_preparation.py
- wren-ai-service/docs/config_examples/config.google_ai_studio.yaml
- wren-ai-service/tools/config/config.full.yaml
- wren-ai-service/src/web/v1/services/init.py
- wren-ai-service/docs/config_examples/config.groq.yaml
- deployment/kustomizations/base/cm.yaml
- wren-ai-service/docs/config_examples/config.google_ai_studio.yaml
- wren-ai-service/src/pipelines/retrieval/init.py
- wren-ai-service/src/pipelines/generation/utils/sql.py
- wren-ai-service/src/pipelines/generation/followup_sql_generation.py
- wren-ai-service/src/pipelines/generation/utils/sql.py
- wren-ai-service/docs/config_examples/config.azure.yaml
- wren-ai-service/src/pipelines/indexing/instructions.py
- wren-ai-service/src/web/v1/services/ask.py
- wren-ai-service/src/pipelines/retrieval/instructions.py
- wren-ai-service/src/pipelines/generation/sql_generation.py
- wren-ai-service/src/pipelines/retrieval/instructions.py
- wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py
- wren-ai-service/src/web/v1/services/instructions.py
- wren-ai-service/src/pipelines/indexing/init.py
- wren-ai-service/src/web/v1/services/semantics_preparation.py
- wren-ai-service/src/globals.py
- wren-ai-service/src/pipelines/indexing/sql_pairs.py
- docker/config.example.yaml
- wren-ai-service/src/web/v1/services/instructions.py
- wren-ai-service/src/web/v1/routers/instructions.py
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: pytest
🔇 Additional comments (10)
wren-ai-service/docs/config_examples/config.ollama.yaml (1)
127-129: Verify consistent embedder usage for indexing and retrieval.You are currently using
litellm_embedder.defaultforinstructions_indexingandlitellm_embedder.openai/nomic-embed-textforinstructions_retrieval. If these two embedders produce vectors with different dimensions or representations, it may reduce retrieval accuracy. Confirm that both embedder configurations are compatible with the sameembedding_model_dim(768) and that it is your intention to index with one embedder and retrieve with another.wren-ai-service/src/globals.py (1)
29-29: No concerns for the new field addition.Introducing
instructions_serviceas a first-class property inServiceContaineris straightforward and consistent with the existing pattern for other services.wren-ai-service/src/web/v1/services/instructions.py (5)
21-26: Looks good for handling event states.Using a Pydantic model with typed statuses and optional error fields is clear and structured.
36-51: Exception handler approach is acceptable.Storing the event state as "failed" and logging the error is straightforward and consistent. No immediate issues found.
52-56: IndexRequest model is clear.The fields are well-defined. No issues found here.
90-94: DeleteRequest model is straightforward.No concerns here.
127-140: Cache-based retrieval logic looks fine.Returning a placeholder event with status=failed when the cache is empty is a reasonable approach for nonexistent IDs. No further issues.
wren-ai-service/src/pipelines/indexing/instructions.py (1)
60-63:Details
❓ Verification inconclusive
Validate 'instruction_id' filter logic.
Here, you’re using the "in" operator for
instruction_id. While Qdrant typically supports this pattern, confirm that large lists of IDs don’t degrade performance, as Qdrant must filter each document individually.
🏁 Script executed:
#!/bin/bash # Searching for potential uses of 'instruction_id' filter across the codebase # to confirm consistent usage. rg -A 5 $'"instruction_id"'Length of output: 2343
Action: Validate the Performance of the 'instruction_id' Filter
The current implementation in
wren-ai-service/src/pipelines/indexing/instructions.py(lines 60–63) uses the "in" operator to filter byinstruction_id, which aligns with Qdrant’s supported filter patterns. However, as noted, please verify that this approach scales well when handling a large list of IDs. While similar usage across the codebase shows consistency (e.g., the filter setup in indexing and metadata handling in retrieval), it’s important to consider one or both of the following:
- Benchmarking: Confirm through performance benchmarks or monitoring that filtering with extensive lists does not lead to degradation in query latency.
- Safeguards or Alternatives: Consider implementing limits, batching, or alternative query strategies if the volume of IDs is expected to grow significantly.
wren-ai-service/tests/pytest/services/test_instructions.py (2)
12-25: Well-structured fixture for test isolationThe fixture correctly sets up the instructions service with a clean document store for each test by using
recreate_index=True. This ensures test isolation and prevents test interference.
225-267: Good use of nested helper functionThe
index_instructionshelper function is a clean way to reduce code duplication while testing cross-project isolation. This approach improves maintainability and readability of the test.
paopa
force-pushed
the
feat/instruction-pipeline-and-endpoints
branch
from
c21cadf to
0f22334
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (3)
wren-ai-service/src/web/v1/services/semantics_preparation.py (1)
145-149:⚠️ Potential issueDuplicate named argument detected.
Within this list comprehension, there is a repeated parameter
project_id=project_idlisted twice. This could introduce a syntax error or unintended overshadowing of the argument. Please remove the duplicate argument.- self._pipelines[name].clean( - project_id=project_id, - project_id=project_id, - delete_all=True, - ) + self._pipelines[name].clean( + project_id=project_id, + delete_all=True, + )wren-ai-service/src/web/v1/services/instructions.py (1)
143-161:⚠️ Potential issueResolve duplicated getitem / setitem methods and the undefined reference to
request.The block from lines 141–143 references
request.event_id, which is not defined in this scope, and immediately after that, lines 145–161 redefine__getitem__and__setitem__again, using parameter names likeevent_idinstead ofid. This strongly indicates a leftover merge conflict or duplication error, leading to ambiguous or overwriting behavior.Consider removing the duplicate methods and ensuring that a single pair of
__getitem__/__setitem__methods exists, using consistent parameter names and references. For example:def __setitem__(self, id: str, value: Event): - return self._cache[request.event_id].with_metadata() + self._cache[id] = value + return self._cache[id].with_metadata() -# Remove the duplicate __getitem__ and __setitem__ definitions below -... +#wren-ai-service/tests/pytest/services/test_instructions.py (1)
1-268: 💡 Verification agent🧩 Analysis chain
Missing test for instruction retrieval functionality
While the indexing and deletion functionality is well tested, there's no test for retrieving instructions, which is a key part of the feature according to the PR description.
🏁 Script executed:
#!/bin/bash # Check if there's a retrieval method in the InstructionsService class cat "wren-ai-service/src/web/v1/services/instructions.py" | grep -n "def retrieve\|def get\|def search" || echo "No retrieval method found in InstructionsService" # Check if there are retrieval tests in other files find wren-ai-service/tests -type f -name "test_*.py" | xargs grep -l "instructions.*retrieve" || echo "No retrieval tests found in other test files"Length of output: 405
Action Required: Implement and Test Retrieval Functionality
Our verification confirms that the retrieval functionality is completely missing:
- The
InstructionsServiceinwren-ai-service/src/web/v1/services/instructions.pydoes not implement any methods for retrieving instructions (e.g., noretrieve,get, orsearchmethods).- No tests exist across the project to cover retrieval scenarios.
Please add the intended retrieval functionality along with the necessary tests to validate it.
🧹 Nitpick comments (12)
wren-ai-service/src/globals.py (2)
96-100: Validate retrieval parameters for instructions.The addition of
"instructions_retrieval": retrieval.Instructions(...)with a similarity threshold and top_k is congruent with other retrieval pipelines. Confirm that these parameters—instructions_similarity_thresholdandinstructions_top_k—are effectively tested and traceable in logs to monitor retrieval performance.
244-251: Indexing pipeline forinstructions_service.The dedicated indexing flow under
instructions_serviceis clear and separates indexing from retrieval. If you intend to unify or share the instructions retrieval pipeline here, consider adding it in a future refactor for simplicity, but this approach remains valid as is.wren-ai-service/src/web/v1/services/semantics_preparation.py (1)
138-138: Leverage**kwargsor remove it if unneeded.The
delete_semanticsmethod now includes**kwargs, but the current implementation does not use or pass them along. If there is no further plan to handle these keyword arguments, consider removing**kwargsfor clarity.wren-ai-service/src/web/v1/routers/instructions.py (2)
80-110: Consider returning a 202 Accepted status for asynchronous indexing.The POST endpoint immediately returns an event ID while indexing happens in the background. By default, this returns an HTTP 200 status, but using 202 (Accepted) more accurately conveys that processing is in progress.
-@router.post("/instructions") +@router.post("/instructions", status_code=202)
112-140: Reassess synchronous vs. asynchronous deletion flow.Unlike the indexing endpoint, this DELETE request awaits completion before returning a final response. For lengthy deletions, consider a background task to mirror the indexing approach. Otherwise, no issues if immediate finalization is desired.
wren-ai-service/src/web/v1/services/instructions.py (2)
34-35: Address the TODO comment regarding method relocation.Lines 36–37 mention a TODO note to move the
_handle_exceptionmethod to a more generic utility or superclass. Moving this logic into a shared helper module (or a dedicated base class) would promote reuse, reduce duplication, and improve maintainability.
31-34: Reassess cache size and TTL values for memory and data consistency.Lines 30–34 initialize a very large cache (maxsize=1,000,000) with a TTL of 120 seconds. Depending on usage patterns and memory constraints, storing up to a million entries might be excessive or require fine-tuning. Consider validating real-world load scenarios and possibly lowering the maxsize or increasing the TTL for more consistent caching strategies.
wren-ai-service/src/pipelines/retrieval/instructions.py (1)
24-34: Avoid overshadowing built-in types by renaming the variablelist.On lines 24 and 25, the variable
listis being used as a local identifier, overshadowing the built-inlisttype. This can lead to confusion and reduce readability.- list = [] + formatted_docs = []wren-ai-service/src/pipelines/indexing/instructions.py (1)
37-47: Consider using the providedinstruction.idas the document identifier.Currently, the
InstructionsConvertergenerates a random UUID for eachDocument(line 38). While this is valid, usinginstruction.idfor the Haystack documentidmight simplify lookups and keep the ID consistent across indexing and retrieval operations. If the instructions and documents share the same identifier, referencing the document byinstruction_idbecomes more straightforward.- Document( - id=str(uuid.uuid4()), - meta={ - "instruction_id": instruction.id, - "question": instruction.question, - **addition, - }, - content=instruction.instruction, - ) + Document( + id=instruction.id, + meta={ + "question": instruction.question, + **addition, + }, + content=instruction.instruction, + )wren-ai-service/tests/pytest/services/test_instructions.py (3)
1-10: Imports look good but consider organizing them by standard/third-party/localThe imports are appropriate for this test file, but consider organizing them into standard library imports (uuid), third-party imports (pytest), and local imports (src modules) for better readability.
import uuid import pytest -from src.config import settings -from src.core.provider import DocumentStoreProvider -from src.globals import create_service_container -from src.providers import generate_components -from src.web.v1.services.instructions import InstructionsService +from src.config import settings +from src.core.provider import DocumentStoreProvider +from src.globals import create_service_container +from src.providers import generate_components +from src.web.v1.services.instructions import InstructionsService
103-120: Test case for empty instructions needs document count verificationThe test verifies the status but doesn't check the document count when empty instructions are provided. For consistency with other tests, add this verification.
assert response.status == "finished" + pipe_components = generate_components(settings.components) + document_store_provider: DocumentStoreProvider = pipe_components[ + "instructions_indexing" + ]["document_store_provider"] + store = document_store_provider.get_store( + dataset_name="instructions", + ) + assert await store.count_documents() == 0
1-268: Consider refactoring repeated code patterns into helper functionsThere's significant code duplication in setting up document store providers and checking document counts. Consider extracting these into helper functions for better maintainability.
async def get_document_count(): """Helper to get the current instruction document count.""" pipe_components = generate_components(settings.components) document_store_provider: DocumentStoreProvider = pipe_components[ "instructions_indexing" ]["document_store_provider"] store = document_store_provider.get_store( dataset_name="instructions", ) return await store.count_documents()
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (56)
wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/tools/config/config.example.yaml(2 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(1 hunks)wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/tools/config/config.example.yaml(1 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)deployment/kustomizations/base/cm.yaml(2 hunks)docker/config.example.yaml(1 hunks)wren-ai-service/src/web/v1/services/instructions.py(1 hunks)wren-ai-service/src/config.py(1 hunks)wren-ai-service/src/globals.py(3 hunks)wren-ai-service/src/pipelines/indexing/__init__.py(2 hunks)wren-ai-service/src/pipelines/retrieval/__init__.py(2 hunks)wren-ai-service/src/web/v1/services/__init__.py(2 hunks)wren-ai-service/src/web/v1/routers/__init__.py(2 hunks)wren-ai-service/src/web/v1/routers/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/instructions.py(6 hunks)wren-ai-service/src/globals.py(1 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/semantics_preparation.py(2 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(1 hunks)wren-ai-service/src/pipelines/indexing/sql_pairs.py(2 hunks)wren-ai-service/src/web/v1/services/semantics_preparation.py(2 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(2 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(3 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(2 hunks)wren-ai-service/src/web/v1/routers/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/instructions.py(3 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(0 hunks)wren-ai-service/src/web/v1/services/ask.py(3 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/tools/config/config.example.yaml(1 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py(1 hunks)wren-ai-service/tests/pytest/services/test_instructions.py(1 hunks)docker/config.example.yaml(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- wren-ai-service/src/pipelines/retrieval/instructions.py
🚧 Files skipped from review as they are similar to previous changes (34)
- wren-ai-service/docs/config_examples/config.deepseek.yaml
- wren-ai-service/docs/config_examples/config.ollama.yaml
- wren-ai-service/docs/config_examples/config.deepseek.yaml
- wren-ai-service/src/pipelines/indexing/instructions.py
- wren-ai-service/src/web/v1/routers/init.py
- wren-ai-service/docs/config_examples/config.ollama.yaml
- wren-ai-service/tools/config/config.full.yaml
- wren-ai-service/docs/config_examples/config.azure.yaml
- wren-ai-service/src/pipelines/retrieval/init.py
- wren-ai-service/tools/config/config.example.yaml
- wren-ai-service/src/config.py
- wren-ai-service/src/pipelines/indexing/init.py
- wren-ai-service/src/web/v1/routers/instructions.py
- wren-ai-service/src/pipelines/retrieval/instructions.py
- wren-ai-service/docs/config_examples/config.ollama.yaml
- wren-ai-service/src/web/v1/services/semantics_preparation.py
- wren-ai-service/src/pipelines/indexing/instructions.py
- wren-ai-service/src/pipelines/generation/followup_sql_generation.py
- wren-ai-service/src/pipelines/generation/utils/sql.py
- wren-ai-service/docs/config_examples/config.azure.yaml
- wren-ai-service/src/globals.py
- wren-ai-service/docs/config_examples/config.google_ai_studio.yaml
- deployment/kustomizations/base/cm.yaml
- wren-ai-service/src/pipelines/generation/utils/sql.py
- wren-ai-service/tools/config/config.full.yaml
- docker/config.example.yaml
- wren-ai-service/src/web/v1/services/ask.py
- wren-ai-service/docs/config_examples/config.deepseek.yaml
- wren-ai-service/src/pipelines/generation/sql_generation.py
- wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py
- wren-ai-service/src/pipelines/indexing/sql_pairs.py
- wren-ai-service/src/web/v1/services/init.py
- wren-ai-service/docs/config_examples/config.groq.yaml
- wren-ai-service/src/web/v1/services/instructions.py
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: pytest
🔇 Additional comments (17)
wren-ai-service/tools/config/config.full.yaml (1)
148-150: New Pipeline Entry: 'instructions_indexing' AdditionThe new pipeline entry for
instructions_indexing—configured withembedder: litellm_embedder.defaultanddocument_store: qdrant—is added correctly. Please double-check that this configuration is consistent with similar entries across other environments.wren-ai-service/docs/config_examples/config.groq.yaml (1)
148-150: New Pipeline Entry: 'instructions_indexing' IntegrationThe addition of the
instructions_indexingentry withlitellm_embedder.defaultanddocument_store: qdrantaligns well with the existing pipeline definitions. Ensure that its behavior matches expectations when indexing instructions.wren-ai-service/docs/config_examples/config.azure.yaml (1)
135-137: Pipeline Expansion: 'instructions_indexing' EntryThe new
instructions_indexingpipe added here uses the correct embedder and document store configuration. Verify that these settings remain consistent with other deployments and that any downstream processes are updated accordingly.wren-ai-service/docs/config_examples/config.google_ai_studio.yaml (1)
137-139: New Pipeline Entry: 'instructions_retrieval' AdditionA dedicated
instructions_retrievalentry has been added usingembedder: litellm_embedder.gemini/text-embedding-004along withdocument_store: qdrant. This enhancement complements the indexing pipeline and should be verified for consistent integration across the system.wren-ai-service/tools/config/config.example.yaml (2)
4-28: Model List Refactor: List Format StandardizationThe refactoring of the LLm models section—shifting to a list notation using hyphens—improves readability and consistency. Please ensure that all model references in downstream configurations or documentation align with this structured format.
148-150: New Pipeline Entry: 'instructions_indexing' ConfigurationThe
instructions_indexingpipeline is now included in the example configuration with the proper embedder (litellm_embedder.default) and document store (qdrant). This addition brings the example config in line with the latest service enhancements.docker/config.example.yaml (1)
132-137: Configuration for instructions pipeline looks consistent.The addition of the
instructions_indexingandinstructions_retrievalsteps aligns with the broader objective to handle AI instructions in the pipeline. No immediate issues are detected, and the embedder assignment mirrors other indexing/retrieval steps. Good job!wren-ai-service/src/globals.py (1)
29-29: Ensure usage context forinstructions_service.Adding
instructions_servicetoServiceContaineris a sensible step, but make sure you confirm that references to this service occur where indexing functionality is required (e.g., router endpoints). This helps avoid dead code or misconfigurations.wren-ai-service/src/web/v1/services/semantics_preparation.py (1)
150-155: Consistently clean theinstructionspipeline.Adding the
instructionspipeline cleanup ensures semantic documents stay in sync. Confirm that downstream code handles calls to clean instructions without issues, especially if partial deletions or project scoping differ from other pipelines.wren-ai-service/src/web/v1/routers/instructions.py (1)
142-155: Graceful error reporting in GET endpoint.The GET endpoint clearly returns an error if the event status is "failed." Just ensure client-side code checks for
errorin the response if status is "failed," and that logs provide enough context for debugging asynchronous issues.wren-ai-service/src/pipelines/retrieval/instructions.py (1)
66-68: This check for empty embeddings addresses a previously raised concern.Lines 66–68 safely return an empty dictionary if embeddings are missing, thus preventing downstream retrieval calls from failing. This aligns with past feedback about handling empty or malformed embeddings.
wren-ai-service/tests/pytest/services/test_instructions.py (6)
12-26: Fixture correctly sets up the test environmentThe fixture properly initializes the
InstructionsServicewith a clean document store for testing. It's good practice to recreate the index for test isolation.
28-67: Test case verifies basic instruction indexing functionalityThis test confirms that instructions with questions are properly indexed. Good job checking both the response status and the document count in the store.
69-101: Edge case handling for empty questions is well testedGood test case that verifies instructions with empty questions aren't indexed, which is an important edge case to handle.
121-172: Deletion test correctly verifies multiple instruction deletionThis test case effectively verifies that multiple instructions can be deleted at once and the document store is empty afterward.
174-223: Single instruction deletion test is well implementedGood test that verifies partial deletion works as expected, leaving other instructions intact.
225-268: Cross-project instruction deletion test is thoroughThe test ensures that deletion is properly scoped to the specified project, which is an important isolation feature. Good use of a nested function to avoid code duplication.
paopa
force-pushed
the
feat/instruction-pipeline-and-endpoints
branch
from
0f22334 to
4124ceb
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🔭 Outside diff range comments (1)
wren-ai-service/src/web/v1/services/semantics_preparation.py (1)
146-148:⚠️ Potential issueFix duplicate parameter in function call.
There's a duplicate
project_idparameter in the clean method call, which is a bug that should be fixed.- self._pipelines[name].clean( - project_id=project_id, - project_id=project_id, - delete_all=True, - ) + self._pipelines[name].clean( + project_id=project_id, + delete_all=True, + )
🧹 Nitpick comments (16)
wren-ai-service/docs/config_examples/config.groq.yaml (1)
127-133: Consistency Check: Embedder Selection for Instructions Pipieval
The configuration adds both aninstructions_indexingentry (usinglitellm_embedder.default) and aninstructions_retrievalentry. Notice that for retrieval the embedder is set tolitellm_embedder.text-embedding-3-large, which differs from the indexing embedder. Please verify whether this divergence is intentional for the GROQ environment or if a standardized embedder (e.g.default) should be used in both cases.wren-ai-service/docs/config_examples/config.ollama.yaml (1)
127-129: Ollama Configuration: Specialized Embedder for Retrieval
Theinstructions_retrievalpipeline in this file useslitellm_embedder.openai/nomic-embed-textwhileinstructions_indexinguseslitellm_embedder.default. Please confirm that employing a specialized embedder for retrieval in the Ollama setup is intentional. If not, consider aligning the embedder with the indexing pipeline for consistency.wren-ai-service/docs/config_examples/config.deepseek.yaml (1)
148-150: DeepSeek Configuration: Verification of Embedder Choice
The newinstructions_retrievalentry is configured withlitellm_embedder.text-embedding-3-large. Please verify that this choice meets the performance and retrieval requirements for the DeepSeek environment. If consistency across various deployment configurations is desired, you may want to consider usinglitellm_embedder.defaulthere as well.docker/config.example.yaml (1)
139-152: Enhancement Suggestion: Include New Instructions Settings
The PR objectives mention adding new settings—instructions_similarity_threshold(default 0.7) andinstructions_top_k(default 10)—for managing retrieval behavior. These settings are not visible in the current settings block. Consider adding these to the configuration (or documenting their defaults elsewhere) to provide clarity and ease of customization for end users.wren-ai-service/tools/config/config.example.yaml (1)
148-150: Add QA coverage for instructions_indexing pipeline.A new pipeline entry “instructions_indexing” is introduced here. Confirm it interfaces properly with the rest of the codebase (e.g., indexing tasks, retrieval logic) by adding or updating integration tests.
wren-ai-service/src/pipelines/retrieval/instructions.py (1)
105-119: Consider making default_instructions retrieval asynchronous.This synchronous function mirrors the structure of the
retrievalfunction but isn't marked async. For consistent concurrency handling, consider using an async approach if QdrantDocumentStore supports async calls.wren-ai-service/src/web/v1/services/instructions.py (4)
16-35: Evaluate nested class approach and TTLCache risks.The nested Error and Event classes are convenient, but consider extracting them for broader reuse. Also confirm that using TTLCache with a large
maxsizeis viable and thread-safe for your workload.
57-89: Add concurrency checks for the indexing method.Your background task usage is effective for asynchronous processing. However, confirm that simultaneous indexing tasks won't overload the pipeline or cause unexpected state changes in shared resources.
96-126: Consider large deletion scenarios.If a high volume of instructions is expected, evaluate batching or chunked processing to avoid potential timeouts or performance bottlenecks when calling
clean.
127-140: Improve not-found event handling.Returning a failed event when no entry is found is helpful, but dev teams might appreciate a unique error code to distinguish missing events from other failures.
wren-ai-service/src/web/v1/routers/instructions.py (3)
1-42: Document default values explicitly.These lines contain helpful docstrings and usage examples. Ensure that optional fields, like
project_id, have their default behavior detailed (e.g., ifproject_idis omitted, the instructions apply globally).
89-109: Use robust error handling for background tasks.The
indexendpoint launches a background task. If the pipeline or background job fails unexpectedly, consider updating the client more proactively, e.g., via a callback or additional logging.
112-115: Use descriptive docstrings for DeleteRequest.Expand the docstring or constraints to clarify behavior when the
instruction_idslist is empty or includes invalid entries.wren-ai-service/tests/pytest/services/test_instructions.py (3)
12-26: Consider using an async fixture with teardown.
Currently, the fixture does not perform any cleanup of test state. If the underlying document store or service container needs disposal, consider using anasync fixturewith a teardown phase (or ayieldfixture) to ensure a clean test environment for subsequent tests.
28-69: Validate stored documents beyond count checks.
This test only checksstore.count_documents()to assert correctness. It may be beneficial to load and inspect the stored documents to ensure that each instruction’s content matches the expectations.
122-172: Check concurrency in deletion operations if needed.
In larger or concurrent use cases, consider verifying whether multipleDeleteRequestcalls might overlap or cause partial deletions. For now, this test is sufficient.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (56)
wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/tools/config/config.example.yaml(2 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(1 hunks)wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/tools/config/config.example.yaml(1 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)deployment/kustomizations/base/cm.yaml(2 hunks)docker/config.example.yaml(1 hunks)wren-ai-service/src/web/v1/services/instructions.py(1 hunks)wren-ai-service/src/config.py(1 hunks)wren-ai-service/src/globals.py(3 hunks)wren-ai-service/src/pipelines/indexing/__init__.py(2 hunks)wren-ai-service/src/pipelines/retrieval/__init__.py(2 hunks)wren-ai-service/src/web/v1/services/__init__.py(2 hunks)wren-ai-service/src/web/v1/routers/__init__.py(2 hunks)wren-ai-service/src/web/v1/routers/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/instructions.py(6 hunks)wren-ai-service/src/globals.py(1 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/semantics_preparation.py(2 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(1 hunks)wren-ai-service/src/pipelines/indexing/sql_pairs.py(2 hunks)wren-ai-service/src/web/v1/services/semantics_preparation.py(2 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(2 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(3 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(2 hunks)wren-ai-service/src/web/v1/routers/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/instructions.py(3 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(0 hunks)wren-ai-service/src/web/v1/services/ask.py(3 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/tools/config/config.example.yaml(1 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py(1 hunks)wren-ai-service/tests/pytest/services/test_instructions.py(1 hunks)docker/config.example.yaml(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- wren-ai-service/src/pipelines/retrieval/instructions.py
🚧 Files skipped from review as they are similar to previous changes (36)
- wren-ai-service/docs/config_examples/config.groq.yaml
- wren-ai-service/src/pipelines/retrieval/instructions.py
- wren-ai-service/src/pipelines/indexing/instructions.py
- wren-ai-service/docs/config_examples/config.deepseek.yaml
- wren-ai-service/src/config.py
- wren-ai-service/docs/config_examples/config.deepseek.yaml
- wren-ai-service/src/pipelines/generation/utils/sql.py
- wren-ai-service/src/pipelines/indexing/instructions.py
- wren-ai-service/src/web/v1/routers/instructions.py
- wren-ai-service/src/web/v1/routers/init.py
- deployment/kustomizations/base/cm.yaml
- wren-ai-service/docs/config_examples/config.groq.yaml
- wren-ai-service/docs/config_examples/config.google_ai_studio.yaml
- wren-ai-service/tools/config/config.example.yaml
- wren-ai-service/src/web/v1/services/ask.py
- wren-ai-service/docs/config_examples/config.azure.yaml
- wren-ai-service/src/pipelines/indexing/init.py
- wren-ai-service/docs/config_examples/config.azure.yaml
- wren-ai-service/src/pipelines/retrieval/init.py
- wren-ai-service/docs/config_examples/config.ollama.yaml
- wren-ai-service/src/pipelines/generation/utils/sql.py
- wren-ai-service/src/pipelines/indexing/instructions.py
- wren-ai-service/src/web/v1/services/semantics_preparation.py
- wren-ai-service/src/pipelines/retrieval/instructions.py
- wren-ai-service/src/web/v1/services/instructions.py
- wren-ai-service/src/pipelines/indexing/sql_pairs.py
- docker/config.example.yaml
- wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py
- wren-ai-service/src/pipelines/generation/followup_sql_generation.py
- wren-ai-service/src/pipelines/generation/sql_generation.py
- wren-ai-service/src/globals.py
- wren-ai-service/src/web/v1/services/instructions.py
- wren-ai-service/src/pipelines/indexing/instructions.py
- wren-ai-service/tools/config/config.full.yaml
- wren-ai-service/docs/config_examples/config.ollama.yaml
- wren-ai-service/src/web/v1/services/init.py
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: pytest
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (go)
🔇 Additional comments (23)
wren-ai-service/docs/config_examples/config.azure.yaml (1)
135-137: Azure Configuration: New Instructions Indexing Pipeline
The new pipeline entry forinstructions_indexingis added withlitellm_embedder.defaultand appears correctly configured. However, there is no correspondinginstructions_retrievalentry in this configuration. Please confirm whether the Azure environment does not require retrieval or if it should be added for uniformity across deployment environments.docker/config.example.yaml (1)
132-137: Docker Configuration: Standardized Embedder Usage
Bothinstructions_indexingandinstructions_retrievalpipelines now uselitellm_embedder.default, ensuring a uniform configuration in the Docker setup. This standardization aids in easier maintenance and predictability across environments. Please ensure that the rationale for differing embedder choices in other configuration files (such as GROQ, Ollama, and DeepSeek) is well documented.wren-ai-service/tools/config/config.full.yaml (1)
148-150: Successfully integrated instructions_indexing pipeline.The addition of the instructions_indexing pipeline follows the same pattern as other indexing pipelines, using the standardized embedder and document store configuration.
wren-ai-service/docs/config_examples/config.google_ai_studio.yaml (1)
137-139: Model inconsistency between instructions_indexing and instructions_retrieval.While instructions_indexing (line 135) uses
litellm_embedder.default, the instructions_retrieval is configured with a specific Google AI modelgemini/text-embedding-004. This inconsistency could lead to embedding space mismatch issues.Consider whether both pipelines should use the same embedder for consistency. Typically, the same embedding model should be used for both indexing and retrieval to ensure vector compatibility.
wren-ai-service/src/web/v1/services/semantics_preparation.py (2)
138-138: Method signature updated correctly to include kwargs.Adding
**kwargsto the method signature makes it more flexible for future extensions.
150-154: Successfully added instructions cleanup.The instructions cleanup follows the same pattern as other pipelines, allowing for complete data cleanup when a project is deleted.
wren-ai-service/src/globals.py (3)
29-29: Added InstructionsService to ServiceContainer.This addition properly extends the ServiceContainer class to include the new instructions service.
96-100: Well-integrated instructions retrieval pipeline.The instructions_retrieval pipeline configuration is properly implemented with appropriate settings for similarity threshold and top_k parameters.
244-251: Properly initialized InstructionsService.The instructions_service is correctly initialized with the instructions_indexing pipeline and cache settings.
Note that only the indexing pipeline is provided to the service, not the retrieval pipeline. Verify whether the instructions_service needs access to both pipelines for complete functionality.
wren-ai-service/tools/config/config.example.yaml (2)
5-14: Ensure the model configuration aligns with usage requirements.These lines introduce a new model with an alias "default" and a
max_tokenslimit of 4096. Confirm that the alias mapping and token limit meet your project's current usage constraints.
15-27: Validate consistency with the first model.This second model definition is structurally similar to the default model. Ensure any differences (e.g., lack of alias) are intentional and won't affect your pipeline behavior.
wren-ai-service/src/pipelines/retrieval/instructions.py (1)
69-79: Clarify default vs. non-default instructions usage.By default, the filter excludes documents with
is_default == True. Validate that this aligns with your retrieval requirements, particularly when noproject_idis supplied.wren-ai-service/src/web/v1/services/instructions.py (4)
37-51: Refine error-handling details.While
_handle_exceptionsets a clear failed status, consider logging more contextual information or using exception-specific error codes to aid in debugging critical production issues.
52-56: Verify that IndexRequest is fully covered by tests.Ensure relevant tests explicitly validate the request data structure and any edge cases (e.g., missing instructions list). This helps guard against future regressions.
90-95: Confirm partial deletes are handled.The
DeleteRequestmodel is straightforward. Double-check that partial failures, if any, are surfaced clearly if certain instruction IDs are invalid or not found.
141-143: Straightforward cache assignment.The
__setitem__implementation is concise and clear, with no immediate concerns.wren-ai-service/src/web/v1/routers/instructions.py (2)
80-87: Validate request structure for indexing.The
PostRequestmodel includes a list of instructions and an optionalproject_id. Ensure that instructions rely on consistent naming across the codebase (e.g., confirm the keys match theInstructionmodel fields).
142-155: Expose event status details.The GET endpoint returns the current event status. Optionally include timestamps or additional metrics (e.g., how many instructions were processed) to improve the client’s visibility into ongoing tasks.
wren-ai-service/tests/pytest/services/test_instructions.py (4)
70-102: LGTM!
This test properly confirms that instructions without any questions do not produce stored documents.
103-121: No documents verification is correct.
Your approach for empty instructions is sound. No issues found.
175-223: Test coverage for single-instruction deletion is comprehensive.
No major issues here; the logic and assertions are clear.
226-268: Cross-project deletion test is well-defined.
The test thoroughly checks that instructions from other projects remain unaffected. No concerns.wren-ai-service/src/pipelines/indexing/instructions.py (1)
95-108:delete_alllogic does not truly remove all instructions ifinstruction_idsis empty.
Whendelete_all=Truebut no instruction IDs are passed, the filter includes'instruction_id' in [], matching no documents. This is the same concern raised previously.
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
deployment/kustomizations/base/cm.yaml(2 hunks)docker/config.example.yaml(2 hunks)wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/eval/__init__.py(1 hunks)wren-ai-service/eval/evaluation.py(1 hunks)wren-ai-service/eval/pipelines.py(9 hunks)wren-ai-service/eval/preparation.py(3 hunks)wren-ai-service/eval/utils.py(1 hunks)wren-ai-service/src/config.py(1 hunks)wren-ai-service/src/globals.py(4 hunks)wren-ai-service/tools/config/config.example.yaml(2 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (7)
- deployment/kustomizations/base/cm.yaml
- wren-ai-service/tools/config/config.example.yaml
- wren-ai-service/docs/config_examples/config.deepseek.yaml
- wren-ai-service/src/config.py
- wren-ai-service/docs/config_examples/config.groq.yaml
- wren-ai-service/tools/config/config.full.yaml
- docker/config.example.yaml
🔇 Additional comments (27)
wren-ai-service/eval/__init__.py (1)
14-14: New configuration flag for instructions feature addedThe added
allow_instructionsflag with a default value ofTrueenables the new instructions handling functionality in the evaluation settings. This aligns with the PR's objective to incorporate custom instructions during SQL query generation.wren-ai-service/eval/utils.py (1)
234-234: Explicit data source specification for WrenEngineAdding the explicit
"data_source": "duckdb"key ensures proper data source configuration when using WrenEngine, which is necessary for the new instructions pipeline to function correctly.wren-ai-service/docs/config_examples/config.ollama.yaml (1)
124-129: Added instruction pipeline components for Ollama configurationThe two new pipeline entries for instructions indexing and retrieval are correctly configured to use the
litellm_embedder.defaultembedder andqdrantdocument store. This implementation enables custom instructions during SQL generation when using Ollama as the LLM provider.wren-ai-service/docs/config_examples/config.azure.yaml (1)
135-140: Added instruction pipeline components for Azure configurationThe new pipeline entries for instructions indexing and retrieval are properly configured with the same components as in other provider configurations, ensuring consistent behavior across different LLM providers.
wren-ai-service/docs/config_examples/config.google_ai_studio.yaml (1)
134-139: New instructions pipeline components added correctlyThe implementation adds two new pipeline components for handling instructions in SQL generation:
instructions_indexing- for storing and indexing instructionsinstructions_retrieval- for retrieving relevant instructionsBoth are correctly configured to use the
litellm_embedder.defaultembedder andqdrantdocument store, maintaining consistency with other retrieval components in the pipeline.wren-ai-service/eval/evaluation.py (1)
174-174: Updated metrics_initiator with settings parameterThe function call to
pipelines.metrics_initiatorhas been updated to include thesettingsparameter, which is necessary for accessing the new instruction-related configuration values likeinstructions_similarity_thresholdandinstructions_top_k.wren-ai-service/eval/preparation.py (3)
432-443: Enhanced dataset preparation with SQL pairs extractionThe code now properly extracts previous ground truths as SQL pairs to provide better context for SQL generation. This implementation supports the new instructions retrieval pipeline by providing rich example pairs.
443-444: Added instructions extraction from evidence fieldThe code extracts the "evidence" field from the ground truth to be used as instructions. This approach enables the system to leverage domain-specific information during SQL generation.
454-455: Updated dataset structure with new fieldsThe candidate evaluation dataset now includes the
samplesfield (using SQL pairs) andinstructionsfield, which properly integrates the new instruction functionality into the evaluation pipeline.wren-ai-service/src/globals.py (4)
29-29: Added instructions_service to ServiceContainerThe
instructions_serviceis properly added to theServiceContainerclass, following the established pattern for service definition in the codebase.
71-73: Added instructions indexing pipelineThe instructions indexing pipeline is correctly integrated into the semantics preparation service, allowing for indexing of instructions content.
100-104: Added instructions retrieval pipeline with configurationThe instructions retrieval pipeline is properly configured with:
- Components from the configuration file
- The new similarity threshold setting
- The new top-k parameter
This implementation follows the pattern established for other retrieval components while adding the specific settings needed for instructions retrieval.
248-255: Initialized instructions service with required componentsThe service initialization follows the established pattern in the codebase, properly configuring the Instructions indexing pipeline and inheriting the query cache settings.
wren-ai-service/eval/pipelines.py (14)
89-94: Good enhancement to extract_units function to handle string inputs directlyThis change improves the function's flexibility by adding support for direct string inputs in addition to dictionary inputs, making the function more versatile in different contexts.
137-138: Parameter signature update to standardize APIGood change. Updating from
query: dicttoparams: dictcreates a more consistent interface and better reflects the actual parameter usage throughout the system.
141-147: API restructuring to support instructionsThe prediction dictionary has been updated to access fields from the params dictionary, including the new
instructionsfield. This change aligns well with the PR objective of supporting instructions in SQL generation.
160-160: Function parameter update for consistencyThe change from passing
predictiondirectly to passing parameters via**paramsensures consistency with the updated method signatures across the codebase.
220-227: Parameter handling refactoring in RetrievalPipelineThe
_processmethod has been updated to work with the newparamsparameter structure. Retrieval results are now set directly onparamsrather than returning a new dictionary, which maintains state better through the pipeline.
229-230: Call signature update in RetrievalPipelineThe
__call__method signature has been updated to acceptparams: dictinstead ofquery: str, ensuring consistency with the changes in theprocessmethod.
261-261: Added support for instructions in GenerationPipelineThe
_allow_instructionsflag is correctly added to control whether instructions should be included in SQL generation, aligning with the PR objectives.
274-281: Well-designed helper method for instructions retrievalThis encapsulated logic correctly formats instructions based on configuration flags and ensures a consistent structure for downstream processing.
282-286: Well-structured helper method for samples retrievalSimilar to the instructions helper, this method provides clean encapsulation of the sample retrieval logic based on configuration flags.
287-306: Updated SQL generation process to support instructionsThe
_processmethod now properly integrates instructions and samples into the SQL generation workflow. The retrieved instructions are formatted appropriately and passed to the generation pipeline.
308-309: Method signature update for consistencyThe
__call__method signature has been updated to match the new parameter structure used throughout the codebase.
399-430: Updated AskPipeline process to support instructionsThe
_processmethod has been properly updated to include instructions in the SQL generation flow. The instructions are retrieved and passed to the generation component.
440-441: Method signature update for consistencyThe
__call__method signature has been updated to acceptparams: dictinstead ofquery: str, ensuring consistency across the codebase.
506-517: Updated metrics_initiator to support settings parameterGood addition of the
settingsparameter with a default value to maintain backward compatibility. This allows for more configuration flexibility while ensuring existing code continues to work.
paopa
force-pushed
the
feat/instruction-pipeline-and-endpoints
branch
from
d89eaa2 to
ec92281
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (7)
wren-ai-service/src/globals.py (1)
249-256: Reduce duplication in pipeline instantiation.The same
instructions_indexingpipeline is defined under bothsemantics_preparation_serviceandinstructions_service. Consider centralizing or reusing the pipeline creation logic to avoid potential drift.- instructions_service=services.InstructionsService( - pipelines={ - "instructions_indexing": indexing.Instructions( - **pipe_components["instructions_indexing"], - ) - }, - **query_cache, - ), + # Reuse the same pipeline object if feasible + instructions_service=services.InstructionsService( + pipelines={ + "instructions_indexing": semantics_preparation_service.pipelines["instructions"] + }, + **query_cache, + ),wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py (2)
1-50: Consider adding negative or error-handling test cases.The test coverage is solid for standard indexing scenarios. Adding a test for invalid documents or incomplete instructions improves robustness.
94-140: Positive coverage of instructions deletion scenarios.You thoroughly test single-instruction, partial, and full deletion. Adding concurrency tests could further validate deletion under load.
wren-ai-service/src/web/v1/routers/instructions.py (1)
79-82: Optional validation for request fields.Consider additional validation (e.g., ensuring unique
idvalues) withinPostRequestfor more robust indexing.wren-ai-service/src/web/v1/services/instructions.py (1)
16-22: Rename the nestedInstructionclass to avoid collisions with the importedInstructionfromsrc.pipelines.indexing.instructions.Currently, the nested
Instructionclass in this service and the imported pipelineInstructionshare the same name but differ in their fields (single-question vs. multi-question). This can be confusing and prone to misuse.-class InstructionsService: - class Instruction(BaseModel): +class InstructionsService: + class ServiceInstruction(BaseModel): id: str instruction: str questions: List[str] is_default: bool = Falsewren-ai-service/src/pipelines/retrieval/instructions.py (1)
24-25: Avoid overshadowing the built-inlisttype.Here,
list = []shadows the built-in name. Though not syntactically incorrect, it can lead to confusion.- list = [] + documents_list = []wren-ai-service/eval/pipelines.py (1)
392-404: Helper methods duplicated from GenerationPipelineThe helper methods are identical to those in GenerationPipeline. This is appropriate given their current implementation, but consider refactoring to a common base class or utility function if more pipelines need these methods in the future.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (32)
deployment/kustomizations/base/cm.yaml(2 hunks)docker/config.example.yaml(2 hunks)wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/eval/__init__.py(1 hunks)wren-ai-service/eval/evaluation.py(1 hunks)wren-ai-service/eval/pipelines.py(10 hunks)wren-ai-service/eval/preparation.py(3 hunks)wren-ai-service/eval/utils.py(1 hunks)wren-ai-service/src/config.py(1 hunks)wren-ai-service/src/globals.py(4 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(1 hunks)wren-ai-service/src/pipelines/indexing/__init__.py(2 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(1 hunks)wren-ai-service/src/pipelines/indexing/sql_pairs.py(2 hunks)wren-ai-service/src/pipelines/retrieval/__init__.py(2 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)wren-ai-service/src/web/v1/routers/__init__.py(2 hunks)wren-ai-service/src/web/v1/routers/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/__init__.py(2 hunks)wren-ai-service/src/web/v1/services/ask.py(3 hunks)wren-ai-service/src/web/v1/services/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/semantics_preparation.py(2 hunks)wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py(1 hunks)wren-ai-service/tests/pytest/services/test_instructions.py(1 hunks)wren-ai-service/tools/config/config.example.yaml(2 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (25)
- wren-ai-service/eval/init.py
- wren-ai-service/docs/config_examples/config.deepseek.yaml
- wren-ai-service/src/web/v1/routers/init.py
- wren-ai-service/docs/config_examples/config.azure.yaml
- wren-ai-service/src/web/v1/services/init.py
- wren-ai-service/docs/config_examples/config.google_ai_studio.yaml
- wren-ai-service/docs/config_examples/config.ollama.yaml
- wren-ai-service/tools/config/config.full.yaml
- wren-ai-service/eval/utils.py
- wren-ai-service/docs/config_examples/config.groq.yaml
- wren-ai-service/tools/config/config.example.yaml
- deployment/kustomizations/base/cm.yaml
- wren-ai-service/src/config.py
- wren-ai-service/src/pipelines/generation/utils/sql.py
- wren-ai-service/eval/preparation.py
- wren-ai-service/src/pipelines/retrieval/init.py
- wren-ai-service/src/pipelines/generation/followup_sql_generation.py
- wren-ai-service/eval/evaluation.py
- wren-ai-service/src/web/v1/services/semantics_preparation.py
- wren-ai-service/src/pipelines/generation/sql_generation.py
- wren-ai-service/src/pipelines/indexing/sql_pairs.py
- wren-ai-service/src/pipelines/indexing/init.py
- wren-ai-service/src/web/v1/services/ask.py
- docker/config.example.yaml
- wren-ai-service/tests/pytest/services/test_instructions.py
⏰ Context from checks skipped due to timeout of 90000ms (2)
- GitHub Check: pytest
- GitHub Check: Analyze (go)
🔇 Additional comments (27)
wren-ai-service/src/globals.py (3)
29-29: Good addition ofinstructions_serviceto ServiceContainer.This property neatly integrates the new instructions functionality into the overall service container.
71-73:Details
❓ Verification inconclusive
Verify potential performance impact of indexing new instructions pipeline.
Adding
"instructions": indexing.Instructions(...)can lead to increased memory usage for large instruction sets. Consider verifying pipeline performance for large data.
🏁 Script executed:
#!/bin/bash # Investigate pipeline performance with large instruction sets rg -A3 'Instructions('Length of output: 101
Action: Re-run the verification with a corrected regex pattern
The initial regex failed due to an unescaped parenthesis. To gather the needed information, please run the updated script below that escapes the parenthesis:
#!/bin/bash # Re-run the search with an escaped parenthesis to correctly locate usages of Instructions( rg -A3 'Instructions\('This will help verify if the new instructions pipeline is used in contexts that might impact performance (especially with large instruction sets).
100-104:Details
❓ Verification inconclusive
Validate the retrieval thresholds and top_k settings.
Ensure that
similarity_thresholdandtop_kare not set too high or low. This might affect retrieval accuracy or performance.
🏁 Script executed:
#!/bin/bash # Check usage references for instructions_retrieval to confirm threshold & top_k correctness rg -A5 'instructions_retrieval'Length of output: 5264
Action Required: Validate 'instructions_retrieval' Parameter Values
Please verify that the
similarity_thresholdandtop_kvalues passed toretrieval.Instructions(inwren-ai-service/src/globals.pyat lines 100–104) are appropriately calibrated—not too high or too low—to maintain an optimal balance between retrieval accuracy and performance. In addition, cross-check these settings against similar configurations found in the example YAML files (such as those indocker/config.example.yaml,wren-ai-service/tools/config/config.full.yaml, etc.) to ensure consistency across the codebase.wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py (1)
52-92: Comprehensive test for multi-project instructions indexing.This test ensures that separate project IDs don’t overwrite each other’s instructions. Good coverage.
wren-ai-service/src/web/v1/routers/instructions.py (7)
19-76: Docstring mismatch with theInstructionmodel fields.The docstring references
"content"and"metadata", but theInstructionmodel uses"instruction"and"question". This may cause confusion for API consumers.
84-86: Response model is straightforward.Returning only
event_idis sufficient for the initial indexing request.
111-115: Approve theDeleteRequestmodel.The structure is clear and consistent with common API patterns.
116-139: Confirm if synchronous deletion is intended.
await service.delete(...)may block the request until deletion completes. If these calls can be lengthy, consider background tasks for consistency with indexing.
141-146: Simple and concise status model.
GetResponseeffectively communicates the event status and potential errors.
148-155: The GET endpoint is well-defined.Exposing event status via
GET /instructions/{event_id}is straightforward and easy to integrate.
88-109:Details
❓ Verification inconclusive
Ensure concurrency safety for event-based state.
Storing event data in
service[event_id]is convenient, but confirm it won’t cause concurrency issues under high request loads.
🏁 Script executed:
#!/bin/bash # Attempt concurrent test calls to the POST endpoint to see if stored events conflict # Pseudocode approach: curl -X POST http://localhost:8000/instructions -d '{"instructions": [{"id": "1", "instruction": "Test"}]}'Length of output: 516
Action: Investigate and confirm concurrency handling for event state.
The current approach uses a simple dictionary assignment (
service[event_id] = InstructionsService.Event(...)) which may be prone to race conditions under high concurrent loads. Note that our automated test attempt failed to reach the endpoint, so its output was inconclusive. Please ensure the following:
- Verify that
instructions_serviceis thread-safe or has appropriate synchronization (e.g., using locks or asynchronous-safe data structures).- Manually test the endpoint under concurrent load when the service is running to confirm that there are no state conflicts or lost events.
- Consider refactoring the code if necessary to handle concurrent access safely.
wren-ai-service/src/web/v1/services/instructions.py (2)
34-41: Confirm thread-safety ofTTLCacheunder concurrent/async usage.
TTLCachefromcachetoolsis not always thread-safe by default, and usage under asynchronous contexts may require additional synchronization. Ensure that calls to_cachewon't introduce race conditions in high-concurrency scenarios.Please verify how your FastAPI or other concurrency model handles shared caches. If uncertain, consider using a thread-safe alternative or add synchronization primitives.
116-129: Verify project-scoped deletion logic for instructions.When
request.project_idis omitted,deletepassesNoneas theproject_idto the pipeline. This could inadvertently lead to deleting instructions across all projects, depending on how the pipeline interprets aNoneproject ID.Please confirm that the pipeline’s
cleanmethod safely handlesNonefor project-specific logic. If you require a strict project-based filter, you may need an explicit check before calling the pipeline.wren-ai-service/src/pipelines/indexing/instructions.py (2)
20-25: Confirm the intended usage of single-question instructions.This
Instructionmodel only holds a singlequestionstring, whereas upstream logic inInstructionsServicesplits a list of questions into separate Instruction objects. Ensure that this design is intentional and tested thoroughly for multi-question scenarios.
99-108:delete_allis unimplemented inInstructionsCleaner.As previously noted, the check
if instruction_ids or delete_all:callscleaner.run(instruction_ids=instruction_ids, ...)even ifdelete_allisTruebutinstruction_idsis empty, resulting in no documents being deleted. Consider adding a parameter and logic inInstructionsCleaner.run()to handle unconditional deletion properly.wren-ai-service/src/pipelines/retrieval/instructions.py (1)
136-137: Ensure default instructions are formatted consistently before merging.Currently,
documents["documents"]is a list of dict, whereasdefault_instructionsis a list ofDocumentobjects. Appending them directly mixes different data structures. Either convertdefault_instructionsto dicts viaOutputFormatteror store structured docs consistently.wren-ai-service/eval/pipelines.py (11)
89-94: Good improvement to handle string input directly in extract_unitsThis enhancement makes the function more flexible by allowing it to process both string inputs and dictionary objects. The additional check with
isinstance(doc, str)provides better compatibility with different input formats.
137-148: Effective parameter naming update and support for instructionsThe method signature change from
querytoparamsimproves the flexibility of the API while maintaining backward compatibility. The addition of instructions and samples fields aligns perfectly with the PR objectives for supporting SQL generation with custom instructions.
220-227: Parameter signature update appropriately applied to RetrievalPipelineThe method now correctly uses and modifies the
paramsdictionary instead ofprediction. This maintains consistency with the changes in the parent class.
261-261: Property initialization correctly implementedThe GenerationPipeline properly initializes the new
_allow_instructionsproperty from the settings, which will control whether instructions are included in SQL generation.
274-286: Well-structured helper methods for retrieving instructions and samplesThese helper methods encapsulate the logic for retrieving and formatting instructions and samples, which promotes code reuse and cleaner implementation. They also respect the feature flags
_allow_instructionsand_allow_sql_samples.
287-307: SQL generation updated to support instructionsThe _process method now properly integrates instructions into the SQL generation pipeline, fulfilling a key objective of this PR. The instructions are correctly retrieved using the helper method and passed to the generation run method.
308-310: API signature consistently updated for call methodThe method signature change from
querytoparamsmaintains consistency with other method updates throughout the pipeline classes.
379-379: Property initialization for AskPipelineSimilar to the GenerationPipeline, AskPipeline now correctly initializes the
_allow_instructionsproperty.
405-439: Instructions integration in AskPipeline's process methodThe AskPipeline _process method has been updated to support instructions similar to the GenerationPipeline. The addition of instructions to the SQL generation run method ensures that custom instructions can influence the generated SQL queries.
441-443: API signature consistently updated for AskPipeline's call methodThe update maintains consistency with the parameter naming changes throughout the file.
507-518: Enhanced metrics_initiator with settings parameterThe addition of the settings parameter to metrics_initiator provides more flexibility in configuration, particularly for setting the database path for DuckDB. This is a good improvement that makes the function more configurable.
… and retrieval pipe
…ocument content for instruction
paopa
force-pushed
the
feat/instruction-pipeline-and-endpoints
branch
from
ec92281 to
2338ef0
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (5)
wren-ai-service/src/web/v1/services/instructions.py (1)
8-10: Consider renaming or removing the shadowedInstructionimport.The file defines a nested
class Instruction(BaseModel)(line 17) while also importing anInstructionfromsrc.pipelines.indexing.instructions(lines 8–10). This can be confusing for future maintainers since two distinct classes share the same name.-from src.pipelines.indexing.instructions import Instruction +from src.pipelines.indexing.instructions import Instruction as PipelineInstructionwren-ai-service/src/web/v1/routers/instructions.py (1)
123-139: Return a consistent error format forDELETE /instructions.Currently, on failure, a 500 status code is returned with
event.error. For partial failures or non-existent instruction IDs, confirm if the pipeline accurately populateserrordetails for consistent client handling.wren-ai-service/tests/pytest/services/test_instructions.py (2)
28-46: Consider adding a test to handle partial failure or nonexistent IDs during indexing.While the logic is tested with valid instructions, exploring edge cases (e.g., invalid IDs or partial writes) would increase confidence in error handling.
116-119: Add test coverage for ingestion errors.No tests appear to simulate pipeline exceptions during indexing. Including such a scenario would ensure that
_handle_exceptionlogic in the service is verified thoroughly.wren-ai-service/src/pipelines/retrieval/instructions.py (1)
127-139: Consider applying formatter to default_instructions for consistent output.The formatted_output function directly appends default_instructions to the documents without formatting them using output_formatter. This could lead to inconsistent data structures in the output.
def formatted_output( default_instructions: list[Document], filtered_documents: dict, output_formatter: OutputFormatter, ) -> dict: if not filtered_documents and not default_instructions: return {"documents": []} documents = output_formatter.run(documents=filtered_documents.get("documents")) - documents["documents"].extend(default_instructions) + if default_instructions: + formatted_defaults = output_formatter.run(documents=default_instructions) + documents["documents"].extend(formatted_defaults.get("documents", [])) return documents
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (32)
deployment/kustomizations/base/cm.yaml(2 hunks)docker/config.example.yaml(2 hunks)wren-ai-service/docs/config_examples/config.azure.yaml(1 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(1 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(1 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(1 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(1 hunks)wren-ai-service/eval/__init__.py(1 hunks)wren-ai-service/eval/evaluation.py(1 hunks)wren-ai-service/eval/pipelines.py(10 hunks)wren-ai-service/eval/preparation.py(3 hunks)wren-ai-service/eval/utils.py(1 hunks)wren-ai-service/src/config.py(1 hunks)wren-ai-service/src/globals.py(4 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(4 hunks)wren-ai-service/src/pipelines/generation/utils/sql.py(1 hunks)wren-ai-service/src/pipelines/indexing/__init__.py(2 hunks)wren-ai-service/src/pipelines/indexing/instructions.py(1 hunks)wren-ai-service/src/pipelines/indexing/sql_pairs.py(2 hunks)wren-ai-service/src/pipelines/retrieval/__init__.py(2 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)wren-ai-service/src/web/v1/routers/__init__.py(2 hunks)wren-ai-service/src/web/v1/routers/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/__init__.py(2 hunks)wren-ai-service/src/web/v1/services/ask.py(3 hunks)wren-ai-service/src/web/v1/services/instructions.py(1 hunks)wren-ai-service/src/web/v1/services/semantics_preparation.py(2 hunks)wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py(1 hunks)wren-ai-service/tests/pytest/services/test_instructions.py(1 hunks)wren-ai-service/tools/config/config.example.yaml(2 hunks)wren-ai-service/tools/config/config.full.yaml(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (26)
- wren-ai-service/eval/init.py
- wren-ai-service/docs/config_examples/config.deepseek.yaml
- wren-ai-service/docs/config_examples/config.azure.yaml
- wren-ai-service/docs/config_examples/config.ollama.yaml
- wren-ai-service/src/web/v1/routers/init.py
- wren-ai-service/src/web/v1/services/init.py
- wren-ai-service/docs/config_examples/config.groq.yaml
- wren-ai-service/src/pipelines/retrieval/init.py
- wren-ai-service/src/config.py
- wren-ai-service/tools/config/config.full.yaml
- wren-ai-service/eval/utils.py
- deployment/kustomizations/base/cm.yaml
- wren-ai-service/src/pipelines/indexing/init.py
- wren-ai-service/src/pipelines/generation/utils/sql.py
- wren-ai-service/tools/config/config.example.yaml
- wren-ai-service/src/web/v1/services/ask.py
- wren-ai-service/docs/config_examples/config.google_ai_studio.yaml
- wren-ai-service/eval/evaluation.py
- wren-ai-service/src/pipelines/generation/sql_generation.py
- wren-ai-service/src/pipelines/generation/followup_sql_generation.py
- wren-ai-service/src/web/v1/services/semantics_preparation.py
- wren-ai-service/eval/preparation.py
- wren-ai-service/src/globals.py
- wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py
- wren-ai-service/src/pipelines/indexing/sql_pairs.py
- docker/config.example.yaml
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: pytest
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (go)
🔇 Additional comments (17)
wren-ai-service/src/web/v1/services/instructions.py (2)
76-86:Details
❓ Verification inconclusive
Check consistency between the service’s
IndexRequestand the pipeline’sInstruction.The loop converts each instruction into multiple pipeline instructions by splitting out each question. Verify that the rest of the pipeline and any consumers handle this multi-document logic correctly, especially if the pipeline’s
Instructionmodel differs from the service’s.
🏁 Script executed:
#!/bin/bash # Using ast-grep to confirm correct usage of multi-document logic across the codebase ast-grep --pattern $'for $_ in $.questions'Length of output: 44
Manual Verification Required: Confirm Multi-Document Handling Consistency
The code at
wren-ai-service/src/web/v1/services/instructions.pycorrectly splits each instruction into multiple pipeline instructions by iterating over eachquestionin the service’s instructions. However, our automated check using theast-greppattern yielded no output, so we couldn’t confirm whether all related pipeline components and consumers are handling this multi-document logic as intended.Please manually verify that:
- The downstream processing (including the pipeline’s
Instructionmodel consumers) correctly handles multiple documents generated from a single instruction.- The structural consistency between the service’s
IndexRequestand the pipeline’sInstructionis maintained.
145-158:Details
✅ Verification successful
Evaluate concurrency impact when multiple concurrent calls share the same
event_id.Access to
_cachevia__getitem__and__setitem__in a concurrent environment might cause race conditions if the sameevent_idis reused. Confirm if the upstream code ensures unique event IDs or otherwise serializes access.
🏁 Script executed:
#!/bin/bash # Checking for potential concurrency patterns around event_id usage rg -A 5 "event_id"Length of output: 22214
Action: Confirm Event ID Uniqueness & Concurrency Considerations
After reviewing the code and search results, it’s clear that upstream endpoints (both in
/instructionsand/sql-pairs) are consistently generating unique event IDs usinguuid.uuid4(). This design minimizes the risk of concurrent requests sharing the same key in_cacheduring__getitem__and__setitem__operations.
Current Behavior:
- New event IDs are generated per transaction, ensuring that each cache entry is associated with a unique event.
- The code paths observed in both routers and service implementations support this unique identifier strategy.
Points to Consider:
- If, in any future scenario, event IDs might come from an external source or be reused, then there could be potential race conditions that would warrant additional synchronization (e.g., locks) around
_cacheaccesses.- It might be beneficial to document that the current design assumes all event IDs are unique, to avoid confusion if changes are made in the future.
Overall, the current implementation appears robust in its handling of event IDs in a concurrent environment.
wren-ai-service/src/web/v1/routers/instructions.py (2)
29-33: Docstring mismatch: replace "content" with "instruction" and remove "metadata" to match the actual model fields.
88-108:Details
❓ Verification inconclusive
Validate background task error handling.
Although the POST endpoint offloads indexing to a background task, any pipeline error is surfaced only via the GET status endpoint. Ensure you handle error states or instruction collisions properly, especially if multiple indexing requests share overlapping instructions.
🏁 Script executed:
#!/bin/bash rg -A 6 "background_tasks.add_task"Length of output: 11065
Validate background task error handling
While the current implementation offloads the indexing operation to a background task (using
background_tasks.add_task), errors occurring there are only made visible via the GET status endpoint. Please ensure that the background task—which callsservice.index—has robust error management to address potential issues such as instruction collisions due to overlapping requests. Consider verifying that:
- Exceptions during the indexing process are properly logged or captured.
- There’s a mechanism (e.g., retries or conflict resolution) to deal with overlapping instructions.
- Any error state from the background task is clearly communicated (or surfaced) so that downstream handling (like GET status queries) can reliably detect and respond to issues.
wren-ai-service/src/pipelines/indexing/instructions.py (2)
20-24: Switch to aquestions: List[str]field to align with usage.
96-108:delete_allparameter logic does not trigger a true “delete all.”wren-ai-service/src/pipelines/retrieval/instructions.py (4)
56-61: Good implementation of the embedding function with short-circuit optimization.The function efficiently short-circuits when no documents are found, preventing unnecessary embedding operations and returning an empty dictionary. This is a performance optimization that avoids processing when there's no data to work with.
64-86: Well-implemented retrieval function with proper error handling.The retrieval function correctly handles the case when embedding is empty by returning an empty dictionary at the beginning, preventing potential errors in downstream operations. The filter construction is also clean and properly handles the optional project_id parameter.
144-171: Well-structured pipeline class with appropriate defaults.The Instructions class is well-designed with clear initialization of components and configuration parameters. The default values for similarity_threshold (0.7) and top_k (10) are sensible starting points for instruction retrieval.
173-183: Good implementation of the run method with proper logging.The run method properly passes all required components and configurations to the pipeline executor. The logging statement provides appropriate visibility into pipeline execution.
wren-ai-service/eval/pipelines.py (7)
88-94: Good enhancement to extract_units for handling string inputs.The function now correctly handles both string and dictionary inputs by directly parsing DDL strings. This makes the function more flexible and reusable.
137-148: Well-structured update to process method for supporting instructions.The method has been updated to use params instead of query, and now properly handles instructions as part of the prediction dictionary. The change maintains backward compatibility while extending functionality.
274-281: Good encapsulation of instruction retrieval logic.The _get_instructions method efficiently encapsulates the logic for retrieving and formatting instructions based on the _allow_instructions flag, which improves code maintainability and readability.
289-301: Well-implemented integration of instructions into the generation pipeline.The _process method has been updated to smoothly incorporate instructions into the SQL generation process. The code correctly uses the helper methods to get instructions and samples.
379-379: Fixed initialization of _allow_instructions in AskPipeline.The initialization of _allow_instructions has been correctly added to the AskPipeline class, addressing a previously identified issue.
413-431: Clean integration of instructions into the SQL generation process.The updated _process method now correctly retrieves and passes instructions to the generation.run method. This implementation maintains the existing functionality while adding support for instruction-guided SQL generation.
507-518: Enhanced metrics_initiator with settings support.The function now properly accepts and uses settings for configuring engine_info, which provides more flexibility and configuration options for metrics initialization.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py (3)
10-50: The indexing test provides good coverage of basic functionality.The test properly validates document creation and the presence of required metadata fields. Consider enhancing it with additional assertions to verify the actual content values match what was provided in the test instructions (not just checking for existence).
For example:
# Validate actual content values match input docs_by_id = {doc.meta["instruction_id"]: doc for doc in documents} assert docs_by_id["test-id-1"].meta["instruction"] == "This is a test instruction" assert docs_by_id["test-id-1"].meta["is_default"] is False assert docs_by_id["test-id-2"].meta["instruction"] == "This is another test instruction" assert docs_by_id["test-id-2"].meta["is_default"] is True
52-92: The multi-project test validates isolation between project IDs.The test verifies documents are properly separated by project ID. For completeness, consider testing both project IDs separately to ensure complete bidirectional isolation:
# Check documents for fake-id-2 documents_second_project = store.filter_documents( filters={ "operator": "AND", "conditions": [ {"field": "project_id", "operator": "==", "value": "fake-id-2"}, ], } ) assert len(documents_second_project) == 1
94-142: The deletion test verifies all deletion scenarios correctly.The test appropriately validates three key scenarios:
- No deletion when project ID doesn't match
- Targeted deletion of a specific instruction
- Bulk deletion of all instructions
Consider adding an assertion to verify which specific instruction remains after the first deletion:
# After deleting test-id-1, verify that test-id-2 remains remaining_docs = store.filter_documents() assert len(remaining_docs) == 1 assert remaining_docs[0].meta["instruction_id"] == "test-id-2"
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: pytest
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (go)
🔇 Additional comments (2)
wren-ai-service/tests/pytest/pipelines/indexing/test_instructions.py (2)
1-8: LGTM! Good imports and structure for the test file.The imports are well-organized and include all necessary components for testing the Instructions pipeline functionality.
1-142: Overall test coverage looks good for the Instructions pipeline.The tests cover all critical operations (indexing, filtering, deletion) across various scenarios. One minor suggestion would be to add a teardown method or fixture that ensures cleanup even if tests fail, although the
recreate_index=Trueapproach works well for test isolation.For the PR objectives, these tests adequately verify the new instructions indexing functionality that will support SQL generation with custom instructions.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py (2)
192-197: Consider updating the dry run test.The dry run pipeline test at the bottom of the file doesn't include the new
instructionsparameter. Consider updating it to test this new functionality.dry_run_pipeline( SQLGenerationReasoning, "sql_generation_reasoning", query="this is a test query", contexts=[], + instructions=["Use concise language", "Focus on numeric aggregations"] )
75-75: Add documentation for the instructions parameter.Consider adding a brief comment or docstring to explain the purpose and expected format of the
instructionsparameter. This would help other developers understand how to use this new feature correctly.def prompt( query: str, documents: List[str], sql_samples: List[str], - instructions: List[str], + instructions: List[str], # Custom instructions to guide SQL generation prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), ) -> dict:async def run( self, query: str, contexts: List[str], sql_samples: Optional[List[str]] = None, - instructions: Optional[List[str]] = None, + instructions: Optional[List[str]] = None, # Optional custom instructions to guide SQL generation configuration: Configuration = Configuration(), query_id: Optional[str] = None, ):Also applies to: 170-171
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py(6 hunks)wren-ai-service/src/web/v1/services/ask.py(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- wren-ai-service/src/web/v1/services/ask.py
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: pytest
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (go)
🔇 Additional comments (4)
wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py (4)
53-58: Good implementation of conditional template section for instructions.The template now properly handles instructions with a conditional block that only renders when instructions are provided. The Jinja syntax is correctly implemented with appropriate iteration over the instructions list.
75-75: Parameter successfully added to function signature.The
instructionsparameter is properly added to thepromptfunction signature and correctly passed to theprompt_builder.runmethod. This ensures the instructions are incorporated into the generated prompt.Also applies to: 83-83
170-171: Good handling of optional instructions parameter.The
runmethod correctly accepts an optionalinstructionsparameter with a default value ofNone, and the execution logic properly handles this by defaulting to an empty list whenNoneis provided. This ensures compatibility with existing code that doesn't pass instructions.Also applies to: 181-181
131-133: Clean code formatting improvement.The formatting changes for the queue assignment improve readability without altering functionality.
Also applies to: 144-146
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (13)
wren-ai-service/docs/config_examples/config.groq.yaml (2)
127-132: Instructions Pipeline Entries Added
The new pipeline entriesinstructions_indexingandinstructions_retrievalare correctly added with the expected embedder and document store. Verify that these configurations meet the required indexing/retrieval behavior for instructions.
152-152: Newline at End of File Missing
YAMLlint flagged that a newline character is missing at the end of the file. Please add an extra newline to satisfy the linter.- instructions_top_k: 10 \ No newline at end of file + instructions_top_k: 10 +🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 152-152: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/docs/config_examples/config.google_ai_studio.yaml (2)
134-139: Instructions Pipeline Entries Integration
The addition of theinstructions_indexingandinstructions_retrievalentries in the pipelines section (lines 134–139) is implemented consistently with other pipeline definitions. Double-check that these entries are referenced correctly by the indexing and retrieval modules.
159-159: End-of-File Newline Required
A newline character is missing at the end of the file as reported by YAMLlint. Please add an extra blank line at the end.- instructions_top_k: 10 \ No newline at end of file + instructions_top_k: 10 +🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 159-159: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/docs/config_examples/config.ollama.yaml (2)
124-129: Pipeline Updates for Instructions
The new entriesinstructions_indexingandinstructions_retrieval(lines 124–129) are now part of the pipelines. They correctly use thelitellm_embedder.defaultfor embedding andqdrantas the document store. Ensure that these settings are harmonized with other pipeline components.
149-149: Add Newline at EOF
YAMLlint indicates that a newline is missing at the end of the file. Please append a newline to comply with best practices.- instructions_top_k: 10 \ No newline at end of file + instructions_top_k: 10 +🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 149-149: no new line character at the end of file
(new-line-at-end-of-file)
deployment/kustomizations/base/cm.yaml (2)
180-185: New Instructions Pipeline in ConfigMap
The new pipeline entries forinstructions_indexingandinstructions_retrieval(lines 180–185) have been added to the ConfigMap for the Wren AI Service. These entries appear well integrated. Confirm that related services and API routers reference these new pipelines appropriately.
206-206: Missing Newline at EOF
A newline is missing at the end of the file. Please add a newline character after line 206 to resolve the YAMLlint error.- instructions_top_k: 10 \ No newline at end of file + instructions_top_k: 10 +🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 206-206: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/tools/config/config.full.yaml (2)
146-152: Updated Instructions Pipeline Configuration
The new pipeline entries forinstructions_indexingandinstructions_retrieval(lines 146–152) are added appropriately. They maintain consistency with the other pipeline definitions by specifying the same embedder and document store.
176-176: Ensure Final Newline at File End
YAMLlint flagged the absence of a newline at the very end of this file. Please add an extra newline to comply with best practices.- instructions_top_k: 10 \ No newline at end of file + instructions_top_k: 10 +🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 176-176: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/docs/config_examples/config.azure.yaml (1)
160-160: Add a newline at the end of file.YAML files should end with a newline character to comply with YAML standards and avoid linting warnings.
instructions_top_k: 10 +🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 160-160: no new line character at the end of file
(new-line-at-end-of-file)
docker/config.example.yaml (1)
158-158: Add a newline at the end of file.YAML files should end with a newline character to comply with YAML standards and avoid linting warnings.
instructions_top_k: 10 +🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 158-158: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/docs/config_examples/config.deepseek.yaml (1)
170-170: Add a newline at the end of file.YAML files should end with a newline character to comply with YAML standards and avoid linting warnings.
instructions_top_k: 10 +🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 170-170: no new line character at the end of file
(new-line-at-end-of-file)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
deployment/kustomizations/base/cm.yaml(2 hunks)docker/config.example.yaml(3 hunks)wren-ai-service/docs/config_examples/config.azure.yaml(2 hunks)wren-ai-service/docs/config_examples/config.deepseek.yaml(2 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(2 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(2 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(2 hunks)wren-ai-service/src/pipelines/retrieval/instructions.py(1 hunks)wren-ai-service/tools/config/config.example.yaml(3 hunks)wren-ai-service/tools/config/config.full.yaml(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- wren-ai-service/tools/config/config.example.yaml
🧰 Additional context used
🪛 YAMLlint (1.35.1)
wren-ai-service/docs/config_examples/config.groq.yaml
[error] 152-152: no new line character at the end of file
(new-line-at-end-of-file)
deployment/kustomizations/base/cm.yaml
[error] 206-206: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/docs/config_examples/config.google_ai_studio.yaml
[error] 159-159: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/docs/config_examples/config.azure.yaml
[error] 160-160: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/docs/config_examples/config.ollama.yaml
[error] 149-149: no new line character at the end of file
(new-line-at-end-of-file)
docker/config.example.yaml
[error] 158-158: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/tools/config/config.full.yaml
[error] 176-176: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/docs/config_examples/config.deepseek.yaml
[error] 170-170: no new line character at the end of file
(new-line-at-end-of-file)
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: pytest
🔇 Additional comments (18)
wren-ai-service/docs/config_examples/config.groq.yaml (1)
148-152: New Instructions Settings Added
The settings forinstructions_similarity_threshold(0.7) andinstructions_top_k(10) have been integrated, which aligns well with similar configurations in other environments.🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 152-152: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/docs/config_examples/config.google_ai_studio.yaml (1)
158-159: New Instructions Settings in Configuration
The new settingsinstructions_similarity_threshold: 0.7andinstructions_top_k: 10have been appended to the settings block. They are consistent with the other configuration files and appear to be set to the intended default values.🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 159-159: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/docs/config_examples/config.ollama.yaml (1)
148-149: Incorporation of New Instructions Settings
The settingsinstructions_similarity_thresholdandinstructions_top_kadded at lines 148–149 are correctly configured. Their values match the ones in other configuration examples.🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 149-149: no new line character at the end of file
(new-line-at-end-of-file)
deployment/kustomizations/base/cm.yaml (1)
202-206: Extended Settings with Instructions Parameters
New settings forinstructions_similarity_thresholdandinstructions_top_k(lines 205–206), along with the related thresholds, are added to the settings block. They are consistent with the changes in other YAML configuration files.🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 206-206: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/tools/config/config.full.yaml (1)
172-176: New Instructions Settings Added to Global Config
The settings block now includesinstructions_similarity_threshold: 0.7andinstructions_top_k: 10(lines 175–176) as part of the extended configurations. These values are consistent across all configuration files.🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 176-176: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/docs/config_examples/config.azure.yaml (2)
135-140: Implementation of instructions indexing and retrieval pipelines looks good.The new pipeline entries for instructions indexing and retrieval are properly configured with the required embedder and document store components.
156-160: New settings for similarity thresholds and retrieval limits are well-defined.The configuration introduces appropriate default values:
- Instructions similarity threshold of 0.7 aligns with the sql_pairs_similarity_threshold
- Top-k value of 10 is a reasonable retrieval limit
🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 160-160: no new line character at the end of file
(new-line-at-end-of-file)
docker/config.example.yaml (4)
5-28: Model definitions are now properly formatted as a list.The change to use a hyphen-prefixed list format for model definitions improves readability and follows YAML best practices.
34-37: Embedder model definition is properly formatted.The embedding model definition now follows the same list format as the LLM models.
132-137: Instructions pipeline entries are properly configured.The new pipeline entries for instructions indexing and retrieval match the pattern of other similar pipelines in the configuration.
154-158: New settings are consistent across config files.The added settings for similarity thresholds and retrieval limits are consistent with the values in other configuration files.
🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 158-158: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/docs/config_examples/config.deepseek.yaml (2)
145-150: Instructions pipeline entries are properly added.The new pipeline entries for instructions indexing and retrieval are configured with the correct components.
166-170: Settings are consistently defined across configurations.The similarity thresholds and retrieval size settings are defined with appropriate default values that match other configuration files.
🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 170-170: no new line character at the end of file
(new-line-at-end-of-file)
wren-ai-service/src/pipelines/retrieval/instructions.py (5)
18-35: Output formatter provides clear structure for instructions.The
OutputFormatterclass neatly transforms Document objects into a consistent dictionary format with relevant instruction metadata.
56-62: Empty document handling is properly implemented.The embedding function correctly handles the case when no documents are found by returning an empty dictionary, which is gracefully handled by downstream functions.
65-86: Embedding existence check is properly implemented.The code now checks if the embedding exists before proceeding with retrieval, addressing a previous review comment.
148-175: Pipeline configuration with sensible defaults.The
Instructionsclass is well-structured with reasonable default values for similarity threshold (0.7) and top-k (10) parameters, matching the configuration in YAML files.
176-188: Pipeline execution properly handles optional project_id.The
runmethod correctly handles the case when project_id is not provided by defaulting to an empty string, ensuring compatibility with the pipeline functions.
| def formatted_output( | ||
| default_instructions: list[Document], | ||
| filtered_documents: dict, | ||
| output_formatter: OutputFormatter, | ||
| ) -> dict: | ||
| if not filtered_documents and not default_instructions: | ||
| return {"documents": []} | ||
|
|
||
| merged = default_instructions.get("documents") + filtered_documents.get("documents") | ||
| documents = output_formatter.run(documents=merged) | ||
| return documents | ||
|
|
There was a problem hiding this comment.
Fix potential runtime error in formatted_output function.
The default_instructions variable is treated as a dictionary with a "documents" key, but if the function's return type is list[Document] as declared, this will cause a runtime error.
@observe(capture_input=False)
def formatted_output(
default_instructions: list[Document],
filtered_documents: dict,
output_formatter: OutputFormatter,
) -> dict:
if not filtered_documents and not default_instructions:
return {"documents": []}
- merged = default_instructions.get("documents") + filtered_documents.get("documents")
+ documents_list = default_instructions if isinstance(default_instructions, list) else default_instructions.get("documents", [])
+ filtered_list = filtered_documents.get("documents", [])
+ merged = documents_list + filtered_list
documents = output_formatter.run(documents=merged)
return documents📝 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 formatted_output( | |
| default_instructions: list[Document], | |
| filtered_documents: dict, | |
| output_formatter: OutputFormatter, | |
| ) -> dict: | |
| if not filtered_documents and not default_instructions: | |
| return {"documents": []} | |
| merged = default_instructions.get("documents") + filtered_documents.get("documents") | |
| documents = output_formatter.run(documents=merged) | |
| return documents | |
| @observe(capture_input=False) | |
| def formatted_output( | |
| default_instructions: list[Document], | |
| filtered_documents: dict, | |
| output_formatter: OutputFormatter, | |
| ) -> dict: | |
| if not filtered_documents and not default_instructions: | |
| return {"documents": []} | |
| documents_list = default_instructions if isinstance(default_instructions, list) else default_instructions.get("documents", []) | |
| filtered_list = filtered_documents.get("documents", []) | |
| merged = documents_list + filtered_list | |
| documents = output_formatter.run(documents=merged) | |
| return documents |
| async def default_instructions( | ||
| count_documents: int, retriever: Any, project_id: str | ||
| ) -> list[Document]: | ||
| if not count_documents: | ||
| return [] | ||
|
|
||
| filters = { | ||
| "operator": "AND", | ||
| "conditions": [ | ||
| {"field": "is_default", "operator": "==", "value": True}, | ||
| ], | ||
| } | ||
|
|
||
| if project_id: | ||
| filters["conditions"].append( | ||
| {"field": "project_id", "operator": "==", "value": project_id} | ||
| ) | ||
|
|
||
| res = await retriever.run( | ||
| query_embedding=None, | ||
| filters=filters, | ||
| ) | ||
| return dict(documents=res.get("documents")) | ||
|
|
There was a problem hiding this comment.
Fix return type inconsistency in default_instructions function.
The function signature indicates it returns list[Document] but the implementation returns a dictionary.
@observe(capture_input=False)
async def default_instructions(
count_documents: int, retriever: Any, project_id: str
-) -> list[Document]:
+) -> dict:
if not count_documents:
- return []
+ return {"documents": []}
filters = {
"operator": "AND",
"conditions": [
{"field": "is_default", "operator": "==", "value": True},
],
}
if project_id:
filters["conditions"].append(
{"field": "project_id", "operator": "==", "value": project_id}
)
res = await retriever.run(
query_embedding=None,
filters=filters,
)
return dict(documents=res.get("documents"))📝 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.
| async def default_instructions( | |
| count_documents: int, retriever: Any, project_id: str | |
| ) -> list[Document]: | |
| if not count_documents: | |
| return [] | |
| filters = { | |
| "operator": "AND", | |
| "conditions": [ | |
| {"field": "is_default", "operator": "==", "value": True}, | |
| ], | |
| } | |
| if project_id: | |
| filters["conditions"].append( | |
| {"field": "project_id", "operator": "==", "value": project_id} | |
| ) | |
| res = await retriever.run( | |
| query_embedding=None, | |
| filters=filters, | |
| ) | |
| return dict(documents=res.get("documents")) | |
| @observe(capture_input=False) | |
| async def default_instructions( | |
| count_documents: int, retriever: Any, project_id: str | |
| ) -> dict: | |
| if not count_documents: | |
| return {"documents": []} | |
| filters = { | |
| "operator": "AND", | |
| "conditions": [ | |
| {"field": "is_default", "operator": "==", "value": True}, | |
| ], | |
| } | |
| if project_id: | |
| filters["conditions"].append( | |
| {"field": "project_id", "operator": "==", "value": project_id} | |
| ) | |
| res = await retriever.run( | |
| query_embedding=None, | |
| filters=filters, | |
| ) | |
| return dict(documents=res.get("documents")) |
|
let's merge this PR after current candidate oss release |
2 participants
Overview
This PR introduces a new instructions indexing and retrieval pipeline and endpoints to enhance SQL generation by incorporating custom instructions during query generation.
Key Changes
instructions_indexingandinstructions_retrievalacross different LLM providers (OpenAI, Azure, Deepseek, Google AI, Groq, Ollama)Configuration Updates
instructions_similarity_threshold(default: 0.7)instructions_top_k(default: 10)Pipeline Changes
Testing Recommendations
Related Documentation
Summary by CodeRabbit
New Features
InstructionsServiceto streamline instruction management.Bug Fixes
instructions_retrievalacross multiple configurations to ensure consistent performance.Tests