diff --git a/src/models/graph_entities.rs b/src/models/graph_entities.rs index 77e159f..866b277 100644 --- a/src/models/graph_entities.rs +++ b/src/models/graph_entities.rs @@ -65,6 +65,12 @@ pub struct GraphMapper { pub key_to_id: HashMap, } +impl Default for GraphMapper { + fn default() -> Self { + GraphMapper::new() + } +} + impl GraphMapper { pub fn new() -> Self { GraphMapper { diff --git a/src/models/text_content.rs b/src/models/text_content.rs index 6bd2772..5d02db8 100644 --- a/src/models/text_content.rs +++ b/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, relationships: Vec, - db_client: &SurrealDbClient, + db_client: &Surreal, ) -> Result<(), ProcessingError> { for entity in entities { info!("{:?}", entity); let _created: Option = 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 = db_client - .client .insert(("knowledge_relationship", &relationship.id.to_string())) .content(relationship) .await?; diff --git a/src/surrealdb/mod.rs b/src/surrealdb/mod.rs index 5f47b5f..6c7eca1 100644 --- a/src/surrealdb/mod.rs +++ b/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::("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; + + fn deref(&self) -> &Self::Target { + &self.client + } +} diff --git a/src/utils/llm.rs b/src/utils/llm.rs index f0a6617..c98a51e 100644 --- a/src/utils/llm.rs +++ b/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, ) -> Result { // Get the nodes from the database - let entities: Vec = db_client.client.select("knowledge_entity").await?; + let entities: Vec = db_client.select("knowledge_entity").await?; info!("{:?}", entities); - let deleted: Vec = db_client.client.delete("knowledge_entity").await?; - info! {"{:?}", deleted}; + let deleted: Vec = db_client.delete("knowledge_entity").await?; + info! {"{:?} KnowledgeEntities deleted", deleted.len()}; + + let relationships: Vec = db_client.select("knowledge_relationship").await?; + info!("{:?}", relationships); + + let relationships_deleted: Vec = db_client.delete("knowledge_relationship").await?; + info!("{:?} Relationships deleted", relationships_deleted.len()); let client = async_openai::Client::new(); let schema = json!({