146 lines
5.7 KiB
Python
146 lines
5.7 KiB
Python
# Copyright (c) 2025 Patrick Motsch
|
|
# All rights reserved.
|
|
|
|
"""
|
|
Query Database action for Chatbot operations.
|
|
Executes SQL queries via the preprocessor connector.
|
|
"""
|
|
|
|
import logging
|
|
import json
|
|
import time
|
|
from typing import Dict, Any
|
|
from modules.workflows.methods.methodBase import action
|
|
from modules.datamodels.datamodelChatbot import ActionResult, ActionDocument
|
|
from modules.connectors.connectorPreprocessor import PreprocessorConnector
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
@action
|
|
async def queryDatabase(self, parameters: Dict[str, Any]) -> ActionResult:
|
|
"""
|
|
Execute a SQL query via the preprocessor connector.
|
|
|
|
Parameters:
|
|
- sqlQuery (str, required): SQL SELECT query to execute. Can also be extracted from analysis_result document if provided in documentList.
|
|
"""
|
|
try:
|
|
# Init progress logger
|
|
workflowId = self.services.workflow.id if self.services.workflow else f"no-workflow-{int(time.time())}"
|
|
operationId = f"chatbot_query_db_{workflowId}_{int(time.time())}"
|
|
|
|
# Start progress tracking
|
|
parentOperationId = parameters.get('parentOperationId')
|
|
self.services.chat.progressLogStart(
|
|
operationId,
|
|
"Database Query",
|
|
"Executing SQL Query",
|
|
"Preprocessing API",
|
|
parentOperationId=parentOperationId
|
|
)
|
|
|
|
# Get SQL query from parameters or extract from documentList
|
|
sqlQuery = parameters.get("sqlQuery")
|
|
|
|
# If sqlQuery not provided, try to extract from documentList (analysis_result)
|
|
if not sqlQuery:
|
|
documentListParam = parameters.get("documentList")
|
|
if documentListParam:
|
|
# Get documents from previous task
|
|
from modules.datamodels.datamodelDocref import DocumentReferenceList
|
|
if isinstance(documentListParam, str):
|
|
docList = DocumentReferenceList.from_string_list([documentListParam])
|
|
elif isinstance(documentListParam, list):
|
|
docList = DocumentReferenceList.from_string_list(documentListParam)
|
|
else:
|
|
docList = documentListParam
|
|
|
|
# Get documents from workflow
|
|
documents = self.services.chat.getChatDocumentsFromDocumentList(docList)
|
|
|
|
# Try to extract SQL query from JSON document
|
|
for doc in documents:
|
|
try:
|
|
# ChatDocument objects have fileId - get file data from database
|
|
if hasattr(doc, 'fileId') and doc.fileId:
|
|
# Get file data from database
|
|
fileData = self.services.interfaceDbComponent.getFileData(doc.fileId)
|
|
if fileData:
|
|
# Decode bytes if needed
|
|
if isinstance(fileData, bytes):
|
|
docData = fileData.decode('utf-8')
|
|
else:
|
|
docData = str(fileData)
|
|
|
|
# Try to parse as JSON
|
|
analysisData = json.loads(docData)
|
|
sqlQuery = analysisData.get("sqlQuery")
|
|
|
|
if sqlQuery:
|
|
logger.info(f"Extracted SQL query from analysis_result document: {sqlQuery[:100]}...")
|
|
break
|
|
except (json.JSONDecodeError, AttributeError, KeyError, TypeError) as e:
|
|
logger.debug(f"Could not parse document as JSON: {e}")
|
|
continue
|
|
|
|
if not sqlQuery:
|
|
return ActionResult.isFailure(error="SQL query is required. Provide sqlQuery parameter or analysis_result document with sqlQuery field.")
|
|
|
|
# Update progress
|
|
self.services.chat.progressLogUpdate(operationId, 0.3, "Validating query")
|
|
|
|
# Initialize connector
|
|
connector = PreprocessorConnector()
|
|
|
|
# Update progress
|
|
self.services.chat.progressLogUpdate(operationId, 0.5, "Executing query")
|
|
|
|
# Execute query
|
|
result = await connector.executeQuery(sqlQuery)
|
|
|
|
# Update progress
|
|
self.services.chat.progressLogUpdate(operationId, 0.8, "Formatting results")
|
|
|
|
# Generate meaningful filename
|
|
meaningful_name = self._generateMeaningfulFileName(
|
|
base_name="database_query",
|
|
extension="txt",
|
|
action_name="queryDatabase"
|
|
)
|
|
|
|
# Create validation metadata
|
|
validationMetadata = self._createValidationMetadata(
|
|
"queryDatabase",
|
|
sqlQuery=sqlQuery[:200] if len(sqlQuery) > 200 else sqlQuery, # Truncate for metadata
|
|
resultLength=len(result)
|
|
)
|
|
|
|
# Create action document
|
|
document = ActionDocument(
|
|
documentName=meaningful_name,
|
|
documentData=result,
|
|
mimeType="text/plain",
|
|
validationMetadata=validationMetadata
|
|
)
|
|
|
|
# Complete progress tracking
|
|
self.services.chat.progressLogFinish(operationId, True)
|
|
|
|
# Close connector
|
|
await connector.close()
|
|
|
|
return ActionResult.isSuccess(documents=[document])
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error executing database query: {str(e)}")
|
|
|
|
# Complete progress tracking with failure
|
|
try:
|
|
self.services.chat.progressLogFinish(operationId, False)
|
|
except:
|
|
pass
|
|
|
|
return ActionResult.isFailure(
|
|
error=str(e)
|
|
)
|
|
|