Skip to main content

vv_agent/app_server/
test_support.rs

1use std::collections::BTreeMap;
2
3use serde_json::json;
4
5use crate::app_server::client::AppServerClient;
6use crate::app_server::processor::MessageProcessor;
7use crate::app_server::thread_store::SqliteThreadStore;
8use crate::app_server::transport::ConnectionId;
9use crate::{
10    Agent, FunctionTool, LLMResponse, ModelRef, Runner, ScriptedModelProvider, ToolCall, ToolOutput,
11};
12
13pub fn scripted_app_server_client(responses: Vec<LLMResponse>) -> AppServerClient {
14    let runner = Runner::builder()
15        .model_provider(ScriptedModelProvider::new(
16            "scripted",
17            "demo-model",
18            responses,
19        ))
20        .workspace(".")
21        .build()
22        .expect("runner");
23    let agent = Agent::builder("assistant")
24        .instructions("Answer the user, then finish.")
25        .model(ModelRef::named("demo-model"))
26        .build()
27        .expect("agent");
28    client_from_runner(runner, agent)
29}
30
31pub fn approval_app_server_client() -> AppServerClient {
32    let dangerous = FunctionTool::builder("dangerous")
33        .description("Requires approval.")
34        .json_schema(json!({"type":"object","properties":{},"required":[]}))
35        .handler(|_ctx, _args: serde_json::Value| async move { Ok(ToolOutput::text("allowed")) })
36        .build()
37        .expect("tool");
38    let runner = Runner::builder()
39        .model_provider(ScriptedModelProvider::new(
40            "scripted",
41            "approval-model",
42            vec![
43                LLMResponse::with_tool_calls(
44                    "",
45                    vec![ToolCall::from_raw_arguments(
46                        "call_1",
47                        "dangerous",
48                        json!({}),
49                    )],
50                ),
51                finish_response("done"),
52            ],
53        ))
54        .workspace(".")
55        .build()
56        .expect("runner");
57    let agent = Agent::builder("approver")
58        .instructions("Call dangerous, then finish.")
59        .model(ModelRef::named("approval-model"))
60        .tool(dangerous)
61        .build()
62        .expect("agent");
63    client_from_runner(runner, agent)
64}
65
66pub fn finish_response(message: &str) -> LLMResponse {
67    let mut args = BTreeMap::new();
68    args.insert("message".to_string(), json!(message));
69    LLMResponse::with_tool_calls(message, vec![ToolCall::new("finish", "task_finish", args)])
70}
71
72fn client_from_runner(runner: Runner, agent: Agent) -> AppServerClient {
73    let (processor, outgoing) = MessageProcessor::new_for_tests_with_runtime(
74        128,
75        runner,
76        agent,
77        SqliteThreadStore::in_memory().expect("store"),
78    );
79    AppServerClient::new_for_processor(processor, outgoing, ConnectionId::new(1))
80}