Add search and upsert records

codegen/build-oas.sh

@@ -134,6 +134,44 @@ remove_shared_classes() {
 	done
 }
 
+# Generated Python code attempts to internally map OpenAPI fields that begin
+# with "_" to a non-underscored alternative. Along with a polymorphic object,
+# this causes collisions and headaches. We massage the generated models to
+# maintain the original field names from the OpenAPI spec and circumvent
+# the remapping behavior as this is simpler for now than creating a fully
+# custom java generator class.
+clean_oas_underscore_manipulation() {
+	temp_file="$(mktemp)"
+
+	db_data_destination="${destination}/db_data"
+
+	# echo "Cleaning up upsert_record.py"
+	sed -i '' \
+	-e "s/'id'/'_id'/g" \
+	-e 's/self.id/self._id/g' \
+	-e 's/id (/_id (/g' \
+	-e 's/= id/= _id/g' \
+	-e 's/id,/_id,/g' \
+	-e "s/'vector\'/'_vector'/g" \
+	-e "s/'embed\'/'_embed'/g" \
+	-e 's/vector (/_vector (/g' \
+	-e 's/embed (/_embed (/g' \
+	"${db_data_destination}/model/upsert_record.py"
+
+	# echo "Cleaning up hit.py"
+	sed -i '' \
+	-e "s/'id'/'_id'/g" \
+	-e "s/'score'/'_score'/g" \
+	-e 's/ id, score,/ _id, _score,/g' \
+	-e 's/id (/_id (/g' \
+	-e 's/score (/_score (/g' \
+	-e 's/self.id/self._id/g' \
+	-e 's/self.score/self._score/g' \
+	-e 's/= id/= _id/g' \
+	-e 's/= score/= _score/g' \
+	"${db_data_destination}/model/hit.py"
+}
+
 update_apis_repo
 update_templates_repo
 verify_spec_version $version
@@ -144,6 +182,7 @@ mkdir -p "${destination}"
 for module in "${modules[@]}"; do
 	generate_client $module
 done
+clean_oas_underscore_manipulation
 
 # This also exists in the generated module code, but we need to reference it
 # in the pinecone.openapi_support package as well without creating a circular

pinecone/config/__init__.py

@@ -5,4 +5,4 @@
 from .pinecone_config import PineconeConfig
 
 if os.getenv("PINECONE_DEBUG") is not None:
-    logging.basicConfig(level=logging.DEBUG)
+    logging.getLogger("pinecone").setLevel(level=logging.DEBUG)

pinecone/control/pinecone.py

@@ -12,7 +12,12 @@
 from pinecone.openapi_support.api_client import ApiClient
 
 
