From 2ad1c4305924dddb37e83842ac437abbd1f75b2a Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Sat, 4 Oct 2025 13:47:59 +0530 Subject: [PATCH] feat: implement Elasticsearch connector with indexing capabilities and validation --- .../app/connectors/elasticsearch_connector.py | 64 +-- surfsense_backend/app/db.py | 4 +- .../elasticsearch_add_connector_route.py | 182 ++------- .../routes/search_source_connectors_routes.py | 74 +++- .../app/schemas/search_source_connector.py | 36 ++ .../app/services/connector_service.py | 8 +- .../app/tasks/connector_indexers/__init__.py | 3 + .../elasticsearch_indexer.py | 363 +++++++++++++++++- 8 files changed, 522 insertions(+), 212 deletions(-) diff --git a/surfsense_backend/app/connectors/elasticsearch_connector.py b/surfsense_backend/app/connectors/elasticsearch_connector.py index 62193fe8d..930d140ce 100644 --- a/surfsense_backend/app/connectors/elasticsearch_connector.py +++ b/surfsense_backend/app/connectors/elasticsearch_connector.py @@ -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}") diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py index 47b3fb02c..457d95544 100644 --- a/surfsense_backend/app/db.py +++ b/surfsense_backend/app/db.py @@ -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): diff --git a/surfsense_backend/app/routes/elasticsearch_add_connector_route.py b/surfsense_backend/app/routes/elasticsearch_add_connector_route.py index 24e4bc3b7..4a6f654cb 100644 --- a/surfsense_backend/app/routes/elasticsearch_add_connector_route.py +++ b/surfsense_backend/app/routes/elasticsearch_add_connector_route.py @@ -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 diff --git a/surfsense_backend/app/routes/search_source_connectors_routes.py b/surfsense_backend/app/routes/search_source_connectors_routes.py index 6250e0c07..50db97668 100644 --- a/surfsense_backend/app/routes/search_source_connectors_routes.py +++ b/surfsense_backend/app/routes/search_source_connectors_routes.py @@ -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}") diff --git a/surfsense_backend/app/schemas/search_source_connector.py b/surfsense_backend/app/schemas/search_source_connector.py index 0ae84cc81..022be65e4 100644 --- a/surfsense_backend/app/schemas/search_source_connector.py +++ b/surfsense_backend/app/schemas/search_source_connector.py @@ -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 diff --git a/surfsense_backend/app/services/connector_service.py b/surfsense_backend/app/services/connector_service.py index d157b8d96..4227fee43 100644 --- a/surfsense_backend/app/services/connector_service.py +++ b/surfsense_backend/app/services/connector_service.py @@ -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, } diff --git a/surfsense_backend/app/tasks/connector_indexers/__init__.py b/surfsense_backend/app/tasks/connector_indexers/__init__.py index fb7936126..766506f70 100644 --- a/surfsense_backend/app/tasks/connector_indexers/__init__.py +++ b/surfsense_backend/app/tasks/connector_indexers/__init__.py @@ -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", diff --git a/surfsense_backend/app/tasks/connector_indexers/elasticsearch_indexer.py b/surfsense_backend/app/tasks/connector_indexers/elasticsearch_indexer.py index f0d16879b..1e4b3d7f0 100644 --- a/surfsense_backend/app/tasks/connector_indexers/elasticsearch_indexer.py +++ b/surfsense_backend/app/tasks/connector_indexers/elasticsearch_indexer.py @@ -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}"