phago_llm/
types.rs

1//! Core types for LLM integration.
2
3use serde::{Deserialize, Serialize};
4
5/// Type of concept extracted by LLM.
6#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
7pub enum ConceptType {
8    /// A named entity (person, place, organization).
9    Entity,
10    /// A scientific or technical concept.
11    Concept,
12    /// A process or action.
13    Process,
14    /// A property or attribute.
15    Property,
16    /// A relationship or connection.
17    Relationship,
18    /// Unknown or unclassified.
19    Other,
20}
21
22impl Default for ConceptType {
23    fn default() -> Self {
24        Self::Concept
25    }
26}
27
28/// A concept extracted from text by an LLM.
29#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct Concept {
31    /// The concept label/name.
32    pub label: String,
33    /// Type of concept.
34    pub concept_type: ConceptType,
35    /// Confidence score (0.0-1.0).
36    pub confidence: f32,
37    /// Optional description or definition.
38    pub description: Option<String>,
39    /// Related concepts mentioned in the same context.
40    pub related: Vec<String>,
41}
42
43impl Concept {
44    /// Create a new concept with just a label.
45    pub fn new(label: impl Into<String>) -> Self {
46        Self {
47            label: label.into(),
48            concept_type: ConceptType::default(),
49            confidence: 1.0,
50            description: None,
51            related: Vec::new(),
52        }
53    }
54
55    /// Set the concept type.
56    pub fn with_type(mut self, concept_type: ConceptType) -> Self {
57        self.concept_type = concept_type;
58        self
59    }
60
61    /// Set the confidence score.
62    pub fn with_confidence(mut self, confidence: f32) -> Self {
63        self.confidence = confidence.clamp(0.0, 1.0);
64        self
65    }
66
67    /// Set the description.
68    pub fn with_description(mut self, description: impl Into<String>) -> Self {
69        self.description = Some(description.into());
70        self
71    }
72
73    /// Add related concepts.
74    pub fn with_related(mut self, related: Vec<String>) -> Self {
75        self.related = related;
76        self
77    }
78}
79
80/// Type of relationship between concepts.
81#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
82pub enum RelationType {
83    /// A is a type of B (hypernym).
84    IsA,
85    /// A is part of B (meronym).
86    PartOf,
87    /// A causes B.
88    Causes,
89    /// A enables B.
90    Enables,
91    /// A requires B.
92    Requires,
93    /// A is related to B (general).
94    RelatedTo,
95    /// A produces B.
96    Produces,
97    /// A regulates B.
98    Regulates,
99    /// A interacts with B.
100    InteractsWith,
101    /// A is located in B.
102    LocatedIn,
103    /// Custom relationship type.
104    Custom(String),
105}
106
107impl Default for RelationType {
108    fn default() -> Self {
109        Self::RelatedTo
110    }
111}
112
113/// A relationship between two concepts.
114#[derive(Debug, Clone, Serialize, Deserialize)]
115pub struct Relationship {
116    /// Source concept.
117    pub source: String,
118    /// Target concept.
119    pub target: String,
120    /// Type of relationship.
121    pub relation_type: RelationType,
122    /// Relationship label (human-readable).
123    pub label: String,
124    /// Confidence score (0.0-1.0).
125    pub confidence: f32,
126    /// Whether the relationship is bidirectional.
127    pub bidirectional: bool,
128}
129
130impl Relationship {
131    /// Create a new relationship.
132    pub fn new(source: impl Into<String>, target: impl Into<String>, label: impl Into<String>) -> Self {
133        Self {
134            source: source.into(),
135            target: target.into(),
136            relation_type: RelationType::default(),
137            label: label.into(),
138            confidence: 1.0,
139            bidirectional: false,
140        }
141    }
142
143    /// Set the relationship type.
144    pub fn with_type(mut self, relation_type: RelationType) -> Self {
145        self.relation_type = relation_type;
146        self
147    }
148
149    /// Set the confidence score.
150    pub fn with_confidence(mut self, confidence: f32) -> Self {
151        self.confidence = confidence.clamp(0.0, 1.0);
152        self
153    }
154
155    /// Mark as bidirectional.
156    pub fn bidirectional(mut self) -> Self {
157        self.bidirectional = true;
158        self
159    }
160}
161
162/// Response from concept extraction.
163#[derive(Debug, Clone, Serialize, Deserialize)]
164pub struct ExtractionResponse {
165    /// Extracted concepts.
166    pub concepts: Vec<Concept>,
167    /// Identified relationships.
168    pub relationships: Vec<Relationship>,
169    /// Raw LLM response (for debugging).
170    pub raw_response: Option<String>,
171    /// Tokens used in the request.
172    pub tokens_used: Option<u32>,
173}
174
175impl ExtractionResponse {
176    /// Create an empty response.
177    pub fn empty() -> Self {
178        Self {
179            concepts: Vec::new(),
180            relationships: Vec::new(),
181            raw_response: None,
182            tokens_used: None,
183        }
184    }
185
186    /// Create from concepts only.
187    pub fn from_concepts(concepts: Vec<Concept>) -> Self {
188        Self {
189            concepts,
190            relationships: Vec::new(),
191            raw_response: None,
192            tokens_used: None,
193        }
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200
201    #[test]
202    fn test_concept_builder() {
203        let concept = Concept::new("mitochondria")
204            .with_type(ConceptType::Concept)
205            .with_confidence(0.95)
206            .with_description("Powerhouse of the cell")
207            .with_related(vec!["ATP".into(), "cell".into()]);
208
209        assert_eq!(concept.label, "mitochondria");
210        assert_eq!(concept.concept_type, ConceptType::Concept);
211        assert!((concept.confidence - 0.95).abs() < 0.001);
212        assert_eq!(concept.description, Some("Powerhouse of the cell".into()));
213        assert_eq!(concept.related.len(), 2);
214    }
215
216    #[test]
217    fn test_relationship_builder() {
218        let rel = Relationship::new("mitochondria", "ATP", "produces")
219            .with_type(RelationType::Produces)
220            .with_confidence(0.9);
221
222        assert_eq!(rel.source, "mitochondria");
223        assert_eq!(rel.target, "ATP");
224        assert_eq!(rel.relation_type, RelationType::Produces);
225    }
226}
phago_llm/types.rs

phago_llm/
types.rs
