mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-20 23:21:06 +02:00
feat: implement Elasticsearch connector with indexing capabilities and validation
This commit is contained in:
parent
0b32b66888
commit
2ad1c43059
8 changed files with 522 additions and 212 deletions
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue