import faiss
import os
from sentence_transformers import SentenceTransformer
from dotenv import load_dotenv
from typing import List, Dict, AsyncGenerator
import numpy as np
# 加载环境变量
load_dotenv()

# DeepSeek API配置
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL")

class FaissHelper:
    def __init__(self, dimension: int):
        self.model = SentenceTransformer(EMBEDDING_MODEL)       
    def generate_embeddings(self,texts: List[str]) -> np.ndarray:
        return self.model.encode(texts)

    def create_index(self,dimension: int):
        index = faiss.IndexFlatL2(dimension)
        return index
    
    def get_index_size(self,index):
        return index.ntotal
    def index_add(self,index, embeddings):
        index.add(embeddings)
    def load_faiss_index(self,index_path):
        index = faiss.read_index(index_path)
        return index

    def search_faiss_index(self, index, query, top_k):
        query_vector = query.reshape(1, -1)
        _, indices = index.search(query_vector, top_k)
        return indices[0]