Use correct k.

definitive-io · Sep 8, 2023 · fd9d4a6 · fd9d4a6
1 parent e41a463
commit fd9d4a6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/code_indexer_loop/__init__.py b/code_indexer_loop/__init__.py
@@ -1 +1 @@
-__version__ = "0.2.0"
+__version__ = "0.2.1"
diff --git a/code_indexer_loop/api.py b/code_indexer_loop/api.py
@@ -51,14 +51,14 @@ def __init__(
 
     def query(self, query: str, k=10) -> str:
         return "\n".join(
-            [node_with_score.node.text for node_with_score in self.index.as_retriever(k=k).retrieve(query)]
+            [node_with_score.node.text for node_with_score in self.index.as_retriever(similarity_top_k=k).retrieve(query)]
         )
 
     def query_nodes(self, query: str, k=10) -> list[NodeWithScore]:
-        return self.index.as_retriever(k=k).retrieve(query)
+        return self.index.as_retriever(similarity_top_k=k).retrieve(query)
 
     def query_documents(self, query: str, k=10) -> list[dict[str, str]]:
-        nodes = self.index.as_retriever(k=k).retrieve(query)
+        nodes = self.index.as_retriever(similarity_top_k=k).retrieve(query)
         files = [node_with_score.node.metadata["file"] for node_with_score in nodes]
         # Deduplicate files, preserving order
         files = list(dict.fromkeys(files))