Shrook21/Code-Assistant

1# import os
2# from datasets import load_dataset
3# from langchain_core.documents import Document
4# from langchain_chroma import Chroma
5# from vectorstore.retriever import get_embedding_model
6# from utils.code_splitter import split_code_by_function
7# from config.settings import PERSIST_DIR
8
9# def build_vectorstore():
10#     if os.path.exists(PERSIST_DIR):
11#         print("Loading existing Chroma DB...")
12#         return Chroma(persist_directory=PERSIST_DIR, embedding_function=get_embedding_model())
13
14#     print("Building new Chroma DB from dataset...")
15#     ds = load_dataset("openai_humaneval")["test"]
16#     examples = [{"id": row["task_id"], "prompt": row["prompt"], "solution": row["canonical_solution"]} for row in ds]
17
18#     raw_documents = [
19#         Document(
20#             page_content=f"{ex['prompt']}\n\n# Solution:\n{ex['solution']}",
21#             metadata={"id": ex["id"]}
22#         ) for ex in examples
23#     ]
24
25#     split_documents = [d for doc in raw_documents for d in split_code_by_function(doc)]
26
27#     vectorstore = Chroma.from_documents(
28#         documents=split_documents,
29#         embedding=get_embedding_model(),
30#         persist_directory=PERSIST_DIR
31#     )
32
33#     print("Chroma DB created and persisted.")
34#     return vectorstore
35import os
36from datasets import load_dataset
37from langchain_core.documents import Document
38from langchain_chroma import Chroma
39from vectorstore.retriever import get_embedding_model
40from utils.code_splitter import split_code_by_function
41from config.settings import PERSIST_DIR
42
43def build_vectorstore(force_rebuild=False):
44    """Build or load the vectorstore"""
45    
46    # Check if vectorstore exists and is not empty
47    if os.path.exists(PERSIST_DIR) and not force_rebuild:
48        print("Loading existing Chroma DB...")
49        try:
50            vectorstore = Chroma(
51                persist_directory=PERSIST_DIR, 
52                embedding_function=get_embedding_model()
53            )
54            
55            # Test if the vectorstore has documents
56            test_results = vectorstore.similarity_search("def", k=1)
57            if test_results:
58                print(f"Loaded existing vectorstore with {len(test_results)} sample documents")
59                return vectorstore
60            else:
61                print("Existing vectorstore appears empty, rebuilding...")
62        except Exception as e:
63            print(f"Error loading existing vectorstore: {e}")
64            print("Rebuilding vectorstore...")
65    
66    print("Building new Chroma DB from dataset...")
67    
68    # Load dataset
69    ds = load_dataset("openai_humaneval")["test"]
70    examples = [
71        {
72            "id": row["task_id"], 
73            "prompt": row["prompt"], 
74            "solution": row["canonical_solution"]
75        } 
76        for row in ds
77    ]
78    
79    print(f"Loaded {len(examples)} examples from dataset")
80    
81    # Create raw documents
82    raw_documents = [
83        Document(
84            page_content=f"{ex['prompt']}\n\n# Solution:\n{ex['solution']}",
85            metadata={"id": ex["id"], "type": "code_example"}
86        ) 
87        for ex in examples
88    ]
89    
90    # Split documents by function
91    split_documents = []
92    for doc in raw_documents:
93        split_docs = split_code_by_function(doc)
94        split_documents.extend(split_docs)
95    
96    print(f"Split into {len(split_documents)} document chunks")
97    
98    # Remove existing directory if rebuilding
99    if force_rebuild and os.path.exists(PERSIST_DIR):
100        import shutil
101        shutil.rmtree(PERSIST_DIR)
102        os.makedirs(PERSIST_DIR, exist_ok=True)
103    
104    # Create vectorstore
105    embedding_model = get_embedding_model()
106    vectorstore = Chroma.from_documents(
107        documents=split_documents,
108        embedding=embedding_model,
109        persist_directory=PERSIST_DIR
110    )
111    
112    print("Chroma DB created and persisted.")
113    
114    # Verify the vectorstore was created successfully
115    test_results = vectorstore.similarity_search("def", k=1)
116    print(f"Verification: Found {len(test_results)} test results")
117    
118    return vectorstore