Update fn.py

fn.py

@@ -151,7 +151,7 @@ def load_vectors():
                     vector = np.array([struct.unpack('>d', binascii.unhexlify(x+'0000000'))[0] for x in row['vector'].split(',')])
                     vectors[name].append([row['chunk'], row['output'], vector])
 
-def search(name, query_text):
+def search(name, query_text, num = 3):
     dt = datetime.datetime.now()
 
     # クエリテキストをエンベディング
@@ -169,8 +169,8 @@ def search(name, query_text):
         similarity = cosine_similarity(query_embeddings, row[2])
         similarities.append((row, similarity))
 
-    # 類似度でソートし、上位3つの結果を取得
-    top_matches = sorted(similarities, key=lambda x: x[1], reverse=True)[:3]
+    # 類似度でソートし、上位num件の結果を取得
+    top_matches = sorted(similarities, key=lambda x: x[1], reverse=True)[:num]
 
     result = ''
     for i, (row, similarity) in enumerate(top_matches, 1):