Skip to content

Commit

Permalink
WIP
Browse files Browse the repository at this point in the history
  • Loading branch information
@lferran
lferran committed Nov 30, 2023
1 parent 0d5c4d0 commit 23c3f46
Show file tree
Hide file tree
Showing 4 changed files with 123 additions and 19 deletions.
17 changes: 8 additions & 9 deletions nucliadb_performance/nucliadb_performance/standalone/test_search.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
save_benchmark_json_results,
)
from nucliadb_performance.utils.misc import (
get_fake_word,
get_kb,
get_request,
make_kbid_request,
Expand All @@ -33,28 +32,28 @@ def init_test(args):
@scenario(weight=1)
async def test_find(session):
kbid, slug = get_test_kb()
request = get_request(slug, with_tags=settings.request_tags)
request = get_request(slug, endpoint="find", with_tags=settings.request_tags)
request.payload["vector"] = compute_vector(request.payload["query"])
await make_kbid_request(
resp = await make_kbid_request(
session,
kbid,
request.method.upper(),
request.url.format(kbid=kbid),
json=request.payload,
)
#TODO: assert that there are results!


@scenario(weight=2)
async def test_suggest(session):
kbid, _ = get_test_kb()
url = "/v1/kb/{kbid}/suggest".format(kbid=kbid)
method = "GET"
kbid, slug = get_test_kb()
request = get_request(slug, endpoint="suggest")
await make_kbid_request(
session,
kbid,
method,
url,
params={"query": get_fake_word()},
request.method.upper(),
request.url.format(kbid=kbid),
params=request.params,
)


Expand Down
13 changes: 7 additions & 6 deletions nucliadb_performance/nucliadb_performance/utils/misc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,10 +96,10 @@ async def make_request(self, method: str, path: str, *args, **kwargs):
kwargs_headers.update(base_headers)
kwargs["headers"] = kwargs_headers
async with func(url, *args, **kwargs) as resp:
if resp.status == 200:
record_request_process_time(resp)
return
await self.handle_search_error(resp)
if resp.status != 200:
await self.handle_search_error(resp)
record_request_process_time(resp)
return await resp.json()

async def handle_search_error(self, resp):
content = None
Expand Down Expand Up @@ -212,7 +212,7 @@ async def make_kbid_request(session, kbid, method, path, params=None, json=None)
global ERRORS
try:
client = get_search_client(session)
await client.make_request(method, path, params=params, json=json)
return await client.make_request(method, path, params=params, json=json)
except RequestError as err:
# Store error info so we can inspect after the script runs
detail = (
Expand All @@ -235,13 +235,14 @@ def print_errors():
print("=" * 50)


def get_request(kbid_or_slug: str, with_tags=None) -> Request:
def get_request(kbid_or_slug: str, endpoint: str, with_tags=None) -> Request:
if settings.saved_requests_file is None:
raise AttributeError("SAVED_REQUESTS_FILE env var is not set!")
saved_requests_file = CURRENT_DIR + "/" + settings.saved_requests_file
requests = load_saved_request(
saved_requests_file,
kbid_or_slug,
endpoint=endpoint,
with_tags=tuple(with_tags) if with_tags else None,
)
if len(requests) == 0:
Expand Down
10 changes: 6 additions & 4 deletions nucliadb_performance/nucliadb_performance/utils/saved_requests.py
View file
Original file line number Diff line number Diff line change
@@ -1,18 +1,21 @@
from functools import cache
from typing import Any, Optional
from typing import Any, Optional, Union

from pydantic import BaseModel

Payload = dict[str, Any]
Params = dict[str, Union[str, list[str]]]


class Request(BaseModel):
url: str
method: str
payload: Optional[Payload] = None
params: Optional[Params] = None


class SavedRequest(BaseModel):
endpoint: str
request: Request
tags: list[str] = []
description: Optional[str] = None
Expand All @@ -29,16 +32,15 @@ class SavedRequests(BaseModel):

@cache
def load_saved_request(
saved_requests_file: str, kbid_or_slug: str, with_tags=None
saved_requests_file: str, kbid_or_slug: str, endpoint: str, with_tags=None
) -> list[Request]:
try:
saved_requests = SavedRequests.parse_file(saved_requests_file)
kb_requests = []
for rs in saved_requests.sets.values():
if kbid_or_slug not in rs.kbs:
continue
kb_requests.extend(rs.requests)

kb_requests.extend([r for r in rs.requests if r.endpoint == endpoint])
if with_tags is None:
return [kb_req.request for kb_req in kb_requests]
else:
Expand Down
102 changes: 102 additions & 0 deletions nucliadb_performance/standalone_search_requests.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
],
"requests": [
{
"endpoint": "find",
"description": "Filtered hybrid search",
"tags": [
"filtered",
Expand All @@ -25,6 +26,7 @@
}
},
{
"endpoint": "find",
"description": "Filtered vector search",
"tags": [
"filtered",
Expand All @@ -45,6 +47,7 @@
}
},
{
"endpoint": "find",
"description": "Hybrid search with date range filter",
"tags": [
"date_range",
Expand All @@ -59,6 +62,37 @@
"range_creation_end": "3020-1-1T00:00:00"
}
}
},
{
"endpoint": "suggest",
"description": "Suggest some paragraphs",
"tags": [],
"request": {
"url": "/v1/kb/{kbid}/suggest",
"method": "GET",
"params": {
"query": "united kingdom",
"features": [
"paragraph"
],
"fields": ["a/title"]
}
}
},
{
"endpoint": "suggest",
"description": "Suggest some paragraphs",
"tags": [],
"request": {
"url": "/v1/kb/{kbid}/suggest",
"method": "GET",
"params": {
"query": "united kingdom",
"features": [
"paragraph", "entities"
]
}
}
}
]
},
Expand All @@ -69,6 +103,7 @@
],
"requests": [
{
"endpoint": "find",
"description": "Filtered hybrid search",
"tags": [
"filtered",
Expand All @@ -86,6 +121,7 @@
}
},
{
"endpoint": "find",
"description": "Filtered vector search",
"tags": [
"filtered",
Expand All @@ -106,6 +142,7 @@
}
},
{
"endpoint": "find",
"description": "Hybrid search with date range filter",
"tags": [
"date_range",
Expand All @@ -120,6 +157,37 @@
"range_creation_end": "3020-1-1T00:00:00"
}
}
},
{
"endpoint": "suggest",
"description": "Suggest some paragraphs",
"tags": [],
"request": {
"url": "/v1/kb/{kbid}/suggest",
"method": "GET",
"params": {
"query": "studies quarterly",
"features": [
"paragraph"
],
"fields": ["a/title"]
}
}
},
{
"endpoint": "suggest",
"description": "Suggest some paragraphs",
"tags": [],
"request": {
"url": "/v1/kb/{kbid}/suggest",
"method": "GET",
"params": {
"query": "studies quarterly",
"features": [
"paragraph", "entities"
]
}
}
}
]
},
Expand All @@ -130,6 +198,7 @@
],
"requests": [
{
"endpoint": "find",
"description": "Filtered hybrid search",
"tags": [
"filtered",
Expand All @@ -148,6 +217,7 @@
}
},
{
"endpoint": "find",
"description": "Filtered vector search",
"tags": [
"filtered",
Expand All @@ -168,6 +238,7 @@
}
},
{
"endpoint": "find",
"description": "Hybrid search with date range filter",
"tags": [
"date_range",
Expand All @@ -182,6 +253,37 @@
"range_creation_end": "3020-1-1T00:00:00"
}
}
},
{
"endpoint": "suggest",
"description": "Suggest some paragraphs",
"tags": [],
"request": {
"url": "/v1/kb/{kbid}/suggest",
"method": "GET",
"params": {
"query": "Radiation protection",
"features": [
"paragraph"
],
"fields": ["a/title"]
}
}
},
{
"endpoint": "suggest",
"description": "Suggest some paragraphs",
"tags": [],
"request": {
"url": "/v1/kb/{kbid}/suggest",
"method": "GET",
"params": {
"query": "Radiation protection",
"features": [
"paragraph", "entities"
]
}
}
}
]
}
Expand Down

0 comments on commit 23c3f46

Please sign in to comment.