-from pinecone.utils import normalize_host, setup_openapi_client, build_plugin_setup_client
+from pinecone.utils import (
+    normalize_host,
+    setup_openapi_client,
+    build_plugin_setup_client,
+    convert_enum_to_string,
+)
 from pinecone.core.openapi.db_control.models import (
     CreateCollectionRequest,
     CreateIndexForModelRequest,
@@ -156,17 +161,12 @@ def __parse_tags(self, tags: Optional[Dict[str, str]]) -> IndexTags:
     def __parse_deletion_protection(
         self, deletion_protection: Union[DeletionProtection, str]
     ) -> DeletionProtectionModel:
-        deletion_protection = self.__parse_enum_to_string(deletion_protection)
+        deletion_protection = convert_enum_to_string(deletion_protection)
         if deletion_protection in ["enabled", "disabled"]:
             return DeletionProtectionModel(deletion_protection)
         else:
             raise ValueError("deletion_protection must be either 'enabled' or 'disabled'")
 
-    def __parse_enum_to_string(self, value: Union[Enum, str]) -> str:
-        if isinstance(value, Enum):
-            return value.value
-        return value
-
     def __parse_index_spec(self, spec: Union[Dict, ServerlessSpec, PodSpec]) -> IndexSpec:
         if isinstance(spec, dict):
             if "serverless" in spec:
@@ -227,9 +227,9 @@ def create_index(
         tags: Optional[Dict[str, str]] = None,
     ) -> IndexModel:
         if metric is not None:
-            metric = self.__parse_enum_to_string(metric)
+            metric = convert_enum_to_string(metric)
         if vector_type is not None:
-            vector_type = self.__parse_enum_to_string(vector_type)
+            vector_type = convert_enum_to_string(vector_type)
         if deletion_protection is not None:
             dp = self.__parse_deletion_protection(deletion_protection)
         else:
@@ -270,8 +270,8 @@ def create_index_for_model(
         deletion_protection: Optional[Union[DeletionProtection, str]] = DeletionProtection.DISABLED,
         timeout: Optional[int] = None,
     ) -> IndexModel:
-        cloud = self.__parse_enum_to_string(cloud)
-        region = self.__parse_enum_to_string(region)
+        cloud = convert_enum_to_string(cloud)
+        region = convert_enum_to_string(region)
         if deletion_protection is not None:
             dp = self.__parse_deletion_protection(deletion_protection)
         else:
@@ -290,7 +290,7 @@ def create_index_for_model(
             parsed_embed = {}
             for key, value in embed.items():
                 if isinstance(value, Enum):
-                    parsed_embed[key] = value.value
+                    parsed_embed[key] = convert_enum_to_string(value)
                 else:
                     parsed_embed[key] = value
 
@@ -325,7 +325,7 @@ def is_ready():
             # Wait indefinitely
             while not is_ready():
                 logger.debug(
-                    f"Waiting for index {name} to be ready. Total wait time: {total_wait_time}"
+                    f"Waiting for index {name} to be ready. Total wait time {total_wait_time} seconds."
                 )
                 total_wait_time += 5
                 time.sleep(5)

pinecone/core/openapi/db_data/model/hit.py

@@ -78,8 +78,8 @@ def openapi_types():
                 and the value is attribute type.
         """
         return {
-            "id": (str,),  # noqa: E501
-            "score": (float,),  # noqa: E501
+            "_id": (str,),  # noqa: E501
+            "_score": (float,),  # noqa: E501
             "fields": ({str: (bool, dict, float, int, list, str, none_type)},),  # noqa: E501
         }
 
@@ -88,8 +88,8 @@ def discriminator():
         return None
 
     attribute_map = {
-        "id": "_id",  # noqa: E501
-        "score": "_score",  # noqa: E501
+        "_id": "_id",  # noqa: E501
+        "_score": "_score",  # noqa: E501
         "fields": "fields",  # noqa: E501
     }
 
@@ -99,12 +99,12 @@ def discriminator():
 
     @classmethod
     @convert_js_args_to_python_args
-    def _from_openapi_data(cls, id, score, fields, *args, **kwargs):  # noqa: E501
+    def _from_openapi_data(cls, _id, _score, fields, *args, **kwargs):  # noqa: E501
         """Hit - a model defined in OpenAPI
 
         Args:
-            id (str): The record id of the search hit.
-            score (float): The similarity score of the returned record.
+            _id (str): The record id of the search hit.
+            _score (float): The similarity score of the returned record.
             fields ({str: (bool, dict, float, int, list, str, none_type)}): The selected record fields associated with the search hit.
 
         Keyword Args:
@@ -163,8 +163,8 @@ def _from_openapi_data(cls, id, score, fields, *args, **kwargs):  # noqa: E501
         self._configuration = _configuration
         self._visited_composed_classes = _visited_composed_classes + (self.__class__,)
 
-        self.id = id
-        self.score = score
+        self._id = _id
+        self._score = _score
         self.fields = fields
         for var_name, var_value in kwargs.items():
             if (
@@ -190,12 +190,12 @@ def _from_openapi_data(cls, id, score, fields, *args, **kwargs):  # noqa: E501
     )
 
     @convert_js_args_to_python_args
-    def __init__(self, id, score, fields, *args, **kwargs):  # noqa: E501
+    def __init__(self, _id, _score, fields, *args, **kwargs):  # noqa: E501
         """Hit - a model defined in OpenAPI
 
         Args:
-            id (str): The record id of the search hit.
-            score (float): The similarity score of the returned record.
+            _id (str): The record id of the search hit.
+            _score (float): The similarity score of the returned record.
             fields ({str: (bool, dict, float, int, list, str, none_type)}): The selected record fields associated with the search hit.
 
         Keyword Args:
@@ -252,8 +252,8 @@ def __init__(self, id, score, fields, *args, **kwargs):  # noqa: E501
         self._configuration = _configuration
         self._visited_composed_classes = _visited_composed_classes + (self.__class__,)
 
-        self.id = id
-        self.score = score
+        self._id = _id
+        self._score = _score
         self.fields = fields
         for var_name, var_value in kwargs.items():
             if (

pinecone/core/openapi/db_data/model/upsert_record.py

@@ -78,15 +78,15 @@ def openapi_types():
                 and the value is attribute type.
         """
         return {
-            "id": (str,)  # noqa: E501
+            "_id": (str,)  # noqa: E501
         }
 
     @cached_property
     def discriminator():
         return None
 
     attribute_map = {
-        "id": "_id"  # noqa: E501
+        "_id": "_id"  # noqa: E501
     }
 
     read_only_vars = {}
@@ -95,11 +95,11 @@ def discriminator():
 
     @classmethod
     @convert_js_args_to_python_args
-    def _from_openapi_data(cls, id, *args, **kwargs):  # noqa: E501
+    def _from_openapi_data(cls, _id, *args, **kwargs):  # noqa: E501
         """UpsertRecord - a model defined in OpenAPI
 
         Args:
-            id (str): The unique ID of the record to upsert. Note that `id` can be used as an alias for `_id`.
+            _id (str): The unique ID of the record to upsert. Note that `id` can be used as an alias for `_id`.
 
         Keyword Args:
             _check_type (bool): if True, values for parameters in openapi_types
@@ -157,7 +157,7 @@ def _from_openapi_data(cls, id, *args, **kwargs):  # noqa: E501
         self._configuration = _configuration
         self._visited_composed_classes = _visited_composed_classes + (self.__class__,)
 
-        self.id = id
+        self._id = _id
         for var_name, var_value in kwargs.items():
             if (
                 var_name not in self.attribute_map
@@ -182,11 +182,11 @@ def _from_openapi_data(cls, id, *args, **kwargs):  # noqa: E501
     )
 
     @convert_js_args_to_python_args
-    def __init__(self, id, *args, **kwargs):  # noqa: E501
+    def __init__(self, _id, *args, **kwargs):  # noqa: E501
         """UpsertRecord - a model defined in OpenAPI
 
         Args:
-            id (str): The unique ID of the record to upsert. Note that `id` can be used as an alias for `_id`.
+            _id (str): The unique ID of the record to upsert. Note that `id` can be used as an alias for `_id`.
 
         Keyword Args:
             _check_type (bool): if True, values for parameters in openapi_types
@@ -242,7 +242,7 @@ def __init__(self, id, *args, **kwargs):  # noqa: E501
         self._configuration = _configuration
         self._visited_composed_classes = _visited_composed_classes + (self.__class__,)
 
-        self.id = id
+        self._id = _id
         for var_name, var_value in kwargs.items():
             if (
                 var_name not in self.attribute_map

pinecone/data/dataclasses/__init__.py

@@ -1,3 +1,6 @@
 from .sparse_values import SparseValues
 from .vector import Vector
 from .fetch_response import FetchResponse
+from .search_query import SearchQuery
+from .search_query_vector import SearchQueryVector
+from .search_rerank import SearchRerank

pinecone/data/dataclasses/search_query.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass
+from typing import Optional, Any, Dict, Union
+from .search_query_vector import SearchQueryVector
+from ..types.search_query_vector_typed_dict import SearchQueryVectorTypedDict
+
+
+@dataclass
+class SearchQuery:
+    """
+    SearchQuery represents the query when searching within a specific namespace.
+    """
+
+    inputs: Dict[str, Any]
+    """
+    The input data to search with.
+    Required.
+    """
+
+    top_k: int
+    """
+    The number of results to return with each search.
+    Required.
+    """
+
+    filter: Optional[Dict[str, Any]] = None
+    """
+    The filter to apply to the search.
+    Optional.
+    """
+
+    vector: Optional[Union[SearchQueryVectorTypedDict, SearchQueryVector]] = None
+    """
+    The vector values to search with. If provided, it overwrites the inputs.
+    """
+
+    id: Optional[str] = None
+    """
+    The unique ID of the vector to be used as a query vector.
+    """
+
+    def __post_init__(self):
+        """
+        Converts `vector` to a `SearchQueryVectorTypedDict` instance if an enum is provided.
+        """
+        if isinstance(self.vector, SearchQueryVector):
+            self.vector = self.vector.as_dict()
+
+    def as_dict(self) -> Dict[str, Any]:
+        """
+        Returns the SearchQuery as a dictionary.
+        """
+        d = {
+            "inputs": self.inputs,
+            "top_k": self.top_k,
+            "filter": self.filter,
+            "vector": self.vector,
+            "id": self.id,
+        }
+        return {k: v for k, v in d.items() if v is not None}

pinecone/data/dataclasses/search_query_vector.py

@@ -0,0 +1,38 @@
+from dataclasses import dataclass
+from typing import Optional, List
+
+
+@dataclass
+class SearchQueryVector:
+    """
+    SearchQueryVector represents the vector values used to query.
+    """
+
+    values: Optional[List[float]] = None
+    """
+    The vector data included in the search request.
+    Optional.
+    """
+
+    sparse_values: Optional[List[float]] = None
+    """
+    The sparse embedding values to search with.
+    Optional.
+    """
+
+    sparse_indices: Optional[List[int]] = None
+    """
+    The sparse embedding indices to search with.
+    Optional.
+    """
+
+    def as_dict(self) -> dict:
+        """
+        Returns the SearchQueryVector as a dictionary.
+        """
+        d = {
+            "values": self.values,
+            "sparse_values": self.sparse_values,
+            "sparse_indices": self.sparse_indices,
+        }
+        return {k: v for k, v in d.items() if v is not None}