From 2acd00ff8ea7366610437319d31b6c32b3ed1b2d Mon Sep 17 00:00:00 2001 From: Per Stark Date: Wed, 23 Oct 2024 14:20:03 +0200 Subject: [PATCH] storing nodes --- src/consumer.rs | 2 +- src/models/file_info.rs | 1 + src/models/graph_entities.rs | 36 +++--------------------------------- src/models/text_content.rs | 35 +++++++++++++++++++++++++++++------ src/utils/llm.rs | 6 +++--- 5 files changed, 37 insertions(+), 43 deletions(-) diff --git a/src/consumer.rs b/src/consumer.rs index 2183188..a6c8ca0 100644 --- a/src/consumer.rs +++ b/src/consumer.rs @@ -14,7 +14,7 @@ async fn main() -> Result<(), Box> { info!("Starting RabbitMQ consumer"); // Set up RabbitMQ config - let config = RabbitMQConfig { + let config = RabbitMQConfig { amqp_addr: "amqp://localhost".to_string(), exchange: "my_exchange".to_string(), queue: "my_queue".to_string(), diff --git a/src/models/file_info.rs b/src/models/file_info.rs index 3b28811..75789d2 100644 --- a/src/models/file_info.rs +++ b/src/models/file_info.rs @@ -322,6 +322,7 @@ impl FileInfo { /// /// # Returns /// * `Result<(), FileError>` - Empty result or an error. + async fn create_record(file_info: &FileInfo, db_client: &SurrealDbClient) -> Result<(), FileError> { // Create the record let _created: Option = db_client diff --git a/src/models/graph_entities.rs b/src/models/graph_entities.rs index 903f222..04fe74c 100644 --- a/src/models/graph_entities.rs +++ b/src/models/graph_entities.rs @@ -1,15 +1,16 @@ use serde::Deserialize; use serde::Serialize; +use std::collections::HashMap; use uuid::Uuid; /// Represents a generic knowledge entity in the graph. #[derive(Debug, Serialize, Deserialize, Clone)] pub struct KnowledgeEntity { - pub id: Uuid, // Generated in Rust + pub id: String, // Generated in Rust pub name: String, pub description: String, pub entity_type: KnowledgeEntityType, - pub source_id: Option, // Links to FileInfo or TextContent + pub source_id: String, // Links to FileInfo or TextContent pub metadata: Option, // Additional metadata } @@ -47,12 +48,6 @@ pub struct KnowledgeRelationship { pub metadata: Option, // Additional metadata } -use std::collections::HashMap; - -use crate::utils::llm::LLMGraphAnalysisResult; -use crate::utils::llm::LLMKnowledgeEntity; -use crate::utils::llm::LLMRelationship; - /// Intermediate struct to hold mapping between LLM keys and generated IDs. pub struct GraphMapper { pub key_to_id: HashMap, @@ -77,28 +72,3 @@ impl GraphMapper { self.key_to_id.get(key) } } - -impl From<&LLMKnowledgeEntity> for KnowledgeEntity { - fn from(llm_entity: &LLMKnowledgeEntity) -> Self { - KnowledgeEntity { - id: Uuid::new_v4(), - name: llm_entity.name.clone(), - description: llm_entity.description.clone(), - entity_type: KnowledgeEntityType::from(llm_entity.entity_type.clone()), - source_id: None, // To be linked externally if needed - metadata: None, // Populate if metadata is provided - } - } -} - -impl From<&LLMRelationship> for KnowledgeRelationship { - fn from(llm_rel: &LLMRelationship) -> Self { - KnowledgeRelationship { - id: Uuid::new_v4(), - in_: Uuid::nil(), // Placeholder; to be set after mapping - out: Uuid::nil(), // Placeholder; to be set after mapping - relationship_type: llm_rel.type_.clone(), - metadata: None, // Populate if metadata is provided - } - } -} diff --git a/src/models/text_content.rs b/src/models/text_content.rs index 4cfabc9..fe98544 100644 --- a/src/models/text_content.rs +++ b/src/models/text_content.rs @@ -1,7 +1,8 @@ use serde::{Deserialize, Serialize}; -use tracing::info; +use surrealdb::RecordId; +use tracing::{debug, info}; use uuid::Uuid; -use crate::{models::file_info::FileInfo, utils::llm::create_json_ld}; +use crate::{models::file_info::FileInfo, surrealdb::{SurrealDbClient, SurrealError}, utils::llm::create_json_ld}; use thiserror::Error; use super::graph_entities::{KnowledgeEntity, KnowledgeRelationship}; @@ -16,11 +17,23 @@ pub struct TextContent { pub category: String, } +#[derive(Debug,Deserialize)] +struct Record { + #[allow(dead_code)] + id: RecordId, +} + /// Error types for processing `TextContent`. #[derive(Error, Debug)] pub enum ProcessingError { #[error("LLM processing error: {0}")] LLMError(String), + + #[error("SurrealDB error: {0}")] + SurrealError(#[from] SurrealError), + + #[error("SurrealDb error: {0}")] + SurrealDbError(#[from] surrealdb::Error), #[error("Graph DB storage error: {0}")] GraphDBError(String), @@ -36,15 +49,15 @@ pub enum ProcessingError { impl TextContent { /// Processes the `TextContent` by sending it to an LLM, storing in a graph DB, and vector DB. pub async fn process(&self) -> Result<(), ProcessingError> { - // let client = Neo4jClient::new("127.0.0.1:7687", "neo4j", "neo4j").await.expect("Failed to create Neo4j client"); - + // Store TextContent + // Step 1: Send to LLM for analysis let analysis = create_json_ld(&self.category, &self.instructions, &self.text).await?; // info!("{:#?}", &analysis); // Step 2: Convert LLM analysis to database entities - let (entities, relationships) = analysis.to_database_entities(); + let (entities, relationships) = analysis.to_database_entities(&self.id); // Step 3: Store in database self.store_in_graph_db(entities, relationships).await?; @@ -61,8 +74,18 @@ impl TextContent { entities: Vec, relationships: Vec ) -> Result<(), ProcessingError> { - for entity in entities { + let db_client = SurrealDbClient::new().await?; + + for entity in entities { info!("{:?}", entity); + + let _created: Option = db_client + .client + .create(("knowledge_entity", &entity.id.to_string())) + .content(entity) + .await?; + + info!("{:?}",_created); } for relationship in relationships { diff --git a/src/utils/llm.rs b/src/utils/llm.rs index cacd9f4..2ae3577 100644 --- a/src/utils/llm.rs +++ b/src/utils/llm.rs @@ -38,7 +38,7 @@ pub struct LLMGraphAnalysisResult { } impl LLMGraphAnalysisResult { - pub fn to_database_entities(&self) -> (Vec, Vec) { + pub fn to_database_entities(&self, source_id: &Uuid) -> (Vec, Vec) { let mut mapper = GraphMapper::new(); // First pass: Create all entities and map their keys to UUIDs @@ -47,11 +47,11 @@ impl LLMGraphAnalysisResult { .map(|llm_entity| { let id = mapper.assign_id(&llm_entity.key); KnowledgeEntity { - id, + id: id.to_string(), name: llm_entity.name.clone(), description: llm_entity.description.clone(), entity_type: KnowledgeEntityType::from(llm_entity.entity_type.clone()), - source_id: None, + source_id: source_id.to_string(), metadata: None, } })