"""
测试提示词脚本
使用requests库测试语法审查API
支持本地模型和在线模型两种调用方式
"""
import requests
import json
from openai import OpenAI


class PromptTester:
    """提示词测试类"""
    
    def __init__(self, api_url, api_key=None, is_local=True):
        """
        初始化测试器
        
        Args:
            api_url: API地址
            api_key: API密钥（在线模型需要）
            is_local: 是否使用本地模型（True: 本地模型使用requests, False: 在线模型使用OpenAI SDK）
        """
        self.api_url = api_url
        self.api_key = api_key
        self.is_local = is_local
        
        # 如果是在线模型，初始化OpenAI客户端
        if not is_local and api_key:
            self.client = OpenAI(
                base_url=api_url.replace('/chat/completions', ''),  # 移除末尾的路径
                api_key=api_key
            )
        else:
            self.client = None
        self.system_prompt = """
system
# role
你是词句语法审查专家

## workflow
- 负责检查文本中的错别字和重复字词等语法问题。
- 检查待审查的词句语法标点格式是否符合规则。
- 检查错别字、多字、少字、重复字词等语法错误。
- 如果发现一个句子有错误，请正确的定位其错误位置，而不是想当然的去找问题。如句子末尾未加句号不要定位到句中去了。
- 给出了最终建议后，需要再次考虑如果按照建议去做了，原文是否通顺合理。
- 请减少对标点符号的严格审查，如果有问题请三思你所说的位置有没有符号，符号对不对。
- 仅检查句尾的标点问题，不要检查句中标点问题。
- 请**不要为文章新增标点符号**如句号、逗号，只需要审查标点符号是否正确即可。

## example
1. 出现了明显错字，出现了的地得使用混乱。
2. 如果出现了错字，大概率是因为拼音拼写错误导致的错字，请根据其相似读音推测正确的字，如果拿不准请不给出修改为xx的建议。
3. 如“供气”错打为了“供器”，应当结合上下文推测“供qi”应该为什么。
4. 对于条款编号而言，'一）'这样的结构是正确的，符合中文规范

## rule
- 不需要强求要输出问题，除非是非常明显的错误。小问题可以忽略。
- 仅检查句尾的标点问题，不要检查句中标点问题。
- 不要为长句切分、添加分句标点。
- 请着重对错别字的审查上，大量减少对标点符号的审查力度。
- 务必遵循<example>中的规则。
- 如果没有错误请不要添加新的issue。
- 遵循中文的语法规范。
- 遵循以下**强制**规范：
    1. 务必结合语境进行分析检查。
    2. 对于**表格制表符**、**表格内容**不需要检查。
    3. 对于术语概念不得曲解。
    4. 没有明显词句语法错误、标点错误的内容不予检查，输出无明显问题。
    5. 已检查出的问题项仅输出一次检查结果，禁止对同一内容重复检查。
    6. 统一解释：如果表格中出现了多列相同的表头标题，不是错误，而是解析时这几个是合并的表头。

## output
```json
    {{
        "issue_point": "问题标题描述",
        "location": "当前问题对应的原始条款内容及位置,如六、验收标准 (页码: 85)，以及其语境上下文",
        "suggestion": "具体的修改建议内容",
        "reason": "问题的原因分析和依据说明",
        "risk_level": ""
    }}
```
# 风险等级分类risk_level:
- 高风险:影响施工安全、可能导致严重后果的词句错误。
- 中风险:影响理解表达、可能导致一般性问题的词句错误。
- 低风险:形式问题、不影响实质内容的词句错误。
"""
        self.user_prompt_template = """
请审查以下内容的词句语法错误，重点关注错别字、重复字词和语法结构：
【待检查文本】
{review_content}
输出格式:务必须严格按照以下标准json格式输出审查结果:
如果未发现明显的词句语法错误，请输出:无明显问题。
如果发现问题，请按以下格式输出:
location字段直接输出原字段内容，不得猜测。
## 示例
```json
{{
  "issue_point": "问题标题描述",
  "location": "当前问题对应的原始条款内容及位置,如六、验收标准 (页码: 85)，以及其语境上下文",
  "suggestion": "具体的修改建议内容",
  "reason": "问题的原因分析和依据说明",
  "risk_level": ""
}}
```
/no_think
"""

    def test_prompt(self, review_content, model="Qwen3-8B", temperature=0.3, max_tokens=2000):
        """
        测试提示词
        
        Args:
            review_content: 待审查的文本内容
            model: 模型名称
            temperature: 温度参数
            max_tokens: 最大token数
            
        Returns:
            API响应结果
        """
        # 填充用户提示词
        user_prompt = self.user_prompt_template.format(review_content=review_content)
        
        # 打印请求信息
        print("=" * 80)
        print(f"调用方式: {'本地模型 (requests)' if self.is_local else '在线模型 (OpenAI SDK)'}")
        print("发送请求到:", self.api_url)
        print("=" * 80)
        print("System Prompt:")
        print("-" * 80)
        print(self.system_prompt)
        print("=" * 80)
        print("User Prompt:")
        print("-" * 80)
        print(user_prompt)
        print("=" * 80)
        
        # 根据is_local选择不同的调用方式
        if self.is_local:
            # 本地模型：使用requests
            return self._call_with_requests(user_prompt, model, temperature, max_tokens)
        else:
            # 在线模型：使用OpenAI SDK
            return self._call_with_openai(user_prompt, model, temperature, max_tokens)
    
    def _call_with_requests(self, user_prompt, model, temperature, max_tokens):
        """使用requests库调用本地模型"""
        # 构建请求数据
        payload = {
            "model": model,
            "messages": [
                {
                    "role": "system",
                    "content": self.system_prompt
                },
                {
                    "role": "user",
                    "content": user_prompt
                }
            ],
            "temperature": temperature,
            "max_tokens": max_tokens
        }
        
        try:
            # 构建请求头
            headers = {"Content-Type": "application/json"}
            if self.api_key:
                headers["Authorization"] = f"Bearer {self.api_key}"
            
            # 发送POST请求
            response = requests.post(
                self.api_url,
                headers=headers,
                json=payload,
                timeout=60
            )
            
            # 检查响应状态
            response.raise_for_status()
            
            # 解析响应
            result = response.json()
            
            print("响应状态码:", response.status_code)
            print("=" * 80)
            print("API响应结果:")
            print("-" * 80)
            print(json.dumps(result, ensure_ascii=False, indent=2))
            print("=" * 80)
            
            # 提取AI回复内容
            if "choices" in result and len(result["choices"]) > 0:
                ai_response = result["choices"][0].get("message", {}).get("content", "")
                print("AI回复内容:")
                print("-" * 80)
                print(ai_response)
                print("=" * 80)
            
            return result
            
        except requests.exceptions.RequestException as e:
            print(f"请求失败: {e}")
            return None
        except json.JSONDecodeError as e:
            print(f"JSON解析失败: {e}")
            print("原始响应:", response.text)
            return None
    
    def _call_with_openai(self, user_prompt, model, temperature, max_tokens):
        """使用OpenAI SDK调用在线模型"""
        if not self.client:
            print("错误: OpenAI客户端未初始化")
            return None
        
        try:
            # 使用OpenAI SDK调用
            extra_body = {
                # enable thinking, set to False to disable test
                "enable_thinking": False,
                # use thinking_budget to contorl num of tokens used for thinking
                # "thinking_budget": 4096
            }
            response = self.client.chat.completions.create(
                model=model,
                messages=[
                    {
                        "role": "system",
                        "content": self.system_prompt
                    },
                    {
                        "role": "user",
                        "content": user_prompt
                    }
                ],
                temperature=temperature,
                max_tokens=max_tokens,
                extra_body=extra_body
            )
            
            # 转换为字典格式
            result = {
                "id": response.id,
                "object": response.object,
                "created": response.created,
                "model": response.model,
                "choices": [
                    {
                        "index": choice.index,
                        "message": {
                            "role": choice.message.role,
                            "content": choice.message.content
                        },
                        "finish_reason": choice.finish_reason
                    }
                    for choice in response.choices
                ],
                "usage": {
                    "prompt_tokens": response.usage.prompt_tokens,
                    "completion_tokens": response.usage.completion_tokens,
                    "total_tokens": response.usage.total_tokens
                }
            }
            
            print("=" * 80)
            print("API响应结果:")
            print("-" * 80)
            print(json.dumps(result, ensure_ascii=False, indent=2))
            print("=" * 80)
            
            # 提取AI回复内容
            if result["choices"] and len(result["choices"]) > 0:
                ai_response = result["choices"][0]["message"]["content"]
                print("AI回复内容:")
                print("-" * 80)
                print(ai_response)
                print("=" * 80)
            
            return result
            
        except Exception as e:
            print(f"OpenAI SDK调用失败: {e}")
            return None


