puuuuush

lib.py

@@ -59,11 +59,20 @@ def prompt_embedding(prompt):
     )
     cur = conn.cursor()
 
-    cur.execute(f"select id, filepath, embedding <-> %s::vector as distance from dbtable order by distance limit 3;", (embedding["embedding"],))
+    cur.execute(f"select id, filepath, embedding <-> %s::vector as distance from dbtable order by distance limit 20;", (embedding["embedding"],))
 
     result = cur.fetchall()
 
-    res = get_data_from_mongo(result)
+    initial_distance = result[0][2]
+    new_result = []
+
+    for res in result:
+        current_distance = res[2]
+        if abs(current_distance - initial_distance) < 1:
+            new_result.append(res)
+
+
+    res = get_data_from_mongo(new_result)
 
     print(ollama_conn.ask_ollama(prompt, res))