chore: cleaning & deref

src/models/graph_entities.rs

@@ -65,6 +65,12 @@ pub struct GraphMapper {
     pub key_to_id: HashMap<String, Uuid>,
 }
 
+impl Default for GraphMapper {
+    fn default() -> Self {
+        GraphMapper::new()
+    }
+}
+
 impl GraphMapper {
     pub fn new() -> Self {
         GraphMapper {

src/models/text_content.rs

@@ -1,4 +1,5 @@
 use serde::{Deserialize, Serialize};
+use surrealdb::{engine::remote::ws::Client, Surreal};
 use tracing::{debug, info};
 use uuid::Uuid;
 use crate::{models::file_info::FileInfo, surrealdb::{SurrealDbClient, SurrealError}, utils::llm::create_json_ld};
@@ -67,13 +68,12 @@ impl TextContent {
         &self,
         entities: Vec<KnowledgeEntity>,
         relationships: Vec<KnowledgeRelationship>,
-        db_client: &SurrealDbClient,
+        db_client: &Surreal<Client>,
     ) -> Result<(), ProcessingError> {
          for entity in entities {
             info!("{:?}", entity);
             
             let _created: Option<KnowledgeEntity> = db_client
-                .client
                 .create(("knowledge_entity", &entity.id.to_string()))
                 .content(entity)
                 .await?;
@@ -82,10 +82,9 @@ impl TextContent {
         }
 
         for relationship in relationships {
-            // info!("{:?}", relationship);
+            info!("{:?}", relationship);
 
             let _created: Option<KnowledgeRelationship> = db_client
-                .client
                 .insert(("knowledge_relationship", &relationship.id.to_string()))
                 .content(relationship)
                 .await?;

src/surrealdb/mod.rs

@@ -1,4 +1,9 @@
-use surrealdb::{engine::remote::ws::{Client, Ws}, opt::auth::Root, Surreal};
+use std::ops::Deref;
+use surrealdb::{
+    engine::remote::ws::{Client, Ws},
+    opt::auth::Root,
+    Surreal,
+};
 use thiserror::Error;
 
 #[derive(Clone)]
@@ -10,11 +15,9 @@ pub struct SurrealDbClient {
 pub enum SurrealError {
     #[error("SurrealDb error: {0}")]
     SurrealDbError(#[from] surrealdb::Error),
-
     // Add more error variants as needed.
 }
 
-
 impl SurrealDbClient {
     /// # Initialize a new datbase client
     ///
@@ -26,7 +29,11 @@ impl SurrealDbClient {
         let db = Surreal::new::<Ws>("127.0.0.1:8000").await?;
 
         // Sign in to database
-        db.signin(Root{username: "root_user", password: "root_password"}).await?;
+        db.signin(Root {
+            username: "root_user",
+            password: "root_password",
+        })
+        .await?;
 
         // Set namespace
         db.use_ns("test").use_db("test").await?;
@@ -34,3 +41,11 @@ impl SurrealDbClient {
         Ok(SurrealDbClient { client: db })
     }
 }
+
+impl Deref for SurrealDbClient {
+    type Target = Surreal<Client>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.client
+    }
+}

src/utils/llm.rs

@@ -1,12 +1,10 @@
-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::graph_entities::{GraphMapper, KnowledgeEntity, KnowledgeEntityType, KnowledgeRelationship};
 use crate::models::text_content::ProcessingError;
-use crate::surrealdb::SurrealDbClient;
 use async_openai::types::{CreateChatCompletionRequestArgs, ChatCompletionRequestUserMessage, ChatCompletionRequestSystemMessage };
 use serde::{Deserialize, Serialize};
 use serde_json::json;
+use surrealdb::engine::remote::ws::Client;
+use surrealdb::Surreal;
 use tracing::{info,debug};
 use uuid::Uuid;
 
@@ -86,14 +84,20 @@ pub async fn create_json_ld(
     category: &str,
     instructions: &str,
     text: &str,
-    db_client: &SurrealDbClient,
+    db_client: &Surreal<Client>,
 ) -> Result<LLMGraphAnalysisResult, ProcessingError> {
     // Get the nodes from the database
-    let entities: Vec<KnowledgeEntity> = db_client.client.select("knowledge_entity").await?;
+    let entities: Vec<KnowledgeEntity> = db_client.select("knowledge_entity").await?;
     info!("{:?}", entities);
 
-    let deleted: Vec<KnowledgeEntity> = db_client.client.delete("knowledge_entity").await?;
-    info! {"{:?}", deleted};
+    let deleted: Vec<KnowledgeEntity> = db_client.delete("knowledge_entity").await?;
+    info! {"{:?} KnowledgeEntities deleted", deleted.len()};
+    
+    let relationships: Vec<KnowledgeRelationship> = db_client.select("knowledge_relationship").await?;
+    info!("{:?}", relationships);
+    
+    let relationships_deleted: Vec<KnowledgeRelationship> = db_client.delete("knowledge_relationship").await?;
+    info!("{:?} Relationships deleted", relationships_deleted.len());
 
     let client = async_openai::Client::new();
     let schema = json!({