enable doc search

app/main.py

@@ -80,15 +80,13 @@ async def process_prompt(
             group_res = agrupar_etiquetas_repetidas(prompt_res)
             procesar_entidades(group_res)
             
-            # Código comentado original mantenido
-            # docs = buscar_documentos(tag_res, collection)
-            # documentos_transformados = []
-            # for documento in docs:
-            #     documento["_id"] = str(documento["_id"])  # Convertir ObjectId a string
-            #     documentos_transformados.append(documento)
-            # json_response = {"data": documentos_transformados, "query": tag_res, "model": prompt_res}
+            docs = buscar_documentos(filter, collection)
+            documentos_transformados = []
+            for documento in docs:
+                documento["_id"] = str(documento["_id"])  # Convertir ObjectId a string
+                documentos_transformados.append(documento)
             
-            response_data = {"model": group_res, "filter": filter}
+            response_data = {"data": documentos_transformados, "model": prompt_res}
             return response_data
             
         except Exception as e:
@@ -118,18 +116,29 @@ def agrupar_etiquetas_repetidas(entidades):
             resultado.append((textos[0], etiqueta))
     return resultado
 
-def procesar_entidades(entidades, filter):
+def procesar_entidades(entidades):
+    """
+    Procesa las entidades extraídas, validando y extrayendo información específica.
+
+    Args:
+        entidades (list): Una lista de tuplas (texto, etiqueta).
+
+    Returns:
+        dict: Un diccionario con los datos procesados y validados.
+    """
+
     datos_procesados = {}
     datos_procesados["ubicacion"] = []
     op1 = ["minima", "menor", "mínima"]
-    op2 = ["maxima", "mayor", "máxima"]
+    op2 = [ "maxima", "mayor", "máxima"]
+
     for texto, etiqueta in entidades:
         if etiqueta == "PROPERTY_TYPE":
             if isinstance(texto, tuple):
                 for t in texto:
                     if t in ["apartamento", "casa", "local", "finca", "oficina", "bodega", "lote"]:
                         filter["type"].append(t)
-            else:
+            else:      
                 if texto in ["apartamento", "casa", "local", "finca", "oficina", "bodega", "lote"]:
                     filter["type"].append(texto)
         elif etiqueta == "OPERATION":
@@ -139,11 +148,14 @@ def procesar_entidades(entidades, filter):
                 filter["operation"] = "Venta"
             else:
                 filter["operation"] = "Arriendo"
-        elif etiqueta == "LOCATION":
+        # elif etiqueta == "LOCATION_TYPE":
+        #     if texto in ["barrio", "vereda", "centro", "zona"]:
+        #         datos_procesados["tipo_ubicacion"] = texto
+        elif etiqueta == "LOCATION":  
             if isinstance(texto, tuple):
                 for t in texto:
                     datos_procesados["ubicacion"].append(t)
-            else:
+            else:           
                 datos_procesados["ubicacion"].append(texto)
         elif etiqueta == "BEDROOMS":
             if isinstance(texto, tuple):
@@ -151,11 +163,12 @@ def procesar_entidades(entidades, filter):
                     numero_habitaciones = re.search(r"\d+", t)
                     if numero_habitaciones:
                         filter["habitaciones"].append(int(numero_habitaciones.group(0)))
-            else:
+                 
+            else:       
                 numero_habitaciones = re.search(r"\d+", texto)
                 if numero_habitaciones:
                     filter["habitaciones"].append(int(numero_habitaciones.group(0)))
-                else:
+                else :
                     numero_habitaciones = re.search(r"^un|una", texto)
                     if numero_habitaciones:
                         filter["habitaciones"].append(1)
@@ -165,14 +178,16 @@ def procesar_entidades(entidades, filter):
                     numero_garajes = re.search(r"\d+", t)
                     if numero_garajes:
                         filter["garajes"].append(int(numero_garajes.group(0)))
+              
             else:
                 numero_garajes = re.search(r"\d+", texto)
                 if numero_garajes:
                     filter["garajes"].append(int(numero_garajes.group(0)))
-                else:
+                else :
                     numero_garajes = re.search(r"^un|una", texto)
                     if numero_garajes:
                         filter["garajes"].append(1)
+             
         elif etiqueta == "AREA_SIZE":
             medida_area = re.search(r"(\d+)\s*?(m2|hectáreas|metros|m)", texto, re.IGNORECASE)
             area = 50
@@ -188,6 +203,7 @@ def procesar_entidades(entidades, filter):
             else:
                 filter["areaMin"] = area - area*0.2
                 filter["areaMax"] = area + area*0.2
+
         elif etiqueta == "COST":
             if isinstance(texto, tuple):
                 for t in texto:
@@ -195,11 +211,14 @@ def procesar_entidades(entidades, filter):
                     valor = 0
                     if valor_costo:
                         valor = int(valor_costo.group(0))
+                        
+                      
                         if filter["min"] < valor:
-                            filter["min"] = filter["max"]
+                            filter["min"] = filter["max"] 
                             filter["max"] = valor
+                           
                         else:
-                            filter["min"] = filter["max"]
+                            filter["min"] = filter["max"] 
                             filter["max"] = valor
             else:
                 valor_costo = re.search(r"\d+", texto)
@@ -223,7 +242,64 @@ def procesar_entidades(entidades, filter):
                 numero_banos = re.search(r"\d+", texto)
                 if numero_banos:
                     filter["banos"].append(int(numero_banos.group(0)))
-                else:
+                else :
                     numero_banos = re.search(r"^un|una", texto)
                     if numero_banos:
-                        filter["banos"].append(1)
+                        filter["banos"].append(1)
+
+
+
+
+def buscar_documentos(datos_procesados, collection):
+    """
+    Busca documentos en MongoDB basados en los datos procesados.
+
+    Args:
+        datos_procesados (dict): Diccionario con los datos procesados.
+        db_nombre (str): Nombre de la base de datos.
+        collection_nombre (str): Nombre de la colección.
+
+    Returns:
+        list: Una lista de documentos que coinciden con la consulta.
+    """
+
+    query = {}
+
+    if "tipo_propiedad" in datos_procesados:
+        query["technicalSheet.property_type_name"] = datos_procesados["tipo_propiedad"]
+    if "operacion" in datos_procesados:
+        query["operation"] = datos_procesados["operacion"]
+    if "ubicacion" in datos_procesados:
+        ubicacion = datos_procesados["ubicacion"]
+        if len(ubicacion) > 0:
+            if len(ubicacion) == 1:
+                query["url_frag"] = {"$regex": ubicacion[0], "$options": "i"}
+            else:
+                regex_pattern = "|".join(datos_procesados["ubicacion"])
+                query["url_frag"] = {"$regex": regex_pattern, "$options": "i"}
+    if "habitaciones" in datos_procesados:
+        query["technicalSheet.bedrooms"] = datos_procesados["habitaciones"]
+    if "garajes" in datos_procesados:
+        query["technicalSheet.garage"] = datos_procesados["garajes"]
+    if "area" in datos_procesados:
+        if datos_procesados["operacion_area"] is not None:
+            if datos_procesados["operacion_area"] == "minima":
+                query["area"] = {"$gte": datos_procesados["area"]["valor"]}
+            elif datos_procesados["operacion_area"] == "maxima":
+                query["area"] = {"$lte": datos_procesados["area"]["valor"]}
+        else:
+            area_value = datos_procesados["area"]["valor"]
+            query["area"] = {"$lte": area_value + area_value*0.2, "$gte": area_value - area_value*0.2}
+
+    if "costos" in datos_procesados:
+        if datos_procesados["operacion_costo"] == "precio máximo":
+            query["costs.sale_amount"] = {"$lte": datos_procesados["costos"][0]}
+        elif datos_procesados["operacion_costo"] == "valor de arriendo":
+            query["costs.rental_amount"] = {"$gte": datos_procesados["costos"][0]}
+        else:
+            query["costs.rental_amount"] = {"$gte": datos_procesados["costos"][0], "$lte": datos_procesados["costos"][-1]}
+    if "banos" in datos_procesados:
+        query["technicalSheet.bathrooms"] = datos_procesados["banos"]
+
+    documentos = list(collection.find(query).limit(10))
+    return documentos