47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
# app/agents/receipt_agent.py
|
|
from langchain_core.messages import HumanMessage
|
|
from langchain_core.output_parsers import PydanticOutputParser
|
|
from langchain.prompts import PromptTemplate
|
|
from ..core.llm import llm
|
|
from ..schemas import ReceiptInfo
|
|
|
|
parser = PydanticOutputParser(pydantic_object=ReceiptInfo)
|
|
|
|
# 更新Prompt模板以包含语言信息
|
|
receipt_template = """
|
|
You are a highly accurate receipt information extraction robot.
|
|
The document's primary language is '{language}'.
|
|
Please extract all key information from the following receipt image.
|
|
If some information is not present in the image, leave it as null.
|
|
Please strictly follow the JSON format below, without adding any extra explanations or comments.
|
|
|
|
{format_instructions}
|
|
"""
|
|
|
|
receipt_prompt = PromptTemplate(
|
|
template=receipt_template,
|
|
input_variables=["language"],
|
|
partial_variables={"format_instructions": parser.get_format_instructions()},
|
|
)
|
|
|
|
|
|
async def agent_extract_receipt_info(image_base64: str, language: str) -> ReceiptInfo:
|
|
"""Agent 2: Extracts receipt information from an image, aware of the document's language."""
|
|
print(f"--- [Agent 2] Calling multimodal LLM to extract receipt info (Language: {language})...")
|
|
|
|
prompt_text = await receipt_prompt.aformat(language=language)
|
|
|
|
msg = HumanMessage(
|
|
content=[
|
|
{"type": "text", "text": prompt_text},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": f"data:image/png;base64,{image_base64}",
|
|
},
|
|
]
|
|
)
|
|
|
|
chain = llm | parser
|
|
receipt_info = await chain.ainvoke([msg])
|
|
return receipt_info
|