Update doc_searcher.py

doc_searcher.py

@@ -1,4 +1,5 @@
 from qdrant_client import QdrantClient
+from qdrant_client.models import Filter, FieldCondition, MatchValue
 from fastembed import SparseTextEmbedding, LateInteractionTextEmbedding
 from qdrant_client import QdrantClient, models
 from sentence_transformers import SentenceTransformer
@@ -13,7 +14,7 @@ class DocSearcher:
         self.late_interaction_model = LateInteractionTextEmbedding(LATE_INTERACTION_MODEL)
         self.qdrant_client = QdrantClient(QDRANT_URL,api_key=QDRANT_API_KEY,timeout=30)
     
-    async def search(self, text: str):
+    async def search(self, text: str,type:int, law_type: str | None = None, offset: int = 0):
         
         dense_query = self.dense_model.encode(text).tolist()
         sparse_query = next(self.sparse_model.query_embed(text))
@@ -22,39 +23,52 @@ class DocSearcher:
                 models.Prefetch(
                     query=dense_query,
                     using=DENSE_MODEL,
-                    params=models.SearchParams(
-                    quantization=models.QuantizationSearchParams(
-                            rescore=False,
-                        ),
-                    ),
-                    limit=200
+                    limit=100
                 ),
                 models.Prefetch(
                     query=models.SparseVector(**sparse_query.as_object()),
                     using=SPARSE_MODEL,
-                    params=models.SearchParams(
-                    quantization=models.QuantizationSearchParams(
-                            rescore=False,
-                        ),
-                    ),
-                    limit=200
+                    limit=100
                 )
         ]
 
+        if type == 2:
+            filter = None
+        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)
+                    ),
+                ],
+                must_not=[
+                    FieldCondition(key="status", match=MatchValue(value="Nevažeći")),
+                ]
+            )
+        else:
+            filter = Filter(
+                must=[
+                    FieldCondition(
+                        key="tip_dokumenta",
+                        match=MatchValue(value=type)
+                    ),
+                ]
+            )
         search_result = self.qdrant_client.query_points(
             collection_name= self.collection_name,
-            search_params=models.SearchParams(
-                hnsw_ef=128,
-                quantization=models.QuantizationSearchParams(
-                        rescore=True,
-                    ),
-                ),
+            query_filter=filter,
             prefetch=prefetch,
             query=models.FusionQuery(
                 fusion=models.Fusion.RRF,
             ),
             with_payload=True,
-            limit = 10
+            limit = 10,
+            offset = offset
         ).points
 
         data = []