From 973621494a7bd0ca766f3de8b6820d76ec6ab066 Mon Sep 17 00:00:00 2001
From: Vojtech Hyza <vhyza@vhyza.eu>
Date: Fri, 14 Jun 2024 14:08:16 +0200
Subject: [PATCH] Fix typespecs

---
 lib/hnswlib_bfindex.ex | 9 +++++----
 lib/hnswlib_index.ex   | 4 ++--
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/lib/hnswlib_bfindex.ex b/lib/hnswlib_bfindex.ex
index 2fbc125..8a73db4 100644
--- a/lib/hnswlib_bfindex.ex
+++ b/lib/hnswlib_bfindex.ex
@@ -65,7 +65,7 @@ defmodule HNSWLib.BFIndex do
   """
   @spec knn_query(%T{}, Nx.Tensor.t() | binary() | [binary()], [
           {:k, pos_integer()}
-        ]) :: :ok | {:error, String.t()}
+        ]) :: {:ok, Nx.Tensor.t(), Nx.Tensor.t()} | {:error, String.t()}
   def knn_query(self, query, opts \\ [])
 
   def knn_query(self = %T{}, query, opts) when is_binary(query) do
@@ -148,8 +148,9 @@ defmodule HNSWLib.BFIndex do
   @doc """
   Get the current number of threads to use in the index.
   """
-  @spec set_num_threads(%T{}, pos_integer()) :: {:ok, integer()} | {:error, String.t()}
-  def set_num_threads(self = %T{}, num_threads) when is_integer(num_threads) and num_threads > 0 do
+  @spec set_num_threads(%T{}, pos_integer()) :: :ok | {:error, String.t()}
+  def set_num_threads(self = %T{}, num_threads)
+      when is_integer(num_threads) and num_threads > 0 do
     HNSWLib.Nif.bfindex_set_num_threads(self.reference, num_threads)
   end
 
@@ -162,7 +163,7 @@ defmodule HNSWLib.BFIndex do
 
     Path to save the index to.
   """
-  @spec save_index(%T{}, Path.t()) :: {:ok, integer()} | {:error, String.t()}
+  @spec save_index(%T{}, Path.t()) :: :ok | {:error, String.t()}
   def save_index(self = %T{}, path) when is_binary(path) do
     HNSWLib.Nif.bfindex_save_index(self.reference, path)
   end
diff --git a/lib/hnswlib_index.ex b/lib/hnswlib_index.ex
index f9aec9f..d626ab2 100644
--- a/lib/hnswlib_index.ex
+++ b/lib/hnswlib_index.ex
@@ -188,7 +188,7 @@ defmodule HNSWLib.Index do
   @doc """
   Set the number of threads to use.
   """
-  @spec set_num_threads(%T{}, integer()) :: {:ok, integer()} | {:error, String.t()}
+  @spec set_num_threads(%T{}, integer()) :: :ok | {:error, String.t()}
   def set_num_threads(self = %T{}, new_num_threads) do
     HNSWLib.Nif.index_set_num_threads(self.reference, new_num_threads)
   end
@@ -210,7 +210,7 @@ defmodule HNSWLib.Index do
 
     Path to save the index to.
   """
-  @spec save_index(%T{}, Path.t()) :: :ok
+  @spec save_index(%T{}, Path.t()) :: :ok | {:error, String.t()}
   def save_index(self = %T{}, path) when is_binary(path) do
     HNSWLib.Nif.index_save_index(self.reference, path)
   end