rag-solution/services/rag/langchain/cli.py

import click
from loguru import logger
import os
from pathlib import Path


# Configure logging to output to both file and stdout as specified in requirements
def setup_logging():
    # Create logs directory if it doesn't exist
    logs_dir = Path("logs")
    logs_dir.mkdir(exist_ok=True)

    # Add file logging with rotation
    logger.add("logs/dev.log", rotation="10 MB", retention="10 days")


@click.group()
def cli():
    """Main CLI group"""
    setup_logging()
    pass


@cli.command(name="ping", help="Ping command that outputs pong")
def ping():
    """Ping command that outputs pong"""
    logger.info("Ping command executed")
    click.echo("pong")


@cli.command(name="enrich", help="Load documents from data directory and store in vector database")
@click.option('--data-dir', default="../../../data", help="Path to the data directory")
@click.option('--collection-name', default="documents", help="Name of the vector store collection")
def enrich(data_dir, collection_name):
    """Load documents from data directory and store in vector database"""
    logger.info(f"Starting enrichment process for directory: {data_dir}")

    try:
        # Import here to avoid circular dependencies
        from vector_storage import initialize_vector_store
        from enrichment import run_enrichment_process

        # Initialize vector store
        vector_store = initialize_vector_store(collection_name=collection_name)

        # Run enrichment process
        run_enrichment_process(vector_store, data_dir=data_dir)

        logger.info("Enrichment process completed successfully!")
        click.echo("Documents have been successfully loaded into the vector store.")

    except Exception as e:
        logger.error(f"Error during enrichment process: {str(e)}")
        click.echo(f"Error: {str(e)}")


if __name__ == "__main__":
    cli()