update: resolve comments and support any directory zip file

This commit is contained in:
Stitch-z 2023-10-24 14:01:33 +08:00
parent d8afe9ea86
commit 5e27ee4bad
7 changed files with 62 additions and 102 deletions

View file

@ -8,10 +8,10 @@
"""
import os
from pathlib import Path
from typing import List
import pytest
from pathlib import Path
from metagpt.actions.invoice_ocr import InvoiceOCR, GenerateTable, ReplyQuestion
@ -48,7 +48,7 @@ async def test_invoice_ocr(invoice_path: str):
),
]
)
async def test_generate_table(invoice_path: str, expected_result: List[dict]):
async def test_generate_table(invoice_path: str, expected_result: list[dict]):
invoice_path = os.path.abspath(os.path.join(os.getcwd(), invoice_path))
filename = os.path.basename(invoice_path)
ocr_result = await InvoiceOCR().run(file_path=Path(invoice_path), filename=filename)

View file

@ -7,11 +7,10 @@
@File : test_invoice_ocr_assistant.py
"""
import os
import pandas as pd
from typing import List
from pathlib import Path
import pytest
import pandas as pd
from metagpt.roles.invoice_ocr_assistant import InvoiceOCRAssistant
from metagpt.schema import Message
@ -23,8 +22,8 @@ from metagpt.schema import Message
[
(
"Invoicing date",
"../../data/invoices/invoice-1.pdf",
"../../../data/invoice_table/invoice-1.xlsx",
Path("../../data/invoices/invoice-1.pdf"),
Path("../../../data/invoice_table/invoice-1.xlsx"),
[
{
"收款人": "小明",
@ -36,8 +35,8 @@ from metagpt.schema import Message
),
(
"Invoicing date",
"../../data/invoices/invoice-2.png",
"../../../data/invoice_table/invoice-2.xlsx",
Path("../../data/invoices/invoice-2.png"),
Path("../../../data/invoice_table/invoice-2.xlsx"),
[
{
"收款人": "铁头",
@ -49,8 +48,8 @@ from metagpt.schema import Message
),
(
"Invoicing date",
"../../data/invoices/invoice-3.jpg",
"../../../data/invoice_table/invoice-3.xlsx",
Path("../../data/invoices/invoice-3.jpg"),
Path("../../../data/invoice_table/invoice-3.xlsx"),
[
{
"收款人": "夏天",
@ -62,8 +61,8 @@ from metagpt.schema import Message
),
(
"Invoicing date",
"../../data/invoices/invoice-4.zip",
"../../../data/invoice_table/invoice-4.xlsx",
Path("../../data/invoices/invoice-4.zip"),
Path("../../../data/invoice_table/invoice-4.xlsx"),
[
{
"收款人": "小明",
@ -89,17 +88,17 @@ from metagpt.schema import Message
)
async def test_invoice_ocr_assistant(
query: str,
invoice_path: str,
invoice_table_path: str,
expected_result: List[dict]
invoice_path: Path,
invoice_table_path: Path,
expected_result: list[dict]
):
invoice_path = os.path.abspath(os.path.join(os.getcwd(), invoice_path))
invoice_path = Path.cwd() / invoice_path
role = InvoiceOCRAssistant()
await role.run(Message(
content=query,
instruct_content={"file_path": invoice_path}
))
invoice_table_path = os.path.abspath(os.path.join(os.getcwd(), invoice_table_path))
invoice_table_path = Path.cwd() / invoice_table_path
df = pd.read_excel(invoice_table_path)
dict_result = df.to_dict(orient='records')
assert dict_result == expected_result