Feature: Replay from HistoryStore Cache.#204
Feature: Replay from HistoryStore Cache.#204bryannolen wants to merge 15 commits intobeelzebub-labs:mainfrom
Conversation
This feature will use a previously observed reply, rather than ask the LLM for a new response if the input command has already been seen and is present in the HistoryStore.
Only used by SSH strategy currently, enables replying to requests from HistoryStore instead of asking the LLM again, significantly saving costs.
WalkthroughThe changes introduce structured conversation tracking to the history store, including new methods for appending and querying conversations. Caching logic is integrated into the SSH protocol handler, enabling replay of cached command outputs when configured. The configuration parser and its tests are updated to support a new Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant SSHHandler
participant HistoryStore
participant Cache
participant LLM
User->>SSHHandler: Send command input
SSHHandler->>Cache: Check EnableCacheReplay & lookup input
alt Cache hit
Cache-->>SSHHandler: Return cached output
SSHHandler->>HistoryStore: AppendConversation(input, cached output)
SSHHandler-->>User: Return cached output
else Cache miss
SSHHandler->>LLM: Process input (if plugin specified)
LLM-->>SSHHandler: Return output
SSHHandler->>HistoryStore: AppendConversation(input, output)
SSHHandler-->>User: Return output
end
Estimated code review effort3 (30–60 minutes) Suggested reviewers
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. 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 (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
♻️ Duplicate comments (1)
protocols/strategies/SSH/ssh.go (1)
155-177: Same code duplication issue as inline commands.This is the same LLM invocation logic that appears in the inline command handler. The previously suggested helper method would eliminate this duplication.
🧹 Nitpick comments (2)
historystore/history_store.go (1)
57-57: Consider using a more visible deprecation approach.While the comment deprecation is clear, consider using Go's conventional
// Deprecated:format for better tooling support.-// DEPRECATED: Use AppendConversation() instead. +// Deprecated: Use AppendConversation() instead.protocols/strategies/SSH/ssh.go (1)
62-85: Extract LLM invocation logic to reduce code duplication.The LLM plugin invocation logic is duplicated between inline commands and interactive sessions. Consider extracting this into a helper method.
func (sshStrategy *SSHStrategy) invokeLLMPlugin(servConf parser.BeelzebubServiceConfiguration, histories []plugins.Message, commandInput string) (string, error) { llmProvider, err := plugins.FromStringToLLMProvider(servConf.Plugin.LLMProvider) if err != nil { log.Errorf("error: %s, fallback OpenAI", err.Error()) llmProvider = plugins.OpenAI } llmHoneypot := plugins.LLMHoneypot{ Histories: histories, OpenAIKey: servConf.Plugin.OpenAISecretKey, Protocol: tracer.SSH, Host: servConf.Plugin.Host, Model: servConf.Plugin.LLMModel, Provider: llmProvider, CustomPrompt: servConf.Plugin.Prompt, } llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot) return llmHoneypotInstance.ExecuteModel(commandInput) }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
historystore/history_store.go(3 hunks)historystore/history_store_test.go(2 hunks)parser/configurations_parser.go(1 hunks)parser/configurations_parser_test.go(1 hunks)protocols/strategies/SSH/ssh.go(4 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (2)
historystore/history_store_test.go (2)
historystore/history_store.go (4)
NewHistoryStore(34-38)Conversation(29-31)HistoryStore(16-19)HistoryEvent(22-26)plugins/llm-integration.go (2)
Message(60-63)Role(65-65)
historystore/history_store.go (1)
plugins/llm-integration.go (1)
Message(60-63)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Analyze (go)
🔇 Additional comments (9)
parser/configurations_parser.go (1)
77-77: LGTM! Configuration field properly added.The
EnableCacheReplayfield is correctly added with appropriate YAML tagging and follows the existing struct pattern.historystore/history_store.go (3)
28-31: LGTM! Well-structured conversation type.The
Conversationstruct appropriately pairs input and output messages, providing a clean abstraction for conversation tracking.
23-26: LGTM! HistoryEvent extension looks good.The addition of the
Conversationsfield maintains backward compatibility while adding the new conversation tracking capability.
103-114: LGTM! Query method is well-implemented.The
QueryConversationsmethod properly handles thread safety and follows a clear search pattern. The logic correctly returns a pointer to the found conversation or nil if not found.protocols/strategies/SSH/ssh.go (4)
45-58: LGTM! Cache checking logic is well-implemented.The cache replay logic correctly checks for existing conversations before proceeding with command execution. The implementation properly handles the case when no cached answer exists.
113-113: LGTM! Good code organization.The comment clearly marks the end of the raw command handler block, improving code readability.
145-152: LGTM! Consistent caching logic in interactive sessions.The cache checking logic for interactive sessions mirrors the inline command logic correctly, maintaining consistency across both execution paths.
186-189: LGTM! Consistent handler name modification.The handler name modification to append " cached" when using cached responses is consistent with the inline command implementation and provides good traceability.
historystore/history_store_test.go (1)
99-134: LGTM! Comprehensive query testing.Excellent test coverage for the
QueryConversationsmethod, including all edge cases:
- Found conversation
- Not found conversation
- Non-existent key
- Empty conversations list
The test structure using subtests is clean and well-organized.
|
Embarrasing typo fixed in commit 1f9687c 😄 |
|
Welcome back mate 🙂 I'll watch it in the next few days |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #204 +/- ##
==========================================
+ Coverage 82.94% 84.02% +1.08%
==========================================
Files 7 7
Lines 428 457 +29
==========================================
+ Hits 355 384 +29
Misses 63 63
Partials 10 10 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
Great job, @bryannolen , I really like the new codebase. Below is a use case that shows the honeypot with rather low interaction(missing test-1 from second ls): Anyway, while testing your branch, I noticed you’re working on a rate limit feature, fantastic, this could be really useful 🙂 Thanks again for your time and effort. Cheers Mario |
2 participants
Extension to HistoryStore, and new (optional) config flag for SSH strategy.
New Service Config flag
enableCacheReplay- defaultfalse: When enabled, the SSH strategy will prefer to use previous successful responses stored in the HistoryStore instead of asking the LLM for a response. This significantly reduces the number of calls needed to be made to the LLM and reduces costs.Currently only supported by SSH strategy.
There are also a number of small incremental improvements to the HistoryStore, inproving efficiency and speed.
Code changes have been tested in a local honeypot for the last few weeks.
New Feature Submissions:
Changes to Core Features:
Summary by CodeRabbit
New Features
Bug Fixes
Tests