AIyaodianzhushou/backend-python/app/core/llm_client.py

import base64
from typing import AsyncIterator, Optional

from openai import AsyncOpenAI

from app.core.config import get_settings


class LLMClient:
    def __init__(self):
        self._client = None

    def _ensure_client(self):
        if self._client is not None:
            return
        settings = get_settings()
        is_local = settings.qwen_base_url.startswith("http://localhost")
        api_key = "local" if is_local else settings.qwen_api_key
        if not api_key and not is_local:
            raise ValueError("QWEN_API_KEY 未设置，请在 .env 文件中填入 API Key")
        self._client = AsyncOpenAI(
            api_key=api_key,
            base_url=settings.qwen_base_url,
        )
        self.model = settings.qwen_model
        self.vl_model = settings.qwen_vl_model
        self.enable_web_search = settings.enable_web_search

    # ============================================================
    # 纯文本对话
    # ============================================================

    async def chat(
        self,
        messages: list[dict],
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        enable_search: bool = False,
    ) -> str:
        self._ensure_client()
        settings = get_settings()
        extra = {}
        if enable_search:
            extra["enable_search"] = True
        response = await self._client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=temperature or settings.qwen_temperature,
            max_tokens=max_tokens or settings.qwen_max_tokens,
            extra_body=extra if extra else None,
        )
        return response.choices[0].message.content or ""

    async def chat_stream(
        self,
        messages: list[dict],
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        enable_search: bool = False,
    ) -> AsyncIterator[str]:
        self._ensure_client()
        settings = get_settings()
        extra = {}
        if enable_search:
            extra["enable_search"] = True
        stream = await self._client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=temperature or settings.qwen_temperature,
            max_tokens=max_tokens or settings.qwen_max_tokens,
            stream=True,
            extra_body=extra if extra else None,
        )
        async for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content

    # ============================================================
    # 图片分析 + OCR（Qwen VL 视觉模型）
    # ============================================================

    async def analyze_image(
        self,
        image_base64: str,
        mime_type: str = "image/jpeg",
        prompt: str = "",
    ) -> str:
        """
        调用 Qwen VL 模型分析图片：
        - 描述图片内容（药品包装、说明书、处方等）
        - 提取图中文字（OCR）
        - 返回可用于 RAG 检索的文本
        """
        self._ensure_client()

        if not prompt:
            prompt = """请仔细分析这张图片，完成以下任务：
1. 描述图片内容（药品包装、说明书、处方单、症状照片等）
2. 提取图片中所有可见的文字（OCR），特别是药品名称、成分、用法用量、批号、有效期等关键信息
3. 如果图片是药品包装/说明书，提取：药品通用名、规格、生产企业、批准文号
4. 如果图片是处方，提取：患者信息、药品名称、用法用量、开具日期

请按以下格式输出：

【图片描述】
图片内容的简要描述。

【OCR 提取文字】
逐条列出提取到的文字内容。
"""

        response = await self._client.chat.completions.create(
            model=self.vl_model,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{mime_type};base64,{image_base64}"
                            },
                        },
                        {"type": "text", "text": prompt},
                    ],
                }
            ],
            max_tokens=2048,
        )
        return response.choices[0].message.content or ""

    async def analyze_image_stream(
        self,
        image_base64: str,
        mime_type: str = "image/jpeg",
        prompt: str = "",
    ) -> AsyncIterator[str]:
        """流式版本：Qwen VL 分析图片"""
        self._ensure_client()

        if not prompt:
            prompt = """请仔细分析这张图片，完成以下任务：
1. 描述图片内容（药品包装、说明书、处方单、症状照片等）
2. 提取图片中所有可见的文字（OCR），特别是药品名称、成分、用法用量等关键信息
3. 如果图片是药品包装/说明书，提取：药品通用名、规格、生产企业、批准文号
请简要输出分析结果。"""

        stream = await self._client.chat.completions.create(
            model=self.vl_model,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{mime_type};base64,{image_base64}"
                            },
                        },
                        {"type": "text", "text": prompt},
                    ],
                }
            ],
            max_tokens=2048,
            stream=True,
        )
        async for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content

    # ============================================================
    # 视频分析 + OCR（Qwen VL 模型支持视频帧提取分析）
    # ============================================================

    async def analyze_video(
        self,
        video_base64: str,
        mime_type: str = "video/mp4",
        prompt: str = "",
    ) -> str:
        """
        调用 Qwen VL 模型分析视频：
        - 提取关键帧并描述视频内容
        - OCR 提取帧中所有可见文字
        - 返回可用于 RAG 检索的文本
        """
        self._ensure_client()

        if not prompt:
            prompt = """请仔细分析这段视频，完成以下任务：
1. 描述视频内容（药品展示、用药指导、症状表现等）
2. 提取视频中所有可见的文字（OCR），特别是药品名称、成分、用法用量、批号等关键信息
3. 如果视频中有药品包装/说明书，提取全部文字信息
4. 总结视频传达的关键药学信息

请按以下格式输出：

【视频描述】
视频内容的简要描述。

【OCR 提取文字】
逐条列出从视频帧中提取到的文字内容。

【关键信息总结】
与药学相关的关键信息摘要。"""

        response = await self._client.chat.completions.create(
            model=self.vl_model,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "video_url",
                            "video_url": {
                                "url": f"data:{mime_type};base64,{video_base64}"
                            },
                        },
                        {"type": "text", "text": prompt},
                    ],
                }
            ],
            max_tokens=2048,
        )
        return response.choices[0].message.content or ""

    async def analyze_video_stream(
        self,
        video_base64: str,
        mime_type: str = "video/mp4",
        prompt: str = "",
    ) -> AsyncIterator[str]:
        """流式版本：Qwen VL 分析视频"""
        self._ensure_client()

        if not prompt:
            prompt = """请分析这段视频，提取关键帧中的文字信息（OCR），特别是药品名称、成分、用法用量等。简要输出。"""

        stream = await self._client.chat.completions.create(
            model=self.vl_model,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "video_url",
                            "video_url": {
                                "url": f"data:{mime_type};base64,{video_base64}"
                            },
                        },
                        {"type": "text", "text": prompt},
                    ],
                }
            ],
            max_tokens=2048,
            stream=True,
        )
        async for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content

    # ============================================================
    # 统一多模态分析（自动判断 image / video）
    # ============================================================

    async def analyze_media(
        self,
        media_base64: str,
        media_type: str,  # "image" or "video"
        mime_type: str = "",
        prompt: str = "",
    ) -> str:
        """统一入口：根据 media_type 自动路由到图片或视频分析"""
        if media_type == "video":
            return await self.analyze_video(
                media_base64,
                mime_type=mime_type or "video/mp4",
                prompt=prompt,
            )
        else:
            return await self.analyze_image(
                media_base64,
                mime_type=mime_type or "image/jpeg",
                prompt=prompt,
            )

    async def analyze_media_stream(
        self,
        media_base64: str,
        media_type: str,
        mime_type: str = "",
        prompt: str = "",
    ) -> AsyncIterator[str]:
        """流式统一入口"""
        if media_type == "video":
            async for token in self.analyze_video_stream(
                media_base64,
                mime_type=mime_type or "video/mp4",
                prompt=prompt,
            ):
                yield token
        else:
            async for token in self.analyze_image_stream(
                media_base64,
                mime_type=mime_type or "image/jpeg",
                prompt=prompt,
            ):
                yield token


llm_client = LLMClient()