VRK1 commited on
Commit
563893d
·
verified ·
1 Parent(s): 99e08ef

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +126 -0
main.py ADDED
@@ -0,0 +1,126 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, Form
2
+ from fastapi.responses import HTMLResponse
3
+ import time
4
+ from ingest import ingest
5
+ from query import retrieve, rerank, answer
6
+
7
+ app = FastAPI()
8
+
9
+
10
+ @app.get("/", response_class=HTMLResponse)
11
+ def home():
12
+ return """
13
+ <!DOCTYPE html>
14
+ <html>
15
+ <head>
16
+ <title>Mini RAG App</title>
17
+ <style>
18
+ body { font-family: Arial; margin: 40px; }
19
+ textarea { width: 100%; height: 120px; }
20
+ input[type=text] { width: 100%; padding: 8px; }
21
+ button { padding: 10px 20px; margin-top: 10px; }
22
+ .box { border: 1px solid #ccc; padding: 15px; margin-top: 20px; }
23
+ .source { font-size: 14px; color: #555; margin-top: 5px; }
24
+ #upload-status { color: green; margin-top: 10px; }
25
+ </style>
26
+ </head>
27
+ <body>
28
+
29
+ <h2>Mini RAG Demo</h2>
30
+
31
+ <div class="box">
32
+ <h3>1. Upload / Paste Text</h3>
33
+ <textarea id="upload-text" placeholder="Paste document text here"></textarea>
34
+ <br>
35
+ <button onclick="uploadText()">Upload</button>
36
+ <div id="upload-status"></div>
37
+ </div>
38
+
39
+ <div class="box">
40
+ <h3>2. Ask Question</h3>
41
+ <form action="/ask" method="post">
42
+ <input type="text" name="question" placeholder="Ask a question">
43
+ <br>
44
+ <button type="submit">Ask</button>
45
+ </form>
46
+ </div>
47
+
48
+ <script>
49
+ async function uploadText() {
50
+ const text = document.getElementById('upload-text').value;
51
+ if (!text.trim()) {
52
+ alert('Please enter text to upload.');
53
+ return;
54
+ }
55
+
56
+ const formData = new FormData();
57
+ formData.append('text', text);
58
+
59
+ try {
60
+ const response = await fetch('/upload', {
61
+ method: 'POST',
62
+ body: formData
63
+ });
64
+
65
+ if (response.ok) {
66
+ document.getElementById('upload-status').innerText = '✅ Text stored in vector DB';
67
+ document.getElementById('upload-text').value = ''; // optional: clear text area
68
+ } else {
69
+ document.getElementById('upload-status').innerText = '❌ Upload failed';
70
+ }
71
+ } catch (err) {
72
+ document.getElementById('upload-status').innerText = '❌ Error: ' + err;
73
+ }
74
+ }
75
+ </script>
76
+
77
+ </body>
78
+ </html>
79
+ """
80
+
81
+
82
+ # Upload endpoint
83
+
84
+ @app.post("/upload")
85
+ def upload(text: str = Form(...)):
86
+ ingest(text)
87
+ return {"status": "success"}
88
+
89
+
90
+
91
+ # Ask endpoint
92
+
93
+ @app.post("/ask", response_class=HTMLResponse)
94
+ def ask(question: str = Form(...)):
95
+ start_time = time.time()
96
+ docs = retrieve(question)
97
+ reranked_docs = rerank(question, docs)
98
+
99
+ if not reranked_docs:
100
+ return "<h3>No relevant context found.</h3><a href='/'>Back</a>"
101
+
102
+
103
+ context_text = "\n\n".join([d["text"] for d in reranked_docs])
104
+ ans = answer(question, context_text)
105
+ elapsed = round(time.time() - start_time, 2)
106
+
107
+
108
+ html = f"""
109
+ <h2>Answer</h2>
110
+ <div class="answer-box" style="white-space: pre-wrap;">{ans}</div>
111
+
112
+
113
+ <h3>Sources</h3>
114
+ """
115
+ for i, d in enumerate(reranked_docs):
116
+ html += f"""
117
+ <div class="source">
118
+ [{i+1}] {d['metadata'].get('source', 'unknown')} | position {d['metadata'].get('position')}
119
+ </div>
120
+ """
121
+
122
+ html += f"""
123
+ <p><b>Time taken:</b> {elapsed} seconds</p>
124
+ <a href="/">Ask another question</a>
125
+ """
126
+ return html