File size: 5,619 Bytes
2edb151
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
from __future__ import annotations

import hashlib
import shutil
import threading
from pathlib import Path

from app.config import Settings
from app.db import (
    get_by_sha,
    insert_receipt,
    open_db,
    set_line_match,
    update_receipt_status,
    upsert_vector,
)
from app.embed import embed_texts
from app.extract import extract_receipt
from app.match import match_receipt
from app.media import to_jpeg_bytes
from app.ocr import maybe_ocr
from app.schemas import ProcessResult, ReceiptStatus
from backends import build_embed, build_llm, build_ocr

_LOCK = threading.Lock()


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _safe_move(src: Path, dest_dir: Path) -> Path:
    dest_dir.mkdir(parents=True, exist_ok=True)
    dest = dest_dir / src.name
    stem, suffix = dest.stem, dest.suffix
    n = 1
    while dest.exists():
        dest = dest_dir / f"{stem}-{n}{suffix}"
        n += 1
    shutil.move(str(src), str(dest))
    return dest


def process_file(path: Path, settings: Settings) -> ProcessResult:
    with _LOCK:
        return _process_file_locked(path, settings)


def _process_file_locked(path: Path, settings: Settings) -> ProcessResult:
    settings.ensure_dirs()
    src = Path(path)
    digest = sha256_file(src)
    con = open_db(settings)
    try:
        existing = get_by_sha(con, digest)
        if existing is not None:
            if src.parent.resolve() == settings.inbox_dir.resolve():
                _safe_move(src, settings.processed_dir)
            return ProcessResult(
                receipt_id=int(existing["id"]),
                status=ReceiptStatus.duplicate,
                source_path=str(src),
                error="duplicate sha256",
            )
        working = src
        if src.parent.resolve() == settings.inbox_dir.resolve():
            working = _safe_move(src, settings.processing_dir)

        ocr_backend = build_ocr(settings)
        llm = build_llm(settings)
        embed = build_embed(settings)
        ocr = maybe_ocr(ocr_backend, working)
        ocr_text = None if ocr is None else ocr.text
        try:
            image = to_jpeg_bytes(working, settings)
        except Exception as exc:
            failed = _safe_move(working, settings.failed_dir)
            rid = insert_receipt(
                con,
                source_path=str(failed),
                sha256=digest,
                status=ReceiptStatus.failed,
                ocr_text=ocr_text,
                error=str(exc),
            )
            return ProcessResult(
                receipt_id=rid,
                status=ReceiptStatus.failed,
                source_path=str(failed),
                error=str(exc),
            )

        if image is None and not ocr_text:
            rid = insert_receipt(
                con,
                source_path=str(working),
                sha256=digest,
                status=ReceiptStatus.needs_ocr,
            )
            return ProcessResult(
                receipt_id=rid,
                status=ReceiptStatus.needs_ocr,
                source_path=str(working),
                error="no image/text for extract",
            )

        try:
            extract = extract_receipt(
                llm, settings=settings, image_jpeg=image, ocr_text=ocr_text
            )
        except Exception as exc:
            final = _safe_move(working, settings.processed_dir)
            rid = insert_receipt(
                con,
                source_path=str(final),
                sha256=digest,
                status=ReceiptStatus.needs_extract,
                ocr_text=ocr_text,
                error=str(exc),
            )
            return ProcessResult(
                receipt_id=rid,
                status=ReceiptStatus.needs_extract,
                source_path=str(final),
                error=str(exc),
            )

        final = _safe_move(working, settings.processed_dir)
        rid = insert_receipt(
            con,
            source_path=str(final),
            sha256=digest,
            status=ReceiptStatus.needs_review,
            extract=extract,
            ocr_text=ocr_text,
        )
        try:
            if extract.vendor or extract.line_items:
                blob = " | ".join(
                    [
                        extract.doc_kind,
                        extract.category,
                        extract.vendor or "",
                        extract.date.isoformat() if extract.date else "",
                        *(item.description for item in extract.line_items[:12]),
                    ]
                )
                vec = embed_texts(embed, [blob], input_type="passage", settings=settings)[0]
                upsert_vector(con, "receipt_vec", "receipt_id", rid, vec)
            matches = match_receipt(con, settings, embed, extract)
            line_rows = con.execute(
                "SELECT id FROM line_items WHERE receipt_id = ? ORDER BY id", (rid,)
            ).fetchall()
            for row, hit in zip(line_rows, matches, strict=False):
                if hit.reason != "vendor knn":
                    set_line_match(con, int(row["id"]), hit)
        except Exception:
            matches = []
        return ProcessResult(
            receipt_id=rid,
            status=ReceiptStatus.needs_review,
            source_path=str(final),
            extract=extract,
            matches=matches,
        )
    finally:
        con.close()