CareerBot/app/services/llm_service.py

import json
from typing import AsyncGenerator

import httpx
from sqlalchemy.orm import Session

from app.models import LLMConfig


def get_active_config(db: Session) -> LLMConfig | None:
    return db.query(LLMConfig).filter(LLMConfig.is_active == True).first()


def _build_url_and_headers(config: LLMConfig) -> tuple[str, dict]:
    url = config.api_url.rstrip("/")
    if not url.endswith("/v1/chat/completions"):
        url += "/v1/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {config.api_key}",
    }
    return url, headers


def _build_body(config: LLMConfig, messages: list[dict], stream: bool = False) -> dict:
    body = {
        "model": config.model_name,
        "messages": messages,
        "max_tokens": config.max_tokens,
        "temperature": config.temperature,
    }
    if stream:
        body["stream"] = True
    return body


def _has_image_content(messages: list[dict]) -> bool:
    for msg in messages:
        content = msg.get("content")
        if isinstance(content, list):
            for part in content:
                if isinstance(part, dict) and part.get("type") == "image_url":
                    return True
    return False


def _strip_images(messages: list[dict]) -> list[dict]:
    """Remove image content, keep text, add a note about the limitation."""
    result = []
    for msg in messages:
        content = msg.get("content")
        if isinstance(content, list):
            text_parts = [p["text"] for p in content if isinstance(p, dict) and p.get("type") == "text"]
            text = "\n".join(text_parts) if text_parts else ""
            text += "\n\n[系统提示：用户上传了一张图片作为职位描述(JD)，但当前配置的模型不支持图片识别。请告知用户改为上传文本(.txt)、Word(.docx)或PDF(.pdf)格式的JD文件，以便进行详细的匹配度分析。]"
            result.append({"role": msg["role"], "content": text})
        else:
            result.append(msg)
    return result


async def _do_stream(client: httpx.AsyncClient, url: str, headers: dict, body: dict):
    """Execute a streaming request and yield content chunks."""
    async with client.stream("POST", url, json=body, headers=headers) as resp:
        resp.raise_for_status()
        async for line in resp.aiter_lines():
            if not line.startswith("data: "):
                continue
            data_str = line[6:].strip()
            if data_str == "[DONE]":
                break
            try:
                chunk = json.loads(data_str)
                content = chunk["choices"][0]["delta"].get("content", "")
                if content:
                    yield content
            except (json.JSONDecodeError, KeyError, IndexError):
                continue


async def chat_completion(messages: list[dict], db: Session) -> str:
    """Non-streaming chat completion with auto vision fallback."""
    config = get_active_config(db)
    if not config:
        return "AI助手尚未配置，请联系管理员。"

    url, headers = _build_url_and_headers(config)
    body = _build_body(config, messages)

    try:
        async with httpx.AsyncClient(timeout=60.0) as client:
            resp = await client.post(url, json=body, headers=headers)
            resp.raise_for_status()
            return resp.json()["choices"][0]["message"]["content"]
    except httpx.HTTPStatusError as e:
        if e.response.status_code == 400 and _has_image_content(messages):
            return await _fallback_no_image(config, messages)
        return f"AI服务调用失败: {str(e)}"
    except Exception as e:
        return f"AI服务调用失败: {str(e)}"


async def _fallback_no_image(config: LLMConfig, messages: list[dict]) -> str:
    url, headers = _build_url_and_headers(config)
    body = _build_body(config, _strip_images(messages))
    try:
        async with httpx.AsyncClient(timeout=60.0) as client:
            resp = await client.post(url, json=body, headers=headers)
            resp.raise_for_status()
            return resp.json()["choices"][0]["message"]["content"]
    except Exception as e:
        return f"AI服务调用失败: {str(e)}"


async def chat_completion_stream(
    messages: list[dict], db: Session
) -> AsyncGenerator[str, None]:
    """Streaming chat completion with auto vision fallback."""
    config = get_active_config(db)
    if not config:
        yield "AI助手尚未配置，请联系管理员。"
        return

    url, headers = _build_url_and_headers(config)
    body = _build_body(config, messages, stream=True)

    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            async for chunk in _do_stream(client, url, headers, body):
                yield chunk
            return
    except httpx.HTTPStatusError as e:
        if e.response.status_code == 400 and _has_image_content(messages):
            pass  # Fall through to retry without images
        else:
            yield f"\n\nAI服务调用失败: {str(e)}"
            return
    except Exception as e:
        if _has_image_content(messages):
            pass  # Fall through to retry
        else:
            yield f"\n\nAI服务调用失败: {str(e)}"
            return

    # Fallback: retry without image content
    fallback_body = _build_body(config, _strip_images(messages), stream=True)
    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            async for chunk in _do_stream(client, url, headers, fallback_body):
                yield chunk
    except Exception as e:
        yield f"\n\nAI服务调用失败: {str(e)}"