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]:
    """
    从研究主题中提取关键词
    
    Args:
        research_topic: 用户输入的研究主题
        llm_client: LLM客户端实例
        
    Returns:
        关键词列表
    """
    prompt = f"""
    作为一个科研助手，请从以下研究主题中提取5-8个关键的检索词，这些词应该能够用于学术文献检索。
    请考虑主题的核心概念、方法论、应用领域等方面。返回格式为逗号分隔的关键词列表，不要有编号或其他文本。

    研究主题: {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

async def expand_search_queries(
    keywords: List[str],
    llm_client: LLMClient
) -> List[str]:
    """
    扩展搜索查询，生成更全面的检索语句
    
    Args:
        keywords: 关键词列表
        llm_client: LLM客户端实例
        
    Returns:
        扩展后的搜索查询列表
    """
    kw_text = ", ".join(keywords)
    
    prompt = f"""
    作为科研检索专家，我需要你帮助构建有效的学术文献检索语句。
    基于以下关键词：{kw_text}
    
    请生成3个不同的arXiv搜索查询语句，这些查询应当：
    1. 使用布尔运算符(AND, OR)来组合关键词
    2. 考虑同义词和相关概念
    3. 针对不同的研究角度
    
    只需返回这些查询语句，每行一个，不要有编号或其他文本。
    """
    
    response = await llm_client.generate_text(prompt, temperature=0.2)
    
    # 清理和解析响应
    queries = [q.strip() for q in response.splitlines() if q.strip()]
    
    # 确保至少有一个查询
    if not queries and keywords:
        queries = [" AND ".join(keywords[:3])]
    
    return queries