import os from pathlib import Path import click from dotenv import load_dotenv from loguru import logger # Load environment variables load_dotenv() # 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_langchain", 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 enrichment import run_enrichment_process from vector_storage import initialize_vector_store # 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)}") @cli.command( name="retrieve", help="Retrieve documents from vector database based on a query", ) @click.argument("query") @click.option( "--collection-name", default="documents_langchain", help="Name of the vector store collection", ) @click.option( "--top-k", default=5, help="Number of documents to retrieve", ) def retrieve(query, collection_name, top_k): """Retrieve documents from vector database based on a query""" logger.info(f"Starting retrieval process for query: {query}") click.echo( "WARNING: Retrieval disabled, since it is no longer relevant for the testing of the retrieving feature. Use chat with agent instead. xoxo" ) @cli.command( name="chat", help="Start an interactive chat session with the RAG agent", ) @click.option( "--collection-name", default="documents_langchain", help="Name of the vector store collection", ) @click.option( "--model", default=None, help="Name of the Ollama model to use for chat", ) def chat(collection_name, model): """Start an interactive chat session with the RAG agent""" logger.info("Starting chat session with RAG agent") try: # Import here to avoid circular dependencies and only when needed from agent import run_chat_loop click.echo("Initializing chat agent...") click.echo("Type 'quit' or 'exit' to end the conversation.\n") # Run the interactive chat loop run_chat_loop(collection_name=collection_name, llm_model=model) logger.info("Chat session ended") except Exception as e: logger.error(f"Error during chat session: {str(e)}") click.echo(f"Error: {str(e)}") if __name__ == "__main__": cli()