add top_k

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "starpoint"
-version = "0.6.1"
+version = "0.6.2"
 authors = [{ name = "pointable", email = "package-maintainers@pointable.ai" }]
 description = "SDK for Starpoint DB"
 readme = "README.md"

starpoint/db.py

@@ -131,6 +131,7 @@ def query(
         text_search_query: Optional[List[str]] = None,
         text_search_weight: Optional[float] = None,
         tokenizer_type: Optional[reader.TokenizerType] = None,
+        top_k: Optional[int] = None,
     ) -> Dict[Any, Any]:
         """Queries a collection. This could be by sql or query embeddings.
         `query()` method from [`Reader`](#reader-objects).
@@ -144,6 +145,10 @@ def query(
             query_embedding: An embedding to query against the collection using similarity search.
                 This is of the shape {"values": List[float], "dimensionality": int}
             params: values for parameterized sql
+            text_search_query: a list of strings to search for in the text column
+            text_search_weight: the weight to apply to the text search
+            tokenizer_type: the tokenizer to use for the text search
+            top_k: the number of results to return
 
         Returns:
             dict: query response json
@@ -163,6 +168,7 @@ def query(
             text_search_query=text_search_query,
             text_search_weight=text_search_weight,
             tokenizer_type=tokenizer_type,
+            top_k=top_k,
         )
 
     def infer_schema(

starpoint/reader.py

@@ -30,6 +30,7 @@
 class TokenizerType(Enum):
     LLAMA2 = "llama2"
     ENSTEM = "en_stem"
+    NAIVE = "naive"
 
 
 class Reader(object):
@@ -54,6 +55,7 @@ def query(
         text_search_query: Optional[List[str]] = None,
         text_search_weight: Optional[float] = None,
         tokenizer_type: Optional[TokenizerType] = None,
+        top_k: Optional[int] = None,
     ) -> Dict[Any, Any]:
         """Queries a collection. This could be by sql or query embeddings.
 
@@ -67,6 +69,7 @@ def query(
             params: values for parameterized sql
             text_search_weight: weight for text search
             tokenizer_type: the type of tokenizer used to perform full text search
+            top_k: the number of results to return
 
         Returns:
             dict: query response json
@@ -101,6 +104,7 @@ def query(
             text_search_query=text_search_query,
             text_search_weight=text_search_weight,
             tokenizer_type=tokenizer_type,
+            top_k=top_k,
         )
         try:
             response = requests.post(