目录

http 流式推理

直接返回

http 流式推理

ocr_stream_client.py

# coding=utf-8
import base64
import json
import os
import urllib.request
import logging
import logging.handlers
from pathlib import Path
from typing import Optional, List, Dict, Any, AsyncGenerator

import requests

LOG_DIR = Path(__file__).resolve().parent / "logs"
LOG_DIR.mkdir(exist_ok=True)

base_url = "http://192.168.100.203:17890"   # 统一使用第二个地址

# 主日志处理器（geogebra_api.log）
main_handler = logging.handlers.TimedRotatingFileHandler(
    filename=LOG_DIR / "geogebra_api.log",
    when="midnight",
    interval=1,
    backupCount=30,
    encoding="utf-8"
)
main_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
stream_handler = logging.StreamHandler()
stream_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))

logging.basicConfig(level=logging.INFO,
                    handlers=[stream_handler, main_handler])
logger = logging.getLogger(__name__)

_MODEL_ID = None

def get_model_id() -> str:
    """获取 vLLM 服务的模型 ID，并缓存"""
    global _MODEL_ID
    if _MODEL_ID is not None:
        return _MODEL_ID
    try:
        resp = urllib.request.urlopen(f"{base_url}/v1/models", timeout=30)
        data = json.loads(resp.read().decode("utf-8"))
        _MODEL_ID = data["data"][0]["id"]
        logger.info(f"获取到模型 ID: {_MODEL_ID}")
        return _MODEL_ID
    except Exception as e:
        logger.error(f"获取模型 ID 失败: {e}")
        # raise HTTPException(status_code=500, detail="无法获取 vLLM 模型 ID")


model_id = get_model_id()

def _http_stream(url: str, payload: dict, timeout: int = 300):
    """流式请求vLLM接口，逐行返回SSE数据"""
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        for line in resp:
            line = line.decode("utf-8").strip()
            if not line:
                continue
            if line.startswith("data: "):
                data_str = line[6:]  # 去掉 "data: " 前缀
                if data_str == "[DONE]":
                    break
                try:
                    chunk = json.loads(data_str)
                    yield chunk
                except json.JSONDecodeError:
                    continue

prompt_file = "ocr_prompt.md"

with open(prompt_file, "r", encoding="utf-8") as f:
    system_prompt = f.read()  # 一行搞定，包含所有换行

