Typing fixes

Author: cevian

nbs/00_vector.ipynb

@@ -68,8 +68,9 @@
     "import asyncpg\n",
     "import uuid\n",
     "from pgvector.asyncpg import register_vector\n",
-    "from typing import (List, Optional, Union, Dict, Tuple)\n",
-    "import json "
+    "from typing import (List, Optional, Union, Dict, Tuple, Any)\n",
+    "import json\n",
+    "import numpy as np "
    ]
   },
   {
@@ -192,7 +193,7 @@
     "       return (query, [id])\n",
     "\n",
     "    def delete_by_metadata_query (self, filter: Union[Dict[str, str], List[Dict[str, str]]]) -> Tuple[str, List]:\n",
-    "        params = []\n",
+    "        params: List[Any] = []\n",
     "        (where, params) = self._where_clause_for_filter(params, filter)\n",
     "        query = \"DELETE FROM {table_name} WHERE {where};\".format(table_name=self._quote_ident(self.table_name), where=where)\n",
     "        return (query, params) \n",
@@ -247,7 +248,7 @@
     "\n",
     "        return (where, params) \n",
     "\n",
-    "    def search_query(self, query_embedding: Optional[List[float]], limit: int=10, filter: Optional[Union[Dict[str, str], List[Dict[str, str]]]] = None) -> Tuple[str, List]:\n",
+    "    def search_query(self, query_embedding: Optional[Union[List[float], np.ndarray]], limit: int=10, filter: Optional[Union[Dict[str, str], List[Dict[str, str]]]] = None) -> Tuple[str, List]:\n",
     "        \"\"\"\n",
     "        Generates a similarity query.\n",
     "\n",
@@ -259,7 +260,7 @@
     "        Returns:\n",
     "            Tuple[str, List]: A tuple containing the query and parameters.\n",
     "        \"\"\"\n",
-    "        params = []\n",
+    "        params: List[Any] = []\n",
     "        if query_embedding is not None:\n",
     "            distance = \"embedding {op} ${index}\".format(op=self.distance_type, index=len(params)+1)\n",
     "            params = params + [query_embedding]\n",
@@ -816,7 +817,7 @@
    "source": [
     "#| export\n",
     "class Sync:\n",
-    "    translated_queries = {}\n",
+    "    translated_queries: Dict[str, str] = {}\n",
     "    \n",
     "    def __init__(\n",
     "        self,\n",
@@ -1033,9 +1034,11 @@
     "            List: List of similar records.\n",
     "        \"\"\"\n",
     "        if query_embedding is not None:\n",
-    "            query_embedding = np.array(query_embedding)\n",
+    "            query_embedding_np = np.array(query_embedding)\n",
+    "        else:\n",
+    "           query_embedding_np = None \n",
     "            \n",
-    "        (query, params) = self.builder.search_query(query_embedding, limit, filter)\n",
+    "        (query, params) = self.builder.search_query(query_embedding_np, limit, filter)\n",
     "        query, params = self._translate_to_pyformat(query, params)\n",
     "        with self.connect() as conn:\n",
     "            with conn.cursor() as cur:\n",

timescale_vector/client.py

@@ -8,8 +8,9 @@
 import asyncpg
 import uuid
 from pgvector.asyncpg import register_vector
-from typing import (List, Optional, Union, Dict, Tuple)
-import json 
+from typing import (List, Optional, Union, Dict, Tuple, Any)
+import json
+import numpy as np 
 
 # %% ../nbs/00_vector.ipynb 7
 SEARCH_RESULT_ID_IDX = 0
@@ -117,7 +118,7 @@ def delete_by_ids_query(self, id: List[uuid.UUID]) ->  Tuple[str, List]:
        return (query, [id])
 
     def delete_by_metadata_query (self, filter: Union[Dict[str, str], List[Dict[str, str]]]) -> Tuple[str, List]:
-        params = []
+        params: List[Any] = []
         (where, params) = self._where_clause_for_filter(params, filter)
         query = "DELETE FROM {table_name} WHERE {where};".format(table_name=self._quote_ident(self.table_name), where=where)
         return (query, params) 
@@ -172,7 +173,7 @@ def _where_clause_for_filter(self, params: List, filter: Optional[Union[Dict[str
 
         return (where, params) 
 
-    def search_query(self, query_embedding: Optional[List[float]], limit: int=10, filter: Optional[Union[Dict[str, str], List[Dict[str, str]]]] = None) -> Tuple[str, List]:
+    def search_query(self, query_embedding: Optional[Union[List[float], np.ndarray]], limit: int=10, filter: Optional[Union[Dict[str, str], List[Dict[str, str]]]] = None) -> Tuple[str, List]:
         """
         Generates a similarity query.
 
@@ -184,7 +185,7 @@ def search_query(self, query_embedding: Optional[List[float]], limit: int=10, fi
         Returns:
             Tuple[str, List]: A tuple containing the query and parameters.
         """
-        params = []
+        params: List[Any] = []
         if query_embedding is not None:
             distance = "embedding {op} ${index}".format(op=self.distance_type, index=len(params)+1)
             params = params + [query_embedding]
@@ -397,7 +398,7 @@ async def search(self,
 
 # %% ../nbs/00_vector.ipynb 20
 class Sync:
-    translated_queries = {}
+    translated_queries: Dict[str, str] = {}
 
     def __init__(
         self,
@@ -614,9 +615,11 @@ def search(self, query_embedding: Optional[List[float]]=None, limit: int=10, fil
             List: List of similar records.
         """
         if query_embedding is not None:
-            query_embedding = np.array(query_embedding)
+            query_embedding_np = np.array(query_embedding)
+        else:
+           query_embedding_np = None 
 
-        (query, params) = self.builder.search_query(query_embedding, limit, filter)
+        (query, params) = self.builder.search_query(query_embedding_np, limit, filter)
         query, params = self._translate_to_pyformat(query, params)
         with self.connect() as conn:
             with conn.cursor() as cur: