Merge pull request #49 from epsilla-cloud/dev

add facets support

Author: richard-epsilla

pyepsilla/__init__.py

@@ -1,7 +1,4 @@
 #!/usr/bin/env python
 # -*- coding:utf-8 -*-
 
-from . import vectordb
-from . import cloud
-from . import enterprise
-from . import utils
+from . import cloud, enterprise, utils, vectordb

pyepsilla/cloud/client.py

@@ -7,12 +7,13 @@
 import pprint
 import socket
 from typing import Optional, Union
-from ..utils.search_engine import SearchEngine
 
 import requests
 import sentry_sdk
 from pydantic import BaseModel, Field, constr
 
+from ..utils.search_engine import SearchEngine
+
 requests.packages.urllib3.disable_warnings()
 
 
@@ -33,53 +34,57 @@ def __init__(self, project_id: str, api_key: str, headers: dict = None):
             self._header.update(headers)
 
     def validate(self):
-        res = requests.get(
+        resp = requests.get(
             url=self._baseurl + "/vectordb/list",
             data=None,
             headers=self._header,
             verify=False,
         )
-        data = res.json()
-        res.close()
+        data = resp.json()
+        resp.close()
+        del resp
         return data
 
     def get_db_list(self):
         db_list = []
         req_url = "{}/vectordb/list".format(self._baseurl)
-        res = requests.get(url=req_url, data=None, headers=self._header, verify=False)
-        status_code = res.status_code
-        body = res.json()
+        resp = requests.get(url=req_url, data=None, headers=self._header, verify=False)
+        status_code = resp.status_code
+        body = resp.json()
         if status_code == 200 and body["statusCode"] == 200:
-            db_list = [db_id for db_id in res.json()["result"]]
-        res.close()
+            db_list = [db_id for db_id in body["result"]]
+        resp.close()
+        del resp
         return db_list
 
     def get_db_info(self, db_id: str):
         req_url = "{}/vectordb/{}".format(self._baseurl, db_id)
-        res = requests.get(url=req_url, data=None, headers=self._header, verify=False)
-        status_code = res.status_code
-        body = res.json()
-        res.close()
+        resp = requests.get(url=req_url, data=None, headers=self._header, verify=False)
+        status_code = resp.status_code
+        body = resp.json()
+        resp.close()
+        del resp
         return status_code, body
 
     def get_db_statistics(self, db_id: str):
         req_url = "{}/vectordb/{}/statistics".format(self._baseurl, db_id)
         req_data = None
-        res = requests.get(
+        resp = requests.get(
             url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
         )
-        status_code = res.status_code
-        body = res.json()
-        res.close()
+        status_code = resp.status_code
+        body = resp.json()
+        resp.close()
+        del resp
         return status_code, body
 
     def vectordb(self, db_id: str):
         # validate project_id and api_key
-        res = self.validate()
-        if res["statusCode"] != 200:
-            if res["statusCode"] == 404:
+        resp = self.validate()
+        if resp["statusCode"] != 200:
+            if resp["statusCode"] == 404:
                 raise Exception("Invalid project_id")
-            if res["statusCode"] == 401:
+            if resp["statusCode"] == 401:
                 raise Exception("Invalid api_key")
 
         # validate db_id
@@ -95,6 +100,7 @@ def vectordb(self, db_id: str):
             )
         else:
             print(resp)
+            del resp
             raise Exception("Failed to get db info")
 
 
@@ -123,10 +129,11 @@ def list_tables(self):
         if self._db_id is None:
             raise Exception("[ERROR] db_id is None!")
         req_url = "{}/table/list".format(self._baseurl)
-        res = requests.get(url=req_url, headers=self._header, verify=False)
-        status_code = res.status_code
-        body = res.json()
-        res.close()
+        resp = requests.get(url=req_url, headers=self._header, verify=False)
+        status_code = resp.status_code
+        body = resp.json()
+        resp.close()
+        del resp
         return status_code, body
 
     # Create table
@@ -144,12 +151,13 @@ def create_table(
         req_data = {"name": table_name, "fields": table_fields}
         if indices is not None:
             req_data["indices"] = indices
-        res = requests.post(
+        resp = requests.post(
             url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
         )
-        status_code = res.status_code
-        body = res.json()
-        res.close()
+        status_code = resp.status_code
+        body = resp.json()
+        resp.close()
+        del resp
         return status_code, body
 
     # Drop table
@@ -158,35 +166,38 @@ def drop_table(self, table_name: str):
             raise Exception("[ERROR] db_id is None!")
         req_url = "{}/table/delete?table_name={}".format(self._baseurl, table_name)
         req_data = {}
-        res = requests.delete(
+        resp = requests.delete(
             url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
         )
-        status_code = res.status_code
-        body = res.json()
-        res.close()
+        status_code = resp.status_code
+        body = resp.json()
+        resp.close()
+        del resp
         return status_code, body
 
     # Insert data into table
     def insert(self, table_name: str, records: list[dict]):
         req_url = "{}/data/insert".format(self._baseurl)
         req_data = {"table": table_name, "data": records}
-        res = requests.post(
+        resp = requests.post(
             url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
         )
-        status_code = res.status_code
-        body = res.json()
-        res.close()
+        status_code = resp.status_code
+        body = resp.json()
+        resp.close()
+        del resp
         return status_code, body
 
     def upsert(self, table_name: str, records: list[dict]):
         req_url = "{}/data/insert".format(self._baseurl)
         req_data = {"table": table_name, "data": records, "upsert": True}
-        res = requests.post(
+        resp = requests.post(
             url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
         )
-        status_code = res.status_code
-        body = res.json()
-        res.close()
+        status_code = resp.status_code
+        body = resp.json()
+        resp.close()
+        del resp
         return status_code, body
 
     # Query data from table
@@ -201,32 +212,46 @@ def query(
         limit: int = 2,
         filter: Optional[str] = None,
         with_distance: Optional[bool] = False,
+        facets: Optional[list[dict]] = None,
     ):
         req_url = "{}/data/query".format(self._baseurl)
-        req_data = {"table": table_name}
+        req_data = {"table": table_name, "limit": limit}
+
+        if response_fields is None:
+            response_fields = []
+
         if query_text is not None:
             req_data["query"] = query_text
         if query_index is not None:
             req_data["queryIndex"] = query_index
-        if query_field != None:
+        if query_field is not None:
             req_data["queryField"] = query_field
-        if query_vector != None:
+        if query_vector is not None:
             req_data["queryVector"] = query_vector
-        if response_fields != None:
+        if response_fields is not None:
             req_data["response"] = response_fields
-        if limit != None:
-            req_data["limit"] = limit
-        if filter != None:
+        if filter is not None:
             req_data["filter"] = filter
-        if with_distance != None:
+        if with_distance is not False:
             req_data["withDistance"] = with_distance
 
-        res = requests.post(
+        if facets is not None and len(facets) > 0:
+            aggregate_not_existing = 0
+            for facet in facets:
+                if "aggregate" not in facet:
+                    aggregate_not_existing += 1
+            if aggregate_not_existing > 0:
+                raise Exception("[ERROR] key aggregate is a must in facets!")
+            else:
+                req_data["facets"] = facets
+
+        resp = requests.post(
             url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
         )
-        status_code = res.status_code
-        body = res.json()
-        res.close()
+        status_code = resp.status_code
+        body = resp.json()
+        resp.close()
+        del resp
         return status_code, body
 
     # Delete data from table
@@ -261,12 +286,13 @@ def delete(
         if filter is not None:
             req_data["filter"] = filter
 
-        res = requests.post(
+        resp = requests.post(
             url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
         )
-        status_code = res.status_code
-        body = res.json()
-        res.close()
+        status_code = resp.status_code
+        body = resp.json()
+        resp.close()
+        del resp
         return status_code, body
 
     # Get data from table
@@ -279,6 +305,7 @@ def get(
         filter: Optional[str] = None,
         skip: Optional[int] = None,
         limit: Optional[int] = None,
+        facets: Optional[list[dict]] = None,
     ):
         """Epsilla supports get records by primary keys as default for now."""
         if primary_keys is not None and ids is not None:
@@ -293,6 +320,7 @@ def get(
             primary_keys = ids
 
         req_data = {"table": table_name}
+
         if response_fields is not None:
             req_data["response"] = response_fields
         if primary_keys is not None:
@@ -304,13 +332,24 @@ def get(
         if limit is not None:
             req_data["limit"] = limit
 
+        if facets is not None and len(facets) > 0:
+            aggregate_not_existing = 0
+            for facet in facets:
+                if "aggregate" not in facet:
+                    aggregate_not_existing += 1
+            if aggregate_not_existing > 0:
+                raise Exception("[ERROR] key aggregate is a must in facets!")
+            else:
+                req_data["facets"] = facets
+
         req_url = "{}/data/get".format(self._baseurl)
-        res = requests.post(
+        resp = requests.post(
             url=req_url, data=json.dumps(req_data), headers=self._header, verify=False
         )
-        status_code = res.status_code
-        body = res.json()
-        res.close()
+        status_code = resp.status_code
+        body = resp.json()
+        resp.close()
+        del resp
         return status_code, body
 
     def as_search_engine(self):

pyepsilla/cloud/sentry.py

@@ -4,14 +4,20 @@
 # Sentry collects crash reports and performance numbers
 # It is possible to turn off data collection using an environment variable named "SENTRY_DISABLE"
 
-import os, sys, platform, uuid, hashlib, socket, requests
+import hashlib
+import os
+import platform
+import socket
+import sys
+import uuid
+
+import requests
 import sentry_sdk
-from sentry_sdk.integrations.atexit import AtexitIntegration
 from pyepsilla.vectordb.version import __version__
-
+from sentry_sdk.integrations.atexit import AtexitIntegration
 
 CONFIG_URL = "https://config.epsilla.com/candidate.json"
-SENTRY_DSN = "https://c705adb9ba9a5750ab5719c69021e3b0@o4505952171917312.ingest.sentry.io/4506031458746368"
+SENTRY_DSN = "https://3f89b94a4a2e7620c8ecce81cb302d43@o4507288359862272.ingest.us.sentry.io/4507288364908545"
 
 try:
     r = requests.get(CONFIG_URL, headers={"Agent": "PyEpsilla Cloud Client"}, timeout=2)
@@ -20,28 +26,39 @@
 except Exception:
     pass
 
+
 def callback(pending, timeout):
     sys.stderr.flush()
-        
+
+
 def get_external_ip() -> str:
     try:
         return requests.get("https://api.ipify.org", timeout=2).text
     except Exception:
         return "NA"
 
+
 def init_sentry():
-    if "SENTRY_DISABLE" not in os.environ:
+    if os.getenv("SENTRY_DISABLE", None) is not None:
         try:
             uid = hashlib.sha256(str(uuid.getnode()).encode()).hexdigest()
             internal_ip = socket.gethostbyname(socket.gethostname())
             external_ip = get_external_ip()
             sentry_sdk.set_tag("uid", uid)
             sentry_sdk.set_tag("internal_ip", internal_ip)
             sentry_sdk.set_tag("external_ip", external_ip)
-            sentry_sdk.set_user({'ip_address': '{{auto}}'})
-            sentry_sdk.set_user({"username": "{}-{}-{}".format(socket.gethostname(), internal_ip, external_ip)})
+            sentry_sdk.set_user({"ip_address": "{{auto}}"})
+            sentry_sdk.set_user(
+                {
+                    "username": "{}-{}-{}".format(
+                        socket.gethostname(), internal_ip, external_ip
+                    )
+                }
+            )
             sentry_sdk.set_tag("version", platform.version())
-            sentry_sdk.set_tag("platform", "{}-{}".format(sys.platform, platform.machine()))
+            sentry_sdk.set_tag(
+                "platform", "{}-{}".format(sys.platform, platform.machine())
+            )
             sentry_sdk.init(
                 dsn=SENTRY_DSN,
                 release=__version__,
@@ -53,3 +70,5 @@ def init_sentry():
         except Exception:
             sentry_sdk.flush()
             pass
+    else:
+        return None