From 40cf2e3b5bb853c65286851718ab2a0419719f6d Mon Sep 17 00:00:00 2001
From: Per Stark <perstark.se@gmail.com>
Date: Mon, 4 Nov 2024 20:58:08 +0100
Subject: [PATCH] storing relationships

---
 src/models/graph_entities.rs |  8 ++++----
 src/models/text_content.rs   | 21 +++++++++++++++------
 src/utils/llm.rs             | 17 +++++++++++++----
 3 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/src/models/graph_entities.rs b/src/models/graph_entities.rs
index 04fe74c..b665776 100644
--- a/src/models/graph_entities.rs
+++ b/src/models/graph_entities.rs
@@ -40,11 +40,11 @@ impl From<String> for KnowledgeEntityType {
 /// Represents a relationship between two knowledge entities.
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct KnowledgeRelationship {
-    pub id: Uuid, // Generated in Rust
+    pub id: String,
     #[serde(rename = "in")]
-    pub in_: Uuid, // Target KnowledgeEntity ID
-    pub out: Uuid, // Source KnowledgeEntity ID
-    pub relationship_type: String, // e.g., RelatedTo, RelevantTo
+    pub in_: String, // Target KnowledgeEntity ID
+    pub out: String,                         // Source KnowledgeEntity ID
+    pub relationship_type: String,           // e.g., RelatedTo, RelevantTo
     pub metadata: Option<serde_json::Value>, // Additional metadata
 }
 
diff --git a/src/models/text_content.rs b/src/models/text_content.rs
index fe98544..ba26132 100644
--- a/src/models/text_content.rs
+++ b/src/models/text_content.rs
@@ -50,9 +50,10 @@ 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> {
         // Store TextContent
+        let db_client = SurrealDbClient::new().await?;
         
         // Step 1: Send to LLM for analysis
-        let analysis = create_json_ld(&self.category, &self.instructions, &self.text).await?;
+        let analysis = create_json_ld(&self.category, &self.instructions, &self.text, &db_client).await?;
         // info!("{:#?}", &analysis);
 
 
@@ -60,7 +61,7 @@ impl TextContent {
         let (entities, relationships) = analysis.to_database_entities(&self.id);
         
         // Step 3: Store in database
-        self.store_in_graph_db(entities, relationships).await?;
+        self.store_in_graph_db(entities, relationships, &db_client).await?;
         
 
         // Step 4: Split text and store in Vector DB
@@ -72,10 +73,9 @@ impl TextContent {
     async fn store_in_graph_db(
         &self,
         entities: Vec<KnowledgeEntity>,
-        relationships: Vec<KnowledgeRelationship>
+        relationships: Vec<KnowledgeRelationship>,
+        db_client: &SurrealDbClient,
     ) -> Result<(), ProcessingError> {
-        let db_client = SurrealDbClient::new().await?;
-
          for entity in entities {
             info!("{:?}", entity);
             
@@ -85,11 +85,20 @@ impl TextContent {
                 .content(entity)
                 .await?;
 
-            info!("{:?}",_created);
+            debug!("{:?}",_created);
         }
 
         for relationship in relationships {
             info!("{:?}", relationship);
+
+            let _created: Option<Record> = db_client
+                .client
+                .insert(("knowledge_relationship", &relationship.id.to_string()))
+                .content(relationship)
+                .await?;
+
+            debug!("{:?}",_created);
+        
         }
 
         Ok(())
diff --git a/src/utils/llm.rs b/src/utils/llm.rs
index 2ae3577..d8b6140 100644
--- a/src/utils/llm.rs
+++ b/src/utils/llm.rs
@@ -4,12 +4,14 @@ use async_openai::types::CreateChatCompletionRequestArgs;
 use serde::Deserialize;
 use serde::Serialize;
 use tracing::debug;
+use tracing::info;
 use uuid::Uuid;
 use crate::models::graph_entities::GraphMapper;
 use crate::models::graph_entities::KnowledgeEntity;
 use crate::models::graph_entities::KnowledgeEntityType;
 use crate::models::graph_entities::KnowledgeRelationship;
 use crate::models::text_content::ProcessingError;
+use crate::surrealdb::SurrealDbClient;
 use serde_json::json;
 
 /// Represents a single knowledge entity from the LLM.
@@ -65,9 +67,9 @@ impl LLMGraphAnalysisResult {
                 let target_id = mapper.get_id(&llm_rel.target)?;
                 
                 Some(KnowledgeRelationship {
-                    id: Uuid::new_v4(),
-                    out: *source_id,
-                    in_: *target_id,
+                    id: Uuid::new_v4().to_string(),
+                    out: source_id.to_string(),
+                    in_: target_id.to_string(),
                     relationship_type: llm_rel.type_.clone(),
                     metadata: None,
                 })
@@ -79,7 +81,14 @@ impl LLMGraphAnalysisResult {
 }
 
 /// Sends text to an LLM for analysis.
-pub async fn create_json_ld(category: &str, instructions: &str, text: &str) -> Result<LLMGraphAnalysisResult, ProcessingError> {
+pub async fn create_json_ld(category: &str, instructions: &str, text: &str, db_client: &SurrealDbClient) -> Result<LLMGraphAnalysisResult, ProcessingError> {
+    // Get the nodes from the database
+    let mut result = db_client.client.query("SELECT * FROM knowledge_entity").await?;
+    info!("{:?}", result.num_statements());
+
+    let db_representation: Vec<KnowledgeEntity> = result.take(1)?;
+    info!("{:?}", db_representation);
+    
         let client = async_openai::Client::new();
         let schema = json!({
           "type": "object",