92 lines
2.9 KiB
Python
92 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
CLI entry point for the RAG solution using LlamaIndex and Qdrant.
|
|
"""
|
|
|
|
import click
|
|
from loguru import logger
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
def setup_logging():
|
|
"""Setup logging with loguru to file and stdout."""
|
|
# Create logs directory if it doesn't exist
|
|
logs_dir = Path("logs")
|
|
logs_dir.mkdir(exist_ok=True)
|
|
|
|
# Remove default logger to customize it
|
|
logger.remove()
|
|
|
|
# Add file handler with rotation
|
|
logger.add(
|
|
"logs/dev.log",
|
|
rotation="10 MB",
|
|
retention="10 days",
|
|
level="INFO",
|
|
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {message}"
|
|
)
|
|
|
|
# Add stdout handler
|
|
logger.add(
|
|
sys.stdout,
|
|
level="INFO",
|
|
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}",
|
|
colorize=True
|
|
)
|
|
|
|
|
|
@click.group()
|
|
@click.version_option(version='1.0.0')
|
|
def main():
|
|
"""Main CLI entry point for the RAG solution."""
|
|
setup_logging()
|
|
logger.info("Starting RAG solution CLI")
|
|
|
|
|
|
@main.command(help="Basic connectivity test that returns 'pong'")
|
|
@click.option('--verbose', '-v', is_flag=True, help="Enable verbose output")
|
|
def ping(verbose):
|
|
"""Ping command that outputs 'pong'."""
|
|
if verbose:
|
|
logger.info("Executing ping command")
|
|
click.echo("pong")
|
|
logger.info("Ping command completed successfully")
|
|
else:
|
|
click.echo("pong")
|
|
logger.info("Ping command executed")
|
|
|
|
|
|
@main.command(help="Load and process documents from the data folder into vector storage")
|
|
@click.option('--data-path', '-d', default="../../../data", help="Path to the data folder relative to current directory")
|
|
@click.option('--recursive', '-r', default=True, is_flag=True, help="Process subdirectories recursively")
|
|
@click.option('--verbose', '-v', is_flag=True, help="Enable verbose output")
|
|
def enrich(data_path, recursive, verbose):
|
|
"""Load and process documents from the data folder into vector storage."""
|
|
if verbose:
|
|
logger.enable("__main__")
|
|
|
|
logger.info(f"Starting document enrichment from: {data_path}")
|
|
logger.info(f"Recursive processing: {recursive}")
|
|
|
|
try:
|
|
# Import the enrichment module
|
|
from enrichment import enrich_documents, process_documents_from_data_folder
|
|
logger.info("Enrichment module imported successfully")
|
|
|
|
# Call the enrichment function
|
|
process_documents_from_data_folder(data_path=data_path, recursive=recursive)
|
|
|
|
logger.info("Document enrichment completed successfully")
|
|
click.echo("Document enrichment completed successfully")
|
|
except ImportError as e:
|
|
logger.error(f"Failed to import enrichment module: {e}")
|
|
click.echo(f"Error: Could not import enrichment module: {e}")
|
|
except Exception as e:
|
|
logger.error(f"Error during document enrichment: {e}")
|
|
click.echo(f"Error during document enrichment: {e}")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|