File size: 5,101 Bytes
c87117f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | import asyncio
import logging
from pathlib import Path
import pandas as pd
logger = logging.getLogger(__name__)
class GraphRAGEngine:
def __init__(self, root_dir: str):
self.root_dir = Path(root_dir)
self.global_search = None
self.local_search = None
self._loaded = False
self._load_error = None
def load(self):
try:
self._build_search_engines()
self._loaded = True
logger.info("GraphRAG engine loaded into memory")
except Exception as e:
self._load_error = str(e)
logger.error(f"Engine load failed: {e}")
raise
def query_sync(self, question: str, method: str) -> str:
if not self._loaded:
raise RuntimeError(f"Engine not loaded: {self._load_error}")
async def _run():
if method == "global":
result = await self.global_search.search(question)
else:
result = await self.local_search.search(question)
return result.response
return asyncio.run(_run())
def _build_search_engines(self):
from graphrag.config.load_config import load_config
from graphrag.query.factory import get_global_search_engine, get_local_search_engine
from graphrag.query.indexer_adapters import (
read_indexer_communities,
read_indexer_entities,
read_indexer_relationships,
read_indexer_reports,
read_indexer_text_units,
)
from graphrag.vector_stores.lancedb import LanceDBVectorStore
config = load_config(root_dir=self.root_dir)
reduce_prompt_path = self.root_dir / "prompts" / "global_search_reduce_system_prompt.txt"
reduce_prompt = reduce_prompt_path.read_text(encoding="utf-8") if reduce_prompt_path.exists() else None
map_prompt_path = self.root_dir / "prompts" / "global_search_map_system_prompt.txt"
map_prompt = map_prompt_path.read_text(encoding="utf-8") if map_prompt_path.exists() else None
out = self.root_dir / "output"
entities_df = pd.read_parquet(out / "entities.parquet")
communities_df = pd.read_parquet(out / "communities.parquet")
community_reports_df = pd.read_parquet(out / "community_reports.parquet")
relationships_df = pd.read_parquet(out / "relationships.parquet")
text_units_df = pd.read_parquet(out / "text_units.parquet")
COMMUNITY_LEVEL = 2
entities = read_indexer_entities(entities_df, communities_df, COMMUNITY_LEVEL)
relationships = read_indexer_relationships(relationships_df)
communities = read_indexer_communities(communities_df, community_reports_df)
reports = read_indexer_reports(community_reports_df, communities_df, COMMUNITY_LEVEL)
text_units = read_indexer_text_units(text_units_df)
logger.info(
f"Loaded: {len(entities)} entities, {len(relationships)} relationships, "
f"{len(communities)} communities, {len(reports)} reports, {len(text_units)} text units"
)
lancedb_uri = str(self.root_dir / "output" / "lancedb")
collection = self._detect_entity_collection(lancedb_uri)
logger.info(f"Using LanceDB collection: {collection}")
entity_embedding_store = LanceDBVectorStore(collection_name=collection)
entity_embedding_store.connect(db_uri=lancedb_uri)
self.global_search = get_global_search_engine(
config=config,
reports=reports,
entities=entities,
communities=communities,
response_type="multiple paragraphs",
map_system_prompt=map_prompt,
reduce_system_prompt=reduce_prompt,
)
self.local_search = get_local_search_engine(
config=config,
reports=reports,
text_units=text_units,
entities=entities,
relationships=relationships,
covariates={},
description_embedding_store=entity_embedding_store,
response_type="multiple paragraphs",
)
@staticmethod
def _detect_entity_collection(lancedb_uri: str) -> str:
try:
import lancedb
db = lancedb.connect(lancedb_uri)
tables = db.table_names()
logger.info(f"LanceDB tables found: {tables}")
for candidate in tables:
if "entity" in candidate.lower():
return candidate
if tables:
return tables[0]
except Exception as e:
logger.warning(f"LanceDB detection failed: {e}")
return "default-entity-description"
_engine = None
def get_engine():
global _engine
if _engine is None:
raise RuntimeError("Engine not initialised. Call init_engine() first.")
return _engine
def init_engine(root_dir: str):
global _engine
if _engine is None:
_engine = GraphRAGEngine(root_dir)
_engine.load()
return _engine |