Skip to content

Commit

Permalink
feat(Besoins): Filtre par type de structure dans le ciblage des struc…
Browse files Browse the repository at this point in the history
…tures par recherche sémantique (#1152)
  • Loading branch information
SebastienReuiller authored Apr 8, 2024
1 parent fe2ae2d commit d3a0e1a
Show file tree
Hide file tree
Showing 3 changed files with 12 additions and 6 deletions.
1 change: 1 addition & 0 deletions lemarche/siaes/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -1186,6 +1186,7 @@ def elasticsearch_index_metadata(self):
"id": self.id,
"name": self.name,
"website": self.website if self.website else "",
"kind": self.kind,
}
if self.latitude and self.longitude:
metadata["geo_location"] = {
Expand Down
3 changes: 2 additions & 1 deletion lemarche/tenders/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -719,9 +719,10 @@ def set_siae_found_list(self):
geo_distance=self.distance_location,
geo_lat=self.location.coords.y,
geo_lon=self.location.coords.x,
siae_kinds=self.siae_kind,
)
else:
siae_ids = api_elasticsearch.siaes_similarity_search(self.description)
siae_ids = api_elasticsearch.siaes_similarity_search(self.description, siae_kinds=self.siae_kind)

siaes_had_found_by_ia = Siae.objects.filter(id__in=siae_ids)
siaes_had_found_by_ia_too = []
Expand Down
14 changes: 9 additions & 5 deletions lemarche/utils/apis/api_elasticsearch.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
)


def siaes_similarity_search(search_text: str, search_filter: dict = {}):
def siaes_similarity_search(search_text: str, search_filter: list = [], siae_kinds: list = []):
"""Performs semantic search with Elasticsearch as a vector db
Args:
Expand All @@ -21,6 +21,10 @@ def siaes_similarity_search(search_text: str, search_filter: dict = {}):
Returns:
list: list of siaes id that match the search query
"""

if siae_kinds:
search_filter.append({"terms": {"metadata.kind.keyword": siae_kinds}})

db = ElasticsearchStore(
embedding=OpenAIEmbeddings(),
es_user=settings.ELASTICSEARCH_USERNAME,
Expand All @@ -40,7 +44,7 @@ def siaes_similarity_search(search_text: str, search_filter: dict = {}):


def siaes_similarity_search_with_geo_distance(
search_text: str, geo_distance: int = None, geo_lat: float = None, geo_lon: float = None
search_text: str, geo_distance: int = None, geo_lat: float = None, geo_lon: float = None, siae_kinds: list = []
):
search_filter = []
if geo_distance and geo_lat and geo_lon:
Expand All @@ -56,10 +60,10 @@ def siaes_similarity_search_with_geo_distance(
}
]

return siaes_similarity_search(search_text, search_filter)
return siaes_similarity_search(search_text, search_filter, siae_kinds)


def siaes_similarity_search_with_city(search_text: str, city: Perimeter):
def siaes_similarity_search_with_city(search_text: str, city: Perimeter, siae_kinds: list = []):
search_filter = [
{
"bool": {
Expand Down Expand Up @@ -88,4 +92,4 @@ def siaes_similarity_search_with_city(search_text: str, city: Perimeter):
}
}
]
return siaes_similarity_search(search_text, search_filter)
return siaes_similarity_search(search_text, search_filter, siae_kinds)

0 comments on commit d3a0e1a

Please sign in to comment.