ai_agent.py - main - Paff/declarant

ai_agent.py
203 строки8 KB
1pash1985-gif
Декларант - веб-версия для Timeweb Cloud
15 июл 2026, 09:36
15 июл 2026, 09:36a6b08b8
# -*- coding: utf-8 -*-
"""
AI-агент для анализа расхождений (Groq / Gemini — бесплатные API).

Единый интерфейс, провайдер переключается в настройках.
Использует openai-совместимый SDK (base_url override).
"""
from __future__ import annotations

import json
import os
from dataclasses import dataclass
from typing import Optional

try:
    from openai import OpenAI
except ImportError:
    OpenAI = None  # type: ignore


# --- Конфигурация провайдеров ---

PROVIDERS = {
    'groq': {
        'base_url': 'https://api.groq.com/openai/v1',
        'model': 'llama-3.3-70b-versatile',
        'env_key': 'GROQ_API_KEY',
    },
    'gemini': {
        'base_url': 'https://generativelanguage.googleapis.com/v1beta/openai/',
        'model': 'gemini-2.5-flash',
        'env_key': 'GEMINI_API_KEY',
    },
}

SYSTEM_PROMPT = """\
Ты — помощник таможенного декларанта. Работаешь с данными о цветах (NL/Ecuador/Colombia).

Контекст:
- NL-маркировки — коды клиентов в колонках шаблона "Таблица европа" (IZUM, BORT20, PRADZ, B-BOOM55...)
- Формат названий NL: "Chr G" = Chrysanthemum Grande, "Chr T" = Tros/spray, "Fr Du" = Freesia Double
- Семейства: ROSA, CHRYSANTHEMUM, LISIANTHUS, ALSTROEMERIA, GERBERA, LILIUM, GYPSOPHILA, HYPERICUM, FREESIA, BOUVARDIA, DIANTHUS, MATTHIOLA, PEONY...

Отвечай ТОЛЬКО в формате JSON. Без markdown, без пояснений вне JSON.
"""


@dataclass
class AISuggestion:
    """Одна подсказка от AI."""
    field: str           # 'marking', 'column', 'family'
    original: str        # исходное значение
    suggestion: str      # предложенная коррекция
    confidence: float    # 0.0 .. 1.0
    reason: str          # пояснение


class AIAgent:
    """Обёртка над LLM-провайдером."""

    def __init__(self, provider: str = 'groq', api_key: str = ''):
        self.provider_name = provider
        cfg = PROVIDERS.get(provider, PROVIDERS['groq'])
        self.model = cfg['model']
        self.base_url = cfg['base_url']
        self.api_key = api_key or os.environ.get(cfg['env_key'], '')
        self._client: Optional[OpenAI] = None

    @property
    def is_available(self) -> bool:
        return bool(OpenAI and self.api_key)

    def _get_client(self) -> OpenAI:
        if self._client is None:
            self._client = OpenAI(base_url=self.base_url, api_key=self.api_key)
        return self._client

    def _call(self, user_prompt: str) -> str:
        """Отправить запрос, вернуть текст ответа."""
        client = self._get_client()
        resp = client.chat.completions.create(
            model=self.model,
            messages=[
                {'role': 'system', 'content': SYSTEM_PROMPT},
                {'role': 'user', 'content': user_prompt},
            ],
            temperature=0.2,
            max_tokens=2000,
        )
        return resp.choices[0].message.content or ''

    # --- Публичные методы ---

    def verify_markings(
        self,
        unmapped_markings: list[str],
        template_columns: list[str],
    ) -> list[AISuggestion]:
        """Предложить колонки для не-найденных маркировок."""
        if not self.is_available or not unmapped_markings:
            return []
        prompt = json.dumps({
            'task': 'match_markings_to_columns',
            'unmapped_markings': unmapped_markings[:20],
            'template_columns': template_columns,
            'instruction': (
                'Для каждого unmapped marking найди наиболее вероятную колонку из template_columns. '
                'Ответь JSON-массивом: [{"marking": "...", "column": "...", "confidence": 0.0-1.0, "reason": "..."}]. '
                'Если уверенности нет — confidence < 0.3.'
            ),
        }, ensure_ascii=False)
        try:
            raw = self._call(prompt)
            # Извлечь JSON из ответа
            raw = raw.strip()
            if raw.startswith('```'):
                raw = raw.split('\n', 1)[1].rsplit('```', 1)[0]
            items = json.loads(raw)
            results = []
            for it in items:
                results.append(AISuggestion(
                    field='column',
                    original=it.get('marking', ''),
                    suggestion=it.get('column', ''),
                    confidence=float(it.get('confidence', 0)),
                    reason=it.get('reason', ''),
                ))
            return results
        except Exception:
            return []

    def classify_unknown(self, varieties: list[str]) -> list[AISuggestion]:
        """Предложить семейства для нераспознанных сортов."""
        if not self.is_available or not varieties:
            return []
        prompt = json.dumps({
            'task': 'classify_flower_varieties',
            'unknown_varieties': varieties[:30],
            'known_families': [
                'ROSA', 'CHRYSANTHEMUM', 'LISIANTHUS', 'ALSTROEMERIA',
                'GERBERA', 'LILIUM', 'GYPSOPHILA', 'HYPERICUM', 'FREESIA',
                'BOUVARDIA', 'DIANTHUS', 'MATTHIOLA', 'PEONY', 'HYDRANGEA',
                'DELPHINIUM', 'SOLIDAGO', 'LIMONIUM', 'ANEMONE', 'ANTIRRHINUM',
            ],
            'instruction': (
                'Определи семейство для каждого сорта. '
                'Ответь JSON-массивом: [{"variety": "...", "family": "...", "confidence": 0.0-1.0, "reason": "..."}]. '
                'Hint: "Chr" = Chrysanthemum, "Fr" = Freesia, "Bou" = Bouvardia, "Ger"/"Ge" = Gerbera, "Li" = Lilium.'
            ),
        }, ensure_ascii=False)
        try:
            raw = self._call(prompt)
            raw = raw.strip()
            if raw.startswith('```'):
                raw = raw.split('\n', 1)[1].rsplit('```', 1)[0]
            items = json.loads(raw)
            results = []
            for it in items:
                results.append(AISuggestion(
                    field='family',
                    original=it.get('variety', ''),
                    suggestion=it.get('family', ''),
                    confidence=float(it.get('confidence', 0)),
                    reason=it.get('reason', ''),
                ))
            return results
        except Exception:
            return []

    def suggest_corrections(
        self,
        problems: list[dict],
    ) -> list[AISuggestion]:
        """Батч-анализ проблем: маркировки + классификация."""
        if not self.is_available or not problems:
            return []
        prompt = json.dumps({
            'task': 'suggest_corrections',
            'problems': problems[:30],
            'instruction': (
                'Проанализируй проблемы и предложи решения. '
                'Ответь JSON-массивом: [{"field": "marking|column|family", '
                '"original": "...", "suggestion": "...", "confidence": 0.0-1.0, "reason": "..."}]'
            ),
        }, ensure_ascii=False)
        try:
            raw = self._call(prompt)
            raw = raw.strip()
            if raw.startswith('```'):
                raw = raw.split('\n', 1)[1].rsplit('```', 1)[0]
            items = json.loads(raw)
            results = []
            for it in items:
                results.append(AISuggestion(
                    field=it.get('field', 'unknown'),
                    original=it.get('original', ''),
                    suggestion=it.get('suggestion', ''),
                    confidence=float(it.get('confidence', 0)),
                    reason=it.get('reason', ''),
                ))
            return results
        except Exception:
            return []