rag-techniques/app.py at main · Rhythamtech/rag-techniques · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
import logging
from rag_techniques.rag import RAG
from rag_techniques.utils import convert_json_to_toon

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def main():
    rag = RAG()

    urls = ["https://raw.githubusercontent.com/asinghcsu/AgenticRAG-Survey/refs/heads/main/README.md",
            "https://docs.cohere.com/page/agentic-rag-mixed-data",
            "https://docs.cohere.com/docs/routing-queries-to-data-sources",
            "https://docs.cohere.com/docs/generating-parallel-queries"]

    while True:
        print("\nSelect an option:")
        print("1. Create Q&A index")
        print("2. Create summary index")
        print("3. Create chunking index")
        print("4. Query using chunking index")
        print("5. Query using Q&A index")
        print("6. Query using summary index")
        print("7. Exit")

        choice = input("Enter your choice: ")

        if choice == "1":
            docs = rag.load_documents(urls)
            chunk_docs = rag.split_documents(docs)
            rag.create_qna_index(chunk_docs)
        elif choice == "2":
            docs = rag.load_documents(urls)
            rag.create_summary_index(docs)

        elif choice == "3":
            docs = rag.load_documents(urls)
            chunk_docs = rag.split_documents(docs)
            rag.create_chunking_index(chunk_docs)

        elif choice == "4":
          user_query = input("\nEnter your query: ")
          results = rag.query_chunking_index(user_query)
          context = ""
          for res in results:
              context += convert_json_to_toon(res.page_content)
              context += "\n" + "-" * 10
          answer = rag.get_answer(user_query, context)
          print("\nAnswer:")
          print(answer)

        elif choice == "5":
            user_query = input("\nEnter your query: ")
            results = rag.query_qna_index(user_query)
            context = ""
            for res in results:
                context += convert_json_to_toon(res.page_content)
                context += "\n" + "-" * 10
            answer = rag.get_answer(user_query, context)
            print("\nAnswer:")
            print(answer)
        elif choice == "6":
            user_query = input("\nEnter your query: ")
            results = rag.query_summary_index(user_query)
            context = ""
            for res in results:
                context += convert_json_to_toon(res.page_content)
                context += "\n" + "-" * 10
            answer = rag.get_answer(user_query, context)
            print("\nAnswer:")
            print(answer)
        elif choice == "7":
            break
        else:
            print("Invalid choice. Please try again.")

if __name__ == "__main__":
    main()