Skip to main content

systemprompt_models/brain/
extraction.rs

1//! Extraction inputs and outputs, the source text an extractor reads, and
2//! the model request the brain sends through the gateway.
3//!
4//! Copyright (c) systemprompt.io — Business Source License 1.1.
5//! See <https://systemprompt.io> for licensing details.
6
7use serde::{Deserialize, Serialize};
8use systemprompt_identifiers::{AiRequestId, DocumentNodeId};
9
10use super::candidate::FactCandidate;
11use super::envelope::BrainEnvelope;
12use super::kind_decl::KindDecl;
13use super::source::{FactSource, SourceSpan};
14
15pub const BRAIN_EXTRACT_PURPOSE: &str = "brain.extract";
16pub const BRAIN_CONSOLIDATE_PURPOSE: &str = "brain.consolidate";
17
18/// One addressable node of a source's text.
19#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
20pub struct SourceNode {
21    pub node_id: DocumentNodeId,
22    pub hash8: String,
23    pub text: String,
24}
25
26/// The readable text of a source and its nodes, so a cited span can be
27/// checked.
28#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
29pub struct SourceText {
30    pub text: String,
31    pub nodes: Vec<SourceNode>,
32}
33
34impl SourceText {
35    #[must_use]
36    pub fn has_span(&self, span: &SourceSpan) -> bool {
37        self.nodes.iter().any(|node| {
38            node.node_id == span.node_id
39                && node.hash8 == span.hash8
40                && node.text.matches(span.find.as_str()).count() > span.nth as usize
41        })
42    }
43}
44
45/// What one extractor is asked to read for one kind.
46#[derive(Debug, Clone, PartialEq, Serialize)]
47pub struct ExtractionInput {
48    pub envelope: BrainEnvelope,
49    pub kind: KindDecl,
50    pub source: FactSource,
51    pub text: Option<SourceText>,
52}
53
54/// The audit record of one extraction call.
55#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
56pub struct ExtractionRecord {
57    pub extractor: String,
58    pub model: Option<String>,
59    pub prompt_version: Option<String>,
60    // JSON: the extractor's raw output is kept verbatim for audit and replay.
61    pub raw_output: serde_json::Value,
62    pub input_tokens: u32,
63    pub output_tokens: u32,
64    pub ai_request_id: Option<AiRequestId>,
65}
66
67/// The candidates an extraction produced and its record.
68#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
69pub struct ExtractionOutput {
70    pub candidates: Vec<FactCandidate>,
71    pub record: ExtractionRecord,
72}
73
74/// The system and user turns of a model call.
75#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
76pub struct BrainPrompt {
77    pub system: String,
78    pub user: String,
79}
80
81/// A named strict JSON schema the model's output must satisfy.
82#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
83pub struct StrictJsonSchema {
84    pub name: String,
85    // JSON: a JSON Schema document is passed to the provider unchanged.
86    pub schema: serde_json::Value,
87}
88
89/// A model call the brain makes through the in-process gateway client.
90#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
91pub struct BrainModelRequest {
92    pub system: String,
93    pub user: String,
94    pub json_schema: StrictJsonSchema,
95    pub max_output_tokens: u32,
96    pub prompt_version: String,
97    pub purpose: String,
98}
99
100impl BrainModelRequest {
101    pub const DEFAULT_MAX_OUTPUT_TOKENS: u32 = 2048;
102
103    #[must_use]
104    pub fn new(prompt: BrainPrompt, json_schema: StrictJsonSchema) -> Self {
105        Self {
106            system: prompt.system,
107            user: prompt.user,
108            json_schema,
109            max_output_tokens: Self::DEFAULT_MAX_OUTPUT_TOKENS,
110            prompt_version: String::new(),
111            purpose: BRAIN_EXTRACT_PURPOSE.to_owned(),
112        }
113    }
114
115    #[must_use]
116    pub const fn with_max_output_tokens(mut self, max_output_tokens: u32) -> Self {
117        self.max_output_tokens = max_output_tokens;
118        self
119    }
120
121    #[must_use]
122    pub fn with_prompt_version(mut self, prompt_version: impl Into<String>) -> Self {
123        self.prompt_version = prompt_version.into();
124        self
125    }
126
127    #[must_use]
128    pub fn with_purpose(mut self, purpose: impl Into<String>) -> Self {
129        self.purpose = purpose.into();
130        self
131    }
132}
133
134/// The model's answer and its metering.
135#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
136pub struct BrainModelResponse {
137    pub model: String,
138    pub text: String,
139    // JSON: structured output parsed against the request's strict schema.
140    pub structured: Option<serde_json::Value>,
141    pub input_tokens: u32,
142    pub output_tokens: u32,
143    pub ai_request_id: Option<AiRequestId>,
144}