gangyan/langchain-chat/server/chat/gen_abstract.py

from fastapi import Body
from configs import LLM_MODELS, TEMPERATURE, MAX_TOKENS
from server.chat.policy_fun_iast import get_llm_model_response
from typing import Optional
from langchain.chains import LLMChain
from langchain.prompts import ChatPromptTemplate
from server.chat.utils import History
from server.utils import wrap_done, get_ChatOpenAI, get_prompt_template
from langchain.callbacks import AsyncIteratorCallbackHandler
import asyncio
from server.knowledge_base.kb_service.base import TextRank
from configs.basic_config import *

async def gen_abstract(
    data: dict = Body(..., description="包含撰写全文的JSON对象", example={"context": "当前已撰写的全文"})
):
    """
    接收当前已撰写的全文生成摘要。
    """
    context: Optional[str] = data.get("context")
    if not context:
        return {"error": "缺少必要参数 'context'"}

    logger.info(f"开始生成摘要...")

    # 定义生成摘要的函数
    def generate_summary(text: str) -> str:
        """使用 TextRank 生成文本摘要"""
        summary = text if len(text)<30000 else TextRank(text,num_sentences=1)
        return summary

    # 根据正文长度决定是否生成摘要
    # if len(context) >= 30000:
    #     context_summary = generate_summary(context)
    #     logger.info(f"生成撰写文稿的摘要: %s", context_summary)
    # else:
    #     context_summary = context  # 直接使用原文
    #     logger.info(f"撰写文稿小于30000字符，使用原文")
    if len(context) >= 40000:
        context_summary = context[:40000]
    else:
        context_summary = context
    # 调用模型生成生成摘要内容
    try:
        article_abstract = get_llm_model_response(
            strategy_name="gen_abstract",
            llm_model_name=LLM_MODELS[0],
            template_prompt_name="gen_abstract",
            prompt_param_dict={
                "context": context_summary,  # 使用摘要或原文
            },
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS
        )
    except Exception as e:
        logger.error("生成摘要时出错: %s", e)
        return {"error": "生成摘要时发生错误，请重试"}

    logger.info("生成摘要内容: %s", article_abstract)
    return article_abstract