File size: 1,708 Bytes
500c1ba
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
import logging
from qdrant_client import QdrantClient
from qdrant_client.http.models import Filter, FieldCondition

class QdrantSearcher:
    def __init__(self, encoder, qdrant_url, access_token):
        self.encoder = encoder
        self.client = QdrantClient(url=qdrant_url, api_key=access_token)

    def search_documents(self, collection_name, query, user_id, limit=3):
        logging.info("Starting document search")
        query_vector = self.encoder.encode(query).tolist()
        query_filter = Filter(must=[FieldCondition(key="user_id", match={"value": user_id})])

        try:
            hits = self.client.search(
                collection_name=collection_name,
                query_vector=query_vector,
                limit=limit,
                query_filter=query_filter
            )
        except Exception as e:
            logging.error(f"Error during Qdrant search: {e}")
            return None, str(e)

        if not hits:
            logging.info("No documents found for the given query")
            return None, "No documents found for the given query."

        hits_list = []
        for hit in hits:
            hit_info = {
                "id": hit.id,
                "score": hit.score,
                "file_id": hit.payload.get('file_id'),
                "organization_id": hit.payload.get('organization_id'),
                "chunk_index": hit.payload.get('chunk_index'),
                "chunk_text": hit.payload.get('chunk_text'),
                "s3_bucket_key": hit.payload.get('s3_bucket_key')
            }
            hits_list.append(hit_info)

        logging.info(f"Document search completed with {len(hits_list)} hits")
        return hits_list, None