Merge pull request #47 from epsilla-cloud/dev

add facets

Author: eric-epsilla

pyepsilla/enterprise/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()  # type: ignore
 
 
@@ -26,7 +27,9 @@ class DbModel(BaseModel):
 
 
 class Client(cloud.Client):
-    def __init__(self, base_url: str, project_id: Optional[str] = "default", headers: dict = None):
+    def __init__(
+        self, base_url: str, project_id: Optional[str] = "default", headers: dict = None
+    ):
         self._project_id = project_id
         self._baseurl = f"{base_url}/api/v3/project/{project_id}"
         self._timeout = 10
@@ -176,7 +179,12 @@ def list_tables(self):
         return status_code, body
 
     # Create table
-    def create_table(self, table_name: str, table_fields: list[dict] = None, indices: list[dict] = None):
+    def create_table(
+        self,
+        table_name: str,
+        table_fields: list[dict] = None,
+        indices: list[dict] = None,
+    ):
         if self._db_id is None:
             raise Exception("[ERROR] db_id is None!")
         if table_fields is None:
@@ -250,32 +258,43 @@ 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}
         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:
+        if limit is not None:
             req_data["limit"] = limit
-        if filter != None:
+        if filter is not None:
             req_data["filter"] = filter
         if with_distance is not False:
             req_data["withDistance"] = with_distance
+        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
 
         res = 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()
+        del res
         return status_code, body
 
     # Delete data from table

pyepsilla/vectordb/client.py

@@ -7,12 +7,13 @@
 import socket
 import time
 from typing import Optional, Union
-from ..utils.search_engine import SearchEngine
 
 import requests
 import sentry_sdk
 from requests.packages.urllib3.exceptions import InsecureRequestWarning
 
+from ..utils.search_engine import SearchEngine
+
 requests.packages.urllib3.disable_warnings(InsecureRequestWarning)
 
 
@@ -249,10 +250,11 @@ def query(
         query_index: str = None,
         query_field: str = None,
         query_vector: Union[list, dict] = None,
-        response_fields: list = None,
+        response_fields: Optional[list] = None,
         limit: int = 2,
-        filter: str = "",
-        with_distance: bool = False,
+        filter: Optional[str] = None,
+        with_distance: Optional[bool] = False,
+        facets: Optional[list[dict]] = None,
     ):
         if self._db is None:
             raise Exception("[ERROR] Please use_db() first!")
@@ -266,6 +268,15 @@ def query(
             "filter": filter,
             "withDistance": with_distance,
         }
+        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
         if query_text is not None:
             req_data["query"] = query_text
         if query_index is not None:
@@ -280,6 +291,7 @@ def query(
         status_code = res.status_code
         body = res.json()
         res.close()
+        del res
         return status_code, body
 
     def get(
@@ -291,6 +303,7 @@ def get(
         filter: Optional[str] = None,
         skip: Optional[int] = None,
         limit: Optional[int] = None,
+        facets: Optional[list[dict]] = None,
     ):
         if self._db is None:
             raise Exception("[ERROR] Please use_db() first!")
@@ -302,26 +315,36 @@ def get(
             print(
                 "[WARN] Both primary_keys and ids are prvoided, will use primary keys by default!"
             )
-        if primary_keys == None and ids != None:
+        if primary_keys is None and ids is not None:
             primary_keys = ids
 
         req_data = {"table": table_name}
 
-        if response_fields != None:
+        if response_fields is not None:
             req_data["response"] = response_fields
 
-        if primary_keys != None:
+        if primary_keys is not None:
             req_data["primaryKeys"] = primary_keys
 
-        if filter != None:
+        if filter is not None:
             req_data["filter"] = filter
 
-        if skip != None:
+        if skip is not None:
             req_data["skip"] = filter
 
-        if limit != None:
+        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 = "{}/api/{}/data/get".format(self._baseurl, self._db)
         res = requests.post(
             url=req_url, data=json.dumps(req_data), headers=self._header, verify=False

pyepsilla/vectordb/sentry.py

@@ -4,13 +4,21 @@
 # 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 .version import __version__
 
 CONFIG_URL = "https://config.epsilla.com/candidate.json"
-SENTRY_DSN = "https://c705adb9ba9a5750ab5719c69021e3b0@o4505952171917312.ingest.sentry.io/4506031458746368"
+SENTRY_DSN = "https://7b7043e213409f7125c511127697d668@o4506949201428480.ingest.us.sentry.io/4506949333680128"
 
 try:
     r = requests.get(CONFIG_URL, headers={"Agent": "PyEpsilla"}, timeout=2)
@@ -19,15 +27,18 @@
 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:
         try:
@@ -37,10 +48,18 @@ def init_sentry():
             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__,

pyepsilla/vectordb/version.py

@@ -1 +1 @@
-__version__ = "0.3.4"
+__version__ = "0.3.5"