pdf_title_image/celery_tasks/LLMAgent.py

# -*- coding: utf-8 -*-
# @Author: privacy
# @Date:   2024-06-11 13:43:14
# @Last Modified by:   privacy
# @Last Modified time: 2024-12-23 18:17:53
import re
import json
from enum import Enum
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Union

import instructor
from openai import OpenAI
from pydantic import BaseModel


class BaseLlmConfig(ABC):
    def __init__(
        self,
        model: Optional[str] = None,
        base_url: Optional[str] = None,
        temperature: float = 0.0,
        max_tokens: int = 3000,
        top_p: float = 1.0
    ):
        self.model = model
        self.base_url = base_url
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.top_p = top_p


class LLMBase(ABC):
    def __init__(self, config: Optional[BaseLlmConfig] = None):
        """Initialize a base LLM class

        :param config: LLM configuration option class, defaults to None
        :type config: Optional[BaseLlmConfig], optional
        """
        if config is None:
            self.config = BaseLlmConfig()
        else:
            self.config = config

    @abstractmethod
    def generate_response(self, messages: List[dict]):
        """
        Generate a response based on the given messages.

        Args:
            messages (list): List of message dicts containing 'role' and 'content'.

        Returns:
            str: The generated response.
        """
        pass


class LLMAgent(LLMBase):
    def __init__(self, config: Optional[BaseLlmConfig] = None):
        super().__init__(config)

        if not self.config.model:
            self.config.model = "gpt-4o"

        self.client = OpenAI(
            base_url=self.config.base_url,
            api_key='ollama'
        )

    def _parse_response(self, response, tools):
        """
        Process the response based on whether tools are used or not.

        Args:
            response: The raw response from API.
            tools: The list of tools provided in the request.

        Returns:
            str or dict: The processed response.
        """
        if tools:
            processed_response = {
                "content": response.choices[0].message.content,
                "tool_calls": []
            }

            if response.choices[0].message.tool_calls:
                for tool_call in response.choices[0].message.tool_calls:
                    processed_response["tool_calls"].append({
                        "name": tool_call.function.name,
                        "arguments": json.loads(tool_call.function.arguments)
                    })

            return processed_response
        else:
            return response.choices[0].message.content

    def generate_response(
        self,
        messages: List[Dict[str, str]],
        response_format=None,
        tools: Optional[List[Dict]] = None,
        tool_choice: str = "auto",
    ):
        """
        Generate a response based on the given messages using OpenAI.

        Args:
            messages (list): List of message dicts containing 'role' and 'content'.
            response_format (str or object, optional): Format of the response. Defaults to "text".
            tools (list, optional): List of tools that the model can call. Defaults to None.
            tool_choice (str, optional): Tool choice method. Defaults to "auto".

        Returns:
            str: The generated response.
        """
        params = {
            "model": self.config.model,
            "messages": messages,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "top_p": self.config.top_p
        }
        if response_format:
            params["response_format"] = response_format
        if tools:
            params["tools"] = tools
            params["tool_choice"] = tool_choice

        response = self.client.chat.completions.create(**params)
        return self._parse_response(response, tools)


