253 lines
11 KiB
Python
253 lines
11 KiB
Python
"""
|
|
AI processing method module.
|
|
Handles direct AI calls for any type of task.
|
|
"""
|
|
|
|
import time
|
|
import logging
|
|
from typing import Dict, Any, List, Optional
|
|
from datetime import datetime, UTC
|
|
|
|
from modules.workflows.methods.methodBase import MethodBase, action
|
|
from modules.datamodels.datamodelChat import ActionResult
|
|
from modules.datamodels.datamodelAi import AiCallOptions
|
|
from modules.datamodels.datamodelChat import ChatDocument
|
|
from modules.aicore.aicorePluginTavily import WebResearchRequest
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class MethodAi(MethodBase):
|
|
"""AI processing methods."""
|
|
|
|
def __init__(self, services):
|
|
super().__init__(services)
|
|
self.name = "ai"
|
|
self.description = "AI processing methods"
|
|
|
|
def _format_timestamp_for_filename(self) -> str:
|
|
"""Format current timestamp as YYYYMMDD-hhmmss for filenames."""
|
|
return datetime.now(UTC).strftime("%Y%m%d-%H%M%S")
|
|
|
|
@action
|
|
async def process(self, parameters: Dict[str, Any]) -> ActionResult:
|
|
"""
|
|
GENERAL:
|
|
- Purpose: Process a user prompt with optional unlimited input documents to produce one or many output documents of the SAME format.
|
|
- Input requirements: aiPrompt (required); optional documentList.
|
|
- Output format: Exactly one file format to select. For multiple output file formats you need to do different calls.
|
|
|
|
Parameters:
|
|
- aiPrompt (str, required): Instruction for the AI.
|
|
- documentList (list, optional): Document reference(s) for context.
|
|
- resultType (str, optional): Output file extension - only one extension allowed (e.g. txt, json, md, csv, xml, html, pdf, docx, xlsx, png, ...). Default: txt.
|
|
"""
|
|
try:
|
|
# Init progress logger
|
|
operationId = f"ai_process_{self.services.currentWorkflow.id}_{int(time.time())}"
|
|
|
|
# Start progress tracking
|
|
self.services.workflow.progressLogStart(
|
|
operationId,
|
|
"Generate",
|
|
"AI Processing",
|
|
f"Format: {parameters.get('resultType', 'txt')}"
|
|
)
|
|
|
|
# Debug logging to see what parameters are received
|
|
logger.info(f"MethodAi.process received parameters: {parameters}")
|
|
logger.info(f"Parameters type: {type(parameters)}")
|
|
logger.info(f"Parameters keys: {list(parameters.keys()) if isinstance(parameters, dict) else 'Not a dict'}")
|
|
|
|
aiPrompt = parameters.get("aiPrompt")
|
|
logger.info(f"aiPrompt extracted: '{aiPrompt}' (type: {type(aiPrompt)})")
|
|
|
|
# Update progress - preparing parameters
|
|
self.services.workflow.progressLogUpdate(operationId, 0.2, "Preparing parameters")
|
|
|
|
documentList = parameters.get("documentList", [])
|
|
if isinstance(documentList, str):
|
|
documentList = [documentList]
|
|
resultType = parameters.get("resultType", "txt")
|
|
|
|
|
|
if not aiPrompt:
|
|
logger.error(f"aiPrompt is missing or empty. Parameters: {parameters}")
|
|
return ActionResult.isFailure(
|
|
error="AI prompt is required"
|
|
)
|
|
|
|
# Determine output extension and default MIME type without duplicating service logic
|
|
normalized_result_type = (str(resultType).strip().lstrip('.').lower() or "txt")
|
|
output_extension = f".{normalized_result_type}"
|
|
output_mime_type = "application/octet-stream" # Prefer service-provided mimeType when available
|
|
logger.info(f"Using result type: {resultType} -> {output_extension}")
|
|
|
|
# Update progress - preparing documents
|
|
self.services.workflow.progressLogUpdate(operationId, 0.3, "Preparing documents")
|
|
|
|
# Get ChatDocuments for AI service - let AI service handle all document processing
|
|
chatDocuments = []
|
|
if documentList:
|
|
chatDocuments = self.services.workflow.getChatDocumentsFromDocumentList(documentList)
|
|
if chatDocuments:
|
|
logger.info(f"Prepared {len(chatDocuments)} documents for AI processing")
|
|
|
|
# Update progress - preparing AI call
|
|
self.services.workflow.progressLogUpdate(operationId, 0.4, "Preparing AI call")
|
|
|
|
# Build options with only resultFormat - let service layer handle all other parameters
|
|
output_format = output_extension.replace('.', '') or 'txt'
|
|
options = AiCallOptions(
|
|
resultFormat=output_format
|
|
# Removed all model parameters - service layer will analyze prompt and determine optimal parameters
|
|
)
|
|
|
|
# Update progress - calling AI
|
|
self.services.workflow.progressLogUpdate(operationId, 0.6, "Calling AI")
|
|
|
|
result = await self.services.ai.callAiDocuments(
|
|
prompt=aiPrompt,
|
|
documents=chatDocuments if chatDocuments else None,
|
|
options=options,
|
|
outputFormat=output_format
|
|
)
|
|
|
|
# Update progress - processing result
|
|
self.services.workflow.progressLogUpdate(operationId, 0.8, "Processing result")
|
|
|
|
from modules.datamodels.datamodelChat import ActionDocument
|
|
|
|
if isinstance(result, dict) and isinstance(result.get("documents"), list):
|
|
action_documents = []
|
|
for d in result["documents"]:
|
|
action_documents.append(ActionDocument(
|
|
documentName=d.get("documentName"),
|
|
documentData=d.get("documentData"),
|
|
mimeType=d.get("mimeType") or output_mime_type
|
|
))
|
|
|
|
# Complete progress tracking
|
|
self.services.workflow.progressLogFinish(operationId, True)
|
|
|
|
return ActionResult.isSuccess(documents=action_documents)
|
|
|
|
extension = output_extension.lstrip('.')
|
|
meaningful_name = self._generateMeaningfulFileName(
|
|
base_name="ai",
|
|
extension=extension,
|
|
action_name="result"
|
|
)
|
|
action_document = ActionDocument(
|
|
documentName=meaningful_name,
|
|
documentData=result,
|
|
mimeType=output_mime_type
|
|
)
|
|
|
|
# Complete progress tracking
|
|
self.services.workflow.progressLogFinish(operationId, True)
|
|
|
|
return ActionResult.isSuccess(documents=[action_document])
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in AI processing: {str(e)}")
|
|
|
|
# Complete progress tracking with failure
|
|
try:
|
|
self.services.workflow.progressLogFinish(operationId, False)
|
|
except:
|
|
pass # Don't fail on progress logging errors
|
|
|
|
return ActionResult.isFailure(
|
|
error=str(e)
|
|
)
|
|
|
|
@action
|
|
async def webResearch(self, parameters: Dict[str, Any]) -> ActionResult:
|
|
"""
|
|
GENERAL:
|
|
- Purpose: Web research and information gathering with basic analysis and sources.
|
|
- Input requirements: user_prompt (required); optional urls, max_results, max_pages, search_depth, extract_depth, pages_search_depth, country, time_range, topic, language.
|
|
- Output format: JSON with results and sources.
|
|
|
|
Parameters:
|
|
- user_prompt (str, required): Research question or topic.
|
|
- urls (list, optional): Specific URLs to crawl.
|
|
- max_results (int, optional): Max search results. Default: 5.
|
|
- max_pages (int, optional): Max pages to crawl per site. Default: 5.
|
|
- extract_depth (str, optional): basic | advanced. Default: advanced.
|
|
- search_depth (int, optional): Crawl depth level - how many times to follow sublinks of a page. Default: 2.
|
|
- country (str, optional): Full English country name (ISO-3166; map codes via pycountry/i18n-iso-countries).
|
|
- time_range (str, optional): d | w | m | y.
|
|
- topic (str, optional): general | news | academic.
|
|
- language (str, optional): Language code (e.g., de, en, fr).
|
|
"""
|
|
try:
|
|
user_prompt = parameters.get("user_prompt")
|
|
urls = parameters.get("urls")
|
|
max_results = parameters.get("max_results", 5)
|
|
max_pages = parameters.get("max_pages", 5)
|
|
extract_depth = parameters.get("extract_depth", "advanced")
|
|
search_depth = parameters.get("pages_search_depth", 2)
|
|
country = parameters.get("country")
|
|
time_range = parameters.get("time_range")
|
|
topic = parameters.get("topic")
|
|
language = parameters.get("language")
|
|
|
|
if not user_prompt:
|
|
return ActionResult.isFailure(
|
|
error="Search query is required"
|
|
)
|
|
|
|
# Build WebResearchRequest (simplified dataclass)
|
|
request = WebResearchRequest(
|
|
user_prompt=user_prompt,
|
|
urls=urls,
|
|
max_results=max_results,
|
|
max_pages=max_pages,
|
|
search_depth=search_depth,
|
|
extract_depth=extract_depth,
|
|
country=country,
|
|
time_range=time_range,
|
|
topic=topic,
|
|
language=language
|
|
)
|
|
|
|
# Call web research service
|
|
logger.info(f"Performing comprehensive web research for: {user_prompt}")
|
|
logger.info(f"Max results: {max_results}, Max pages: {max_pages}")
|
|
if urls:
|
|
logger.info(f"Using provided URLs: {len(urls)}")
|
|
|
|
result = await self.services.ai.webResearch(request)
|
|
|
|
if not result.success:
|
|
return ActionResult.isFailure(error=result.error)
|
|
|
|
# Convert WebResearchResult to ActionResult format
|
|
documents = []
|
|
for doc in result.documents:
|
|
documents.append({
|
|
"documentName": doc.documentName,
|
|
"documentData": {
|
|
"user_prompt": doc.documentData.user_prompt,
|
|
"websites_analyzed": doc.documentData.websites_analyzed,
|
|
"additional_links_found": doc.documentData.additional_links_found,
|
|
"analysis_result": doc.documentData.analysis_result,
|
|
"sources": [{"title": s.title, "url": str(s.url)} for s in doc.documentData.sources],
|
|
"additional_links": doc.documentData.additional_links,
|
|
"debug_info": doc.documentData.debug_info
|
|
},
|
|
"mimeType": doc.mimeType
|
|
})
|
|
|
|
# Return result in the standard ActionResult format
|
|
return ActionResult.isSuccess(
|
|
documents=documents
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in web research: {str(e)}")
|
|
return ActionResult.isFailure(
|
|
error=str(e)
|
|
)
|
|
|