prernajeet01/Reasoning_AI_Agent
0
1import gradio as gr2import os3import tempfile4import pandas as pd5import boto36from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader, UnstructuredPowerPointLoader, UnstructuredExcelLoader, TextLoader7from langchain.text_splitter import RecursiveCharacterTextSplitter8from langchain_community.embeddings import OpenAIEmbeddings9from langchain_community.vectorstores import FAISS10from langchain.chains import RetrievalQA11from langchain_community.chat_models import BedrockChat12from langchain_openai import ChatOpenAI13from langchain.schema import Document14from pathlib import Path15from typing import List, Union16import logging17 18# Optional OCR support19try:20 from pdf2image import convert_from_path21 import pytesseract22 OCR_AVAILABLE = True23except ImportError:24 OCR_AVAILABLE = False25 26# Set up logging27logging.basicConfig(28 level=logging.INFO,29 format='%(asctime)s - %(levelname)s - %(message)s'30)31 32def get_api_keys():33 """Get API keys from Hugging Face Spaces secrets."""34 aws_access_key = os.environ.get("AWS_ACCESS_KEY_ID")35 aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")36 aws_region = os.environ.get("AWS_REGION", "us-east-1") # Default to us-east-1 if not specified37 openai_key = os.environ.get("OPENAI_API_KEY")38 39 if not aws_access_key or not aws_secret_key or not openai_key:40 return {41 "status": "error",42 "message": "Please set AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, and OPENAI_API_KEY in your Hugging Face Space secrets."43 }44 45 return {46 "status": "success",47 "aws_access_key": aws_access_key,48 "aws_secret_key": aws_secret_key,49 "aws_region": aws_region,50 "openai_key": openai_key51 }52 53class AuditAgent:54 def __init__(self, model_name, provider):55 self.model_name = model_name56 self.provider = provider57 self.document_store = None58 59 # Initialize text splitter60 self.text_splitter = RecursiveCharacterTextSplitter(61 chunk_size=1000,62 chunk_overlap=20063 )64 65 # Get API keys66 api_keys = get_api_keys()67 if api_keys["status"] == "error":68 raise ValueError(api_keys["message"])69 70 # Initialize embeddings71 self.embeddings = OpenAIEmbeddings(openai_api_key=api_keys["openai_key"])72 73 if provider == "bedrock":74 # Initialize AWS Bedrock client75 try:76 self.bedrock_client = boto3.client(77 service_name="bedrock-runtime",78 aws_access_key_id=api_keys["aws_access_key"],79 aws_secret_access_key=api_keys["aws_secret_key"],80 region_name=api_keys["aws_region"]81 )82 83 # Use BedrockChat with the same interface84 self.llm = BedrockChat(85 client=self.bedrock_client,86 model_id="anthropic.claude-3-sonnet-20240229-v1:0",87 model_kwargs={"temperature": 0.2}88 )89 except Exception as e:90 logging.error(f"Bedrock initialization error: {str(e)}")91 raise ValueError(f"Bedrock initialization error: {str(e)}")92 elif provider == "openai":93 self.llm = ChatOpenAI(94 model_name=model_name,95 openai_api_key=api_keys["openai_key"],96 temperature=0.297 )98 else:99 raise ValueError(f"Unsupported provider: {provider}")100 101 def process_query(self, query):102 """Process a general query or numerical problem."""103 if not query.strip():104 return "Please provide a non-empty query."105 106 system_prompt = """You are an expert auditor assistant. Provide clear, detailed responses to audit-related queries. 107 For numerical problems, show your calculations step by step. Always consider relevant accounting standards and auditing principles."""108 109 try:110 if self.provider == "bedrock":111 # Handle the response format for BedrockChat112 response = self.llm.invoke(113 f"{system_prompt}\n\nUser: {query}\nAssistant:"114 )115 # Extract the content based on response structure116 return response.content if hasattr(response, 'content') else str(response)117 elif self.provider == "openai":118 response = self.llm.invoke(119 [120 {"role": "system", "content": system_prompt},121 {"role": "user", "content": query}122 ]123 )124 return response.content125 else:126 raise ValueError(f"Unsupported provider: {self.provider}")127 except Exception as e:128 return f"Error processing query: {str(e)}"129 130 def process_documents(self, file_paths):131 """Process multiple documents and return results."""132 results = {}133 134 for file_path in file_paths:135 try:136 # Get file extension137 file_ext = os.path.splitext(file_path.lower())[1]138 139 # Validate file extension140 supported_exts = ['.pdf', '.docx', '.pptx', '.xlsx', '.xls', '.txt']141 if file_ext not in supported_exts:142 results[file_path] = f"Unsupported file type: {file_ext}"143 continue144 145 # Read file content146 with open(file_path, 'rb') as f:147 content = f.read()148 149 # Process document based on type150 documents = self.process_document(content, file_ext)151 152 # Create vector store with the documents153 if documents:154 if not self.document_store:155 self.document_store = FAISS.from_documents(documents, self.embeddings)156 else:157 # Add to existing store158 self.document_store.add_documents(documents)159 160 num_chunks = len(documents)161 results[file_path] = f"Success ({num_chunks} chunks extracted)"162 else:163 results[file_path] = "No content could be extracted"164 except Exception as e:165 logging.error(f"Error processing document {file_path}: {str(e)}")166 results[file_path] = str(e)167 168 return results169 170 def process_document(self, content, doc_type):171 """Process document content based on type."""172 with tempfile.NamedTemporaryFile(delete=False, suffix=doc_type) as temp_file:173 temp_file.write(content)174 temp_file_path = temp_file.name175 176 try:177 documents = self.load_document(temp_file_path)178 return self.split_documents(documents)179 finally:180 if os.path.exists(temp_file_path):181 os.unlink(temp_file_path)182 183 def load_document(self, file_path):184 """Load document using appropriate loader with OCR fallback for PDFs."""185 file_path = Path(file_path)186 suffix = file_path.suffix.lower()187 188 if suffix == '.pdf':189 # Try normal PDF loading first190 try:191 loader = PyPDFLoader(str(file_path))192 documents = loader.load()193 if not any(doc.page_content.strip() for doc in documents):194 raise ValueError("No text content found")195 return documents196 except Exception as e:197 logging.warning(f"Standard PDF extraction failed: {str(e)}")198 # If normal loading fails, try OCR199 if OCR_AVAILABLE:200 logging.info("Attempting PDF extraction with OCR")201 return self._process_pdf_with_ocr(file_path)202 else:203 raise ValueError("PDF extraction failed and OCR is not available")204 elif suffix == '.docx':205 try:206 # Enhanced error handling for Word documents207 loader = Docx2txtLoader(str(file_path))208 documents = loader.load()209 210 # Verify content was extracted211 if not documents or not any(doc.page_content.strip() for doc in documents):212 raise ValueError("No content extracted from Word document")213 214 return documents215 except Exception as e:216 logging.error(f"Word document processing error: {str(e)}")217 raise ValueError(f"Failed to process Word document: {str(e)}")218 elif suffix == '.pptx':219 loader = UnstructuredPowerPointLoader(str(file_path))220 return loader.load()221 elif suffix in ['.xlsx', '.xls']:222 loader = UnstructuredExcelLoader(str(file_path))223 return loader.load()224 elif suffix == '.txt':225 loader = TextLoader(str(file_path))226 return loader.load()227 else:228 raise ValueError(f"Unsupported file type: {suffix}")229 230 def _process_pdf_with_ocr(self, file_path):231 """Process PDF with OCR using Tesseract."""232 if not OCR_AVAILABLE:233 raise ImportError("pdf2image and pytesseract required for OCR processing")234 235 documents = []236 images = convert_from_path(str(file_path))237 238 for i, image in enumerate(images):239 text = pytesseract.image_to_string(image)240 if text.strip():241 documents.append(Document(242 page_content=text,243 metadata={"source": str(file_path), "page": i + 1}244 ))245 246 return documents247 248 def split_documents(self, documents):249 """Split documents into chunks."""250 return self.text_splitter.split_documents(documents)251 252 def query_documents(self, query):253 """Query the processed documents."""254 if not self.document_store:255 return "Please upload and process documents first"256 257 if not query.strip():258 return "Please provide a non-empty query."259 260 try:261 qa_chain = RetrievalQA.from_chain_type(262 llm=self.llm,263 chain_type="stuff",264 retriever=self.document_store.as_retriever(),265 return_source_documents=True266 )267 268 response = qa_chain({"query": query})269 270 result = response['result']271 source_docs = response.get('source_documents', [])272 273 if source_docs:274 result += "\n\n**Sources:**\n"275 for i, doc in enumerate(source_docs, 1):276 result += f"{i}. {doc.metadata.get('source', 'Unknown source')}, page {doc.metadata.get('page', 'N/A')}\n"277 278 return result279 except Exception as e:280 return f"Error querying documents: {str(e)}"281 282# Updated LLM configurations - replaced openorca-mini with o3-mini283llm_configs = {284 "claude-3-sonnet": {285 "name": "anthropic.claude-3-sonnet-20240229-v1:0",286 "provider": "bedrock",287 "description": "Balanced performance (AWS Bedrock)"288 },289 "gpt-4": {290 "name": "gpt-4",291 "provider": "openai",292 "description": "Advanced reasoning"293 },294 "gpt-3.5-turbo": {295 "name": "gpt-3.5-turbo",296 "provider": "openai",297 "description": "Fast responses"298 },299 "o3-mini": {300 "name": "o3-mini",301 "provider": "openai",302 "description": "Compact OpenAI model"303 }304}305 306def create_interface():307 # Check API keys first308 api_keys = get_api_keys()309 if api_keys["status"] == "error":310 with gr.Blocks(theme=gr.themes.Base()) as demo:311 gr.Markdown("# ⚠️ Configuration Error")312 gr.Markdown(api_keys["message"])313 gr.Markdown("""314 To set up your Hugging Face Space:315 1. Go to your Space's Settings316 2. Add your API keys as secrets:317 - AWS_ACCESS_KEY_ID318 - AWS_SECRET_ACCESS_KEY319 - AWS_REGION320 - OPENAI_API_KEY321 3. Restart your Space322 """)323 return demo324 325 # Initialize agents dictionary - will be initialized on demand326 audit_agents = {}327 328 with gr.Blocks(theme=gr.themes.Base()) as demo:329 gr.Markdown("# 🔍 Amy - Your Audit Copilot")330 331 # Status indicator for initialization and operations332 status_message = gr.Textbox(label="Status", value="Ready")333 334 # Document processing section - moved above model selection335 gr.Markdown("## 📑 Document Processing")336 with gr.Row():337 file_upload = gr.File(338 file_count="multiple", 339 label="Upload Audit Documents (PDF, DOCX, PPTX, TXT, XLSX)",340 type="filepath"341 )342 upload_button = gr.Button("Process Documents")343 upload_output = gr.Textbox(label="Processing Status", lines=10)344 345 # Use tabs for model selection346 with gr.Tabs() as model_tabs:347 model_tab_dict = {}348 for model_id, config in llm_configs.items():349 with gr.Tab(f"{model_id} - {config['description']}") as tab:350 model_tab_dict[model_id] = tab351 352 with gr.Tabs() as feature_tabs:353 # Chat interface with history354 with gr.Tab("💬 Conversation"):355 chat_history = gr.Chatbot(height=400)356 chat_input = gr.Textbox(357 lines=3, 358 label="Ask your audit question",359 placeholder="Enter your question here..."360 )361 chat_clear = gr.Button("Clear Chat")362 chat_button = gr.Button("Send")363 364 with gr.Tab("🔢 Numerical Problem"):365 problem_input = gr.Textbox(366 lines=5,367 label="Describe the Problem",368 placeholder="Enter your numerical audit problem..."369 )370 solve_button = gr.Button("Solve")371 solution_output = gr.Markdown(label="Solution")372 373 # Document query tab374 with gr.Tab("🔍 Document Query"):375 query_input = gr.Textbox(376 lines=3,377 label="Query Documents",378 placeholder="Ask about your uploaded documents..."379 )380 query_button = gr.Button("Query")381 query_output = gr.Markdown(label="Response")382 383 # Track the selected model384 selected_model = gr.State("claude-3-sonnet")385 386 # Update selected model when tabs change387 def update_selected_model(evt: gr.SelectData):388 model_ids = list(llm_configs.keys())389 if evt.index < len(model_ids):390 return model_ids[evt.index]391 return "claude-3-sonnet" # Default392 393 model_tabs.select(update_selected_model, outputs=[selected_model])394 395 # Get or initialize agent and return both agent and status message396 def get_or_initialize_agent(model_name):397 """Initialize an agent if not already initialized and return a status message"""398 init_message = f"Initializing {model_name}..."399 400 # If agent already exists, return it with a status message401 if model_name in audit_agents:402 return audit_agents[model_name], f"{model_name} ready"403 404 # Try to initialize the agent405 try:406 config = llm_configs[model_name]407 logging.info(init_message)408 agent = AuditAgent(config["name"], config["provider"])409 audit_agents[model_name] = agent410 success_message = f"{model_name} initialized successfully"411 logging.info(success_message)412 return agent, success_message413 except Exception as e:414 error_message = f"Error initializing {model_name}: {str(e)}"415 logging.error(error_message)416 return None, error_message417 418 # Handle chat with history419 def respond_to_chat(message, history, model_name):420 if not message.strip():421 return "", history422 423 # Get or initialize agent424 agent, init_status = get_or_initialize_agent(model_name)425 426 # If initialization failed427 if agent is None:428 history.append((message, f"Could not initialize {model_name}. {init_status}"))429 return "", history, f"Error: {init_status}"430 431 # Process the query432 try:433 result = agent.process_query(message)434 history.append((message, result))435 return "", history, f"Response from {model_name}"436 except Exception as e:437 error_msg = f"Error: {str(e)}"438 history.append((message, error_msg))439 return "", history, error_msg440 441 # Clear chat history442 def clear_chat_history():443 return [], "Chat history cleared"444 445 # Handle numerical problem446 def handle_problem(problem, model_name):447 if not problem.strip():448 return "Please provide a problem description", "No problem entered"449 450 status = f"Solving problem with {model_name}..."451 452 # Get or initialize agent453 agent, init_status = get_or_initialize_agent(model_name)454 455 # If initialization failed456 if agent is None:457 return f"Could not initialize {model_name}. {init_status}", init_status458 459 # Process the problem460 try:461 result = agent.process_query(problem)462 return result, f"Problem solved with {model_name}"463 except Exception as e:464 error_msg = f"Error solving problem: {str(e)}"465 return error_msg, error_msg466 467 # Improved file upload handler for multiple files468 def handle_file_upload(file_paths, model_name):469 if not file_paths:470 return "No files uploaded. Please upload files."471 472 # Get or initialize agent473 agent, init_status = get_or_initialize_agent(model_name)474 475 # If initialization failed476 if agent is None:477 return init_status478 479 logging.info(f"Processing {len(file_paths)} files")480 481 # Process all documents482 try:483 results = agent.process_documents(file_paths)484 485 # Format results486 output_lines = ["## Document Processing Results"]487 for file_path, status in results.items():488 file_name = os.path.basename(file_path)489 if "Success" in status:490 output_lines.append(f"✓ {file_name}: {status}")491 else:492 output_lines.append(f"❌ {file_name}: {status}")493 494 if any("Success" in status for status in results.values()):495 output_lines.append("\n✅ Documents are ready for querying!")496 497 return "\n".join(output_lines)498 except Exception as e:499 logging.error(f"File upload error: {str(e)}")500 return f"Error processing files: {str(e)}"501 502 # Handle document query503 def handle_query(query, model_name):504 if not query.strip():505 return "Please provide a query", "No query entered"506 507 status = f"Querying documents with {model_name}..."508 509 # Get or initialize agent510 agent, init_status = get_or_initialize_agent(model_name)511 512 # If initialization failed513 if agent is None:514 return f"Could not initialize {model_name}. {init_status}", init_status515 516 # Query the documents517 try:518 result = agent.query_documents(query)519 return result, f"Documents queried with {model_name}"520 except Exception as e:521 error_msg = f"Error querying documents: {str(e)}"522 return error_msg, error_msg523 524 # Set up event handlers525 chat_button.click(526 respond_to_chat,527 inputs=[chat_input, chat_history, selected_model],528 outputs=[chat_input, chat_history, status_message]529 )530 531 chat_clear.click(532 clear_chat_history,533 outputs=[chat_history, status_message]534 )535 536 solve_button.click(537 handle_problem,538 inputs=[problem_input, selected_model],539 outputs=[solution_output, status_message]540 )541 542 upload_button.click(543 handle_file_upload,544 inputs=[file_upload, selected_model],545 outputs=[upload_output]546 )547 548 query_button.click(549 handle_query,550 inputs=[query_input, selected_model],551 outputs=[query_output, status_message]552 )553 554 return demo555 556if __name__ == "__main__":557 demo = create_interface()558 demo.launch(share=True)