siwei_ai/app/services/embed_service.py

import time
import os

from llm_model.embed import embed
from app.common.res import res_success, res_error
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores.chroma import Chroma

# 解析文件到向量


def parse_file_to_embed(file):
    start_time = time.time()
    embedded = embed(file)
    end_time = time.time()
    print("Time taken for embedding: ", end_time - start_time)

    if embedded:
        return res_success(msg="File embedded successfully")
    else:
        return res_error(msg="File embedded unsuccessfully")

# 删除向量


def delete_embed():
    db = get_vector_db()
    db.delete_collection()

    return res_success(msg="Collection deleted successfully")


CHROMA_PATH = os.getenv('CHROMA_PATH', 'chroma')
COLLECTION_NAME = os.getenv('COLLECTION_NAME', 'siwei_ai')
TEXT_EMBEDDING_MODEL = os.getenv('TEXT_EMBEDDING_MODEL', 'nomic-embed-text')


def get_vector_db():
    embedding = OllamaEmbeddings(
        model=TEXT_EMBEDDING_MODEL, show_progress=True, num_gpu=0, num_thread=4)

    db = Chroma(
        collection_name=COLLECTION_NAME,
        persist_directory=CHROMA_PATH,
        embedding_function=embedding
    )

    return db