def get_proj(input_json: dict, standard: str):
    # client = LLMAgent(
    #     config=BaseLlmConfig(
    #         base_url='http://180.76.147.97:11434/v1',
    #         # model='qwen2:7b',
    #         model='qwen2.5:7b',
    #         # model='sam4096/qwen2tools:latest',
    #         # model='wangshenzhi/llama3.1_8b_chinese_chat:latest',
    #         temperature=1.0,
    #         max_tokens=32 * 1024
    #     )
    # )

    class LevelEnum(str, Enum):
        A = 'A'
        B = 'B'
        C = 'C'
        D = 'D'

    class ResInfo(BaseModel):
        最终得分: int
        最终等级: LevelEnum
        评价原因: str

    client = instructor.from_openai(
        OpenAI(
            # base_url='http://180.76.147.97:11434/v1',
            base_url='https://dashscope.aliyuncs.com/compatible-mode/v1',
            # api_key='ollama'
            api_key='sk-45971a6af6d94ccd89321f8f6d370b38'
        ),
        mode=instructor.Mode.JSON,
    )

    messages = [
        {"role": "system", "content": "你是一位优秀的数据分析师"},
        {"role": "user", "content": "Q: 现在有这样一个数据 input_json: %s 数据集以JSON形式呈现，在数据集 input_json 上分析结果，%s，请一步步进行推理并得出结论。如果无法得出结论，则返回一个中等的评价。" % (input_json, standard)}
    ]
    # 最终等级A、B、C、D由高到低返回单个英文字母表示即可。

    # try:
    response = client.chat.completions.create(
        # model='qwen2.5:7b',
        model='qwen-turbo',
        # model='wangshenzhi/llama3.1_8b_chinese_chat:latest',
        response_model=ResInfo,
        messages=messages,
        max_retries=3
    )
    # except Exception as e:
    #     print("client: ", client)
    #     print("input_json: ", input_json)
    #     print("standard: ", standard)
    #     raise str(e)

    return response


