use rucora::agent::extractor::Extractor;
use rucora::provider::OpenAiProvider;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use tracing::{Level, info};
use tracing_subscriber::FmtSubscriber;
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
struct Person {
#[schemars(required)]
pub name: Option<String>,
#[schemars(required)]
pub age: Option<u8>,
#[schemars(required)]
pub profession: Option<String>,
#[schemars(required)]
pub city: Option<String>,
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
dotenv::dotenv().ok();
let subscriber = FmtSubscriber::builder()
.with_max_level(Level::INFO)
.with_target(false)
.finish();
tracing::subscriber::set_global_default(subscriber)?;
info!("╔════════════════════════════════════════╗");
info!("║ rucora 结构化数据提取示例 ║");
info!("╚════════════════════════════════════════╝\n");
if std::env::var("OPENAI_API_KEY").is_err() && std::env::var("OPENAI_BASE_URL").is_err() {
info!("⚠ 未设置 API 配置");
info!(" 使用 OpenAI: export OPENAI_API_KEY=sk-your-key");
info!(" 使用 Ollama: export OPENAI_BASE_URL=http://localhost:11434");
return Ok(());
}
let model_name = std::env::var("MODEL_NAME").expect("没有设置环境变量MODEL_NAME");
info!("1. 创建 Provider...");
let provider = OpenAiProvider::from_env()?;
info!("✓ Provider 创建成功\n");
info!("═══════════════════════════════════════");
info!("示例 1: 提取人物信息");
info!("═══════════════════════════════════════\n");
let person_extractor = Extractor::<Person>::builder(provider, model_name)
.preamble("只提取明确提到的信息,不要推测。")
.retries(2)
.build();
let text1 = "John Doe 是 30 岁的软件工程师,住在纽约。";
info!("输入文本:{}", text1);
match person_extractor.extract(text1).await {
Ok(person) => {
info!("✓ 提取成功!");
info!(" 姓名:{:?}", person.name);
info!(" 年龄:{:?}", person.age);
info!(" 职业:{:?}", person.profession);
info!(" 城市:{:?}\n", person.city);
}
Err(e) => {
info!("✗ 提取失败:{}\n", e);
}
}
info!("═══════════════════════════════════════");
info!("示例 2: 带 Usage 追踪");
info!("═══════════════════════════════════════\n");
let text2 = "Jane Smith 是 25 岁的数据科学家,在旧金山工作。";
info!("输入文本:{}", text2);
match person_extractor.extract_with_usage(text2).await {
Ok(response) => {
info!("✓ 提取成功!");
info!(" 姓名:{:?}", response.data.name);
info!(" 年龄:{:?}", response.data.age);
info!(" 职业:{:?}", response.data.profession);
info!(" 城市:{:?}", response.data.city);
if let Some(usage) = response.usage {
info!("\n Token 使用统计:");
info!(" 输入 token: {}", usage.input_tokens);
info!(" 输出 token: {}", usage.output_tokens);
info!(" 总计 token: {}\n", usage.total_tokens);
}
}
Err(e) => {
info!("✗ 提取失败:{}\n", e);
}
}
info!("═══════════════════════════════════════");
info!("示例完成!");
info!("═══════════════════════════════════════");
Ok(())
}