library-desk/src/services/entity_linking_utils.py

"""
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
        }