feat: save all system prompt in conversation #552
Conversation
…message if needed
…without attachment
…sage block update methods
clementb49
force-pushed
the
saveAllSystems
branch
from
ffce9aa to
7fe94c0
Compare
WalkthroughThis update introduces a new file version constant and integrates system message handling into various components. Changes include adding validation methods for file attachments, refactoring message models with a new base class and validators, and implementing migration functions for conversation formats. Provider engine methods and GUI completion workflows now expect a specialized system message type. New custom Pydantic types have been added, and several test suites have been created or refactored to cover these updates and ensure robust validation. Changes
Sequence Diagram(s)sequenceDiagram
participant U as User
participant CT as ConversationTab
participant C as Conversation
participant PE as ProviderEngine
U->>CT: Initiate completion request
CT->>CT: Retrieve SystemMessage via get_system_message()
CT->>C: Add new message block with system message
CT->>PE: Call completion(new_block, conversation, system_message)
PE->>PE: Process completion using SystemMessage
PE-->>CT: Return completion result
CT-->>U: Deliver completed message
Suggested reviewers
✨ Finishing Touches
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 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
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🔭 Outside diff range comments (2)
basilisk/gui/conversation_tab.py (2)
1281-1293: 🧹 Nitpick (assertive)Inconsistent handling of system messages when adding blocks.
_pre_handle_completion_with_streampassessystem_messagetoadd_block, but_post_completion_without_streamomits it. Align these flows if the system prompt must always be captured.
1427-1434:⚠️ Potential issueMethod name mismatch breaks removal logic.
TheConversationmodel definesremove_block, but the code callsremove_message_block, which doesn't exist.- self.conversation.remove_message_block(message_block) + self.conversation.remove_block(message_block)
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
⛔ Files ignored due to path filters (1)
poetry.lockis excluded by!**/*.lock
📒 Files selected for processing (19)
basilisk/consts.py(1 hunks)basilisk/conversation/__init__.py(2 hunks)basilisk/conversation/attached_file.py(6 hunks)basilisk/conversation/conversation_helper.py(3 hunks)basilisk/conversation/conversation_model.py(3 hunks)basilisk/gui/conversation_tab.py(7 hunks)basilisk/provider_engine/anthropic_engine.py(2 hunks)basilisk/provider_engine/ollama_engine.py(3 hunks)basilisk/provider_engine/openai_engine.py(1 hunks)basilisk/types.py(1 hunks)pyproject.toml(2 hunks)tests/conftest.py(1 hunks)tests/conversation/__init__.py(1 hunks)tests/conversation/test_attachment.py(1 hunks)tests/conversation/test_migration.py(1 hunks)tests/conversation/test_model.py(1 hunks)tests/conversation/test_save_restore.py(1 hunks)tests/test_message_position_manager.py(2 hunks)tests/test_types.py(1 hunks)
🔇 Additional comments (34)
tests/conversation/__init__.py (1)
1-1: Appropriate linting suppression.Suppressing the D104 warning (missing docstring in package) makes sense for test module initialization files that don't contain actual code.
basilisk/consts.py (1)
49-50: Version constant added correctly for the new file format.The version increment to
2suggests this PR implements a breaking change to accommodate storing system prompts in conversations. This constant will be important for migration and backward compatibility.basilisk/conversation/__init__.py (2)
18-18: SystemMessage export properly added.The SystemMessage class has been correctly imported from the conversation_model module to support the new feature of saving system prompts.
34-34: SystemMessage properly exposed in all.Adding SystemMessage to the all list ensures it can be imported directly from the conversation module, which is the appropriate pattern.
pyproject.toml (2)
42-42: Good addition of ordered-set dependency.The ordered-set package will likely be used for maintaining insertion order while ensuring uniqueness, which makes sense for storing system prompts in conversations.
62-62: Appropriate test dependency for HTTP mocking.Adding pytest-httpx as a development dependency is beneficial for mocking HTTP requests in tests. This suggests proper test coverage for the new features that may interact with remote resources.
basilisk/provider_engine/openai_engine.py (1)
307-307: Improved robustness with defensive attribute accessThe change from direct attribute access to using
getattrwith a default value is a good defensive programming practice. This prevents potentialAttributeErrorexceptions if the message object doesn't have an "attachments" attribute.basilisk/provider_engine/ollama_engine.py (3)
16-16: Added SystemMessage import for type annotationThis import is necessary to support the type annotation change in the
completionmethod.
124-124: Improved robustness with defensive attribute accessSimilar to the change in OpenAIEngine, this uses
getattrwith a default value to prevent potentialAttributeErrorexceptions if the message object doesn't have an "attachments" attribute.
88-88:❓ Verification inconclusive
Parameter type restriction to SystemMessage
The parameter type has been restricted from
Message | NonetoSystemMessage | None, indicating a deliberate design decision to use a more specialized message type for system prompts.Run the following script to check if all callers of this method are updated to pass SystemMessage:
🏁 Script executed:
#!/bin/bash # Description: Verify all calls to OllamaEngine.completion use SystemMessage ast-grep --pattern 'OllamaEngine().completion($$$, $$$, $system_message, $$$)'Length of output: 80
Action Required: Manually verify all calls to OllamaEngine.completion.
The parameter type has indeed been specialized from
Message | NonetoSystemMessage | None, which is the intended design. However, the automated verification script using ast-grep did not produce any output. This result is inconclusive regarding whether all call sites correctly pass aSystemMessagefor thesystem_messageparameter. Please manually inspect the codebase to confirm that every invocation ofOllamaEngine().completionuses aSystemMessage(orNone) for this parameter.
- Checklist for manual verification:
- Review all call sites of the
OllamaEngine().completionmethod.- Ensure that the argument corresponding to
system_messageis either an instance ofSystemMessageorNone.- Confirm no outdated use of a generic
Messageis present.basilisk/types.py (2)
10-12: Well-defined custom path type with validationThe
PydanticUPathtype is well-designed, using Pydantic'sAnnotatedtype to create a validated UPath type. The plain validator ensures that values are properly converted to UPath instances.
15-49: Well-implemented custom OrderedSet with Pydantic integrationThe
PydanticOrderedSetclass properly extends OrderedSet and implements the required Pydantic core schema generation method. The implementation handles both instance and non-instance schemas, and provides proper serialization to JSON.The schema generation is thorough, handling type arguments correctly and providing appropriate validation and serialization.
tests/test_types.py (2)
10-82: Comprehensive test suite for PydanticOrderedSetThe test suite for PydanticOrderedSet is thorough and covers various aspects:
- Basic set operations
- Order preservation
- Pydantic model integration
- JSON serialization/deserialization
- Type validation
- Schema generation
This ensures the custom type functions correctly in all expected scenarios.
84-95: Complete tests for PydanticUPathThe PydanticUPath tests cover both valid and invalid path cases, ensuring that:
- Valid paths are correctly converted to UPath objects
- Invalid inputs raise appropriate exceptions
This provides good coverage for this simpler custom type.
basilisk/provider_engine/anthropic_engine.py (2)
26-26: Introduce specialized SystemMessage for better type safety.Good enhancement to use the more specific
SystemMessagetype rather than genericMessagefor system-level instructions.
251-251:✅ Verification successful
Updated parameter type enhances type safety.
Changing the type annotation from
Message | NonetoSystemMessage | Noneenforces stricter type checking and better represents the expected input. Make sure this change is consistently applied across all engine implementations.
🏁 Script executed:
#!/bin/bash # Check if the type annotation is consistent across all engine implementations echo "Checking for SystemMessage type annotations in engine files..." rg "system_message: SystemMessage \| None" basilisk/provider_engine/Length of output: 372
Type Annotation Consistency Confirmed
The change fromMessage | NonetoSystemMessage | Noneinbasilisk/provider_engine/anthropic_engine.pyhas been verified, and the same annotation is found inbasilisk/provider_engine/ollama_engine.py. No further changes are required.basilisk/conversation/attached_file.py (2)
261-273: Add validation to ensure file location exists.This is a good addition that will prevent downstream errors when attempting to use non-existent files. The validator appropriately skips URL files and provides a clear error message for missing files.
498-500: Add directory creation for improved robustness.Good enhancement to create the directory structure if it doesn't exist, preventing potential file operation errors.
tests/conftest.py (1)
1-115: Well-structured test fixtures for better test organization.The new test fixtures provide a clean, centralized way to create test objects, reducing duplication across test files. The
system_messagefixture aligns well with the changes in the engine implementations.Consider adding docstrings to clarify the purpose and usage of these fixtures, particularly for more complex ones.
tests/conversation/test_save_restore.py (2)
132-217: Comprehensive tests for system message save/restore functionality.These tests thoroughly verify that system messages are correctly saved and restored, including both unique and shared system messages across message blocks. This is essential to ensure the new
SystemMessagetype works correctly with the conversation persistence mechanism.
219-462: Complete coverage of attachment-related save/restore scenarios.The tests cover different types of attachments (text, image, URL) and various scenarios (single attachment, multiple attachments, citations). This provides good confidence in the robustness of the save/restore functionality with different content types.
tests/conversation/test_model.py (1)
1-329: Comprehensive test suite with good structural organizationThis new file introduces a well-structured test suite for the Conversation module, with clear separation of concerns across multiple test classes. The tests appropriately verify both happy paths and error conditions for:
- AI model and message validation
- Message block validation
- Basic conversation functionality
- Complex scenarios with multiple blocks and system messages
The fixture usage is clean, and the systematic verification of system message handling across different operations provides thorough coverage for the new functionality saving all system prompts in conversations.
basilisk/conversation/conversation_helper.py (3)
101-103: Good defensive programming with type checkingThe added type check for
ImageFilebefore attempting to resize attachments is a good practice, ensuring that only image attachments are processed by the resize operation.
196-208: Correct implementation for v0 to v1 migrationThe migration function correctly sets the version number to 1 while preserving all other conversation data, facilitating proper versioning support.
236-237: Clear migration mechanism with explicit step listThe
migration_stepslist provides a clean, maintainable way to organize migration functions, making it easy to add future migrations as the conversation format evolves.tests/conversation/test_migration.py (2)
17-120: Well-structured tests for migration functionsThe test class provides comprehensive coverage of the migration functions, testing:
- v0 to v1 migration
- v1 to v2 migration with system message
- v1 to v2 migration without system message
- v1 to v2 migration with empty messages
All tests verify the key aspects of each migration step including proper handling of system messages, which is central to this PR's objective.
123-280: Thorough file migration tests for all versionsThe tests effectively verify that:
- Files in older formats (v0, v1) are properly migrated to the latest version
- Files already in the current format (v2) are loaded correctly
- Files with invalid versions are rejected with appropriate errors
This ensures robust handling of conversation files across different versions.
tests/test_message_position_manager.py (1)
1-254: Well-executed test refactoring to pytest styleThe refactoring from unittest to pytest improves readability and maintainability by:
- Using pytest fixtures instead of setup methods
- Organizing tests into logical classes by functionality
- Using modern pytest assertion style
- Maintaining complete test coverage during the transition
The changes support the larger PR goal by ensuring all components are consistently tested with modern practices.
tests/conversation/test_attachment.py (3)
1-2: Initial test setup looks solid.
No issues spotted with the file-level docstring.
21-34: Robust test coverage for non-existent files.
These parametrized tests correctly handle local and in-memory scenarios withpytest.raises(FileNotFoundError).
279-285: Good parameterization for MIME type wildcard tests.
The test covers single, multiple, unknown, and empty MIME sets comprehensively.basilisk/conversation/conversation_model.py (2)
108-110: Strict system role enforcement is correct.
This ensures that onlysystemrole messages use theSystemMessageclass.
201-229: Incremental migration logic looks effective.
The while-loop steps through each version, applying transformations. Ensure coverage for a scenario whereversion == BSKC_VERSIONinitially, so no migrations are needed.basilisk/gui/conversation_tab.py (1)
43-43: ImportingSystemMessageis a welcome addition.
It clearly distinguishes between user/assistant messages and system-level instructions.
There was a problem hiding this comment.
Actionable comments posted: 3
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (4)
basilisk/conversation/attached_file.py(9 hunks)basilisk/conversation/conversation_helper.py(3 hunks)basilisk/gui/conversation_tab.py(8 hunks)tests/conversation/test_attachment.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: build_app / build_windows (x64)
🔇 Additional comments (29)
basilisk/conversation/conversation_helper.py (5)
8-8: Imports look good.
No issues found with addingAnyto the imported typing constructs.
11-11: Valid usage of ValidationInfo.
ImportingValidationInfofor the migration functions is appropriate.
196-209: Migration from v0 to v1.
Straightforward version assignment is correct. Ensure you handle any existing “version” field gracefully in upstream logic, if present.
211-235: Verify partial assignment of system_index.
The code setssystem_indexfor only the last message. If multiple messages might require references to the same system, confirm that this partial assignment is intentional and won’t cause confusion downstream.
238-239: Organized migration steps.
Storing these functions in a list is a clean approach to versioned migrations.basilisk/gui/conversation_tab.py (10)
43-43: Looks great calling out SystemMessage class
The direct inclusion ofSystemMessageclearly aligns with the new system message handling approach.
1089-1090: Clearly specializing the return type
Changing the return signature toSystemMessage | Noneensures a distinct type for system messages.
1098-1098: Explicit system message construction
ReturningSystemMessage(content=system_prompt)is consistent with the specialized system message structure.
1235-1235: Consistent kwargs retrieval
Grabbingsystem_messagefromkwargsis a straightforward pattern that keeps the code flexible.
1240-1244: Passing system_message into streaming
Passingsystem_messageto_pre_handle_completion_with_streamensures the updated logic applies to streamed completions as well.
1270-1272: Incorporating system_message in non-stream flows
Includingsystem_messagein_post_completion_without_streammaintains consistency with the streaming workflow.
1283-1283: Unified add_block call
Adding bothnew_blockandsystem_messageensures the conversation model stays aware of any system context.
1320-1322: New parameter in post-completion
Acceptingsystem_messagehere keeps the method interface uniform with the streaming variant.
1329-1329: Propagating system context
Insertingsystem_messageintoadd_blockhelps maintain the correct conversation state throughout.
1433-1433: Centralized block removal
Callingself.conversation.remove_block(message_block)centralizes logic and avoids direct list manipulation, improving maintainability.tests/conversation/test_attachment.py (6)
84-91: Method call uses the correct name.The test correctly uses the fixed method name
get_display_info()rather than the previously misspelledget_dispay_info(). This aligns with the renamed method in the implementation file.
153-159: Good test for truncating data URLs.This test ensures that data URLs are properly truncated for display purposes, which is important for maintaining readable user interfaces when dealing with potentially long base64-encoded image data.
183-187: Good test setup for image resize output folder.Creating a dedicated folder for resized images in the tests matches the implementation's behavior, allowing for better isolation and cleanup of test artifacts.
189-215: Comprehensive test coverage for image resizing.The parametrized test covers an excellent range of edge cases for image resizing, including handling of same size, smaller size, larger size, various quality settings, and negative dimensions.
247-268: Well-structured mock for HTTP responses in URL tests.The test effectively mocks HTTP responses for the image URL test, providing proper content type and length headers to simulate a real HTTP response.
269-285: Good coverage of MIME type parsing scenarios.The test parametrization for
parse_supported_attachment_formatscovers important cases like single MIME type, multiple MIME types, unknown MIME type, and empty MIME type set.basilisk/conversation/attached_file.py (8)
261-273: Good addition of location existence validation.Adding this validator improves robustness by ensuring file locations are valid early in the object lifecycle, with appropriate handling for URL-type attachments.
One minor suggestion: Consider adding debug logging when bypassing the validation for URL types or when successfully confirming a file exists, which would make troubleshooting easier.
327-327: Use of getattr for fallback improves robustness.Using
getattrwith a default value is a clean way to handle the optionalresize_locationattribute, ensuring the code works correctly even when the attribute isn't present.
372-372: Consistency improvement by using send_location.Good change to use
self.send_locationinstead ofself.location, which ensures that the resized image is used if available.
381-381: Consistency improvement by using send_location.Good change to use
self.send_locationinstead ofself.location, which ensures that the resized image is used if available.
391-391: Fixed method name typo.Corrected the method name from
get_dispay_infotoget_display_info, which improves code consistency and readability.
498-500: Added directory creation to prevent failures.Good improvement to create the optimized_images directory if it doesn't exist, preventing potential failures when saving resized images.
Consider adding error handling in case directory creation fails for any reason:
resize_folder = conv_folder.joinpath("optimized_images") -resize_folder.mkdir(parents=True, exist_ok=True) +try: + resize_folder.mkdir(parents=True, exist_ok=True) +except Exception as e: + log.error(f"Failed to create resize folder: {e}") + return resize_location = resize_folder / self.location.name
524-524: Simplify code by reusing parent method.Good refactoring to use
super().encode_base64()instead of duplicating the implementation. This improves maintainability and reduces the chance of inconsistent behavior.
537-537: Ensure proper MIME type handling for encoded images.The method now correctly uses the updated
encode_image()method, maintaining consistent behavior throughout the codebase.
| def _pre_handle_completion_with_stream( | ||
| self, new_block: MessageBlock, system_message: Message | None | ||
| ): |
There was a problem hiding this comment.
🧹 Nitpick (assertive)
Extended method signature
Consider enhancing the docstring to reflect that system_message is specifically of type SystemMessage | None.
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (2)
basilisk/conversation/conversation_helper.py(4 hunks)basilisk/gui/conversation_tab.py(8 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: build_app / build_windows (x64)
🔇 Additional comments (16)
basilisk/conversation/conversation_helper.py (6)
8-8: No issues with importing Any.
This import correctly supports type annotations in subsequent code segments.
11-11: Importing ValidationInfo is appropriate.
This aligns well with the new migration functions’ signatures.
84-92: Expanded signature handles both attachments and images.
AcceptingAttachmentFile | ImageFileis consistent with the updated docstring describing restore behavior for various attachment types.
103-104: Skipping non-image attachments is correct.
This ensures onlyImageFileinstances proceed to image-specific resizing.
198-211: Minimal migration to v1 is straightforward.
Settingvalue["version"] = 1matches the intended outcome. Theinfoparameter is part of Pydantic’s validator signature, so it’s fine to keep it even if not explicitly used.
240-241: Migration steps array is neatly structured.
Listing each step in sequence is a clear, extensible approach for future versions.basilisk/gui/conversation_tab.py (10)
43-43: Import ofSystemMessagelooks correct
No issues found with this addition.
1090-1090: Updated return type toSystemMessage | None
This change clarifies the return type, aligning it with the newSystemMessageusage.
1098-1098: InstantiatingSystemMessagedirectly is consistent
Ensure that theSystemMessageclass defaults to a system role if required.
1235-1235: Retrievingsystem_messagefromkwargsis valid
This optional argument handling appears correct.
1240-1244: Passingsystem_messageto_pre_handle_completion_with_stream
Use ofwx.CallAfterwith both arguments is appropriate and poses no issues.
1270-1272: Supplyingsystem_messageto_post_completion_without_stream
No concerns; the parameters align well.
1274-1283: Consider aligning type annotation with actualSystemMessageusage
Currently, the parameter is annotated asMessage | None, whileget_system_messagereturnsSystemMessage | None. For clarity and consistency, unify these if the parameter must strictly be a system message.
1320-1322: Signature acceptsMessage | Nonebut referencing a system message
If the function expects only aSystemMessage, update the annotation for consistency.
1327-1329: Docstring refers to a system message specifically
As above, ensuring the type (SystemMessage | None) matches the docstring would remove ambiguity.
1433-1433: Switching toremove_blockmethod
Using the conversation’s own method for block removal centralizes logic.
Summary by CodeRabbit
New Features
SystemMessageentity to enhance module exports.Improvements
pytestfor improved clarity and modularity.