feat(library-desk): add entity linking endpoint
- Add /entity-linking/link-page endpoint to find and link entity mentions - Creates both MENTIONS relationships in Neo4j and hyperlinks in wiki content - Supports automatic re-indexing after linking - Returns detailed statistics on entities found and linked - Protects existing markdown links from being nested - Idempotent: safe to run multiple times Implements dual entity linking: 1. Graph relationships (MENTIONS) for knowledge graph traversal 2. Wiki content hyperlinks for user navigation
This commit is contained in:
@@ -0,0 +1,354 @@
|
||||
"""
|
||||
Entity Linking Router
|
||||
|
||||
Finds and links mentions of existing entities in wiki pages to the knowledge graph.
|
||||
Creates both:
|
||||
1. Graph relationships (MENTIONS) in Neo4j
|
||||
2. Hyperlinks in wiki page content (markdown links)
|
||||
"""
|
||||
import logging
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
import re
|
||||
|
||||
from src.core.dependencies import (
|
||||
get_wiki_service,
|
||||
get_graph_service,
|
||||
get_ingestion_service,
|
||||
verify_api_key
|
||||
)
|
||||
from src.services.wiki_service import WikiService
|
||||
from src.services.graph_service import GraphService
|
||||
from src.services.ingestion_service import IngestionService
|
||||
from src.models.wiki import WikiPageUpdate
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/entity-linking", tags=["Entity Linking"])
|
||||
|
||||
|
||||
class EntityLinkingResult(BaseModel):
|
||||
"""Result of entity linking operation."""
|
||||
page_id: int
|
||||
page_title: str
|
||||
entities_found: List[Dict[str, Any]]
|
||||
new_graph_links_created: int # New MENTIONS relationships in Neo4j
|
||||
total_graph_links: int
|
||||
content_links_added: int # New hyperlinks added to wiki page content
|
||||
content_updated: bool # Whether wiki page content was modified
|
||||
re_indexed: bool
|
||||
processing_time_ms: float
|
||||
|
||||
|
||||
class EntityLinkingRequest(BaseModel):
|
||||
"""Request for entity linking."""
|
||||
user: str
|
||||
page_id: int
|
||||
create_relationships: bool = True
|
||||
re_index_if_changed: bool = True
|
||||
|
||||
|
||||
@router.post("/link-page", response_model=EntityLinkingResult)
|
||||
async def link_entities_in_page(
|
||||
request: EntityLinkingRequest,
|
||||
wiki_service: WikiService = Depends(get_wiki_service),
|
||||
graph_service: GraphService = Depends(get_graph_service),
|
||||
ingestion_service: IngestionService = Depends(get_ingestion_service),
|
||||
api_key: str = Depends(verify_api_key)
|
||||
) -> EntityLinkingResult:
|
||||
"""
|
||||
Find and link entities mentioned in a wiki page.
|
||||
|
||||
This endpoint:
|
||||
1. Retrieves the page content from Wiki.js
|
||||
2. Fetches all existing entities from the knowledge graph
|
||||
3. Finds mentions of those entities in the page text
|
||||
4. Creates MENTIONS relationships in Neo4j
|
||||
5. Optionally re-indexes the page if new links were created
|
||||
|
||||
Args:
|
||||
request: EntityLinkingRequest with page_id and user
|
||||
|
||||
Returns:
|
||||
EntityLinkingResult with statistics about linked entities
|
||||
"""
|
||||
import time
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
logger.info(f"Entity linking for page {request.page_id} (user: {request.user})")
|
||||
|
||||
# Step 1: Get page content
|
||||
page = await wiki_service.get_page(request.page_id, request.user)
|
||||
if not page:
|
||||
raise HTTPException(status_code=404, detail=f"Page {request.page_id} not found")
|
||||
|
||||
page_title = page.title
|
||||
page_content = page.content
|
||||
original_content = page_content # Keep original for comparison
|
||||
|
||||
logger.info(f"Processing page: {page_title}")
|
||||
|
||||
# Step 2: Get all entities from knowledge graph with their wiki page paths
|
||||
entities = await get_entities_with_paths(graph_service, request.user)
|
||||
logger.info(f"Found {len(entities)} existing entities in graph")
|
||||
|
||||
# Step 3: Find entity mentions in page content
|
||||
found_entities = find_entity_mentions(page_content, entities)
|
||||
logger.info(f"Found {len(found_entities)} entity mentions in page")
|
||||
|
||||
# Step 4: Add hyperlinks to wiki content for entities with pages
|
||||
content_links_added = 0
|
||||
content_updated = False
|
||||
|
||||
if found_entities:
|
||||
updated_content, content_links_added = add_entity_links_to_content(
|
||||
page_content,
|
||||
found_entities
|
||||
)
|
||||
|
||||
if updated_content != original_content:
|
||||
content_updated = True
|
||||
logger.info(f"Added {content_links_added} hyperlinks to page content")
|
||||
|
||||
# Update the wiki page (preserve existing title, description, tags)
|
||||
try:
|
||||
update_data = WikiPageUpdate(
|
||||
content=updated_content,
|
||||
title=page.title,
|
||||
description=page.description if hasattr(page, 'description') else None,
|
||||
tags=page.tags if hasattr(page, 'tags') else None
|
||||
)
|
||||
await wiki_service.update_page(
|
||||
page_id=request.page_id,
|
||||
page_data=update_data,
|
||||
user=request.user
|
||||
)
|
||||
logger.info(f"Updated wiki page {request.page_id} with entity links")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update wiki page content: {e}")
|
||||
# Continue anyway - graph links can still be created
|
||||
|
||||
# Step 5: Create graph relationships if requested
|
||||
new_graph_links = 0
|
||||
total_graph_links = 0
|
||||
|
||||
if request.create_relationships and found_entities:
|
||||
new_graph_links = await graph_service.create_entity_mentions(
|
||||
page_id=request.page_id,
|
||||
user=request.user,
|
||||
entity_names=found_entities
|
||||
)
|
||||
total_graph_links = len(found_entities)
|
||||
logger.info(f"Created {new_graph_links} new MENTIONS relationships")
|
||||
|
||||
# Step 6: Re-index if changes were made (content or graph)
|
||||
re_indexed = False
|
||||
if request.re_index_if_changed and (content_updated or new_graph_links > 0):
|
||||
logger.info(f"Re-indexing page {request.page_id} due to entity linking changes")
|
||||
try:
|
||||
await ingestion_service.ingest_page(request.page_id, request.user)
|
||||
re_indexed = True
|
||||
except Exception as e:
|
||||
logger.error(f"Re-indexing failed: {e}")
|
||||
# Don't fail the whole operation if re-indexing fails
|
||||
|
||||
processing_time = (time.time() - start_time) * 1000
|
||||
|
||||
return EntityLinkingResult(
|
||||
page_id=request.page_id,
|
||||
page_title=page_title,
|
||||
entities_found=found_entities,
|
||||
new_graph_links_created=new_graph_links,
|
||||
total_graph_links=total_graph_links,
|
||||
content_links_added=content_links_added,
|
||||
content_updated=content_updated,
|
||||
re_indexed=re_indexed,
|
||||
processing_time_ms=processing_time
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Entity linking failed: {e}", exc_info=True)
|
||||
raise HTTPException(status_code=500, detail=f"Entity linking failed: {str(e)}")
|
||||
|
||||
|
||||
def find_entity_mentions(content: str, entities: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Find mentions of entities in the page content.
|
||||
|
||||
Uses case-insensitive regex matching to find entity names in the text.
|
||||
|
||||
Args:
|
||||
content: Page content to search
|
||||
entities: List of entities with 'name', 'type', and optionally 'path' fields
|
||||
|
||||
Returns:
|
||||
List of found entities with additional 'mentions' field
|
||||
"""
|
||||
found = []
|
||||
content_lower = content.lower()
|
||||
|
||||
for entity in entities:
|
||||
entity_name = entity.get("name", "")
|
||||
if not entity_name or len(entity_name) < 3:
|
||||
continue
|
||||
|
||||
# Create regex pattern for whole word matching
|
||||
# This avoids matching "John" in "Johnson"
|
||||
pattern = r'\b' + re.escape(entity_name.lower()) + r'\b'
|
||||
|
||||
# Find all matches
|
||||
matches = list(re.finditer(pattern, content_lower))
|
||||
|
||||
if matches:
|
||||
found.append({
|
||||
"name": entity_name,
|
||||
"type": entity.get("type", "unknown"),
|
||||
"mentions": len(matches),
|
||||
"entity_id": entity.get("id"),
|
||||
"path": entity.get("path") # Include path if available
|
||||
})
|
||||
|
||||
# Sort by number of mentions (descending)
|
||||
found.sort(key=lambda x: x["mentions"], reverse=True)
|
||||
|
||||
return found
|
||||
|
||||
|
||||
async def get_entities_with_paths(graph_service: GraphService, user: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get all entities and check which ones have corresponding wiki pages.
|
||||
|
||||
Returns entities with their names, types, and paths (if they have wiki pages).
|
||||
"""
|
||||
from src.core.multi_tenancy import get_neo4j_user_base_label
|
||||
|
||||
user_base_label = get_neo4j_user_base_label(user)
|
||||
|
||||
# Query to get entities and optionally join with Document nodes by matching title
|
||||
query = f"""
|
||||
// Get all entities
|
||||
MATCH (e:{user_base_label})
|
||||
WHERE NOT e:Document
|
||||
|
||||
// Optionally match Document with same name (case-insensitive)
|
||||
OPTIONAL MATCH (d:Document)
|
||||
WHERE toLower(d.title) = toLower(e.name)
|
||||
|
||||
RETURN e.name as name,
|
||||
e.type as type,
|
||||
e.id as id,
|
||||
d.path as path,
|
||||
d.page_id as doc_page_id
|
||||
"""
|
||||
|
||||
try:
|
||||
results = await graph_service.neo4j.execute_query(query, {})
|
||||
return [
|
||||
{
|
||||
"name": r["name"],
|
||||
"type": r.get("type", "unknown"),
|
||||
"id": r.get("id"),
|
||||
"path": r.get("path"), # Will be None if no matching Document
|
||||
"doc_page_id": r.get("doc_page_id")
|
||||
}
|
||||
for r in results
|
||||
]
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get entities with paths: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def add_entity_links_to_content(
|
||||
content: str,
|
||||
entities_with_paths: List[Dict[str, Any]]
|
||||
) -> Tuple[str, int]:
|
||||
"""
|
||||
Add markdown hyperlinks for entities in the content.
|
||||
|
||||
Only links entities that:
|
||||
1. Have a wiki page (path is not None)
|
||||
2. Are not already inside existing links
|
||||
3. Are not already linked in the content
|
||||
|
||||
Returns:
|
||||
Tuple of (updated_content, links_added_count)
|
||||
"""
|
||||
if not entities_with_paths:
|
||||
return content, 0
|
||||
|
||||
# Filter to only entities with paths
|
||||
linkable_entities = [e for e in entities_with_paths if e.get("path")]
|
||||
if not linkable_entities:
|
||||
return content, 0
|
||||
|
||||
# Sort by length (longest first) to avoid partial replacements
|
||||
# e.g., "Machine Learning" before "Machine"
|
||||
linkable_entities.sort(key=lambda x: len(x["name"]), reverse=True)
|
||||
|
||||
# Find all existing markdown links to protect them
|
||||
link_pattern = r'\[([^\]]+)\]\([^\)]+\)'
|
||||
existing_links = list(re.finditer(link_pattern, content))
|
||||
|
||||
# Create a list of protected ranges (start, end) for ENTIRE links (text + URL)
|
||||
# This prevents linking entity names inside existing link URLs
|
||||
protected_ranges = [(m.start(), m.end()) for m in existing_links]
|
||||
|
||||
updated_content = content
|
||||
links_added = 0
|
||||
|
||||
for entity in linkable_entities:
|
||||
entity_name = entity["name"]
|
||||
entity_path = entity["path"]
|
||||
|
||||
# Skip short entity names to avoid false positives
|
||||
if len(entity_name) < 3:
|
||||
continue
|
||||
|
||||
# Create markdown link
|
||||
# We'll link to the relative path without the "users/jpmschweitzer/" prefix
|
||||
clean_path = entity_path.replace("users/jpmschweitzer/", "")
|
||||
markdown_link = f"[{entity_name}](/{clean_path})"
|
||||
|
||||
# Find all potential matches
|
||||
pattern = r'\b(' + re.escape(entity_name) + r')\b'
|
||||
matches = list(re.finditer(pattern, updated_content, flags=re.IGNORECASE))
|
||||
|
||||
# Filter out matches that are inside existing links
|
||||
valid_matches = []
|
||||
for match in matches:
|
||||
match_start = match.start()
|
||||
match_end = match.end()
|
||||
|
||||
# Check if this match is inside any protected range
|
||||
inside_link = False
|
||||
for prot_start, prot_end in protected_ranges:
|
||||
if prot_start <= match_start < prot_end or prot_start < match_end <= prot_end:
|
||||
inside_link = True
|
||||
break
|
||||
|
||||
# Also check if already a link (pattern like [entity_name](...))
|
||||
if match_end < len(updated_content) - 1:
|
||||
next_chars = updated_content[match_end:match_end+2]
|
||||
if next_chars == '](':
|
||||
inside_link = True
|
||||
|
||||
if not inside_link:
|
||||
valid_matches.append(match)
|
||||
|
||||
if not valid_matches:
|
||||
continue
|
||||
|
||||
# Replace valid matches in reverse order (to preserve positions)
|
||||
for match in reversed(valid_matches):
|
||||
updated_content = (
|
||||
updated_content[:match.start()] +
|
||||
markdown_link +
|
||||
updated_content[match.end():]
|
||||
)
|
||||
links_added += 1
|
||||
|
||||
return updated_content, links_added
|
||||
Reference in New Issue
Block a user