Ehtasham08/Procurement
0
1import os2import faiss3import numpy as np4from groq import Groq5from sentence_transformers import SentenceTransformer6import streamlit as st7from langchain.text_splitter import RecursiveCharacterTextSplitter8import pickle9from typing import List, Dict10import docx11import PyPDF212from io import BytesIO13import tempfile14import shutil15 16# Page configuration17st.set_page_config(18 page_title="Procurement RAG Assistant",19 page_icon="🏢",20 layout="wide",21 initial_sidebar_state="expanded"22)23 24# Custom CSS for better UI25st.markdown("""26<style>27 .main-header {28 background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);29 padding: 1rem;30 border-radius: 10px;31 color: white;32 text-align: center;33 margin-bottom: 2rem;34 }35 .upload-section {36 background-color: #f0f2f6;37 padding: 1rem;38 border-radius: 10px;39 margin-bottom: 1rem;40 }41 .stats-container {42 background-color: #e8f4fd;43 padding: 1rem;44 border-radius: 10px;45 margin-top: 1rem;46 }47</style>48""", unsafe_allow_html=True)49 50class DocumentProcessor:51 def __init__(self):52 self.text_splitter = RecursiveCharacterTextSplitter(53 chunk_size=1000,54 chunk_overlap=200,55 length_function=len,56 separators=["\n\n", "\n", " ", ""]57 )58 59 def extract_text_from_file(self, file) -> str:60 """Extract text from uploaded file"""61 try:62 if file.type == "text/plain":63 return str(file.read(), "utf-8")64 65 elif file.type == "application/pdf":66 pdf_reader = PyPDF2.PdfReader(BytesIO(file.read()))67 text = ""68 for page in pdf_reader.pages:69 text += page.extract_text() + "\n"70 return text71 72 elif file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":73 doc = docx.Document(BytesIO(file.read()))74 text = ""75 for paragraph in doc.paragraphs:76 text += paragraph.text + "\n"77 return text78 79 else:80 st.error(f"Unsupported file type: {file.type}")81 return ""82 83 except Exception as e:84 st.error(f"Error processing file {file.name}: {str(e)}")85 return ""86 87 def chunk_documents(self, text: str) -> List[str]:88 """Split text into chunks"""89 if not text.strip():90 return []91 92 chunks = self.text_splitter.split_text(text)93 # Filter out very short chunks94 chunks = [chunk for chunk in chunks if len(chunk.strip()) > 50]95 return chunks96 97class VectorDatabase:98 def __init__(self, model_name="all-MiniLM-L6-v2"):99 self.model_name = model_name100 self.model = None101 self.dimension = 384 # Dimension for all-MiniLM-L6-v2102 self.index = None103 self.chunks = []104 self.metadata = []105 self.initialize_model()106 107 def initialize_model(self):108 """Initialize the sentence transformer model"""109 try:110 with st.spinner("Loading embedding model..."):111 self.model = SentenceTransformer(self.model_name)112 self.index = faiss.IndexFlatIP(self.dimension)113 except Exception as e:114 st.error(f"Error initializing model: {str(e)}")115 116 def add_documents(self, chunks: List[str], filename: str):117 """Add document chunks to vector database"""118 if not chunks:119 return120 121 try:122 # Generate embeddings123 embeddings = self.model.encode(chunks, show_progress_bar=True)124 125 # Normalize embeddings for cosine similarity126 faiss.normalize_L2(embeddings)127 128 # Add to FAISS index129 self.index.add(embeddings)130 131 # Store chunks and metadata132 self.chunks.extend(chunks)133 self.metadata.extend([134 {135 "filename": filename, 136 "chunk_id": i,137 "content_preview": chunk[:100] + "..." if len(chunk) > 100 else chunk138 } 139 for i, chunk in enumerate(chunks)140 ])141 142 except Exception as e:143 st.error(f"Error adding documents: {str(e)}")144 145 def search(self, query: str, k: int = 3) -> List[Dict]:146 """Search for similar chunks"""147 if not self.chunks or not query.strip():148 return []149 150 try:151 # Generate query embedding152 query_embedding = self.model.encode([query])153 faiss.normalize_L2(query_embedding)154 155 # Search in FAISS156 scores, indices = self.index.search(query_embedding, k)157 158 # Return results159 results = []160 for i, (score, idx) in enumerate(zip(scores[0], indices[0])):161 if idx < len(self.chunks) and score > 0.1: # Threshold for relevance162 results.append({163 "chunk": self.chunks[idx],164 "score": float(score),165 "metadata": self.metadata[idx]166 })167 return results168 169 except Exception as e:170 st.error(f"Error searching: {str(e)}")171 return []172 173 def get_stats(self):174 """Get database statistics"""175 if not self.chunks:176 return {"total_chunks": 0, "total_documents": 0}177 178 unique_docs = set([m['filename'] for m in self.metadata])179 return {180 "total_chunks": len(self.chunks),181 "total_documents": len(unique_docs),182 "documents": list(unique_docs)183 }184 185class ProcurementRAG:186 def __init__(self, vector_db: VectorDatabase):187 self.vector_db = vector_db188 self.client = None189 self.initialize_groq()190 191 def initialize_groq(self):192 """Initialize Groq client"""193 groq_api_key = st.secrets.get("GROQ_API_KEY") or os.getenv("GROQ_API_KEY")194 195 if not groq_api_key:196 st.error("🔑 Please set your Groq API key in Hugging Face Secrets or environment variables")197 st.info("Get your free API key from: https://console.groq.com/")198 return199 200 try:201 self.client = Groq(api_key=groq_api_key)202 except Exception as e:203 st.error(f"Error initializing Groq client: {str(e)}")204 205 def generate_response(self, query: str) -> str:206 """Generate response using RAG with Groq"""207 if not self.client:208 return "❌ Groq client not initialized. Please check your API key."209 210 # Retrieve relevant chunks211 relevant_chunks = self.vector_db.search(query, k=3)212 213 if not relevant_chunks:214 return "❌ No relevant information found in the uploaded documents. Please upload procurement-related documents first."215 216 # Create context from retrieved chunks217 context = "\n\n".join([218 f"Document: {chunk['metadata']['filename']}\nContent: {chunk['chunk']}" 219 for chunk in relevant_chunks220 ])221 222 # Create procurement-specific prompt223 prompt = f"""224 You are an expert procurement assistant with deep knowledge of procurement processes, policies, and best practices.225 226 Context from uploaded documents:227 {context}228 229 User Question: {query}230 231 Instructions:232 1. Provide a comprehensive answer based PRIMARILY on the context provided233 2. If the context is insufficient, supplement with general procurement knowledge234 3. Structure your response clearly with bullet points when appropriate235 4. Include specific references to the documents when possible236 5. Focus on actionable insights and practical guidance237 238 Answer:239 """240 241 try:242 response = self.client.chat.completions.create(243 model="llama3-8b-8192",244 messages=[245 {"role": "system", "content": "You are an expert procurement assistant providing accurate, helpful guidance on procurement processes and policies."},246 {"role": "user", "content": prompt}247 ],248 temperature=0.3,249 max_tokens=1500,250 top_p=0.9251 )252 253 return response.choices[0].message.content254 255 except Exception as e:256 return f"❌ Error generating response: {str(e)}"257 258def main():259 # Header260 st.markdown("""261 <div class="main-header">262 <h1>🏢 Procurement RAG Assistant</h1>263 <p>Upload your procurement documents and get instant answers!</p>264 </div>265 """, unsafe_allow_html=True)266 267 # Initialize session state268 if 'vector_db' not in st.session_state:269 st.session_state.vector_db = VectorDatabase()270 st.session_state.doc_processor = DocumentProcessor()271 st.session_state.rag_system = ProcurementRAG(st.session_state.vector_db)272 273 if 'messages' not in st.session_state:274 st.session_state.messages = []275 276 # Sidebar for document upload277 with st.sidebar:278 st.header("📄 Document Management")279 280 # API Key Status281 groq_api_key = st.secrets.get("GROQ_API_KEY") or os.getenv("GROQ_API_KEY")282 if groq_api_key:283 st.success("✅ Groq API Key configured")284 else:285 st.error("❌ Groq API Key not found")286 st.info("Add GROQ_API_KEY to your Hugging Face Space secrets")287 288 st.markdown("---")289 290 # File upload section291 st.markdown('<div class="upload-section">', unsafe_allow_html=True)292 st.subheader("Upload Documents")293 294 uploaded_files = st.file_uploader(295 "Choose procurement documents",296 accept_multiple_files=True,297 type=['txt', 'pdf', 'docx'],298 help="Upload procurement policies, procedures, contracts, or guidelines"299 )300 301 if uploaded_files:302 for file in uploaded_files:303 col1, col2 = st.columns([3, 1])304 with col1:305 st.text(f"📄 {file.name}")306 with col2:307 if st.button("Process", key=f"process_{file.name}"):308 process_file(file)309 310 st.markdown('</div>', unsafe_allow_html=True)311 312 # Database statistics313 stats = st.session_state.vector_db.get_stats()314 if stats["total_chunks"] > 0:315 st.markdown("---")316 st.markdown('<div class="stats-container">', unsafe_allow_html=True)317 st.subheader("📊 Database Stats")318 st.metric("Total Documents", stats["total_documents"])319 st.metric("Total Chunks", stats["total_chunks"])320 321 if stats["documents"]:322 st.subheader("📋 Uploaded Documents")323 for doc in stats["documents"]:324 st.text(f"• {doc}")325 326 st.markdown('</div>', unsafe_allow_html=True)327 328 # Main chat interface329 st.header("💬 Ask Questions")330 331 # Sample questions332 if not st.session_state.messages:333 st.info("💡 **Sample Questions:**")334 sample_questions = [335 "What is our vendor approval process?",336 "What are the spending limits for different approval levels?",337 "How should we handle contract negotiations?",338 "What documentation is required for purchases over $10,000?",339 "What are our sustainability requirements for suppliers?"340 ]341 342 cols = st.columns(len(sample_questions))343 for i, question in enumerate(sample_questions):344 with cols[i]:345 if st.button(f"💡 {question[:30]}...", key=f"sample_{i}"):346 st.session_state.messages.append({"role": "user", "content": question})347 st.rerun()348 349 # Display chat history350 for message in st.session_state.messages:351 with st.chat_message(message["role"]):352 st.markdown(message["content"])353 354 # Chat input355 if prompt := st.chat_input("Ask about procurement policies, procedures, or processes..."):356 # Add user message357 st.session_state.messages.append({"role": "user", "content": prompt})358 with st.chat_message("user"):359 st.markdown(prompt)360 361 # Generate response362 with st.chat_message("assistant"):363 with st.spinner("🤔 Analyzing documents and generating response..."):364 response = st.session_state.rag_system.generate_response(prompt)365 st.markdown(response)366 st.session_state.messages.append({"role": "assistant", "content": response})367 368 # Clear chat button369 if st.session_state.messages:370 if st.button("🗑️ Clear Chat History"):371 st.session_state.messages = []372 st.rerun()373 374def process_file(file):375 """Process uploaded file"""376 with st.spinner(f"Processing {file.name}..."):377 try:378 # Extract text379 text = st.session_state.doc_processor.extract_text_from_file(file)380 381 if not text.strip():382 st.error(f"No text could be extracted from {file.name}")383 return384 385 # Chunk documents386 chunks = st.session_state.doc_processor.chunk_documents(text)387 388 if not chunks:389 st.error(f"No valid chunks created from {file.name}")390 return391 392 # Add to vector database393 st.session_state.vector_db.add_documents(chunks, file.name)394 395 st.success(f"✅ Successfully processed {file.name}")396 st.info(f"📊 Added {len(chunks)} chunks to knowledge base")397 398 except Exception as e:399 st.error(f"❌ Error processing {file.name}: {str(e)}")400 401if __name__ == "__main__":402 main()