Skip to main content

synaptic_lark/tools/
doc_process.rs

1use async_trait::async_trait;
2use serde_json::{json, Value};
3use synaptic_core::{SynapticError, Tool};
4
5use crate::{auth::TokenCache, LarkConfig};
6
7/// Extract structured data from documents using the Lark intelligent document processing API.
8///
9/// Supports invoice, receipt, ID card, and general document extraction.
10pub struct LarkDocProcessTool {
11    token_cache: TokenCache,
12    base_url: String,
13    client: reqwest::Client,
14}
15
16impl LarkDocProcessTool {
17    /// Create a new tool using the given config.
18    pub fn new(config: LarkConfig) -> Self {
19        let base_url = config.base_url.clone();
20        Self {
21            token_cache: config.token_cache(),
22            base_url,
23            client: reqwest::Client::new(),
24        }
25    }
26}
27
28#[async_trait]
29impl Tool for LarkDocProcessTool {
30    fn name(&self) -> &'static str {
31        "lark_doc_process"
32    }
33
34    fn description(&self) -> &'static str {
35        "Extract structured data from documents (PDF, images) using Lark intelligent document \
36         processing. Supports invoice, receipt, ID card, and general document extraction."
37    }
38
39    fn parameters(&self) -> Option<Value> {
40        Some(json!({
41            "type": "object",
42            "properties": {
43                "file_key": {
44                    "type": "string",
45                    "description": "Feishu file_key of the document"
46                },
47                "task_type": {
48                    "type": "string",
49                    "description": "Extraction type: invoice | receipt | id_card | general",
50                    "enum": ["invoice", "receipt", "id_card", "general"]
51                }
52            },
53            "required": ["file_key", "task_type"]
54        }))
55    }
56
57    async fn call(&self, args: Value) -> Result<Value, SynapticError> {
58        let file_key = args["file_key"].as_str().ok_or_else(|| {
59            SynapticError::Tool("lark_doc_process: missing 'file_key'".to_string())
60        })?;
61        let task_type = args["task_type"].as_str().ok_or_else(|| {
62            SynapticError::Tool("lark_doc_process: missing 'task_type'".to_string())
63        })?;
64        let token = self.token_cache.get_token().await?;
65        let body = json!({ "file_key": file_key, "task_type": task_type });
66        let url = format!("{}/document_ai/v1/entity/recognize", self.base_url);
67        let resp = self
68            .client
69            .post(&url)
70            .bearer_auth(&token)
71            .json(&body)
72            .send()
73            .await
74            .map_err(|e| SynapticError::Tool(format!("lark_doc_process: {e}")))?;
75        let rb: Value = resp
76            .json()
77            .await
78            .map_err(|e| SynapticError::Tool(format!("lark_doc_process parse: {e}")))?;
79        if rb["code"].as_i64().unwrap_or(-1) != 0 {
80            return Err(SynapticError::Tool(format!(
81                "lark_doc_process API error: {}",
82                rb["msg"].as_str().unwrap_or("unknown")
83            )));
84        }
85        Ok(rb["data"].clone())
86    }
87}