if __name__ == '__main__':
    input_json = """
    [{'page_numbers': [351],
    'table': [['项目名称', '七一七所实验楼操作间改造工程项目施工合同'],
                ['项目所在地', '武汉市江夏区阳光大道717号'],
                ['发包人名称', '中国船舶重工集团公司第七一七研究所'],
                ['发包人地址', '中国船舶重工集团公司第七一七研究所'],
                ['发包人联系人及电话', '刘伟13521523315'],
                ['合同价格', '315.1995万'],
                ['开工日期', '2019.06.10'],
                ['竣工日期', '2019.06.30'],
                ['承担的工作', '在现有值班工作区修建噪音隔离工作间，主体框架结构，天花板采用加厚保温隔热隔音防火板，室内安装空调柜机1台。'],
                ['工程质量', '合格'],
                ['项目经理', '游洋'],
                ['技术负责人', '魏娟'],
                ['总监理工程师及电话', '蔡工13521523215'],
                ['项目描述', '在现有值班工作区修建噪音隔离工作间，主体框架结构，天花板采用加厚保温隔热隔音防火板，室内安装空调柜机1台。'],
                ['备注', '']]},
    {'page_numbers': [370],
    'table': [['项目名称', '南京医药湖北工业园（续建）项目高位立体库地坪施工'],
                ['项目所在地', '武汉江夏区藏龙岛科技园凤凰大道18号'],
                ['发包人名称', '南京医药湖北有限公司'],
                ['发包人地址', '南京医药园湖北分部'],
                ['发包人联系人及电话', '张工18911231255'],
                ['合同价格', '538万'],
                ['开工日期', '2020.12.18'],
                ['竣工日期', '2021.04.18'],
                ['承担的工作', '详见工程量清单'],
                ['工程质量', '合格'],
                ['项目经理', '魏娟'],
                ['技术负责人', '张宏兵'],
                ['总监理工程师及电话', '秦工13521525155'],
                ['项目描述', '详见工程量清单'],
                ['备注', '']]},
    {'page_numbers': [384],
    'table': [['项目名称', '武汉国家生物产业创新基地D1-3a国家人类遗传资源样本库无机魔石地坪施工项目'],
                ['项目所在地', '光谷三路和高新大道交汇处。武汉国家生物产业创新基地园区内'],
                ['发包人名称', '山河建设集团有限公司'],
                ['发包人地址', '武汉市光谷三路和高新大道'],
                ['发包人联系人及电话', '张华15527812012'],
                ['合同价格', '347.003万元'],
                ['开工日期', '2019.05.01'],
                ['竣工日期', '2019.07.31'],
                ['承担的工作',
                '本次工程建设内容主要为武汉国家生物产业创新基地D1-3a国家人类遗传资源样本库，建筑面积13182.28平方米，其中无机磨石地坪施工面积为10547.36平方米。'],
                ['工程质量', '合格'],
                ['项目经理', '游洋'],
                ['技术负责人', '秦晓罡'],
                ['总监理工程师及电话', '张工18125152152'],
                ['项目描述',
                '本次工程建设内容主要为武汉国家生物产业创新基地D1-3a国家人类遗传资源样本库，建筑面积13182.28平方米，其中无机磨石地坪施工面积为10547.36平方米。'],
                ['备注', '']]},
    {'page_numbers': [395],
    'table': [['项目名称', '汉江新集水电站地下室地坪整治'],
                ['项目所在地', '湖北襄阳太平店镇新集水电站'],
                ['发包人名称', '中国水利水电第四工程局有限公司'],
                ['发包人地址', '中国水利水电第四工程局有限公司武汉分部'],
                ['发包人联系人及电话', '吴华镇15231522525'],
                ['合同价格', '598.4116万元'],
                ['开工日期', '2021.12.31'],
                ['竣工日期', '2022.04.02'],
                ['承担的工作', '新集水电站地下室原有地坪为地砖，空鼓破损现象严重，线凿除原地面，改为无机磨石地坪。'],
                ['工程质量', '合格'],
                ['项目经理', '董娜'],
                ['技术负责人', '秦晓罡'],
                ['总监理工程师及电话', '袁工13213251225'],
                ['项目描述', '新集水电站地下室原有地坪为地砖，空鼓破损现象严重，线凿除原地面，改为无机磨石地坪。'],
                ['备注', '']]},
    {'page_numbers': [424],
    'table': [['项目名称', '武汉金光电子器械有限公司车间（二期）工程'],
                ['项目所在地', '武汉市江夏区大桥新区民营工业园'],
                ['发包人名称', '武汉金光电子器械有限公司'],
                ['发包人地址', '武汉市江夏区经济开发区'],
                ['发包人联系人及电话', '刘江15231522525'],
                ['合同价格', '2550万元'],
                ['开工日期', '2019.06.26'],
                ['竣工日期', '2020.03.20'],
                ['承担的工作', '详见工程量清单。'],
                ['工程质量', '合格'],
                ['项目经理', '辛正军'],
                ['技术负责人', '张宏兵'],
                ['总监理工程师及电话', '魏工18911251325'],
                ['项目描述', '详见工程量清单。'],
                ['备注', '']]},
    {'page_numbers': [455],
    'table': [['项目名称', '宇通客车VMI仓储中心扩建项目总承包工程'],
                ['项目所在地', '郑州市郑州经济技术开发区兰心西路以北、龙兴南街以东'],
                ['发包人名称', '郑州宇通客车股份有限公司'],
                ['发包人地址', '郑州市郑州市郑州经济技术开发区'],
                ['发包人联系人及电话', '王江伟15528012500'],
                ['合同价格', '4428万元'],
                ['开工日期', '2020.12.25'],
                ['竣工日期', '2021.06.20'],
                ['承担的工作', '详见工程量清单。'],
                ['工程质量', '合格'],
                ['项目经理', '秦亚飞'],
                ['技术负责人', '张宏兵'],
                ['总监理工程师及电话', '李鹏18124215525'],
                ['项目描述', '该工程施工图纸范围内的土建、钢结构、公用安装等全部工程，总建筑面积28652.02平方米。'],
                ['备注', '']]},
    {'page_numbers': [677],
    'table': [['项目名称', '/'],
                ['项目所在地', '/'],
                ['发包人名称', '/'],
                ['发包人地址', '/'],
                ['发包人电话', '/'],
                ['签约合同价', '/'],
                ['开工日期', '/'],
                ['计划竣工日期', '/'],
                ['承担的工作', '/'],
                ['工程质量', '/'],
                ['项目经理', '/'],
                ['技术负责人', '/'],
                ['总监理工程师及电话', '/'],
                ['项目描述', '/'],
                ['备注', '/']]}]
    """

    standard = """满足招标文件投标人资格要求的业绩得 50 分，每增加 1 项单项合同面积在 8000 平方及以上无机磨石地坪工程的施工业绩加 10 分，加至 100 分为止。"""

    print(get_proj(input_json=input_json, standard=standard))