hello-agents/Co-creation-projects/Apricity-InnocoreAI/utils/embedding.py

"""
InnoCore AI 向量生成工具
"""

import asyncio
from typing import List, Dict, Optional, Any
import numpy as np
from openai import AsyncOpenAI
import hashlib
import json

from ..core.config import get_config
from ..core.exceptions import AgentException

class EmbeddingGenerator:
    """向量生成器"""
    
    def __init__(self):
        self.config = get_config()
        self.client = None
        self.embedding_model = self.config.vector_db.embedding_model
        self.cache = {}  # 简单的内存缓存
    
    async def initialize(self):
        """初始化向量生成器"""
        try:
            self.client = AsyncOpenAI(
                api_key=self.config.llm.api_key,
                base_url=self.config.llm.base_url
            )
        except Exception as e:
            raise AgentException(f"向量生成器初始化失败: {str(e)}")
    
    async def generate_embedding(self, text: str, use_cache: bool = True) -> List[float]:
        """生成文本向量"""
        if not text:
            return [0.0] * 1536  # 返回零向量
        
        # 检查缓存
        if use_cache:
            cache_key = self._get_cache_key(text)
            if cache_key in self.cache:
                return self.cache[cache_key]
        
        try:
            # 清理文本
            cleaned_text = self._clean_text(text)
            
            # 调用OpenAI API
            response = await self.client.embeddings.create(
                model=self.embedding_model,
                input=cleaned_text
            )
            
            embedding = response.data[0].embedding
            
            # 缓存结果
            if use_cache:
                cache_key = self._get_cache_key(text)
                self.cache[cache_key] = embedding
            
            return embedding
            
        except Exception as e:
            raise AgentException(f"向量生成失败: {str(e)}")
    
    async def generate_batch_embeddings(self, texts: List[str], 
                                       batch_size: int = 10) -> List[List[float]]:
        """批量生成向量"""
        embeddings = []
        
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            
            try:
                # 批量调用API
                cleaned_texts = [self._clean_text(text) for text in batch]
                
                response = await self.client.embeddings.create(
                    model=self.embedding_model,
                    input=cleaned_texts
                )
                
                batch_embeddings = [item.embedding for item in response.data]
                embeddings.extend(batch_embeddings)
                
            except Exception as e:
                # 如果批量失败，逐个生成
                for text in batch:
                    try:
                        embedding = await self.generate_embedding(text)
                        embeddings.append(embedding)
                    except Exception as single_error:
                        print(f"单个向量生成失败: {str(single_error)}")
                        embeddings.append([0.0] * 1536)  # 零向量
        
        return embeddings
    
    async def generate_paper_embedding(self, paper_info: Dict[str, Any]) -> List[float]:
        """为论文生成综合向量"""
        # 组合论文的关键信息
        title = paper_info.get("title", "")
        abstract = paper_info.get("abstract", "")
        authors = " ".join(paper_info.get("authors", []))
        
        # 构建综合文本
        combined_text = f"{title} {abstract} {authors}"
        
        # 如果有结构化内容，也包含进来
        sections = paper_info.get("sections", {})
        if sections:
            section_text = " ".join(sections.values())
            combined_text += " " + section_text
        
        return await self.generate_embedding(combined_text)
    
    async def generate_section_embeddings(self, sections: Dict[str, str]) -> Dict[str, List[float]]:
        """为各个章节生成向量"""
        section_embeddings = {}
        
        for section_name, section_content in sections.items():
            if section_content.strip():
                try:
                    embedding = await self.generate_embedding(section_content)
                    section_embeddings[section_name] = embedding
                except Exception as e:
                    print(f"章节 {section_name} 向量生成失败: {str(e)}")
                    section_embeddings[section_name] = [0.0] * 1536
        
        return section_embeddings
    
    def _clean_text(self, text: str) -> str:
        """清理文本"""
        if not text:
            return ""
        
        # 移除多余的空白字符
        text = ' '.join(text.split())
        
        # 截断过长的文本（OpenAI有token限制）
        max_length = 8000  # 保守估计
        if len(text) > max_length:
            text = text[:max_length]
        
        return text
    
    def _get_cache_key(self, text: str) -> str:
        """生成缓存键"""
        return hashlib.md5(text.encode()).hexdigest()
    
    def clear_cache(self):
        """清空缓存"""
        self.cache.clear()
    
    def get_cache_size(self) -> int:
        """获取缓存大小"""
        return len(self.cache)
    
    async def calculate_similarity(self, text1: str, text2: str) -> float:
        """计算两个文本的相似度"""
        try:
            embedding1 = await self.generate_embedding(text1)
            embedding2 = await self.generate_embedding(text2)
            
            return self._cosine_similarity(embedding1, embedding2)
            
        except Exception as e:
            print(f"相似度计算失败: {str(e)}")
            return 0.0
    
    def _cosine_similarity(self, vec1: List[float], vec2: List[float]) -> float:
        """计算余弦相似度"""
        if len(vec1) != len(vec2):
            return 0.0
        
        try:
            vec1_np = np.array(vec1)
            vec2_np = np.array(vec2)
            
            dot_product = np.dot(vec1_np, vec2_np)
            norm1 = np.linalg.norm(vec1_np)
            norm2 = np.linalg.norm(vec2_np)
            
            if norm1 == 0 or norm2 == 0:
                return 0.0
            
            return dot_product / (norm1 * norm2)
            
        except Exception:
            return 0.0
    
    async def find_most_similar(self, query_text: str, 
                               candidate_texts: List[str],
                               top_k: int = 5) -> List[Dict[str, Any]]:
        """找到最相似的文本"""
        if not candidate_texts:
            return []
        
        try:
            # 生成查询向量
            query_embedding = await self.generate_embedding(query_text)
            
            # 生成候选文本向量
            candidate_embeddings = await self.generate_batch_embeddings(candidate_texts)
            
            # 计算相似度
            similarities = []
            for i, candidate_embedding in enumerate(candidate_embeddings):
                similarity = self._cosine_similarity(query_embedding, candidate_embedding)
                similarities.append({
                    "text": candidate_texts[i],
                    "similarity": similarity,
                    "index": i
                })
            
            # 按相似度排序
            similarities.sort(key=lambda x: x["similarity"], reverse=True)
            
            return similarities[:top_k]
            
        except Exception as e:
            print(f"相似文本查找失败: {str(e)}")
            return []
    
    async def cluster_texts(self, texts: List[str], 
                          num_clusters: int = 3) -> Dict[str, Any]:
        """文本聚类（简化实现）"""
        try:
            # 生成所有文本的向量
            embeddings = await self.generate_batch_embeddings(texts)
            
            # 简单的聚类逻辑（基于相似度阈值）
            clusters = {}
            cluster_id = 0
            used_indices = set()
            
            for i, embedding in enumerate(embeddings):
                if i in used_indices:
                    continue
                
                # 创建新聚类
                clusters[f"cluster_{cluster_id}"] = {
                    "texts": [texts[i]],
                    "indices": [i],
                    "center": embedding
                }
                used_indices.add(i)
                
                # 查找相似文本加入同一聚类
                for j, other_embedding in enumerate(embeddings):
                    if j in used_indices:
                        continue
                    
                    similarity = self._cosine_similarity(embedding, other_embedding)
                    if similarity > 0.8:  # 相似度阈值
                        clusters[f"cluster_{cluster_id}"]["texts"].append(texts[j])
                        clusters[f"cluster_{cluster_id}"]["indices"].append(j)
                        used_indices.add(j)
                
                cluster_id += 1
            
            return {
                "clusters": clusters,
                "num_clusters": len(clusters),
                "total_texts": len(texts)
            }
            
        except Exception as e:
            print(f"文本聚类失败: {str(e)}")
            return {"clusters": {}, "num_clusters": 0, "total_texts": len(texts)}
    
    async def extract_keywords(self, text: str, max_keywords: int = 10) -> List[str]:
        """提取关键词（基于TF-IDF的简化实现）"""
        try:
            # 分词
            words = text.lower().split()
            
            # 过滤停用词（简化版）
            stop_words = {
                'the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for',
                'of', 'with', 'by', 'is', 'are', 'was', 'were', 'be', 'been', 'have',
                'has', 'had', 'do', 'does', 'did', 'will', 'would', 'could', 'should'
            }
            
            filtered_words = [word for word in words if word not in stop_words and len(word) > 2]
            
            # 计算词频
            word_freq = {}
            for word in filtered_words:
                word_freq[word] = word_freq.get(word, 0) + 1
            
            # 按频率排序
            sorted_words = sorted(word_freq.items(), key=lambda x: x[1], reverse=True)
            
            # 返回前N个关键词
            return [word for word, freq in sorted_words[:max_keywords]]
            
        except Exception as e:
            print(f"关键词提取失败: {str(e)}")
            return []
    
    def get_embedding_info(self) -> Dict[str, Any]:
        """获取向量生成器信息"""
        return {
            "model": self.embedding_model,
            "cache_size": len(self.cache),
            "vector_dimension": 1536,  # OpenAI embedding维度
            "provider": "openai"
        }