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        .needs_approval(true)
36        .handler(|_ctx, _args: serde_json::Value| async move { Ok(ToolOutput::text("allowed")) })
37        .build()
38        .expect("tool");
39    let runner = Runner::builder()
40        .model_provider(ScriptedModelProvider::new(
41            "scripted",
42            "approval-model",
43            vec![
44                LLMResponse::with_tool_calls(
45                    "",
46                    vec![ToolCall::from_raw_arguments(
47                        "call_1",
48                        "dangerous",
49                        json!({}),
50                    )],
51                ),
52                finish_response("done"),
53            ],
54        ))
55        .workspace(".")
56        .build()
57        .expect("runner");
58    let agent = Agent::builder("approver")
59        .instructions("Call dangerous, then finish.")
60        .model(ModelRef::named("approval-model"))
61        .tool(dangerous)
62        .build()
63        .expect("agent");
64    client_from_runner(runner, agent)
65}
66
67pub fn finish_response(message: &str) -> LLMResponse {
68    let mut args = BTreeMap::new();
69    args.insert("message".to_string(), json!(message));
70    LLMResponse::with_tool_calls(message, vec![ToolCall::new("finish", "task_finish", args)])
71}
72
73fn client_from_runner(runner: Runner, agent: Agent) -> AppServerClient {
74    let (processor, outgoing) = MessageProcessor::new_for_tests_with_runtime(
75        128,
76        runner,
77        agent,
78        SqliteThreadStore::in_memory().expect("store"),
79    );
80    AppServerClient::new_for_processor(processor, outgoing, ConnectionId::new(1))
81}