Query.py

import os
import asyncio
from ylightRag.yllm.openai import gpt_4o_mini_complete, openai_embed
from ylightRag.ykg.shared_storage import initialize_pipeline_status
from dotenv import load_dotenv
from ylightRag.ylightrag import LightRAG
from ylightRag.ybase import QueryParam
from ylightRag.yutils import EmbeddingFunc

load_dotenv()

WORKING_DIR = "Cache"

async def main():
    # Initialize RAG system
    rag = LightRAG(
        llm_model_func=gpt_4o_mini_complete,
        embedding_func=EmbeddingFunc(
            embedding_dim=1536,
            max_token_size=8192,
            func=openai_embed
        ),
        working_dir=WORKING_DIR
    )

    await rag.initialize_storages()
    await initialize_pipeline_status()
    
    print("=" * 60)
    print("🤖 Interactive RAG Query System")
    print("=" * 60)
    print("Welcome! You can now ask questions about your documents.")
    print("Type 'exit' to quit the application.")
    print("Type 'help' for available commands.")
    print("-" * 60)
    
    while True:
        try:
            # Get user input
            user_query = input("\n💬 Enter your query: ").strip()
            
            # Check for exit condition
            if user_query.lower() in ['exit', 'quit', 'q']:
                print("\n👋 Goodbye! Thank you for using the RAG system.")
                break
            
            # Check for empty input
            if not user_query:
                print("⚠️ Please enter a valid query.")
                continue
            
            # Process the query
            print(f"\n🔍 Processing query: '{user_query}'")
            print("⏳ Please wait...")
            
            # Execute the query
            response = await rag.aquery(
                user_query, 
                param=QueryParam(mode="local")
            )
            
            # Display the response
            print("\n" + "=" * 60)
            print("📝 Response:")
            print("=" * 60)
            print(response)
            print("=" * 60)
            
        except KeyboardInterrupt:
            print("\n\n⚠️ Interrupted by user. Exiting...")
            break
        except Exception as e:
            print(f"\n❌ An error occurred: {str(e)}")
            print("Please try again with a different query.")
            continue

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nApplication terminated by user.")
    except Exception as e:
        print(f"Failed to start application: {str(e)}")