From 1212e330989db957bca6bc18d1a9f2f467010985 Mon Sep 17 00:00:00 2001
From: Per Stark <perstark.se@gmail.com>
Date: Mon, 18 Nov 2024 12:37:04 +0100
Subject: [PATCH] chore: cleaning and clarifying

---
 src/models/text_content.rs | 20 ++++++++++----------
 src/utils/llm.rs           | 13 +++++++++----
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/src/models/text_content.rs b/src/models/text_content.rs
index 5cfad4e..8b03cf6 100644
--- a/src/models/text_content.rs
+++ b/src/models/text_content.rs
@@ -108,17 +108,17 @@ impl TextContent {
             debug!("{:?}",_created);
         }
 
-        for relationship in &relationships {
-            let in_entity: Option<KnowledgeEntity> = db_client.select(("knowledge_entity",relationship.in_.to_string())).await?;
-            let out_entity: Option<KnowledgeEntity> = db_client.select(("knowledge_entity", relationship.out.to_string())).await?;
+        // for relationship in &relationships {
+        //     let in_entity: Option<KnowledgeEntity> = db_client.select(("knowledge_entity",relationship.in_.to_string())).await?;
+        //     let out_entity: Option<KnowledgeEntity> = db_client.select(("knowledge_entity", relationship.out.to_string())).await?;
             
-            if let (Some(in_), Some(out)) = (in_entity, out_entity) {
-            info!("{} - {} is {} to {} - {}", in_.id, in_.name, relationship.relationship_type, out.id, out.name);
-            }
-            else {
-                info!("No in or out entities found");
-            }
-        }
+        //     if let (Some(in_), Some(out)) = (in_entity, out_entity) {
+        //     info!("{} - {} is {} to {} - {}", in_.id, in_.name, relationship.relationship_type, out.id, out.name);
+        //     }
+        //     else {
+        //         info!("No in or out entities found");
+        //     }
+        // }
 
         info!("Inserted to database: {:?} entities, {:?} relationships", entities.len(), relationships.len());
 
diff --git a/src/utils/llm.rs b/src/utils/llm.rs
index 3f8d5c7..dd92883 100644
--- a/src/utils/llm.rs
+++ b/src/utils/llm.rs
@@ -10,7 +10,7 @@ use serde::{Deserialize, Serialize};
 use serde_json::json;
 use surrealdb::engine::remote::ws::Client;
 use surrealdb::Surreal;
-use tracing::debug;
+use tracing::{debug, info};
 use uuid::Uuid;
 
 /// Represents a single knowledge entity from the LLM.
@@ -164,8 +164,10 @@ pub async fn create_json_ld(
     // Generate embedding of the input
     let input_embedding = generate_embedding(&client, input_text).await?;
 
+    let number_of_entities_to_get = 10;
+    
     // Construct the query
-    let closest_query = format!("SELECT *, vector::distance::knn() AS distance FROM knowledge_entity WHERE embedding <|3,40|> {:?} ORDER BY distance", input_embedding);
+    let closest_query = format!("SELECT *, vector::distance::knn() AS distance FROM knowledge_entity WHERE embedding <|{},40|> {:?} ORDER BY distance",number_of_entities_to_get, input_embedding);
 
     // Perform query and deserialize to struct    
     let closest_entities: Vec<KnowledgeEntity> = db_client.query(closest_query).await?.take(0)?;
@@ -176,6 +178,8 @@ pub async fn create_json_ld(
         name: String,
         description: String
     }
+
+    info!("Number of KnowledgeEntities sent as context: {}", closest_entities.len());
     
     // Only keep most relevant information
     let closest_entities_to_llm: Vec<KnowledgeEntityToLLM> = closest_entities.clone().into_iter().map(|entity| KnowledgeEntityToLLM {
@@ -239,7 +243,7 @@ pub async fn create_json_ld(
 
     // Construct the system and user messages
     let system_message = r#"
-            You are an expert document analyzer. You will receive a document's text content, along with user instructions and a category. Your task is to provide a structured JSON object representing the content in a graph format suitable for a graph database. You will also be presented with some existing knowledge_entities, do not replicate these!
+            You are an expert document analyzer. You will receive a document's text content, along with user instructions and a category. Your task is to provide a structured JSON object representing the content in a graph format suitable for a graph database. You will also be presented with some existing knowledge_entities from the database, do not replicate these!
             
             The JSON should have the following structure:
             
@@ -271,10 +275,11 @@ pub async fn create_json_ld(
             5. Use the `source` key to indicate the originating entity and the `target` key to indicate the related entity"
             6. You will be presented with a few existing KnowledgeEntities that are similar to the current ones. They will have an existing UUID. When creating relationships to these entities, use their UUID.
             7. Only create relationships between existing KnowledgeEntities.
+            8. Entities that exist already in the database should NOT be created again. If there is only a minor overlap, skip creating a new entity.
             "#;
 
     let user_message = format!(
-        "Category: {}\nInstructions: {}\nContent:\n{}\nExisting KnowledgeEntities:{:?}",
+        "Category: {}\nInstructions: {}\nContent:\n{}\nExisting KnowledgeEntities in database:{:?}",
         category, instructions, text, closest_entities_to_llm
     );