gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix event loop conflict

Browse Source
This commit is contained in:
Ken Wiltshire
2024-11-06 11:18:14 -05:00
parent a73ef6ec5a
commit ac0b7fd775
9 changed files with 185 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
get_all_edges_nx.py
View File

@@ -1,22 +1,28 @@
import networkx as nx import networkx as nx
G = nx.read_graphml('./dickensTestEmbedcall/graph_chunk_entity_relation.graphml') G = nx.read_graphml("./dickensTestEmbedcall/graph_chunk_entity_relation.graphml")
def get_all_edges_and_nodes(G): def get_all_edges_and_nodes(G):
# Get all edges and their properties # Get all edges and their properties
edges_with_properties = [] edges_with_properties = []
for u, v, data in G.edges(data=True): for u, v, data in G.edges(data=True):
edges_with_properties.append({ edges_with_properties.append(
'start': u, {
'end': v, "start": u,
'label': data.get('label', ''), # Assuming 'label' is used for edge type "end": v,
'properties': data, "label": data.get(
'start_node_properties': G.nodes[u], "label", ""
'end_node_properties': G.nodes[v] ), # Assuming 'label' is used for edge type
}) "properties": data,
"start_node_properties": G.nodes[u],
"end_node_properties": G.nodes[v],
}
)
return edges_with_properties return edges_with_properties
# Example usage # Example usage
if __name__ == "__main__": if __name__ == "__main__":
# Assume G is your NetworkX graph loaded from Neo4j # Assume G is your NetworkX graph loaded from Neo4j

126
lightrag/kg/neo4j_impl.py
View File

