BBIT-Kai
80 lines
2.2 KiB
Python
80 lines
2.2 KiB
Python
import json
|
|
import re
|
|
from typing import TypedDict
|
|
|
|
from langchain_core.messages import HumanMessage
|
|
from langgraph.graph import StateGraph, END
|
|
|
|
from config.llm import llmVision
|
|
|
|
|
|
# -------- 定义状态 --------
|
|
class State(TypedDict):
|
|
image_url: str # 图像
|
|
content: str # 最终内容
|
|
|
|
|
|
def send_analyze(state: State, prompt_text: str):
|
|
messages = [
|
|
HumanMessage(
|
|
content=[
|
|
{"type": "text", "text": prompt_text},
|
|
{"type": "image_url", "image_url": {"url": state["image_url"]}},
|
|
]
|
|
)
|
|
]
|
|
return llmVision.invoke(messages).content
|
|
|
|
|
|
def analysis(state: State):
|
|
state["content"] = send_analyze(
|
|
state,
|
|
"""
|
|
提示词示例
|
|
你是一个图像分析助手。现在给你一张车的侧身照片，请你从图中分析车上运输的牲畜种类。
|
|
|
|
要求：
|
|
1. 牲畜种类可能是：牛、羊、猪、鸡、鸭、鹅。
|
|
2. 如果图中无法判断牲畜类型，请在备注字段 remark 中写明“无法识别”或你观察到的情况。
|
|
3. 不允许输出多余文字，直接返回 JSON。
|
|
|
|
JSON 示例格式：
|
|
{
|
|
"livestock_type": "<牲畜种类>", // 如果能识别就填牛/羊/猪/鸡/鸭/鹅
|
|
"remark": "<备注>" // 如果无法识别，写明原因；否则可留空
|
|
}
|
|
请确保输出的 JSON 可以被严格解析。
|
|
""",
|
|
)
|
|
return state
|
|
|
|
|
|
# ------------------------------------------------------------------------ 构建有向图 --------
|
|
workflow = StateGraph(State)
|
|
# 必须先从 START 指向 analysis
|
|
workflow.add_node("analysis", analysis)
|
|
workflow.set_entry_point("analysis")
|
|
workflow.add_edge("analysis", END)
|
|
graph = workflow.compile()
|
|
|
|
|
|
# 执行函数
|
|
|
|
|
|
async def get_vehicle_response(image_url: str):
|
|
final_state = graph.invoke(
|
|
{
|
|
"image_url": image_url,
|
|
}
|
|
)
|
|
# 去掉 ```json 和 ``` 包裹
|
|
content_str = re.sub(r"^```json\s*|\s*```$", "", final_state["content"].strip())
|
|
# 把 JSON 字符串转为字典
|
|
try:
|
|
content_dict = json.loads(content_str)
|
|
except json.JSONDecodeError:
|
|
print("JSON解析失败")
|
|
content_dict = {}
|
|
|
|
return content_dict
|