bharathmunakala/Role_Base_Access_Control

1import os
2import chromadb
3from pathlib import Path
4from dotenv import load_dotenv
5from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext, Settings
6from llama_index.vector_stores.chroma import ChromaVectorStore
7from llama_index.embeddings.cohere import CohereEmbedding
8
9# Load environment variables
10load_dotenv()
11
12# Configure the embedding model
13cohere_api_key = os.getenv("COHERE_API_KEY")
14if not cohere_api_key:
15    raise ValueError("COHERE_API_KEY not found in environment variables")
16
17# Initialize the embedding model
18embed_model = CohereEmbedding(
19    cohere_api_key=cohere_api_key,
20    model_name="embed-english-v3.0",
21    input_type="search_document"
22)
23
24# Set the global embedding model
25Settings.embed_model = embed_model
26
27def process_documents(department: str, base_dir: str = "./resources/data"):
28    """
29    Process and index documents for a specific department
30    
31    Args:
32        department: The department name (e.g., 'hr', 'engineering')
33        base_dir: Base directory containing department folders
34    """
35    print(f"Processing documents for {department} department...")
36    
37    # Define paths
38    dept_path = Path(base_dir) / department
39    general_path = Path(base_dir) / "general"
40    persist_dir = f"./chroma_db/{department}"
41    
42    # Create directory if it doesn't exist
43    os.makedirs(persist_dir, exist_ok=True)
44    
45    # Initialize Chroma client
46    chroma_client = chromadb.PersistentClient(path=persist_dir)
47    
48    # Clear existing collection if it exists
49    try:
50        chroma_client.delete_collection("documents")
51    except:
52        pass
53    
54    # Create a new collection
55    chroma_collection = chroma_client.get_or_create_collection("documents")
56    
57    # Create vector store
58    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
59    storage_context = StorageContext.from_defaults(vector_store=vector_store)
60    
61    # Load department-specific documents
62    documents = []
63    
64    # Add department-specific files
65    if dept_path.exists() and dept_path.is_dir():
66        for file_path in dept_path.glob("*"):
67            if file_path.is_file() and file_path.suffix in ['.md', '.txt', '.csv']:
68                print(f"Processing {file_path.name}...")
69                try:
70                    # Read the file content
71                    with open(file_path, 'r', encoding='utf-8') as f:
72                        content = f.read()
73                    
74                    # Create a document with metadata
75                    from llama_index.core import Document
76                    doc = Document(
77                        text=content,
78                        metadata={
79                            "source": str(file_path.name),
80                            "department": department,
81                            "type": "department_specific"
82                        }
83                    )
84                    documents.append(doc)
85                except Exception as e:
86                    print(f"Error processing {file_path}: {str(e)}")
87    
88    # Add general documents
89    if general_path.exists() and general_path.is_dir():
90        for file_path in general_path.glob("*"):
91            if file_path.is_file() and file_path.suffix in ['.md', '.txt', '.csv']:
92                print(f"Processing general document: {file_path.name}...")
93                try:
94                    # Read the file content
95                    with open(file_path, 'r', encoding='utf-8') as f:
96                        content = f.read()
97                    
98                    # Create a document with metadata
99                    from llama_index.core import Document
100                    doc = Document(
101                        text=content,
102                        metadata={
103                            "source": str(file_path.name),
104                            "department": "general",
105                            "type": "general"
106                        }
107                    )
108                    documents.append(doc)
109                except Exception as e:
110                    print(f"Error processing general document {file_path}: {str(e)}")
111    
112    if not documents:
113        print(f"No documents found for {department} department.")
114        return
115    
116    print(f"Indexing {len(documents)} documents...")
117    
118    # Create index with the documents
119    index = VectorStoreIndex.from_documents(
120        documents,
121        storage_context=storage_context,
122        show_progress=True,
123        embed_model=embed_model
124    )
125    
126    print(f"✅ Successfully indexed {len(documents)} documents for {department} department")
127    print(f"Index stored in: {persist_dir}")
128
129def main():
130    """Main function to process documents for all departments"""
131    departments = ["hr", "engineering", "finance", "marketing"]
132    
133    for dept in departments:
134        print(f"\n{'='*50}")
135        print(f"Processing {dept.upper()} department")
136        print(f"{'='*50}")
137        process_documents(dept)
138    
139    print("\n✅ Document processing completed for all departments!")
140
141if __name__ == "__main__":
142    main()
143