feat: implement Elasticsearch connector with indexing capabilities and validation

This commit is contained in:
Anish Sarkar 2025-10-04 13:47:59 +05:30
parent 0b32b66888
commit 2ad1c43059
8 changed files with 522 additions and 212 deletions

View file

@ -1,4 +1,3 @@
# import asyncio
import json
import logging
from collections.abc import AsyncGenerator
@ -7,18 +6,27 @@ from typing import Any
from elasticsearch import AsyncElasticsearch
from elasticsearch.exceptions import AuthenticationException, ConnectionError
from app.schemas.connector import ConnectorConfig
from app.schemas.documents import DocumentCreate, DocumentType
from app.utils.document_converter import DocumentConverter
from app.db import DocumentType
logger = logging.getLogger(__name__)
# Create a simple document structure for Elasticsearch results
class ElasticsearchDocument:
def __init__(
self, title: str, content: str, url: str, metadata: dict | None = None
):
self.title = title
self.content = content
self.url = url
self.metadata = metadata or {}
self.document_type = DocumentType.ELASTICSEARCH_CONNECTOR
class ElasticsearchConnector:
def __init__(self, config: ConnectorConfig):
def __init__(self, config: dict):
self.config = config
self.client: AsyncElasticsearch | None = None
self.document_converter = DocumentConverter()
async def connect(self) -> bool:
"""Establish connection to Elasticsearch"""
@ -106,8 +114,8 @@ class ElasticsearchConnector:
indices: list[str] | None = None,
size: int = 100,
fields: list[str] | None = None,
) -> AsyncGenerator[DocumentCreate, None]:
"""Search documents in Elasticsearch and yield them as DocumentCreate objects"""
) -> AsyncGenerator[ElasticsearchDocument, None]:
"""Search documents in Elasticsearch and yield them as ElasticsearchDocument objects"""
try:
if not self.client:
await self.connect()
@ -137,7 +145,7 @@ class ElasticsearchConnector:
for hit in response["hits"]["hits"]:
try:
# Convert Elasticsearch document to SurfSense document
# Convert Elasticsearch document to ElasticsearchDocument
document = await self._convert_es_document(hit)
if document:
yield document
@ -150,7 +158,7 @@ class ElasticsearchConnector:
async def get_document_by_id(
self, index: str, doc_id: str
) -> DocumentCreate | None:
) -> ElasticsearchDocument | None:
"""Get a specific document by ID"""
try:
if not self.client:
@ -164,8 +172,8 @@ class ElasticsearchConnector:
async def _convert_es_document(
self, es_doc: dict[str, Any]
) -> DocumentCreate | None:
"""Convert Elasticsearch document to SurfSense DocumentCreate"""
) -> ElasticsearchDocument | None:
"""Convert Elasticsearch document to ElasticsearchDocument"""
try:
source = es_doc.get("_source", {})
@ -190,27 +198,25 @@ class ElasticsearchConnector:
# Generate URL pointing back to Elasticsearch
url = f"elasticsearch://{es_doc.get('_index')}/{es_doc.get('_id')}"
# Convert content to markdown if it's HTML
if content and isinstance(content, str):
content = await self.document_converter.convert_to_markdown(content)
# Create metadata
metadata = {
"elasticsearch_index": es_doc.get("_index"),
"elasticsearch_id": es_doc.get("_id"),
"elasticsearch_score": es_doc.get("_score"),
"source_fields": list(source.keys()),
**{
k: v
for k, v in source.items()
if k not in ["content", "text", "body", "message"]
and isinstance(v, str | int | float | bool)
},
}
return DocumentCreate(
return ElasticsearchDocument(
title=title[:500], # Limit title length
content=content,
url=url,
document_type=DocumentType.ELASTICSEARCH,
metadata={
"elasticsearch_index": es_doc.get("_index"),
"elasticsearch_id": es_doc.get("_id"),
"elasticsearch_score": es_doc.get("_score"),
"source_fields": list(source.keys()),
**{
k: v
for k, v in source.items()
if k not in ["content", "text", "body", "message"]
and isinstance(v, str | int | float | bool)
},
},
metadata=metadata,
)
except Exception as e:
logger.error(f"Error converting Elasticsearch document: {e}")

View file

@ -50,7 +50,7 @@ class DocumentType(str, Enum):
GOOGLE_GMAIL_CONNECTOR = "GOOGLE_GMAIL_CONNECTOR"
AIRTABLE_CONNECTOR = "AIRTABLE_CONNECTOR"
LUMA_CONNECTOR = "LUMA_CONNECTOR"
ELASTICSEARCH = "ELASTICSEARCH"
ELASTICSEARCH_CONNECTOR = "ELASTICSEARCH_CONNECTOR"
class SearchSourceConnectorType(str, Enum):
@ -69,7 +69,7 @@ class SearchSourceConnectorType(str, Enum):
GOOGLE_GMAIL_CONNECTOR = "GOOGLE_GMAIL_CONNECTOR"
AIRTABLE_CONNECTOR = "AIRTABLE_CONNECTOR"
LUMA_CONNECTOR = "LUMA_CONNECTOR"
ELASTICSEARCH = "ELASTICSEARCH"
ELASTICSEARCH_CONNECTOR = "ELASTICSEARCH_CONNECTOR"
class ChatType(str, Enum):

View file

@ -1,25 +1,20 @@
import hashlib
import contextlib
import logging
from datetime import UTC, datetime
from typing import Any
from fastapi import APIRouter, Depends, HTTPException, status
from pydantic import BaseModel
from sqlalchemy import select, update
from sqlalchemy import select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from app.config import config
from app.connectors.elasticsearch_connector import ElasticsearchConnector
from app.db import (
Chunk,
Document,
SearchSourceConnector,
SearchSourceConnectorType,
User,
get_async_session,
)
from app.schemas.users import User
from app.tasks.connector_indexers.elasticsearch_indexer import ElasticsearchIndexer
from app.users import current_active_user
logger = logging.getLogger(__name__)
@ -48,10 +43,8 @@ async def add_elasticsearch_connector(
) -> dict[str, Any]:
"""Add a new Elasticsearch connector for the current user"""
# Test the connection before saving
elasticsearch_connector = ElasticsearchConnector(connector_data)
try:
elasticsearch_connector = ElasticsearchConnector(connector_data)
connection_test = await elasticsearch_connector.test_connection()
if not connection_test.get("success"):
raise HTTPException(
@ -59,13 +52,14 @@ async def add_elasticsearch_connector(
detail=f"Failed to connect to Elasticsearch: {connection_test.get('error')}",
)
except Exception as e:
logger.error(f"Error testing Elasticsearch connection: {e}")
logger.error(f"Error initializing or testing Elasticsearch connector: {e}")
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=f"Failed to test Elasticsearch connection: {e!s}",
detail=f"Failed to initialize or test Elasticsearch connector: {e!s}",
) from e
finally:
await elasticsearch_connector.disconnect()
with contextlib.suppress(Exception):
await elasticsearch_connector.disconnect()
try:
# Check if connector already exists for this user
@ -73,7 +67,7 @@ async def add_elasticsearch_connector(
select(SearchSourceConnector).filter(
SearchSourceConnector.user_id == current_user.id,
SearchSourceConnector.connector_type
== SearchSourceConnectorType.ELASTICSEARCH,
== SearchSourceConnectorType.ELASTICSEARCH_CONNECTOR,
)
)
existing_connector = existing_connector_result.scalars().first()
@ -92,7 +86,7 @@ async def add_elasticsearch_connector(
# Create new connector
new_connector = SearchSourceConnector(
name=f"Elasticsearch - {connector_data.hostname}:{connector_data.port}",
connector_type=SearchSourceConnectorType.ELASTICSEARCH,
connector_type=SearchSourceConnectorType.ELASTICSEARCH_CONNECTOR,
is_indexable=True,
config=connector_data.model_dump(),
user_id=current_user.id,
@ -159,7 +153,7 @@ async def get_elasticsearch_connectors(
select(SearchSourceConnector).filter(
SearchSourceConnector.user_id == current_user.id,
SearchSourceConnector.connector_type
== SearchSourceConnectorType.ELASTICSEARCH,
== SearchSourceConnectorType.ELASTICSEARCH_CONNECTOR,
)
)
connectors = result.scalars().all()
@ -213,7 +207,7 @@ async def update_elasticsearch_connector(
SearchSourceConnector.id == connector_id,
SearchSourceConnector.user_id == current_user.id,
SearchSourceConnector.connector_type
== SearchSourceConnectorType.ELASTICSEARCH,
== SearchSourceConnectorType.ELASTICSEARCH_CONNECTOR,
)
)
connector = result.scalar_one_or_none()
@ -224,18 +218,24 @@ async def update_elasticsearch_connector(
detail="Elasticsearch connector not found",
)
# Test the new connection
elasticsearch_connector = ElasticsearchConnector(connector_data)
# Test the new connection with robust error handling
try:
elasticsearch_connector = ElasticsearchConnector(connector_data)
connection_test = await elasticsearch_connector.test_connection()
if not connection_test.get("success"):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=f"Failed to connect to Elasticsearch: {connection_test.get('error')}",
)
except Exception as e:
logger.error(f"Error initializing or testing Elasticsearch connector: {e}")
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=f"Failed to initialize or test Elasticsearch connector: {e!s}",
) from e
finally:
await elasticsearch_connector.disconnect()
with contextlib.suppress(Exception):
await elasticsearch_connector.disconnect()
# Update the connector
connector.config = connector_data.model_dump()
@ -289,7 +289,7 @@ async def delete_elasticsearch_connector(
SearchSourceConnector.id == connector_id,
SearchSourceConnector.user_id == current_user.id,
SearchSourceConnector.connector_type
== SearchSourceConnectorType.ELASTICSEARCH,
== SearchSourceConnectorType.ELASTICSEARCH_CONNECTOR,
)
)
connector = result.scalar_one_or_none()
@ -389,138 +389,4 @@ async def get_elasticsearch_indices(
await elasticsearch_connector.disconnect()
@router.post("/index-connector/{connector_id}")
async def index_elasticsearch_connector(
connector_id: int,
search_space_id: int,
current_user: User = Depends(current_active_user),
session: AsyncSession = Depends(get_async_session),
) -> dict[str, Any]:
"""
Index documents from an Elasticsearch connector into a search space
"""
try:
# Get the connector directly from database
stmt = select(SearchSourceConnector).where(
SearchSourceConnector.id == connector_id,
SearchSourceConnector.user_id == current_user.id,
)
result = await session.execute(stmt)
connector = result.scalar_one_or_none()
if not connector:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Connector {connector_id} not found",
)
if connector.connector_type != SearchSourceConnectorType.ELASTICSEARCH:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Connector is not an Elasticsearch connector",
)
# Initialize the indexer
indexer = ElasticsearchIndexer(connector.config)
# Test connection first
logger.info(f"Testing Elasticsearch connection for connector {connector_id}")
connection_test = await indexer.test_connection()
if not connection_test.get("success"):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=f"Connection test failed: {connection_test.get('error')}",
)
# Index documents
logger.info(f"Starting indexing for Elasticsearch connector {connector_id}")
indexed_count = 0
error_count = 0
async for document in indexer.get_documents():
try:
# Create content hash
content_hash = hashlib.sha256(document.content.encode()).hexdigest()
# Check if document already exists
existing_doc_stmt = select(Document).where(
Document.content_hash == content_hash,
Document.search_space_id == search_space_id,
)
existing_doc_result = await session.execute(existing_doc_stmt)
if existing_doc_result.scalar_one_or_none():
logger.info(f"Document already exists, skipping: {document.title}")
continue
# Generate embedding for document
embedding = await config.embedding_model_instance.encode(
document.content
)
# Create document
db_document = Document(
title=document.title,
content=document.content,
content_hash=content_hash,
document_type=document.document_type,
document_metadata=document.metadata,
search_space_id=search_space_id,
embedding=embedding,
)
session.add(db_document)
await session.flush() # Get the document ID
# Create a single chunk with the entire document content
chunk_embedding = await config.embedding_model_instance.encode(
document.content
)
db_chunk = Chunk(
content=document.content,
embedding=chunk_embedding,
document_id=db_document.id,
)
session.add(db_chunk)
indexed_count += 1
if indexed_count % 10 == 0:
logger.info(f"Indexed {indexed_count} Elasticsearch documents")
await session.commit() # Periodic commit
except Exception as e:
logger.error(f"Error indexing document: {e}")
error_count += 1
await session.commit()
# Update last_indexed_at timestamp
update_stmt = (
update(SearchSourceConnector)
.where(SearchSourceConnector.id == connector_id)
.values(last_indexed_at=datetime.now(UTC))
)
await session.execute(update_stmt)
await session.commit()
logger.info(
f"Elasticsearch indexing completed: {indexed_count} documents indexed, {error_count} errors"
)
return {
"success": True,
"indexed_documents": indexed_count,
"errors": error_count,
"connector_id": connector_id,
"search_space_id": search_space_id,
"message": f"Successfully indexed {indexed_count} documents from Elasticsearch",
}
except HTTPException:
raise
except Exception as e:
logger.error(f"Error during Elasticsearch indexing: {e}")
await session.rollback()
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error during indexing: {e!s}",
) from e
# Note: Removed the index_elasticsearch_connector endpoint - now handled by universal system

