Lightstar/backend/utils/api_client.py

"""
API客户端工具
用于与外部API（如LLM API、文献检索API）进行交互
"""

import httpx
import json
import os
import time
from typing import Dict, List, Any, Optional
import logging

from backend.config import DEEPSEEK_API_KEY, DEEPSEEK_API_URL, DEEPSEEK_MODEL

logger = logging.getLogger(__name__)

class LLMClient:
    """大型语言模型API客户端"""
    
    def __init__(
        self, 
        api_key: Optional[str] = None,
        model: Optional[str] = None,
        api_url: Optional[str] = None
    ):
        self.api_key = api_key or DEEPSEEK_API_KEY
        self.model = model or DEEPSEEK_MODEL
        self.api_url = api_url or DEEPSEEK_API_URL
        
        if not self.api_key:
            logger.warning("No LLM API key provided. API calls will fail.")
    
    async def generate_text(
        self, 
        prompt: str, 
        temperature: float = 0.3,
        max_tokens: int = 1000
    ) -> str:
        """
        调用LLM生成文本
        
        Args:
            prompt: 提示文本
            temperature: 温度参数，控制随机性
            max_tokens: 最大生成token数
            
        Returns:
            生成的文本
        """
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        payload = {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": temperature,
            "max_tokens": max_tokens
        }
        
        try:
            async with httpx.AsyncClient(timeout=60.0) as client:
                response = await client.post(
                    f"{self.api_url}/chat/completions",
                    headers=headers,
                    json=payload
                )
                response.raise_for_status()
                result = response.json()
                return result["choices"][0]["message"]["content"]
        except Exception as e:
            logger.error(f"Error calling LLM API: {str(e)}")
            return f"Error generating text: {str(e)}"


class ArxivClient:
    """arXiv API客户端"""
    
    async def search_papers(
        self, 
        query: str, 
        max_results: int = 10
    ) -> List[Dict[str, Any]]:
        """
        搜索arXiv文献
        
        Args:
            query: 搜索查询
            max_results: 最大结果数量
            
        Returns:
            文献列表
        """
        try:
            import arxiv
            
            search = arxiv.Search(
                query=query,
                max_results=max_results,
                sort_by=arxiv.SortCriterion.Relevance
            )
            
            results = []
            for paper in await search.results():
                results.append({
                    "id": paper.get_short_id(),
                    "title": paper.title,
                    "authors": [author.name for author in paper.authors],
                    "summary": paper.summary,
                    "published": paper.published.isoformat() if paper.published else None,
                    "updated": paper.updated.isoformat() if paper.updated else None,
                    "link": paper.pdf_url,
                    "source": "arxiv"
                })
            
            return results
        except Exception as e:
            logger.error(f"Error searching arXiv: {str(e)}")
            return []