Spaces:
Paused
Paused
| import os | |
| from langchain_community.vectorstores import FAISS | |
| from langchain_huggingface import HuggingFaceEmbeddings | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| def get_embeddings_model(): | |
| return HuggingFaceEmbeddings( | |
| model_name="sentence-transformers/all-MiniLM-L6-v2" | |
| ) | |
| def get_vector_store_path(user_id: str) -> str: | |
| base_path = os.path.join("data", "users", user_id, "vector_store") | |
| os.makedirs(base_path, exist_ok=True) | |
| return base_path | |
| def split_text_into_chunks(text: str, chunk_size: int = 500, chunk_overlap: int = 50): | |
| text_splitter = RecursiveCharacterTextSplitter( | |
| chunk_size=chunk_size, | |
| chunk_overlap=chunk_overlap, | |
| length_function=len, | |
| ) | |
| chunks = text_splitter.split_text(text) | |
| return chunks | |
| def save_chunks_to_vectorstore(user_id: str, chunks: list): | |
| embeddings = get_embeddings_model() | |
| vector_store_path = get_vector_store_path(user_id) | |
| vector_store_file = os.path.join(vector_store_path, "index.faiss") | |
| if os.path.exists(vector_store_file): | |
| vectorstore = FAISS.load_local( | |
| vector_store_path, | |
| embeddings, | |
| allow_dangerous_deserialization=True | |
| ) | |
| vectorstore.add_texts(chunks) | |
| else: | |
| vectorstore = FAISS.from_texts(chunks, embeddings) | |
| vectorstore.save_local(vector_store_path) | |
| return f"Added {len(chunks)} chunks to knowledge base" | |
| def load_vectorstore(user_id: str): | |
| embeddings = get_embeddings_model() | |
| vector_store_path = get_vector_store_path(user_id) | |
| vector_store_file = os.path.join(vector_store_path, "index.faiss") | |
| if not os.path.exists(vector_store_file): | |
| raise FileNotFoundError(f"No knowledge base found for user {user_id}") | |
| vectorstore = FAISS.load_local( | |
| vector_store_path, | |
| embeddings, | |
| allow_dangerous_deserialization=True | |
| ) | |
| return vectorstore | |
| def retrieve_relevant_chunks(user_id: str, query: str, k: int = 3): | |
| vectorstore = load_vectorstore(user_id) | |
| docs = vectorstore.similarity_search(query, k=k) | |
| return [doc.page_content for doc in docs] | |