feat(library-desk): add entity management to graph and ingestion services
Graph Service: - Add get_all_entities() to retrieve entities with wiki page paths - Add create_entity_mentions() for MENTIONS relationship creation - Support entity-to-document linking via title matching Ingestion Service: - Add _link_existing_entities() for automatic entity linking - Auto-link entities during page ingestion - Support skip_entity_linking parameter for granular control
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,415 @@
|
||||
"""
|
||||
Document Ingestion Service
|
||||
|
||||
Orchestrates the ingestion of wiki pages into the knowledge base:
|
||||
1. Fetches page content from Wiki.js
|
||||
2. Generates vector embeddings (Qdrant)
|
||||
3. Extracts entities and updates knowledge graph (Neo4j)
|
||||
|
||||
This service is called by:
|
||||
- Consolidation service (after creating/updating pages)
|
||||
- Manual ingestion endpoints
|
||||
- Batch ingestion jobs
|
||||
"""
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import List, Optional
|
||||
from datetime import datetime
|
||||
import time
|
||||
|
||||
from src.services.vector_service import VectorService
|
||||
from src.services.graph_service import GraphService
|
||||
from src.clients.wikijs_client import WikiJSClient
|
||||
from src.models.ingestion import (
|
||||
IngestionRequest,
|
||||
IngestionResult,
|
||||
BatchIngestionRequest,
|
||||
BatchIngestionResult
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class IngestionService:
|
||||
"""
|
||||
Service for ingesting wiki pages into the knowledge base.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vector_service: VectorService,
|
||||
graph_service: GraphService,
|
||||
wiki_client: WikiJSClient
|
||||
):
|
||||
self.vector = vector_service
|
||||
self.graph = graph_service
|
||||
self.wiki = wiki_client
|
||||
|
||||
async def ingest_page(
|
||||
self,
|
||||
page_id: int,
|
||||
user: str,
|
||||
force_refresh: bool = False,
|
||||
skip_vectors: bool = False,
|
||||
skip_graph: bool = False,
|
||||
skip_entity_linking: bool = False
|
||||
) -> IngestionResult:
|
||||
"""
|
||||
Ingest a single wiki page into the knowledge base.
|
||||
|
||||
Args:
|
||||
page_id: Wiki page ID
|
||||
user: User identifier
|
||||
force_refresh: Force re-ingestion even if unchanged
|
||||
skip_vectors: Skip vector embedding generation
|
||||
skip_graph: Skip graph entity extraction
|
||||
skip_entity_linking: Skip automatic entity linking
|
||||
|
||||
Returns:
|
||||
IngestionResult with operation details
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
logger.info(f"Starting ingestion for page {page_id} (user: {user})")
|
||||
|
||||
try:
|
||||
# Fetch page to get metadata
|
||||
page = await self.wiki.get_page(page_id)
|
||||
if not page:
|
||||
return IngestionResult(
|
||||
page_id=page_id,
|
||||
page_title=f"Page {page_id}",
|
||||
success=False,
|
||||
error="Page not found in Wiki.js",
|
||||
processing_time_ms=(time.time() - start_time) * 1000
|
||||
)
|
||||
|
||||
page_title = page.get("title", f"Page {page_id}")
|
||||
page_path = page.get("path", "")
|
||||
|
||||
# Ingest vectors and graph in parallel
|
||||
tasks = []
|
||||
|
||||
if not skip_vectors:
|
||||
tasks.append(self._ingest_vectors(page_id, user, force_refresh))
|
||||
else:
|
||||
tasks.append(asyncio.create_task(asyncio.sleep(0))) # Dummy task
|
||||
|
||||
if not skip_graph:
|
||||
tasks.append(self._ingest_graph(page_id, user, force_refresh))
|
||||
else:
|
||||
tasks.append(asyncio.create_task(asyncio.sleep(0))) # Dummy task
|
||||
|
||||
# Execute in parallel
|
||||
vector_result, graph_result = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
# Handle errors
|
||||
vector_chunks = 0
|
||||
graph_entities = 0
|
||||
graph_relationships = 0
|
||||
errors = []
|
||||
|
||||
if not skip_vectors:
|
||||
if isinstance(vector_result, Exception):
|
||||
errors.append(f"Vector ingestion failed: {str(vector_result)}")
|
||||
logger.error(f"Vector ingestion failed for page {page_id}: {vector_result}")
|
||||
else:
|
||||
vector_chunks = vector_result.get("chunks_created", 0)
|
||||
|
||||
if not skip_graph:
|
||||
if isinstance(graph_result, Exception):
|
||||
errors.append(f"Graph ingestion failed: {str(graph_result)}")
|
||||
logger.error(f"Graph ingestion failed for page {page_id}: {graph_result}")
|
||||
else:
|
||||
# entities_extracted is a list, get its length
|
||||
entities_list = graph_result.get("entities_extracted", [])
|
||||
graph_entities = len(entities_list) if isinstance(entities_list, list) else 0
|
||||
graph_relationships = graph_result.get("relationships_created", 0)
|
||||
|
||||
# Step 3: Link existing entities in the page content (after graph extraction)
|
||||
entity_links_created = 0
|
||||
if not skip_entity_linking and not skip_graph and not isinstance(graph_result, Exception):
|
||||
try:
|
||||
entity_links_created = await self._link_existing_entities(page_id, user, page)
|
||||
logger.info(f"Created {entity_links_created} entity mention links for page {page_id}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Entity linking failed for page {page_id}: {e}")
|
||||
# Don't fail the whole ingestion if entity linking fails
|
||||
|
||||
processing_time_ms = (time.time() - start_time) * 1000
|
||||
|
||||
result = IngestionResult(
|
||||
page_id=page_id,
|
||||
page_title=page_title,
|
||||
page_path=page_path,
|
||||
success=len(errors) == 0,
|
||||
error="; ".join(errors) if errors else None,
|
||||
vector_chunks_created=vector_chunks,
|
||||
graph_entities_extracted=graph_entities,
|
||||
graph_relationships_created=graph_relationships,
|
||||
processing_time_ms=processing_time_ms
|
||||
)
|
||||
|
||||
if result.success:
|
||||
logger.info(
|
||||
f"Successfully ingested page {page_id}: "
|
||||
f"{vector_chunks} chunks, {graph_entities} entities, "
|
||||
f"{graph_relationships} relationships, {entity_links_created} entity links "
|
||||
f"in {processing_time_ms:.0f}ms"
|
||||
)
|
||||
else:
|
||||
logger.warning(f"Partial ingestion failure for page {page_id}: {result.error}")
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ingestion failed for page {page_id}: {e}", exc_info=True)
|
||||
return IngestionResult(
|
||||
page_id=page_id,
|
||||
page_title=f"Page {page_id}",
|
||||
success=False,
|
||||
error=str(e),
|
||||
processing_time_ms=(time.time() - start_time) * 1000
|
||||
)
|
||||
|
||||
async def _ingest_vectors(
|
||||
self,
|
||||
page_id: int,
|
||||
user: str,
|
||||
force_refresh: bool
|
||||
) -> dict:
|
||||
"""
|
||||
Ingest page into vector database.
|
||||
|
||||
Returns:
|
||||
Dict with chunks_created count
|
||||
"""
|
||||
try:
|
||||
summary = await self.vector.update_from_page(
|
||||
page_id=page_id,
|
||||
user=user,
|
||||
force_refresh=force_refresh
|
||||
)
|
||||
|
||||
return {
|
||||
"chunks_created": summary.chunks_created,
|
||||
"chunks_deleted": summary.chunks_deleted
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Vector ingestion failed for page {page_id}: {e}")
|
||||
raise
|
||||
|
||||
async def _ingest_graph(
|
||||
self,
|
||||
page_id: int,
|
||||
user: str,
|
||||
force_refresh: bool
|
||||
) -> dict:
|
||||
"""
|
||||
Ingest page into knowledge graph.
|
||||
|
||||
Returns:
|
||||
Dict with entities_extracted and relationships_created counts
|
||||
"""
|
||||
try:
|
||||
summary = await self.graph.update_from_page(
|
||||
page_id=page_id,
|
||||
user=user,
|
||||
force_refresh=force_refresh
|
||||
)
|
||||
|
||||
return {
|
||||
"entities_extracted": summary.entities_extracted,
|
||||
"relationships_created": summary.relationships_created
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Graph ingestion failed for page {page_id}: {e}")
|
||||
raise
|
||||
|
||||
async def _link_existing_entities(
|
||||
self,
|
||||
page_id: int,
|
||||
user: str,
|
||||
page: dict
|
||||
) -> int:
|
||||
"""
|
||||
Find and link mentions of existing entities in the page content.
|
||||
|
||||
This runs automatically after graph extraction to create MENTIONS relationships
|
||||
for entities that already exist in the knowledge graph but were mentioned in
|
||||
this page.
|
||||
|
||||
Args:
|
||||
page_id: Wiki page ID
|
||||
user: User identifier
|
||||
page: Page dict with content (from WikiJSClient)
|
||||
|
||||
Returns:
|
||||
Number of new entity mention links created
|
||||
"""
|
||||
import re
|
||||
|
||||
try:
|
||||
page_content = page.get("content", "")
|
||||
if not page_content or len(page_content) < 10:
|
||||
return 0
|
||||
|
||||
# Get all existing entities from the knowledge graph
|
||||
entities = await self.graph.get_all_entities(user)
|
||||
if not entities:
|
||||
logger.debug(f"No existing entities found for user {user}, skipping entity linking")
|
||||
return 0
|
||||
|
||||
# Find entity mentions in page content
|
||||
found_entities = []
|
||||
content_lower = page_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_entities.append({
|
||||
"name": entity_name,
|
||||
"type": entity.get("type", "unknown"),
|
||||
"mentions": len(matches),
|
||||
"entity_id": entity.get("id")
|
||||
})
|
||||
|
||||
if not found_entities:
|
||||
logger.debug(f"No entity mentions found in page {page_id}")
|
||||
return 0
|
||||
|
||||
# Create MENTIONS relationships
|
||||
new_links_created = await self.graph.create_entity_mentions(
|
||||
page_id=page_id,
|
||||
user=user,
|
||||
entity_names=found_entities
|
||||
)
|
||||
|
||||
return new_links_created
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Entity linking failed for page {page_id}: {e}")
|
||||
raise
|
||||
|
||||
async def ingest_batch(
|
||||
self,
|
||||
page_ids: List[int],
|
||||
user: str,
|
||||
force_refresh: bool = False,
|
||||
skip_vectors: bool = False,
|
||||
skip_graph: bool = False,
|
||||
max_concurrent: int = 3
|
||||
) -> BatchIngestionResult:
|
||||
"""
|
||||
Ingest multiple wiki pages concurrently.
|
||||
|
||||
Args:
|
||||
page_ids: List of wiki page IDs to ingest
|
||||
user: User identifier
|
||||
force_refresh: Force re-ingestion
|
||||
skip_vectors: Skip vector embedding generation
|
||||
skip_graph: Skip graph entity extraction
|
||||
max_concurrent: Maximum concurrent ingestion tasks
|
||||
|
||||
Returns:
|
||||
BatchIngestionResult with per-page results
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
logger.info(f"Starting batch ingestion of {len(page_ids)} pages (user: {user})")
|
||||
|
||||
results = []
|
||||
semaphore = asyncio.Semaphore(max_concurrent)
|
||||
|
||||
async def ingest_with_semaphore(page_id: int):
|
||||
async with semaphore:
|
||||
return await self.ingest_page(
|
||||
page_id=page_id,
|
||||
user=user,
|
||||
force_refresh=force_refresh,
|
||||
skip_vectors=skip_vectors,
|
||||
skip_graph=skip_graph
|
||||
)
|
||||
|
||||
# Execute all ingestions with concurrency control
|
||||
tasks = [ingest_with_semaphore(page_id) for page_id in page_ids]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
# Calculate summary
|
||||
successful = sum(1 for r in results if r.success)
|
||||
failed = len(results) - successful
|
||||
total_processing_time_ms = (time.time() - start_time) * 1000
|
||||
|
||||
batch_result = BatchIngestionResult(
|
||||
total_pages=len(page_ids),
|
||||
successful=successful,
|
||||
failed=failed,
|
||||
results=results,
|
||||
total_processing_time_ms=total_processing_time_ms
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Batch ingestion complete: {successful}/{len(page_ids)} successful "
|
||||
f"in {total_processing_time_ms:.0f}ms"
|
||||
)
|
||||
|
||||
return batch_result
|
||||
|
||||
async def ingest_all_pages(
|
||||
self,
|
||||
user: str,
|
||||
path_prefix: Optional[str] = None,
|
||||
force_refresh: bool = False,
|
||||
max_concurrent: int = 3
|
||||
) -> BatchIngestionResult:
|
||||
"""
|
||||
Ingest all wiki pages for a user.
|
||||
|
||||
Args:
|
||||
user: User identifier
|
||||
path_prefix: Optional path prefix filter (e.g., "users/jpmschweitzer")
|
||||
force_refresh: Force re-ingestion
|
||||
max_concurrent: Maximum concurrent ingestion tasks
|
||||
|
||||
Returns:
|
||||
BatchIngestionResult
|
||||
"""
|
||||
logger.info(f"Finding all pages for user {user} (prefix: {path_prefix or 'all'})")
|
||||
|
||||
# Search for all pages
|
||||
pages = await self.wiki.search_pages(
|
||||
query="", # Empty query returns all pages
|
||||
path_prefix=path_prefix or f"users/{user}"
|
||||
)
|
||||
|
||||
if not pages:
|
||||
logger.warning(f"No pages found for user {user}")
|
||||
return BatchIngestionResult(
|
||||
total_pages=0,
|
||||
successful=0,
|
||||
failed=0,
|
||||
results=[],
|
||||
total_processing_time_ms=0
|
||||
)
|
||||
|
||||
page_ids = [p['id'] for p in pages]
|
||||
logger.info(f"Found {len(page_ids)} pages to ingest")
|
||||
|
||||
return await self.ingest_batch(
|
||||
page_ids=page_ids,
|
||||
user=user,
|
||||
force_refresh=force_refresh,
|
||||
max_concurrent=max_concurrent
|
||||
)
|
||||
Reference in New Issue
Block a user