model_client.py

import base64
import io
import json
import os
import re
import requests

API_BASE = "https://api.modelbest.cn/v1/chat/completions"
MODEL = "MiniCPM-V-4.6-Instruct"
API_KEY = os.environ.get("OPENBMB_API_KEY", "")

SARVAM_TRANSLATE = "https://api.sarvam.ai/translate"
SARVAM_KEY = os.environ.get("SARVAM_API_KEY", "")

VALID_CATEGORIES = [
    "pothole", "road_damage", "road_blocking", "footpath", "streetlight",
    "garbage", "drain", "waterlogging", "stray_animal",
    "dead_animal", "other",
]

VISION_PROMPT = (
    "You are a strict civic issue classifier for Bengaluru city. BBMP handles "
    "ONLY these public problems: potholes, damaged roads, roads blocked by debris "
    "or fallen trees, broken footpaths, dead streetlights, garbage or illegal "
    "dumping, blocked or open drains, water leaks, waterlogging, and stray or "
    "dead animals on public land. "
    "Look at this photo and be strict. Set is_civic_issue to true ONLY if the "
    "photo clearly and unambiguously shows one of those specific problems on a "
    "public street or public land. "
    "Set is_civic_issue to FALSE for: people, faces, selfies, animals that are "
    "fine, normal walls, buildings, rooms, indoor scenes, food, plants, vehicles, "
    "screenshots, documents, landscapes, or any image where you are not clearly "
    "seeing a specific civic problem. When in doubt, set it to false. "
    "Respond with a JSON object only, no other text, with these keys: "
    "is_civic_issue (true or false), "
    "category (one of: pothole, road_damage, road_blocking, footpath, "
    "streetlight, garbage, drain, waterlogging, stray_animal, dead_animal, other), "
    "severity (low, medium, high), "
    "description (one factual sentence describing what you see), "
    "visible_text (any text on signs or boards in the image, or empty string)."
)


def _headers():
    return {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json",
    }


def _encode_image(pil_image):
    buffer = io.BytesIO()
    pil_image.convert("RGB").save(buffer, format="JPEG", quality=85)
    encoded = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return f"data:image/jpeg;base64,{encoded}"


def _extract_json(text):
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if not match:
        return None
    try:
        return json.loads(match.group())
    except json.JSONDecodeError:
        return None


def classify_image(pil_image):
    payload = {
        "model": MODEL,
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": VISION_PROMPT},
                    {"type": "image_url",
                     "image_url": {"url": _encode_image(pil_image)}},
                ],
            }
        ],
        "temperature": 0.2,
    }
    resp = requests.post(API_BASE, headers=_headers(),
                         data=json.dumps(payload), timeout=60)
    resp.raise_for_status()
    content = resp.json()["choices"][0]["message"]["content"]
    parsed = _extract_json(content)

    if parsed is None:
        return {
            "is_civic_issue": False,
            "category": "other",
            "severity": "medium",
            "description": content.strip()[:200],
            "visible_text": "",
        }

    category = parsed.get("category", "other")
    if category not in VALID_CATEGORIES:
        category = "other"

    return {
        "is_civic_issue": bool(parsed.get("is_civic_issue", False)),
        "category": category,
        "severity": parsed.get("severity", "medium"),
        "description": parsed.get("description", ""),
        "visible_text": parsed.get("visible_text", ""),
    }


def _build_prompt(context, language):
    if language == "Kannada":
        lang_line = (
            "Write the entire complaint in fluent, natural Kannada (ಕನ್ನಡ) using "
            "Kannada script throughout. Keep officer names, ward names, and phone "
            "numbers in their original form. Do not write any English version."
        )
    else:
        lang_line = "Write the complaint in English."

    return (
        "Write a formal civic complaint to the Bruhat Bengaluru Mahanagara "
        "Palike (BBMP). Keep it under 120 words, polite and factual. "
        f"{lang_line}\n"
        "Use these details:\n"
        f"Issue: {context['category_label']}\n"
        f"Severity: {context['severity']}\n"
        f"Observation: {context['description']}\n"
        f"Ward: {context['ward_name']} (Ward {context['ward_no']}), "
        f"{context['zone']} zone\n"
        f"Addressed to: {context['officer_role']} {context['officer_name']}\n"
        "Start with a subject line. Request a clear timeline for resolution. "
        "Do not invent names, dates, or reference numbers. "
        "Return only the complaint text."
    )


def _translate_to_kannada(text):
    payload = {
        "input": text,
        "source_language_code": "en-IN",
        "target_language_code": "kn-IN",
        "speaker_gender": "Male",
        "mode": "formal",
        "model": "sarvam-translate:v1",
    }
    resp = requests.post(
        SARVAM_TRANSLATE,
        headers={
            "api-subscription-key": SARVAM_KEY,
            "Content-Type": "application/json",
        },
        data=json.dumps(payload),
        timeout=60,
    )
    resp.raise_for_status()
    return resp.json().get("translated_text", "").strip()


def draft_complaint(context, language="English"):
    prompt = _build_prompt(context, "English")

    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.4,
    }
    resp = requests.post(API_BASE, headers=_headers(),
                         data=json.dumps(payload), timeout=90)
    resp.raise_for_status()
    english = (resp.json()["choices"][0]["message"].get("content") or "").strip()

    if not english:
        return "Could not generate the complaint. Please try again."

    if language == "Kannada":
        kannada = _translate_to_kannada(english)
        return kannada if kannada else english

    return english