nestrs_ai/
lib.rs

1//! NestJS [`@nestjs/ai`](https://docs.nestjs.com/ai/overview) analogue —
2//! named [`LanguageModel`]s, [`AiService::generate_text`] / [`AiService::generate_object`],
3//! agent tools via [`AgentToolRegistry`], and [`FakeModel`] for tests.
4//!
5//! Model calls run under [`nestrs_resilience`] (timeout + jittered retry).
6//! [`AiService::generate_object`] **fails closed** on invalid JSON or missing required
7//! schema fields. Tool handlers must enforce the **caller's** authorization.
8//!
9//! Umbrella feature: `nestrs` / `ai` (implies `resilience`) → `nestrs::ai`.
10#![doc(html_root_url = "https://docs.rs/nestrs-ai/1.7.1")]
11
12mod fake;
13mod module;
14mod options;
15mod service;
16mod tools;
17
18pub use fake::FakeModel;
19pub use module::AiModule;
20pub use options::{AiModelConfig, AiOptions, GenerateObjectRequest, GenerateTextRequest};
21pub use service::{AiError, AiService, GenerateResult, LanguageModel};
22pub use tools::{AgentFunction, AgentToolRegistry};
23
24pub use nestrs_core::{DynamicModule, Injectable, ProviderRegistry};
25
26#[cfg(test)]
27mod tests {
28    use super::*;
29    use serde_json::json;
30    use std::sync::Arc;
31
32    #[tokio::test]
33    async fn fake_model_generate_text() {
34        let fake = FakeModel::new(vec![Ok(json!({"text": "hello"}))]);
35        let mut models = std::collections::HashMap::new();
36        models.insert(
37            "default".to_string(),
38            Arc::new(fake) as Arc<dyn LanguageModel>,
39        );
40        let svc = AiService::new(AiOptions {
41            models,
42            defaults_timeout_ms: 5_000,
43            max_steps: 5,
44        });
45        let out = svc
46            .generate_text(GenerateTextRequest {
47                model: None,
48                label: "t".into(),
49                prompt: "hi".into(),
50                instructions: None,
51                tools: vec![],
52            })
53            .await
54            .expect("text");
55        assert_eq!(out.text, "hello");
56    }
57
58    #[tokio::test]
59    async fn agent_loop_calls_tool() {
60        let fake = FakeModel::new(vec![
61            Ok(json!({
62                "tool_calls": [{"name": "echo", "arguments": {"x": 1}}],
63                "text": ""
64            })),
65            Ok(json!({"text": "done"})),
66        ]);
67        let mut models = std::collections::HashMap::new();
68        models.insert(
69            "default".to_string(),
70            Arc::new(fake) as Arc<dyn LanguageModel>,
71        );
72        let tools = AgentToolRegistry::new();
73        tools.register(AgentFunction {
74            name: "echo".into(),
75            description: "echo".into(),
76            handler: Arc::new(|args| Box::pin(async move { Ok(args) })),
77        });
78        let svc = AiService::new(AiOptions {
79            models,
80            defaults_timeout_ms: 5_000,
81            max_steps: 5,
82        })
83        .with_tools(tools);
84        let out = svc
85            .generate_text(GenerateTextRequest {
86                model: None,
87                label: "agent".into(),
88                prompt: "run".into(),
89                instructions: None,
90                tools: vec!["echo".into()],
91            })
92            .await
93            .expect("agent");
94        assert_eq!(out.text, "done");
95        assert_eq!(out.steps.len(), 2);
96    }
97
98    #[tokio::test]
99    async fn generate_object_fails_closed_on_invalid_json() {
100        let fake = FakeModel::new(vec![Ok(json!({"text": "not-json"}))]);
101        let mut models = std::collections::HashMap::new();
102        models.insert(
103            "default".to_string(),
104            Arc::new(fake) as Arc<dyn LanguageModel>,
105        );
106        let svc = AiService::new(AiOptions {
107            models,
108            defaults_timeout_ms: 5_000,
109            max_steps: 1,
110        });
111        let schema = schemars::json_schema!({
112            "type": "object"
113        });
114        let err = svc
115            .generate_object(GenerateObjectRequest {
116                model: None,
117                label: "obj".into(),
118                prompt: "give object".into(),
119                instructions: None,
120                schema_name: "Obj".into(),
121                schema,
122            })
123            .await
124            .expect_err("must fail closed");
125        assert!(matches!(err, AiError::Output(_)));
126    }
127
128    #[tokio::test]
129    async fn generate_object_fails_closed_on_missing_required() {
130        #[derive(serde::Serialize, schemars::JsonSchema)]
131        struct NeedName {
132            name: String,
133        }
134
135        let fake = FakeModel::new(vec![Ok(json!({"text": "{}"}))]);
136        let mut models = std::collections::HashMap::new();
137        models.insert(
138            "default".to_string(),
139            Arc::new(fake) as Arc<dyn LanguageModel>,
140        );
141        let svc = AiService::new(AiOptions {
142            models,
143            defaults_timeout_ms: 5_000,
144            max_steps: 1,
145        });
146        let schema = schemars::schema_for!(NeedName);
147        let err = svc
148            .generate_object(GenerateObjectRequest {
149                model: None,
150                label: "obj".into(),
151                prompt: "give object".into(),
152                instructions: None,
153                schema_name: "NeedName".into(),
154                schema,
155            })
156            .await
157            .expect_err("missing required must fail");
158        assert!(
159            matches!(&err, AiError::Output(m) if m.contains("name")),
160            "unexpected: {err}"
161        );
162    }
163}
nestrs_ai/lib.rs

nestrs_ai/
lib.rs
