llm_utils.py

import os
import json
import httpx
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

# ------------------------------------------------------------------
# Configuration
# ------------------------------------------------------------------
SELLER_ASSISTANT_URL = os.getenv("SELLER_ASSISTANT_URL", "")
# If empty, the external call is skipped.

# ------------------------------------------------------------------
# Local lightweight model (DistilGPT-2) – used only if no Seller Assistant
# ------------------------------------------------------------------
_local_pipeline = None

def get_local_llm():
    """Lazy load a small text generation model (DistilGPT-2, ~300MB)."""
    global _local_pipeline
    if _local_pipeline is None:
        print("Loading local DistilGPT-2 for fallback description generation...")
        _local_pipeline = pipeline(
            "text-generation",
            model="distilgpt2",
            device=-1,  # CPU
            max_new_tokens=80,
            do_sample=True,
            temperature=0.7
        )
    return _local_pipeline

# ------------------------------------------------------------------
# External call to Seller Assistant Service
# ------------------------------------------------------------------
async def call_seller_assistant(prompt: str, api_key: str = None) -> str:
    """Send prompt to the Seller Assistant (Falcon/Phi‑3) and return reply."""
    if not SELLER_ASSISTANT_URL:
        return None
    headers = {}
    if api_key:
        headers["X-API-Key"] = api_key
    async with httpx.AsyncClient(timeout=30.0) as client:
        try:
            resp = await client.post(
                f"{SELLER_ASSISTANT_URL}/api/chat",
                json={"message": prompt},
                headers=headers
            )
            if resp.status_code == 200:
                data = resp.json()
                return data.get("reply", "")
        except Exception as e:
            print(f"Seller Assistant call failed: {e}")
    return None

# ------------------------------------------------------------------
# Generate SEO title + description (main entry point)
# ------------------------------------------------------------------
async def generate_description(category: str, material: str, colors: list, api_key: str = None) -> dict:
    """
    Returns a dictionary with 'title' and 'description'.
    Tries external assistant first, then local model, then template.
    """
    color_list = ", ".join(colors[:3])
    prompt = (
        f"Write an SEO product title (max 60 chars) and description (2 sentences) "
        f"for a {material} {category} in colors {color_list}. "
        f"Return as JSON with keys 'title' and 'description'."
    )

    # 1) Try external Seller Assistant (best quality)
    if SELLER_ASSISTANT_URL:
        reply = await call_seller_assistant(prompt, api_key)
        if reply:
            try:
                # Extract JSON from reply (the assistant may output extra text)
                start = reply.find("{")
                end = reply.rfind("}") + 1
                if start != -1 and end > start:
                    json_str = reply[start:end]
                    data = json.loads(json_str)
                    if "title" in data and "description" in data:
                        return data
            except:
                pass

    # 2) Fallback to local DistilGPT-2 (lower quality but works offline)
    try:
        local_llm = get_local_llm()
        result = local_llm(prompt)[0]["generated_text"]
        # Extract after prompt (crude)
        if result.startswith(prompt):
            result = result[len(prompt):].strip()
        # Attempt to parse JSON from result
        start = result.find("{")
        end = result.rfind("}") + 1
        if start != -1 and end > start:
            data = json.loads(result[start:end])
            return data
        # If not JSON, fallback to template
    except Exception as e:
        print(f"Local LLM failed: {e}")

    # 3) Template fallback
    return {
        "title": f"Premium {material} {category} – {color_list}",
        "description": f"High-quality {material} {category} available in {color_list}. Perfect for your needs."
    }