LightRT/text2sql_backend
0
1from langchain_chroma import Chroma2from langchain_huggingface import HuggingFaceEmbeddings3import uuid4from src.scheme import create_scheme5import os6from dotenv import load_dotenv7import chromadb8 9load_dotenv()10 11COLLECTION_NAME = "Text2SQL"12 13chroma_client = chromadb.CloudClient(14 api_key=os.getenv("CHROMA_API_KEY"),15 tenant=os.getenv("CHROMA_TENANT"),16 database=os.getenv("CHROMA_DATABASE"),17)18 19embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")20vectorstore = Chroma(collection_name=COLLECTION_NAME,embedding_function=embedding_model,client=chroma_client)21 22def create_embeddings(connection_url : str , user_id : str) :23 docs = create_scheme(connection_url)24 25 for doc in docs :26 doc.metadata['user_id'] = user_id27 28 ids = [str(uuid.uuid4()) for _ in docs]29 30 vectorstore.add_documents(documents=docs , ids=ids)