""" Shared entity linking utilities for Library Desk. Provides bidirectional entity linking functionality that can be used by: - Consolidation service (knowledge consolidation) - Wiki router (smart page creation) - Any other service that creates wiki pages """ import logging from typing import Dict, Any, Optional from src.core.multi_tenancy import get_neo4j_user_base_label logger = logging.getLogger(__name__) async def apply_bidirectional_entity_linking( page_id: int, page_title: str, user: str, neo4j_client: "Neo4jClient", wiki_service: "WikiService", ingestion_service: Optional["IngestionService"] = None ) -> Dict[str, int]: """ Apply bidirectional entity linking after page creation/update. This runs AFTER ingestion so entities are extracted and in the graph. Steps: 1. Link entities in the new page (forward links to existing entities) 2. Find pages that mention the new entity (reverse references) 3. Link entities in those pages (backward links to the new entity) Args: page_id: Wiki page ID page_title: Page title (used to find reverse references) user: User identifier neo4j_client: Neo4j client for graph queries wiki_service: Wiki service for page operations ingestion_service: Optional ingestion service for re-indexing Returns: Dict with link counts: { "forward_links": int, # Links added to the new page "backward_links": int, # Links added to other pages pointing to new page "pages_updated": int # Number of other pages updated } """ from src.routers.entity_linking import ( link_entities_in_page, EntityLinkingRequest, get_entities_with_paths, add_entity_links_to_content ) from src.core.dependencies import get_graph_service, get_wiki_service, get_ingestion_service from src.models.wiki import WikiPageUpdate forward_links = 0 backward_links = 0 pages_updated = 0 try: graph_service = get_graph_service() # Use provided services or get defaults wiki_svc = wiki_service ingestion_svc = ingestion_service or get_ingestion_service() # STEP 1: Forward linking - link entities in the new page logger.info(f"Step 1/3: Linking entities in page {page_id} ('{page_title}')") try: forward_result = await link_entities_in_page( request=EntityLinkingRequest( user=user, page_id=page_id, create_relationships=True, re_index_if_changed=False # Already indexed, no need to re-index ), wiki_service=wiki_svc, graph_service=graph_service, ingestion_service=ingestion_svc, api_key="" # Internal call, no auth needed ) forward_links = forward_result.content_links_added logger.info(f"Added {forward_links} forward links in page {page_id}") except Exception as e: logger.error(f"Failed to add forward links: {e}") # STEP 2: Find reverse references - which pages mention this new entity? logger.info(f"Step 2/3: Finding pages that mention '{page_title}'") user_base_label = get_neo4j_user_base_label(user) # Query to find documents that mention entities with this page's title reverse_query = f""" // Find entities with the same name as the page title MATCH (e:{user_base_label}) WHERE toLower(e.name) = toLower($title) AND NOT e:Document // Find documents that mention those entities MATCH (d:Document)-[r:MENTIONS]->(e) WHERE d.page_id <> $page_id // Exclude the page itself RETURN DISTINCT d.page_id as page_id, d.title as title LIMIT 50 """ try: reverse_refs = await neo4j_client.execute_query( reverse_query, {"title": page_title, "page_id": page_id} ) logger.info(f"Found {len(reverse_refs)} pages that mention '{page_title}'") except Exception as e: logger.error(f"Failed to find reverse references: {e}") reverse_refs = [] # STEP 3: Backward linking - add links in those pages to the new entity if reverse_refs: logger.info(f"Step 3/3: Adding backward links in {len(reverse_refs)} pages") for ref in reverse_refs: try: backward_result = await link_entities_in_page( request=EntityLinkingRequest( user=user, page_id=ref['page_id'], create_relationships=False, # Relationships already exist re_index_if_changed=False # Don't re-index for link updates ), wiki_service=wiki_svc, graph_service=graph_service, ingestion_service=ingestion_svc, api_key="" ) if backward_result.content_links_added > 0: backward_links += backward_result.content_links_added pages_updated += 1 logger.info( f"Added {backward_result.content_links_added} links " f"in page {ref['page_id']} ('{ref['title']}')" ) except Exception as e: logger.error(f"Failed to add backward links in page {ref['page_id']}: {e}") else: logger.info("Step 3/3: No reverse references found, skipping backward linking") return { "forward_links": forward_links, "backward_links": backward_links, "pages_updated": pages_updated } except Exception as e: logger.error(f"Bidirectional entity linking failed: {e}", exc_info=True) return { "forward_links": 0, "backward_links": 0, "pages_updated": 0 }