BBIT-Kai
79 lines
2.3 KiB
Python
79 lines
2.3 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, # todo
|
|
"""
|
|
你是一个图像分析助手。现在给你一张道路图片，你需要观察远离你的第二根车道上，画面中心的车辆，请你从分析该车辆。
|
|
|
|
要求：
|
|
1. have_animal 字段中填写 true
|
|
2. livestock_type 字段中填写 “货物种类”，例如 牛、羊、猪、鸡、鸭、鹅、钢管、土渣等任何你观察到的
|
|
3. remark 字段 需要你简短的描述该车辆车身情况，例如什么颜色的车身，你需要完整的尽你所能形容一下。
|
|
3. 不允许输出多余文字，直接返回 JSON。
|
|
|
|
JSON 示例格式：
|
|
{
|
|
"have_animal": true,
|
|
"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
|