File size: 8,828 Bytes
a37e6db | 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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 | """
GraphRAG Redis Worker (Upstash)
Background process to run community detection (Leiden Algorithm) on PostgreSQL Graph.
Consumes jobs from 'graph_reindex_jobs' Redis queue.
"""
import json
import logging
import os
import sys
import time
from typing import Any, Dict, List, Optional
import uuid
from sqlalchemy import text
# Add backend to path
sys.path.append(os.getcwd())
from core.database import SessionLocal, get_db_session
from core.models import CommunityMembership, GraphCommunity, GraphEdge, GraphNode
from core.service_factory import ServiceFactory
import asyncio
try:
import networkx as nx
except ImportError:
class MockGraph:
def __init__(self):
self._nodes = {}
self._edges = {}
@property
def nodes(self): return self._nodes
def add_node(self, id, **attr): self._nodes[id] = attr
def add_edge(self, u, v, **attr): self._edges[(u, v)] = attr
def number_of_nodes(self): return len(self._nodes)
class nx:
Graph = MockGraph
@staticmethod
def connected_components(G):
return [list(G.nodes.keys())]
# Configure Logging
logging.basicConfig(
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
level=logging.INFO
)
logger = logging.getLogger(__name__)
class RedisWorker:
def __init__(self, redis_url: str = None):
self.max_ram_nodes = 50000
self.queue_name = "graph_reindex_jobs"
self.redis_client = None
# Init Redis
redis_url = redis_url or os.getenv("UPSTASH_REDIS_URL") or os.getenv("REDIS_URL")
if redis_url:
try:
import redis
self.redis_client = redis.from_url(redis_url)
logger.info(f"Connected to Redis: {redis_url.split('@')[-1]}") # Log host only
except ImportError:
logger.warning("redis-py not installed.")
except Exception as e:
logger.error(f"Failed to connect to Redis: {e}")
else:
logger.warning("No REDIS_URL provided. Worker will not listen to queue.")
def fetch_graph(self, workspace_id: str) -> nx.Graph:
"""Load entire workspace graph into NetworkX"""
session = SessionLocal()
G = nx.Graph()
try:
logger.info(f"Fetching nodes for workspace {workspace_id}...")
nodes = session.query(GraphNode.id, GraphNode.name).filter_by(workspace_id=workspace_id).all()
for n in nodes:
G.add_node(n.id, name=n.name)
logger.info(f"Fetching edges for workspace {workspace_id}...")
edges = session.query(GraphEdge.source_node_id, GraphEdge.target_node_id, GraphEdge.weight).filter_by(workspace_id=workspace_id).all()
for e in edges:
G.add_edge(e.source_node_id, e.target_node_id, weight=e.weight)
return G
finally:
session.close()
def detect_communities(self, G: nx.Graph) -> List[List[str]]:
"""Run Louvain/Leiden algorithm"""
if G.number_of_nodes() == 0:
return []
try:
from networkx.algorithms.community import louvain_communities
logger.info(f"Running Louvain on {G.number_of_nodes()} nodes...")
communities = louvain_communities(G, seed=42)
return [list(c) for c in communities]
except ImportError:
logger.warning("Louvain not available, falling back to connected components")
return [list(c) for c in nx.connected_components(G)]
async def summarize_community(self, workspace_id: str, G: nx.Graph, community_nodes: List[str]) -> Dict[str, Any]:
"""Generate LLM summary and keywords for a community using unified LLMService"""
llm = ServiceFactory.get_llm_service()
# Prepare context
nodes_list = [f"- {G.nodes[n].get('name', 'Unknown')} ({G.nodes[n].get('type', 'entity')})" for n in community_nodes[:20]]
nodes_str = "\n".join(nodes_list)
prompt = f"""Summarize this knowledge graph community of related entities.
Entities:
{nodes_str}
Respond in valid JSON only with this structure:
{{
"summary": "Short 1-2 sentence description emphasizing the common theme",
"keywords": ["keyword1", "keyword2", "keyword3"]
}}"""
try:
# LLMService handles tenant isolation and usage tracking automatically
result = await llm.generate_response(
prompt=prompt,
tenant_id=workspace_id,
system_prompt="You are a GraphRAG Community Analyst. Categorize and summarize groups of entities.",
json_mode=True
)
data = json.loads(result)
return {
"summary": data.get("summary", f"Community of {len(community_nodes)} entities."),
"keywords": data.get("keywords", [])
}
except Exception as e:
logger.error(f"Failed to summarize community via LLMService: {e}")
node_names = [G.nodes[n].get("name", "Unknown") for n in community_nodes[:3]]
return {
"summary": f"Group related to {', '.join(node_names)}.",
"keywords": node_names
}
def save_communities(self, workspace_id: str, communities: List[List[str]], G: nx.Graph):
"""Persist results to Postgres"""
session = SessionLocal()
try:
session.execute(text("DELETE FROM graph_communities WHERE workspace_id = :ws_id"), {"ws_id": workspace_id})
session.commit()
logger.info(f"Summarizing and saving {len(communities)} communities...")
for i, members in enumerate(communities):
if len(members) < 2: continue
# Perform async LLM summarization
res = asyncio.run(self.summarize_community(workspace_id, G, members))
comm = GraphCommunity(
workspace_id=workspace_id,
level=0,
summary=res["summary"],
keywords=res["keywords"]
)
session.add(comm)
session.flush() # Get ID
for node_id in members:
membership = CommunityMembership(
community_id=comm.id,
node_id=node_id
)
session.add(membership)
session.commit()
logger.info(f"Saved {len(communities)} communities for workspace {workspace_id}")
except Exception as e:
session.rollback()
logger.error(f"Failed to save communities: {e}")
finally:
session.close()
def process_job(self, workspace_id: str):
logger.info(f"WORKER: Starting job for {workspace_id}")
G = self.fetch_graph(workspace_id)
if G.number_of_nodes() > self.max_ram_nodes:
logger.error(f"Graph too large ({G.number_of_nodes()} nodes).")
return
communities = self.detect_communities(G)
self.save_communities(workspace_id, communities, G)
logger.info("WORKER: Job Finished.")
def run(self):
"""Main listening loop"""
if not self.redis_client:
logger.error("Redis not connected. Exiting.")
return
logger.info(f"Listening on queue: {self.queue_name}...")
while True:
# Scale-to-Zero logic: If fetch returns None after timeout, exit
# For now, block indefinitely or use timeout
try:
# brpop returns tuple (queue_name, value)
job = self.redis_client.brpop(self.queue_name, timeout=30)
if job:
_, workspace_id_bytes = job
workspace_id = workspace_id_bytes.decode('utf-8')
self.process_job(workspace_id)
else:
logger.info("Queue empty (timeout). Idle...")
# In production with 'machines on demand', we would exit here
# sys.exit(0)
except Exception as e:
logger.error(f"Worker Error: {e}")
time.sleep(5)
if __name__ == "__main__":
# If run with argument, process single job (Manual/Test mode)
if len(sys.argv) > 1:
workspace_id = sys.argv[1]
worker = RedisWorker(redis_url="mock://") # Skip redis conn
worker.process_job(workspace_id)
else:
# Run in Daemon mode
worker = RedisWorker()
worker.run()
|