atomic_doc_processor.py

import os
import json
import datetime
import argparse
from openai import OpenAI
from pdf2image import convert_from_path
from PIL import Image
import io
import yaml
from difflib import SequenceMatcher
import concurrent.futures
import base64  # For image base64

# Load config
with open('config.yaml', 'r') as f:
    CONFIG = yaml.safe_load(f)

client = OpenAI(api_key=CONFIG['api_key'])

def estimate_tokens(text):
    return len(text) // 4 + 1

def split_content_smart(content, max_tokens, overlap_chars=CONFIG['chunk_overlap']):
    chunks = []
    current_chunk = ""
    paragraphs = content.split('\n\n')
    for para in paragraphs:
        para = para.strip()
        if not para:
            continue
        if estimate_tokens(current_chunk + para) > max_tokens:
            if current_chunk:
                chunks.append(current_chunk)
            overlap = current_chunk[-overlap_chars:] if len(current_chunk) > overlap_chars else current_chunk
            current_chunk = overlap + para
        else:
            current_chunk += ('\n\n' if current_chunk else '') + para
    
    if estimate_tokens(current_chunk) > max_tokens:
        sentences = current_chunk.split('. ')
        current_chunk = ""
        for sent in sentences:
            sent = sent.strip() + ('. ' if sent else '')
            if estimate_tokens(current_chunk + sent) > max_tokens:
                if current_chunk:
                    chunks.append(current_chunk)
                overlap = current_chunk[-overlap_chars:] if len(current_chunk) > overlap_chars else current_chunk
                current_chunk = overlap + sent
            else:
                current_chunk += sent
    if current_chunk:
        chunks.append(current_chunk)
    
    return chunks

def llm_call(prompt, model=CONFIG['base_model'], is_multimodal=False, image=None):
    messages = [{"role": "user", "content": prompt}]
    if is_multimodal:
        content = [{"type": "text", "text": prompt}]
        if image:
            buffered = io.BytesIO()
            image.save(buffered, format="PNG")
            img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
            content.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_base64}"}})
        messages = [{"role": "user", "content": content}]
        model = CONFIG['multimodal_model']
    
    max_tokens = CONFIG.get('max_output_tokens', None)
    kwargs = {'max_tokens': max_tokens} if max_tokens else {}
    response = client.chat.completions.create(model=model, messages=messages, **kwargs)
    return response.choices[0].message.content.strip()

def decompose_content(content="", is_image=False, image=None):
    if is_image:
        prompt = CONFIG['extract_image_prompt']
    else:
        prompt = CONFIG['decompose_prompt'].format(content=content)
    output = llm_call(prompt, is_multimodal=is_image, image=image)
    try:
        return json.loads(output)
    except:
        return []

def reconstruct(facts):
    facts_json = json.dumps(facts)
    prompt = CONFIG['reconstruct_prompt'].format(facts=facts_json)
    return llm_call(prompt)

def validate(original, reconstructed, retries=0, facts=None):
    prompt = CONFIG['validate_prompt'].format(original=original, reconstructed=reconstructed)
    result = llm_call(prompt)
    try:
        res = json.loads(result)
        # Integrate similarity_threshold for double-check
        diff_ratio = SequenceMatcher(None, original, reconstructed).ratio()
        effective_loss = res['loss'] or (diff_ratio < CONFIG['similarity_threshold'])
        if not effective_loss or retries >= CONFIG['max_retries']:
            return reconstructed, effective_loss
        # Append missing to facts for self-healing
        if facts is not None:
            facts.extend(res['missing'])
            return validate(original, reconstruct(facts), retries + 1, facts)
        return reconstructed, True
    except:
        return reconstructed, True

def dedup_facts(facts, threshold=CONFIG['dedup_threshold']):
    unique = []
    for f in facts:
        if all(SequenceMatcher(None, f['fact'], u['fact']).ratio() < threshold for u in unique):
            unique.append(f)
    return unique

def extract_text_from_image(image):
    # Use multimodal model to extract text as OCR alternative (since no pytesseract in env)
    prompt = "Extract all text from this image accurately, including tables and charts. Output as plain text."
    return llm_call(prompt, is_multimodal=True, image=image)

def process_file(file_path, mode):
    if file_path.endswith('.md'):
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        max_tokens = CONFIG['max_token_length']
        all_facts = []
        if estimate_tokens(content) > max_tokens:
            chunks = split_content_smart(content, max_tokens)
        else:
            chunks = [content]
        for chunk in chunks:
            facts = decompose_content(chunk)
            all_facts.extend(facts)
        if mode == 'dedup':
            all_facts = dedup_facts(all_facts)
        reconstructed = reconstruct(all_facts)
        final, loss = validate(content, reconstructed, facts=all_facts)
        if loss:
            raise ValueError(f"Information loss in {file_path}; adjust prompts.")
        return file_path, final
    elif file_path.endswith('.pdf'):
        images = convert_from_path(file_path)
        all_facts = []
        for img in images:
            facts = decompose_content(is_image=True, image=img)
            all_facts.extend(facts)
        if mode == 'dedup':
            all_facts = dedup_facts(all_facts)
        reconstructed = reconstruct(all_facts)
        # Add OCR-like extraction for original using multimodal
        original = " ".join([extract_text_from_image(img) for img in images])
        final, loss = validate(original, reconstructed, facts=all_facts)
        if loss:
            raise ValueError(f"Information loss in {file_path}; adjust prompts.")
        return file_path, final

def main():
    parser = argparse.ArgumentParser(description="Atomic Document Processor")
    parser.add_argument('--input_path', default='.', help='Input path (dir or file)')
    parser.add_argument('--mode', default='high_fidelity', choices=['high_fidelity', 'dedup'], help='Processing mode')
    parser.add_argument('--per_dir', action='store_true', help='Output per directory/file')
    args = parser.parse_args()

    num_threads = max(CONFIG['min_threads'], min(CONFIG['num_threads'], CONFIG['max_threads']))

    timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
    output_dir = f"output_{timestamp}"
    os.makedirs(output_dir, exist_ok=True)

    files_to_process = []
    if os.path.isdir(args.input_path):
        for root, _, files in os.walk(args.input_path):
            for file in files:
                if file.endswith(('.pdf', '.md')):
                    files_to_process.append(os.path.join(root, file))
    else:
        files_to_process = [args.input_path]

    results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor:
        future_to_file = {executor.submit(process_file, file, args.mode): file for file in files_to_process}
        for future in concurrent.futures.as_completed(future_to_file):
            try:
                results.append(future.result())
            except Exception as e:
                print(f"Error processing {future_to_file[future]}: {e}")

    if not args.per_dir:
        all_output = []
        for file, output in sorted(results, key=lambda x: x[0]):  # Sort by file path
            all_output.append(f"# Processed: {file}\n{output}")
        with open(os.path.join(output_dir, 'all_processed.md'), 'w', encoding='utf-8') as f:
            f.write('\n\n'.join(all_output))
    else:
        for file, output in results:
            rel_path = os.path.relpath(file, args.input_path)
            out_path = os.path.join(output_dir, rel_path.replace('.pdf', '.md').replace('.md', '_processed.md'))
            os.makedirs(os.path.dirname(out_path), exist_ok=True)
            with open(out_path, 'w', encoding='utf-8') as f:
                f.write(output)

if __name__ == "__main__":
    main()