use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use crate::client::Client;
use crate::error::Result;
#[derive(Debug, Clone, Serialize, Default)]
pub struct DocumentRequest {
pub file_base64: String,
pub filename: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_format: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct DocumentResponse {
pub content: String,
pub format: String,
#[serde(default)]
pub meta: Option<HashMap<String, serde_json::Value>>,
#[serde(default)]
pub cost_ticks: i64,
#[serde(default)]
pub request_id: String,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct ChunkDocumentRequest {
pub file_base64: String,
pub filename: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_chunk_tokens: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub overlap_tokens: Option<i32>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct DocumentChunk {
pub index: i32,
pub text: String,
#[serde(default)]
pub token_count: Option<i32>,
}
pub type ChunkRequest = ChunkDocumentRequest;
#[derive(Debug, Clone, Deserialize)]
pub struct ChunkDocumentResponse {
pub chunks: Vec<DocumentChunk>,
#[serde(default)]
pub total_chunks: Option<i32>,
#[serde(default)]
pub cost_ticks: i64,
#[serde(default)]
pub request_id: String,
}
pub type ChunkResponse = ChunkDocumentResponse;
#[derive(Debug, Clone, Serialize, Default)]
pub struct ProcessDocumentRequest {
pub file_base64: String,
pub filename: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
}
pub type ProcessRequest = ProcessDocumentRequest;
#[derive(Debug, Clone, Deserialize)]
pub struct ProcessDocumentResponse {
pub content: String,
#[serde(default)]
pub model: Option<String>,
#[serde(default)]
pub cost_ticks: i64,
#[serde(default)]
pub request_id: String,
}
pub type ProcessResponse = ProcessDocumentResponse;
impl Client {
pub async fn extract_document(&self, req: &DocumentRequest) -> Result<DocumentResponse> {
let (mut resp, meta) = self
.post_json::<DocumentRequest, DocumentResponse>("/qai/v1/documents/extract", req)
.await?;
if resp.cost_ticks == 0 {
resp.cost_ticks = meta.cost_ticks;
}
if resp.request_id.is_empty() {
resp.request_id = meta.request_id;
}
Ok(resp)
}
pub async fn chunk_document(&self, req: &ChunkDocumentRequest) -> Result<ChunkDocumentResponse> {
let (mut resp, meta) = self
.post_json::<ChunkDocumentRequest, ChunkDocumentResponse>("/qai/v1/documents/chunk", req)
.await?;
if resp.cost_ticks == 0 {
resp.cost_ticks = meta.cost_ticks;
}
if resp.request_id.is_empty() {
resp.request_id = meta.request_id;
}
Ok(resp)
}
pub async fn process_document(&self, req: &ProcessDocumentRequest) -> Result<ProcessDocumentResponse> {
let (mut resp, meta) = self
.post_json::<ProcessDocumentRequest, ProcessDocumentResponse>("/qai/v1/documents/process", req)
.await?;
if resp.cost_ticks == 0 {
resp.cost_ticks = meta.cost_ticks;
}
if resp.request_id.is_empty() {
resp.request_id = meta.request_id;
}
Ok(resp)
}
}