111 lines
4.6 KiB
Python
Executable File
111 lines
4.6 KiB
Python
Executable File
import os
|
|
import json
|
|
import logging
|
|
import asyncio
|
|
import re
|
|
from typing import Dict, Any, Optional
|
|
from google import genai
|
|
from google.genai import types
|
|
from sqlalchemy import select
|
|
from app.db.session import SessionLocal
|
|
from app.models import SystemParameter
|
|
|
|
logger = logging.getLogger("AI-Service")
|
|
|
|
class AIService:
|
|
"""
|
|
AI Service v1.2.5 - Final Integrated Edition
|
|
- Robot 2: Technikai dúsítás (Search + Regex JSON parsing)
|
|
- Robot 3: OCR (Controlled JSON generation)
|
|
"""
|
|
api_key = os.getenv("GEMINI_API_KEY")
|
|
client = genai.Client(api_key=api_key) if api_key else None
|
|
PRIMARY_MODEL = "gemini-2.0-flash"
|
|
|
|
@classmethod
|
|
async def get_config_delay(cls) -> float:
|
|
try:
|
|
async with SessionLocal() as db:
|
|
stmt = select(SystemParameter).where(SystemParameter.key == "AI_REQUEST_DELAY")
|
|
res = await db.execute(stmt)
|
|
param = res.scalar_one_or_none()
|
|
return float(param.value) if param else 1.0
|
|
except Exception: return 1.0
|
|
|
|
@classmethod
|
|
async def get_clean_vehicle_data(cls, make: str, raw_model: str, v_type: str) -> Optional[Dict[str, Any]]:
|
|
"""Robot 2: Adatbányászat Google Search segítségével."""
|
|
if not cls.client: return None
|
|
await asyncio.sleep(await cls.get_config_delay())
|
|
|
|
search_tool = types.Tool(google_search=types.GoogleSearch())
|
|
|
|
prompt = f"""
|
|
KERESS RÁ az interneten: {make} {raw_model} ({v_type}) pontos gyári modellkódja és technikai adatai.
|
|
Adj választ szigorúan csak egy JSON blokkban:
|
|
{{
|
|
"marketing_name": "tiszta név",
|
|
"synonyms": ["név1", "név2"],
|
|
"technical_code": "gyári kód",
|
|
"year_from": int,
|
|
"year_to": int_vagy_null,
|
|
"ccm": int,
|
|
"kw": int,
|
|
"maintenance": {{ "oil_type": "string", "oil_qty": float, "spark_plug": "string", "coolant": "string" }}
|
|
}}
|
|
FONTOS: A 'technical_code' NEM lehet üres. Ha nem találod, adj 'N/A' értéket!
|
|
"""
|
|
|
|
# Search tool használata esetén a response_mime_type tilos!
|
|
config = types.GenerateContentConfig(
|
|
system_instruction="Profi járműtechnikai adatbányász vagy. Csak tiszta JSON-t válaszolsz markdown kódblokk nélkül.",
|
|
tools=[search_tool],
|
|
temperature=0.1
|
|
)
|
|
|
|
try:
|
|
response = cls.client.models.generate_content(model=cls.PRIMARY_MODEL, contents=prompt, config=config)
|
|
text = response.text
|
|
# Tisztítás: ha az AI mégis tenne bele markdown jeleket
|
|
clean_json = re.sub(r'```json\s*|```', '', text).strip()
|
|
res_json = json.loads(clean_json)
|
|
if isinstance(res_json, list) and len(res_json) > 0: res_json = res_json[0]
|
|
return res_json if isinstance(res_json, dict) else None
|
|
except Exception as e:
|
|
logger.error(f"❌ AI hiba ({make} {raw_model}): {e}")
|
|
return None
|
|
|
|
@classmethod
|
|
async def analyze_document_image(cls, image_data: bytes, doc_type: str) -> Optional[Dict[str, Any]]:
|
|
"""Robot 3: OCR funkció - Forgalmi, Személyi, Számla, Odometer."""
|
|
if not cls.client: return None
|
|
await asyncio.sleep(await cls.get_config_delay())
|
|
|
|
prompts = {
|
|
"identity": "Személyes okmány adatok (név, szám, lejárat).",
|
|
"vehicle_reg": "Forgalmi adatok (rendszám, alvázszám, kW, ccm).",
|
|
"invoice": "Számla adatok (partner, végösszeg, dátum).",
|
|
"odometer": "Csak a kilométeróra állása számként."
|
|
}
|
|
|
|
# Itt maradhat a response_mime_type, mert nem használunk Search-öt
|
|
config = types.GenerateContentConfig(
|
|
system_instruction="Profi OCR dokumentum-elemző vagy. Csak tiszta JSON-t válaszolsz.",
|
|
response_mime_type="application/json"
|
|
)
|
|
|
|
try:
|
|
response = cls.client.models.generate_content(
|
|
model=cls.PRIMARY_MODEL,
|
|
contents=[
|
|
f"Elemezd ezt a képet ({doc_type}): {prompts.get(doc_type, 'OCR')}",
|
|
types.Part.from_bytes(data=image_data, mime_type="image/jpeg")
|
|
],
|
|
config=config
|
|
)
|
|
res_json = json.loads(response.text)
|
|
if isinstance(res_json, list) and len(res_json) > 0: res_json = res_json[0]
|
|
return res_json if isinstance(res_json, dict) else None
|
|
except Exception as e:
|
|
logger.error(f"❌ OCR hiba: {e}")
|
|
return None |