Semantic search (#17)
* ✨(backend) add trivial vector embedding add a trivial vector embedding with constant [0.0, 0.0] * 🙈(core) gitignore ignore files related to sqlite and pdb * ✨(backend) introduce hybrid search handle full-text along with seamntic search * ✨(backend) install basic embedding model I need an embedding to performe semantic search. I need a simple model from hugging face. * ✨(backend) embbed the text embed the text of the query and the document. * 🐛(backend) fix filters and refactor view filter were broken by previous commits. This fixes them. * ♻️(backend) refactor pipeline creation pipeline had to be refactored. This refactors it. * ✨(backend) improve filtering filtering were done once after hybrid computation. For efficiency it should be done of each subquery. * 🔧(setting) move variables to settings NLP_SEARCH_PIPELINE_ID and HYBRID_SEARCH_WEIGHTS is moved to setting file so user can param Find. * ✨(backend) use albert api we choose to rely on Albert API instead of installing a model in local. It is less effort to maintain. * ♻️(backend) hide EMBEDDING_API_KEY EMBEDDING_API_KEY should not be visible. * ✨(backend) move opensearch functions to a service user of find should be able to disable seamntic search. If it is not properly setted it is also turned off without impatcing full-text search. * 🧪(backend) test add test so the app is tested. * 📝(backend) add documentation add documentation so the app is documented. * ♻️(backend) remove local model the local model is no longer useful. Its file must be removed. * 🧪(backend) one more test The app must be tested more. This tests the app more. * ✨(bakckend) handle k value k value must be handled so the user can have a control over the number of results. * ♻️(backend) clean branch beanch had to be cleaned bery very much * 🔧(infra) define EMBEDDING_API_KEY EMBEDDING_API_KEY must be hiden. This hide EMBEDDING_API_KEY * 🚨(backend) fix linters linters must be fixed. This commit fixes them. * 📝(core) add changelog changelog must be updated. This updates the changelog. * 🐛(backend) fix linters more linters had to be fixed more. This fixes linters more. * 🐛(backend) fix tests test must be fixed. This fixes the tests. * ♻️(backend) improve variable managment variable managment must be improve. This improve variable managment. * ♻️(backend) remove embedding from schemas embedding must be removed from schemas. This removes embedding from shemas. * ✨(backend) add reindex_with_embedding command We must be able to enable hybrid search if it was disabled or chnage the embedding model. To do so we must reindex all documents with a new embedding. reindex_with_embedding does that. * ✨(backend) add create_pipeline command We must be able to create the command pipeline once and not check at all request. * 🧪(backend) tests I add a test and fix other tests * 🚨(backend) linters linters must be fixed. This fixes linters. * ✨(backend) remove pagination Semantic search has an impact of pagination. Pagination will be perfomed in services consuming Find API (Doc, Drive etc...) * ✨(backend) improve reindexing we want to handle error case and model change. I introduce a embedding_model field to keep track of the embedding state. * 🧪(backend) test more the command must be tested more. This tests the command more. * 🧪(backend) test concurent update do not lead do data loss updates on a document mught be done by a user while reindexing. I check the latest data is not lost. using if_seq_no and if_primary_term is not only not useful but whould require reruning the command. * ✨(backend) improve reindexing again reindexing must preserve the latest updates. I reintroduce the no_seq update field. * ♻️(backend) various small improvments I make various small improvments.
This commit is contained in:
Charles Englebert
committed by
GitHub
GitHub
parent
11846238f2
commit
3ec95c9edf
@@ -69,6 +69,7 @@ src/frontend/tsclient
|
||||
.pylint.d
|
||||
.pytest_cache
|
||||
db.sqlite3
|
||||
*history.sqlite
|
||||
.mypy_cache
|
||||
|
||||
# Site media
|
||||
@@ -79,3 +80,7 @@ db.sqlite3
|
||||
.vscode/
|
||||
*.iml
|
||||
.devcontainer
|
||||
.vscode-server
|
||||
|
||||
#pdb
|
||||
*.pdbhistory
|
||||
|
||||
@@ -10,6 +10,7 @@ and this project adheres to
|
||||
|
||||
## Added
|
||||
|
||||
- ✨(backend) add semantic search
|
||||
- backend application
|
||||
- helm chart
|
||||
- 🐛(backend) fix missing index creation in 'index/' view
|
||||
|
||||
@@ -9,4 +9,5 @@ flowchart TD
|
||||
Back -- REST API --> Opensearch
|
||||
Back <--> Celery --> DB
|
||||
User -- HTTP --> Dashboard --> Opensearch
|
||||
Back -- REST API --> Embedding Endpoint
|
||||
```
|
||||
|
||||
@@ -103,3 +103,10 @@ These are the environment variables you can set for the `find-backend` container
|
||||
| USER_OIDC_ESSENTIAL_CLAIMS | Essential claims in OIDC token | [] |
|
||||
| Y_PROVIDER_API_BASE_URL | Y Provider url | |
|
||||
| Y_PROVIDER_API_KEY | Y provider API key | |
|
||||
| HYBRID_SEARCH_ENABLED | Flag to enable hybrid (an then semantic) search | True |
|
||||
| HYBRID_SEARCH_WEIGHTS | Weights used in the weighted sum of the hybrid search score | [0.3, 0.7] |
|
||||
| EMBEDDING_API_PATH | URL of the embedding api | https://albert.api.etalab.gouv.fr/v1/embeddings |
|
||||
| EMBEDDING_API_KEY | API key of the embedding api | |
|
||||
| EMBEDDING_REQUEST_TIMEOUT | time out in seconds of the embedding requests | 10 |
|
||||
| EMBEDDING_API_MODEL_NAME | Name of the embedding model used on the api | embeddings-small |
|
||||
| EMBEDDING_DIMENSION | Size of the embedding vector | 1024 |
|
||||
|
||||
+27
-2
@@ -6,9 +6,11 @@ are visible.
|
||||
|
||||
## Setup Opensearch
|
||||
|
||||
Add the following settings to your Django settings for the Find application.
|
||||
### General
|
||||
|
||||
```shell
|
||||
Add the following variables to your Django settings to configure Find and enable full-text search.
|
||||
|
||||
```python
|
||||
# Login for opensearch
|
||||
OPENSEARCH_USER=opensearch-user
|
||||
OPENSEARCH_PASSWORD=your-opensearch-password
|
||||
@@ -21,6 +23,29 @@ OPENSEARCH_PORT=9200
|
||||
OPENSEARCH_USE_SSL=True
|
||||
```
|
||||
|
||||
### Semantic search
|
||||
|
||||
Find offers a semantic search feature. You can either use pure full-text search or a hybrid full-text + semantic search. To enable the hybrid search, add the fallowing settings.
|
||||
|
||||
```python
|
||||
# Enable flag
|
||||
HYBRID_SEARCH_ENABLED = True
|
||||
|
||||
# weighted sum: full_text_weight, semantic_search_weight
|
||||
HYBRID_SEARCH_WEIGHTS = 0.7,0.3
|
||||
|
||||
# Embedding
|
||||
EMBEDDING_API_PATH = https://embedding.api.example.com/full/path/
|
||||
EMBEDDING_API_KEY = your-embedding-api-key
|
||||
EMBEDDING_REQUEST_TIMEOUT = 10
|
||||
EMBEDDING_API_MODEL_NAME = embedding-api-model-name
|
||||
EMBEDDING_DIMENSION = 1024
|
||||
```
|
||||
|
||||
The hybrid search computes a score for full-text and semantic search and combines them through a weighted sum. HYBRID_SEARCH_WEIGHTS contains the weights of full-text and semantic respectively.
|
||||
|
||||
You need to use an embedding api similar to https://albert.api.etalab.gouv.fr/documentation#tag/Embeddings/operation/embeddings_v1_embeddings_post.
|
||||
|
||||
## Setup indexation API
|
||||
|
||||
Other applications can index their files through the **`/index/`** endpoint with a simple token authentication.
|
||||
|
||||
@@ -49,3 +49,8 @@ OIDC_RS_SIGN_ALGO=RS256
|
||||
|
||||
OIDC_RS_BACKEND_CLASS="core.authentication.FinderResourceServerBackend"
|
||||
OIDC_RS_ENCRYPTION_KEY_TYPE="RSA"
|
||||
|
||||
# Hybrid Search settings
|
||||
HYBRID_SEARCH_ENABLED=True
|
||||
EMBEDDING_API_KEY=ThisIsAnExampleKeyForDevPurposeOnly
|
||||
EMBEDDING_API_PATH=https://albert.api.etalab.gouv.fr/v1/embeddings
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Find Core application"""
|
||||
|
||||
from django.apps import AppConfig
|
||||
from django.utils.translation import gettext_lazy as _
|
||||
|
||||
from core.management.commands.create_search_pipeline import (
|
||||
ensure_search_pipeline_exists,
|
||||
)
|
||||
from core.services.opensearch import (
|
||||
check_hybrid_search_enabled,
|
||||
)
|
||||
|
||||
|
||||
class CoreConfig(AppConfig):
|
||||
"""Configuration class for the Find core app."""
|
||||
|
||||
name = "core"
|
||||
app_label = "core"
|
||||
verbose_name = _("Find core application")
|
||||
|
||||
def ready(self):
|
||||
"""
|
||||
Ensure search pipeline exists if hybrid search is enabled.
|
||||
"""
|
||||
if check_hybrid_search_enabled():
|
||||
ensure_search_pipeline_exists()
|
||||
@@ -0,0 +1,53 @@
|
||||
"""
|
||||
Handle create the search pipeline command of the hybrid search.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.management.base import BaseCommand
|
||||
|
||||
from opensearchpy.exceptions import NotFoundError
|
||||
|
||||
from core.services.opensearch import (
|
||||
opensearch_client,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
"""Handle create the search pipeline command of the hybrid search."""
|
||||
|
||||
help = __doc__
|
||||
|
||||
def handle(self, *args, **options):
|
||||
ensure_search_pipeline_exists()
|
||||
|
||||
|
||||
def ensure_search_pipeline_exists():
|
||||
"""Create search pipeline for hybrid search if it does not exist"""
|
||||
try:
|
||||
opensearch_client().search_pipeline.get(settings.HYBRID_SEARCH_PIPELINE_ID)
|
||||
logger.info("Search pipeline exists already")
|
||||
except NotFoundError:
|
||||
logger.info("Creating search pipeline: %s", settings.HYBRID_SEARCH_PIPELINE_ID)
|
||||
opensearch_client().transport.perform_request(
|
||||
method="PUT",
|
||||
url="/_search/pipeline/" + settings.HYBRID_SEARCH_PIPELINE_ID,
|
||||
body={
|
||||
"description": "Post processor for hybrid search",
|
||||
"phase_results_processors": [
|
||||
{
|
||||
"normalization-processor": {
|
||||
"combination": {
|
||||
"technique": "arithmetic_mean",
|
||||
"parameters": {
|
||||
"weights": settings.HYBRID_SEARCH_WEIGHTS
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,128 @@
|
||||
"""
|
||||
Handle reindexing of documents with embeddings in OpenSearch.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.management.base import BaseCommand, CommandError
|
||||
|
||||
from opensearchpy.exceptions import NotFoundError
|
||||
|
||||
from core.services.opensearch import (
|
||||
check_hybrid_search_enabled,
|
||||
embed_text,
|
||||
format_document,
|
||||
opensearch_client,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
"""Reindex all documents with embeddings"""
|
||||
|
||||
help = __doc__
|
||||
opensearch_client_ = opensearch_client()
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument("index_name", type=str)
|
||||
|
||||
def handle(self, *args, **options):
|
||||
"""Launch the reindexing with embedding."""
|
||||
|
||||
index_name = options["index_name"]
|
||||
|
||||
if not check_hybrid_search_enabled():
|
||||
raise CommandError("Hybrid search is not enabled or properly configured.")
|
||||
|
||||
try:
|
||||
self.opensearch_client_.indices.get(index=index_name)
|
||||
except NotFoundError as error:
|
||||
raise CommandError(f"Index {index_name} does not exist.") from error
|
||||
|
||||
self.stdout.write(f"[INFO] Start reindexing {index_name} with embedding.")
|
||||
|
||||
result = reindex_with_embedding(index_name)
|
||||
|
||||
self.stdout.write(
|
||||
f"[INFO] Reindexing of {index_name} is done.\n"
|
||||
f"nb success embedding: {result['nb_success_embedding']}\n"
|
||||
f"nb failed embedding: {result['nb_failed_embedding']} embedding fails\n"
|
||||
)
|
||||
|
||||
|
||||
def reindex_with_embedding(index_name, batch_size=500, scroll="10m"):
|
||||
"""
|
||||
Reindex documents from source index to destination index with embeddings.
|
||||
|
||||
returns a dict with the number of successful embeddings and failed embeddings.
|
||||
"""
|
||||
opensearch_client_ = opensearch_client()
|
||||
page = opensearch_client_.search(
|
||||
index=index_name,
|
||||
scroll=scroll,
|
||||
size=batch_size,
|
||||
seq_no_primary_term=True,
|
||||
body={
|
||||
"query": {
|
||||
"bool": {
|
||||
"should": [
|
||||
{"bool": {"must_not": {"exists": {"field": "embedding"}}}},
|
||||
{
|
||||
"bool": {
|
||||
"must_not": {
|
||||
"term": {
|
||||
"embedding_model": settings.EMBEDDING_API_MODEL_NAME
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
],
|
||||
"minimum_should_match": 1,
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
nb_failed_embedding = 0
|
||||
nb_success_embedding = 0
|
||||
while len(page["hits"]["hits"]) > 0:
|
||||
actions = []
|
||||
for hit in page["hits"]["hits"]:
|
||||
source = hit["_source"]
|
||||
embedding = embed_text(
|
||||
format_document(source.get("title", ""), source.get("content", ""))
|
||||
)
|
||||
if embedding:
|
||||
actions.append(
|
||||
{
|
||||
"update": {
|
||||
"_id": hit["_id"],
|
||||
# if_seq_no and if_primary_term ensure we only update indexes
|
||||
# if the document hasn't changed
|
||||
"if_seq_no": hit["_seq_no"],
|
||||
"if_primary_term": hit["_primary_term"],
|
||||
}
|
||||
}
|
||||
)
|
||||
actions.append(
|
||||
{
|
||||
"doc": {
|
||||
"embedding": embedding,
|
||||
"embedding_model": settings.EMBEDDING_API_MODEL_NAME,
|
||||
}
|
||||
}
|
||||
)
|
||||
nb_success_embedding += 1
|
||||
else:
|
||||
nb_failed_embedding += 1
|
||||
|
||||
opensearch_client_.bulk(index=index_name, body=actions)
|
||||
page = opensearch_client_.scroll(scroll_id=page["_scroll_id"], scroll=scroll)
|
||||
|
||||
opensearch_client_.clear_scroll(scroll_id=page["_scroll_id"])
|
||||
return {
|
||||
"nb_failed_embedding": nb_failed_embedding,
|
||||
"nb_success_embedding": nb_success_embedding,
|
||||
}
|
||||
@@ -1,54 +0,0 @@
|
||||
"""Opensearch related utils."""
|
||||
|
||||
from django.conf import settings
|
||||
|
||||
from opensearchpy import OpenSearch
|
||||
from opensearchpy.exceptions import NotFoundError
|
||||
|
||||
client = OpenSearch(
|
||||
hosts=[{"host": settings.OPENSEARCH_HOST, "port": settings.OPENSEARCH_PORT}],
|
||||
http_auth=(settings.OPENSEARCH_USER, settings.OPENSEARCH_PASSWORD),
|
||||
timeout=50,
|
||||
use_ssl=settings.OPENSEARCH_USE_SSL,
|
||||
verify_certs=False,
|
||||
)
|
||||
|
||||
|
||||
def ensure_index_exists(index_name):
|
||||
"""Create index if it does not exist"""
|
||||
try:
|
||||
client.indices.get(index=index_name)
|
||||
except NotFoundError:
|
||||
client.indices.create(
|
||||
index=index_name,
|
||||
body={
|
||||
"mappings": {
|
||||
"dynamic": "strict",
|
||||
"properties": {
|
||||
"id": {"type": "keyword"},
|
||||
"title": {
|
||||
"type": "keyword", # Primary field for exact matches and sorting
|
||||
"fields": {
|
||||
"text": {
|
||||
"type": "text"
|
||||
} # Sub-field for full-text search
|
||||
},
|
||||
},
|
||||
"depth": {"type": "integer"},
|
||||
"path": {
|
||||
"type": "keyword",
|
||||
"fields": {"text": {"type": "text"}},
|
||||
},
|
||||
"numchild": {"type": "integer"},
|
||||
"content": {"type": "text"},
|
||||
"created_at": {"type": "date"},
|
||||
"updated_at": {"type": "date"},
|
||||
"size": {"type": "long"},
|
||||
"users": {"type": "keyword"},
|
||||
"groups": {"type": "keyword"},
|
||||
"reach": {"type": "keyword"},
|
||||
"is_active": {"type": "boolean"},
|
||||
},
|
||||
}
|
||||
},
|
||||
)
|
||||
@@ -114,5 +114,4 @@ class SearchQueryParametersSchema(BaseModel):
|
||||
reach: Optional[enums.ReachEnum] = None
|
||||
order_by: Optional[Literal[enums.ORDER_BY_OPTIONS]] = Field(default=enums.RELEVANCE)
|
||||
order_direction: Optional[Literal["asc", "desc"]] = Field(default="desc")
|
||||
page_number: Optional[conint(ge=1)] = Field(default=1)
|
||||
page_size: Optional[conint(ge=1, le=100)] = Field(default=50)
|
||||
nb_results: Optional[conint(ge=1, le=300)] = Field(default=50)
|
||||
|
||||
@@ -0,0 +1,333 @@
|
||||
"""Opensearch related utils."""
|
||||
|
||||
import logging
|
||||
from functools import cache
|
||||
|
||||
from django.conf import settings
|
||||
|
||||
import requests
|
||||
from opensearchpy import OpenSearch
|
||||
from opensearchpy.exceptions import NotFoundError
|
||||
from rest_framework.exceptions import ValidationError
|
||||
|
||||
from core import enums
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
REQUIRED_ENV_VARIABLES = [
|
||||
"OPENSEARCH_HOST",
|
||||
"OPENSEARCH_PORT",
|
||||
"OPENSEARCH_USER",
|
||||
"OPENSEARCH_PASSWORD",
|
||||
"OPENSEARCH_USE_SSL",
|
||||
]
|
||||
|
||||
|
||||
@cache
|
||||
def opensearch_client():
|
||||
"""Get OpenSearch client, ensuring required env variables are set"""
|
||||
missing_env_variables = [
|
||||
variable
|
||||
for variable in REQUIRED_ENV_VARIABLES
|
||||
if getattr(settings, variable, None) is None
|
||||
]
|
||||
if missing_env_variables:
|
||||
raise ValidationError(
|
||||
f"Missing required OpenSearch environment variables: {', '.join(missing_env_variables)}"
|
||||
)
|
||||
|
||||
return OpenSearch(
|
||||
hosts=[{"host": settings.OPENSEARCH_HOST, "port": settings.OPENSEARCH_PORT}],
|
||||
http_auth=(settings.OPENSEARCH_USER, settings.OPENSEARCH_PASSWORD),
|
||||
timeout=50,
|
||||
use_ssl=settings.OPENSEARCH_USE_SSL,
|
||||
verify_certs=False,
|
||||
)
|
||||
|
||||
|
||||
# pylint: disable=too-many-arguments, too-many-positional-arguments
|
||||
def search( # noqa : PLR0913
|
||||
q,
|
||||
nb_results,
|
||||
order_by,
|
||||
order_direction,
|
||||
search_indices,
|
||||
reach,
|
||||
visited,
|
||||
user_sub,
|
||||
groups,
|
||||
):
|
||||
"""Perform an OpenSearch search"""
|
||||
query = get_query(
|
||||
q=q,
|
||||
nb_results=nb_results,
|
||||
reach=reach,
|
||||
visited=visited,
|
||||
user_sub=user_sub,
|
||||
groups=groups,
|
||||
)
|
||||
return opensearch_client().search( # pylint: disable=unexpected-keyword-arg
|
||||
index=",".join(search_indices),
|
||||
body={
|
||||
"_source": enums.SOURCE_FIELDS, # limit the fields to return
|
||||
"script_fields": {
|
||||
"number_of_users": {"script": {"source": "doc['users'].size()"}},
|
||||
"number_of_groups": {"script": {"source": "doc['groups'].size()"}},
|
||||
},
|
||||
"sort": get_sort(
|
||||
query_keys=query.keys(),
|
||||
order_by=order_by,
|
||||
order_direction=order_direction,
|
||||
),
|
||||
"size": nb_results,
|
||||
# Compute query
|
||||
"query": query,
|
||||
},
|
||||
params=get_params(query_keys=query.keys()),
|
||||
# disable=unexpected-keyword-arg because
|
||||
# ignore_unavailable is not in the the method declaration
|
||||
ignore_unavailable=True,
|
||||
)
|
||||
|
||||
|
||||
# pylint: disable=too-many-arguments, too-many-positional-arguments
|
||||
def get_query( # noqa : PLR0913
|
||||
q, nb_results, reach, visited, user_sub, groups
|
||||
):
|
||||
"""Build OpenSearch query body based on parameters"""
|
||||
filter_ = get_filter(reach, visited, user_sub, groups)
|
||||
|
||||
if q == "*":
|
||||
logger.info("Performing match_all query")
|
||||
return {
|
||||
"bool": {
|
||||
"must": {"match_all": {}},
|
||||
"filter": {"bool": {"filter": filter_}},
|
||||
},
|
||||
}
|
||||
|
||||
hybrid_search_enabled = check_hybrid_search_enabled()
|
||||
if hybrid_search_enabled:
|
||||
embedding = embed_text(q)
|
||||
else:
|
||||
embedding = None
|
||||
|
||||
if not embedding:
|
||||
logger.info("Performing full-text search without embedding: %s", q)
|
||||
return {
|
||||
"bool": {
|
||||
"must": {
|
||||
"multi_match": {
|
||||
"query": q,
|
||||
# Give title more importance over content by a power of 3
|
||||
"fields": ["title.text^3", "content"],
|
||||
}
|
||||
},
|
||||
"filter": filter_,
|
||||
}
|
||||
}
|
||||
|
||||
logger.info("Performing hybrid search with embedding: %s", q)
|
||||
return {
|
||||
"hybrid": {
|
||||
"queries": [
|
||||
{
|
||||
"bool": {
|
||||
"must": {
|
||||
"multi_match": {
|
||||
"query": q,
|
||||
# Give title more importance over content by a power of 3
|
||||
"fields": ["title.text^3", "content"],
|
||||
}
|
||||
},
|
||||
"filter": filter_,
|
||||
}
|
||||
},
|
||||
{
|
||||
"bool": {
|
||||
"must": {
|
||||
"knn": {
|
||||
"embedding": {
|
||||
"vector": embedding,
|
||||
"k": nb_results,
|
||||
}
|
||||
}
|
||||
},
|
||||
"filter": filter_,
|
||||
}
|
||||
},
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def get_filter(reach, visited, user_sub, groups):
|
||||
"""Build OpenSearch filter"""
|
||||
filters = [
|
||||
{"term": {"is_active": True}}, # filter out inactive documents
|
||||
# Access control filters
|
||||
{
|
||||
"bool": {
|
||||
"should": [
|
||||
# Public or authenticated (not restricted)
|
||||
{
|
||||
"bool": {
|
||||
"must_not": {
|
||||
"term": {enums.REACH: enums.ReachEnum.RESTRICTED},
|
||||
},
|
||||
"must": {
|
||||
"terms": {"_id": sorted(visited)},
|
||||
},
|
||||
}
|
||||
},
|
||||
# Restricted: either user or group must match
|
||||
{"term": {enums.USERS: user_sub}},
|
||||
{"terms": {enums.GROUPS: groups}},
|
||||
],
|
||||
"minimum_should_match": 1,
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
# Optional reach filter
|
||||
if reach is not None:
|
||||
filters.append({"term": {enums.REACH: reach}})
|
||||
|
||||
return filters
|
||||
|
||||
|
||||
def get_sort(query_keys, order_by, order_direction):
|
||||
"""Build OpenSearch sort clause"""
|
||||
# Add sorting logic based on relevance or specified field
|
||||
if "hybrid" in query_keys:
|
||||
# sorting by other field than "_score" is not supported in hybrid search
|
||||
# see: https://github.com/opensearch-project/neural-search/issues/866
|
||||
return {"_score": {"order": order_direction}}
|
||||
if order_by == enums.RELEVANCE:
|
||||
return {"_score": {"order": order_direction}}
|
||||
|
||||
return {order_by: {"order": order_direction}}
|
||||
|
||||
|
||||
def get_params(query_keys):
|
||||
"""Build OpenSearch search parameters"""
|
||||
if "hybrid" in query_keys:
|
||||
return {"search_pipeline": settings.HYBRID_SEARCH_PIPELINE_ID}
|
||||
return {}
|
||||
|
||||
|
||||
def embed_document(document):
|
||||
"""Get embedding vector for a given document"""
|
||||
return embed_text(format_document(document.title, document.content))
|
||||
|
||||
|
||||
def format_document(title, content):
|
||||
"""Get the embedding input format for a document"""
|
||||
return f"<{title}>:<{content}>"
|
||||
|
||||
|
||||
def embed_text(text):
|
||||
"""
|
||||
Get embedding vector for the given text from any OpenAI-compatible embedding API
|
||||
"""
|
||||
response = requests.post(
|
||||
settings.EMBEDDING_API_PATH,
|
||||
headers={"Authorization": f"Bearer {settings.EMBEDDING_API_KEY}>"},
|
||||
json={
|
||||
"input": text,
|
||||
"model": settings.EMBEDDING_API_MODEL_NAME,
|
||||
"dimensions": settings.EMBEDDING_DIMENSION,
|
||||
"encoding_format": "float",
|
||||
},
|
||||
timeout=settings.EMBEDDING_REQUEST_TIMEOUT,
|
||||
)
|
||||
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except requests.HTTPError as e:
|
||||
logger.warning("embedding API request failed: %s", str(e))
|
||||
return None
|
||||
|
||||
try:
|
||||
embedding = response.json()["data"][0]["embedding"]
|
||||
except (KeyError, IndexError, TypeError):
|
||||
logger.warning("unexpected embedding response format: %s", response.text)
|
||||
return None
|
||||
|
||||
return embedding
|
||||
|
||||
|
||||
def ensure_index_exists(index_name):
|
||||
"""Create index if it does not exist"""
|
||||
try:
|
||||
opensearch_client().indices.get(index=index_name)
|
||||
except NotFoundError:
|
||||
logger.info("Creating index: %s", index_name)
|
||||
opensearch_client().indices.create(
|
||||
index=index_name,
|
||||
body={
|
||||
"settings": {"index.knn": True},
|
||||
"mappings": {
|
||||
"dynamic": "strict",
|
||||
"properties": {
|
||||
"id": {"type": "keyword"},
|
||||
"title": {
|
||||
"type": "keyword",
|
||||
"fields": {"text": {"type": "text"}},
|
||||
},
|
||||
"depth": {"type": "integer"},
|
||||
"path": {
|
||||
"type": "keyword",
|
||||
"fields": {"text": {"type": "text"}},
|
||||
},
|
||||
"numchild": {"type": "integer"},
|
||||
"content": {"type": "text"},
|
||||
"created_at": {"type": "date"},
|
||||
"updated_at": {"type": "date"},
|
||||
"size": {"type": "long"},
|
||||
"users": {"type": "keyword"},
|
||||
"groups": {"type": "keyword"},
|
||||
"reach": {"type": "keyword"},
|
||||
"is_active": {"type": "boolean"},
|
||||
"embedding": {
|
||||
# for simplicity, embedding is always present but is empty
|
||||
# when hybrid search is disabled
|
||||
"type": "knn_vector",
|
||||
"dimension": settings.EMBEDDING_DIMENSION,
|
||||
"method": {
|
||||
"engine": "lucene",
|
||||
"space_type": "l2",
|
||||
"name": "hnsw",
|
||||
"parameters": {},
|
||||
},
|
||||
},
|
||||
"embedding_model": {"type": "keyword"},
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@cache
|
||||
def check_hybrid_search_enabled():
|
||||
"""Check that all required environment variables are set for hybrid search."""
|
||||
if settings.HYBRID_SEARCH_ENABLED is not True:
|
||||
logger.info("Hybrid search is disabled via HYBRID_SEARCH_ENABLED setting")
|
||||
return False
|
||||
|
||||
required_vars = [
|
||||
"HYBRID_SEARCH_WEIGHTS",
|
||||
"EMBEDDING_API_PATH",
|
||||
"EMBEDDING_API_KEY",
|
||||
"EMBEDDING_API_MODEL_NAME",
|
||||
"EMBEDDING_DIMENSION",
|
||||
]
|
||||
missing_vars = [var for var in required_vars if not getattr(settings, var, None)]
|
||||
if missing_vars:
|
||||
logger.warning(
|
||||
"Missing variables for hybrid search: %s", ", ".join(missing_vars)
|
||||
)
|
||||
return False
|
||||
|
||||
return True
|
||||
@@ -0,0 +1,78 @@
|
||||
"""
|
||||
Unit test for `create_search_pipeline` command.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from django.core.management import call_command
|
||||
|
||||
import pytest
|
||||
|
||||
from core.services.opensearch import opensearch_client
|
||||
from core.tests.utils import (
|
||||
delete_search_pipeline,
|
||||
enable_hybrid_search,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def before_each():
|
||||
"""Delete search pipeline before each test"""
|
||||
delete_search_pipeline()
|
||||
yield
|
||||
delete_search_pipeline()
|
||||
|
||||
|
||||
def test_create_search_pipeline(settings, caplog):
|
||||
"""Test command create search pipeline"""
|
||||
# create documents and index them with hybrid search disabled
|
||||
|
||||
enable_hybrid_search(settings)
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
call_command("create_search_pipeline")
|
||||
|
||||
assert any(
|
||||
f"Creating search pipeline: {settings.HYBRID_SEARCH_PIPELINE_ID}" in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
|
||||
# calling get works without raising NotFoundError
|
||||
opensearch_client().search_pipeline.get(settings.HYBRID_SEARCH_PIPELINE_ID)
|
||||
|
||||
|
||||
def test_create_search_pipeline_but_it_exists_already(settings, caplog):
|
||||
"""Test command create search pipeline but it already exists"""
|
||||
# create documents and index them with hybrid search disabled
|
||||
|
||||
opensearch_client().transport.perform_request(
|
||||
method="PUT",
|
||||
url="/_search/pipeline/" + settings.HYBRID_SEARCH_PIPELINE_ID,
|
||||
body={
|
||||
"description": "Post processor for hybrid search",
|
||||
"phase_results_processors": [
|
||||
{
|
||||
"normalization-processor": {
|
||||
"combination": {
|
||||
"technique": "arithmetic_mean",
|
||||
"parameters": {"weights": settings.HYBRID_SEARCH_WEIGHTS},
|
||||
}
|
||||
}
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
call_command("create_search_pipeline")
|
||||
|
||||
assert any(
|
||||
"Search pipeline exists already" in message for message in caplog.messages
|
||||
)
|
||||
assert not any(
|
||||
f"Creating search pipeline: {settings.HYBRID_SEARCH_PIPELINE_ID}" in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
|
||||
# the pipeline is still here
|
||||
opensearch_client().search_pipeline.get(settings.HYBRID_SEARCH_PIPELINE_ID)
|
||||
@@ -0,0 +1,287 @@
|
||||
"""
|
||||
Unit test for `reindex_with_embedding` command.
|
||||
"""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from django.core.management import CommandError, call_command
|
||||
|
||||
import pytest
|
||||
import responses
|
||||
|
||||
from core.management.commands.reindex_with_embedding import (
|
||||
check_hybrid_search_enabled as check_hybrid_search_enabled_command,
|
||||
)
|
||||
from core.management.commands.reindex_with_embedding import (
|
||||
reindex_with_embedding,
|
||||
)
|
||||
from core.services.opensearch import check_hybrid_search_enabled, opensearch_client
|
||||
from core.tests.mock import albert_embedding_response
|
||||
from core.tests.utils import (
|
||||
bulk_create_documents,
|
||||
delete_search_pipeline,
|
||||
delete_test_indices,
|
||||
enable_hybrid_search,
|
||||
prepare_index,
|
||||
)
|
||||
|
||||
SERVICE_NAME = "test-index"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def before_each():
|
||||
"""Clear caches and delete search pipeline before each test"""
|
||||
clear_caches()
|
||||
yield
|
||||
clear_caches()
|
||||
|
||||
|
||||
def clear_caches():
|
||||
"""Clear caches used in opensearch service and factories"""
|
||||
check_hybrid_search_enabled.cache_clear()
|
||||
# the instance of check_hybrid_search_enabled used in utils.py
|
||||
# is different and must be cleared separately
|
||||
check_hybrid_search_enabled_command.cache_clear()
|
||||
delete_search_pipeline()
|
||||
delete_test_indices()
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_reindex_with_embedding_command(settings):
|
||||
"""Test command create indexes with embedding and search pipeline"""
|
||||
# create documents and index them with hybrid search disabled
|
||||
opensearch_client_ = opensearch_client()
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "wolf", "content": "wolves live in packs and hunt together"},
|
||||
{"title": "dog", "content": "dogs are loyal domestic animals"},
|
||||
{"title": "cat", "content": "cats are curious and independent pets"},
|
||||
]
|
||||
)
|
||||
prepare_index(SERVICE_NAME, documents)
|
||||
|
||||
# the index has not been embedded in the initial state
|
||||
initial_index = opensearch_client_.search(
|
||||
index=SERVICE_NAME, size=3, body={"query": {"match_all": {}}}
|
||||
)
|
||||
assert len(initial_index["hits"]["hits"]) == 3
|
||||
for embedded_hit in initial_index["hits"]["hits"]:
|
||||
assert embedded_hit["_source"]["embedding"] == None
|
||||
assert embedded_hit["_source"]["embedding_model"] is None
|
||||
|
||||
# enable hybrid search
|
||||
enable_hybrid_search(settings)
|
||||
check_hybrid_search_enabled_command.cache_clear()
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
|
||||
call_command("reindex_with_embedding", SERVICE_NAME)
|
||||
|
||||
opensearch_client_.indices.refresh(index=SERVICE_NAME)
|
||||
embedded_index = opensearch_client_.search(
|
||||
index=SERVICE_NAME, size=3, body={"query": {"match_all": {}}}
|
||||
)
|
||||
|
||||
# the source index has been replaced with embedding version
|
||||
assert len(embedded_index["hits"]["hits"]) == 3
|
||||
for embedded_hit in embedded_index["hits"]["hits"]:
|
||||
embedded_source = embedded_hit["_source"]
|
||||
# the index contains a embedding and embedding_model
|
||||
assert (
|
||||
embedded_source["embedding"]
|
||||
== albert_embedding_response.response["data"][0]["embedding"]
|
||||
)
|
||||
assert embedded_source["embedding_model"] == settings.EMBEDDING_API_MODEL_NAME
|
||||
# assert initial value have not been effected
|
||||
initial_hits = [
|
||||
hit_
|
||||
for hit_ in initial_index["hits"]["hits"]
|
||||
if hit_["_id"] == embedded_hit["_id"]
|
||||
]
|
||||
assert len(initial_hits) == 1
|
||||
initial_source = initial_hits[0]["_source"]
|
||||
assert initial_source["title"] == embedded_source["title"]
|
||||
assert initial_source["content"] == embedded_source["content"]
|
||||
assert initial_source["created_at"] == embedded_source["created_at"]
|
||||
assert initial_source["users"] == embedded_source["users"]
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_reindex_can_fail_and_restart(settings):
|
||||
"""Test command handles embedding errors gracefully and continues processing."""
|
||||
opensearch_client_ = opensearch_client()
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "wolf", "content": "wolves live in packs and hunt together"},
|
||||
{"title": "dog", "content": "dogs are loyal domestic animals"},
|
||||
{"title": "cat", "content": "cats are curious and independent pets"},
|
||||
]
|
||||
)
|
||||
prepare_index(SERVICE_NAME, documents)
|
||||
|
||||
# enable hybrid search after first indexing
|
||||
enable_hybrid_search(settings)
|
||||
check_hybrid_search_enabled_command.cache_clear()
|
||||
|
||||
# First call succeeds, second fails, third succeeds
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json={"error": "rate limit exceeded"},
|
||||
status=429,
|
||||
)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
|
||||
result = reindex_with_embedding(SERVICE_NAME)
|
||||
|
||||
# assert results reflect 2 successes and 1 failure
|
||||
assert result["nb_success_embedding"] == 2
|
||||
assert result["nb_failed_embedding"] == 1
|
||||
|
||||
# assert the index state
|
||||
opensearch_client_.indices.refresh(index=SERVICE_NAME)
|
||||
embedded_index = opensearch_client_.search(
|
||||
index=SERVICE_NAME, size=3, body={"query": {"match_all": {}}}
|
||||
)
|
||||
# Should have 2 documents with embeddings, 1 without due to error
|
||||
embedded_count = 0
|
||||
not_embedded_count = 0
|
||||
for hit in embedded_index["hits"]["hits"]:
|
||||
if hit["_source"].get("embedding"):
|
||||
embedded_count += 1
|
||||
assert (
|
||||
hit["_source"]["embedding_model"] == settings.EMBEDDING_API_MODEL_NAME
|
||||
)
|
||||
else:
|
||||
not_embedded_count += 1
|
||||
assert hit["_source"]["embedding_model"] is None
|
||||
assert embedded_count == 2
|
||||
assert not_embedded_count == 1
|
||||
|
||||
# the command can be run again to index failed items
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
result = reindex_with_embedding(SERVICE_NAME)
|
||||
|
||||
# assert results
|
||||
assert result["nb_success_embedding"] == 1
|
||||
assert result["nb_failed_embedding"] == 0
|
||||
|
||||
# assert there is now 1 more success and 0 failures
|
||||
opensearch_client_.indices.refresh(index=SERVICE_NAME)
|
||||
embedded_index = opensearch_client_.search(
|
||||
index=SERVICE_NAME, size=3, body={"query": {"match_all": {}}}
|
||||
)
|
||||
for hit in embedded_index["hits"]["hits"]:
|
||||
assert (
|
||||
hit["_source"]["embedding"]
|
||||
== albert_embedding_response.response["data"][0]["embedding"]
|
||||
)
|
||||
assert hit["_source"]["embedding_model"] == settings.EMBEDDING_API_MODEL_NAME
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_reindex_preserves_concurrent_updates(settings):
|
||||
"""
|
||||
Test that concurrent document updates don't get overwritten by reindexing.
|
||||
This test simulates the fallowing scenario:
|
||||
• the hybrid search is disabled
|
||||
• documents are created and indexed without indexing
|
||||
• the hybrid search is enabled
|
||||
• the reindexing is triggered
|
||||
• one document is updated while the reindexing is still running
|
||||
Because the updated document is modified after the hybrid search is enabled,
|
||||
it has properly been indexed with embedding, the reindexing command must
|
||||
ignore this document to preserve this latest update.
|
||||
"""
|
||||
opensearch_client_ = opensearch_client()
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "wolf", "content": "wolves live in packs and hunt together"},
|
||||
{"title": "dog", "content": "dogs are loyal domestic animals"},
|
||||
]
|
||||
)
|
||||
prepare_index(SERVICE_NAME, documents)
|
||||
enable_hybrid_search(settings)
|
||||
|
||||
updated_title = "updated dog"
|
||||
updated_embedding = [
|
||||
1.0
|
||||
] * settings.EMBEDDING_DIMENSION # dummy embedding to simulate concurrent update
|
||||
# add a side_effect on the search to simulate a concurrent update
|
||||
patch(
|
||||
"core.services.opensearch.opensearch_client_.search",
|
||||
side_effect=opensearch_client_.update(
|
||||
index=SERVICE_NAME,
|
||||
id=documents[1]["id"],
|
||||
body={
|
||||
"doc": {
|
||||
"title": updated_title,
|
||||
"embedding": updated_embedding,
|
||||
"embedding_model": settings.EMBEDDING_API_MODEL_NAME,
|
||||
}
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
result = reindex_with_embedding(SERVICE_NAME)
|
||||
assert result["nb_success_embedding"] == 2
|
||||
assert result["nb_failed_embedding"] == 0
|
||||
|
||||
opensearch_client_.indices.refresh(index=SERVICE_NAME)
|
||||
embedded_index = opensearch_client_.search(
|
||||
index=SERVICE_NAME, size=2, body={"query": {"match_all": {}}}
|
||||
)
|
||||
# Check that the latest update is preserved
|
||||
dog_doc = [
|
||||
hit
|
||||
for hit in embedded_index["hits"]["hits"]
|
||||
if hit["_source"]["title"] == updated_title
|
||||
]
|
||||
assert len(dog_doc) == 1
|
||||
assert dog_doc[0]["_source"]["embedding"] == updated_embedding
|
||||
assert dog_doc[0]["_source"]["embedding_model"] == settings.EMBEDDING_API_MODEL_NAME
|
||||
|
||||
|
||||
def test_reindex_command_but_hybrid_search_is_disabled():
|
||||
"""Test the `reindex_with_embedding` command fails when hybrid search is disabled."""
|
||||
with pytest.raises(CommandError) as err:
|
||||
call_command("reindex_with_embedding", SERVICE_NAME)
|
||||
|
||||
assert str(err.value) == "Hybrid search is not enabled or properly configured."
|
||||
|
||||
|
||||
def test_reindex_command_but_index_does_not_exist(settings):
|
||||
"""Test the `reindex_with_embedding` command fails when the idex does not exist."""
|
||||
wrong_index = "wrong-index-name"
|
||||
enable_hybrid_search(settings)
|
||||
|
||||
with pytest.raises(CommandError) as err:
|
||||
call_command("reindex_with_embedding", wrong_index)
|
||||
|
||||
assert str(err.value) == f"Index {wrong_index} does not exist."
|
||||
File diff suppressed because it is too large
Load Diff
@@ -8,7 +8,8 @@ from django.utils import timezone
|
||||
import pytest
|
||||
from rest_framework.test import APIClient
|
||||
|
||||
from core import factories, opensearch
|
||||
from core import factories
|
||||
from core.services import opensearch
|
||||
|
||||
pytestmark = pytest.mark.django_db
|
||||
|
||||
@@ -59,14 +60,15 @@ def test_api_documents_index_bulk_success():
|
||||
|
||||
def test_api_documents_index_bulk_ensure_index():
|
||||
"""A registered service should be create the opensearch index if need."""
|
||||
opensearch_client_ = opensearch.opensearch_client()
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = factories.DocumentSchemaFactory.build_batch(3)
|
||||
|
||||
# Delete the index
|
||||
opensearch.client.indices.delete(index="*test*")
|
||||
opensearch_client_.indices.delete(index="*test*")
|
||||
|
||||
with pytest.raises(opensearch.NotFoundError):
|
||||
opensearch.client.indices.get(index="test-service")
|
||||
opensearch_client_.indices.get(index="test-service")
|
||||
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/index/",
|
||||
@@ -81,7 +83,7 @@ def test_api_documents_index_bulk_ensure_index():
|
||||
assert [d["status"] for d in responses] == ["success"] * 3
|
||||
|
||||
# The index has been rebuilt
|
||||
opensearch.client.indices.get(index="test-service")
|
||||
opensearch_client_.indices.get(index="test-service")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -288,7 +290,7 @@ def test_api_documents_index_bulk_default(field, default_value):
|
||||
responses = response.json()
|
||||
assert [d["status"] for d in responses] == ["success"] * 3
|
||||
|
||||
indexed_document = opensearch.client.get(
|
||||
indexed_document = opensearch.opensearch_client().get(
|
||||
index=service.name, id=responses[0]["_id"]
|
||||
)["_source"]
|
||||
assert indexed_document[field] == default_value
|
||||
@@ -389,7 +391,7 @@ def test_api_documents_index_opensearch_errors():
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = factories.DocumentSchemaFactory.build_batch(3)
|
||||
|
||||
with mock.patch.object(opensearch.client, "bulk") as mock_bulk:
|
||||
with mock.patch.object(opensearch.opensearch_client(), "bulk") as mock_bulk:
|
||||
mock_bulk.return_value = {
|
||||
"items": [
|
||||
{"index": {"status": 201}},
|
||||
|
||||
@@ -5,13 +5,23 @@ import datetime
|
||||
from django.utils import timezone
|
||||
|
||||
import pytest
|
||||
import responses
|
||||
from rest_framework.test import APIClient
|
||||
|
||||
from core import factories, opensearch
|
||||
from core import factories
|
||||
from core.services import opensearch
|
||||
from core.tests.mock import albert_embedding_response
|
||||
from core.tests.utils import delete_test_indices, enable_hybrid_search
|
||||
|
||||
pytestmark = pytest.mark.django_db
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_caches():
|
||||
"""Clear caches and delete search pipeline before each test"""
|
||||
opensearch.check_hybrid_search_enabled.cache_clear()
|
||||
|
||||
|
||||
def test_api_documents_index_single_anonymous():
|
||||
"""Anonymous requests should not be allowed to index documents."""
|
||||
document = factories.DocumentSchemaFactory.build()
|
||||
@@ -39,9 +49,21 @@ def test_api_documents_index_single_invalid_token():
|
||||
assert response.json() == {"detail": "Invalid token."}
|
||||
|
||||
|
||||
def test_api_documents_index_single_success():
|
||||
"""A registered service should be able to index document with a valid token."""
|
||||
@responses.activate
|
||||
def test_api_documents_index_single_hybrid_enabled_success(settings):
|
||||
"""
|
||||
A registered service should be able to index document with a valid token.
|
||||
If hybrid search is enabled, the indexing should have embedding of
|
||||
dimension settings.EMBEDDING_DIMENSION.
|
||||
"""
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
enable_hybrid_search(settings)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
document = factories.DocumentSchemaFactory.build()
|
||||
|
||||
response = APIClient().post(
|
||||
@@ -54,17 +76,56 @@ def test_api_documents_index_single_success():
|
||||
assert response.status_code == 201
|
||||
assert response.json()["_id"] == str(document["id"])
|
||||
|
||||
new_indexed_document = opensearch.opensearch_client().get(
|
||||
index=service.name, id=str(document["id"])
|
||||
)
|
||||
assert new_indexed_document["_version"] == 1
|
||||
assert new_indexed_document["_source"]["title"] == document["title"].strip().lower()
|
||||
assert new_indexed_document["_source"]["content"] == document["content"]
|
||||
assert (
|
||||
new_indexed_document["_source"]["embedding"]
|
||||
== albert_embedding_response.response["data"][0]["embedding"]
|
||||
)
|
||||
assert (
|
||||
new_indexed_document["_source"]["embedding_model"]
|
||||
== settings.EMBEDDING_API_MODEL_NAME
|
||||
)
|
||||
|
||||
def test_api_documents_index_bulk_ensure_index():
|
||||
|
||||
def test_api_documents_index_single_hybrid_disabled_success():
|
||||
"""If hybrid search is not enabled, the indexing should have an embedding equal to None."""
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
document = factories.DocumentSchemaFactory.build()
|
||||
opensearch.check_hybrid_search_enabled.cache_clear()
|
||||
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/index/",
|
||||
document,
|
||||
HTTP_AUTHORIZATION=f"Bearer {service.token:s}",
|
||||
format="json",
|
||||
)
|
||||
|
||||
assert response.status_code == 201
|
||||
assert response.json()["_id"] == str(document["id"])
|
||||
|
||||
new_indexed_document = opensearch.opensearch_client().get(
|
||||
index=service.name, id=str(document["id"])
|
||||
)
|
||||
assert new_indexed_document["_version"] == 1
|
||||
assert new_indexed_document["_source"]["title"] == document["title"].strip().lower()
|
||||
assert new_indexed_document["_source"]["content"] == document["content"]
|
||||
assert new_indexed_document["_source"]["embedding"] is None
|
||||
|
||||
|
||||
def test_api_documents_index_bulk_ensure_index(settings):
|
||||
"""A registered service should be create the opensearch index if need."""
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
document = factories.DocumentSchemaFactory.build()
|
||||
|
||||
# Delete the index
|
||||
opensearch.client.indices.delete(index="*test*")
|
||||
opensearch_client_ = opensearch.opensearch_client()
|
||||
delete_test_indices()
|
||||
|
||||
with pytest.raises(opensearch.NotFoundError):
|
||||
opensearch.client.indices.get(index="test-service")
|
||||
opensearch_client_.indices.get(index="test-service")
|
||||
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/index/",
|
||||
@@ -77,7 +138,7 @@ def test_api_documents_index_bulk_ensure_index():
|
||||
assert response.json()["_id"] == str(document["id"])
|
||||
|
||||
# The index has been rebuilt
|
||||
data = opensearch.client.indices.get(index="test-service")
|
||||
data = opensearch_client_.indices.get(index="test-service")
|
||||
|
||||
assert data["test-service"]["mappings"] == {
|
||||
"dynamic": "strict",
|
||||
@@ -103,6 +164,17 @@ def test_api_documents_index_bulk_ensure_index():
|
||||
"groups": {"type": "keyword"},
|
||||
"reach": {"type": "keyword"},
|
||||
"is_active": {"type": "boolean"},
|
||||
"embedding": {
|
||||
"type": "knn_vector",
|
||||
"dimension": settings.EMBEDDING_DIMENSION,
|
||||
"method": {
|
||||
"engine": "lucene",
|
||||
"space_type": "l2",
|
||||
"name": "hnsw",
|
||||
"parameters": {},
|
||||
},
|
||||
},
|
||||
"embedding_model": {"type": "keyword"},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -300,7 +372,7 @@ def test_api_documents_index_single_default(field, default_value):
|
||||
assert response.status_code == 201
|
||||
assert response.json()["_id"] == str(document["id"])
|
||||
|
||||
indexed_document = opensearch.client.get(
|
||||
indexed_document = opensearch.opensearch_client().get(
|
||||
index=service.name, id=str(document["id"])
|
||||
)["_source"]
|
||||
assert indexed_document[field] == default_value
|
||||
|
||||
@@ -13,12 +13,34 @@ import responses
|
||||
from rest_framework.test import APIClient
|
||||
|
||||
from core import enums, factories
|
||||
from core.services.opensearch import check_hybrid_search_enabled, opensearch_client
|
||||
|
||||
from .utils import build_authorization_bearer, prepare_index, setup_oicd_resource_server
|
||||
from .mock import albert_embedding_response
|
||||
from .utils import (
|
||||
build_authorization_bearer,
|
||||
bulk_create_documents,
|
||||
enable_hybrid_search,
|
||||
prepare_index,
|
||||
setup_oicd_resource_server,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.django_db
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def before_each():
|
||||
"""Clear cached functions before each test to avoid side effects"""
|
||||
clear_caches()
|
||||
yield
|
||||
clear_caches()
|
||||
|
||||
|
||||
def clear_caches():
|
||||
"""Clear cached functions before each test to avoid side effects"""
|
||||
opensearch_client.cache_clear()
|
||||
check_hybrid_search_enabled.cache_clear()
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_search_auth_invalid_parameters(settings):
|
||||
"""Invalid service parameters should result in a 401 error"""
|
||||
@@ -41,6 +63,31 @@ def test_api_documents_search_auth_invalid_parameters(settings):
|
||||
assert response.json() == {"detail": "Resource Server is improperly configured"}
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_search_opensearch_env_variables_not_set(settings):
|
||||
"""
|
||||
Missing environment variables for OpenSearch client should
|
||||
result in a 500 internal server error
|
||||
"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
factories.ServiceFactory(name="test-service")
|
||||
|
||||
del settings.OPENSEARCH_HOST # Remove required settings
|
||||
del settings.OPENSEARCH_PASSWORD
|
||||
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{"q": "*"},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {build_authorization_bearer()}",
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert response.json() == [
|
||||
"Missing required OpenSearch environment variables: OPENSEARCH_HOST, OPENSEARCH_PASSWORD"
|
||||
]
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_search_query_unknown_user(settings):
|
||||
"""Searching a document without an existing user should result in a 401 error"""
|
||||
@@ -69,17 +116,22 @@ def test_api_documents_search_query_unknown_user(settings):
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_search_services_invalid_parameters(settings):
|
||||
"""Invalid pagination parameters should result in a 400 error"""
|
||||
"""Invalid services parameter should result in a 400 error"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
factories.ServiceFactory(name="test-service")
|
||||
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
# services should be a list
|
||||
{"q": "a quick fox", "services": {}},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
HTTP_AUTHORIZATION=f"Bearer {build_authorization_bearer()}",
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
@@ -94,17 +146,22 @@ def test_api_documents_search_services_invalid_parameters(settings):
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_search_reached_docs_invalid_parameters(settings):
|
||||
"""Invalid pagination parameters should result in a 400 error"""
|
||||
"""Invalid visited parameters should result in a 400 error"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
factories.ServiceFactory(name="test-service")
|
||||
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
# visited should be a list
|
||||
{"q": "a quick fox", "visited": {}},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
HTTP_AUTHORIZATION=f"Bearer {build_authorization_bearer()}",
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
@@ -118,142 +175,88 @@ def test_api_documents_search_reached_docs_invalid_parameters(settings):
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_search_query_title(settings):
|
||||
"""Searching a document by its title should work as expected"""
|
||||
def test_api_documents_search_match_all(settings):
|
||||
"""Searching a document with q='*' should match all docs"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
nb_documents = 12
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
document = factories.DocumentSchemaFactory.build(
|
||||
title="The quick brown fox",
|
||||
content="the wolf",
|
||||
reach=random.choice(["public", "authenticated"]),
|
||||
documents = factories.DocumentSchemaFactory.build_batch(
|
||||
nb_documents, reach=random.choice(["public", "authenticated"])
|
||||
)
|
||||
|
||||
# Add other documents
|
||||
other_fox_document = factories.DocumentSchemaFactory.build(
|
||||
title="The blue fox",
|
||||
content="the wolf",
|
||||
reach=random.choice(["public", "authenticated"]),
|
||||
)
|
||||
no_fox_document = factories.DocumentSchemaFactory.build(
|
||||
title="The brown goat",
|
||||
content="the wolf",
|
||||
reach=random.choice(["public", "authenticated"]),
|
||||
)
|
||||
documents = [document, other_fox_document, no_fox_document]
|
||||
prepare_index(service.name, documents)
|
||||
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{"q": "a quick fox", "visited": [doc["id"] for doc in documents]},
|
||||
{"q": "*", "visited": [doc["id"] for doc in documents]},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
assert len(response.json()) == nb_documents
|
||||
|
||||
fox_data = response.json()[0]
|
||||
assert list(fox_data.keys()) == ["_index", "_id", "_score", "_source", "fields"]
|
||||
assert fox_data["_id"] == str(document["id"])
|
||||
assert fox_data["_source"] == {
|
||||
"depth": 1,
|
||||
"numchild": 0,
|
||||
"path": document["path"],
|
||||
"size": document["size"],
|
||||
"created_at": document["created_at"].isoformat(),
|
||||
"updated_at": document["updated_at"].isoformat(),
|
||||
"reach": document["reach"],
|
||||
"title": "The quick brown fox",
|
||||
}
|
||||
assert fox_data["fields"] == {"number_of_users": [3], "number_of_groups": [3]}
|
||||
|
||||
other_fox_data = response.json()[1]
|
||||
assert list(other_fox_data.keys()) == [
|
||||
"_index",
|
||||
"_id",
|
||||
"_score",
|
||||
"_source",
|
||||
"fields",
|
||||
]
|
||||
assert other_fox_data["_id"] == str(other_fox_document["id"])
|
||||
assert other_fox_data["_source"] == {
|
||||
"depth": 1,
|
||||
"numchild": 0,
|
||||
"path": other_fox_document["path"],
|
||||
"size": other_fox_document["size"],
|
||||
"created_at": other_fox_document["created_at"].isoformat(),
|
||||
"updated_at": other_fox_document["updated_at"].isoformat(),
|
||||
"reach": other_fox_document["reach"],
|
||||
"title": "The blue fox",
|
||||
}
|
||||
assert other_fox_data["fields"] == {"number_of_users": [3], "number_of_groups": [3]}
|
||||
assert [r["_id"] for r in response.json()] == [str(doc["id"]) for doc in documents]
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_search_query_content(settings):
|
||||
"""Searching a document by its content should work as expected"""
|
||||
def test_api_documents_full_text_search_query_title(settings):
|
||||
"""Searching a document by its title should work as expected"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
document = factories.DocumentSchemaFactory.build(
|
||||
title="the wolf",
|
||||
content="The quick brown fox",
|
||||
reach=random.choice(["public", "authenticated"]),
|
||||
)
|
||||
|
||||
# Add other documents
|
||||
other_fox_document = factories.DocumentSchemaFactory.build(
|
||||
title="the wolf",
|
||||
content="The blue fox",
|
||||
reach=random.choice(["public", "authenticated"]),
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "The quick brown fox", "content": "the wolf"},
|
||||
{"title": "The blue fox", "content": "the wolf"},
|
||||
{"title": "The brown goat", "content": "the wolf"},
|
||||
]
|
||||
)
|
||||
no_fox_document = factories.DocumentSchemaFactory.build(
|
||||
title="the wolf",
|
||||
content="The brown goat",
|
||||
reach=random.choice(["public", "authenticated"]),
|
||||
)
|
||||
|
||||
documents = [document, other_fox_document, no_fox_document]
|
||||
prepare_index(service.name, documents)
|
||||
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{"q": "a quick fox", "visited": [doc["id"] for doc in documents]},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
HTTP_AUTHORIZATION=f"Bearer {build_authorization_bearer()}",
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
|
||||
fox_data = response.json()[0]
|
||||
assert list(fox_data.keys()) == ["_index", "_id", "_score", "_source", "fields"]
|
||||
assert fox_data["_id"] == str(document["id"])
|
||||
assert fox_data["_source"] == {
|
||||
fox_response = response.json()[0]
|
||||
fox_document = documents[0]
|
||||
assert list(fox_response.keys()) == ["_index", "_id", "_score", "_source", "fields"]
|
||||
assert fox_response["_id"] == str(documents[0]["id"])
|
||||
assert fox_response["_source"] == {
|
||||
"depth": 1,
|
||||
"numchild": 0,
|
||||
"path": document["path"],
|
||||
"size": document["size"],
|
||||
"created_at": document["created_at"].isoformat(),
|
||||
"updated_at": document["updated_at"].isoformat(),
|
||||
"reach": document["reach"],
|
||||
"title": document["title"],
|
||||
"path": fox_document["path"],
|
||||
"size": fox_document["size"],
|
||||
"created_at": fox_document["created_at"].isoformat(),
|
||||
"updated_at": fox_document["updated_at"].isoformat(),
|
||||
"reach": fox_document["reach"],
|
||||
"title": fox_document["title"],
|
||||
}
|
||||
assert fox_data["fields"] == {"number_of_users": [3], "number_of_groups": [3]}
|
||||
assert fox_response["fields"] == {"number_of_users": [1], "number_of_groups": [3]}
|
||||
|
||||
other_fox_data = response.json()[1]
|
||||
assert list(other_fox_data.keys()) == [
|
||||
other_fox_response = response.json()[1]
|
||||
other_fox_document = documents[1]
|
||||
assert list(other_fox_response.keys()) == [
|
||||
"_index",
|
||||
"_id",
|
||||
"_score",
|
||||
"_source",
|
||||
"fields",
|
||||
]
|
||||
assert other_fox_data["_id"] == str(other_fox_document["id"])
|
||||
assert other_fox_data["_source"] == {
|
||||
assert other_fox_response["_id"] == str(other_fox_document["id"])
|
||||
assert other_fox_response["_source"] == {
|
||||
"depth": 1,
|
||||
"numchild": 0,
|
||||
"path": other_fox_document["path"],
|
||||
@@ -263,7 +266,185 @@ def test_api_documents_search_query_content(settings):
|
||||
"reach": other_fox_document["reach"],
|
||||
"title": other_fox_document["title"],
|
||||
}
|
||||
assert other_fox_data["fields"] == {"number_of_users": [3], "number_of_groups": [3]}
|
||||
assert other_fox_response["fields"] == {
|
||||
"number_of_users": [1],
|
||||
"number_of_groups": [3],
|
||||
}
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_full_text_search(settings):
|
||||
"""Searching a document by its content should work as expected"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "The quick brown fox", "content": "the wolf"},
|
||||
{"title": "The blue fox", "content": "the wolf"},
|
||||
{"title": "The brown goat", "content": "the wolf"},
|
||||
]
|
||||
)
|
||||
prepare_index(service.name, documents)
|
||||
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{"q": "a quick fox", "visited": [doc["id"] for doc in documents]},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
|
||||
fox_response = response.json()[0]
|
||||
fox_document = documents[0]
|
||||
assert list(fox_response.keys()) == ["_index", "_id", "_score", "_source", "fields"]
|
||||
assert fox_response["_id"] == str(fox_document["id"])
|
||||
assert fox_response["_score"] > 0
|
||||
assert fox_response["_source"] == {
|
||||
"depth": 1,
|
||||
"numchild": 0,
|
||||
"path": fox_document["path"],
|
||||
"size": fox_document["size"],
|
||||
"created_at": fox_document["created_at"].isoformat(),
|
||||
"updated_at": fox_document["updated_at"].isoformat(),
|
||||
"reach": fox_document["reach"],
|
||||
"title": fox_document["title"],
|
||||
}
|
||||
assert fox_response["fields"] == {"number_of_users": [1], "number_of_groups": [3]}
|
||||
|
||||
other_fox_response = response.json()[1]
|
||||
other_fox_document = documents[1]
|
||||
assert list(other_fox_response.keys()) == [
|
||||
"_index",
|
||||
"_id",
|
||||
"_score",
|
||||
"_source",
|
||||
"fields",
|
||||
]
|
||||
assert other_fox_response["_id"] == str(other_fox_document["id"])
|
||||
assert other_fox_response["_score"] > 0
|
||||
assert other_fox_response["_source"] == {
|
||||
"depth": 1,
|
||||
"numchild": 0,
|
||||
"path": other_fox_document["path"],
|
||||
"size": other_fox_document["size"],
|
||||
"created_at": other_fox_document["created_at"].isoformat(),
|
||||
"updated_at": other_fox_document["updated_at"].isoformat(),
|
||||
"reach": other_fox_document["reach"],
|
||||
"title": other_fox_document["title"],
|
||||
}
|
||||
assert other_fox_response["fields"] == {
|
||||
"number_of_users": [1],
|
||||
"number_of_groups": [3],
|
||||
}
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_hybrid_search(settings):
|
||||
"""Searching a document by its content should work as expected"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
# hybrid search is enabled by default
|
||||
enable_hybrid_search(settings)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
) # mock embedding API
|
||||
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "The quick brown fox", "content": "the wolf"},
|
||||
{"title": "The blue fox", "content": "the wolf"},
|
||||
{"title": "The brown goat", "content": "the wolf"},
|
||||
]
|
||||
)
|
||||
prepare_index(service.name, documents)
|
||||
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{"q": "a quick fox", "visited": [doc["id"] for doc in documents]},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert (
|
||||
len(response.json()) == 3
|
||||
) # semantic search always returns a response of size nb_results
|
||||
|
||||
fox_response = response.json()[0]
|
||||
fox_document = documents[0]
|
||||
assert list(fox_response.keys()) == ["_index", "_id", "_score", "_source", "fields"]
|
||||
assert fox_response["_id"] == str(fox_document["id"])
|
||||
assert fox_response["_score"] > 0
|
||||
assert fox_response["_source"] == {
|
||||
"depth": 1,
|
||||
"numchild": 0,
|
||||
"path": fox_document["path"],
|
||||
"size": fox_document["size"],
|
||||
"created_at": fox_document["created_at"].isoformat(),
|
||||
"updated_at": fox_document["updated_at"].isoformat(),
|
||||
"reach": fox_document["reach"],
|
||||
"title": fox_document["title"],
|
||||
}
|
||||
assert fox_response["fields"] == {"number_of_users": [1], "number_of_groups": [3]}
|
||||
|
||||
other_fox_response = response.json()[1]
|
||||
other_fox_document = documents[1]
|
||||
assert list(other_fox_response.keys()) == [
|
||||
"_index",
|
||||
"_id",
|
||||
"_score",
|
||||
"_source",
|
||||
"fields",
|
||||
]
|
||||
assert other_fox_response["_id"] == str(other_fox_document["id"])
|
||||
assert other_fox_response["_score"] > 0
|
||||
assert other_fox_response["_source"] == {
|
||||
"depth": 1,
|
||||
"numchild": 0,
|
||||
"path": other_fox_document["path"],
|
||||
"size": other_fox_document["size"],
|
||||
"created_at": other_fox_document["created_at"].isoformat(),
|
||||
"updated_at": other_fox_document["updated_at"].isoformat(),
|
||||
"reach": other_fox_document["reach"],
|
||||
"title": other_fox_document["title"],
|
||||
}
|
||||
assert other_fox_response["fields"] == {
|
||||
"number_of_users": [1],
|
||||
"number_of_groups": [3],
|
||||
}
|
||||
|
||||
no_fox_response = response.json()[2]
|
||||
no_fox_document = documents[2]
|
||||
assert list(no_fox_response.keys()) == [
|
||||
"_index",
|
||||
"_id",
|
||||
"_score",
|
||||
"_source",
|
||||
"fields",
|
||||
]
|
||||
assert no_fox_response["_id"] == str(no_fox_document["id"])
|
||||
assert no_fox_response["_source"] == {
|
||||
"depth": 1,
|
||||
"numchild": 0,
|
||||
"path": no_fox_document["path"],
|
||||
"size": no_fox_document["size"],
|
||||
"created_at": no_fox_document["created_at"].isoformat(),
|
||||
"updated_at": no_fox_document["updated_at"].isoformat(),
|
||||
"reach": no_fox_document["reach"],
|
||||
"title": no_fox_document["title"],
|
||||
}
|
||||
assert no_fox_response["fields"] == {
|
||||
"number_of_users": [1],
|
||||
"number_of_groups": [3],
|
||||
}
|
||||
|
||||
|
||||
@responses.activate
|
||||
@@ -271,7 +452,12 @@ def test_api_documents_search_ordering_by_fields(settings):
|
||||
"""It should be possible to order by several fields"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = factories.DocumentSchemaFactory.build_batch(
|
||||
4, reach=random.choice(["public", "authenticated"])
|
||||
@@ -319,7 +505,12 @@ def test_api_documents_search_ordering_by_relevance(settings):
|
||||
"""It should be possible to order by relevance (score)"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = factories.DocumentSchemaFactory.build_batch(
|
||||
4, reach=random.choice(["public", "authenticated"])
|
||||
@@ -354,7 +545,12 @@ def test_api_documents_search_ordering_by_unknown_field(settings):
|
||||
"""Trying to sort by an unknown field should return a 400 error"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
# Setup: Initialize the service and documents only once
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = factories.DocumentSchemaFactory.build_batch(
|
||||
@@ -397,7 +593,12 @@ def test_api_documents_search_ordering_by_unknown_direction(settings):
|
||||
"""Trying to sort with an unknown direction should return a 400 error"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = factories.DocumentSchemaFactory.build_batch(
|
||||
2, reach=random.choice(["public", "authenticated"])
|
||||
@@ -432,7 +633,12 @@ def test_api_documents_search_filtering_by_reach(settings):
|
||||
"""It should be possible to filter results by their reach"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = factories.DocumentSchemaFactory.build_batch(
|
||||
4, reach=random.choice(["public", "authenticated"])
|
||||
@@ -458,15 +664,17 @@ def test_api_documents_search_filtering_by_reach(settings):
|
||||
assert reach == result["_source"]["reach"]
|
||||
|
||||
|
||||
# Pagination
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_search_pagination_basic(settings):
|
||||
"""Pagination should correctly return documents for the specified page and page size"""
|
||||
def test_api_documents_search_with_nb_results(settings):
|
||||
"""nb_size should correctly return results of given size"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = factories.DocumentSchemaFactory.build_batch(
|
||||
9, reach=random.choice(["public", "authenticated"])
|
||||
@@ -474,13 +682,12 @@ def test_api_documents_search_pagination_basic(settings):
|
||||
ids = [str(doc["id"]) for doc in documents]
|
||||
prepare_index(service.name, documents)
|
||||
|
||||
# Request the first page with a page size of 3
|
||||
nb_results = 3
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{
|
||||
"q": "*",
|
||||
"page_number": 1,
|
||||
"page_size": 3,
|
||||
"nb_results": nb_results,
|
||||
"visited": [doc["id"] for doc in documents],
|
||||
},
|
||||
format="json",
|
||||
@@ -489,16 +696,14 @@ def test_api_documents_search_pagination_basic(settings):
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 3 # Page size is 3
|
||||
assert [r["_id"] for r in data] == ids[0:3]
|
||||
assert [r["_id"] for r in data] == ids[0:nb_results]
|
||||
|
||||
# Request the second page with a page size of 3
|
||||
nb_results = 6
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{
|
||||
"q": "*",
|
||||
"page_number": 2,
|
||||
"page_size": 3,
|
||||
"nb_results": nb_results,
|
||||
"visited": [doc["id"] for doc in documents],
|
||||
},
|
||||
format="json",
|
||||
@@ -506,113 +711,36 @@ def test_api_documents_search_pagination_basic(settings):
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 3
|
||||
assert [r["_id"] for r in data] == ids[3:6]
|
||||
assert [r["_id"] for r in data] == ids[0:nb_results]
|
||||
|
||||
# Request the third page with a page size of 5 (should contain the remaining 3 documents)
|
||||
nb_results = 10
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{
|
||||
"q": "*",
|
||||
"page_number": 3,
|
||||
"page_size": 3,
|
||||
"nb_results": nb_results,
|
||||
"visited": [doc["id"] for doc in documents],
|
||||
},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 3
|
||||
assert [r["_id"] for r in data] == ids[6:9]
|
||||
# nb_results > total number of documents => returns all documents
|
||||
assert [r["_id"] for r in data] == ids[0:9]
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_search_pagination_last_page_edge_case(settings):
|
||||
"""Requesting the last page should return the correct number of remaining documents"""
|
||||
def test_api_documents_search_nb_results_invalid_parameters(settings):
|
||||
"""Invalid nb_results parameters should result in a 400 error"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = factories.DocumentSchemaFactory.build_batch(
|
||||
8, reach=random.choice(["public", "authenticated"])
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
ids = [str(doc["id"]) for doc in documents]
|
||||
prepare_index(service.name, documents)
|
||||
|
||||
# Request the first page with a page size of 3
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{
|
||||
"q": "*",
|
||||
"page_number": 1,
|
||||
"page_size": 3,
|
||||
"visited": [doc["id"] for doc in documents],
|
||||
},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 3
|
||||
assert [r["_id"] for r in response.json()] == ids[0:3]
|
||||
|
||||
# Request the third page with a page size of 3 (should contain the last 1 document)
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{
|
||||
"q": "*",
|
||||
"page_number": 3,
|
||||
"page_size": 3,
|
||||
"visited": [doc["id"] for doc in documents],
|
||||
},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2 # Only 2 documents should be on the last page
|
||||
assert [r["_id"] for r in response.json()] == ids[6:]
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_search_pagination_out_of_bounds(settings):
|
||||
"""
|
||||
Requesting a page number that exceeds the total number of pages should return an empty list
|
||||
"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = factories.DocumentSchemaFactory.build_batch(
|
||||
4, reach=random.choice(["public", "authenticated"])
|
||||
)
|
||||
prepare_index(service.name, documents)
|
||||
|
||||
# Request the fourth page with a page size of 2 (there are only 2 pages)
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{
|
||||
"q": "*",
|
||||
"page_number": 4,
|
||||
"page_size": 2,
|
||||
"visited": [doc["id"] for doc in documents],
|
||||
},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 0 # No documents should be returned
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_search_pagination_invalid_parameters(settings):
|
||||
"""Invalid pagination parameters should result in a 400 error"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = factories.DocumentSchemaFactory.build_batch(
|
||||
4, reach=random.choice(["public", "authenticated"])
|
||||
@@ -622,26 +750,18 @@ def test_api_documents_search_pagination_invalid_parameters(settings):
|
||||
parameters = [
|
||||
(
|
||||
"invalid",
|
||||
10,
|
||||
"int_parsing",
|
||||
"Input should be a valid integer, unable to parse string as an integer",
|
||||
),
|
||||
(
|
||||
1,
|
||||
"invalid",
|
||||
"int_parsing",
|
||||
"Input should be a valid integer, unable to parse string as an integer",
|
||||
),
|
||||
(-1, 10, "greater_than_equal", "Input should be greater than or equal to 1"),
|
||||
(1, -10, "greater_than_equal", "Input should be greater than or equal to 1"),
|
||||
(0, 10, "greater_than_equal", "Input should be greater than or equal to 1"),
|
||||
(1, 0, "greater_than_equal", "Input should be greater than or equal to 1"),
|
||||
(-1, "greater_than_equal", "Input should be greater than or equal to 1"),
|
||||
(0, "greater_than_equal", "Input should be greater than or equal to 1"),
|
||||
(350, "less_than_equal", "Input should be less than or equal to 300"),
|
||||
]
|
||||
|
||||
for page_number, page_size, error_type, error_message in parameters:
|
||||
for nb_results, error_type, error_message in parameters:
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{"q": "*", "page_number": page_number, "page_size": page_size},
|
||||
{"q": "*", "nb_results": nb_results},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
)
|
||||
@@ -652,11 +772,16 @@ def test_api_documents_search_pagination_invalid_parameters(settings):
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_api_documents_search_pagination_with_filtering(settings):
|
||||
"""Pagination should work correctly when combined with filtering by reach"""
|
||||
def test_api_documents_search_nb_results_with_filtering(settings):
|
||||
"""nb_results should work correctly when combined with filtering by reach"""
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
public_documents = factories.DocumentSchemaFactory.build_batch(3, reach="public")
|
||||
public_ids = [str(doc["id"]) for doc in public_documents]
|
||||
@@ -665,37 +790,17 @@ def test_api_documents_search_pagination_with_filtering(settings):
|
||||
)
|
||||
prepare_index(service.name, public_documents + private_documents)
|
||||
|
||||
# Filter by public documents, request first page
|
||||
nb_results = 3
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{
|
||||
"q": "*",
|
||||
"reach": "public",
|
||||
"page_number": 1,
|
||||
"page_size": 2,
|
||||
"nb_results": nb_results,
|
||||
"visited": public_ids,
|
||||
},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
assert [r["_id"] for r in response.json()] == public_ids[0:2]
|
||||
|
||||
# Request second page for public documents (remaining 1 document)
|
||||
response = APIClient().post(
|
||||
"/api/v1.0/documents/search/",
|
||||
{
|
||||
"q": "*",
|
||||
"reach": "public",
|
||||
"page_number": 2,
|
||||
"page_size": 2,
|
||||
"visited": public_ids,
|
||||
},
|
||||
format="json",
|
||||
HTTP_AUTHORIZATION=f"Bearer {token}",
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 1
|
||||
assert [r["_id"] for r in response.json()] == public_ids[2:]
|
||||
assert [r["_id"] for r in response.json()] == public_ids[0:nb_results]
|
||||
|
||||
@@ -10,7 +10,9 @@ import responses
|
||||
from rest_framework.test import APIClient
|
||||
|
||||
from core import enums, factories
|
||||
from core.services.opensearch import opensearch_client
|
||||
|
||||
from .mock import albert_embedding_response
|
||||
from .utils import (
|
||||
build_authorization_bearer,
|
||||
delete_test_indices,
|
||||
@@ -21,8 +23,15 @@ from .utils import (
|
||||
pytestmark = pytest.mark.django_db
|
||||
|
||||
|
||||
def test_api_documents_search_access_control_anonymous():
|
||||
@responses.activate
|
||||
def test_api_documents_search_access_control_anonymous(settings):
|
||||
"""Anonymous users should not be allowed to search documents even public."""
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
service = factories.ServiceFactory(name="test-service")
|
||||
documents = []
|
||||
for reach in enums.ReachEnum:
|
||||
@@ -43,6 +52,12 @@ def test_api_documents_search_access_control(settings):
|
||||
- only configured services providers are allowed (e.g docs)
|
||||
(groups is not yet implemnted)
|
||||
"""
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
@@ -90,6 +105,12 @@ def test_api_documents_search_access__only_visited_public(
|
||||
Authenticated users should only see documents with reach="public"
|
||||
that are in "visited" list.
|
||||
"""
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub", audience="docs")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
@@ -121,6 +142,12 @@ def test_api_documents_search_access__any_owner_public(settings):
|
||||
Authenticated users should only see documents with reach="public"
|
||||
that are in "visited" list.
|
||||
"""
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub", audience="docs")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
@@ -159,6 +186,12 @@ def test_api_documents_search_access__services(settings):
|
||||
Authenticated users should only see documents of audience
|
||||
service providers (e.g docs)
|
||||
"""
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub", audience="a-client")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
@@ -193,11 +226,18 @@ def test_api_documents_search_access__missing_index(settings):
|
||||
"""
|
||||
When the service has no opensearch index, returns an empty list.
|
||||
"""
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub", audience="a-client")
|
||||
token = build_authorization_bearer()
|
||||
factories.ServiceFactory(name="test-index-a", client_id="a-client")
|
||||
|
||||
delete_test_indices()
|
||||
opensearch_client.cache_clear()
|
||||
|
||||
# a-client has no index. ignore it.
|
||||
response = APIClient().post(
|
||||
@@ -217,6 +257,12 @@ def test_api_documents_search_access__related_services(settings):
|
||||
Authenticated users should only see documents of audience
|
||||
service providers and its related services (e.g drive + docs)
|
||||
"""
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub", audience="c-client")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
@@ -258,6 +304,12 @@ def test_api_documents_search_access__related_missing_index(settings):
|
||||
"""
|
||||
When the service has no opensearch index, returns the related services data.
|
||||
"""
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub", audience="a-client")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
@@ -298,6 +350,12 @@ def test_api_documents_search_access__request_services(settings):
|
||||
from requested services : 'services' parameter.
|
||||
Raise 400 error if not all requested services are authorized.
|
||||
"""
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub", audience="c-client")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
@@ -386,6 +444,12 @@ def test_api_documents_search_access__authenticated(settings):
|
||||
- only configured services providers are allowed (e.g docs)
|
||||
(groups is not yet implemnted)
|
||||
"""
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
setup_oicd_resource_server(responses, settings, sub="user_sub", audience="docs")
|
||||
token = build_authorization_bearer()
|
||||
|
||||
|
||||
@@ -0,0 +1,432 @@
|
||||
"""
|
||||
Test suite for opensearch service
|
||||
"""
|
||||
|
||||
import logging
|
||||
import operator
|
||||
from json import dumps as json_dumps
|
||||
|
||||
import pytest
|
||||
import responses
|
||||
|
||||
from core.services import opensearch
|
||||
|
||||
from ..services.opensearch import (
|
||||
check_hybrid_search_enabled,
|
||||
embed_text,
|
||||
search,
|
||||
)
|
||||
from .mock import albert_embedding_response
|
||||
from .utils import (
|
||||
bulk_create_documents,
|
||||
delete_search_pipeline,
|
||||
enable_hybrid_search,
|
||||
prepare_index,
|
||||
)
|
||||
from .utils import (
|
||||
check_hybrid_search_enabled as check_hybrid_search_enabled_utils,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.django_db
|
||||
|
||||
|
||||
SERVICE_NAME = "test-service"
|
||||
PARAMS = {
|
||||
"nb_results": 20,
|
||||
"order_by": "relevance",
|
||||
"order_direction": "desc",
|
||||
"search_indices": {SERVICE_NAME},
|
||||
"reach": None,
|
||||
"user_sub": "user_sub",
|
||||
"groups": [],
|
||||
"visited": [],
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def before_each():
|
||||
"""Clear caches and delete search pipeline before each test"""
|
||||
clear_caches()
|
||||
yield
|
||||
clear_caches()
|
||||
|
||||
|
||||
def clear_caches():
|
||||
"""Clear caches used in opensearch service and factories"""
|
||||
check_hybrid_search_enabled.cache_clear()
|
||||
# the instance of check_hybrid_search_enabled used in utils.py
|
||||
# is different and must be cleared separately
|
||||
check_hybrid_search_enabled_utils.cache_clear()
|
||||
delete_search_pipeline()
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_hybrid_search_success(settings, caplog):
|
||||
"""Test the hybrid search is successful"""
|
||||
enable_hybrid_search(settings)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "wolf", "content": "wolves live in packs and hunt together"},
|
||||
{"title": "dog", "content": "dogs are loyal domestic animals"},
|
||||
{"title": "cat", "content": "cats are curious and independent pets"},
|
||||
]
|
||||
)
|
||||
prepare_index(SERVICE_NAME, documents)
|
||||
|
||||
q = "canine pet"
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = search(q=q, **PARAMS)
|
||||
|
||||
assert any(
|
||||
f"Performing hybrid search with embedding: {q}" in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
|
||||
assert result["hits"]["max_score"] > 0.0
|
||||
# hybrid search always returns a response of fixed sized sorted and scored by relevance
|
||||
assert {hit["_source"]["title"] for hit in result["hits"]["hits"]} == {
|
||||
doc["title"] for doc in documents
|
||||
}
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_hybrid_search_without_embedded_index(settings, caplog):
|
||||
"""Test the hybrid search is successful"""
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "wolf", "content": "wolves live in packs and hunt together"},
|
||||
{"title": "dog", "content": "dogs are loyal domestic animals"},
|
||||
{"title": "cat", "content": "cats are curious and independent pets"},
|
||||
]
|
||||
)
|
||||
# index is prepared but hybrid search is not yet enable.
|
||||
# they then won't be embedded.
|
||||
prepare_index(SERVICE_NAME, documents)
|
||||
|
||||
# check embedding is None
|
||||
indexed_documents = opensearch.opensearch_client().search(
|
||||
index=SERVICE_NAME, body={"query": {"match_all": {}}}
|
||||
)
|
||||
assert indexed_documents["hits"]["hits"][0]["_source"]["embedding"] is None
|
||||
|
||||
# hybrid search is enabled before to do the first requests
|
||||
enable_hybrid_search(settings)
|
||||
|
||||
q = "canine pet"
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = search(q=q, **PARAMS)
|
||||
|
||||
# the hybrid search is done successfully
|
||||
assert any(
|
||||
f"Performing hybrid search with embedding: {q}" in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
|
||||
# but no match can obviously be found
|
||||
assert result["hits"]["max_score"] == 0.0
|
||||
assert len(result["hits"]["hits"]) == 0
|
||||
|
||||
# The full-text search is still functional
|
||||
q = "wolf"
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = search(q=q, **PARAMS)
|
||||
|
||||
assert any(
|
||||
f"Performing hybrid search with embedding: {q}" in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
|
||||
assert result["hits"]["max_score"] > 0.0
|
||||
assert len(result["hits"]["hits"]) == 1
|
||||
assert result["hits"]["hits"][0]["_source"]["title"] == q
|
||||
|
||||
|
||||
def test_fall_back_on_full_text_search_if_hybrid_search_disabled(settings, caplog):
|
||||
"""Test the full-text search is done when HYBRID_SEARCH_ENABLED=False"""
|
||||
enable_hybrid_search(settings)
|
||||
settings.HYBRID_SEARCH_ENABLED = False
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "wolf", "content": "wolves live in packs and hunt together"},
|
||||
{"title": "dog", "content": "dogs are loyal domestic animals"},
|
||||
{"title": "cat", "content": "cats are curious and independent pets"},
|
||||
]
|
||||
)
|
||||
prepare_index(SERVICE_NAME, documents)
|
||||
|
||||
q = "wolf"
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = search(q=q, **PARAMS)
|
||||
|
||||
assert any(
|
||||
"Hybrid search is disabled via HYBRID_SEARCH_ENABLED setting" in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
assert any(
|
||||
f"Performing full-text search without embedding: {q}" in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
|
||||
assert result["hits"]["max_score"] > 0.0
|
||||
assert len(result["hits"]["hits"]) == 1
|
||||
assert result["hits"]["hits"][0]["_source"]["title"] == "wolf"
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_fall_back_on_full_text_search_if_embedding_api_fails(settings, caplog):
|
||||
"""Test the full-text search is done when the embedding api fails"""
|
||||
enable_hybrid_search(settings)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
status=401,
|
||||
body=json_dumps({"message": "Authentication failed."}),
|
||||
)
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "wolf", "content": "wolves live in packs and hunt together"},
|
||||
{"title": "dog", "content": "dogs are loyal domestic animals"},
|
||||
{"title": "cat", "content": "cats are curious and independent pets"},
|
||||
]
|
||||
)
|
||||
prepare_index(SERVICE_NAME, documents)
|
||||
|
||||
q = "wolf"
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = search(q=q, **PARAMS)
|
||||
|
||||
assert any(
|
||||
"embedding API request failed: 401 Client Error: Unauthorized" in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
assert any(
|
||||
f"Performing full-text search without embedding: {q}" in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
assert result["hits"]["max_score"] > 0.0
|
||||
assert len(result["hits"]["hits"]) == 1
|
||||
assert result["hits"]["hits"][0]["_source"]["title"] == "wolf"
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_fall_back_on_full_text_search_if_variable_are_missing(settings, caplog):
|
||||
"""Test the full-text search is done when variables are missing for hybrid search"""
|
||||
enable_hybrid_search(settings)
|
||||
del settings.HYBRID_SEARCH_WEIGHTS
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "wolf", "content": "wolves live in packs and hunt together"},
|
||||
{"title": "dog", "content": "dogs are loyal domestic animals"},
|
||||
{"title": "cat", "content": "cats are curious and independent pets"},
|
||||
]
|
||||
)
|
||||
prepare_index(SERVICE_NAME, documents)
|
||||
|
||||
q = "wolf"
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = search(q=q, **PARAMS)
|
||||
|
||||
assert any(
|
||||
"Missing variables for hybrid search: HYBRID_SEARCH_WEIGHTS" in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
assert any(
|
||||
f"Performing full-text search without embedding: {q}" in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
assert result["hits"]["max_score"] > 0.0
|
||||
assert len(result["hits"]["hits"]) == 1
|
||||
assert result["hits"]["hits"][0]["_source"]["title"] == "wolf"
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_match_all(settings, caplog):
|
||||
"""Test match all when q='*' and no semantic search is needed"""
|
||||
enable_hybrid_search(settings)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "wolf", "content": "wolves live in packs and hunt together"},
|
||||
{"title": "dog", "content": "dogs are loyal domestic animals"},
|
||||
{"title": "cat", "content": "cats are curious and independent pets"},
|
||||
]
|
||||
)
|
||||
prepare_index(SERVICE_NAME, documents)
|
||||
|
||||
q = "*"
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = search(q=q, **PARAMS)
|
||||
|
||||
assert any("Performing match_all query" in message for message in caplog.messages)
|
||||
assert result["hits"]["max_score"] > 0.0
|
||||
assert len(result["hits"]["hits"]) == 3
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_search_ordering_by_relevance(settings, caplog):
|
||||
"""Test the hybrid supports ordering by relevance asc and desc"""
|
||||
enable_hybrid_search(settings)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "wolf", "content": "wolves live in packs and hunt together"},
|
||||
{"title": "dog", "content": "dogs are loyal domestic animals"},
|
||||
{"title": "cat", "content": "cats are curious and independent pets"},
|
||||
]
|
||||
)
|
||||
q = "canine pet"
|
||||
prepare_index(SERVICE_NAME, documents)
|
||||
|
||||
for direction in ["asc", "desc"]:
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = search(q=q, **{**PARAMS, "order_direction": direction})
|
||||
|
||||
# Check that results are sorted by score as expected
|
||||
hits = result["hits"]["hits"]
|
||||
compare = operator.le if direction == "asc" else operator.ge
|
||||
for i in range(len(hits) - 1):
|
||||
assert compare(hits[i]["_score"], hits[i + 1]["_score"])
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_hybrid_search_number_of_matches(settings):
|
||||
"""
|
||||
In this test full-text search always return 0 documents.
|
||||
The test checks the number of hits returned by hybrid search with different k values.
|
||||
"""
|
||||
enable_hybrid_search(settings)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
|
||||
documents = bulk_create_documents(
|
||||
[
|
||||
{"title": "wolf", "content": "wolves live in packs and hunt together"},
|
||||
{"title": "dog", "content": "dogs are loyal domestic animals"},
|
||||
{"title": "cat", "content": "cats are curious and independent pets"},
|
||||
]
|
||||
)
|
||||
prepare_index(SERVICE_NAME, documents)
|
||||
|
||||
q = "pony" # full-text matches 0 document
|
||||
for nb_results in [1, 2, 3]: # semantic should match k documents
|
||||
result = search(q=q, **{**PARAMS, "nb_results": nb_results})
|
||||
assert len(result["hits"]["hits"]) == nb_results
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_embed_text_success(settings):
|
||||
"""Test embed_text retrieval is successful"""
|
||||
enable_hybrid_search(settings)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json=albert_embedding_response.response,
|
||||
status=200,
|
||||
)
|
||||
text = "canine pet"
|
||||
|
||||
embedding = embed_text(text)
|
||||
|
||||
assert embedding == albert_embedding_response.response["data"][0]["embedding"]
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_embed_401_http_error(settings, caplog):
|
||||
"""Test embed_text does not crash and returns None on 401 error"""
|
||||
enable_hybrid_search(settings)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
status=401,
|
||||
body=json_dumps({"message": "Authentication failed."}),
|
||||
)
|
||||
text = "canine pet"
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
embedding = embed_text(text)
|
||||
|
||||
assert any(
|
||||
"embedding API request failed: 401 Client Error: Unauthorized" in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
|
||||
assert embedding is None
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_embed_500_http_error(settings, caplog):
|
||||
"""Test embed_text does not crash and returns None on 500 error"""
|
||||
enable_hybrid_search(settings)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
status=500,
|
||||
body=json_dumps({"message": "Internal server error."}),
|
||||
)
|
||||
text = "canine pet"
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
embedding = embed_text(text)
|
||||
|
||||
assert any(
|
||||
"embedding API request failed: 500 Server Error: Internal Server Error"
|
||||
in message
|
||||
for message in caplog.messages
|
||||
)
|
||||
|
||||
assert embedding is None
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_embed_wrong_format(settings, caplog):
|
||||
"""Test embed_text does not crash and returns None if api returns a wrong format"""
|
||||
enable_hybrid_search(settings)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
settings.EMBEDDING_API_PATH,
|
||||
json={"wrong": "format"},
|
||||
status=200,
|
||||
)
|
||||
text = "canine pet"
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
embedding = embed_text(text)
|
||||
|
||||
assert any(
|
||||
"unexpected embedding response format" in message for message in caplog.messages
|
||||
)
|
||||
|
||||
assert embedding is None
|
||||
@@ -2,23 +2,65 @@
|
||||
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
from functools import partial
|
||||
from typing import List
|
||||
|
||||
from django.conf import settings as django_settings
|
||||
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
from joserfc import jwe as jose_jwe
|
||||
from joserfc import jwt as jose_jwt
|
||||
from joserfc.jwk import RSAKey
|
||||
from jwt.utils import to_base64url_uint
|
||||
from opensearchpy.exceptions import NotFoundError
|
||||
from opensearchpy.helpers import bulk
|
||||
|
||||
from core import opensearch
|
||||
from core import factories
|
||||
from core.management.commands.create_search_pipeline import (
|
||||
ensure_search_pipeline_exists,
|
||||
)
|
||||
from core.services import opensearch
|
||||
from core.services.opensearch import check_hybrid_search_enabled
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def enable_hybrid_search(settings):
|
||||
"""Enable hybrid search settings for tests."""
|
||||
settings.HYBRID_SEARCH_ENABLED = True
|
||||
settings.HYBRID_SEARCH_WEIGHTS = [0.3, 0.7]
|
||||
settings.EMBEDDING_API_KEY = "test-api-key"
|
||||
settings.EMBEDDING_API_PATH = "https://test.embedding.api/v1/embeddings"
|
||||
settings.EMBEDDING_REQUEST_TIMEOUT = 10
|
||||
settings.EMBEDDING_API_MODEL_NAME = "embeddings-small"
|
||||
settings.EMBEDDING_DIMENSION = 1024
|
||||
ensure_search_pipeline_exists()
|
||||
|
||||
|
||||
def bulk_create_documents(document_payloads):
|
||||
"""Create documents in bulk from payloads"""
|
||||
return [
|
||||
factories.DocumentSchemaFactory.build(**document_payload, users=["user_sub"])
|
||||
for document_payload in document_payloads
|
||||
]
|
||||
|
||||
|
||||
def delete_search_pipeline():
|
||||
"""Delete the hybrid search pipeline if it exists"""
|
||||
try:
|
||||
opensearch.opensearch_client().transport.perform_request(
|
||||
method="DELETE",
|
||||
url=f"/_search/pipeline/{django_settings.HYBRID_SEARCH_PIPELINE_ID}",
|
||||
)
|
||||
except NotFoundError:
|
||||
logger.info("Search pipeline not found, nothing to delete.")
|
||||
|
||||
|
||||
def delete_test_indices():
|
||||
"""Drop all search index containing the 'test' word"""
|
||||
opensearch.client.indices.delete(index="*test*")
|
||||
opensearch.opensearch_client().indices.delete(index="*test*")
|
||||
|
||||
|
||||
def prepare_index(index_name, documents: List, cleanup=True):
|
||||
@@ -33,17 +75,27 @@ def prepare_index(index_name, documents: List, cleanup=True):
|
||||
{
|
||||
"_op_type": "index",
|
||||
"_index": index_name,
|
||||
"_id": doc["id"],
|
||||
"_source": {k: v for k, v in doc.items() if k != "id"},
|
||||
"_id": document["id"],
|
||||
"_source": {
|
||||
**{k: v for k, v in document.items() if k != "id"},
|
||||
"embedding": opensearch.embed_text(
|
||||
opensearch.format_document(document["title"], document["content"])
|
||||
)
|
||||
if check_hybrid_search_enabled()
|
||||
else None,
|
||||
"embedding_model": django_settings.EMBEDDING_API_MODEL_NAME
|
||||
if check_hybrid_search_enabled()
|
||||
else None,
|
||||
},
|
||||
}
|
||||
for doc in documents
|
||||
for document in documents
|
||||
]
|
||||
bulk(opensearch.client, actions)
|
||||
bulk(opensearch.opensearch_client(), actions)
|
||||
|
||||
# Force refresh again so all changes are visible to search
|
||||
opensearch.client.indices.refresh(index=index_name)
|
||||
opensearch.opensearch_client().indices.refresh(index=index_name)
|
||||
|
||||
count = opensearch.client.count(index=index_name)["count"]
|
||||
count = opensearch.opensearch_client().count(index=index_name)["count"]
|
||||
assert count == len(documents), f"Expected {len(documents)}, got {count}"
|
||||
|
||||
|
||||
|
||||
+44
-95
@@ -2,6 +2,7 @@
|
||||
|
||||
import logging
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.exceptions import SuspiciousOperation
|
||||
|
||||
from lasuite.oidc_resource_server.authentication import ResourceServerAuthentication
|
||||
@@ -10,11 +11,17 @@ from pydantic import ValidationError as PydanticValidationError
|
||||
from rest_framework import status, views
|
||||
from rest_framework.response import Response
|
||||
|
||||
from . import enums, schemas
|
||||
from . import schemas
|
||||
from .authentication import ServiceTokenAuthentication
|
||||
from .models import Service
|
||||
from .opensearch import client, ensure_index_exists
|
||||
from .permissions import IsAuthAuthenticated
|
||||
from .services.opensearch import (
|
||||
check_hybrid_search_enabled,
|
||||
embed_document,
|
||||
ensure_index_exists,
|
||||
opensearch_client,
|
||||
search,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -80,6 +87,7 @@ class IndexDocumentView(views.APIView):
|
||||
errors.
|
||||
"""
|
||||
index_name = request.auth.name
|
||||
opensearch_client_ = opensearch_client()
|
||||
|
||||
if isinstance(request.data, list):
|
||||
# Bulk indexing several documents
|
||||
@@ -98,7 +106,15 @@ class IndexDocumentView(views.APIView):
|
||||
results.append({"index": i, "status": "error", "errors": errors})
|
||||
has_errors = True
|
||||
else:
|
||||
document_dict = document.model_dump()
|
||||
document_dict = {
|
||||
**document.model_dump(),
|
||||
"embedding": embed_document(document)
|
||||
if check_hybrid_search_enabled()
|
||||
else None,
|
||||
"embedding_model": settings.EMBEDDING_API_MODEL_NAME
|
||||
if check_hybrid_search_enabled()
|
||||
else None,
|
||||
}
|
||||
_id = document_dict.pop("id")
|
||||
actions.append({"index": {"_id": _id}})
|
||||
actions.append(document_dict)
|
||||
@@ -110,7 +126,7 @@ class IndexDocumentView(views.APIView):
|
||||
# Build index if needed.
|
||||
ensure_index_exists(index_name)
|
||||
|
||||
response = client.bulk(index=index_name, body=actions)
|
||||
response = opensearch_client_.bulk(index=index_name, body=actions)
|
||||
for i, item in enumerate(response["items"]):
|
||||
if item["index"]["status"] != 201:
|
||||
results[i]["status"] = "error"
|
||||
@@ -124,13 +140,21 @@ class IndexDocumentView(views.APIView):
|
||||
|
||||
# Indexing a single document
|
||||
document = schemas.DocumentSchema(**request.data)
|
||||
document_dict = document.model_dump()
|
||||
document_dict = {
|
||||
**document.model_dump(),
|
||||
"embedding": embed_document(document)
|
||||
if check_hybrid_search_enabled()
|
||||
else None,
|
||||
"embedding_model": settings.EMBEDDING_API_MODEL_NAME
|
||||
if check_hybrid_search_enabled()
|
||||
else None,
|
||||
}
|
||||
_id = document_dict.pop("id")
|
||||
|
||||
# Build index if needed.
|
||||
ensure_index_exists(index_name)
|
||||
|
||||
client.index(index=index_name, body=document_dict, id=_id)
|
||||
opensearch_client_.index(index=index_name, body=document_dict, id=_id)
|
||||
|
||||
return Response(
|
||||
{"status": "created", "_id": _id}, status=status.HTTP_201_CREATED
|
||||
@@ -148,7 +172,8 @@ class SearchDocumentView(ResourceServerMixin, views.APIView):
|
||||
authentication_classes = [ResourceServerAuthentication]
|
||||
permission_classes = [IsAuthAuthenticated]
|
||||
|
||||
def _get_opensearch_indices(self, audience, services):
|
||||
@staticmethod
|
||||
def _get_opensearch_indices(audience, services):
|
||||
# Get request user service
|
||||
try:
|
||||
user_service = Service.objects.get(client_id=audience, is_active=True)
|
||||
@@ -191,11 +216,8 @@ class SearchDocumentView(ResourceServerMixin, views.APIView):
|
||||
order_direction : str, optional
|
||||
Order direction, 'asc' for ascending or 'desc' for descending.
|
||||
Defaults to 'desc'.
|
||||
page_number : int, optional
|
||||
The page number to retrieve.
|
||||
Defaults to 1 if not specified.
|
||||
page_size : int, optional
|
||||
The number of results to return per page.
|
||||
nb_results : int, optional
|
||||
The number of results to return.
|
||||
Defaults to 50 if not specified.
|
||||
services: List[str], optional
|
||||
List of services on which we intend to run the query (current service if left empty)
|
||||
@@ -219,10 +241,6 @@ class SearchDocumentView(ResourceServerMixin, views.APIView):
|
||||
# Extract and validate query parameters using Pydantic schema
|
||||
params = schemas.SearchQueryParametersSchema(**request.data)
|
||||
|
||||
# Compute pagination parameters
|
||||
from_value = (params.page_number - 1) * params.page_size
|
||||
size_value = params.page_size
|
||||
|
||||
# Get index list for search query
|
||||
try:
|
||||
search_indices = self._get_opensearch_indices(
|
||||
@@ -231,85 +249,16 @@ class SearchDocumentView(ResourceServerMixin, views.APIView):
|
||||
except SuspiciousOperation as e:
|
||||
return Response({"detail": str(e)}, status=status.HTTP_400_BAD_REQUEST)
|
||||
|
||||
# Prepare the search query
|
||||
search_body = {
|
||||
"_source": enums.SOURCE_FIELDS, # limit the fields to return
|
||||
"script_fields": {
|
||||
"number_of_users": {"script": {"source": "doc['users'].size()"}},
|
||||
"number_of_groups": {"script": {"source": "doc['groups'].size()"}},
|
||||
},
|
||||
"query": {"bool": {"must": [], "filter": []}},
|
||||
"sort": [],
|
||||
"from": from_value,
|
||||
"size": size_value,
|
||||
}
|
||||
|
||||
# Adding the text query
|
||||
if params.q == "*":
|
||||
search_body["query"]["bool"]["must"].append({"match_all": {}})
|
||||
else:
|
||||
search_body["query"]["bool"]["must"].append(
|
||||
{
|
||||
"multi_match": {
|
||||
"query": params.q,
|
||||
# Give title more importance over content by a power of 3
|
||||
"fields": ["title.text^3", "content"],
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
# Add sorting logic based on relevance or specified field
|
||||
if params.order_by == enums.RELEVANCE:
|
||||
search_body["sort"].append({"_score": {"order": params.order_direction}})
|
||||
else:
|
||||
search_body["sort"].append(
|
||||
{params.order_by: {"order": params.order_direction}}
|
||||
)
|
||||
|
||||
# Apply access control based on documents reach
|
||||
search_body["query"]["bool"]["must"].append(
|
||||
{
|
||||
"bool": {
|
||||
"should": [
|
||||
# Access control on public & authenticated reach
|
||||
{
|
||||
"bool": {
|
||||
"must_not": {
|
||||
"term": {enums.REACH: enums.ReachEnum.RESTRICTED},
|
||||
},
|
||||
# Limit search to already visited documents.
|
||||
"must": {
|
||||
"terms": {
|
||||
"_id": sorted(params.visited),
|
||||
}
|
||||
},
|
||||
},
|
||||
},
|
||||
# Access control on restricted search : either user or group should match
|
||||
{"term": {enums.USERS: user_sub}},
|
||||
{"terms": {enums.GROUPS: groups}},
|
||||
],
|
||||
# At least one of the 2 optional should clauses must apply
|
||||
"minimum_should_match": 1,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
# Optional filter by reach if explicitly provided in the query
|
||||
if params.reach is not None:
|
||||
search_body["query"]["bool"]["filter"].append(
|
||||
{"term": {enums.REACH: params.reach}}
|
||||
)
|
||||
|
||||
# Always filter out inactive documents
|
||||
search_body["query"]["bool"]["filter"].append({"term": {"is_active": True}})
|
||||
|
||||
response = client.search( # pylint: disable=unexpected-keyword-arg
|
||||
index=",".join(search_indices),
|
||||
body=search_body,
|
||||
# Argument added by the query_params() decorator of opensearch and
|
||||
# not in the method declaration.
|
||||
ignore_unavailable=True,
|
||||
response = search(
|
||||
q=params.q,
|
||||
nb_results=params.nb_results,
|
||||
order_by=params.order_by,
|
||||
order_direction=params.order_direction,
|
||||
search_indices=search_indices,
|
||||
reach=params.reach,
|
||||
visited=params.visited,
|
||||
user_sub=user_sub,
|
||||
groups=groups,
|
||||
)
|
||||
|
||||
return Response(response["hits"]["hits"], status=status.HTTP_200_OK)
|
||||
|
||||
@@ -14,7 +14,8 @@ from django.utils.text import slugify
|
||||
from faker import Faker
|
||||
from opensearchpy.helpers import bulk
|
||||
|
||||
from core import enums, factories, opensearch
|
||||
from core import enums, factories
|
||||
from core.services.opensearch import ensure_index_exists, opensearch_client
|
||||
|
||||
from demo import defaults
|
||||
|
||||
@@ -36,7 +37,7 @@ class BulkIndexing:
|
||||
|
||||
def bulk_index(self):
|
||||
"""Actually index documents in bulk to OpenSearch."""
|
||||
_success, failed = bulk(opensearch.client, self.actions, stats_only=False)
|
||||
_success, failed = bulk(opensearch_client(), self.actions, stats_only=False)
|
||||
|
||||
if failed:
|
||||
self.handle_failures(failed)
|
||||
@@ -141,7 +142,8 @@ def create_demo(stdout):
|
||||
"""
|
||||
Create a database with demo data for developers to work in a realistic environment.
|
||||
"""
|
||||
opensearch.client.indices.delete("*")
|
||||
opensearch_client_ = opensearch_client()
|
||||
opensearch_client_.indices.delete("*")
|
||||
|
||||
with Timeit(stdout, "Creating services"):
|
||||
services = factories.ServiceFactory.create_batch(
|
||||
@@ -149,8 +151,8 @@ def create_demo(stdout):
|
||||
)
|
||||
|
||||
for service in services:
|
||||
opensearch.ensure_index_exists(service.name)
|
||||
opensearch.client.indices.refresh(index=service.name)
|
||||
ensure_index_exists(service.name)
|
||||
opensearch_client_.indices.refresh(index=service.name)
|
||||
|
||||
with Timeit(stdout, "Creating documents"):
|
||||
actions = BulkIndexing(stdout)
|
||||
@@ -163,14 +165,14 @@ def create_demo(stdout):
|
||||
with Timeit(stdout, "Creating dev services"):
|
||||
for conf in defaults.DEV_SERVICES:
|
||||
service = factories.ServiceFactory(**conf)
|
||||
opensearch.ensure_index_exists(service.name)
|
||||
opensearch.client.indices.refresh(index=service.name)
|
||||
ensure_index_exists(service.name)
|
||||
opensearch_client_.indices.refresh(index=service.name)
|
||||
|
||||
# Check and report on indexed documents
|
||||
total_indexed = 0
|
||||
for service in services:
|
||||
opensearch.client.indices.refresh(index=service.name)
|
||||
indexed = opensearch.client.count(index=service.name)["count"]
|
||||
opensearch_client_.indices.refresh(index=service.name)
|
||||
indexed = opensearch_client_.count(index=service.name)["count"]
|
||||
stdout.write(f" - {service.name:s}: {indexed:d} documents")
|
||||
total_indexed += indexed
|
||||
|
||||
|
||||
@@ -7,7 +7,8 @@ from django.test import override_settings
|
||||
|
||||
import pytest
|
||||
|
||||
from core import models, opensearch
|
||||
from core import models
|
||||
from core.services.opensearch import opensearch_client
|
||||
|
||||
from demo import defaults
|
||||
|
||||
@@ -26,7 +27,7 @@ def test_commands_create_demo():
|
||||
call_command("create_demo")
|
||||
|
||||
assert models.Service.objects.exclude(name="docs").count() == 2
|
||||
assert opensearch.client.count()["count"] == 4
|
||||
assert opensearch_client().count()["count"] == 4
|
||||
|
||||
docs = models.Service.objects.get(name="docs")
|
||||
assert docs.client_id == "impress"
|
||||
|
||||
@@ -262,6 +262,35 @@ class Base(Configuration):
|
||||
|
||||
AUTH_USER_MODEL = "core.User"
|
||||
|
||||
# Hybrid Search settings
|
||||
HYBRID_SEARCH_ENABLED = values.BooleanValue(
|
||||
default=False, environ_name="HYBRID_SEARCH_ENABLED", environ_prefix=None
|
||||
)
|
||||
HYBRID_SEARCH_PIPELINE_ID = "hybrid-search-pipeline"
|
||||
HYBRID_SEARCH_WEIGHTS = values.ListValue(
|
||||
default=[0.3, 0.7], environ_name="HYBRID_SEARCH_WEIGHTS", environ_prefix=None
|
||||
)
|
||||
EMBEDDING_API_PATH = values.Value(
|
||||
# embedding is the vector representation of a document used for semantic search
|
||||
default="None",
|
||||
environment_name="EMBEDDING_API_PATH",
|
||||
environ_prefix=None,
|
||||
)
|
||||
EMBEDDING_API_KEY = values.Value(
|
||||
default=None, environ_name="EMBEDDING_API_KEY", environ_prefix=None
|
||||
)
|
||||
EMBEDDING_REQUEST_TIMEOUT = values.Value(
|
||||
default=10, environ_name="EMBEDDING_REQUEST_TIMEOUT", environ_prefix=None
|
||||
)
|
||||
EMBEDDING_API_MODEL_NAME = values.Value(
|
||||
default="embeddings-small",
|
||||
environ_name="EMBEDDING_API_MODEL_NAME",
|
||||
environ_prefix=None,
|
||||
)
|
||||
EMBEDDING_DIMENSION = values.IntegerValue(
|
||||
default=1024, environ_name="EMBEDDING_DIMENSION", environ_prefix=None
|
||||
)
|
||||
|
||||
# CORS
|
||||
CORS_ALLOW_CREDENTIALS = True
|
||||
CORS_ALLOW_ALL_ORIGINS = values.BooleanValue(True)
|
||||
|
||||
@@ -89,6 +89,7 @@ exclude = [
|
||||
"venv",
|
||||
"__pycache__",
|
||||
"*/migrations/*",
|
||||
".vscode*"
|
||||
]
|
||||
line-length = 88
|
||||
|
||||
|
||||
Reference in New Issue
Block a user