From 6d8cd05c1a28c29ca96f4b66703b88d6bc09674f Mon Sep 17 00:00:00 2001
From: Per Stark <perstark.se@gmail.com>
Date: Tue, 5 Nov 2024 22:00:04 +0100
Subject: [PATCH] working entity storage

---
 src/models/graph_entities.rs | 17 ++++++++++++++---
 src/models/text_content.rs   | 23 ++++++++---------------
 src/utils/llm.rs             | 12 ++++++++----
 3 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/src/models/graph_entities.rs b/src/models/graph_entities.rs
index b665776..f450a21 100644
--- a/src/models/graph_entities.rs
+++ b/src/models/graph_entities.rs
@@ -1,17 +1,28 @@
 use serde::Deserialize;
+use serde::Deserializer;
 use serde::Serialize;
 use std::collections::HashMap;
+use surrealdb::sql::Thing;
 use uuid::Uuid;
 
 /// Represents a generic knowledge entity in the graph.
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct KnowledgeEntity {
-    pub id: String, // Generated in Rust
+    #[serde(deserialize_with = "thing_to_string")]
+    pub id: String,
     pub name: String,
     pub description: String,
     pub entity_type: KnowledgeEntityType,
-    pub source_id: String,                   // Links to FileInfo or TextContent
-    pub metadata: Option<serde_json::Value>, // Additional metadata
+    pub source_id: String,
+    pub metadata: Option<serde_json::Value>,
+}
+
+fn thing_to_string<'de, D>(deserializer: D) -> Result<String, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    let thing = Thing::deserialize(deserializer)?;
+    Ok(thing.id.to_raw())
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
diff --git a/src/models/text_content.rs b/src/models/text_content.rs
index ba26132..df10669 100644
--- a/src/models/text_content.rs
+++ b/src/models/text_content.rs
@@ -1,5 +1,4 @@
 use serde::{Deserialize, Serialize};
-use surrealdb::RecordId;
 use tracing::{debug, info};
 use uuid::Uuid;
 use crate::{models::file_info::FileInfo, surrealdb::{SurrealDbClient, SurrealError}, utils::llm::create_json_ld};
@@ -17,12 +16,6 @@ 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 {
@@ -79,7 +72,7 @@ impl TextContent {
          for entity in entities {
             info!("{:?}", entity);
             
-            let _created: Option<Record> = db_client
+            let _created: Option<KnowledgeEntity> = db_client
                 .client
                 .create(("knowledge_entity", &entity.id.to_string()))
                 .content(entity)
@@ -89,15 +82,15 @@ impl TextContent {
         }
 
         for relationship in relationships {
-            info!("{:?}", relationship);
+            // info!("{:?}", relationship);
 
-            let _created: Option<Record> = db_client
-                .client
-                .insert(("knowledge_relationship", &relationship.id.to_string()))
-                .content(relationship)
-                .await?;
+            // let _created: Option<Record> = db_client
+            //     .client
+            //     .insert(("knowledge_relationship", &relationship.id.to_string()))
+            //     .content(relationship)
+            //     .await?;
 
-            debug!("{:?}",_created);
+            // debug!("{:?}",_created);
         
         }
 
diff --git a/src/utils/llm.rs b/src/utils/llm.rs
index d8b6140..0c8491b 100644
--- a/src/utils/llm.rs
+++ b/src/utils/llm.rs
@@ -3,6 +3,8 @@ use async_openai::types::ChatCompletionRequestUserMessage;
 use async_openai::types::CreateChatCompletionRequestArgs;
 use serde::Deserialize;
 use serde::Serialize;
+use surrealdb::sql::Thing;
+use surrealdb::RecordId;
 use tracing::debug;
 use tracing::info;
 use uuid::Uuid;
@@ -83,11 +85,13 @@ impl LLMGraphAnalysisResult {
 /// Sends text to an LLM for analysis.
 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 entities: Vec<KnowledgeEntity> = db_client.client.select("knowledge_entity").await?;
+    info!("{:?}", entities);
 
-    let db_representation: Vec<KnowledgeEntity> = result.take(1)?;
-    info!("{:?}", db_representation);
+    let deleted: Vec<KnowledgeEntity> = db_client.client.delete("knowledge_entity").await?;
+    info!{"{:?}", deleted};
+
+    
     
         let client = async_openai::Client::new();
         let schema = json!({