Skip to main content

dkp_gen_core/tools/
mod.rs

1pub mod brave;
2pub mod discovered;
3pub mod search;
4pub mod web_fetch;
5
6use std::sync::Arc;
7
8use async_trait::async_trait;
9use serde_json::Value;
10
11use crate::config::GenConfig;
12use crate::error::{GenError, GenResult};
13use brave::BraveSearchProvider;
14use discovered::{DiscoveredLog, DiscoveredSource};
15use search::SearchProvider;
16use web_fetch::WebFetchTool;
17
18/// Construct the configured `SearchProvider`, or `None` if no search API key
19/// is set (tool use may still offer `web_fetch` alone in that case).
20pub fn make_search_provider(config: &GenConfig) -> GenResult<Option<Arc<dyn SearchProvider>>> {
21    if config.search_api_key.is_empty() {
22        return Ok(None);
23    }
24    match config.search_provider.as_str() {
25        "brave" => {
26            let provider = BraveSearchProvider::new(config.search_api_key.clone())?;
27            Ok(Some(Arc::new(provider)))
28        }
29        other => Err(GenError::ToolFailed {
30            name: "web_search".to_string(),
31            message: format!("unknown search_provider '{other}' — supported providers: brave"),
32        }),
33    }
34}
35
36/// A tool advertised to the model, in a backend-agnostic shape.
37pub struct ToolSpec {
38    pub name: String,
39    pub description: String,
40    /// JSON Schema for the tool's arguments.
41    pub parameters: Value,
42}
43
44impl ToolSpec {
45    /// Render as an OpenAI-compatible `tools` array entry.
46    pub fn to_openai(&self) -> Value {
47        serde_json::json!({
48            "type": "function",
49            "function": {
50                "name": self.name,
51                "description": self.description,
52                "parameters": self.parameters,
53            }
54        })
55    }
56}
57
58/// Executes tool calls on behalf of an `LlmClient` tool-use loop.
59#[async_trait]
60pub trait ToolExecutor: Send + Sync {
61    fn specs(&self) -> Vec<ToolSpec>;
62    async fn execute(&self, name: &str, arguments: &Value) -> GenResult<String>;
63}
64
65/// The concrete tool set `dkp-gen-core` offers to generation steps:
66/// always `web_fetch`, plus `web_search` when a search provider is configured.
67/// Every fetched/discovered URL is appended to a staging provenance log
68/// (`build/sources_discovered.jsonl`) — never written directly to `evidence/sources.csv`.
69pub struct GenToolExecutor {
70    web_fetch: WebFetchTool,
71    search: Option<Arc<dyn SearchProvider>>,
72    discovered: Arc<DiscoveredLog>,
73    command: String,
74    verbose: bool,
75}
76
77impl GenToolExecutor {
78    pub fn new(
79        http_timeout_secs: u64,
80        discovered: Arc<DiscoveredLog>,
81        command: impl Into<String>,
82        search: Option<Arc<dyn SearchProvider>>,
83        verbose: bool,
84    ) -> GenResult<Self> {
85        Ok(Self {
86            web_fetch: WebFetchTool::new(http_timeout_secs)?,
87            search,
88            discovered,
89            command: command.into(),
90            verbose,
91        })
92    }
93}
94
95#[async_trait]
96impl ToolExecutor for GenToolExecutor {
97    fn specs(&self) -> Vec<ToolSpec> {
98        let mut specs = vec![ToolSpec {
99            name: "web_fetch".to_string(),
100            description: "Fetch a URL and return its content as plain text.".to_string(),
101            parameters: serde_json::json!({
102                "type": "object",
103                "properties": {
104                    "url": {"type": "string", "description": "The URL to fetch."}
105                },
106                "required": ["url"]
107            }),
108        }];
109        if self.search.is_some() {
110            specs.push(ToolSpec {
111                name: "web_search".to_string(),
112                description: "Search the web and return matching page titles, URLs, and snippets."
113                    .to_string(),
114                parameters: serde_json::json!({
115                    "type": "object",
116                    "properties": {
117                        "query": {"type": "string", "description": "The search query."}
118                    },
119                    "required": ["query"]
120                }),
121            });
122        }
123        specs
124    }
125
126    async fn execute(&self, name: &str, arguments: &Value) -> GenResult<String> {
127        match name {
128            "web_fetch" => {
129                let url = arguments["url"]
130                    .as_str()
131                    .ok_or_else(|| GenError::ToolFailed {
132                        name: name.to_string(),
133                        message: "missing required argument 'url'".to_string(),
134                    })?;
135                if self.verbose {
136                    eprintln!("    🔧 web_fetch({url})");
137                }
138                let result = self.web_fetch.fetch(url).await;
139                if self.verbose {
140                    match &result {
141                        Ok(text) => eprintln!("       ✓ fetched {} chars", text.len()),
142                        Err(e) => eprintln!("       ✗ {e}"),
143                    }
144                }
145                let text = result?;
146                self.discovered.append(&DiscoveredSource::now(
147                    url.to_string(),
148                    None,
149                    "web_fetch",
150                    None,
151                    self.command.clone(),
152                ))?;
153                Ok(text)
154            }
155            "web_search" => {
156                let provider = self.search.as_ref().ok_or_else(|| GenError::ToolFailed {
157                    name: name.to_string(),
158                    message: "web_search is not available: no search provider configured"
159                        .to_string(),
160                })?;
161                let query = arguments["query"]
162                    .as_str()
163                    .ok_or_else(|| GenError::ToolFailed {
164                        name: name.to_string(),
165                        message: "missing required argument 'query'".to_string(),
166                    })?;
167                if self.verbose {
168                    eprintln!("    🔧 web_search({query:?})");
169                }
170                let results = provider.search(query, 5).await?;
171                if self.verbose {
172                    if results.is_empty() {
173                        eprintln!("       ✓ 0 results");
174                    } else {
175                        eprintln!("       ✓ {} result(s):", results.len());
176                        for r in &results {
177                            eprintln!("         - {} ({})", r.title, r.url);
178                        }
179                    }
180                }
181                for r in &results {
182                    self.discovered.append(&DiscoveredSource::now(
183                        r.url.clone(),
184                        Some(r.title.clone()),
185                        "web_search",
186                        Some(query.to_string()),
187                        self.command.clone(),
188                    ))?;
189                }
190                Ok(format_search_results(&results))
191            }
192            other => Err(GenError::ToolFailed {
193                name: other.to_string(),
194                message: "unknown tool".to_string(),
195            }),
196        }
197    }
198}
199
200fn format_search_results(results: &[search::SearchResult]) -> String {
201    if results.is_empty() {
202        return "No results found.".to_string();
203    }
204    results
205        .iter()
206        .enumerate()
207        .map(|(i, r)| format!("{}. {} — {}\n{}", i + 1, r.title, r.url, r.snippet))
208        .collect::<Vec<_>>()
209        .join("\n\n")
210}