systemprompt_models/brain/
extraction.rs1use 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#[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#[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#[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#[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 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#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
69pub struct ExtractionOutput {
70 pub candidates: Vec<FactCandidate>,
71 pub record: ExtractionRecord,
72}
73
74#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
76pub struct BrainPrompt {
77 pub system: String,
78 pub user: String,
79}
80
81#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
83pub struct StrictJsonSchema {
84 pub name: String,
85 pub schema: serde_json::Value,
87}
88
89#[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#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
136pub struct BrainModelResponse {
137 pub model: String,
138 pub text: String,
139 pub structured: Option<serde_json::Value>,
141 pub input_tokens: u32,
142 pub output_tokens: u32,
143 pub ai_request_id: Option<AiRequestId>,
144}