File size: 4,274 Bytes
e5034c3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | use std::marker::PhantomData;
pub trait Transformer: Sized {
type Value;
fn transform(&mut self, value: Self::Value) -> Self::Value;
fn pipe<B>(self, other: B) -> Pipe<Self, B> {
Pipe(self, other)
}
fn when<F: Fn(&Self::Value) -> bool>(self, cond: F) -> Cond<Self, F>
where
Self: Sized,
{
Cond(self, cond)
}
}
pub struct DefaultTransformation<T>(PhantomData<T>);
impl<T> DefaultTransformation<T> {
pub fn new() -> Self {
Self(PhantomData)
}
}
impl<T> Default for DefaultTransformation<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> Transformer for DefaultTransformation<T> {
type Value = T;
fn transform(&mut self, value: Self::Value) -> Self::Value {
value
}
}
pub struct Cond<A, F>(A, F);
impl<A, F> Transformer for Cond<A, F>
where
A: Transformer,
F: Fn(&A::Value) -> bool,
{
type Value = A::Value;
fn transform(&mut self, value: Self::Value) -> Self::Value {
let f = &self.1;
if f(&value) {
self.0.transform(value)
} else {
value
}
}
}
pub struct Pipe<A, B>(A, B);
impl<A, B, V> Transformer for Pipe<A, B>
where
A: Transformer<Value = V>,
B: Transformer<Value = V>,
{
type Value = V;
fn transform(&mut self, value: Self::Value) -> Self::Value {
self.1.transform(self.0.transform(value))
}
}
// Re-export specific transformers
mod drop_reasoning_details;
mod image_handling;
mod normalize_tool_args;
mod reasoning_normalizer;
mod set_model;
mod sort_tools;
mod transform_tool_calls;
pub use drop_reasoning_details::DropReasoningDetails;
pub use image_handling::ImageHandling;
pub use normalize_tool_args::NormalizeToolCallArguments;
pub use reasoning_normalizer::ReasoningNormalizer;
pub use set_model::SetModel;
pub use sort_tools::SortTools;
pub use transform_tool_calls::TransformToolCalls;
#[cfg(test)]
mod tests {
use insta::assert_yaml_snapshot;
use pretty_assertions::assert_eq;
use serde::Serialize;
use super::*;
use crate::{
Context, ContextMessage, ToolCallFull, ToolCallId, ToolName, ToolOutput, ToolResult,
};
#[derive(Serialize)]
struct TransformationSnapshot {
transformation: String,
before: Context,
after: Context,
}
impl TransformationSnapshot {
fn new(transformation: &str, before: Context, after: Context) -> Self {
Self { transformation: transformation.to_string(), before, after }
}
}
fn create_context_with_tool_calls() -> Context {
let tool_call = ToolCallFull {
name: ToolName::new("test_tool"),
call_id: Some(ToolCallId::new("call_123")),
arguments: serde_json::json!({"param": "value"}).into(),
thought_signature: None,
};
Context::default()
.add_message(ContextMessage::system("System message"))
.add_message(ContextMessage::assistant(
"I'll help you",
None,
None,
Some(vec![tool_call]),
))
.add_tool_results(vec![ToolResult {
name: ToolName::new("test_tool"),
call_id: Some(ToolCallId::new("call_123")),
output: ToolOutput::text("Tool result text".to_string()),
}])
}
#[test]
fn test_default_transformation() {
let fixture = Context::default().add_message(ContextMessage::user("Test message", None));
let mut transformer = DefaultTransformation::<Context>::new();
let actual = transformer.transform(fixture.clone());
let expected = fixture;
assert_eq!(actual, expected);
}
#[test]
fn test_transformer_pipe() {
let fixture = create_context_with_tool_calls();
let transform_tool_calls = TransformToolCalls::new();
let image_handling = ImageHandling::new();
let mut combined = transform_tool_calls.pipe(image_handling);
let actual = combined.transform(fixture.clone());
let snapshot =
TransformationSnapshot::new("TransformToolCalls.pipe(ImageHandling)", fixture, actual);
assert_yaml_snapshot!(snapshot);
}
}
|