mixed context, entities, community reports

Author: benolt

Cargo.lock

@@ -7,12 +7,13 @@ edition = "2021"
 anyhow = "1.0.86"
 async-trait = "0.1.74"
 futures = "0.3.30"
-polars = { version = "0.41.3", features = ["dtype-struct", "lazy", "parquet"] }
+polars = { version = "0.41.3", features = ["dtype-struct", "is_in", "lazy", "parquet"] }
 polars-lazy = "0.41.3"
 rand = "0.8.5"
 serde = { version = "1.0.188", features = ["derive"] }
 serde_json = "1.0.117"
 tokio = { version = "1.36", features = ["full"] }
+uuid = { version = "1.6.1", features = ["v4"] }
 
 [dev-dependencies]
 async-openai = "0.23.4"

shinkai-libs/shinkai-graphrag/Cargo.toml

@@ -13,7 +13,23 @@ use rand::prelude::SliceRandom;
 
 use crate::models::{CommunityReport, Entity};
 
-use super::context_builder::GlobalSearchContextBuilderParams;
+#[derive(Debug, Clone)]
+pub struct CommunityContextBuilderParams {
+    pub use_community_summary: bool,
+    pub column_delimiter: String,
+    pub shuffle_data: bool,
+    pub include_community_rank: bool,
+    pub min_community_rank: u32,
+    pub community_rank_name: String,
+    pub include_community_weight: bool,
+    pub community_weight_name: String,
+    pub normalize_community_weight: bool,
+    pub max_tokens: usize,
+    pub context_name: String,
+    // conversation_history: Option<ConversationHistory>,
+    // conversation_history_user_turns_only: bool,
+    // conversation_history_max_turns: Option<i32>,
+}
 
 pub struct GlobalCommunityContext {
     community_reports: Vec<CommunityReport>,
@@ -34,11 +50,11 @@ impl GlobalCommunityContext {
         }
     }
 
-    pub async fn build_context(
+    pub fn build_context(
         &self,
-        context_builder_params: GlobalSearchContextBuilderParams,
+        context_builder_params: CommunityContextBuilderParams,
     ) -> anyhow::Result<(Vec<String>, HashMap<String, DataFrame>)> {
-        let GlobalSearchContextBuilderParams {
+        let CommunityContextBuilderParams {
             use_community_summary,
             column_delimiter,
             shuffle_data,

shinkai-libs/shinkai-graphrag/src/context_builder/community_context.rs

@@ -1,19 +1 @@
-#[derive(Debug, Clone)]
-pub struct GlobalSearchContextBuilderParams {
-    pub use_community_summary: bool,
-    pub column_delimiter: String,
-    pub shuffle_data: bool,
-    pub include_community_rank: bool,
-    pub min_community_rank: u32,
-    pub community_rank_name: String,
-    pub include_community_weight: bool,
-    pub community_weight_name: String,
-    pub normalize_community_weight: bool,
-    pub max_tokens: usize,
-    pub context_name: String,
-    //conversation_history: Option<ConversationHistory>,
-    // conversation_history_user_turns_only: bool,
-    // conversation_history_max_turns: Option<i32>,
-}
-
 pub struct ConversationHistory {}

shinkai-libs/shinkai-graphrag/src/context_builder/context_builder.rs

@@ -1,4 +1,2 @@
 pub mod community_context;
 pub mod context_builder;
-pub mod indexer_entities;
-pub mod indexer_reports;

shinkai-libs/shinkai-graphrag/src/context_builder/mod.rs

@@ -0,0 +1,2 @@
+pub mod indexer_entities;
+pub mod indexer_reports;

…/src/context_builder/indexer_entities.rs …src/indexer_adapters/indexer_entities.rsshinkai-libs/shinkai-graphrag/src/context_builder/indexer_entities.rs renamed to shinkai-libs/shinkai-graphrag/src/indexer_adapters/indexer_entities.rs shinkai-libs/shinkai-graphrag/src/context_builder/indexer_entities.rs renamed to shinkai-libs/shinkai-graphrag/src/indexer_adapters/indexer_entities.rs

@@ -1,5 +1,7 @@
 pub mod context_builder;
+pub mod indexer_adapters;
 pub mod llm;
 pub mod models;
+pub mod retrieval;
 pub mod search;
 pub mod vector_stores;

…g/src/context_builder/indexer_reports.rs …/src/indexer_adapters/indexer_reports.rsshinkai-libs/shinkai-graphrag/src/context_builder/indexer_reports.rs renamed to shinkai-libs/shinkai-graphrag/src/indexer_adapters/indexer_reports.rs shinkai-libs/shinkai-graphrag/src/context_builder/indexer_reports.rs renamed to shinkai-libs/shinkai-graphrag/src/indexer_adapters/indexer_reports.rs

@@ -47,6 +47,7 @@ pub trait BaseLLM {
 #[async_trait]
 pub trait BaseTextEmbedding {
     async fn aembed(&self, text: &str) -> Vec<f64>;
+    fn embed(&self, text: &str) -> Vec<f64>;
 }
 
 pub enum GlobalSearchPhase {