stakpak_api/local/hooks/file_scratchpad_context/
mod.rs

1use crate::local::context_managers::file_scratchpad_context_manager::{
2    FileScratchpadContextManager, FileScratchpadContextManagerOptions,
3};
4use crate::local::{ModelOptions, ModelSet};
5use crate::models::AgentState;
6use stakpak_shared::define_hook;
7use stakpak_shared::hooks::{Hook, HookAction, HookContext, HookError, LifecycleEvent};
8use stakpak_shared::models::integrations::openai::Role;
9use stakpak_shared::models::llm::{LLMInput, LLMMessage, LLMMessageContent};
10
11const SYSTEM_PROMPT: &str = include_str!("./system_prompt.txt");
12const SCRATCHPAD_FILE: &str = ".stakpak/session/scratchpad.md";
13const TODO_FILE: &str = ".stakpak/session/todo.md";
14
15pub struct FileScratchpadContextHook {
16    pub model_set: ModelSet,
17    pub context_manager: FileScratchpadContextManager,
18}
19
20pub struct FileScratchpadContextHookOptions {
21    pub model_options: ModelOptions,
22    pub scratchpad_path: Option<String>,
23    pub todo_path: Option<String>,
24    pub history_action_message_size_limit: Option<usize>,
25    pub history_action_message_keep_last_n: Option<usize>,
26    pub history_action_result_keep_last_n: Option<usize>,
27    pub overwrite_if_different: Option<bool>,
28}
29
30impl FileScratchpadContextHook {
31    pub fn new(options: FileScratchpadContextHookOptions) -> Self {
32        let model_set: ModelSet = options.model_options.into();
33        let context_manager =
34            FileScratchpadContextManager::new(FileScratchpadContextManagerOptions {
35                scratchpad_file_path: options
36                    .scratchpad_path
37                    .unwrap_or(SCRATCHPAD_FILE.to_string())
38                    .into(),
39                todo_file_path: options.todo_path.unwrap_or(TODO_FILE.to_string()).into(),
40                history_action_message_size_limit: options
41                    .history_action_message_size_limit
42                    .unwrap_or(100),
43                history_action_message_keep_last_n: options
44                    .history_action_message_keep_last_n
45                    .unwrap_or(1),
46                history_action_result_keep_last_n: options
47                    .history_action_result_keep_last_n
48                    .unwrap_or(50),
49                overwrite_if_different: options.overwrite_if_different.unwrap_or(true),
50            });
51
52        Self {
53            model_set,
54            context_manager,
55        }
56    }
57}
58
59define_hook!(
60    FileScratchpadContextHook,
61    "file_scratchpad_context",
62    async |&self, ctx: &mut HookContext<AgentState>, event: &LifecycleEvent| {
63        if *event != LifecycleEvent::BeforeInference {
64            return Ok(HookAction::Continue);
65        }
66
67        let model = self.model_set.get_model(&ctx.state.agent_model);
68
69        let tools = ctx
70            .state
71            .tools
72            .clone()
73            .map(|t| t.into_iter().map(Into::into).collect());
74
75        let cwd = std::env::current_dir().unwrap_or_default();
76        let scratchpad_file_path =
77            cwd.join(self.context_manager.get_scratchpad_path(ctx.session_id));
78        let todo_file_path = cwd.join(self.context_manager.get_todo_path(ctx.session_id));
79        let system_prompt = SYSTEM_PROMPT
80            .replace(
81                "{{SCRATCHPAD_PATH}}",
82                &scratchpad_file_path.display().to_string(),
83            )
84            .replace("{{TODO_PATH}}", &todo_file_path.display().to_string());
85
86        let mut messages = Vec::new();
87        messages.push(LLMMessage {
88            role: Role::System.to_string(),
89            content: LLMMessageContent::String(system_prompt),
90        });
91        messages.extend(
92            self.context_manager
93                .reduce_context_with_session(ctx.state.messages.clone(), ctx.session_id),
94        );
95
96        ctx.state.llm_input = Some(LLMInput {
97            model,
98            messages,
99            max_tokens: 16000,
100            tools,
101            provider_options: None,
102            headers: None,
103        });
104
105        Ok(HookAction::Continue)
106    }
107);
stakpak_api/local/hooks/file_scratchpad_context/mod.rs

stakpak_api/local/hooks/file_scratchpad_context/
mod.rs
