Sahil1694/Atlan

1import os
2import json
3from langchain.text_splitter import RecursiveCharacterTextSplitter
4from langchain.docstore.document import Document
5from langchain_huggingface import HuggingFaceEmbeddings
6from langchain_chroma import Chroma
7
8# --- Configuration ---
9KNOWLEDGE_BASE_PATH = "data/knowledge_base.json"
10# Updated save path to be specific to ChromaDB
11VECTOR_STORE_SAVE_PATH = "data/chroma_db"
12
13# Configuration for the text splitter
14CHUNK_SIZE = 1000  # The number of characters in each chunk
15CHUNK_OVERLAP = 100 # The number of characters to overlap between chunks
16
17# We will use a powerful, open-source embedding model.
18# "all-MiniLM-L6-v2" is a great starting point - it's fast and effective.
19EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"
20
21# --- Core Logic ---
22
23def load_knowledge_base(filepath: str) -> list:
24    """Loads the scraped data from the JSON file."""
25    try:
26        with open(filepath, 'r', encoding='utf-8') as f:
27            data = json.load(f)
28        return data
29    except FileNotFoundError:
30        return []
31    except json.JSONDecodeError:
32        return []
33
34def create_documents_from_data(data: list) -> list:
35    """Converts the raw data into LangChain's Document format."""
36    documents = [
37        Document(page_content=item['content'], metadata={'source': item['url']})
38        for item in data if item.get('content')
39    ]
40    return documents
41
42def split_documents(documents: list) -> list:
43    """Splits the documents into smaller chunks using the recursive strategy."""
44    text_splitter = RecursiveCharacterTextSplitter(
45        chunk_size=CHUNK_SIZE,
46        chunk_overlap=CHUNK_OVERLAP,
47        length_function=len,
48        add_start_index=True,
49    )
50    chunked_documents = text_splitter.split_documents(documents)
51    return chunked_documents
52
53def build_and_save_vector_store(chunks: list, save_path: str):
54    """
55    Creates embeddings for the chunks and builds a ChromaDB vector store,
56    then saves it to disk.
57    """
58    if not chunks:
59        return
60
61    # Initialize the embedding model. It might download the model files on the first run.
62    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
63
64    # Chroma will automatically create and save the database in the specified directory.
65    # The `persist_directory` argument tells Chroma to save the data to disk.
66    vector_store = Chroma.from_documents(
67        documents=chunks,
68        embedding=embeddings,
69        persist_directory=save_path
70    )
71
72# --- Main Execution ---
73
74if __name__ == "__main__":
75    # 1. Load the raw scraped data
76    scraped_data = load_knowledge_base(KNOWLEDGE_BASE_PATH)
77
78    if scraped_data:
79        # 2. Convert to LangChain Documents
80        documents = create_documents_from_data(scraped_data)
81
82        # 3. Split the documents into chunks
83        chunked_docs = split_documents(documents)
84
85        # 4. Build and save the vector store
86        build_and_save_vector_store(chunked_docs, VECTOR_STORE_SAVE_PATH)
87
88