258 lines
8.9 KiB
Python
258 lines
8.9 KiB
Python
import asyncio
|
|
import sys
|
|
import os
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
# Add the project root to the sys.path
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
|
|
|
|
from modules.datamodels.datamodelAi import AiCallRequest, AiCallOptions, OperationType
|
|
from modules.datamodels.datamodelChat import ChatDocument
|
|
from modules.services.serviceAi.subCoreAi import SubCoreAi
|
|
|
|
class MockAiObjects:
|
|
def __init__(self, responses):
|
|
self.responses = responses
|
|
self.call_count = 0
|
|
|
|
async def call(self, request: AiCallRequest):
|
|
if self.call_count < len(self.responses):
|
|
response_content = self.responses[self.call_count]
|
|
self.call_count += 1
|
|
mock_response = MagicMock()
|
|
mock_response.content = response_content
|
|
mock_response.modelName = "mock-model"
|
|
mock_response.priceUsd = 0.001
|
|
mock_response.processingTime = 0.1
|
|
print(f" Mock AI Call {self.call_count}: Responding with partial result (length: {len(response_content)})")
|
|
return mock_response
|
|
else:
|
|
print(" Mock AI Call: No more mock responses, returning empty.")
|
|
mock_response = MagicMock()
|
|
mock_response.content = ""
|
|
return mock_response
|
|
|
|
class MockServices:
|
|
def __init__(self):
|
|
self.currentWorkflow = MagicMock()
|
|
self.currentWorkflow.id = "test_workflow_123"
|
|
self.workflow = MagicMock()
|
|
self.workflow.createProgressLogger.return_value = MagicMock()
|
|
self.workflow.storeWorkflowStat = AsyncMock()
|
|
self.ai = MagicMock()
|
|
self.ai.sanitizePromptContent.side_effect = lambda content, type: content
|
|
self.utils = MagicMock()
|
|
self.utils.debugLogToFile.side_effect = lambda msg, tag: print(f" DEBUG ({tag}): {msg}")
|
|
self.utils.configGet.return_value = False # Disable debug files for tests
|
|
|
|
class MockDocumentProcessor:
|
|
async def callAiText(self, prompt, documents, options):
|
|
return "Extracted content from documents: Sample text content"
|
|
|
|
async def test_unified_architecture():
|
|
print("\n=== Testing Unified Architecture ===")
|
|
|
|
# Mock responses: 1 for generation prompt building + 2 for actual generation
|
|
mock_responses = [
|
|
# Response 1: Generation prompt building
|
|
"Generate JSON content that creates a structured document with prime numbers in a table format. Use the canonical JSON format with sections and elements.",
|
|
|
|
# Response 2: First part of generation
|
|
"""{
|
|
"metadata": {
|
|
"title": "Prime Numbers List",
|
|
"splitStrategy": "single_document",
|
|
"source_documents": [],
|
|
"extraction_method": "ai_generation"
|
|
},
|
|
"documents": [
|
|
{
|
|
"id": "doc_primes_1_500",
|
|
"title": "Prime Numbers 1-500",
|
|
"filename": "primes_1_500.docx",
|
|
"sections": [
|
|
{
|
|
"id": "section_1",
|
|
"content_type": "table",
|
|
"elements": [
|
|
{
|
|
"headers": ["Number", "Prime"],
|
|
"rows": [
|
|
["1", "2"], ["2", "3"], ["3", "5"], ["4", "7"], ["5", "11"]
|
|
]
|
|
}
|
|
],
|
|
"order": 1
|
|
}
|
|
]
|
|
}
|
|
]
|
|
} [CONTINUE: Generate remaining prime numbers from 501 to 1000]""",
|
|
|
|
# Response 3: Second part of generation
|
|
"""{
|
|
"metadata": {
|
|
"title": "Prime Numbers List",
|
|
"splitStrategy": "single_document",
|
|
"source_documents": [],
|
|
"extraction_method": "ai_generation"
|
|
},
|
|
"documents": [
|
|
{
|
|
"id": "doc_primes_501_1000",
|
|
"title": "Prime Numbers 501-1000",
|
|
"filename": "primes_501_1000.docx",
|
|
"sections": [
|
|
{
|
|
"id": "section_2",
|
|
"content_type": "table",
|
|
"elements": [
|
|
{
|
|
"headers": ["Number", "Prime"],
|
|
"rows": [
|
|
["501", "3571"], ["502", "3572"], ["503", "3581"]
|
|
]
|
|
}
|
|
],
|
|
"order": 2
|
|
}
|
|
]
|
|
}
|
|
]
|
|
}"""
|
|
]
|
|
|
|
mock_ai_objects = MockAiObjects(mock_responses)
|
|
mock_services = MockServices()
|
|
mock_document_processor = MockDocumentProcessor()
|
|
|
|
core_ai_service = SubCoreAi(mock_services, mock_ai_objects)
|
|
|
|
prompt = "Generate the first 1000 prime numbers and arrange them in a structured table format."
|
|
options = AiCallOptions(operationType=OperationType.GENERATE_CONTENT)
|
|
output_format = "docx"
|
|
title = "Prime Numbers List"
|
|
|
|
print(f"User Prompt: '{prompt}'")
|
|
print("Testing unified architecture with direct generation (no documents)...")
|
|
|
|
# Test the unified generation method directly
|
|
result = await core_ai_service._callAiUnifiedGeneration(prompt, None, options, output_format, title)
|
|
|
|
print("\n--- Generated JSON Result ---")
|
|
print(f"Result length: {len(result)} characters")
|
|
print(f"Result preview: {result[:300]}...")
|
|
|
|
# Verify it's valid JSON
|
|
import json
|
|
try:
|
|
parsed_result = json.loads(result)
|
|
print(f"✅ Valid JSON with {len(parsed_result.get('documents', []))} documents")
|
|
|
|
# Verify it's using the multi-document format
|
|
if "documents" in parsed_result and "metadata" in parsed_result:
|
|
print("✅ Using unified multi-document format")
|
|
print("✅ Architecture is properly unified!")
|
|
return True
|
|
else:
|
|
print("❌ Not using multi-document format")
|
|
return False
|
|
except json.JSONDecodeError as e:
|
|
print(f"❌ Invalid JSON: {str(e)}")
|
|
return False
|
|
|
|
async def test_with_documents():
|
|
print("\n=== Testing Unified Architecture WITH Documents ===")
|
|
|
|
# Mock responses: 1 for generation prompt building + 1 for actual generation
|
|
mock_responses = [
|
|
# Response 1: Generation prompt building
|
|
"Generate JSON content that creates a comprehensive fruit analysis report based on the extracted content. Use the canonical JSON format with sections and elements.",
|
|
|
|
# Response 2: Generation with extracted content
|
|
"""{
|
|
"metadata": {
|
|
"title": "Fruit Analysis Report",
|
|
"splitStrategy": "single_document",
|
|
"source_documents": ["doc1"],
|
|
"extraction_method": "ai_generation"
|
|
},
|
|
"documents": [
|
|
{
|
|
"id": "doc_fruit_analysis",
|
|
"title": "Fruit Analysis Report",
|
|
"filename": "fruit_analysis.docx",
|
|
"sections": [
|
|
{
|
|
"id": "section_1",
|
|
"content_type": "paragraph",
|
|
"elements": [
|
|
{
|
|
"text": "Based on the extracted content, here is a comprehensive fruit analysis..."
|
|
}
|
|
],
|
|
"order": 1
|
|
}
|
|
]
|
|
}
|
|
]
|
|
}"""
|
|
]
|
|
|
|
mock_ai_objects = MockAiObjects(mock_responses)
|
|
mock_services = MockServices()
|
|
mock_document_processor = MockDocumentProcessor()
|
|
|
|
core_ai_service = SubCoreAi(mock_services, mock_ai_objects)
|
|
|
|
prompt = "Extract all fruit information and create a comprehensive analysis report."
|
|
options = AiCallOptions(operationType=OperationType.GENERATE_CONTENT)
|
|
output_format = "docx"
|
|
title = "Fruit Analysis Report"
|
|
|
|
print(f"User Prompt: '{prompt}'")
|
|
print("Testing unified architecture with document extraction...")
|
|
|
|
# Test the unified generation method with extracted content
|
|
result = await core_ai_service._callAiUnifiedGeneration(prompt, "Sample fruit data: apples, oranges, bananas", options, output_format, title)
|
|
|
|
print("\n--- Generated JSON Result ---")
|
|
print(f"Result length: {len(result)} characters")
|
|
print(f"Result preview: {result[:300]}...")
|
|
|
|
# Verify it's valid JSON
|
|
import json
|
|
try:
|
|
parsed_result = json.loads(result)
|
|
print(f"✅ Valid JSON with {len(parsed_result.get('documents', []))} documents")
|
|
|
|
# Verify it's using the multi-document format
|
|
if "documents" in parsed_result and "metadata" in parsed_result:
|
|
print("✅ Using unified multi-document format")
|
|
print("✅ Architecture is properly unified!")
|
|
return True
|
|
else:
|
|
print("❌ Not using multi-document format")
|
|
return False
|
|
except json.JSONDecodeError as e:
|
|
print(f"❌ Invalid JSON: {str(e)}")
|
|
return False
|
|
|
|
async def main():
|
|
print("🚀 Testing Unified Architecture Implementation")
|
|
print("=" * 60)
|
|
|
|
success1 = await test_unified_architecture()
|
|
success2 = await test_with_documents()
|
|
|
|
if success1 and success2:
|
|
print("\n🎉 ALL TESTS PASSED! Unified architecture is properly implemented.")
|
|
print("✅ Single document = multi-document with n=1")
|
|
print("✅ Always uses multi-document JSON format")
|
|
print("✅ Continuation logic works for long responses")
|
|
print("✅ Both scenarios (with/without documents) work")
|
|
else:
|
|
print("\n❌ Some tests failed. Please check the implementation.")
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|