HookのPod側への移動・Interceptorの実装
This commit is contained in:
@@ -8,11 +8,8 @@ use std::time::{Duration, Instant};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_worker::Worker;
|
||||
use llm_worker::hook::{
|
||||
Hook, HookError, PostToolCall, PostToolCallContext, PostToolCallResult, PreToolCall,
|
||||
PreToolCallResult, ToolCallContext,
|
||||
};
|
||||
use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent};
|
||||
use llm_worker::interceptor::{Interceptor, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo};
|
||||
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta};
|
||||
|
||||
mod common;
|
||||
@@ -156,21 +153,21 @@ async fn test_before_tool_call_skip() {
|
||||
worker.register_tool(allowed_tool.definition()).unwrap();
|
||||
worker.register_tool(blocked_tool.definition()).unwrap();
|
||||
|
||||
// Hook to skip "blocked_tool"
|
||||
struct BlockingHook;
|
||||
// Policy to skip "blocked_tool"
|
||||
struct BlockingPolicy;
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreToolCall> for BlockingHook {
|
||||
async fn call(&self, ctx: &mut ToolCallContext) -> Result<PreToolCallResult, HookError> {
|
||||
if ctx.call.name == "blocked_tool" {
|
||||
Ok(PreToolCallResult::Skip)
|
||||
impl Interceptor for BlockingPolicy {
|
||||
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction {
|
||||
if info.call.name == "blocked_tool" {
|
||||
PreToolAction::Skip
|
||||
} else {
|
||||
Ok(PreToolCallResult::Continue)
|
||||
PreToolAction::Continue
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
worker.add_pre_tool_call_hook(BlockingHook);
|
||||
worker.set_interceptor(BlockingPolicy);
|
||||
|
||||
let _result = worker.run("Test hook").await;
|
||||
|
||||
@@ -235,25 +232,22 @@ async fn test_post_tool_call_modification() {
|
||||
|
||||
worker.register_tool(simple_tool_definition()).unwrap();
|
||||
|
||||
// Hook to modify results
|
||||
struct ModifyingHook {
|
||||
// Policy to modify results
|
||||
struct ModifyingPolicy {
|
||||
modified_content: Arc<std::sync::Mutex<Option<String>>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PostToolCall> for ModifyingHook {
|
||||
async fn call(
|
||||
&self,
|
||||
ctx: &mut PostToolCallContext,
|
||||
) -> Result<PostToolCallResult, HookError> {
|
||||
ctx.result.content = format!("[Modified] {}", ctx.result.content);
|
||||
*self.modified_content.lock().unwrap() = Some(ctx.result.content.clone());
|
||||
Ok(PostToolCallResult::Continue)
|
||||
impl Interceptor for ModifyingPolicy {
|
||||
async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction {
|
||||
info.result.content = format!("[Modified] {}", info.result.content);
|
||||
*self.modified_content.lock().unwrap() = Some(info.result.content.clone());
|
||||
PostToolAction::Continue
|
||||
}
|
||||
}
|
||||
|
||||
let modified_content = Arc::new(std::sync::Mutex::new(None));
|
||||
worker.add_post_tool_call_hook(ModifyingHook {
|
||||
worker.set_interceptor(ModifyingPolicy {
|
||||
modified_content: modified_content.clone(),
|
||||
});
|
||||
|
||||
|
||||
@@ -1,194 +0,0 @@
|
||||
//! Streaming hook tests
|
||||
|
||||
mod common;
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use common::MockLlmClient;
|
||||
use llm_worker::hook::{
|
||||
Hook, HookError, OnStreamChunk, OnStreamComplete, OnTextDelta, OnToolCallDelta,
|
||||
StreamChunkContext, StreamCompleteContext, StreamHookResult, TextDeltaContext,
|
||||
ToolCallDeltaContext,
|
||||
};
|
||||
use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent};
|
||||
use llm_worker::{Worker, WorkerError};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_text_delta_hooks_run_in_registration_order() {
|
||||
let events = vec![
|
||||
Event::text_block_start(0),
|
||||
Event::text_delta(0, "A"),
|
||||
Event::text_delta(0, "B"),
|
||||
Event::text_block_stop(0, None),
|
||||
Event::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}),
|
||||
];
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut worker = Worker::new(client);
|
||||
|
||||
struct RecorderHook {
|
||||
label: &'static str,
|
||||
records: Arc<Mutex<Vec<String>>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<OnTextDelta> for RecorderHook {
|
||||
async fn call(&self, input: &mut TextDeltaContext) -> Result<StreamHookResult, HookError> {
|
||||
self.records
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("{}:{}", self.label, input.delta));
|
||||
Ok(StreamHookResult::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
let records = Arc::new(Mutex::new(Vec::new()));
|
||||
worker.add_on_text_delta_hook(RecorderHook {
|
||||
label: "first",
|
||||
records: records.clone(),
|
||||
});
|
||||
worker.add_on_text_delta_hook(RecorderHook {
|
||||
label: "second",
|
||||
records: records.clone(),
|
||||
});
|
||||
|
||||
let result = worker.run("hello").await;
|
||||
assert!(result.is_ok(), "run should succeed: {result:?}");
|
||||
|
||||
let got = records.lock().unwrap().clone();
|
||||
assert_eq!(
|
||||
got,
|
||||
vec![
|
||||
"first:A".to_string(),
|
||||
"second:A".to_string(),
|
||||
"first:B".to_string(),
|
||||
"second:B".to_string(),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_stream_chunk_and_stream_complete_hooks_are_called() {
|
||||
let events = vec![
|
||||
Event::ping(),
|
||||
Event::text_block_start(0),
|
||||
Event::text_delta(0, "hi"),
|
||||
Event::text_block_stop(0, None),
|
||||
Event::usage(10, 5),
|
||||
Event::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}),
|
||||
];
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut worker = Worker::new(client);
|
||||
|
||||
struct ChunkCounter(Arc<Mutex<usize>>);
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<OnStreamChunk> for ChunkCounter {
|
||||
async fn call(
|
||||
&self,
|
||||
_input: &mut StreamChunkContext,
|
||||
) -> Result<StreamHookResult, HookError> {
|
||||
let mut guard = self.0.lock().unwrap();
|
||||
*guard += 1;
|
||||
Ok(StreamHookResult::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
struct CompleteRecorder(Arc<Mutex<Vec<(usize, usize)>>>);
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<OnStreamComplete> for CompleteRecorder {
|
||||
async fn call(
|
||||
&self,
|
||||
input: &mut StreamCompleteContext,
|
||||
) -> Result<StreamHookResult, HookError> {
|
||||
self.0.lock().unwrap().push((input.turn, input.event_count));
|
||||
Ok(StreamHookResult::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
let chunk_count = Arc::new(Mutex::new(0usize));
|
||||
let completes = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
worker.add_on_stream_chunk_hook(ChunkCounter(chunk_count.clone()));
|
||||
worker.add_on_stream_complete_hook(CompleteRecorder(completes.clone()));
|
||||
|
||||
let result = worker.run("hello").await;
|
||||
assert!(result.is_ok(), "run should succeed: {result:?}");
|
||||
|
||||
assert_eq!(*chunk_count.lock().unwrap(), 6);
|
||||
assert_eq!(completes.lock().unwrap().as_slice(), &[(0usize, 6usize)]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tool_call_delta_hook_can_abort_run() {
|
||||
let events = vec![
|
||||
Event::tool_use_start(0, "call_1", "unknown_tool"),
|
||||
Event::tool_input_delta(0, r#"{"x":1}"#),
|
||||
Event::tool_use_stop(0),
|
||||
Event::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}),
|
||||
];
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut worker = Worker::new(client);
|
||||
|
||||
struct AbortToolDelta;
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<OnToolCallDelta> for AbortToolDelta {
|
||||
async fn call(
|
||||
&self,
|
||||
_input: &mut ToolCallDeltaContext,
|
||||
) -> Result<StreamHookResult, HookError> {
|
||||
Ok(StreamHookResult::Abort("blocked by tool delta".to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
worker.add_on_tool_call_delta_hook(AbortToolDelta);
|
||||
|
||||
let result = worker.run("hello").await;
|
||||
match result {
|
||||
Err(WorkerError::Aborted(reason)) => assert_eq!(reason, "blocked by tool delta"),
|
||||
other => panic!("expected aborted result, got: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_stream_hook_pause_is_mapped_to_aborted() {
|
||||
let events = vec![
|
||||
Event::text_block_start(0),
|
||||
Event::text_delta(0, "pause me"),
|
||||
Event::text_block_stop(0, None),
|
||||
Event::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}),
|
||||
];
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut worker = Worker::new(client);
|
||||
|
||||
struct PauseHook;
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<OnTextDelta> for PauseHook {
|
||||
async fn call(&self, _input: &mut TextDeltaContext) -> Result<StreamHookResult, HookError> {
|
||||
Ok(StreamHookResult::Pause)
|
||||
}
|
||||
}
|
||||
|
||||
worker.add_on_text_delta_hook(PauseHook);
|
||||
|
||||
let result = worker.run("hello").await;
|
||||
match result {
|
||||
Err(WorkerError::Aborted(reason)) => assert_eq!(reason, "Paused by stream hook"),
|
||||
other => panic!("expected aborted result, got: {other:?}"),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user