llm_client.py

"""
大模型API客户端 - 支持智谱AI GLM-4-Flash免费API
新策略：每实验只调用一次API，批量生成所有内容
"""

import json
import traceback
from typing import Optional, Callable, Dict, List

try:
    import requests
    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False


class LLMClient:
    """大模型客户端 - 智谱AI GLM-4-Flash"""
    
    DEFAULT_API_URL = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
    DEFAULT_MODEL = "glm-4-flash"
    
    def __init__(self, api_key: Optional[str] = None, api_url: Optional[str] = None):
        self.api_key = api_key
        self.api_url = api_url or self.DEFAULT_API_URL
        self.model = self.DEFAULT_MODEL
        self.last_error = None
        
        if not HAS_REQUESTS:
            self.last_error = "缺少requests库"
    
    def is_available(self) -> bool:
        return HAS_REQUESTS and self.api_key is not None
    
    def generate_full_experiment(self, title: str, template_structure: List[Dict],
                                  rag_context: str = "", word_target: int = 5000) -> Dict[str, str]:
        """
        一次性生成整个实验的所有内容
        
        Args:
            title: 报告标题
            template_structure: 模板结构 [{heading, subsections:[{heading}]}]
            rag_context: RAG检索到的上下文
            word_target: 目标字数
        
        Returns:
            {section_heading: content_text, ...}
        """
        if not self.is_available():
            return {}
        
        # 构建结构描述
        structure_desc = []
        for sec in template_structure:
            structure_desc.append(f"【{sec['heading']}】")
            for sub in sec.get('subsections', []):
                structure_desc.append(f"  - {sub['heading']}")
        
        system_prompt = f"""你是一位专业的学术报告撰写助手。请为实验报告《{title}》撰写完整内容。

核心要求：
1. 按照给定的章节结构，为每个子章节撰写内容
2. 内容要专业、详实、有深度，避免空洞套话
3. 结合具体技术细节和实践经验
4. 语言流畅自然，符合学术写作规范
5. 总字数约{word_target}字

输出格式要求：
- 每个子章节之间用 "---SECTION_BREAK---" 分隔
- 只输出内容，不要输出章节标题
- 段落之间不要有空行
- 每个子章节至少200字
"""
        
        user_prompt = f"""报告标题：{title}
目标字数：{word_target}字

章节结构：
"""
        user_prompt += "\n".join(structure_desc)
        
        if rag_context:
            user_prompt += f"""

参考信息（来自知识库）：
{rag_context}
"""
        
        user_prompt += """

请按照章节结构顺序，依次为每个子章节撰写内容。
用 "---SECTION_BREAK---" 分隔不同子章节的内容。
"""
        
        result = self._call_api(system_prompt, user_prompt, max_tokens=8000)
        
        if result:
            return self._parse_batch_result(result, template_structure)
        return {}
    
    def _parse_batch_result(self, text: str, template_structure: List[Dict]) -> Dict[str, str]:
        """解析批量生成结果"""
        sections = text.split("---SECTION_BREAK---")
        sections = [s.strip() for s in sections if s.strip()]
        
        result = {}
        idx = 0
        for sec in template_structure:
            for sub in sec.get('subsections', []):
                if idx < len(sections):
                    result[sub['heading']] = sections[idx]
                    idx += 1
        
        return result
    
    def _call_api(self, system_prompt: str, user_prompt: str,
                   max_tokens: int = 4096, temperature: float = 0.7) -> str:
        """调用大模型API"""
        if not HAS_REQUESTS:
            return ""
        
        try:
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.api_key}"
            }
            
            data = {
                "model": self.model,
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                "max_tokens": max_tokens,
                "temperature": temperature,
                "stream": False
            }
            
            response = requests.post(
                self.api_url,
                headers=headers,
                json=data,
                timeout=120
            )
            
            if response.status_code == 200:
                result = response.json()
                if isinstance(result, dict) and "error" in result:
                    error_msg = result["error"].get("message", "未知错误") if isinstance(result["error"], dict) else str(result["error"])
                    self.last_error = f"API错误: {error_msg}"
                    return ""
                
                if isinstance(result, dict) and "choices" in result and isinstance(result["choices"], list) and len(result["choices"]) > 0:
                    choice = result["choices"][0]
                    if isinstance(choice, dict) and "message" in choice and isinstance(choice["message"], dict):
                        content = choice["message"].get("content", "")
                        if content:
                            return content.strip()
                    self.last_error = "API返回格式异常"
                else:
                    self.last_error = "API返回格式异常"
            else:
                self.last_error = f"API请求失败: HTTP {response.status_code}"
                try:
                    error_info = response.json()
                    if isinstance(error_info, dict) and "error" in error_info:
                        error_msg = error_info["error"].get("message", "") if isinstance(error_info["error"], dict) else str(error_info["error"])
                        self.last_error += f" - {error_msg}"
                except:
                    pass
        
        except requests.exceptions.Timeout:
            self.last_error = "API请求超时"
        except requests.exceptions.ConnectionError:
            self.last_error = "网络连接失败"
        except Exception as e:
            self.last_error = f"API调用异常: {str(e)}"
            traceback.print_exc()
        
        return ""


class LLMContentGenerator:
    """基于LLM的内容生成器 - 每实验只调用一次API"""
    
    def __init__(self, llm_client: Optional[LLMClient] = None):
        self.llm = llm_client
        self.use_llm = llm_client is not None and llm_client.is_available()
        self.fallback_mode = False
        self.status_callback = None
    
    def set_status_callback(self, callback):
        self.status_callback = callback
    
    def _emit_status(self, status: str, message: str):
        if self.status_callback:
            self.status_callback(status, message)
    
    def generate_experiment_content(self, title: str, template_structure: List[Dict],
                                     rag_context: str = "", word_target: int = 5000) -> Dict[str, List[str]]:
        """
        生成整个实验的所有内容 - 只调用一次API
        
        Returns:
            {section_heading: [paragraphs], ...}
        """
        if self.use_llm and not self.fallback_mode:
            self._emit_status("llm_call", f"🤖 正在为【{title}】生成完整内容...")
            
            result = self.llm.generate_full_experiment(
                title=title,
                template_structure=template_structure,
                rag_context=rag_context,
                word_target=word_target
            )
            
            if result:
                self._emit_status("llm_success", f"✅ 【{title}】生成完成")
                # 将文本分割为段落列表
                parsed = {}
                for heading, text in result.items():
                    paragraphs = [p.strip() for p in text.replace('\r\n', '\n').split('\n') if p.strip()]
                    parsed[heading] = paragraphs if paragraphs else [text]
                return parsed
            else:
                error_msg = self.llm.last_error or "未知错误"
                self._emit_status("llm_error", f"⚠️ LLM调用失败: {error_msg}")
                self._emit_status("llm_fallback", "🔄 回退到模板生成...")
                self.fallback_mode = True
        
        # 回退到模板生成
        return {}
    
    def reset_fallback(self):
        self.fallback_mode = False