NEKReport/cargo_rag_v2.py

302 lines
13 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.

"""
Новый RAG-движок для анализа грузоперевозок.
Использует русскоязычные критерии из shipping_types.json.
Не хранит сессии, только проводит анализ.
"""
import json
import re
import logging
from pathlib import Path
from typing import List, Dict, Optional, Any
import openai
logger = logging.getLogger(__name__)
# Путь к файлу с типами перевозок (лежит рядом с этим файлом)
SHIPPING_TYPES_PATH = Path(__file__).resolve().parent / "shipping_types.json"
DEFAULT_MODEL = "deepseek-v4-flash" # или ваша модель
def load_shipping_types() -> List[Dict]:
"""Загружает список типов перевозок из JSONфайла."""
try:
with open(SHIPPING_TYPES_PATH, "r", encoding="utf-8") as f:
raw = json.load(f)
if not isinstance(raw, list):
raise ValueError("Ожидался список в shipping_types.json")
logger.info(f"Загружено {len(raw)} типов перевозок")
return raw
except Exception as e:
logger.error(f"Ошибка загрузки shipping_types.json: {e}")
return []
class CargoRAGEngineV2:
"""
Новый движок анализа грузоперевозок.
- Определяет тип перевозки по тексту писем.
- Заполняет критерии (на русском) с помощью LLM.
- Нормализует типы значений (числа, булевы).
"""
def __init__(self, openai_client: openai.OpenAI):
self.client = openai_client
self.shipping_types = load_shipping_types()
self._type_by_name = {
t.get("name", "").strip(): t for t in self.shipping_types
}
# ------------------------------------------------------------------
# Вспомогательные методы для совместимости со старыми эндпоинтами
# ------------------------------------------------------------------
def process_shipping_type_criteria(self, criteria_text: str) -> str:
"""Заглушка возвращает исходный текст, обработка не требуется."""
return criteria_text
def reload_shipping_types(self):
"""Перезагружает типы перевозок из файла."""
self.shipping_types = load_shipping_types()
self._type_by_name = {
t.get("name", "").strip(): t for t in self.shipping_types
}
def normalize_result(self, raw_data: Dict, shipping_type: Dict) -> Dict:
criteria_text = shipping_type.get("criteria", "")
criteria_list = self._criteria_text_to_list(criteria_text)
results = []
for idx, crit in enumerate(criteria_list, start=1):
raw_val = raw_data.get(crit)
normalized_val = self._normalize_value(raw_val)
results.append({
"number": idx, # ← добавили
"criterion": crit,
"value": normalized_val
})
return {
"shipping_type": shipping_type.get("name", ""),
"criteria_results": results
}
def record_cargo_learning(self, **kwargs) -> bool:
"""Заглушка для сохранения примеров (можно реализовать позже)."""
logger.info("record_cargo_learning вызван, но не реализован в V2")
return False
# ------------------------------------------------------------------
# Сборка текста писем
# ------------------------------------------------------------------
def _collect_emails_text(self, emails: List[Dict]) -> str:
parts = []
for mail in emails:
subject = mail.get("subject", "")
sender = mail.get("senderName", "") or mail.get("sender", "")
body = mail.get("body", "")
parts.append(f"Subject: {subject}\nFrom: {sender}\n\n{body}")
return "\n\n---\n\n".join(parts)
# ------------------------------------------------------------------
# Определение типа перевозки (первый проход)
# ------------------------------------------------------------------
def _detect_shipping_type(self, context_text: str, query_text: str = "") -> Optional[Dict]:
if not self.shipping_types:
return None
type_names = [t["name"] for t in self.shipping_types if t.get("name")]
names_list = "\n".join(f"- {name}" for name in type_names)
system = (
"Ты эксперт по логистике. Из списка типов перевозок выбери ОДИН, "
"наиболее подходящий для приведённых ниже писем. "
"Ответь ТОЛЬКО точным названием типа (как в списке), без пояснений."
)
user = (
f"Типы перевозок:\n{names_list}\n\n"
f"Запрос менеджера: {query_text}\n\n"
f"Письма:\n{context_text[:8000]}"
)
try:
resp = self.client.chat.completions.create(
model=DEFAULT_MODEL,
messages=[
{"role": "system", "content": system},
{"role": "user", "content": user}
],
temperature=0.0,
)
type_name = resp.choices[0].message.content.strip()
logger.info(f"LLM определил тип: '{type_name}'")
return self._type_by_name.get(type_name)
except Exception as e:
logger.error(f"Ошибка определения типа: {e}")
return None
# ------------------------------------------------------------------
# Разбор текста критериев в список ключей
# ------------------------------------------------------------------
@staticmethod
def _criteria_text_to_list(criteria_text: str) -> List[str]:
if not isinstance(criteria_text, str):
return []
lines = [ln.strip() for ln in criteria_text.splitlines() if ln.strip()]
items = []
for ln in lines:
parts = ln.split("\t", 1)
if len(parts) == 2:
text = parts[1].strip()
else:
text = re.sub(r"^\d+\.?\s*", "", ln).strip()
if text:
items.append(text)
# Убираем дубликаты, сохраняя порядок
seen = set()
unique = []
for it in items:
if it not in seen:
seen.add(it)
unique.append(it)
return unique
# ------------------------------------------------------------------
# Построение динамического JSON и запрос к LLM (второй проход)
# ------------------------------------------------------------------
def _extract_by_criteria(
self,
context_text: str,
shipping_type: Dict,
query_text: str = ""
) -> Dict[str, Any]:
criteria_text = shipping_type.get("criteria", "")
if not criteria_text:
logger.warning("У типа перевозки отсутствуют критерии")
return {}
criteria_list = self._criteria_text_to_list(criteria_text)
template = {c: None for c in criteria_list}
system = (
"Ты ассистент логиста. Твоя задача извлечь из деловой переписки информацию "
"строго по заданным критериям. Верни **только** JSON-объект, где ключи критерии "
"(на русском языке), а значения найденные данные.\n"
"Правила:\n"
"- Числа (вес, количество) возвращай как числа (не строки).\n"
"- Булевы ответы (да/нет) возвращай как true/false.\n"
"- Если по критерию ничего не найдено, оставь null.\n"
"- Не выдумывай данные, опирайся только на текст писем."
)
user = (
f"Тип перевозки: {shipping_type.get('name', '')}\n"
f"Доп. запрос менеджера: {query_text}\n\n"
f"Шаблон JSON (заполни его):\n{json.dumps(template, ensure_ascii=False, indent=2)}\n\n"
f"Текст писем:\n{context_text}"
)
try:
resp = self.client.chat.completions.create(
model=DEFAULT_MODEL,
messages=[
{"role": "system", "content": system},
{"role": "user", "content": user}
],
temperature=0.0,
)
answer = resp.choices[0].message.content
return self._parse_json_response(answer)
except Exception as e:
logger.error(f"Ошибка извлечения критериев: {e}")
return {}
def _parse_json_response(self, text: str) -> Dict:
text = re.sub(r"```json\s*", "", text)
text = re.sub(r"```", "", text)
text = text.strip()
try:
return json.loads(text)
except json.JSONDecodeError:
start = text.find("{")
end = text.rfind("}")
if start != -1 and end != -1:
try:
return json.loads(text[start:end+1])
except json.JSONDecodeError:
pass
logger.warning("Не удалось распарсить JSON-ответ LLM")
return {}
# ------------------------------------------------------------------
# Нормализация значений
# ------------------------------------------------------------------
@staticmethod
def _normalize_value(val: Any) -> Any:
if not isinstance(val, str):
return val
s = val.strip()
low = s.lower()
if low in ("да", "yes", "true", "есть", "присутствует"):
return True
if low in ("нет", "no", "false", "отсутствует"):
return False
numeric_part = re.sub(r"[^\d.,\-]", "", s).replace(",", ".")
if numeric_part:
try:
if "." in numeric_part:
return float(numeric_part)
else:
return int(numeric_part)
except ValueError:
pass
return s
def normalize_result(self, raw_data: Dict, shipping_type: Dict) -> Dict:
criteria_text = shipping_type.get("criteria", "")
criteria_list = self._criteria_text_to_list(criteria_text)
results = []
for crit in criteria_list:
raw_val = raw_data.get(crit)
normalized_val = self._normalize_value(raw_val)
results.append({
"criterion": crit,
"value": normalized_val
})
return {
"shipping_type": shipping_type.get("name", ""),
"criteria_results": results
}
# ------------------------------------------------------------------
# Основной метод анализа (вызывается из API)
# ------------------------------------------------------------------
def analyze(self, emails: List[Dict], query_text: str = "") -> Dict:
"""
Главный метод анализа писем.
Возвращает:
{
"shipping_type": "...",
"criteria_results": [{"criterion": ..., "value": ...}, ...],
"emails_count": int
}
или {"error": "..."} в случае ошибки.
"""
context = self._collect_emails_text(emails)
if not context.strip():
return {"error": "Нет текста для анализа"}
shipping_type = self._detect_shipping_type(context, query_text)
if not shipping_type:
return {"error": "Не удалось определить тип перевозки"}
raw_data = self._extract_by_criteria(context, shipping_type, query_text)
if not raw_data:
return {"error": "Не удалось извлечь данные по критериям"}
report = self.normalize_result(raw_data, shipping_type)
report["emails_count"] = len(emails)
return report