86 lines
3.0 KiB
Python
86 lines
3.0 KiB
Python
import sys
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
|
|
sys.path.append(os.getcwd())
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger("AI_Runner")
|
|
|
|
def main():
|
|
if len(sys.argv) < 4:
|
|
print(json.dumps({"status": "failed", "error": "Invalid arguments. Usage: python ai_runner.py <input_path> <session_id> <user_id>"}))
|
|
sys.exit(1)
|
|
|
|
input_path = sys.argv[1]
|
|
session_id = sys.argv[2]
|
|
user_id = sys.argv[3]
|
|
|
|
try:
|
|
from app.core.model_manager import ModelManager
|
|
from app.modules.documents.processors.pdf.pdf_converter import UniversalContentIntelligence
|
|
from app.modules.documents.processors.pdf.image_extractor import ImageExtractor
|
|
|
|
logger.info(f"🧠 Step 1: Running Marker AI on {input_path}")
|
|
with ModelManager() as ai:
|
|
result = ai.process_document(input_path)
|
|
|
|
if hasattr(result, "markdown"):
|
|
content = result.markdown
|
|
elif isinstance(result, dict) and "markdown" in result:
|
|
content = result["markdown"]
|
|
else:
|
|
try:
|
|
from marker.output import text_from_rendered
|
|
content = text_from_rendered(result)
|
|
except:
|
|
content = str(result)
|
|
|
|
logger.info("🔍 Step 2: Running Universal Content Intelligence")
|
|
metadata = UniversalContentIntelligence.extract_universal_metadata(content)
|
|
|
|
figure_image_map = {}
|
|
pattern1 = r'!\[Image\s+(\d+)\]\((http[s]?://[^)]+)\)\s*\n\s*\*\*Figure\s+(\d+):'
|
|
matches1 = re.findall(pattern1, content, re.IGNORECASE | re.MULTILINE)
|
|
for _, image_url, figure_num in matches1:
|
|
figure_image_map[figure_num] = image_url
|
|
|
|
pattern2 = r'!\[Image\s+(\d+)\]\(([^)]+)\)'
|
|
matches2 = re.findall(pattern2, content, re.IGNORECASE)
|
|
for image_num, image_url in matches2:
|
|
if image_num not in figure_image_map:
|
|
figure_image_map[image_num] = image_url
|
|
|
|
logger.info("🖼️ Step 3: Extracting images")
|
|
images = []
|
|
try:
|
|
extractor = ImageExtractor()
|
|
images = extractor.process_markdown_output_images(result, session_id=session_id, user_id=user_id)
|
|
logger.info(f"✅ Processed {len(images)} images")
|
|
except Exception as img_err:
|
|
logger.error(f"⚠️ Image extraction failed: {img_err}")
|
|
|
|
print(json.dumps({
|
|
"status": "success",
|
|
"markdown_content": content,
|
|
"session_id": session_id,
|
|
"images": images,
|
|
"metadata": {
|
|
"figures": figure_image_map,
|
|
"title": metadata.title if hasattr(metadata, 'title') else ""
|
|
}
|
|
}))
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ AI Runner Failed: {e}")
|
|
print(json.dumps({
|
|
"status": "failed",
|
|
"error": str(e)
|
|
}))
|
|
sys.exit(1)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|