diff --git a/services/extractor.py b/services/extractor.py index ff9081a..5ba52d2 100644 --- a/services/extractor.py +++ b/services/extractor.py @@ -7,10 +7,11 @@ from langchain_core.output_parsers import PydanticOutputParser from langchain_google_genai import ChatGoogleGenerativeAI from pydantic import BaseModel -from utils.parsers import process_profile, process_passport -from validation.from_description import FromDescription +from utils.parsers import process_profile, process_passport, process_account +from validation.from_account import FromAccount from validation.from_passport import FromPassport from validation.from_profile import FromProfile +from validation.from_description import FromDescription def extract_description(client_data: dict[str, Any]) -> FromDescription: passport_data = client_data.get("description") @@ -31,6 +32,26 @@ def extract_description(client_data: dict[str, Any]) -> FromDescription: return result +def extract_account(client_data: dict[str, Any])-> FromAccount: + account_data = client_data.get("account") + + prompt_template = ( + "Extract the following information from the provided text.\n" + "Return only JSON matching this format:\n{format_instructions}\n\n" + "Trim email if needed\n" + "Passport text:\n{processed_text}" + ) + + result = __run_extraction_chain( + raw_file_data=account_data, + file_processor=process_account, + pydantic_model=FromAccount, + prompt_template=prompt_template, + ) + + return result + + def extract_passport(client_data: dict[str, Any]) -> FromPassport: passport_data = client_data.get("passport") diff --git a/tests/dummy.py b/tests/dummy.py index 3922cd9..efd5cb0 100644 --- a/tests/dummy.py +++ b/tests/dummy.py @@ -1,4 +1,7 @@ from datetime import date +import json +from typing import Any +import config from validation.cross_validate import ExtractedData from validation.from_account import FromAccount from validation.from_description import FromDescription @@ -106,4 +109,12 @@ def dummy_data() -> ExtractedData: description=dummy_description(), passport=dummy_passport(), profile=dummy_profile(), - ) \ No newline at end of file + ) + +def dummy_client_data() -> dict[str, Any]: + # TODO make generic + resp_path = f"{config.GAME_FILES_DIR}/65e6ec83-88b1-4d1f-8560-a1418803348b/000000_decision_accept_active/000000_response.json" + out = {} + with open(resp_path, "r") as file: + out = json.loads(str(json.loads(file.read()))) + return out \ No newline at end of file