Lightstar/backend/utils/keywords.py

"""
关键词提取和处理工具
"""

import re
from typing import List, Dict, Any
from .api_client import LLMClient

async def extract_keywords(
    research_topic: str,
    llm_client: LLMClient
) -> List[str]:
    """从研究主题中提取核心关键词"""
    prompt = f"""
    请从以下研究主题中提取4-6个核心关键词。这些关键词应当:
    1. 直接来源于用户文本，不要发散生成
    2. 准确反映研究主题的核心概念
    3. 具有学术检索价值
    4. 优先选择专业术语、技术名词和研究领域
    
    研究主题: {research_topic}
    
    仅返回逗号分隔的关键词列表:
    """
    
    response = await llm_client.generate_text(prompt, temperature=0.1)
    
    # 清理处理响应
    response = re.sub(r'[\n\r]+', ',', response)
    response = re.sub(r'\s*,\s*', ',', response)
    response = re.sub(r'^\s+|\s+$', '', response)
    
    keywords = [kw.strip() for kw in response.split(',') if kw.strip()]
    
    # 限制数量
    return keywords[:6]

async def expand_search_queries(
    keywords: List[str],
    llm_client: LLMClient
) -> List[str]:
    """基于关键词生成研究方向"""
    kw_text = ", ".join(keywords)
    
    prompt = f"""
    基于以下关键词，生成4-6个明确、具体且各不相同的研究方向。
    
    关键词: {kw_text}
    
    每个研究方向应当:
    1. 具体且可研究，而非笼统的话题
    2. 聚焦不同的研究角度或应用领域
    3. 适合作为学术研究问题
    4. 形式为研究问题或研究主题陈述
    
    只需返回研究方向列表，每行一个，不要包含编号或其他解释。
    """
    
    response = await llm_client.generate_text(prompt, temperature=0.7)
    
    # 解析研究方向
    directions = [line.strip() for line in response.splitlines() if line.strip()]
    
    return directions