Update retrieval.py

retrieval.py

@@ -1,20 +1,27 @@
 import json
 import numpy as np
-from db import cursor
+from db import get_db_connection
 from embeddings import embedding_model, cosine_similarity
 
 def retrieve_top_chunks(question, k=3):
     q_embedding = embedding_model.encode(question)
 
-    cursor.execute(
-        "SELECT document, chunk_id, text, embedding FROM chunks"
-    )
+    conn = get_db_connection()
+    cursor = conn.cursor()
+    
+    try:
+        cursor.execute(
+            "SELECT document, chunk_id, text, embedding FROM chunks"
+        )
 
-    results = []
-    for doc, cid, text, emb_json in cursor.fetchall():
-        emb = np.array(json.loads(emb_json))
-        score = cosine_similarity(q_embedding, emb)
-        results.append((score, doc, cid, text))
+        results = []
+        for doc, cid, text, emb_json in cursor.fetchall():
+            emb = np.array(json.loads(emb_json))
+            score = cosine_similarity(q_embedding, emb)
+            results.append((score, doc, cid, text))
 
-    results.sort(key=lambda x: x[0], reverse=True)
-    return results[:k]
+        results.sort(key=lambda x: x[0], reverse=True)
+        return results[:k]
+    finally:
+        cursor.close()
+        conn.close()