def main():
    """主函数"""
    
    # ==================== 配置区域 ====================
    # 是否使用本地模型（True: 本地模型, False: 在线模型）
    is_local = True
    
    # 本地模型配置
    local_config = {
        "api_url": "http://192.168.91.253:9002/v1/chat/completions",
        "api_key": None,  # 本地模型不需要API Key
        "model": "Qwen3-8B",
        "temperature": 0.3,
        "max_tokens": 2000
    }
    
    # 在线模型配置（ModelScope）
    online_config = {
        "api_url": "https://api-inference.modelscope.cn/v1",
        "api_key": "ms-c0349a0a-8f15-466b-96be-4f96d001d8f2",  # ModelScope Token
        "model": "Qwen/Qwen3-14B",
        "temperature": 0.7,
        "max_tokens": 2000
    }
    # ==================== 配置区域结束 ====================
    
    # 根据is_local选择配置
    config = local_config if is_local else online_config
    
    print("=" * 80)
    print(f"当前使用: {'本地模型' if is_local else '在线模型'}")
    print(f"API地址: {config['api_url']}")
    print(f"模型名称: {config['model']}")
    print("=" * 80)
    
    # 创建测试器
    tester = PromptTester(config["api_url"], config["api_key"], is_local)
    
    custom_content = """

上下时不得手持物件；必须从指定的路线上下，不得在高空投掷材料或工具等物；不得将易滚易滑的工具、材料堆放在脚手架上；工作完毕应及时将工具、零星材料、零部件等一切易坠落物件清理干净，以防落下伤人，上下大型零件时，应采用可靠的起吊机具。
12）要处处注意危险标志和危险地方。夜间作业，必须设置足够的照明设施，否则禁止施工。
13）严禁上下同时垂直作业。若特殊情况必须垂直作业，应经有关领导批准，并在上下两层间设备专用的防护棚或者其他隔离设施。
14）严禁坐在高空无遮栏处休息，防止坠落。
15）卷扬机等各种升降材料的设备严禁上下载人。
16）超过3米长的铺板不能同时站两人作业。
17）遇六级以上大风时，禁止露天进行高空作业。
（4）起重作业安全保障措施
1）起重设备在现场初次使用时应进行试吊，试吊严格按照《起重机试验规范和程序》(GB/T 5905-2011）执行，试吊合格后方可使用。
2）起重作业前，检查起重设备的空载运转、回转、起重、变幅等各种机构的制动器、安全限位、防护装置等，确认正常后方可作业。
3）起重作业时，应依次逐级操作，严禁越档操作，在变换运转方向时，应将控制器转到零位，停止转动后，再转向另一方向，严禁急开急停。
4）司机、起重信号司索工等人员应经专业培训，取得特种设备作业人员证后，方可上岗作业。
5）起吊物件旋转或移动，应将起吊物提升至可能遇到的障碍物上方0.5m 以上。
6）起重作业严格执行安全操作规程，严格遵守“十不吊”的要求。
7）吊装人员必须戴安全帽、安全带；
8）吊装工作区设有明显标志，并设专人警戒，与吊装无关人员严禁入内。起重机工作时，起重臂杆旋转半径范围内，严禁站人或通过。
9）运输、吊装构件时，严禁在被运输、吊装的构件上指挥和放置材料、工具
10）吊装时，必须用吊笼或钢丝绳、保险绳绑扎牢固后才能吊运和传递，不得随意抛掷材料物体、工具，防止滑脱伤人或意外事故。
11）构件必须绑扎牢固，起吊点应通过构件的重心位置，吊升时应平稳，避免振动或摆动，以防构件脱落。
12）起吊构件时，速度不应太快，不得在高空停留过久，严禁猛升猛降，以防构件脱落。
13）构件就位后临时固定前，不得松钩、解开吊装索具。构件固定后，应检查连接牢固和稳定情况，当确定连接安全可靠，才可拆除临时固定工具和进行下一步吊装。
14）严禁夜间施工作业。15）起重机行驶的道路必须平整、坚实、可靠，停放地点必须平坦。
16）起重机不得停放在斜坡道上工作，不允许起重机两条履带或支腿停留部位一高一低或土质一硬一软。

"""
    tester.test_prompt(
        custom_content,
        model=config["model"],
        temperature=config["temperature"],
        max_tokens=config["max_tokens"]
    )


if __name__ == "__main__":
    main()