Agent Zero
3d9b76cea4
Check Cross-Plugin Imports / check (push) Has been cancelled
- SPIKE-E: FTS+Vector+Permission benchmark on 10k records (all <30ms) - E-PROV: supports_fts/vector/rag/graph capability flags on all providers - E-FTS/VEC: All 11 providers refactored to BaseSearchProvider with permission filtering - E-PERM: Over-fetch strategy for vector+permission (15x faster than ANY() filter) - E-FUSE: rrf_fusion_multi() for N-way RRF over FTS+Vector+RAG+Graph - E-LLM: Query understanding cleaned up to use central llm_complete() - E-CHUNK: Document chunking module + document_chunks table with HNSW index - E-EMB: Chunk embedding ARQ jobs (index_file_chunks, reindex_chunks) - E-RAG: RAG retrieval via FileSearchProvider.search_rag() - E-GRAPH: GraphRAG BFS traversal via GraphRAGSearchProvider.search_graph() - E-IX-EVT: Auto-indexing via outbox events + delete/cleanup handlers - E-IX-RE: Batch reindex with progress tracking + reindex_all job - E-DATA-LIFE: Lifecycle module (remove/rebuild/restore/correct) + API endpoints - E-K-MEM: AgentMemorySearchProvider - E-P-AI: AIChatSearchProvider - E-P-WF: WorkflowSearchProvider - E-P-COMM: ConversationSearchProvider verified (already on BaseSearchProvider) - E-API: Filter params (date_from/to, tags, sort) + /facets endpoint - E-TOOL: unified_search AI tool registered in ToolRegistry - E-MCP: Search tool in MCP server with normal RBAC/tenant checks - E-UI-CMD: CommandPalette (Cmd+K) with debounced search + recent searches - E-UI-FAC: SearchFacets, SearchResultCard, SavedSearches components - E-TEST: 40 new tests in test_unified_search_phase_e.py (105 total green) - E-DOC: api-documentation.md, plugin-development-guide.md, test-strategy.md updated 105 tests passing, TypeScript clean.
300 lines
11 KiB
Python
300 lines
11 KiB
Python
"""GraphRAGSearchProvider — registers GraphRAG as a search provider in unified_search."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import uuid
|
|
from typing import Any
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.plugins.builtins.unified_search.contracts import BaseSearchProvider
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class GraphRAGSearchProvider(BaseSearchProvider):
|
|
supports_fts: bool = False
|
|
supports_vector: bool = False
|
|
supports_rag: bool = False
|
|
supports_graph: bool = True
|
|
"""Search provider for GraphRAG entity relationships.
|
|
|
|
Enables full-text and semantic search over relationship metadata
|
|
and entity types in the knowledge graph.
|
|
"""
|
|
|
|
entity_type = "graph_relationship"
|
|
|
|
async def _search_fts_filtered(
|
|
self,
|
|
db: AsyncSession,
|
|
tsquery: str,
|
|
tenant_id: uuid.UUID,
|
|
limit: int,
|
|
visible_ids: set[uuid.UUID] | None,
|
|
) -> list[dict[str, Any]]:
|
|
"""Full-text search on relationship metadata and types."""
|
|
if visible_ids is not None:
|
|
sql = text(
|
|
"""
|
|
SELECT r.*, ts_rank(
|
|
to_tsvector('pg_catalog.german',
|
|
coalesce(r.relationship_type, '') || ' ' ||
|
|
coalesce(r.source_type, '') || ' ' ||
|
|
coalesce(r.target_type, '') || ' ' ||
|
|
coalesce(r.meta::text, '')
|
|
),
|
|
to_tsquery('pg_catalog.german', :q)
|
|
) AS rank
|
|
FROM entity_relationships r
|
|
WHERE r.tenant_id = :tid
|
|
AND r.deleted_at IS NULL
|
|
AND to_tsvector('pg_catalog.german',
|
|
coalesce(r.relationship_type, '') || ' ' ||
|
|
coalesce(r.source_type, '') || ' ' ||
|
|
coalesce(r.target_type, '') || ' ' ||
|
|
coalesce(r.meta::text, '')
|
|
) @@ to_tsquery('pg_catalog.german', :q)
|
|
AND r.id = ANY(:visible_ids)
|
|
ORDER BY rank DESC
|
|
LIMIT :lim
|
|
"""
|
|
)
|
|
result = await db.execute(
|
|
sql,
|
|
{
|
|
"q": tsquery,
|
|
"tid": tenant_id,
|
|
"lim": limit,
|
|
"visible_ids": list(visible_ids),
|
|
},
|
|
)
|
|
else:
|
|
sql = text(
|
|
"""
|
|
SELECT r.*, ts_rank(
|
|
to_tsvector('pg_catalog.german',
|
|
coalesce(r.relationship_type, '') || ' ' ||
|
|
coalesce(r.source_type, '') || ' ' ||
|
|
coalesce(r.target_type, '') || ' ' ||
|
|
coalesce(r.meta::text, '')
|
|
),
|
|
to_tsquery('pg_catalog.german', :q)
|
|
) AS rank
|
|
FROM entity_relationships r
|
|
WHERE r.tenant_id = :tid
|
|
AND r.deleted_at IS NULL
|
|
AND to_tsvector('pg_catalog.german',
|
|
coalesce(r.relationship_type, '') || ' ' ||
|
|
coalesce(r.source_type, '') || ' ' ||
|
|
coalesce(r.target_type, '') || ' ' ||
|
|
coalesce(r.meta::text, '')
|
|
) @@ to_tsquery('pg_catalog.german', :q)
|
|
ORDER BY rank DESC
|
|
LIMIT :lim
|
|
"""
|
|
)
|
|
result = await db.execute(
|
|
sql,
|
|
{"q": tsquery, "tid": tenant_id, "lim": limit},
|
|
)
|
|
rows = result.mappings().all()
|
|
return [dict(r) for r in rows]
|
|
|
|
async def _search_vector_filtered(
|
|
self,
|
|
db: AsyncSession,
|
|
embedding: list[float],
|
|
tenant_id: uuid.UUID,
|
|
limit: int,
|
|
visible_ids: set[uuid.UUID] | None,
|
|
) -> list[dict[str, Any]]:
|
|
"""Semantic search is not yet supported for graph relationships.
|
|
|
|
Returns empty list — relationships are searched via FTS on metadata.
|
|
"""
|
|
return []
|
|
|
|
async def get_embedding_text(
|
|
self, db: AsyncSession, entity_id: uuid.UUID, tenant_id: uuid.UUID
|
|
) -> str:
|
|
"""Get text for embedding generation."""
|
|
sql = text(
|
|
"""
|
|
SELECT relationship_type, source_type, source_id, target_type, target_id, metadata
|
|
FROM entity_relationships
|
|
WHERE id = :eid AND tenant_id = :tid
|
|
"""
|
|
)
|
|
result = await db.execute(sql, {"eid": entity_id, "tid": tenant_id})
|
|
row = result.mappings().first()
|
|
if not row:
|
|
return ""
|
|
parts = [
|
|
row.get("relationship_type", ""),
|
|
row.get("source_type", ""),
|
|
str(row.get("source_id", "")),
|
|
row.get("target_type", ""),
|
|
str(row.get("target_id", "")),
|
|
str(row.get("metadata", {})),
|
|
]
|
|
return " ".join(str(p) for p in parts if p)
|
|
|
|
async def search_graph(
|
|
self,
|
|
db: AsyncSession,
|
|
query_analysis: dict[str, Any],
|
|
tenant_id: uuid.UUID,
|
|
limit: int,
|
|
user_id: uuid.UUID | None = None,
|
|
is_system_admin: bool = False,
|
|
) -> list[dict[str, Any]]:
|
|
"""Graph traversal search using BFS on entity_relationships.
|
|
|
|
Starts from entities found by FTS on relationship metadata, then
|
|
expands via BFS to find related entities up to depth 2.
|
|
Returns related entities with relationship metadata.
|
|
"""
|
|
normalized_query = query_analysis.get("normalized_query", "")
|
|
semantic_terms = query_analysis.get("semantic_terms", [])
|
|
tsquery_parts = [normalized_query] + semantic_terms
|
|
tsquery = " & ".join(
|
|
part.strip().replace(" ", " & ")
|
|
for part in tsquery_parts
|
|
if part and part.strip()
|
|
)
|
|
if not tsquery:
|
|
tsquery = normalized_query
|
|
if not tsquery:
|
|
return []
|
|
|
|
# Step 1: Find seed relationships via FTS
|
|
seed_sql = text(
|
|
"""
|
|
SELECT r.id, r.source_type, r.source_id, r.target_type, r.target_id,
|
|
r.relationship_type, r.metadata
|
|
FROM entity_relationships r
|
|
WHERE r.tenant_id = :tid
|
|
AND r.deleted_at IS NULL
|
|
AND to_tsvector('pg_catalog.german',
|
|
coalesce(r.relationship_type, '') || ' ' ||
|
|
coalesce(r.source_type, '') || ' ' ||
|
|
coalesce(r.target_type, '') || ' ' ||
|
|
coalesce(r.metadata::text, '')
|
|
) @@ to_tsquery('pg_catalog.german', :q)
|
|
LIMIT :lim
|
|
"""
|
|
)
|
|
result = await db.execute(seed_sql, {"tid": tenant_id, "q": tsquery, "lim": limit})
|
|
seed_rows = result.mappings().all()
|
|
if not seed_rows:
|
|
return []
|
|
|
|
# Collect seed entity IDs (both source and target)
|
|
seed_entities: dict[str, set[str]] = {}
|
|
for row in seed_rows:
|
|
for role in ("source", "target"):
|
|
etype = row[f"{role}_type"]
|
|
eid = str(row[f"{role}_id"])
|
|
seed_entities.setdefault(etype, set()).add(eid)
|
|
|
|
# Step 2: BFS traversal — find relationships connected to seed entities
|
|
all_entity_ids: dict[str, set[str]] = {}
|
|
for etype, eids in seed_entities.items():
|
|
all_entity_ids.setdefault(etype, set()).update(eids)
|
|
|
|
visited_rel_ids: set[str] = {str(r["id"]) for r in seed_rows}
|
|
results: list[dict[str, Any]] = []
|
|
|
|
# Add seed results first
|
|
for row in seed_rows:
|
|
results.append(self._relationship_to_result(dict(row)))
|
|
|
|
# BFS: expand from seed entities (depth 1)
|
|
if len(results) < limit:
|
|
remaining = limit - len(results)
|
|
bfs_sql = text(
|
|
"""
|
|
SELECT r.id, r.source_type, r.source_id, r.target_type, r.target_id,
|
|
r.relationship_type, r.metadata
|
|
FROM entity_relationships r
|
|
WHERE r.tenant_id = :tid
|
|
AND r.deleted_at IS NULL
|
|
AND (
|
|
(r.source_type = :etype AND r.source_id = ANY(:eids))
|
|
OR (r.target_type = :etype AND r.target_id = ANY(:eids))
|
|
)
|
|
AND r.id <> ALL(:exclude_ids)
|
|
LIMIT :lim
|
|
"""
|
|
)
|
|
for etype, eids in seed_entities.items():
|
|
if len(results) >= limit:
|
|
break
|
|
bfs_result = await db.execute(
|
|
bfs_sql,
|
|
{
|
|
"tid": tenant_id,
|
|
"etype": etype,
|
|
"eids": list(eids),
|
|
"exclude_ids": list(visited_rel_ids),
|
|
"lim": remaining,
|
|
},
|
|
)
|
|
bfs_rows = bfs_result.mappings().all()
|
|
for row in bfs_rows:
|
|
rid = str(row["id"])
|
|
if rid not in visited_rel_ids:
|
|
visited_rel_ids.add(rid)
|
|
results.append(self._relationship_to_result(dict(row)))
|
|
if len(results) >= limit:
|
|
break
|
|
|
|
return results[:limit]
|
|
|
|
def _relationship_to_result(self, row: dict[str, Any]) -> dict[str, Any]:
|
|
"""Convert a relationship row to a search result dict."""
|
|
return {
|
|
"id": str(row.get("id", "")),
|
|
"entity_type": self.entity_type,
|
|
"entity_id": str(row.get("id", "")),
|
|
"title": f"{row.get('source_type', '?')} --[{row.get('relationship_type', '?')}]--> {row.get('target_type', '?')}",
|
|
"snippet": str(row.get("metadata", {})),
|
|
"score": 0.0,
|
|
"data": {
|
|
"source_type": row.get("source_type"),
|
|
"source_id": str(row.get("source_id", "")),
|
|
"target_type": row.get("target_type"),
|
|
"target_id": str(row.get("target_id", "")),
|
|
"relationship_type": row.get("relationship_type"),
|
|
},
|
|
}
|
|
|
|
def to_search_result(self, entity: object) -> dict[str, Any]:
|
|
"""Convert relationship to search result dict."""
|
|
if isinstance(entity, dict):
|
|
return {
|
|
"entity_type": self.entity_type,
|
|
"entity_id": str(entity.get("id", "")),
|
|
"title": f"{entity.get('source_type', '?')} --[{entity.get('relationship_type', '?')}]--> {entity.get('target_type', '?')}",
|
|
"snippet": str(entity.get("metadata", {})),
|
|
"score": float(entity.get("rank", 0.0)),
|
|
"data": {
|
|
"source_type": entity.get("source_type"),
|
|
"source_id": str(entity.get("source_id", "")),
|
|
"target_type": entity.get("target_type"),
|
|
"target_id": str(entity.get("target_id", "")),
|
|
"relationship_type": entity.get("relationship_type"),
|
|
},
|
|
}
|
|
return {
|
|
"entity_type": self.entity_type,
|
|
"entity_id": str(getattr(entity, "id", "")),
|
|
"title": f"{getattr(entity, 'source_type', '?')} --[{getattr(entity, 'relationship_type', '?')}]--> {getattr(entity, 'target_type', '?')}",
|
|
"snippet": str(getattr(entity, "metadata", {})),
|
|
"score": 0.0,
|
|
"data": {},
|
|
}
|