SachetAI-pathway/run.py at main · saanidhi-git/SachetAI-pathway · 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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
#!/usr/bin/env python3
import argparse
import sys

from config import settings


def banner():
    print("""
========================================
  NDRF DISASTER RESPONSE ASSISTANT
  Maharashtra Real-time Information
  Powered by Gemini AI
========================================
    """)


def cli_mode():
    from rag import RAGEngine

    banner()

    if not settings.check():
        print("Set GEMINI_API_KEY in .env file")
        print("Get free key:  https://aistudio.google.com/app/apikey")
        sys.exit(1)

    print("Starting...")
    eng = RAGEngine()
    eng.start()

    print("\nReady.  Type question or 'quit' to exit.")
    print("\nSamples:")
    print("  - Current Sholapur flood status?")
    print("  - Landslide evacuation routes?")
    print("  - NDRF helpline?\n")

    while True:
        try:
            inp = input("> ").strip()

            if not inp:
                continue

            if inp.lower() in ["quit", "exit", "q"]:
                break

            if inp.lower() == "status":
                s = eng.status()
                print(f"Docs: {s['docs']} | Updated: {s['updated']}")
                continue

            if inp.lower() == "refresh":
                r = eng.refresh()
                print(f"Refreshed: {r['docs']} docs")
                continue

            print("\nProcessing.. .\n")
            r = eng.ask(inp)

            print("-" * 40)
            print(r["answer"])
            print("-" * 40)
            print(f"Time: {r['time']}\n")

        except KeyboardInterrupt:
            break
        except Exception as e:
            print(f"Error: {e}")

    eng.stop()
    print("Bye!")


def server_mode():
    from api. server import serve

    banner()
    settings.check()

    print(f"Server:  http://{settings.HOST}:{settings. PORT}")
    print(f"Docs: http://{settings.HOST}:{settings.PORT}/docs")
    print("Ctrl+C to stop\n")

    serve()


def test_mode():
    from rag import RAGEngine

    print("Testing...")

    if not settings. GEMINI_API_KEY:
        print("FAIL: No API key")
        print("Get free key: https://aistudio.google.com/app/apikey")
        sys.exit(1)

    eng = RAGEngine()
    eng.pipe.refresh()

    q = "Flood status in Sholapur?"
    r = eng.ask(q)

    print(f"Query: {q}")
    print(f"Response: {len(r['answer'])} chars")
    print(f"Docs: {r['docs']}")
    print("PASS")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--cli", action="store_true")
    parser.add_argument("--server", action="store_true")
    parser.add_argument("--test", action="store_true")

    args = parser.parse_args()

    if args.test:
        test_mode()
    elif args.cli:
        cli_mode()
    else:
        server_mode()


if __name__ == "__main__":
    main()