NEKReport/precheck/prompt_builder.py

131 lines
5.6 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# precheck/prompt_builder.py
import json
from typing import Dict, List, Optional
from .scanner import ScanReport, FieldStatus
from .registry import FIELD_REGISTRY, FieldDefinition
class DynamicPromptBuilder:
"""
На основе ScanReport создаёт промпт для LLM, запрашивая только недостающие поля.
"""
def __init__(self, registry: List[FieldDefinition] = None):
self.registry = registry or FIELD_REGISTRY
def build(self, report: ScanReport,
shipping_type_name: Optional[str] = None,
context_text: str = "",
query_text: str = "") -> Dict:
"""
Возвращает словарь с ключами:
- system_prompt: str
- user_prompt: str
- json_schema: dict (схема ожидаемого ответа)
"""
# Определим, какие поля нужно запросить:
# 1) Все поля со статусом NOT_FOUND или UNCERTAIN, которые required_by_default
# 2) Дополнительно можно запросить UNCERTAIN даже если не required
fields_to_request: List[FieldDefinition] = []
for field_def in self.registry:
status = report.field_statuses.get(field_def.name)
if not status:
continue
if status.status == "NOT_FOUND" or status.status == "UNCERTAIN":
fields_to_request.append(field_def)
if not fields_to_request:
return {
"system_prompt": "",
"user_prompt": "",
"json_schema": {},
}
# Строим JSON-схему с обёрткой shipments
schema = {
"type": "object",
"properties": {
"shipments": {
"type": "array",
"items": {
"type": "object",
"properties": {},
"additionalProperties": False
}
}
},
"required": ["shipments"],
"additionalProperties": False
}
shipment_props = schema["properties"]["shipments"]["items"]["properties"]
shipment_required = []
for fd in fields_to_request:
prop = {}
if fd.type == "string":
prop["type"] = "string"
elif fd.type == "number":
prop["type"] = "number"
elif fd.type == "boolean":
prop["type"] = "boolean"
elif fd.type == "array":
prop["type"] = "array"
if fd.name == "dimensions":
prop["items"] = {
"type": "object",
"properties": {
"length_cm": {"type": "number"},
"width_cm": {"type": "number"},
"height_cm": {"type": "number"},
"weight_kg": {"type": "number"},
"volume_cbm": {"type": "number"}
}
}
elif fd.name == "vehicle_dimensions":
prop["items"] = {
"type": "object",
"properties": {
"length_cm": {"type": "number"},
"width_cm": {"type": "number"},
"height_cm": {"type": "number"}
}
}
else:
prop["items"] = {"type": "string"}
elif fd.type == "object":
prop["type"] = "object"
if fd.name == "dangerous_goods":
prop["properties"] = {
"batteries": {"type": "boolean"},
"gases": {"type": "boolean"},
"liquids": {"type": "boolean"},
"dry_ice": {"type": "boolean"}
}
shipment_props[fd.name] = prop
if fd.required_by_default:
shipment_required.append(fd.name)
if shipment_required:
schema["properties"]["shipments"]["items"]["required"] = shipment_required
# Системный промпт
system_prompt = (
"Ты — эксперт по логистике. Из предоставленного текста письма извлеки все описанные грузоперевозки. "
"Для каждой перевозки заполни только указанные ниже поля. "
"Если в письме несколько независимых партий/маршрутов — верни массив с несколькими объектами. "
"Если поле не найдено, оставь null. Ответь строго в формате JSON-объекта с единственным ключом \"shipments\". "
"Не добавляй лишних данных."
)
# Пользовательский промпт
user_prompt = (
f"Контекст письма:\n{context_text}\n\n"
f"Дополнительный запрос: {query_text}\n\n"
f"Ожидаемая JSON-схема ответа:\n{json.dumps(schema, ensure_ascii=False, indent=2)}"
)
return {
"system_prompt": system_prompt,
"user_prompt": user_prompt,
"json_schema": schema
}