def encode_image_to_base64(image_path: str) -> str:
    """将图片编码为 base64"""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def ocr_image_processing(img_path: str, model_id: str, base_url: str, session_id: str, rep_id: str, timeout: int = 300) -> AsyncGenerator[str, None]:
    try:
        # 读取图片数据
        with open(img_path, 'rb') as f:
            image_data = f.read()

        # 获取图片文件名
        image_filename = os.path.basename(img_path)
        json_log_path = image_filename[:-4] + ".json"

        # 构建消息
        ocr_messages = [{"role": "user", "content": [{"type": "text", "text": "请识别图片中的文字和数学公式"},
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64.b64encode(image_data).decode('utf-8')}"}}]}]

        full_ocr_messages = [{"role": "system", "content": system_prompt}] + ocr_messages

        # 构建请求负载
        ocr_payload = {"model": model_id, "messages": full_ocr_messages, "temperature": 0.1, "max_tokens": 2048, "stream": True}

        # 直接使用 requests 发送请求
        response = requests.post(f"{base_url}/v1/chat/completions", json=ocr_payload, stream=True, timeout=timeout)

        if response.status_code != 200:
            error_msg = f"HTTP错误: {response.status_code}"
            logger.error(error_msg)
            yield f"data: {json.dumps({'type': 'ocr_error', 'error': error_msg, 'session_id': session_id, 'rep_id': rep_id},ensure_ascii=False)}\n\n"
            return

        ocr_full_content = ""
        ocr_success = False

        # 处理流式响应
        for line in response.iter_lines():
            if not line:
                continue

            line = line.decode('utf-8').strip()
            if not line.startswith("data: "):
                continue

            data_str = line[6:]  # 去掉 "data: " 前缀
            if data_str == "[DONE]":
                break

            try:
                chunk = json.loads(data_str)
                choices = chunk.get("choices", [])
                if not choices:
                    continue

                delta = choices[0].get("delta", {})
                content = delta.get("content", "")

                if content:
                    ocr_full_content += content
                    print(content)
                    yield f"data: {json.dumps({'type': 'ocr_chunk', 'content': content, 'session_id': session_id, 'rep_id': rep_id},ensure_ascii=False)}\n\n"
            except json.JSONDecodeError:
                continue

        ocr_text = ocr_full_content.strip()

        if ocr_text:
            ocr_success = True
            logger.info(f"OCR识别成功，长度: {len(ocr_text)}")

            # 保存记录
            record = {"session_id": session_id, "image_filename": image_filename, "ocr_text": ocr_text, "success": ocr_success}

            with open(json_log_path, "w", encoding="utf-8") as f:
                json.dump(record, f, ensure_ascii=False, indent=2)

            logger.info(f"OCR记录已保存: {json_log_path}")
            logger.info(f"OCR识别文本内容:\n{ocr_text}")

            yield f"data: {json.dumps({'type': 'ocr_done', 'content': ocr_text, 'session_id': session_id, 'rep_id': rep_id},ensure_ascii=False)}\n\n"
        else:
            logger.warning("OCR识别结果为空")
            yield f"data: {json.dumps({'type': 'ocr_error', 'error': 'OCR识别结果为空', 'session_id': session_id, 'rep_id': rep_id},ensure_ascii=False)}\n\n"

    except FileNotFoundError:
        error_msg = f"图片文件不存在: {img_path}"
        logger.error(error_msg)
        yield f"data: {json.dumps({'type': 'ocr_error', 'error': error_msg, 'session_id': session_id, 'rep_id': rep_id},ensure_ascii=False)}\n\n"
    except Exception as e:
        error_msg = f"OCR处理异常: {str(e)}"
        logger.error(error_msg)
        yield f"data: {json.dumps({'type': 'ocr_error', 'error': error_msg, 'session_id': session_id, 'rep_id': rep_id},ensure_ascii=False)}\n\n"

if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    # 配置参数
    img_path = r"C:\Users\ChanJing-01\Documents\math\07\hanshu.png"
    session_id = "session_id"
    rep_id = "rep_id"

    # 模拟的http_stream_function（实际使用时需要替换为真实的请求函数）
    def mock_http_stream_function(url, payload, timeout):
        # 模拟返回数据
        yield {"choices": [{"delta": {"content": "这是一个测试"}}]}
        yield {"choices": [{"delta": {"content": " OCR识别结果"}}]}

    # 调用函数
    # for chunk in ocr_image_processing(img_path=img_path, model_id=model_id, base_url=base_url, session_id=session_id, rep_id=rep_id, http_stream_function=mock_http_stream_function):
    for chunk in ocr_image_processing(img_path=img_path, model_id=model_id, base_url=base_url, session_id=session_id, rep_id=rep_id):
        print(chunk)

直接返回

ocr_from_server.py

import urllib
from http.client import HTTPException
from pathlib import Path

import requests
import json
import base64
from typing import List, Dict, Any
import time

def encode_image_to_base64(image_path: str) -> str:
    """将图片编码为 base64"""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

_MODEL_ID = None

def get_model_id() -> str:
    """获取 vLLM 服务的模型 ID，并缓存"""
    global _MODEL_ID
    if _MODEL_ID is not None:
        return _MODEL_ID
    try:
        resp = urllib.request.urlopen(f"{BASE_URL}/v1/models", timeout=30)
        data = json.loads(resp.read().decode("utf-8"))
        _MODEL_ID = data["data"][0]["id"]
        print(f"获取到模型 ID: {_MODEL_ID}")
        return _MODEL_ID
    except Exception as e:
        print(f"获取模型 ID 失败: {e}")
        raise HTTPException(status_code=500, detail="无法获取 vLLM 模型 ID")

def chat_with_images(text: str, image_paths: List[str], max_tokens: int = 1024, temperature: float = 0.7) -> str:

    model_id = get_model_id()
    # 构建 content 列表
    content = []

    # 1. 添加文本内容
    content.append({"type": "text", "text": text})

    for image_path in image_paths:
        # 如果是本地图片，转换为 base64
        if image_path.startswith(("http://", "https://")):
            # URL 图片
            content.append({"type": "image_url", "image_url": {"url": image_path}})
        else:
            # 本地图片
            base64_image = encode_image_to_base64(image_path)
            content.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}})

    prompt_file = "ocr_prompt.md"

    with open(prompt_file, "r", encoding="utf-8") as f:
        system_prompt = f.read()  # 一行搞定，包含所有换行

    payload = {
        "model": model_id,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": content},
        ],
        "temperature": 0.0,
        "max_tokens": 4096, "stream": True,
    }

    print(f"📤 Sending request with {len(image_paths)} images...")
    print(f"📝 Text: {text}")
    print("-" * 80)

    full_response = ""
    try:
        response = requests.post(f"{BASE_URL}/v1/chat/completions", json=payload, stream=True, timeout=300)

        if response.status_code != 200:
            print(f"❌ Error: {response.status_code}")
            print(response.text)
            return ""

        # 5. 处理流式响应
        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith("data: "):
                    data = line[6:]  # 去掉 "data: " 前缀
                    if data == "[DONE]":
                        break
                    try:
                        chunk = json.loads(data)
                        choices = chunk.get("choices", [])
                        if choices:
                            delta = choices[0].get("delta", {})
                            content_delta = delta.get("content", "")
                            if content_delta:
                                print(content_delta, end="", flush=True)
                                full_response += content_delta
                    except json.JSONDecodeError:
                        continue

        print("\n" + "-" * 80)
        return full_response

    except requests.exceptions.Timeout:
        print("❌ Request timeout")
        return ""
    except Exception as e:
        print(f"❌ Error: {e}")
        return ""


if __name__ == "__main__":
    BASE_URL = "http://192.168.100.203:17890"
    image_paths = [r"C:\Users\ChanJing-01\Documents\math\07\a12.jpg"]
    image_paths = [r"C:\Users\ChanJing-01\Documents\math\07\hanshu.png"]
    # response = chat_with_images(text="图中有两个矩形，把几何图复现一下", image_paths=[r"C:\Users\ChanJing-01\Pictures\wav_sucai\ca1e63ed-1560-4ade-9e45-57aa500b7123.png"])
    response = chat_with_images(text="请识别图片中的文字和数学公式", image_paths=image_paths)
    print(response)