diff --git a/src/main/kotlin/com/android/trisolarisserver/component/DocumentExtractionService.kt b/src/main/kotlin/com/android/trisolarisserver/component/DocumentExtractionService.kt new file mode 100644 index 0000000..e856d67 --- /dev/null +++ b/src/main/kotlin/com/android/trisolarisserver/component/DocumentExtractionService.kt @@ -0,0 +1,310 @@ +package com.android.trisolarisserver.component + +import com.android.trisolarisserver.controller.DocumentPrompts +import com.android.trisolarisserver.db.repo.GuestRepo +import com.android.trisolarisserver.models.booking.GuestDocument +import com.android.trisolarisserver.models.booking.GuestVehicle +import com.android.trisolarisserver.repo.GuestVehicleRepo +import com.android.trisolarisserver.repo.PropertyRepo +import java.time.OffsetDateTime +import java.util.UUID + +@org.springframework.stereotype.Component +class DocumentExtractionService( + private val llamaClient: LlamaClient, + private val guestRepo: GuestRepo, + private val guestVehicleRepo: GuestVehicleRepo, + private val propertyRepo: PropertyRepo +) { + fun extractAndApply(imageUrl: String, document: GuestDocument, propertyId: UUID): ExtractionResult { + val results = linkedMapOf() + val detections = listOf( + Detection( + detect = { + results["isVehiclePhoto"] = llamaClient.ask( + imageUrl, + "IS THIS A VEHICLE NUMBER PLATE PHOTO? Answer YES or NO only." + ) + isYes(results["isVehiclePhoto"]) + }, + handle = { + results["vehicleNumber"] = llamaClient.ask( + imageUrl, + "VEHICLE NUMBER PLATE? Reply only number or NONE." + ) + } + ), + Detection( + detect = { + results["hasAadhar"] = llamaClient.ask( + imageUrl, + "CONTAINS AADHAAR? Answer YES or NO only." + ) + results["hasUidai"] = llamaClient.ask( + imageUrl, + "CONTAINS UIDAI? Answer YES or NO only." + ) + isYes(results["hasAadhar"]) || isYes(results["hasUidai"]) + }, + handle = { + val aadharQuestions = linkedMapOf( + "hasAddress" to "POSTAL ADDRESS PRESENT? Answer YES or NO only.", + "hasDob" to "DOB? Reply YES or NO.", + "hasGenderMentioned" to "GENDER MENTIONED? Reply YES or NO." + ) + for ((key, question) in aadharQuestions) { + results[key] = llamaClient.ask(imageUrl, question) + } + val hasAddress = isYes(results["hasAddress"]) + if (hasAddress) { + val addressQuestions = linkedMapOf( + DocumentPrompts.PIN_CODE, + DocumentPrompts.ADDRESS + ) + for ((key, question) in addressQuestions) { + results[key] = llamaClient.ask(imageUrl, question) + } + } + val hasDob = isYes(results["hasDob"]) + val hasGender = isYes(results["hasGenderMentioned"]) + if (hasDob && hasGender) { + val aadharFrontQuestions = linkedMapOf( + DocumentPrompts.NAME, + DocumentPrompts.DOB, + DocumentPrompts.ID_NUMBER, + DocumentPrompts.GENDER + ) + for ((key, question) in aadharFrontQuestions) { + results[key] = llamaClient.ask(imageUrl, question) + } + } + } + ), + Detection( + detect = { + results["hasDrivingLicence"] = llamaClient.ask( + imageUrl, + "CONTAINS DRIVING LICENCE? Answer YES or NO only." + ) + results["hasTransportDept"] = llamaClient.ask( + imageUrl, + "CONTAINS TRANSPORT DEPARTMENT? Answer YES or NO only." + ) + isYes(results["hasDrivingLicence"]) || isYes(results["hasTransportDept"]) + }, + handle = { + val drivingQuestions = linkedMapOf( + DocumentPrompts.NAME, + DocumentPrompts.DOB, + "idNumber" to "DL NUMBER? Reply only number or NONE.", + DocumentPrompts.ADDRESS, + DocumentPrompts.PIN_CODE, + DocumentPrompts.CITY, + DocumentPrompts.GENDER, + DocumentPrompts.NATIONALITY + ) + for ((key, question) in drivingQuestions) { + results[key] = llamaClient.ask(imageUrl, question) + } + } + ), + Detection( + detect = { + results["hasElectionCommission"] = llamaClient.ask( + imageUrl, + "CONTAINS ELECTION COMMISSION OF INDIA? Answer YES or NO only." + ) + isYes(results["hasElectionCommission"]) + }, + handle = { + val voterQuestions = linkedMapOf( + DocumentPrompts.NAME, + DocumentPrompts.DOB, + "idNumber" to "VOTER ID NUMBER? Reply only number or NONE.", + DocumentPrompts.ADDRESS, + DocumentPrompts.PIN_CODE, + DocumentPrompts.CITY, + DocumentPrompts.GENDER, + DocumentPrompts.NATIONALITY + ) + for ((key, question) in voterQuestions) { + results[key] = llamaClient.ask(imageUrl, question) + } + } + ), + Detection( + detect = { + results["hasIncomeTaxDept"] = llamaClient.ask( + imageUrl, + "CONTAINS INCOME TAX DEPARTMENT? Answer YES or NO only." + ) + isYes(results["hasIncomeTaxDept"]) + }, + handle = { + val panQuestions = linkedMapOf( + DocumentPrompts.NAME, + DocumentPrompts.DOB, + "idNumber" to "PAN NUMBER? Reply only number or NONE.", + DocumentPrompts.ADDRESS, + DocumentPrompts.PIN_CODE, + DocumentPrompts.CITY, + DocumentPrompts.GENDER, + DocumentPrompts.NATIONALITY + ) + for ((key, question) in panQuestions) { + results[key] = llamaClient.ask(imageUrl, question) + } + } + ), + Detection( + detect = { + results["hasPassport"] = llamaClient.ask( + imageUrl, + "CONTAINS PASSPORT? Answer YES or NO only." + ) + isYes(results["hasPassport"]) + }, + handle = { + val passportQuestions = linkedMapOf( + DocumentPrompts.NAME, + DocumentPrompts.DOB, + "idNumber" to "PASSPORT NUMBER? Reply only number or NONE.", + DocumentPrompts.ADDRESS, + DocumentPrompts.PIN_CODE, + DocumentPrompts.CITY, + DocumentPrompts.GENDER, + DocumentPrompts.NATIONALITY + ) + for ((key, question) in passportQuestions) { + results[key] = llamaClient.ask(imageUrl, question) + } + } + ) + ) + + var handled = false + for (detection in detections) { + if (detection.detect()) { + detection.handle() + handled = true + break + } + } + + if (!handled) { + val generalQuestions = linkedMapOf( + DocumentPrompts.NAME, + DocumentPrompts.DOB, + DocumentPrompts.ID_NUMBER, + DocumentPrompts.ADDRESS, + DocumentPrompts.VEHICLE_NUMBER, + DocumentPrompts.PIN_CODE, + DocumentPrompts.CITY, + DocumentPrompts.GENDER, + DocumentPrompts.NATIONALITY + ) + for ((key, question) in generalQuestions) { + results[key] = llamaClient.ask(imageUrl, question) + } + } + + results["docType"] = computeDocType(results, handled) + applyGuestUpdates(document, propertyId, results) + return ExtractionResult(results, handled) + } + + private fun isYes(value: String?): Boolean { + return value.orEmpty().contains("YES", ignoreCase = true) + } + + private fun computeDocType(results: Map, handled: Boolean): String { + if (!handled) return "GENERAL" + return when { + isYes(results["hasCourt"]) || + isYes(results["hasHighCourt"]) || + isYes(results["hasSupremeCourt"]) || + isYes(results["hasJudiciary"]) -> "COURT_ID" + isYes(results["hasPolice"]) -> "POLICE_ID" + isYes(results["hasPassport"]) -> "PASSPORT" + isYes(results["hasTransportDept"]) || + isYes(results["hasDrivingLicence"]) -> "TRANSPORT" + isYes(results["hasIncomeTaxDept"]) -> "PAN" + isYes(results["hasElectionCommission"]) -> "VOTER_ID" + isYes(results["hasAadhar"]) || + isYes(results["hasUidai"]) -> { + if (isYes(results["hasAddress"])) "AADHAR_BACK" else "AADHAR_FRONT" + } + results["vehicleNumber"].orEmpty().isNotBlank() && + !results["vehicleNumber"]!!.contains("NONE", true) -> "VEHICLE" + isYes(results["isVehiclePhoto"]) -> "VEHICLE_PHOTO" + else -> "UNKNOWN" + } + } + + private fun applyGuestUpdates( + document: GuestDocument, + propertyId: UUID, + results: Map + ) { + val extractedName = cleanedValue(results[DocumentPrompts.NAME.first]) + val extractedAddress = cleanedValue(results[DocumentPrompts.ADDRESS.first]) + val guestIdValue = document.guest.id + if (guestIdValue != null && (extractedName != null || extractedAddress != null)) { + val guestEntity = guestRepo.findById(guestIdValue).orElse(null) + if (guestEntity != null) { + var updated = false + if (guestEntity.name.isNullOrBlank() && extractedName != null) { + guestEntity.name = extractedName + updated = true + } + if (guestEntity.addressText.isNullOrBlank() && extractedAddress != null) { + guestEntity.addressText = extractedAddress + updated = true + } + if (updated) { + guestEntity.updatedAt = OffsetDateTime.now() + guestRepo.save(guestEntity) + } + } + } + + val extractedVehicle = cleanedValue(results["vehicleNumber"]) + if (isYes(results["isVehiclePhoto"]) && extractedVehicle != null) { + val guestIdSafe = document.guest.id + if (guestIdSafe != null && + !guestVehicleRepo.existsByPropertyIdAndVehicleNumberIgnoreCase(propertyId, extractedVehicle) + ) { + val property = propertyRepo.findById(propertyId).orElse(null) + val guestEntity = guestRepo.findById(guestIdSafe).orElse(null) + if (property != null && guestEntity != null) { + guestVehicleRepo.save( + GuestVehicle( + property = property, + guest = guestEntity, + booking = document.booking, + vehicleNumber = extractedVehicle + ) + ) + } + } + } + } +} + +data class ExtractionResult( + val results: LinkedHashMap, + val handled: Boolean +) + +private data class Detection( + val detect: () -> Boolean, + val handle: () -> Unit +) + +private fun cleanedValue(value: String?): String? { + val trimmed = value?.trim().orEmpty() + if (trimmed.isBlank()) return null + val upper = trimmed.uppercase() + if (upper == "NONE" || upper == "N/A" || upper == "NA" || upper == "NULL") return null + return trimmed +} diff --git a/src/main/kotlin/com/android/trisolarisserver/controller/GuestDocumentResponses.kt b/src/main/kotlin/com/android/trisolarisserver/controller/GuestDocumentResponses.kt new file mode 100644 index 0000000..abb7e4d --- /dev/null +++ b/src/main/kotlin/com/android/trisolarisserver/controller/GuestDocumentResponses.kt @@ -0,0 +1,46 @@ +package com.android.trisolarisserver.controller + +import com.android.trisolarisserver.models.booking.GuestDocument +import com.fasterxml.jackson.databind.ObjectMapper +import java.util.UUID + +data class GuestDocumentResponse( + val id: UUID, + val propertyId: UUID, + val guestId: UUID, + val bookingId: UUID, + val uploadedByUserId: UUID, + val uploadedAt: String, + val originalFilename: String, + val contentType: String?, + val sizeBytes: Long, + val extractedData: Map?, + val extractedAt: String? +) + +internal fun GuestDocument.toResponse(objectMapper: ObjectMapper): GuestDocumentResponse { + val id = id ?: throw IllegalStateException("Document id missing") + val extracted: Map? = extractedData?.let { + try { + val raw = objectMapper.readValue(it, Map::class.java) + raw.entries.associate { entry -> + entry.key.toString() to (entry.value?.toString() ?: "") + } + } catch (_: Exception) { + null + } + } + return GuestDocumentResponse( + id = id, + propertyId = property.id!!, + guestId = guest.id!!, + bookingId = booking.id!!, + uploadedByUserId = uploadedBy.id!!, + uploadedAt = uploadedAt.toString(), + originalFilename = originalFilename, + contentType = contentType, + sizeBytes = sizeBytes, + extractedData = extracted, + extractedAt = extractedAt?.toString() + ) +}