View file

@ -7,7 +7,7 @@ PUT /search-source-connectors/{connector_id} - Update a specific connector
DELETE /search-source-connectors/{connector_id} - Delete a specific connector
POST /search-source-connectors/{connector_id}/index - Index content from a connector to a search space
Note: Each user can have only one connector of each type (SERPER_API, TAVILY_API, SLACK_CONNECTOR, NOTION_CONNECTOR, GITHUB_CONNECTOR, LINEAR_CONNECTOR, DISCORD_CONNECTOR, LUMA_CONNECTOR).
Note: Each user can have only one connector of each type (SERPER_API, TAVILY_API, SLACK_CONNECTOR, NOTION_CONNECTOR, GITHUB_CONNECTOR, LINEAR_CONNECTOR, DISCORD_CONNECTOR, LUMA_CONNECTOR, ELASTICSEARCH_CONNECTOR).
"""
import logging
@ -40,6 +40,7 @@ from app.tasks.connector_indexers import (
index_clickup_tasks,
index_confluence_pages,
index_discord_messages,
index_elasticsearch_documents,
index_github_repos,
index_google_calendar_events,
index_google_gmail_messages,
@ -346,6 +347,7 @@ async def index_connector_content(
- JIRA_CONNECTOR: Indexes issues and comments from Jira
- DISCORD_CONNECTOR: Indexes messages from all accessible Discord channels
- LUMA_CONNECTOR: Indexes events from Luma
- ELASTICSEARCH_CONNECTOR: Indexes documents from Elasticsearch instances
Args:
connector_id: ID of the connector to use
@ -572,6 +574,24 @@ async def index_connector_content(
)
response_message = "Luma indexing started in the background."
elif (
connector.connector_type
== SearchSourceConnectorType.ELASTICSEARCH_CONNECTOR
):
# Run indexing in background
logger.info(
f"Triggering Elasticsearch indexing for connector {connector_id} into search space {search_space_id} from {indexing_from} to {indexing_to}"
)
background_tasks.add_task(
run_elasticsearch_indexing_with_new_session,
connector_id,
search_space_id,
str(user.id),
indexing_from,
indexing_to,
)
response_message = "Elasticsearch indexing started in the background."
else:
raise HTTPException(
status_code=400,
@ -1341,3 +1361,55 @@ async def run_luma_indexing(
)
except Exception as e:
logger.error(f"Error in background Luma indexing task: {e!s}")
async def run_elasticsearch_indexing_with_new_session(
connector_id: int,
search_space_id: int,
user_id: str,
start_date: str,
end_date: str,
):
"""
Create a new session and run the Elasticsearch indexing task.
This prevents session leaks by creating a dedicated session for the background task.
"""
async with async_session_maker() as session:
await run_elasticsearch_indexing(
session, connector_id, search_space_id, user_id, start_date, end_date
)
async def run_elasticsearch_indexing(
session: AsyncSession,
connector_id: int,
search_space_id: int,
user_id: str,
start_date: str,
end_date: str,
):
"""
Background task to run Elasticsearch indexing.
"""
try:
indexed_count, error_message = await index_elasticsearch_documents(
session=session,
connector_id=connector_id,
search_space_id=search_space_id,
user_id=user_id,
start_date=start_date,
end_date=end_date,
update_last_indexed=False, # Don't update timestamp in the indexing function
)
if error_message:
logger.error(f"Elasticsearch indexing failed: {error_message}")
else:
logger.info(
f"Elasticsearch indexing completed successfully: {indexed_count} documents indexed"
)
# Update the last indexed timestamp only on success
await update_connector_last_indexed(session, connector_id)
except Exception as e:
logger.error(f"Error in background Elasticsearch indexing task: {e!s}")

View file

@ -208,6 +208,42 @@ class SearchSourceConnectorBase(BaseModel):
if not config.get("LUMA_API_KEY"):
raise ValueError("LUMA_API_KEY cannot be empty")
elif connector_type == SearchSourceConnectorType.ELASTICSEARCH_CONNECTOR:
# For ELASTICSEARCH_CONNECTOR, allow various connection parameters
allowed_keys = [
"hostname",
"port",
"username",
"password",
"api_key",
"ssl_enabled",
"indices",
"query",
"search_fields",
"max_documents",
]
# Check if hostname is present (required)
if not config.get("hostname"):
raise ValueError("Elasticsearch connector must have a hostname")
# Validate that all config keys are allowed
for key in config:
if key not in allowed_keys:
raise ValueError(
f"Invalid config key for Elasticsearch connector: {key}"
)
# Validate port if provided
if "port" in config and not isinstance(config["port"], int):
raise ValueError("Elasticsearch port must be an integer")
# Validate max_documents if provided
if "max_documents" in config and not isinstance(
config["max_documents"], int
):
raise ValueError("max_documents must be an integer")
return config

View file

@ -2041,7 +2041,7 @@ class ConnectorService:
top_k=top_k,
user_id=user_id,
search_space_id=search_space_id,
document_type="ELASTICSEARCH",
document_type="ELASTICSEARCH_CONNECTOR",
)
elif search_mode == SearchMode.DOCUMENTS:
elasticsearch_chunks = await self.document_retriever.hybrid_search(
@ -2049,7 +2049,7 @@ class ConnectorService:
top_k=top_k,
user_id=user_id,
search_space_id=search_space_id,
document_type="ELASTICSEARCH",
document_type="ELASTICSEARCH_CONNECTOR",
)
# Transform document retriever results to match expected format
elasticsearch_chunks = self._transform_document_results(
@ -2061,7 +2061,7 @@ class ConnectorService:
return {
"id": 12,
"name": "Elasticsearch",
"type": "ELASTICSEARCH",
"type": "ELASTICSEARCH_CONNECTOR",
"sources": [],
}, []
@ -2109,7 +2109,7 @@ class ConnectorService:
result_object = {
"id": 12,
"name": "Elasticsearch",
"type": "ELASTICSEARCH",
"type": "ELASTICSEARCH_CONNECTOR",
"sources": sources_list,
}

View file

@ -17,6 +17,7 @@ Available indexers:
- Google Gmail: Index messages from Google Gmail
- Google Calendar: Index events from Google Calendar
- Luma: Index events from Luma
- Elasticsearch: Index documents from Elasticsearch instances
"""
# Communication platforms
@ -27,6 +28,7 @@ from .confluence_indexer import index_confluence_pages
from .discord_indexer import index_discord_messages
# Development platforms
from .elasticsearch_indexer import index_elasticsearch_documents
from .github_indexer import index_github_repos
from .google_calendar_indexer import index_google_calendar_events
from .google_gmail_indexer import index_google_gmail_messages
@ -46,6 +48,7 @@ __all__ = [ # noqa: RUF022
"index_confluence_pages",
"index_discord_messages",
# Development platforms
"index_elasticsearch_documents",
"index_github_repos",
# Calendar and scheduling
"index_google_calendar_events",

View file

@ -1,31 +1,46 @@
"""
Elasticsearch connector indexer.
"""
import asyncio
import logging
from collections.abc import AsyncGenerator
from typing import Any
from datetime import datetime
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.asyncio import AsyncSession
from app.config import config
from app.connectors.elasticsearch_connector import ElasticsearchConnector
from app.schemas.documents import DocumentCreate
from app.tasks.connector_indexers.base import BaseConnectorIndexer
from app.db import Document, DocumentType, SearchSourceConnectorType
from app.services.llm_service import get_user_long_context_llm
from app.services.task_logging_service import TaskLoggingService
from app.utils.document_converters import (
create_document_chunks,
generate_content_hash,
generate_document_summary,
)
logger = logging.getLogger(__name__)
from .base import (
check_duplicate_document_by_hash,
get_connector_by_id,
logger,
update_connector_last_indexed,
)
class ElasticsearchIndexer(BaseConnectorIndexer):
"""Indexer for Elasticsearch data sources"""
class ElasticsearchIndexer:
"""Indexer class for Elasticsearch data sources (kept for backward compatibility)"""
def __init__(self, connector_config: dict[str, Any]):
super().__init__(connector_config)
def __init__(self, connector_config: dict):
self.connector_config = connector_config
self.connector = ElasticsearchConnector(self.config)
async def get_documents(self) -> AsyncGenerator[DocumentCreate, None]:
"""Fetch documents from Elasticsearch based on configuration"""
async def get_documents(self):
"""Get documents from Elasticsearch (kept for backward compatibility)"""
try:
# Connect to Elasticsearch
if not await self.connector.connect():
logger.error("Failed to connect to Elasticsearch")
return
# Get configuration parameters
query = self.connector_config.get("query", "*")
indices = self.connector_config.get("indices", None)
max_docs = self.connector_config.get("max_documents", 1000)
@ -33,7 +48,6 @@ class ElasticsearchIndexer(BaseConnectorIndexer):
logger.info(f"Searching Elasticsearch with query: {query}")
# Search documents
doc_count = 0
async for document in self.connector.search_documents(
query=query, indices=indices, size=max_docs, fields=fields
@ -44,7 +58,6 @@ class ElasticsearchIndexer(BaseConnectorIndexer):
yield document
doc_count += 1
# Add small delay to prevent overwhelming Elasticsearch
if doc_count % 10 == 0:
await asyncio.sleep(0.1)
@ -55,6 +68,320 @@ class ElasticsearchIndexer(BaseConnectorIndexer):
finally:
await self.connector.disconnect()
async def test_connection(self) -> dict[str, Any]:
"""Test the Elasticsearch connection"""
async def test_connection(self):
"""Test the Elasticsearch connection (kept for backward compatibility)"""
return await self.connector.test_connection()
async def index_elasticsearch_documents(
session: AsyncSession,
connector_id: int,
search_space_id: int,
user_id: str,
start_date: str
| None = None, # Not used for Elasticsearch but kept for consistency
end_date: str | None = None, # Not used for Elasticsearch but kept for consistency
update_last_indexed: bool = True,
) -> tuple[int, str | None]:
"""
Index Elasticsearch documents.
Args:
session: Database session
connector_id: ID of the Elasticsearch connector
search_space_id: ID of the search space to store documents in
user_id: User ID
start_date: Start date for indexing (not used for Elasticsearch but kept for consistency)
end_date: End date for indexing (not used for Elasticsearch but kept for consistency)
update_last_indexed: Whether to update the last_indexed_at timestamp (default: True)
Returns:
Tuple containing (number of documents indexed, error message or None)
"""
task_logger = TaskLoggingService(session, search_space_id)
# Log task start
log_entry = await task_logger.log_task_start(
task_name="elasticsearch_documents_indexing",
source="connector_indexing_task",
message=f"Starting Elasticsearch documents indexing for connector {connector_id}",
metadata={
"connector_id": connector_id,
"user_id": str(user_id),
"start_date": start_date,
"end_date": end_date,
},
)
try:
# Get the connector
await task_logger.log_task_progress(
log_entry,
f"Retrieving Elasticsearch connector {connector_id} from database",
{"stage": "connector_retrieval"},
)
# Get the connector from the database
connector = await get_connector_by_id(
session, connector_id, SearchSourceConnectorType.ELASTICSEARCH_CONNECTOR
)
if not connector:
await task_logger.log_task_failure(
log_entry,
f"Connector with ID {connector_id} not found or is not an Elasticsearch connector",
"Connector not found",
{"error_type": "ConnectorNotFound"},
)
return (
0,
f"Connector with ID {connector_id} not found or is not an Elasticsearch connector",
)
logger.info(f"Starting Elasticsearch indexing for connector {connector_id}")
# Initialize Elasticsearch connector
await task_logger.log_task_progress(
log_entry,
f"Initializing Elasticsearch connector for connector {connector_id}",
{"stage": "connector_initialization"},
)
try:
elasticsearch_connector = ElasticsearchConnector(connector.config)
except Exception as e:
await task_logger.log_task_failure(
log_entry,
f"Invalid Elasticsearch configuration for connector {connector_id}: {e}",
"Invalid configuration",
{"error_type": "InvalidConfiguration"},
)
return 0, f"Invalid Elasticsearch configuration: {e}"
# Test connection first
await task_logger.log_task_progress(
log_entry,
f"Testing Elasticsearch connection for connector {connector_id}",
{"stage": "connection_test"},
)
connection_test = await elasticsearch_connector.test_connection()
if not connection_test.get("success"):
await task_logger.log_task_failure(
log_entry,
f"Connection test failed: {connection_test.get('error')}",
"Connection failed",
{"error_type": "ConnectionError"},
)
return 0, f"Connection test failed: {connection_test.get('error')}"
# Get configuration parameters
query = connector.config.get("query", "*")
indices = connector.config.get("indices", None)
max_docs = connector.config.get("max_documents", 1000)
fields = connector.config.get("search_fields", None)
await task_logger.log_task_progress(
log_entry,
f"Fetching Elasticsearch documents with query: {query}",
{
"stage": "fetching_documents",
"query": query,
"indices": indices,
"max_documents": max_docs,
},
)
# Index documents
documents_indexed = 0
documents_skipped = 0
skipped_documents = []
try:
# Connect to Elasticsearch
if not await elasticsearch_connector.connect():
await task_logger.log_task_failure(
log_entry,
"Failed to connect to Elasticsearch",
"Connection failed",
{"error_type": "ConnectionError"},
)
return 0, "Failed to connect to Elasticsearch"
# Search documents
doc_count = 0
async for elasticsearch_doc in elasticsearch_connector.search_documents(
query=query, indices=indices, size=max_docs, fields=fields
):
try:
if doc_count >= max_docs:
break
# Convert ElasticsearchDocument to markdown content
document_content = elasticsearch_doc.content
document_title = elasticsearch_doc.title
if not document_content.strip():
logger.warning(
f"Skipping document with no content: {document_title}"
)
skipped_documents.append(f"{document_title} (no content)")
documents_skipped += 1
continue
# Generate content hash
content_hash = generate_content_hash(
document_content, search_space_id
)
# Check if document already exists
existing_document_by_hash = await check_duplicate_document_by_hash(
session, content_hash
)
if existing_document_by_hash:
logger.info(
f"Document with content hash {content_hash} already exists for document {document_title}. Skipping processing."
)
documents_skipped += 1
continue
# Generate summary with metadata
user_llm = await get_user_long_context_llm(session, user_id)
if user_llm:
# Create metadata for summary generation
document_metadata = {
"document_title": document_title,
"document_type": "Elasticsearch Document",
"connector_type": "Elasticsearch",
"source_index": elasticsearch_doc.metadata.get(
"elasticsearch_index", "unknown"
),
}
document_metadata.update(elasticsearch_doc.metadata)
(
summary_content,
summary_embedding,
) = await generate_document_summary(
document_content, user_llm, document_metadata
)
else:
# Fallback to simple summary if no LLM configured
summary_content = (
f"Elasticsearch Document: {document_title}\n\n"
)
if elasticsearch_doc.metadata.get("elasticsearch_index"):
summary_content += f"Index: {elasticsearch_doc.metadata['elasticsearch_index']}\n"
# Add content preview
content_preview = document_content[:300]
if len(document_content) > 300:
content_preview += "..."
summary_content += f"Content: {content_preview}\n"
summary_embedding = config.embedding_model_instance.embed(
summary_content
)
# Create document chunks
chunks = await create_document_chunks(document_content)
# Create document
document = Document(
search_space_id=search_space_id,
title=f"Elasticsearch Document - {document_title}",
document_type=DocumentType.ELASTICSEARCH_CONNECTOR,
document_metadata={
"document_title": document_title,
"source_index": elasticsearch_doc.metadata.get(
"elasticsearch_index", "unknown"
),
"document_id": elasticsearch_doc.metadata.get(
"elasticsearch_id", "unknown"
),
"indexed_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
**elasticsearch_doc.metadata,
},
content=summary_content,
content_hash=content_hash,
embedding=summary_embedding,
chunks=chunks,
)
session.add(document)
documents_indexed += 1
doc_count += 1
logger.info(f"Successfully indexed new document {document_title}")
# Add small delay to prevent overwhelming Elasticsearch
if doc_count % 10 == 0:
await asyncio.sleep(0.1)
except Exception as e:
logger.error(
f"Error processing document {elasticsearch_doc.title if 'elasticsearch_doc' in locals() else 'Unknown'}: {e!s}",
exc_info=True,
)
skipped_documents.append(
f"{elasticsearch_doc.title if 'elasticsearch_doc' in locals() else 'Unknown'} (processing error)"
)
documents_skipped += 1
continue
logger.info(f"Fetched {doc_count} documents from Elasticsearch")
except Exception as e:
logger.error(f"Error fetching documents from Elasticsearch: {e}")
await task_logger.log_task_failure(
log_entry,
f"Error fetching documents from Elasticsearch: {e}",
"Fetch error",
{"error_type": "FetchError"},
)
return 0, f"Error fetching documents from Elasticsearch: {e}"
finally:
await elasticsearch_connector.disconnect()
total_processed = documents_indexed
if total_processed > 0:
await update_connector_last_indexed(session, connector, update_last_indexed)
await session.commit()
await task_logger.log_task_success(
log_entry,
f"Successfully completed Elasticsearch indexing for connector {connector_id}",
{
"documents_processed": total_processed,
"documents_indexed": documents_indexed,
"documents_skipped": documents_skipped,
"skipped_documents_count": len(skipped_documents),
},
)
logger.info(
f"Elasticsearch indexing completed: {documents_indexed} new documents, {documents_skipped} skipped"
)
return total_processed, None
except SQLAlchemyError as db_error:
await session.rollback()
await task_logger.log_task_failure(
log_entry,
f"Database error during Elasticsearch indexing for connector {connector_id}",
str(db_error),
{"error_type": "SQLAlchemyError"},
)
logger.error(f"Database error: {db_error!s}", exc_info=True)
return 0, f"Database error: {db_error!s}"
except Exception as e:
await session.rollback()
await task_logger.log_task_failure(
log_entry,
f"Failed to index Elasticsearch documents for connector {connector_id}",
str(e),
{"error_type": type(e).__name__},
)
logger.error(f"Failed to index Elasticsearch documents: {e!s}", exc_info=True)
return 0, f"Failed to index Elasticsearch documents: {e!s}"