Skip to content

Commit a42e834

Browse files
krishnemdkrishnemd
authored
Merge pull request #12 from mahendra-dhakal/pagination
pagination added for query results
2 parents ca54bf9 + 0da4472 commit a42e834

4 files changed

Lines changed: 112 additions & 89 deletions

File tree

‎src/database.py‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -240,18 +240,18 @@ async def insert_embeddings(self, table_name: str, rows: List[Tuple[str, str, li
240240
logger.error(f"Error inserting embeddings into '{table_name}': {str(e)}")
241241
raise
242242

243-
async def get_similar_rows(self, query_embedding: str, num_of_rows: int) -> List[Any]:
244-
"""Retrieve similar rows based on vector embedding similarity."""
243+
async def get_similar_rows(self, query_embedding: str, num_of_rows: int, offset: int = 0) -> List[Any]:
244+
"""Retrieve similar rows based on vector embedding similarity with pagination."""
245245
sql = f"""
246246
SELECT table_name, row_data, embedding <=> CAST($1 AS vector) AS similarity
247247
FROM text_embeddings
248248
ORDER BY similarity ASC
249-
LIMIT $2
249+
LIMIT $2 OFFSET $3
250250
"""
251251
try:
252252
async with self._conn.acquire() as conn:
253253
async with conn.transaction():
254-
results = await conn.fetch(sql, query_embedding, num_of_rows)
254+
results = await conn.fetch(sql, query_embedding, num_of_rows, offset)
255255
return results
256256
except Exception as e:
257257
logger.error(f"Error retrieving similar rows: {str(e)}")
@@ -317,4 +317,3 @@ async def get_schema(self) -> dict:
317317
except Exception as e:
318318
logger.error(f"Error fetching schema: {str(e)}")
319319
raise
320-

‎src/routes.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import re
22
import logging
33
import uuid
4-
from fastapi import FastAPI, Request, Form, HTTPException, Depends
4+
from fastapi import FastAPI, Request, Form, HTTPException, Depends, Query
55
from fastapi.responses import HTMLResponse
66
from fastapi.templating import Jinja2Templates
77
from typing import List, Any
@@ -182,4 +182,26 @@ async def submit_feedback(request: Request, query_token: str = Form(...), feedba
182182
"message": f"Error submitting feedback: {str(e)}",
183183
"type": "error"
184184
}
185+
)
186+
187+
@app.post("/get-similar-rows", response_class=HTMLResponse)
188+
async def get_similar_rows_endpoint(
189+
request: Request,
190+
user_query: str = Form(...),
191+
page: int = Query(1, ge=1),
192+
page_size: int = Query(10, ge=1, le=100),
193+
db: DatabaseManager = Depends(get_db)
194+
):
195+
"""Endpoint to fetch similar rows with pagination."""
196+
try:
197+
formatted_rows, _ = await get_similar_rows_from_vector(db, user_query, VECTOR_ROWS_IN_PROMPT, page, page_size)
198+
return templates.TemplateResponse(
199+
"text-to-sql.html",
200+
{"request": request, "similar_rows": formatted_rows, "page": page, "page_size": page_size}
201+
)
202+
except Exception as e:
203+
logger.error(f"Error in get_similar_rows_endpoint: {e}")
204+
return templates.TemplateResponse(
205+
"text-to-sql.html",
206+
{"request": request, "message": f"Error fetching similar rows: {e}", "type": "error"}
185207
)

‎src/vector.py‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,22 @@
55

66
logger = logging.getLogger(__name__)
77

8-
async def get_similar_rows_from_vector(db: DatabaseManager, user_query: str, num_of_rows: int = VECTOR_ROWS_IN_PROMPT) -> tuple:
9-
"""Fetch similar rows using vector embeddings synchronously."""
8+
async def get_similar_rows_from_vector(db: DatabaseManager, user_query: str, num_of_rows: int = VECTOR_ROWS_IN_PROMPT, page: int = 1, page_size: int = 10) -> tuple:
9+
"""Fetch similar rows using vector embeddings synchronously with pagination."""
1010
try:
1111
embed_model = SentenceTransformer(SENTENCE_TRANSFORMER_MODEL)
1212
query_embedding = embed_model.encode(user_query)
1313
logger.info("Query embedding created")
1414
embedding_str = '[' + ','.join(map(str, query_embedding)) + ']'
1515
results = await db.get_similar_rows(embedding_str, num_of_rows)
1616
logger.info("Similar rows retrieved")
17-
18-
formatted_rows = "".join([f"Table: {row[0]}, Data: {row[1]}\n" for row in results])
17+
18+
# Implement pagination
19+
start_index = (page - 1) * page_size
20+
end_index = start_index + page_size
21+
paginated_results = results[start_index:end_index]
22+
23+
formatted_rows = "".join([f"Table: {row[0]}, Data: {row[1]}\n" for row in paginated_results])
1924
return formatted_rows, user_query
2025
except Exception as e:
2126
logger.error(f"Error in vector search: {e}")

‎templates/text-to-sql.html‎

Lines changed: 76 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -1,93 +1,90 @@
11
{% if message %}
2-
<div class="message {% if type == 'error' %}error{% else %}success{% endif %}">
3-
{{ message }}
4-
</div>
2+
<div class="message {% if type == 'error' %}error{% else %}success{% endif %}">
3+
{{ message }}
4+
</div>
55
{% endif %}
66

77
{% if sql_query %}
8-
<div class="query-section">
9-
<h3>Generated SQL Query</h3>
10-
<div class="query-box">{{ sql_query }}</div>
11-
12-
{% if query_token %}
13-
<div class="actions-container">
14-
<div class="action-group">
15-
<form class="action-form">
16-
<input type="hidden" name="query_token" value="{{ query_token }}">
17-
<button
18-
hx-post="/execute-sql"
19-
hx-trigger="click"
20-
hx-target="#query-result"
21-
hx-include="[name='query_token']"
22-
type="button"
23-
class="primary-button"
24-
>
25-
<span>Execute Query</span>
26-
<span class="spinner htmx-indicator"></span>
27-
</button>
28-
</form>
29-
</div>
8+
<div class="query-section">
9+
<h3>Generated SQL Query</h3>
10+
<div class="query-box">{{ sql_query }}</div>
3011

31-
<div class="feedback-group" id="feedback-group">
32-
<span class="feedback-label">Was this query correct?</span>
33-
<div class="feedback-buttons">
34-
<button
35-
hx-post="/submit-feedback"
36-
hx-target="#feedback-result"
37-
hx-include="[name='query_token']"
38-
hx-vals='{"feedback": "yes"}'
39-
class="feedback-button yes"
40-
>
41-
<svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
42-
<polyline points="20 6 9 17 4 12"></polyline>
43-
</svg>
44-
Yes
45-
</button>
46-
<button
47-
hx-post="/submit-feedback"
48-
hx-target="#feedback-result"
49-
hx-include="[name='query_token']"
50-
hx-vals='{"feedback": "no"}'
51-
class="feedback-button no"
52-
>
53-
<svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
54-
<line x1="18" y1="6" x2="6" y2="18"></line>
55-
<line x1="6" y1="6" x2="18" y2="18"></line>
56-
</svg>
57-
No
58-
</button>
59-
</div>
60-
</div>
61-
</div>
12+
{% if query_token %}
13+
<div class="actions-container">
14+
<div class="action-group">
15+
<form class="action-form">
16+
<input type="hidden" name="query_token" value="{{ query_token }}">
17+
<button hx-post="/execute-sql" hx-trigger="click" hx-target="#query-result"
18+
hx-include="[name='query_token']" type="button" class="primary-button">
19+
<span>Execute Query</span>
20+
<span class="spinner htmx-indicator"></span>
21+
</button>
22+
</form>
23+
</div>
6224

63-
<div id="feedback-result"></div>
64-
{% endif %}
25+
<div class="feedback-group" id="feedback-group">
26+
<span class="feedback-label">Was this query correct?</span>
27+
<div class="feedback-buttons">
28+
<button hx-post="/submit-feedback" hx-target="#feedback-result" hx-include="[name='query_token']"
29+
hx-vals='{"feedback": "yes"}' class="feedback-button yes">
30+
<svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
31+
<polyline points="20 6 9 17 4 12"></polyline>
32+
</svg>
33+
Yes
34+
</button>
35+
<button hx-post="/submit-feedback" hx-target="#feedback-result" hx-include="[name='query_token']"
36+
hx-vals='{"feedback": "no"}' class="feedback-button no">
37+
<svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
38+
<line x1="18" y1="6" x2="6" y2="18"></line>
39+
<line x1="6" y1="6" x2="18" y2="18"></line>
40+
</svg>
41+
No
42+
</button>
43+
</div>
44+
</div>
6545
</div>
46+
47+
<div id="feedback-result"></div>
48+
{% endif %}
49+
</div>
6650
{% endif %}
6751

6852
{% if column_names and results %}
69-
<div class="results-section">
70-
<h3>Query Results</h3>
71-
<div class="table-container">
72-
<table class="results-table">
73-
<thead>
74-
<tr>
75-
{% for column in column_names %}
76-
<th data-column="{{ column }}">{{ column }}</th>
77-
{% endfor %}
78-
</tr>
79-
</thead>
80-
<tbody>
81-
{% for row in results %}
82-
<tr>
83-
{% for value in row %}
84-
<td title="{{ value }}">{{ value }}</td>
85-
{% endfor %}
86-
</tr>
53+
<div class="results-section">
54+
<h3>Query Results</h3>
55+
<div class="table-container">
56+
<table class="results-table">
57+
<thead>
58+
<tr>
59+
{% for column in column_names %}
60+
<th data-column="{{ column }}">{{ column }}</th>
8761
{% endfor %}
88-
</tbody>
89-
</table>
90-
</div>
62+
</tr>
63+
</thead>
64+
<tbody>
65+
{% for row in results %}
66+
<tr>
67+
{% for value in row %}
68+
<td title="{{ value }}">{{ value }}</td>
69+
{% endfor %}
70+
</tr>
71+
{% endfor %}
72+
</tbody>
73+
</table>
9174
</div>
75+
</div>
9276
{% endif %}
9377

78+
{% if similar_rows %}
79+
<div class="results-section">
80+
<h3>Similar Rows</h3>
81+
<pre>{{ similar_rows }}</pre>
82+
<div class="pagination">
83+
{% if page > 1 %}
84+
<a href="?page={{ page - 1 }}&page_size={{ page_size }}" class="pagination-button">Previous</a>
85+
{% endif %}
86+
<span>Page {{ page }}</span>
87+
<a href="?page={{ page + 1 }}&page_size={{ page_size }}" class="pagination-button">Next</a>
88+
</div>
89+
</div>
90+
{% endif %}

0 commit comments

Comments
 (0)