diff --git a/utils/mongo_utils.py b/utils/mongo_utils.py index 39edb9f..c207801 100644 --- a/utils/mongo_utils.py +++ b/utils/mongo_utils.py @@ -1,6 +1,6 @@ from pymongo import MongoClient - -from utils.config import DB_NAME, MONGO_URI, TENDERS_COLLECTION, VECTOR_COLLECTION +from bson.objectid import ObjectId +from utils.config import MONGO_URI, DB_NAME, VECTOR_COLLECTION, TENDERS_COLLECTION mongo = MongoClient(MONGO_URI) db = mongo[DB_NAME] @@ -8,6 +8,7 @@ vector_collection = db[VECTOR_COLLECTION] tenders_collection = db[TENDERS_COLLECTION] +ALLOWED_INDUSTRIES = ["Water & Sanitation", "Power & Energy"] def store_embeddings_in_db(embeddings, document_name, tender_id): try: @@ -15,7 +16,12 @@ def store_embeddings_in_db(embeddings, document_name, tender_id): except Exception as e: print(f"❌ Mongo Insert Error: {e}") - def get_tender_ids(min_value): - cursor = tenders_collection.find({"tender_value": {"$gte": min_value}}, {"_id": 1}) + cursor = tenders_collection.find( + { + "tender_value": {"$gte": min_value}, + "industries": {"$in": ALLOWED_INDUSTRIES} # filter by industries + }, + {"_id": 1} + ) return [str(doc["_id"]) for doc in cursor]