Convert Equipment Manuals to Searchable Knowledge Bases with Granite-Docling and LlamaIndex

                      
                     
"""
Production implementation for converting equipment manuals into searchable knowledge bases.
Uses Granite-Docling for document handling and LlamaIndex for indexing.
"""

from typing import Dict, Any, List
import os
import logging
import requests
from time import sleep
from sqlalchemy import create_engine, Table, MetaData, Column, String, Integer
from sqlalchemy.orm import sessionmaker, scoped_session

# Logging setup for monitoring
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Database connection pooling setup
DATABASE_URL = os.getenv('DATABASE_URL', 'sqlite:///manuals.db')
engine = create_engine(DATABASE_URL, pool_size=10, max_overflow=20)
session_factory = sessionmaker(bind=engine)
session = scoped_session(session_factory)

class Config:
    """
    Configuration class for environment variables.
    Holds critical configuration parameters.
    """
    def __init__(self):
        self.granite_api_url = os.getenv('GRANITE_API_URL')
        self.llama_api_url = os.getenv('LLAMA_API_URL')

config = Config()

def validate_input(data: Dict[str, Any]) -> bool:
    """Validate request data.
    
    Args:
        data: Input to validate
    Returns:
        True if valid
    Raises:
        ValueError: If validation fails
    """
    if 'manual_id' not in data:
        raise ValueError('Missing manual_id')  # Check for required field
    if not isinstance(data['manual_id'], str):
        raise ValueError('manual_id must be a string')  # Ensure type is correct
    return True

def sanitize_fields(data: Dict[str, Any]) -> Dict[str, Any]:
    """Sanitize input fields to prevent injection attacks.
    
    Args:
        data: Dictionary containing fields to sanitize
    Returns:
        Sanitized dictionary
    """
    return {k: v.strip() for k, v in data.items()}  # Strip whitespace

def fetch_data(manual_id: str) -> Dict[str, Any]:
    """Fetch equipment manual data from external API.
    
    Args:
        manual_id: Unique identifier for the manual
    Returns:
        JSON response with manual details
    Raises:
        Exception: If API call fails
    """
    try:
        response = requests.get(f"{config.granite_api_url}/{manual_id}")
        response.raise_for_status()
        return response.json()  # Return JSON data
    except requests.exceptions.RequestException as e:
        logger.error(f"Failed to fetch data for {manual_id}: {e}")
        raise Exception("API call failed")

def transform_records(data: Dict[str, Any]) -> Dict[str, Any]:
    """Transform raw data into a suitable format for indexing.
    
    Args:
        data: Raw data from the API
    Returns:
        Transformed data ready for LlamaIndex
    """
    # Example transformation logic
    return {
        'title': data['title'],  # Keep title
        'content': data['body'],  # Extract content
        'id': data['id'],  # Include ID
    }

def save_to_db(data: Dict[str, Any]) -> None:
    """Save transformed data to the database.
    
    Args:
        data: Data to save
    Raises:
        Exception: If database operation fails
    """
    try:
        with session.begin():  # Use context manager for transaction
            # Insert data into the database
            session.execute(
                Table('manuals', MetaData(),
                    Column('id', Integer, primary_key=True),
                    Column('title', String),
                    Column('content', String)
                ).insert().values(data)
            )
        logger.info(f"Data saved for {data['id']}")  # Log the save operation
    except Exception as e:
        logger.error(f"Failed to save data: {e}")
        raise Exception("Database save failed")

def call_api_for_indexing(data: Dict[str, Any]) -> None:
    """Call LlamaIndex API to index the data.
    
    Args:
        data: Data to be indexed
    Raises:
        Exception: If indexing fails
    """
    try:
        response = requests.post(config.llama_api_url, json=data)
        response.raise_for_status()
        logger.info(f"Indexed {data['id']} successfully")
    except requests.exceptions.RequestException as e:
        logger.error(f"Failed to index {data['id']}: {e}")
        raise Exception("Indexing API call failed")

def process_batch(manual_ids: List[str]) -> None:
    """Process a batch of manuals for conversion and indexing.
    
    Args:
        manual_ids: List of manual IDs to process
    """
    for manual_id in manual_ids:
        try:
            # Fetch and validate data
            raw_data = fetch_data(manual_id)
            validate_input(raw_data)
            sanitized_data = sanitize_fields(raw_data)
            transformed_data = transform_records(sanitized_data)
            # Save and index
            save_to_db(transformed_data)
            call_api_for_indexing(transformed_data)
        except Exception as e:
            logger.error(f"Error processing {manual_id}: {e}")
            continue  # Continue with the next manual

if __name__ == '__main__':
    # Example usage
    manual_ids = ['MNL001', 'MNL002', 'MNL003']  # List of manual IDs
    process_batch(manual_ids)  # Process manuals