60 lines
2.3 KiB
Python
60 lines
2.3 KiB
Python
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_keywords(
|
||
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 = TextRank(text, num_sentences=1) # 生成80句话的摘要
|
||
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_keywords = get_llm_model_response(
|
||
strategy_name="gen_keywords",
|
||
llm_model_name=LLM_MODELS[0],
|
||
template_prompt_name="gen_keywords",
|
||
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_keywords)
|
||
return article_keywords
|