src/embedding.py

from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
import uuid
from src.scheme import create_scheme
import os
from dotenv import load_dotenv
import chromadb

load_dotenv()

COLLECTION_NAME = "Text2SQL"

chroma_client = chromadb.CloudClient(
    api_key=os.getenv("CHROMA_API_KEY"),
    tenant=os.getenv("CHROMA_TENANT"),
    database=os.getenv("CHROMA_DATABASE"),
)

embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = Chroma(collection_name=COLLECTION_NAME,embedding_function=embedding_model,client=chroma_client)

def create_embeddings(connection_url : str , user_id : str) :
    docs = create_scheme(connection_url)

    for doc in docs :
        doc.metadata['user_id'] = user_id

    ids = [str(uuid.uuid4()) for _ in docs]

    vectorstore.add_documents(documents=docs , ids=ids)