fix author matching and handling of filtering

Author: WonderPG

pyproject.toml

@@ -16,7 +16,6 @@ dependencies = [
     "boto3",
     "cohere",
     "elasticsearch >= 8.5",
-    "elasticsearch-dsl",
     "fastapi <= 0.112.0",
     "fastapi-pagination",
     "httpx",

src/scholarag/app/dependencies.py

@@ -4,9 +4,8 @@
 import logging
 from enum import Enum
 from functools import cache
-from typing import Annotated, AsyncIterator
+from typing import Annotated, Any, AsyncIterator
 
-from elasticsearch_dsl import Q, Search
 from fastapi import Depends, HTTPException, Query, Request
 from fastapi.security import HTTPBearer
 from httpx import AsyncClient, HTTPStatusError
@@ -202,24 +201,22 @@ def get_query_from_params(
             pattern=r"^\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\d|3[01])$",
         ),
     ] = None,
-) -> dict[str, str] | None:
+) -> dict[str, dict[str, list[dict[str, Any]]]] | None:
     """Get the query parameters and generate an ES query for filtering."""
-    search = Search()
+    query: dict[str, dict[str, list[dict[str, Any]]]] = {"bool": {"must": []}}
     if article_types:
-        search = search.query(Q("terms", article_type=article_types))
+        query["bool"]["must"].append({"terms": {"article_type": article_types}})
     if authors:
-        search = search.query(Q("terms", authors=authors))
+        query["bool"]["must"].append({"terms": {"authors.keyword": authors}})
     if journals:
-        search = search.query(Q("terms", journal=journals))
+        query["bool"]["must"].append({"terms": {"journal": journals}})
     if date_from:
-        search = search.query(Q("range", date={"gte": date_from}))
+        query["bool"]["must"].append({"range": {"date": {"gte": date_from}}})
     if date_to:
-        search = search.query(Q("range", date={"lte": date_to}))
+        query["bool"]["must"].append({"range": {"date": {"lte": date_to}}})
 
-    logger.info(
-        f"Searching the database with the query {json.dumps(search.to_dict())}."
-    )
-    return None if not search.to_dict() else search.to_dict()["query"]
+    logger.info(f"Searching the database with the query {json.dumps(query)}.")
+    return None if not query["bool"]["must"] else query
 
 
 class ErrorCode(Enum):

src/scholarag/app/routers/retrieval.py

@@ -250,7 +250,7 @@ async def article_count(
 
     # If further filters, append them
     if filter_query:
-        query["query"]["bool"]["must"].append(filter_query)
+        query["query"]["bool"]["must"].extend(filter_query["bool"]["must"])
 
     # Aggregation query.
     aggs = {
@@ -390,7 +390,7 @@ async def article_listing(
 
     # If further filters, append them
     if filter_query:
-        query["query"]["bool"]["must"].append(filter_query)
+        query["query"]["bool"]["must"].extend(filter_query["bool"]["must"])
 
     aggs: dict[str, Any] = {
         "relevant_ids": {
@@ -410,8 +410,9 @@ async def article_listing(
         aggs["relevant_ids"]["aggs"]["score"]["max"] = {"field": "date"}
 
     results = await ds_client.search(
-        index=settings.db.index_paragraphs, query=query, size=78, aggs=aggs
+        index=settings.db.index_paragraphs, query=query, size=0, aggs=aggs
     )
+
     logger.info(f"unique article retrieval took: {time.time() - start}s")
 
     docs = [

tests/app/test_dependencies.py

@@ -184,7 +184,7 @@ def test_get_query_from_params():
         "bool": {
             "must": [
                 {"terms": {"article_type": ["publication", "review"]}},
-                {"terms": {"authors": ["Guy Manderson", "Joe Guy"]}},
+                {"terms": {"authors.keyword": ["Guy Manderson", "Joe Guy"]}},
                 {"terms": {"journal": ["1111-1111"]}},
                 {"range": {"date": {"gte": "2020-01-01"}}},
                 {"range": {"date": {"lte": "2020-01-02"}}},

tests/app/test_qa.py

@@ -41,7 +41,7 @@ def test_generative_qa(app_client, mock_http_calls):
         "bool": {
             "must": [
                 {"terms": {"article_type": params["article_types"]}},
-                {"terms": {"authors": params["authors"]}},
+                {"terms": {"authors.keyword": params["authors"]}},
                 {"range": {"date": {"gte": params["date_from"]}}},
                 {"range": {"date": {"lte": params["date_to"]}}},
             ]