Update doc_searcher.py

doc_searcher.py

@@ -19,18 +19,27 @@ class DocSearcher:
         dense_query = self.dense_model.encode(text).tolist()
         sparse_query = next(self.sparse_model.query_embed(text))
 
-        prefetch = [
-                models.Prefetch(
-                    query=dense_query,
-                    using=DENSE_MODEL,
-                    limit=100
-                ),
-                models.Prefetch(
-                    query=models.SparseVector(**sparse_query.as_object()),
-                    using=SPARSE_MODEL,
-                    limit=100
-                )
-        ]
+        if type == 0:
+            prefetch = [
+                    models.Prefetch(
+                        query=dense_query,
+                        using=DENSE_MODEL,
+                        limit=100
+                    ),
+                    models.Prefetch(
+                        query=models.SparseVector(**sparse_query.as_object()),
+                        using=SPARSE_MODEL,
+                        limit=100
+                    )
+            ]
+        else:
+            prefetch = [
+                    models.Prefetch(
+                        query=dense_query,
+                        using=DENSE_MODEL,
+                        limit=100
+                    )
+            ]
 
         if type == 0 and law_type is not None:
             filter = Filter(
@@ -54,6 +63,19 @@ class DocSearcher:
                     )
                 ]
             )
+        elif type == 1 and law_type is not None:
+            filter = Filter(
+                must=[
+                    FieldCondition(
+                        key="tip_dokumenta",
+                        match=MatchValue(value=type)
+                    ),
+                    FieldCondition(
+                        key="vrsta_akta",
+                        match=MatchValue(value=law_type)
+                    ),
+                ]
+            )
         else:
             return []