@@ -1,17 +1,16 @@
import asyncio import asyncio
import html
import os import os
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any, Union, cast, Tuple, List, Dict from typing import Any, Union, Tuple, List, Dict
import numpy as np
import inspect import inspect
from lightrag.utils import load_json, logger, write_json from lightrag.utils import logger
from ..base import ( from ..base import BaseGraphStorage
BaseGraphStorage from neo4j import (
AsyncGraphDatabase,
exceptions as neo4jExceptions,
AsyncDriver,
AsyncManagedTransaction,
) )
from neo4j import AsyncGraphDatabase,exceptions as neo4jExceptions,AsyncDriver,AsyncSession, AsyncManagedTransaction
from contextlib import asynccontextmanager
from tenacity import ( from tenacity import (
@@ -35,7 +34,9 @@ class Neo4JStorage(BaseGraphStorage):
URI = os.environ["NEO4J_URI"] URI = os.environ["NEO4J_URI"]
USERNAME = os.environ["NEO4J_USERNAME"] USERNAME = os.environ["NEO4J_USERNAME"]
PASSWORD = os.environ["NEO4J_PASSWORD"] PASSWORD = os.environ["NEO4J_PASSWORD"]
self._driver: AsyncDriver = AsyncGraphDatabase.driver(URI, auth=(USERNAME, PASSWORD)) self._driver: AsyncDriver = AsyncGraphDatabase.driver(
URI, auth=(USERNAME, PASSWORD)
)
return None return None
def __post_init__(self): def __post_init__(self):
@@ -43,14 +44,11 @@ class Neo4JStorage(BaseGraphStorage):
"node2vec": self._node2vec_embed, "node2vec": self._node2vec_embed,
} }
async def close(self): async def close(self):
if self._driver: if self._driver:
await self._driver.close() await self._driver.close()
self._driver = None self._driver = None
async def __aexit__(self, exc_type, exc, tb): async def __aexit__(self, exc_type, exc, tb):
if self._driver: if self._driver:
await self._driver.close() await self._driver.close()
@@ -58,12 +56,13 @@ class Neo4JStorage(BaseGraphStorage):
async def index_done_callback(self): async def index_done_callback(self):
print("KG successfully indexed.") print("KG successfully indexed.")
async def has_node(self, node_id: str) -> bool: async def has_node(self, node_id: str) -> bool:
entity_name_label = node_id.strip('\"') entity_name_label = node_id.strip('"')
async with self._driver.session() as session: async with self._driver.session() as session:
query = f"MATCH (n:`{entity_name_label}`) RETURN count(n) > 0 AS node_exists" query = (
f"MATCH (n:`{entity_name_label}`) RETURN count(n) > 0 AS node_exists"
)
result = await session.run(query) result = await session.run(query)
single_result = await result.single() single_result = await result.single()
logger.debug( logger.debug(
@@ -72,8 +71,8 @@ class Neo4JStorage(BaseGraphStorage):
return single_result["node_exists"] return single_result["node_exists"]
async def has_edge(self, source_node_id: str, target_node_id: str) -> bool: async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
entity_name_label_source = source_node_id.strip('\"') entity_name_label_source = source_node_id.strip('"')
entity_name_label_target = target_node_id.strip('\"') entity_name_label_target = target_node_id.strip('"')
async with self._driver.session() as session: async with self._driver.session() as session:
query = ( query = (
@@ -90,12 +89,9 @@ class Neo4JStorage(BaseGraphStorage):
def close(self): def close(self):
self._driver.close() self._driver.close()
async def get_node(self, node_id: str) -> Union[dict, None]: async def get_node(self, node_id: str) -> Union[dict, None]:
async with self._driver.session() as session: async with self._driver.session() as session:
entity_name_label = node_id.strip('\"') entity_name_label = node_id.strip('"')
query = f"MATCH (n:`{entity_name_label}`) RETURN n" query = f"MATCH (n:`{entity_name_label}`) RETURN n"
result = await session.run(query) result = await session.run(query)
record = await result.single() record = await result.single()
@@ -103,15 +99,13 @@ class Neo4JStorage(BaseGraphStorage):
node = record["n"] node = record["n"]
node_dict = dict(node) node_dict = dict(node)
logger.debug( logger.debug(
f'{inspect.currentframe().f_code.co_name}: query: {query}, result: {node_dict}' f"{inspect.currentframe().f_code.co_name}: query: {query}, result: {node_dict}"
) )
return node_dict return node_dict
return None return None
async def node_degree(self, node_id: str) -> int: async def node_degree(self, node_id: str) -> int:
entity_name_label = node_id.strip('\"') entity_name_label = node_id.strip('"')
async with self._driver.session() as session: async with self._driver.session() as session:
query = f""" query = f"""
@@ -123,16 +117,15 @@ class Neo4JStorage(BaseGraphStorage):
if record: if record:
edge_count = record["totalEdgeCount"] edge_count = record["totalEdgeCount"]
logger.debug( logger.debug(
f'{inspect.currentframe().f_code.co_name}:query:{query}:result:{edge_count}' f"{inspect.currentframe().f_code.co_name}:query:{query}:result:{edge_count}"
) )
return edge_count return edge_count
else: else:
return None return None
async def edge_degree(self, src_id: str, tgt_id: str) -> int: async def edge_degree(self, src_id: str, tgt_id: str) -> int:
entity_name_label_source = src_id.strip('\"') entity_name_label_source = src_id.strip('"')
entity_name_label_target = tgt_id.strip('\"') entity_name_label_target = tgt_id.strip('"')
src_degree = await self.node_degree(entity_name_label_source) src_degree = await self.node_degree(entity_name_label_source)
trg_degree = await self.node_degree(entity_name_label_target) trg_degree = await self.node_degree(entity_name_label_target)
@@ -142,15 +135,15 @@ class Neo4JStorage(BaseGraphStorage):
degrees = int(src_degree) + int(trg_degree) degrees = int(src_degree) + int(trg_degree)
logger.debug( logger.debug(
f'{inspect.currentframe().f_code.co_name}:query:src_Degree+trg_degree:result:{degrees}' f"{inspect.currentframe().f_code.co_name}:query:src_Degree+trg_degree:result:{degrees}"
) )
return degrees return degrees
async def get_edge(
self, source_node_id: str, target_node_id: str
async def get_edge(self, source_node_id: str, target_node_id: str) -> Union[dict, None]: ) -> Union[dict, None]:
entity_name_label_source = source_node_id.strip('\"') entity_name_label_source = source_node_id.strip('"')
entity_name_label_target = target_node_id.strip('\"') entity_name_label_target = target_node_id.strip('"')
""" """
Find all edges between nodes of two given labels Find all edges between nodes of two given labels
@@ -166,22 +159,24 @@ class Neo4JStorage(BaseGraphStorage):
MATCH (start:`{entity_name_label_source}`)-[r]->(end:`{entity_name_label_target}`) MATCH (start:`{entity_name_label_source}`)-[r]->(end:`{entity_name_label_target}`)
RETURN properties(r) as edge_properties RETURN properties(r) as edge_properties
LIMIT 1 LIMIT 1
""".format(entity_name_label_source=entity_name_label_source, entity_name_label_target=entity_name_label_target) """.format(
entity_name_label_source=entity_name_label_source,
entity_name_label_target=entity_name_label_target,
)
result = await session.run(query) result = await session.run(query)
record = await result.single() record = await result.single()
if record: if record:
result = dict(record["edge_properties"]) result = dict(record["edge_properties"])
logger.debug( logger.debug(
f'{inspect.currentframe().f_code.co_name}:query:{query}:result:{result}' f"{inspect.currentframe().f_code.co_name}:query:{query}:result:{result}"
) )
return result return result
else: else:
return None return None
async def get_node_edges(self, source_node_id: str) -> List[Tuple[str, str]]: async def get_node_edges(self, source_node_id: str) -> List[Tuple[str, str]]:
node_label = source_node_id.strip('\"') node_label = source_node_id.strip('"')
""" """
Retrieves all edges (relationships) for a particular node identified by its label. Retrieves all edges (relationships) for a particular node identified by its label.
@@ -194,22 +189,33 @@ class Neo4JStorage(BaseGraphStorage):
results = await session.run(query) results = await session.run(query)
edges = [] edges = []
async for record in results: async for record in results:
source_node = record['n'] source_node = record["n"]
connected_node = record['connected'] connected_node = record["connected"]
source_label = list(source_node.labels)[0] if source_node.labels else None source_label = (
target_label = list(connected_node.labels)[0] if connected_node and connected_node.labels else None list(source_node.labels)[0] if source_node.labels else None
)
target_label = (
list(connected_node.labels)[0]
if connected_node and connected_node.labels
else None
)
if source_label and target_label: if source_label and target_label:
edges.append((source_label, target_label)) edges.append((source_label, target_label))
return edges return edges
@retry( @retry(
stop=stop_after_attempt(3), stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10), wait=wait_exponential(multiplier=1, min=4, max=10),
retry=retry_if_exception_type((neo4jExceptions.ServiceUnavailable, neo4jExceptions.TransientError, neo4jExceptions.WriteServiceUnavailable)), retry=retry_if_exception_type(
(
neo4jExceptions.ServiceUnavailable,
neo4jExceptions.TransientError,
neo4jExceptions.WriteServiceUnavailable,
)
),
) )
async def upsert_node(self, node_id: str, node_data: Dict[str, Any]): async def upsert_node(self, node_id: str, node_data: Dict[str, Any]):
""" """
@@ -219,7 +225,7 @@ class Neo4JStorage(BaseGraphStorage):
node_id: The unique identifier for the node (used as label) node_id: The unique identifier for the node (used as label)
node_data: Dictionary of node properties node_data: Dictionary of node properties
""" """
label = node_id.strip('\"') label = node_id.strip('"')
properties = node_data properties = node_data
async def _do_upsert(tx: AsyncManagedTransaction): async def _do_upsert(tx: AsyncManagedTransaction):
@@ -228,7 +234,9 @@ class Neo4JStorage(BaseGraphStorage):
SET n += $properties SET n += $properties
""" """
await tx.run(query, properties=properties) await tx.run(query, properties=properties)
logger.debug(f"Upserted node with label '{label}' and properties: {properties}") logger.debug(
f"Upserted node with label '{label}' and properties: {properties}"
)
try: try:
async with self._driver.session() as session: async with self._driver.session() as session:
@@ -240,9 +248,17 @@ class Neo4JStorage(BaseGraphStorage):
@retry( @retry(
stop=stop_after_attempt(3), stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10), wait=wait_exponential(multiplier=1, min=4, max=10),
retry=retry_if_exception_type((neo4jExceptions.ServiceUnavailable, neo4jExceptions.TransientError, neo4jExceptions.WriteServiceUnavailable)), retry=retry_if_exception_type(
(
neo4jExceptions.ServiceUnavailable,
neo4jExceptions.TransientError,
neo4jExceptions.WriteServiceUnavailable,
) )
async def upsert_edge(self, source_node_id: str, target_node_id: str, edge_data: Dict[str, Any]): ),
)
async def upsert_edge(
self, source_node_id: str, target_node_id: str, edge_data: Dict[str, Any]
):
""" """
Upsert an edge and its properties between two nodes identified by their labels. Upsert an edge and its properties between two nodes identified by their labels.
@@ -251,8 +267,8 @@ class Neo4JStorage(BaseGraphStorage):
target_node_id (str): Label of the target node (used as identifier) target_node_id (str): Label of the target node (used as identifier)
edge_data (dict): Dictionary of properties to set on the edge edge_data (dict): Dictionary of properties to set on the edge
""" """
source_node_label = source_node_id.strip('\"') source_node_label = source_node_id.strip('"')
target_node_label = target_node_id.strip('\"') target_node_label = target_node_id.strip('"')
edge_properties = edge_data edge_properties = edge_data
async def _do_upsert_edge(tx: AsyncManagedTransaction): async def _do_upsert_edge(tx: AsyncManagedTransaction):
@@ -265,7 +281,9 @@ class Neo4JStorage(BaseGraphStorage):
RETURN r RETURN r
""" """
await tx.run(query, properties=edge_properties) await tx.run(query, properties=edge_properties)
logger.debug(f"Upserted edge from '{source_node_label}' to '{target_node_label}' with properties: {edge_properties}") logger.debug(
f"Upserted edge from '{source_node_label}' to '{target_node_label}' with properties: {edge_properties}"
)
try: try:
async with self._driver.session() as session: async with self._driver.session() as session:
@@ -273,6 +291,6 @@ class Neo4JStorage(BaseGraphStorage):
except Exception as e: except Exception as e:
logger.error(f"Error during edge upsert: {str(e)}") logger.error(f"Error during edge upsert: {str(e)}")
raise raise
async def _node2vec_embed(self): async def _node2vec_embed(self):
print("Implemented but never called.") print("Implemented but never called.")

15
lightrag/lightrag.py
View File

@@ -1,6 +1,5 @@
import asyncio import asyncio
import os import os
import importlib
from dataclasses import asdict, dataclass, field from dataclasses import asdict, dataclass, field
from datetime import datetime from datetime import datetime
from functools import partial from functools import partial
@@ -25,9 +24,7 @@ from .storage import (
NetworkXStorage, NetworkXStorage,
) )
from .kg.neo4j_impl import ( from .kg.neo4j_impl import Neo4JStorage
Neo4JStorage
)
# future KG integrations # future KG integrations
# from .kg.ArangoDB_impl import ( # from .kg.ArangoDB_impl import (
@@ -35,7 +32,6 @@ from .kg.neo4j_impl import (
# ) # )
from .utils import ( from .utils import (
EmbeddingFunc, EmbeddingFunc,
compute_mdhash_id, compute_mdhash_id,
@@ -52,6 +48,7 @@ from .base import (
QueryParam, QueryParam,
) )
def always_get_an_event_loop() -> asyncio.AbstractEventLoop: def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
try: try:
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
@@ -64,7 +61,6 @@ def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
@dataclass @dataclass
class LightRAG: class LightRAG:
working_dir: str = field( working_dir: str = field(
default_factory=lambda: f"./lightrag_cache_{datetime.now().strftime('%Y-%m-%d-%H:%M:%S')}" default_factory=lambda: f"./lightrag_cache_{datetime.now().strftime('%Y-%m-%d-%H:%M:%S')}"
) )
@@ -74,8 +70,6 @@ class LightRAG:
current_log_level = logger.level current_log_level = logger.level
log_level: str = field(default=current_log_level) log_level: str = field(default=current_log_level)
# text chunking # text chunking
chunk_token_size: int = 1200 chunk_token_size: int = 1200
chunk_overlap_token_size: int = 100 chunk_overlap_token_size: int = 100
@@ -131,7 +125,9 @@ class LightRAG:
logger.debug(f"LightRAG init with param:\n {_print_config}\n") logger.debug(f"LightRAG init with param:\n {_print_config}\n")
# @TODO: should move all storage setup here to leverage initial start params attached to self. # @TODO: should move all storage setup here to leverage initial start params attached to self.
self.graph_storage_cls: Type[BaseGraphStorage] = self._get_storage_class()[self.kg] self.graph_storage_cls: Type[BaseGraphStorage] = self._get_storage_class()[
self.kg
]
if not os.path.exists(self.working_dir): if not os.path.exists(self.working_dir):
logger.info(f"Creating working directory {self.working_dir}") logger.info(f"Creating working directory {self.working_dir}")
@@ -185,6 +181,7 @@ class LightRAG:
**self.llm_model_kwargs, **self.llm_model_kwargs,
) )
) )
def _get_storage_class(self) -> Type[BaseGraphStorage]: def _get_storage_class(self) -> Type[BaseGraphStorage]:
return { return {
"Neo4JStorage": Neo4JStorage, "Neo4JStorage": Neo4JStorage,

8
lightrag/operate.py
View File

@@ -466,7 +466,6 @@ async def _build_local_query_context(
text_chunks_db: BaseKVStorage[TextChunkSchema], text_chunks_db: BaseKVStorage[TextChunkSchema],
query_param: QueryParam, query_param: QueryParam,
): ):
results = await entities_vdb.query(query, top_k=query_param.top_k) results = await entities_vdb.query(query, top_k=query_param.top_k)
if not len(results): if not len(results):
@@ -928,7 +927,6 @@ async def hybrid_query(
query_param, query_param,
) )
if hl_keywords: if hl_keywords:
high_level_context = await _build_global_query_context( high_level_context = await _build_global_query_context(
hl_keywords, hl_keywords,
@@ -939,7 +937,6 @@ async def hybrid_query(
query_param, query_param,
) )
context = combine_contexts(high_level_context, low_level_context) context = combine_contexts(high_level_context, low_level_context)
if query_param.only_need_context: if query_param.only_need_context:
@@ -1010,7 +1007,9 @@ def combine_contexts(high_level_context, low_level_context):
combined_entities = process_combine_contexts(hl_entities, ll_entities) combined_entities = process_combine_contexts(hl_entities, ll_entities)
# Combine and deduplicate the relationships # Combine and deduplicate the relationships
combined_relationships = process_combine_contexts(hl_relationships, ll_relationships) combined_relationships = process_combine_contexts(
hl_relationships, ll_relationships
)
# Combine and deduplicate the sources # Combine and deduplicate the sources
combined_sources = process_combine_contexts(hl_sources, ll_sources) combined_sources = process_combine_contexts(hl_sources, ll_sources)
@@ -1046,7 +1045,6 @@ async def naive_query(
chunks_ids = [r["id"] for r in results] chunks_ids = [r["id"] for r in results]
chunks = await text_chunks_db.get_by_ids(chunks_ids) chunks = await text_chunks_db.get_by_ids(chunks_ids)
maybe_trun_chunks = truncate_list_by_token_size( maybe_trun_chunks = truncate_list_by_token_size(
chunks, chunks,
key=lambda x: x["content"], key=lambda x: x["content"],

1
lightrag/storage.py
View File

@@ -233,7 +233,6 @@ class NetworkXStorage(BaseGraphStorage):
raise ValueError(f"Node embedding algorithm {algorithm} not supported") raise ValueError(f"Node embedding algorithm {algorithm} not supported")
return await self._node_embed_algorithms[algorithm]() return await self._node_embed_algorithms[algorithm]()
# @TODO: NOT USED # @TODO: NOT USED
async def _node2vec_embed(self): async def _node2vec_embed(self):
from graspologic import embed from graspologic import embed

14
lightrag/utils.py
View File

@@ -176,19 +176,20 @@ def truncate_list_by_token_size(list_data: list, key: callable, max_token_size:
return list_data[:i] return list_data[:i]
return list_data return list_data
def list_of_list_to_csv(data: List[List[str]]) -> str: def list_of_list_to_csv(data: List[List[str]]) -> str:
output = io.StringIO() output = io.StringIO()
writer = csv.writer(output) writer = csv.writer(output)
writer.writerows(data) writer.writerows(data)
return output.getvalue() return output.getvalue()
def csv_string_to_list(csv_string: str) -> List[List[str]]: def csv_string_to_list(csv_string: str) -> List[List[str]]:
output = io.StringIO(csv_string) output = io.StringIO(csv_string)
reader = csv.reader(output) reader = csv.reader(output)
return [row for row in reader] return [row for row in reader]
def save_data_to_file(data, file_name): def save_data_to_file(data, file_name):
with open(file_name, "w", encoding="utf-8") as f: with open(file_name, "w", encoding="utf-8") as f:
json.dump(data, f, ensure_ascii=False, indent=4) json.dump(data, f, ensure_ascii=False, indent=4)
@@ -253,6 +254,7 @@ def xml_to_json(xml_file):
print(f"An error occurred: {e}") print(f"An error occurred: {e}")
return None return None
def process_combine_contexts(hl, ll): def process_combine_contexts(hl, ll):
header = None header = None
list_hl = csv_string_to_list(hl.strip()) list_hl = csv_string_to_list(hl.strip())
@@ -268,13 +270,11 @@ def process_combine_contexts(hl, ll):
return "" return ""
if list_hl: if list_hl:
list_hl = [','.join(item[1:]) for item in list_hl if item] list_hl = [",".join(item[1:]) for item in list_hl if item]
if list_ll: if list_ll:
list_ll = [','.join(item[1:]) for item in list_ll if item] list_ll = [",".join(item[1:]) for item in list_ll if item]
combined_sources_set = set( combined_sources_set = set(filter(None, list_hl + list_ll))
filter(None, list_hl + list_ll)
)
combined_sources = [",\t".join(header)] combined_sources = [",\t".join(header)]

19
test.py
View File

@@ -1,7 +1,6 @@
import os import os
from lightrag import LightRAG, QueryParam from lightrag import LightRAG, QueryParam
from lightrag.llm import gpt_4o_mini_complete, gpt_4o_complete from lightrag.llm import gpt_4o_mini_complete, gpt_4o_complete
from pprint import pprint
######### #########
# Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert() # Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert()
# import nest_asyncio # import nest_asyncio
@@ -15,7 +14,7 @@ if not os.path.exists(WORKING_DIR):
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
llm_model_func=gpt_4o_mini_complete # Use gpt_4o_mini_complete LLM model llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model
# llm_model_func=gpt_4o_complete # Optionally, use a stronger model # llm_model_func=gpt_4o_complete # Optionally, use a stronger model
) )
@@ -23,13 +22,21 @@ with open("./book.txt") as f:
rag.insert(f.read()) rag.insert(f.read())
# Perform naive search # Perform naive search
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))) print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
)
# Perform local search # Perform local search
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))) print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
)
# Perform global search # Perform global search
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))) print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
)
# Perform hybrid search # Perform hybrid search
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))) print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
)

18
test_neo4j.py
View File

@@ -18,7 +18,7 @@ rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model
kg="Neo4JStorage", kg="Neo4JStorage",
log_level="INFO" log_level="INFO",
# llm_model_func=gpt_4o_complete # Optionally, use a stronger model # llm_model_func=gpt_4o_complete # Optionally, use a stronger model
) )
@@ -26,13 +26,21 @@ with open("./book.txt") as f:
rag.insert(f.read()) rag.insert(f.read())
# Perform naive search # Perform naive search
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))) print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
)
# Perform local search # Perform local search
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))) print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
)
# Perform global search # Perform global search
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))) print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
)
# Perform hybrid search # Perform hybrid search
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))) print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
)