GraphRAG Local Search

Author: benolt

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

@@ -11,7 +11,9 @@ use polars::{
 };
 use rand::prelude::SliceRandom;
 
-use super::{context_builder::ContextBuilderParams, indexer_entities::Entity, indexer_reports::CommunityReport};
+use crate::models::{CommunityReport, Entity};
+
+use super::context_builder::GlobalSearchContextBuilderParams;
 
 pub struct GlobalCommunityContext {
     community_reports: Vec<CommunityReport>,
@@ -34,9 +36,9 @@ impl GlobalCommunityContext {
 
     pub async fn build_context(
         &self,
-        context_builder_params: ContextBuilderParams,
+        context_builder_params: GlobalSearchContextBuilderParams,
     ) -> anyhow::Result<(Vec<String>, HashMap<String, DataFrame>)> {
-        let ContextBuilderParams {
+        let GlobalSearchContextBuilderParams {
             use_community_summary,
             column_delimiter,
             shuffle_data,

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

@@ -1,6 +1,5 @@
 #[derive(Debug, Clone)]
-pub struct ContextBuilderParams {
-    //conversation_history: Option<ConversationHistory>,
+pub struct GlobalSearchContextBuilderParams {
     pub use_community_summary: bool,
     pub column_delimiter: String,
     pub shuffle_data: bool,
@@ -12,6 +11,7 @@ pub struct ContextBuilderParams {
     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>,
 }

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

@@ -1,8 +1,9 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::HashSet;
 
 use polars::prelude::*;
 use polars_lazy::dsl::col;
-use serde::{Deserialize, Serialize};
+
+use crate::models::Entity;
 
 use super::indexer_reports::filter_under_community_level;
 
@@ -52,23 +53,6 @@ pub fn read_indexer_entities(
     Ok(entities)
 }
 
-#[derive(Debug, Clone, Deserialize, Serialize)]
-pub struct Entity {
-    pub id: String,
-    pub short_id: Option<String>,
-    pub title: String,
-    pub entity_type: Option<String>,
-    pub description: Option<String>,
-    pub description_embedding: Option<Vec<f64>>,
-    pub name_embedding: Option<Vec<f64>>,
-    pub graph_embedding: Option<Vec<f64>>,
-    pub community_ids: Option<Vec<String>>,
-    pub text_unit_ids: Option<Vec<String>>,
-    pub document_ids: Option<Vec<String>>,
-    pub rank: Option<i32>,
-    pub attributes: Option<HashMap<String, String>>,
-}
-
 pub fn read_entities(
     df: DataFrame,
     id_col: &str,

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

@@ -1,8 +1,9 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::HashSet;
 
 use polars::prelude::*;
 use polars_lazy::dsl::col;
-use serde::{Deserialize, Serialize};
+
+use crate::models::CommunityReport;
 
 use super::indexer_entities::get_field;
 
@@ -59,20 +60,6 @@ pub fn filter_under_community_level(df: &DataFrame, community_level: u32) -> any
     Ok(result)
 }
 
-#[derive(Debug, Clone, Deserialize, Serialize)]
-pub struct CommunityReport {
-    pub id: String,
-    pub short_id: Option<String>,
-    pub title: String,
-    pub community_id: String,
-    pub summary: String,
-    pub full_content: String,
-    pub rank: Option<f64>,
-    pub summary_embedding: Option<Vec<f64>>,
-    pub full_content_embedding: Option<Vec<f64>>,
-    pub attributes: Option<HashMap<String, String>>,
-}
-
 pub fn read_community_reports(
     df: DataFrame,
     id_col: &str,

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

@@ -1,3 +1,5 @@
 pub mod context_builder;
 pub mod llm;
+pub mod models;
 pub mod search;
+pub mod vector_stores;

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

@@ -0,0 +1,59 @@
+use std::collections::HashMap;
+
+#[derive(Debug, Clone)]
+pub struct CommunityReport {
+    pub id: String,
+    pub short_id: Option<String>,
+    pub title: String,
+    pub community_id: String,
+    pub summary: String,
+    pub full_content: String,
+    pub rank: Option<f64>,
+    pub summary_embedding: Option<Vec<f64>>,
+    pub full_content_embedding: Option<Vec<f64>>,
+    pub attributes: Option<HashMap<String, String>>,
+}
+
+#[derive(Debug, Clone)]
+pub struct Entity {
+    pub id: String,
+    pub short_id: Option<String>,
+    pub title: String,
+    pub entity_type: Option<String>,
+    pub description: Option<String>,
+    pub description_embedding: Option<Vec<f64>>,
+    pub name_embedding: Option<Vec<f64>>,
+    pub graph_embedding: Option<Vec<f64>>,
+    pub community_ids: Option<Vec<String>>,
+    pub text_unit_ids: Option<Vec<String>>,
+    pub document_ids: Option<Vec<String>>,
+    pub rank: Option<i32>,
+    pub attributes: Option<HashMap<String, String>>,
+}
+
+#[derive(Debug, Clone)]
+pub struct Relationship {
+    pub id: String,
+    pub short_id: Option<String>,
+    pub source: String,
+    pub target: String,
+    pub weight: Option<f64>,
+    pub description: Option<String>,
+    pub description_embedding: Option<Vec<f64>>,
+    pub text_unit_ids: Option<Vec<String>>,
+    pub document_ids: Option<Vec<String>>,
+    pub attributes: Option<HashMap<String, String>>,
+}
+
+#[derive(Debug, Clone)]
+pub struct TextUnit {
+    pub id: String,
+    pub short_id: Option<String>,
+    pub text: String,
+    pub text_embedding: Option<Vec<f64>>,
+    pub entity_ids: Option<Vec<String>>,
+    pub relationship_ids: Option<Vec<String>>,
+    pub n_tokens: Option<i32>,
+    pub document_ids: Option<Vec<String>>,
+    pub attributes: Option<HashMap<String, String>>,
+}

shinkai-libs/shinkai-graphrag/src/search/base.rs

@@ -0,0 +1,29 @@
+use std::collections::HashMap;
+
+use polars::frame::DataFrame;
+
+#[derive(Debug, Clone)]
+pub enum ResponseType {
+    String(String),
+    KeyPoints(Vec<KeyPoint>),
+}
+
+#[derive(Debug, Clone)]
+pub enum ContextData {
+    String(String),
+    DataFrames(Vec<DataFrame>),
+    Dictionary(HashMap<String, DataFrame>),
+}
+
+#[derive(Debug, Clone)]
+pub enum ContextText {
+    String(String),
+    Strings(Vec<String>),
+    Dictionary(HashMap<String, String>),
+}
+
+#[derive(Debug, Clone)]
+pub struct KeyPoint {
+    pub answer: String,
+    pub score: i32,
+}

shinkai-libs/shinkai-graphrag/src/search/global_search/global_search.rs

@@ -1,12 +1,12 @@
 use futures::future::join_all;
-use polars::frame::DataFrame;
 use serde_json::Value;
 use std::collections::HashMap;
 use std::time::Instant;
 
 use crate::context_builder::community_context::GlobalCommunityContext;
-use crate::context_builder::context_builder::{ContextBuilderParams, ConversationHistory};
+use crate::context_builder::context_builder::{ConversationHistory, GlobalSearchContextBuilderParams};
 use crate::llm::llm::{BaseLLM, BaseLLMCallback, LLMParams, MessageType};
+use crate::search::base::{ContextData, ContextText, KeyPoint, ResponseType};
 use crate::search::global_search::prompts::NO_DATA_ANSWER;
 
 use super::prompts::{GENERAL_KNOWLEDGE_INSTRUCTION, MAP_SYSTEM_PROMPT, REDUCE_SYSTEM_PROMPT};
@@ -21,32 +21,6 @@ pub struct SearchResult {
     pub prompt_tokens: usize,
 }
 
-#[derive(Debug, Clone)]
-pub enum ResponseType {
-    String(String),
-    KeyPoints(Vec<KeyPoint>),
-}
-
-#[derive(Debug, Clone)]
-pub enum ContextData {
-    String(String),
-    DataFrames(Vec<DataFrame>),
-    Dictionary(HashMap<String, DataFrame>),
-}
-
-#[derive(Debug, Clone)]
-pub enum ContextText {
-    String(String),
-    Strings(Vec<String>),
-    Dictionary(HashMap<String, String>),
-}
-
-#[derive(Debug, Clone)]
-pub struct KeyPoint {
-    pub answer: String,
-    pub score: i32,
-}
-
 pub struct GlobalSearchResult {
     pub response: ResponseType,
     pub context_data: ContextData,
@@ -88,7 +62,7 @@ pub struct GlobalSearch {
     llm: Box<dyn BaseLLM>,
     context_builder: GlobalCommunityContext,
     num_tokens_fn: fn(&str) -> usize,
-    context_builder_params: ContextBuilderParams,
+    context_builder_params: GlobalSearchContextBuilderParams,
     map_system_prompt: String,
     reduce_system_prompt: String,
     response_type: String,
@@ -114,7 +88,7 @@ pub struct GlobalSearchParams {
     pub max_data_tokens: usize,
     pub map_llm_params: LLMParams,
     pub reduce_llm_params: LLMParams,
-    pub context_builder_params: ContextBuilderParams,
+    pub context_builder_params: GlobalSearchContextBuilderParams,
 }
 
 impl GlobalSearch {

shinkai-libs/shinkai-graphrag/src/search/local_search/local_search.rs

@@ -0,0 +1,96 @@
+use std::{collections::HashMap, time::Instant};
+
+use crate::{
+    llm::llm::{BaseLLM, LLMParams, MessageType},
+    search::base::{ContextData, ContextText, ResponseType},
+};
+
+use super::{
+    mixed_context::{LocalSearchContextBuilderParams, LocalSearchMixedContext},
+    prompts::LOCAL_SEARCH_SYSTEM_PROMPT,
+};
+
+pub struct LocalSearchResult {
+    pub response: ResponseType,
+    pub context_data: ContextData,
+    pub context_text: ContextText,
+    pub completion_time: f64,
+    pub llm_calls: usize,
+    pub prompt_tokens: usize,
+}
+
+pub struct LocalSearch {
+    llm: Box<dyn BaseLLM>,
+    context_builder: LocalSearchMixedContext,
+    num_tokens_fn: fn(&str) -> usize,
+    system_prompt: String,
+    response_type: String,
+    llm_params: LLMParams,
+    context_builder_params: LocalSearchContextBuilderParams,
+}
+
+impl LocalSearch {
+    pub fn new(
+        llm: Box<dyn BaseLLM>,
+        context_builder: LocalSearchMixedContext,
+        num_tokens_fn: fn(&str) -> usize,
+        llm_params: LLMParams,
+        context_builder_params: LocalSearchContextBuilderParams,
+        response_type: String,
+        system_prompt: Option<String>,
+    ) -> Self {
+        let system_prompt = system_prompt.unwrap_or(LOCAL_SEARCH_SYSTEM_PROMPT.to_string());
+
+        LocalSearch {
+            llm,
+            context_builder,
+            num_tokens_fn,
+            system_prompt,
+            response_type,
+            llm_params,
+            context_builder_params,
+        }
+    }
+
+    pub async fn asearch(&self, query: String) -> anyhow::Result<LocalSearchResult> {
+        let start_time = Instant::now();
+        let (context_text, context_records) = self
+            .context_builder
+            .build_context(self.context_builder_params.clone())
+            .await?;
+
+        let search_prompt = self
+            .system_prompt
+            .replace("{context_data}", &context_text)
+            .replace("{response_type}", &self.response_type);
+
+        let mut search_messages = Vec::new();
+        search_messages.push(HashMap::from([
+            ("role".to_string(), "system".to_string()),
+            ("content".to_string(), search_prompt.clone()),
+        ]));
+        search_messages.push(HashMap::from([
+            ("role".to_string(), "user".to_string()),
+            ("content".to_string(), query.to_string()),
+        ]));
+
+        let search_response = self
+            .llm
+            .agenerate(
+                MessageType::Dictionary(search_messages),
+                false,
+                None,
+                self.llm_params.clone(),
+            )
+            .await?;
+
+        Ok(LocalSearchResult {
+            response: ResponseType::String(search_response),
+            context_data: ContextData::Dictionary(context_records),
+            context_text: ContextText::String(context_text),
+            completion_time: start_time.elapsed().as_secs_f64(),
+            llm_calls: 1,
+            prompt_tokens: (self.num_tokens_fn)(&search_prompt),
+        })
+    }
+}