initial commit

README.md

@@ -1,13 +1,73 @@
----
-title: TopKimages
-emoji: 🏢
-colorFrom: blue
-colorTo: purple
-sdk: gradio
-sdk_version: 5.16.0
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Sustainable Content Moderation
+emoji: 🌍
+colorFrom: green
+colorTo: blue
+sdk: gradio
+sdk_version: 5.16.0
+app_file: app.py
+pinned: false
+---
+# Multi-Model multi-agent Analysis System
+
+A multi-agent system for analyzing pump equipment images using state-of-the-art AI models.
+
+## Overview
+This application uses multiple AI agents to:
+- Process user queries about pump equipment
+- Analyze equipment images
+- Search relevant technical context
+- Generate comprehensive analysis reports
+
+## Usage
+1. Enter your query about pump equipment
+2. (Optional) Add any specific constraints
+3. Upload equipment images (supported formats: JPG, JPEG, PNG)
+4. Select number of top results to show
+5. Choose report format (summary/detailed)
+6. Click "Analyze" to process
+
+## Technical Details
+- Built with Gradio 5.16.0
+- Uses multi-model approach:
+  - Lightweight models for initial processing
+  - Advanced models for detailed analysis
+- Implements memory-efficient processing
+- Supports batch image processing
+
+## Limitations
+- Maximum image size: 5MB
+- Maximum resolution: 2048x2048
+- Maximum images per request: 10
+- Query length limit: 500 characters
+
+## Models Used
+- Image Captioning: BLIP and BLIP-2
+- Image Classification: ResNet and ViT
+- Text Processing: LaMini-Flan-T5
+
+## Requirements
+- Python 3.8+
+- See requirements.txt for full dependencies
+
+## Local Development
+`pip install -r requirements.txt python app.py`
+
+## Deployment
+This app is optimized for HuggingFace Spaces deployment.
+
+## Error Handling
+- Input validation with clear error messages
+- Resource monitoring and automatic cleanup
+- Graceful error recovery
+
+## Credits
+- Built using HuggingFace's model hub
+- Powered by Gradio interface
+- Uses Wikipedia API for context gathering
+
+## License
+MIT License
+
+## Support
+For issues or questions, please open a GitHub issue.

agents/assembler_agent.py

@@ -0,0 +1,204 @@
+from typing import Dict, List
+from models.model_config import ModelConfig
+from models.data_models import AssemblerInput
+from .base_agent import BaseAgent
+import datetime 
+
+class AssemblerAgent(BaseAgent):
+    def __init__(self, name: str = "AssemblerAgent"):
+        super().__init__(name)
+        self.final_report: Dict = {}
+
+
+    def reason(self, input_data: AssemblerInput) -> List[str]:
+        """
+        Plan how to assemble the final report from all agent results
+        """
+        thoughts = []
+        
+        try:
+            # Analyze available inputs
+            thoughts.append("Analyzing inputs from all agents:")
+            thoughts.append(f"- User input processing results available: {bool(input_data.user_input_results)}")
+            thoughts.append(f"- Context learning results available: {bool(input_data.context_results)}")
+            thoughts.append(f"- Image analysis results available: {bool(input_data.image_results)}")
+            
+            # Plan report structure
+            thoughts.append("\nPlanning report structure:")
+            thoughts.append("1. User Query Summary")
+            thoughts.append("2. Context Analysis")
+            thoughts.append("3. Image Analysis Results")
+            thoughts.append("4. Model Performance Metrics")
+            thoughts.append("5. Final Recommendations")
+            
+            # Consider report format
+            thoughts.append(f"\nReport Format: {input_data.report_format}")
+            if input_data.report_format == "detailed":
+                thoughts.append("- Will include full model decisions")
+                thoughts.append("- Will include confidence scores")
+                thoughts.append("- Will include processing statistics")
+            else:
+                thoughts.append("- Will provide condensed summary")
+                thoughts.append("- Will focus on key findings")
+            
+            # Store thoughts in state
+            self.state.thoughts.extend(thoughts)
+            self.logger.info("Reasoning complete for report assembly")
+            
+            return thoughts
+            
+        except Exception as e:
+            error_msg = f"Error during assembly reasoning: {str(e)}"
+            self.state.errors.append(error_msg)
+            self.logger.error(error_msg)
+            return thoughts
+        
+    def execute(self, input_data: AssemblerInput) -> Dict:
+        """
+        Assemble final report from all agent results
+        """
+        try:
+            if not self.validate(input_data):
+                return {
+                    'status': 'error',
+                    'error': self.state.errors[-1]
+                }
+
+            report = {
+                'summary': {
+                    'user_query': {},
+                    'context_analysis': {},
+                    'image_analysis': {},
+                    'recommendations': []
+                },
+                'details': {
+                    'model_decisions': {},
+                    'processing_stats': {},
+                    'confidence_scores': {}
+                },
+                'metadata': {
+                    'report_format': input_data.report_format,
+                    'timestamp': datetime.now().isoformat()
+                },
+                'status': 'processing'
+            }
+
+            # Process user input results
+            if input_data.user_input_results:
+                report['summary']['user_query'] = {
+                    'original_query': input_data.user_input_results.get('query', ''),
+                    'constraints': input_data.user_input_results.get('constraints', []),
+                    'intent': input_data.user_input_results.get('intent', '')
+                }
+
+            # Process context results
+            if input_data.context_results:
+                report['summary']['context_analysis'] = {
+                    'key_findings': input_data.context_results.get('summaries', {}),
+                    'relevant_keywords': list(input_data.context_results.get('keywords', set())),
+                    'sources': list(input_data.context_results.get('gathered_context', {}).keys())
+                }
+
+            # Process image analysis results
+            if input_data.image_results:
+                report['summary']['image_analysis'] = {
+                    'selected_images': input_data.image_results.get('selected_images', []),
+                    'analysis_summary': {
+                        path: results['caption']['text'] 
+                        for path, results in input_data.image_results.get('analyzed_images', {}).items()
+                    }
+                }
+
+            # Add detailed information if requested
+            if input_data.report_format == "detailed":
+                report['details']['model_decisions'] = {
+                    'context_models': input_data.context_results.get('model_decisions', {}),
+                    'image_models': input_data.image_results.get('model_decisions', {})
+                }
+                report['details']['processing_stats'] = {
+                    'context_processing': input_data.context_results.get('model_decisions', {}).get('processing_stats', {}),
+                    'image_processing': input_data.image_results.get('model_decisions', {}).get('processing_stats', {})
+                }
+
+            # Generate recommendations
+            report['summary']['recommendations'] = self._generate_recommendations(
+                report['summary']['context_analysis'],
+                report['summary']['image_analysis']
+            )
+
+            report['status'] = 'success'
+            self.final_report = report
+            
+            # Log decision
+            self.log_decision({
+                'action': 'report_assembly',
+                'format': input_data.report_format,
+                'sections_completed': list(report['summary'].keys())
+            })
+
+            return report
+
+        except Exception as e:
+            error_msg = f"Error executing report assembly: {str(e)}"
+            self.state.errors.append(error_msg)
+            self.logger.error(error_msg)
+            return {'status': 'error', 'error': error_msg}
+
+
+    def _generate_recommendations(self, context_analysis: Dict, image_analysis: Dict) -> List[str]:
+        """
+        Generate recommendations based on context and image analysis
+        
+        Args:
+            context_analysis: Results from context learning
+            image_analysis: Results from image analysis
+            
+        Returns:
+            List[str]: List of recommendations
+        """
+        try:
+            recommendations = []
+            
+            # Check if we have sufficient data
+            if not context_analysis or not image_analysis:
+                return ["Insufficient data to generate recommendations"]
+            
+            # Analyze context findings
+            if context_analysis.get('key_findings'):
+                recommendations.append("Based on context analysis:")
+                for source, finding in context_analysis['key_findings'].items():
+                    if finding:  # Check if finding exists
+                        recommendations.append(f"- {finding}")
+            
+            # Analyze image findings
+            if image_analysis.get('selected_images'):
+                recommendations.append("\nBased on image analysis:")
+                recommendations.append(f"- Found {len(image_analysis['selected_images'])} relevant images")
+                
+                # Add specific image recommendations
+                if image_analysis.get('analysis_summary'):
+                    for img_path, caption in image_analysis['analysis_summary'].items():
+                        if caption:  # Check if caption exists
+                            recommendations.append(f"- {caption}")
+            
+            # Combine findings for final recommendations
+            recommendations.append("\nKey Recommendations:")
+            if context_analysis.get('relevant_keywords'):
+                keywords = context_analysis['relevant_keywords'][:5]  # Top 5 keywords
+                recommendations.append(f"- Focus areas identified: {', '.join(keywords)}")
+            
+            # Add source credibility note
+            if context_analysis.get('sources'):
+                recommendations.append(f"- Analysis based on {len(context_analysis['sources'])} credible sources")
+            
+            # Add confidence note
+            recommendations.append("- Regular monitoring and updates recommended")
+            
+            self.logger.info("Generated recommendations successfully")
+            return recommendations
+            
+        except Exception as e:
+            error_msg = f"Error generating recommendations: {str(e)}"
+            self.logger.error(error_msg)
+            return ["Error generating recommendations. Please check the detailed report."]
+

agents/base_agent.py

@@ -0,0 +1,99 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+from dataclasses import dataclass, field
+from datetime import datetime
+import logging
+
+from models.data_models import AgentState
+
+
+    
+class BaseAgent(ABC):
+    def __init__(self, name: str):
+        self.name = name
+        self.state = AgentState()
+        self.logger = logging.getLogger(name)
+    
+    @abstractmethod
+    def reason(self, input_data: Any) -> List[str]:
+        """Implement step-by-step reasoning process"""
+        pass
+    
+    @abstractmethod
+    def execute(self, input_data: Any) -> Any:
+        """Implement main execution logic"""
+        pass
+    
+    def validate(self, input_data: Any) -> bool:
+        """
+        Validate input data with basic checks and logging
+        
+        Args:
+            input_data: Any - Data to validate
+            
+        Returns:
+            bool - True if valid, False otherwise
+            
+        Note: Child classes should override this method 
+        with additional specific validation rules
+        """
+        try:
+            # Basic validation: check if input exists
+            if input_data is None:
+                self.state.errors.append("Input data is None")
+                return False
+                
+            # Check if input is empty
+            if isinstance(input_data, (str, list, dict)) and not input_data:
+                self.state.errors.append("Input data is empty")
+                return False
+                
+            self.logger.debug(f"Input validation successful for {self.name}")
+            return True
+            
+        except Exception as e:
+            error_msg = f"Validation error: {str(e)}"
+            self.state.errors.append(error_msg)
+            self.logger.error(error_msg)
+            return False
+
+    
+    def log_decision(self, decision: Dict) -> None:
+        """
+        Track agent decisions with timestamps and metadata
+        
+        Args:
+            decision: Dictionary containing:
+                - action: str - What action was taken
+                - reason: str - Why this decision was made
+                - metadata: Dict - Any additional information
+        """
+        timestamped_decision = {
+            'timestamp': datetime.now().isoformat(),
+            'agent_name': self.name,
+            **decision
+        }
+        self.state.decisions.append(timestamped_decision)
+        self.logger.info(f"Decision logged: {timestamped_decision}")
+
+    
+    def get_state(self) -> Dict:
+        """
+        Return current agent state as a dictionary
+        
+        Returns:
+            Dict containing:
+                - intent: str - Current agent intent
+                - thoughts: List[str] - Reasoning steps
+                - decisions: List[Dict] - History of decisions
+                - errors: List[str] - Any errors encountered
+                - last_updated: str - Timestamp of last state change
+        """
+        return {
+            'intent': self.state.intent,
+            'thoughts': self.state.thoughts.copy(),
+            'decisions': self.state.decisions.copy(),
+            'errors': self.state.errors.copy(),
+            'last_updated': datetime.now().isoformat()
+        }
+

agents/context_agent.py

@@ -0,0 +1,246 @@
+from typing import Dict, List, Optional
+import wikipedia
+from transformers import pipeline
+from models.model_config import ModelConfig
+from models.data_models import ContextInput
+from .base_agent import BaseAgent
+
+class ContextLearnAgent(BaseAgent):
+
+    def __init__(self, name: str = "ContextLearnAgent"):
+        super().__init__(name)
+        self.learned_context: Dict = {}
+        self.models: Dict[str, ModelConfig] = {
+            'lightweight_summarizer': ModelConfig(
+                name='facebook/bart-large-cnn',
+                type='lightweight',
+                task='summarization'
+            ),
+            'advanced_summarizer': ModelConfig(
+                name='facebook/bart-large-xsum',
+                type='advanced',
+                task='summarization'
+            )
+        }
+
+    def reason(self, input_data: ContextInput) -> List[str]:
+        """
+        Plan the context gathering and learning process
+        """
+        thoughts = []
+        
+        try:
+            # Analyze query requirements
+            thoughts.append(f"Analyzing query: '{input_data.processed_query}'")
+            
+            # Plan Wikipedia search strategy
+            thoughts.append("\nContext Gathering Strategy:")
+            thoughts.append("1. Extract key terms from query")
+            thoughts.append("2. Search Wikipedia for relevant articles")
+            thoughts.append("3. Analyze article relevance")
+            
+            # Plan summarization approach
+            thoughts.append("\nSummarization Strategy:")
+            thoughts.append(f"- Start with {self.models['lightweight_summarizer'].name}")
+            thoughts.append(f"- Use {self.models['advanced_summarizer'].name} for complex content")
+            
+            if input_data.constraints:
+                thoughts.append("\nConstraint Handling:")
+                for constraint in input_data.constraints:
+                    thoughts.append(f"- Will filter content against: {constraint}")
+            
+            # Store thoughts in state
+            self.state.thoughts.extend(thoughts)
+            self.logger.info("Reasoning complete for context learning")
+            
+            return thoughts
+            
+        except Exception as e:
+            error_msg = f"Error during context reasoning: {str(e)}"
+            self.state.errors.append(error_msg)
+            self.logger.error(error_msg)
+            return thoughts
+
+
+    def execute(self, input_data: ContextInput) -> Dict:
+        """
+        Gather and process context information
+        """
+        try:
+            if not self.validate(input_data):
+                return {
+                    'status': 'error',
+                    'error': self.state.errors[-1]
+                }
+
+            results = {
+                'gathered_context': {},
+                'keywords': set(),
+                'summaries': {},
+                'model_decisions': {
+                    'summarizer_used': [],
+                    'processing_stats': {
+                        'articles_found': 0,
+                        'articles_processed': 0
+                    }
+                },
+                'status': 'processing'
+            }
+
+            # Search Wikipedia
+            try:
+                search_results = wikipedia.search(input_data.processed_query, results=5)
+                results['model_decisions']['processing_stats']['articles_found'] = len(search_results)
+                
+                for title in search_results:
+                    try:
+                        page = wikipedia.page(title)
+                        
+                        # Start with lightweight summarization
+                        summary = self._get_summary(page.content, 'lightweight')
+                        
+                        # Use advanced summarizer if content is complex
+                        if len(page.content.split()) > 1000:  # Long article
+                            summary = self._get_summary(page.content, 'advanced')
+                        
+                        results['gathered_context'][title] = {
+                            'url': page.url,
+                            'summary': summary['text'],
+                            'confidence': summary['confidence'],
+                            'model_used': summary['model_used']
+                        }
+                        
+                        # Extract keywords
+                        results['keywords'].update(self._extract_keywords(page.content))
+                        results['model_decisions']['summarizer_used'].append(summary['model_used'])
+                        results['model_decisions']['processing_stats']['articles_processed'] += 1
+                        
+                    except wikipedia.exceptions.DisambiguationError as e:
+                        self.logger.warning(f"Disambiguation for {title}: {str(e)}")
+                    except wikipedia.exceptions.PageError as e:
+                        self.logger.warning(f"Page error for {title}: {str(e)}")
+                        
+            except Exception as e:
+                self.logger.error(f"Wikipedia search error: {str(e)}")
+
+            results['status'] = 'success'
+            self.learned_context = results
+            
+            # Log decision
+            self.log_decision({
+                'action': 'context_gathering',
+                'articles_processed': results['model_decisions']['processing_stats']['articles_processed'],
+                'keywords_found': len(results['keywords'])
+            })
+
+            return results
+
+        except Exception as e:
+            error_msg = f"Error executing context gathering: {str(e)}"
+            self.state.errors.append(error_msg)
+            self.logger.error(error_msg)
+            return {'status': 'error', 'error': error_msg}
+
+
+
+    # def __init__(self, name: str = "ContextLearnAgent"):
+    #     super().__init__(name)
+    #     self.learned_context: Dict = {}
+        
+    # def reason(self, input_data: ContextInput) -> List[str]:
+    #     """
+    #     Analyze what information needs to be gathered about pumps
+        
+    #     Args:
+    #         input_data: Processed user query and constraints
+    #     Returns:
+    #         List[str]: Reasoning steps about required information
+    #     """
+    #     thoughts = []
+        
+    #     try:
+    #         # Analyze search requirements
+    #         thoughts.append(f"Analyzing search query: '{input_data.processed_query}'")
+            
+    #         # Consider search sources
+    #         thoughts.append(f"Planning to search in: {', '.join(input_data.search_sources)}")
+            
+    #         # Consider constraints
+    #         if input_data.constraints:
+    #             thoughts.append(f"Will filter results based on constraints: {input_data.constraints}")
+                
+    #         # Consider result limit
+    #         thoughts.append(f"Will gather up to {input_data.max_results} relevant results")
+            
+    #         # Store thoughts in state
+    #         self.state.thoughts.extend(thoughts)
+    #         self.logger.info("Reasoning complete for context gathering")
+            
+    #         return thoughts
+            
+    #     except Exception as e:
+    #         error_msg = f"Error during context reasoning: {str(e)}"
+    #         self.state.errors.append(error_msg)
+    #         self.logger.error(error_msg)
+    #         return thoughts
+
+            
+    # def execute(self, input_data: ContextInput) -> Dict:
+    #     """
+    #     Gather and process information from specified sources
+        
+    #     Args:
+    #         input_data: Search parameters and constraints
+    #     Returns:
+    #         Dict containing:
+    #             - sources: List[str] - Sources used
+    #             - gathered_info: Dict - Information gathered by source
+    #             - summary: str - Brief summary of findings
+    #             - status: str - Processing status
+    #     """
+    #     try:
+    #         # First validate the input
+    #         if not self.validate(input_data):
+    #             return {
+    #                 'status': 'error',
+    #                 'error': self.state.errors[-1]
+    #             }
+            
+    #         # Initialize results structure
+    #         gathered_info = {
+    #             'sources': input_data.search_sources,
+    #             'gathered_info': {},
+    #             'summary': '',
+    #             'status': 'processing'
+    #         }
+            
+    #         # Process each source (placeholder for actual API calls)
+    #         for source in input_data.search_sources:
+    #             gathered_info['gathered_info'][source] = {
+    #                 'status': 'pending',
+    #                 'content': [],
+    #                 'metadata': {
+    #                     'timestamp': datetime.now().isoformat(),
+    #                     'query': input_data.processed_query
+    #                 }
+    #             }
+            
+    #         # Log the decision
+    #         self.log_decision({
+    #             'action': 'gather_context',
+    #             'sources': input_data.search_sources,
+    #             'query': input_data.processed_query
+    #         })
+            
+    #         # Store in agent's state
+    #         self.learned_context = gathered_info
+    #         gathered_info['status'] = 'success'
+            
+    #         return gathered_info
+            
+    #     except Exception as e:
+    #         error_msg = f"Error executing context gathering: {str(e)}"
+    #         self.state.errors.append(error_msg)
+    #         self.logger.error(error_msg)
+    #         return {'status': 'error', 'error': error_msg}
+

agents/image_agent.py

@@ -0,0 +1,1065 @@
+from typing import Dict, List
+import torch
+from PIL import Image
+from transformers import (
+    BlipProcessor, BlipForConditionalGeneration,
+    Blip2Processor, Blip2ForConditionalGeneration,
+    AutoImageProcessor, ResNetForImageClassification,
+    ViTImageProcessor, ViTForImageClassification
+)
+
+from models.model_config import ModelConfig
+from models.data_models import ImageAnalysisInput
+from config.settings import config
+from .base_agent import BaseAgent
+
+from langchain.llms import HuggingFacePipeline
+from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
+
+
+
+class ImageAnalyzerAgent(BaseAgent):
+    def __init__(self, name: str = "ImageAnalyzerAgent"):
+        super().__init__(name)
+        self.analyzed_images: Dict = {}
+        # self.models: Dict[str, ModelConfig] = {
+        #     'lightweight_caption': ModelConfig('clip-vit-base', 'lightweight', 'caption'),
+        #     'advanced_caption': ModelConfig('blip2-opt', 'advanced', 'caption'),
+        #     'lightweight_classifier': ModelConfig('resnet18', 'lightweight', 'classification'),
+        #     #'advanced_classifier': ModelConfig('vit-large', 'advanced', 'classification')
+            
+        # }
+        self.models: Dict[str, ModelConfig] = {
+            'lightweight_caption': ModelConfig(
+                name='Salesforce/blip-image-captioning-base',
+                type='lightweight',
+                task='caption'
+            ),
+            'advanced_caption': ModelConfig(
+                name='Salesforce/blip2-opt-2.7b',
+                type='advanced',
+                task='caption'
+            ),
+            'lightweight_classifier': ModelConfig(
+                name='microsoft/resnet-50',
+                type='lightweight',
+                task='classification'
+            ),
+            'advanced_classifier': ModelConfig(
+                name='google/vit-base-patch16-224',
+                type='lightweight',
+                task='classification'
+            )
+        }
+
+        # Add LangChain setup
+        self.llm_model_name = "MBZUAI/LaMini-Flan-T5-783M"
+        self._initialize_llm()
+        
+    # def reason(self, input_data: ImageAnalysisInput) -> List[str]:
+    #     """
+    #     Plan the image analysis process for pump-related images
+        
+    #     Args:
+    #         input_data: Contains uploaded images, context, and constraints
+    #     Returns:
+    #         List[str]: Reasoning steps about image analysis plan
+    #     """
+    #     thoughts = []
+        
+    #     try:
+    #         # Consider number of images
+    #         thoughts.append(f"Processing {len(input_data.images)} uploaded images")
+            
+    #         # Consider context from ContextLearnAgent
+    #         if input_data.context:
+    #             thoughts.append("Using learned context to guide image analysis")
+    #             thoughts.append(f"Looking for pump-related features based on context")
+            
+    #         # Consider constraints
+    #         if input_data.constraints:
+    #             thoughts.append(f"Will apply constraints: {input_data.constraints}")
+            
+    #         # Plan analysis steps
+    #         thoughts.append(f"Analysis plan:")
+    #         thoughts.append("1. Verify image accessibility")
+    #         thoughts.append("2. Perform image captioning")
+    #         thoughts.append("3. Match captions with context")
+    #         thoughts.append(f"4. Select top {input_data.top_k} relevant images")
+            
+    #         # Store thoughts in state
+    #         self.state.thoughts.extend(thoughts)
+    #         self.logger.info("Reasoning complete for image analysis")
+            
+    #         return thoughts
+            
+    #     except Exception as e:
+    #         error_msg = f"Error during image analysis reasoning: {str(e)}"
+    #         self.state.errors.append(error_msg)
+    #         self.logger.error(error_msg)
+    #         return thoughts
+
+    # def reason(self, input_data: ImageAnalysisInput) -> List[str]:
+    #     """
+    #     Plan the image analysis process with model selection strategy
+        
+    #     Args:
+    #         input_data: Contains uploaded images, context, and constraints
+    #     Returns:
+    #         List[str]: Reasoning steps about image analysis plan
+    #     """
+    #     thoughts = []
+        
+    #     try:
+    #         # Consider input volume and resources
+    #         thoughts.append(f"Processing {len(input_data.images)} uploaded images")
+            
+    #         # Model selection strategy
+    #         thoughts.append("Planning model selection strategy:")
+    #         thoughts.append("- Will start with lightweight models for efficiency")
+    #         thoughts.append(f"- Using {self.models['lightweight_caption'].name} for initial caption generation")
+    #         thoughts.append(f"- Will switch to {self.models['advanced_caption'].name} if confidence below {self.models['lightweight_caption'].threshold}")
+            
+    #         # Context consideration
+    #         if input_data.context:
+    #             thoughts.append("Using learned context to guide image analysis:")
+    #             thoughts.append("- Will match image captions against context")
+    #             thoughts.append("- Will use context for relevance scoring")
+            
+    #         # Constraints handling
+    #         if input_data.constraints:
+    #             thoughts.append(f"Will apply constraints during analysis:")
+    #             for constraint in input_data.constraints:
+    #                 thoughts.append(f"- {constraint}")
+            
+    #         # Analysis pipeline
+    #         thoughts.append("Analysis pipeline:")
+    #         thoughts.append("1. Initial lightweight model caption generation")
+    #         thoughts.append("2. Confidence check and model escalation if needed")
+    #         thoughts.append("3. Context matching and constraint application")
+    #         thoughts.append(f"4. Selection of top {input_data.top_k} relevant images")
+            
+    #         # Store thoughts in state
+    #         self.state.thoughts.extend(thoughts)
+    #         self.logger.info("Reasoning complete for image analysis")
+            
+    #         return thoughts
+            
+    #     except Exception as e:
+    #         error_msg = f"Error during image analysis reasoning: {str(e)}"
+    #         self.state.errors.append(error_msg)
+    #         self.logger.error(error_msg)
+    #         return thoughts
+    
+    def _initialize_llm(self):
+        """Initialize LangChain components"""
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(self.llm_model_name)
+            model = AutoModelForSeq2SeqLM.from_pretrained(self.llm_model_name)
+            
+            pipe = pipeline(
+                "text2text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                max_length=512,
+                temperature=0.3
+            )
+            
+            self.llm = HuggingFacePipeline(pipeline=pipe)
+            
+            # Create relevance analysis chain
+            relevance_template = """
+            Analyze the relevance between image content and context:
+            
+            Image Caption: {caption}
+            Image Classification: {classification}
+            Context Keywords: {context_keywords}
+            Domain Context: {domain_context}
+            
+            Provide a detailed analysis of:
+            1. Content Relevance
+            2. Domain Alignment
+            3. Context Matching Score (0-1)
+            
+            Analysis:
+            """
+            
+            self.relevance_chain = LLMChain(
+                llm=self.llm,
+                prompt=PromptTemplate(
+                    template=relevance_template,
+                    input_variables=["caption", "classification", 
+                                   "context_keywords", "domain_context"]
+                )
+            )
+            
+        except Exception as e:
+            self.logger.error(f"Error initializing LLM: {str(e)}")
+            raise
+    
+    
+    def reason(self, input_data: ImageAnalysisInput) -> List[str]:
+        """
+        Plan the image analysis process with LLM-enhanced reasoning
+        """
+        thoughts = []
+        
+        try:
+            # Initial assessment
+            thoughts.append(f"Processing {len(input_data.images)} images")
+            
+            # Context understanding
+            if input_data.context:
+                thoughts.append("\nAnalyzing Context:")
+                thoughts.append(f"- Domain context available: {bool(input_data.context.get('domain_context'))}")
+                thoughts.append(f"- Keywords identified: {len(input_data.context.get('keywords', []))}")
+                
+                # Use LLM to analyze context requirements
+                context_analysis = self.relevance_chain.run(
+                    caption="Context analysis phase",
+                    classification="Initial planning",
+                    context_keywords=str(input_data.context.get('keywords', [])),
+                    domain_context=input_data.context.get('domain_context', 'Not specified')
+                )
+                
+                thoughts.append("\nLLM Context Analysis:")
+                for line in context_analysis.split('\n'):
+                    if line.strip():
+                        thoughts.append(f"- {line.strip()}")
+            
+            # Model selection strategy
+            thoughts.append("\nModel Selection Strategy:")
+            thoughts.append("1. Image Processing Pipeline:")
+            thoughts.append(f"   - Initial caption: {self.models['lightweight_caption'].name}")
+            thoughts.append(f"   - Initial classification: {self.models['lightweight_classifier'].name}")
+            thoughts.append("   - Will escalate to advanced models if confidence below threshold")
+            
+            # Analysis plan
+            thoughts.append("\nAnalysis Pipeline:")
+            thoughts.append("1. Generate captions and classifications")
+            thoughts.append("2. Perform LLM-based relevance analysis")
+            thoughts.append("3. Apply context matching")
+            thoughts.append(f"4. Select top {input_data.top_k} relevant images")
+            
+            if input_data.constraints:
+                thoughts.append("\nConstraint Handling:")
+                for constraint in input_data.constraints:
+                    thoughts.append(f"- Will verify: {constraint}")
+            
+            # Store thoughts in state
+            self.state.thoughts.extend(thoughts)
+            self.logger.info("Reasoning complete for image analysis")
+            self._cleanup_llm()
+            
+            return thoughts
+            
+        except Exception as e:
+            error_msg = f"Error during reasoning: {str(e)}"
+            self.state.errors.append(error_msg)
+            self.logger.error(error_msg)
+            self._cleanup_llm()
+            return thoughts
+
+    
+    def _analyze_image_context_relevance(self, 
+                                    caption: str, 
+                                    classification: str, 
+                                    context: Dict) -> Dict:
+        """
+        Analyze relevance between image content and context using LLM
+        
+        Args:
+            caption: Generated image caption
+            classification: Image classification result
+            context: Context information from ContextLearnAgent
+            
+        Returns:
+            Dict containing:
+                - relevance_score: float (0-1)
+                - content_analysis: str
+                - domain_alignment: str
+                - confidence: float
+        """
+        try:
+            # Prepare context information
+            context_keywords = context.get('keywords', [])
+            domain_context = context.get('domain_context', '')
+            
+            # Get LLM analysis
+            llm_analysis = self.relevance_chain.run(
+                caption=caption,
+                classification=classification,
+                context_keywords=str(context_keywords),
+                domain_context=domain_context
+            )
+            
+            # Parse LLM output
+            analysis_result = {
+                'relevance_score': 0.0,
+                'content_analysis': '',
+                'domain_alignment': '',
+                'confidence': 0.0
+            }
+            
+            # Extract information from LLM output
+            current_section = ''
+            for line in llm_analysis.split('\n'):
+                line = line.strip()
+                if 'Content Relevance' in line:
+                    current_section = 'content'
+                elif 'Domain Alignment' in line:
+                    current_section = 'domain'
+                elif 'Context Matching Score' in line:
+                    try:
+                        # Extract score (0-1) from text
+                        score = float([s for s in line.split() if s.replace('.','').isdigit()][0])
+                        analysis_result['relevance_score'] = min(1.0, max(0.0, score))
+                    except:
+                        pass
+                elif line and current_section:
+                    if current_section == 'content':
+                        analysis_result['content_analysis'] += line + ' '
+                    elif current_section == 'domain':
+                        analysis_result['domain_alignment'] += line + ' '
+            
+            # Calculate confidence based on clarity of analysis
+            analysis_result['confidence'] = min(1.0, 
+                (len(analysis_result['content_analysis']) + 
+                len(analysis_result['domain_alignment'])) / 200)
+            
+            self.logger.debug(f"Context relevance analysis completed with "
+                            f"score: {analysis_result['relevance_score']}")
+            
+            return analysis_result
+            
+        except Exception as e:
+            self.logger.error(f"Error in context relevance analysis: {str(e)}")
+            return {
+                'relevance_score': 0.0,
+                'content_analysis': '',
+                'domain_alignment': '',
+                'confidence': 0.0
+            }
+
+    
+    
+    def reason1(self, input_data: ImageAnalysisInput) -> List[str]:
+        """
+        Plan the image analysis process with both captioning and classification
+        """
+        thoughts = []
+        
+        try:
+            # Input volume assessment
+            thoughts.append(f"Processing {len(input_data.images)} images")
+            
+            # Model strategy explanation
+            thoughts.append("\nModel Selection Strategy:")
+            thoughts.append("1. Captioning Pipeline:")
+            thoughts.append(f"   - Start with {self.models['lightweight_caption'].name} for efficient processing")
+            thoughts.append(f"   - Escalate to {self.models['advanced_caption'].name} if confidence < {self.models['lightweight_caption'].threshold}")
+            
+            thoughts.append("\n2. Classification Pipeline:")
+            thoughts.append(f"   - Begin with {self.models['lightweight_classifier'].name} for initial classification")
+            thoughts.append(f"   - Switch to {self.models['advanced_classifier'].name} if confidence < {self.models['lightweight_classifier'].threshold}")
+            
+            # Context consideration
+            if input_data.context:
+                thoughts.append("\nContext Integration:")
+                thoughts.append("- Will match image captions against provided context")
+                thoughts.append("- Will verify classifications against context requirements")
+            
+            # Constraints handling
+            if input_data.constraints:
+                thoughts.append("\nConstraints Application:")
+                for constraint in input_data.constraints:
+                    thoughts.append(f"- {constraint}")
+            
+            # Process outline
+            thoughts.append("\nAnalysis Pipeline:")
+            thoughts.append("1. Initial lightweight model processing (caption + classification)")
+            thoughts.append("2. Advanced model processing where needed")
+            thoughts.append("3. Context matching and constraint verification")
+            thoughts.append(f"4. Selection of top {input_data.top_k} most relevant images")
+            thoughts.append("5. Memory cleanup after processing")
+            
+            # Store thoughts in state
+            self.state.thoughts.extend(thoughts)
+            self.logger.info("Reasoning complete for image analysis")
+            
+            return thoughts
+            
+        except Exception as e:
+            error_msg = f"Error during reasoning: {str(e)}"
+            self.state.errors.append(error_msg)
+            self.logger.error(error_msg)
+            return thoughts
+
+        
+    # def execute(self, input_data: ImageAnalysisInput) -> Dict:
+    #     """
+    #     Analyze images using multiple models with tiered approach
+        
+    #     Args:
+    #         input_data: Contains images and analysis parameters
+    #     Returns:
+    #         Dict containing:
+    #             - analyzed_images: Dict of image analysis results
+    #             - selected_images: List of top-k relevant images
+    #             - model_decisions: Dict of model choices and confidence
+    #             - status: Processing status
+    #     """
+    #     try:
+    #         if not self.validate(input_data):
+    #             return {
+    #                 'status': 'error',
+    #                 'error': self.state.errors[-1]
+    #             }
+
+    #         results = {
+    #             'analyzed_images': {},
+    #             'selected_images': [],
+    #             'model_decisions': {},
+    #             'status': 'processing'
+    #         }
+
+    #         # Process each image
+    #         for img_path in input_data.images:
+    #             # Start with lightweight models
+    #             caption_result = self._get_image_caption(img_path, 'lightweight')
+                
+    #             # If confidence is low, use advanced model
+    #             if caption_result['confidence'] < self.models['lightweight_caption'].threshold:
+    #                 caption_result = self._get_image_caption(img_path, 'advanced')
+                
+    #             results['analyzed_images'][img_path] = {
+    #                 'caption': caption_result['caption'],
+    #                 'confidence': caption_result['confidence'],
+    #                 'model_used': caption_result['model_used']
+    #             }
+
+    #         # Select top-k relevant images based on context matching
+    #         results['selected_images'] = self._select_relevant_images(
+    #             results['analyzed_images'],
+    #             input_data.context,
+    #             input_data.top_k
+    #         )
+
+    #         results['status'] = 'success'
+    #         self.analyzed_images = results
+            
+    #         # Log decision
+    #         self.log_decision({
+    #             'action': 'analyze_images',
+    #             'num_images': len(input_data.images),
+    #             'selected_images': len(results['selected_images'])
+    #         })
+
+    #         return results
+
+    #     except Exception as e:
+    #         error_msg = f"Error executing image analysis: {str(e)}"
+    #         self.state.errors.append(error_msg)
+    #         self.logger.error(error_msg)
+    #         return {'status': 'error', 'error': error_msg}
+
+    def execute(self, input_data: ImageAnalysisInput) -> Dict:
+        """
+        Execute image analysis using tiered model approach
+        """
+        try:
+            if not self.validate(input_data):
+                return {
+                    'status': 'error',
+                    'error': self.state.errors[-1]
+                }
+
+            results = {
+                'analyzed_images': {},
+                'selected_images': [],
+                'model_decisions': {
+                    'caption_models': set(),
+                    'classifier_models': set(),
+                    'processing_stats': {
+                        'lightweight_usage': 0,
+                        'advanced_usage': 0
+                    }
+                },
+                'status': 'processing'
+            }
+
+            # Process each image
+            for img_path in input_data.images:
+                self.logger.info(f"Processing image: {img_path}")
+                
+                # Initial lightweight processing
+                caption_result = self._get_image_caption(img_path, 'lightweight')
+                classify_result = self._get_classification(img_path, 'lightweight')
+                results['model_decisions']['processing_stats']['lightweight_usage'] += 1
+                
+                # Advanced processing if needed
+                if caption_result['confidence'] < self.models['lightweight_caption'].threshold:
+                    self.logger.info("Escalating to advanced caption model")
+                    caption_result = self._get_image_caption(img_path, 'advanced')
+                    results['model_decisions']['processing_stats']['advanced_usage'] += 1
+                
+                if classify_result['confidence'] < self.models['lightweight_classifier'].threshold:
+                    self.logger.info("Escalating to advanced classification model")
+                    classify_result = self._get_classification(img_path, 'advanced')
+                    results['model_decisions']['processing_stats']['advanced_usage'] += 1
+                
+                # Store results for this image
+                results['analyzed_images'][img_path] = {
+                    'caption': {
+                        'text': caption_result['caption'],
+                        'confidence': caption_result['confidence'],
+                        'model_used': caption_result['model_used']
+                    },
+                    'classification': {
+                        'label': classify_result['class'],
+                        'confidence': classify_result['confidence'],
+                        'model_used': classify_result['model_used']
+                    },
+                    'combined_confidence': (caption_result['confidence'] + classify_result['confidence']) / 2
+                }
+                
+                # Track models used
+                results['model_decisions']['caption_models'].add(caption_result['model_used'])
+                results['model_decisions']['classifier_models'].add(classify_result['model_used'])
+                
+                # Cleanup after each image to manage memory
+                self._cleanup_models()
+
+            # Select top-k relevant images
+            results['selected_images'] = self._select_relevant_images(
+                results['analyzed_images'],
+                input_data.context,
+                input_data.top_k
+            )
+
+            results['status'] = 'success'
+            self.analyzed_images = results
+            
+            # Log final decision
+            self.log_decision({
+                'action': 'complete_image_analysis',
+                'num_images': len(input_data.images),
+                'selected_images': len(results['selected_images']),
+                'model_usage': results['model_decisions']['processing_stats']
+            })
+
+            return results
+
+        except Exception as e:
+            error_msg = f"Error executing image analysis: {str(e)}"
+            self.state.errors.append(error_msg)
+            self.logger.error(error_msg)
+            return {'status': 'error', 'error': error_msg}
+
+
+    # def execute(self, input_data: ImageAnalysisInput) -> Dict:
+    #     """
+    #     Analyze images using both captioning and classification models
+        
+    #     Args:
+    #         input_data: Contains images and analysis parameters
+    #     Returns:
+    #         Dict containing:
+    #             - analyzed_images: Dict of image analysis results
+    #             - selected_images: List of top-k relevant images
+    #             - model_decisions: Dict of model choices and confidence
+    #             - status: Processing status
+    #     """
+    #     try:
+    #         if not self.validate(input_data):
+    #             return {
+    #                 'status': 'error',
+    #                 'error': self.state.errors[-1]
+    #             }
+
+    #         results = {
+    #             'analyzed_images': {},
+    #             'selected_images': [],
+    #             'model_decisions': {},
+    #             'status': 'processing'
+    #         }
+
+    #         # Process each image
+    #         for img_path in input_data.images:
+    #             # Start with lightweight models
+    #             caption_result = self._get_image_caption(img_path, 'lightweight')
+    #             classify_result = self._get_classification(img_path, 'lightweight')
+                
+    #             # If either confidence is low, use advanced models
+    #             if caption_result['confidence'] < self.models['lightweight_caption'].threshold:
+    #                 caption_result = self._get_image_caption(img_path, 'advanced')
+                
+    #             if classify_result['confidence'] < self.models['lightweight_classifier'].threshold:
+    #                 classify_result = self._get_classification(img_path, 'advanced')
+                
+    #             # Store results for this image
+    #             results['analyzed_images'][img_path] = {
+    #                 'caption': caption_result['caption'],
+    #                 'caption_confidence': caption_result['confidence'],
+    #                 'caption_model': caption_result['model_used'],
+    #                 'classification': classify_result['class'],
+    #                 'class_confidence': classify_result['confidence'],
+    #                 'class_model': classify_result['model_used']
+    #             }
+
+    #         # Select top-k relevant images considering both caption and classification
+    #         results['selected_images'] = self._select_relevant_images(
+    #             results['analyzed_images'],
+    #             input_data.context,
+    #             input_data.top_k
+    #         )
+
+    #         # Track model decisions
+    #         results['model_decisions'] = {
+    #             'caption_models_used': set(img['caption_model'] for img in results['analyzed_images'].values()),
+    #             'classifier_models_used': set(img['class_model'] for img in results['analyzed_images'].values())
+    #         }
+
+    #         results['status'] = 'success'
+    #         self.analyzed_images = results
+            
+    #         # Log decision
+    #         self.log_decision({
+    #             'action': 'analyze_images',
+    #             'num_images': len(input_data.images),
+    #             'selected_images': len(results['selected_images']),
+    #             'models_used': results['model_decisions']
+    #         })
+
+    #         # Cleanup to manage memory
+    #         self._cleanup_models()
+
+    #         return results
+
+    #     except Exception as e:
+    #         error_msg = f"Error executing image analysis: {str(e)}"
+    #         self.state.errors.append(error_msg)
+    #         self.logger.error(error_msg)
+    #         return {'status': 'error', 'error': error_msg}
+
+
+
+    # def _get_image_caption(self, img_path: str, model_type: str) -> Dict:
+    #     """
+    #     Generate caption for an image using specified model type
+        
+    #     Args:
+    #         img_path: Path to the image
+    #         model_type: Either 'lightweight' or 'advanced'
+            
+    #     Returns:
+    #         Dict containing:
+    #             - caption: str - Generated caption
+    #             - confidence: float - Confidence score
+    #             - model_used: str - Name of the model used
+    #     """
+    #     try:
+    #         # Select appropriate model
+    #         model_key = f"{model_type}_caption"
+    #         model_config = self.models[model_key]
+            
+    #         # TODO: Implement actual model loading and inference
+    #         # For now, return placeholder result
+    #         return {
+    #             'caption': f"Placeholder caption for {img_path}",
+    #             'confidence': 0.8 if model_type == 'advanced' else 0.6,
+    #             'model_used': model_config.name
+    #         }
+            
+    #     except Exception as e:
+    #         self.logger.error(f"Error in image captioning: {str(e)}")
+    #         return {
+    #             'caption': '',
+    #             'confidence': 0.0,
+    #             'model_used': 'none'
+    #         }
+
+    def _get_image_caption(self, img_path: str, model_type: str) -> Dict:
+        """
+        Generate caption for an image using BLIP models
+        
+        Args:
+            img_path: Path to the image
+            model_type: Either 'lightweight' or 'advanced'
+            
+        Returns:
+            Dict containing:
+                - caption: str - Generated caption
+                - confidence: float - Confidence score
+                - model_used: str - Name of model used
+        """
+        try:
+            model_key = f"{model_type}_caption"
+            model_config = self.models[model_key]
+            
+            # Load image
+            image = Image.open(img_path).convert('RGB')
+            
+            # Initialize model and processor based on type
+            if model_type == 'lightweight':
+                if 'lightweight_processor' not in self.__dict__:
+                    self.lightweight_processor = BlipProcessor.from_pretrained(model_config.name)
+                    self.lightweight_model = BlipForConditionalGeneration.from_pretrained(model_config.name)
+                processor = self.lightweight_processor
+                model = self.lightweight_model
+            else:
+                if 'advanced_processor' not in self.__dict__:
+                    self.advanced_processor = Blip2Processor.from_pretrained(model_config.name)
+                    self.advanced_model = Blip2ForConditionalGeneration.from_pretrained(model_config.name)
+                processor = self.advanced_processor
+                model = self.advanced_model
+            
+            # Process image
+            inputs = processor(image, return_tensors="pt")
+            
+            # Generate caption
+            outputs = model.generate(**inputs, max_new_tokens=50)
+            caption = processor.decode(outputs[0], skip_special_tokens=True)
+            
+            # Calculate confidence (using max probability as proxy)
+            with torch.no_grad():
+                logits = model(**inputs).logits
+                confidence = float(torch.max(torch.softmax(logits[0], dim=-1)).item())
+            
+            return {
+                'caption': caption,
+                'confidence': confidence,
+                'model_used': model_config.name
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error in image captioning: {str(e)}")
+            return {
+                'caption': '',
+                'confidence': 0.0,
+                'model_used': 'none'
+            }
+        
+    # def _select_relevant_images(self, analyzed_images: Dict, context: Dict, top_k: int) -> List[str]:
+    #     """
+    #     Select most relevant images based on context matching and caption analysis
+        
+    #     Args:
+    #         analyzed_images: Dict of image analysis results
+    #         context: Context information from ContextLearnAgent
+    #         top_k: Number of top images to return
+            
+    #     Returns:
+    #         List[str]: Paths of top-k relevant images
+    #     """
+    #     try:
+    #         relevance_scores = {}
+            
+    #         for img_path, analysis in analyzed_images.items():
+    #             # Calculate relevance score based on:
+    #             # 1. Caption confidence
+    #             # 2. Context matching
+    #             # 3. Model reliability (advanced models given higher weight)
+                
+    #             base_score = analysis['confidence']
+                
+    #             # Adjust score based on model type
+    #             model_weight = 1.2 if 'advanced' in analysis['model_used'] else 1.0
+                
+    #             # TODO: Implement context matching using embedding similarity
+    #             context_match_score = 0.5  # Placeholder
+                
+    #             final_score = base_score * model_weight * context_match_score
+    #             relevance_scores[img_path] = final_score
+                
+    #             self.logger.debug(f"Score for {img_path}: {final_score}")
+            
+    #         # Sort by score and get top-k
+    #         selected_images = sorted(
+    #             relevance_scores.items(),
+    #             key=lambda x: x[1],
+    #             reverse=True
+    #         )[:top_k]
+            
+    #         return [img_path for img_path, _ in selected_images]
+            
+    #     except Exception as e:
+    #         self.logger.error(f"Error in selecting relevant images: {str(e)}")
+    #         return []
+
+    
+    
+    def _cleanup_llm(self):
+        """
+        Cleanup LLM and model resources to manage memory
+        """
+        try:
+            # Clear LangChain resources
+            if hasattr(self, 'relevance_chain'):
+                if hasattr(self.relevance_chain, 'clear'):
+                    self.relevance_chain.clear()
+            
+            # Clear any cached data
+            if hasattr(self, 'llm'):
+                if hasattr(self.llm, 'clear_cache'):
+                    self.llm.clear_cache()
+            
+            # Force garbage collection
+            import gc
+            gc.collect()
+            
+            # Clear CUDA cache if available
+            import torch
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                
+            self.logger.info("LLM and model resources cleaned up")
+            
+        except Exception as e:
+            self.logger.error(f"Error in LLM cleanup: {str(e)}")
+
+    
+    def _select_relevant_images(self, analyzed_images: Dict, context: Dict, top_k: int = None) -> List[str]:
+        """
+        Select most relevant images using LLM-enhanced analysis
+        """
+        try:
+            # Use provided top_k or fall back to global config
+            top_k = top_k if top_k is not None else config.top_k
+            
+            relevance_scores = {}
+            
+            for img_path, analysis in analyzed_images.items():
+                # Get base scores from caption and classification
+                base_score = analysis['combined_confidence']
+                
+                # Get LLM-based relevance analysis
+                llm_relevance = self._analyze_image_context_relevance(
+                    caption=analysis['caption']['text'],
+                    classification=analysis['classification']['label'],
+                    context=context
+                )
+                
+                # Calculate final score combining:
+                # 1. Base confidence from models
+                # 2. LLM relevance score
+                # 3. Model weights (advanced models given higher weight)
+                model_weight = 1.2 if 'advanced' in analysis['caption']['model_used'] else 1.0
+                
+                final_score = (
+                    base_score * 0.3 +                    # Original confidence (30%)
+                    llm_relevance['relevance_score'] * 0.5 +  # LLM relevance (50%)
+                    llm_relevance['confidence'] * 0.2     # LLM confidence (20%)
+                ) * model_weight
+                
+                # Store results
+                relevance_scores[img_path] = {
+                    'score': final_score,
+                    'analysis': llm_relevance
+                }
+                
+                self.logger.debug(f"Score for {img_path}: {final_score} "
+                                f"(LLM relevance: {llm_relevance['relevance_score']})")
+            
+            # Sort by score and get top-k
+            selected_images = sorted(
+                relevance_scores.items(),
+                key=lambda x: x[1]['score'],
+                reverse=True
+            )[:top_k]
+            
+            # Log selection decisions
+            self.log_decision({
+                'action': 'image_selection',
+                'selected_count': len(selected_images),
+                'selection_criteria': {
+                    'llm_analysis_used': True,
+                    'top_scores': {img: data['score'] 
+                                for img, data in selected_images}
+                }
+            })
+
+            self._cleanup_llm()
+            
+            return [img_path for img_path, _ in selected_images]
+            
+        except Exception as e:
+            self.logger.error(f"Error in selecting relevant images: {str(e)}")
+            self._cleanup_llm()
+            return []
+
+    
+    def _select_relevant_images1(self, analyzed_images: Dict, context: Dict, top_k: int = None) -> List[str]:
+        """
+        Select most relevant images based on caption, classification, and context
+        
+        Args:
+            analyzed_images: Dict of image analysis results
+            context: Context information from ContextLearnAgent
+            top_k: Number of top images to return
+            
+        Returns:
+            List[str]: Paths of top-k most relevant images
+        """
+        try:
+            top_k = top_k if top_k is not None else config.top_k
+
+            relevance_scores = {}
+            
+            for img_path, analysis in analyzed_images.items():
+                # Initialize base score from combined confidence
+                base_score = analysis['combined_confidence']
+                
+                # Adjust score based on model types used
+                caption_model_weight = 1.2 if 'advanced' in analysis['caption']['model_used'] else 1.0
+                classify_model_weight = 1.2 if 'advanced' in analysis['classification']['model_used'] else 1.0
+                
+                # Context matching for captions
+                caption_context_score = 0.5  # Default score
+                if context and 'keywords' in context:
+                    caption_matches = sum(1 for keyword in context['keywords'] 
+                                    if keyword.lower() in analysis['caption']['text'].lower())
+                    caption_context_score = min(1.0, caption_matches / len(context['keywords']))
+                
+                # Context matching for classification
+                classification_context_score = 0.5  # Default score
+                if context and 'expected_classes' in context:
+                    if analysis['classification']['label'].lower() in [cls.lower() for cls in context['expected_classes']]:
+                        classification_context_score = 1.0
+                
+                # Calculate final score
+                final_score = (
+                    base_score * 
+                    caption_model_weight * 
+                    classify_model_weight * 
+                    (caption_context_score + classification_context_score) / 2
+                )
+                
+                relevance_scores[img_path] = final_score
+                
+                self.logger.debug(f"Score for {img_path}: {final_score} "
+                                f"(caption_score: {caption_context_score}, "
+                                f"class_score: {classification_context_score})")
+            
+            # Sort by score and get top-k
+            selected_images = sorted(
+                relevance_scores.items(),
+                key=lambda x: x[1],
+                reverse=True
+            )[:top_k]
+            
+            # Log selection decisions
+            self.log_decision({
+                'action': 'image_selection',
+                'selected_count': len(selected_images),
+                'top_scores': {img: score for img, score in selected_images}
+            })
+            
+            return [img_path for img_path, _ in selected_images]
+            
+        except Exception as e:
+            self.logger.error(f"Error in selecting relevant images: {str(e)}")
+            return []
+
+
+
+
+
+    def _get_classification(self, img_path: str, model_type: str) -> Dict:
+        """
+        Classify image using ResNet (lightweight) or ViT (advanced)
+        
+        Args:
+            img_path: Path to the image
+            model_type: Either 'lightweight' or 'advanced'
+            
+        Returns:
+            Dict containing:
+                - class: str - Predicted class
+                - confidence: float - Confidence score
+                - model_used: str - Name of model used
+        """
+        try:
+            model_key = f"{model_type}_classifier"
+            model_config = self.models[model_key]
+            
+            # Load image
+            image = Image.open(img_path).convert('RGB')
+            
+            # Initialize model and processor based on type
+            if model_type == 'lightweight':
+                if 'lightweight_clf_processor' not in self.__dict__:
+                    self.lightweight_clf_processor = AutoImageProcessor.from_pretrained(model_config.name)
+                    self.lightweight_clf_model = ResNetForImageClassification.from_pretrained(model_config.name)
+                processor = self.lightweight_clf_processor
+                model = self.lightweight_clf_model
+            else:
+                if 'advanced_clf_processor' not in self.__dict__:
+                    self.advanced_clf_processor = ViTImageProcessor.from_pretrained(model_config.name)
+                    self.advanced_clf_model = ViTForImageClassification.from_pretrained(model_config.name)
+                processor = self.advanced_clf_processor
+                model = self.advanced_clf_model
+            
+            # Process image
+            inputs = processor(image, return_tensors="pt")
+            
+            # Get predictions
+            with torch.no_grad():
+                outputs = model(**inputs)
+                logits = outputs.logits
+                probs = torch.softmax(logits, dim=-1)
+                
+                # Get highest probability class
+                confidence, predicted_idx = torch.max(probs, dim=-1)
+                predicted_label = model.config.id2label[predicted_idx.item()]
+                
+            return {
+                'class': predicted_label,
+                'confidence': float(confidence.item()),
+                'model_used': model_config.name
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error in image classification: {str(e)}")
+            return {
+                'class': '',
+                'confidence': 0.0,
+                'model_used': 'none'
+            }
+
+
+    def _cleanup_models(self) -> None:
+        """
+        Clean up loaded models to free memory
+        Strategy: Remove models but keep processors (smaller memory footprint)
+        """
+        try:
+            # Caption models cleanup
+            if hasattr(self, 'lightweight_model'):
+                del self.lightweight_model
+                torch.cuda.empty_cache() if torch.cuda.is_available() else None
+                self.logger.info("Cleaned up lightweight caption model")
+                
+            if hasattr(self, 'advanced_model'):
+                del self.advanced_model
+                torch.cuda.empty_cache() if torch.cuda.is_available() else None
+                self.logger.info("Cleaned up advanced caption model")
+                
+            # Classification models cleanup
+            if hasattr(self, 'lightweight_clf_model'):
+                del self.lightweight_clf_model
+                torch.cuda.empty_cache() if torch.cuda.is_available() else None
+                self.logger.info("Cleaned up lightweight classification model")
+                
+            if hasattr(self, 'advanced_clf_model'):
+                del self.advanced_clf_model
+                torch.cuda.empty_cache() if torch.cuda.is_available() else None
+                self.logger.info("Cleaned up advanced classification model")
+                
+            # Keep processors as they're smaller and faster to reload
+            self.logger.info("Model cleanup completed")
+            
+        except Exception as e:
+            self.logger.error(f"Error during model cleanup: {str(e)}")

agents/user_input_agent.py

@@ -0,0 +1,321 @@
+from dataclasses import dataclass
+from typing import Dict, List, Any
+
+from langchain.llms import HuggingFacePipeline
+from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
+
+import datetime
+
+from .base_agent import BaseAgent
+from models.data_models import UserInput 
+
+
+class UserInputAgent(BaseAgent):
+    def __init__(self, name: str = "UserInputAgent"):
+        super().__init__(name)
+        self.model_name = "MBZUAI/LaMini-Flan-T5-783M"
+        self._initialize_llm()
+        
+
+    def _initialize_llm(self):
+        """Initialize the language model pipeline"""
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
+            
+            # Create pipeline
+            pipe = pipeline(
+                "text2text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                max_length=512,
+                temperature=0.3
+            )
+            
+            # Create LangChain HF pipeline
+            self.llm = HuggingFacePipeline(pipeline=pipe)
+            
+            # Create intent analysis chain
+            intent_template = """
+            Analyze this image analysis task request:
+            Query: {query}
+            Constraints: {constraints}
+            
+            Provide a structured analysis including:
+            1. Main purpose of the request
+            2. Key requirements
+            3. Important constraints to consider
+            
+            Analysis:
+            """
+            
+            self.intent_chain = LLMChain(
+                llm=self.llm,
+                prompt=PromptTemplate(
+                    template=intent_template,
+                    input_variables=["query", "constraints"]
+                )
+            )
+            
+        except Exception as e:
+            self.logger.error(f"Error initializing LLM: {str(e)}")
+            raise
+    
+    def reason(self, input_data: UserInput) -> List[str]:
+        """
+        Analyze user input using LangChain for better understanding
+        """
+        thoughts = []
+        
+        try:
+            # Initial analysis
+            thoughts.append(f"Processing user query: '{input_data.query}'")
+            if input_data.constraints:
+                thoughts.append(f"With constraints: '{input_data.constraints}'")
+            
+            # Use LangChain for intent analysis
+            llm_analysis = self.intent_chain.run(
+                query=input_data.query,
+                constraints=input_data.constraints if input_data.constraints else "None specified"
+            )
+            
+            # Parse and add LLM insights
+            thoughts.append("\nLLM Analysis:")
+            for line in llm_analysis.split('\n'):
+                if line.strip():
+                    thoughts.append(f"- {line.strip()}")
+            
+            # Add reasoning about next steps
+            thoughts.append("\nPlanned Actions:")
+            thoughts.append("1. Extract key terms for context search")
+            thoughts.append("2. Prepare constraints for filtering")
+            thoughts.append("3. Format requirements for image analysis")
+            
+            # Store thoughts in state
+            self.state.thoughts.extend(thoughts)
+            self.logger.info("Reasoning complete for user input")
+            
+            return thoughts
+            
+        except Exception as e:
+            error_msg = f"Error during reasoning: {str(e)}"
+            self.state.errors.append(error_msg)
+            self.logger.error(error_msg)
+            return thoughts
+
+    # def reason1(self, input_data: UserInput) -> List[str]:
+    #     """
+    #     Analyze user input and formulate reasoning steps about their intent
+        
+    #     Args:
+    #         input_data: UserInput containing query, constraints, and top_k
+            
+    #     Returns:
+    #         List[str]: Reasoning steps about user's intent
+    #     """
+    #     thoughts = []
+        
+    #     try:
+    #         # Analyze main query
+    #         thoughts.append(f"Analyzing user query: '{input_data.query}'")
+            
+    #         # Analyze constraints if provided
+    #         if input_data.constraints:
+    #             thoughts.append(f"Considering constraints: '{input_data.constraints}'")
+    #         else:
+    #             thoughts.append("No specific constraints provided")
+                
+    #         # Consider result limit
+    #         thoughts.append(f"User requests top {input_data.top_k} results")
+            
+    #         # Store thoughts in state
+    #         self.state.thoughts.extend(thoughts)
+    #         self.logger.info("Reasoning complete for user input")
+            
+    #         return thoughts
+            
+    #     except Exception as e:
+    #         error_msg = f"Error during reasoning: {str(e)}"
+    #         self.state.errors.append(error_msg)
+    #         self.logger.error(error_msg)
+    #         return thoughts
+
+    # def execute(self, input_data: UserInput) -> Dict:
+    #     """
+    #     Process user input and prepare structured data for other agents
+        
+    #     Args:
+    #         input_data: UserInput containing query, constraints, and top_k
+            
+    #     Returns:
+    #         Dict containing:
+    #             - processed_query: str - Cleaned and formatted query
+    #             - constraints: List[str] - List of parsed constraints
+    #             - parameters: Dict - Additional parameters including top_k
+    #             - status: str - Processing status
+    #     """
+    #     try:
+    #         # First validate the input
+    #         if not self.validate(input_data):
+    #             return {
+    #                 'status': 'error',
+    #                 'error': self.state.errors[-1]
+    #             }
+                
+    #         # Process the input
+    #         processed_data = {
+    #             'processed_query': input_data.query.strip().lower(),
+    #             'constraints': [c.strip() for c in input_data.constraints.split(';')] if input_data.constraints else [],
+    #             'parameters': {
+    #                 'top_k': input_data.top_k,
+    #                 'timestamp': datetime.now().isoformat()
+    #             },
+    #             'status': 'success'
+    #         }
+            
+    #         # Log the decision
+    #         self.log_decision({
+    #             'action': 'process_user_input',
+    #             'input': str(input_data),
+    #             'output': processed_data
+    #         })
+            
+    #         return processed_data
+            
+    #     except Exception as e:
+    #         error_msg = f"Error executing user input processing: {str(e)}"
+    #         self.state.errors.append(error_msg)
+    #         self.logger.error(error_msg)
+    #         return {'status': 'error', 'error': error_msg}
+
+    def execute(self, input_data: UserInput) -> Dict:
+        """
+        Process user input and prepare it for other agents
+        """
+        try:
+            if not self.validate(input_data):
+                return {
+                    'status': 'error',
+                    'error': self.state.errors[-1]
+                }
+
+            # Process input with LLM analysis
+            llm_analysis = self._analyze_intent(
+                query=input_data.query,
+                constraints=input_data.constraints
+            )
+
+            results = {
+                'processed_input': {
+                    'original_query': input_data.query,
+                    'constraints': input_data.constraints.split(';') if input_data.constraints else [],
+                    'llm_analysis': {
+                        'main_purpose': llm_analysis.get('purpose', ''),
+                        'key_requirements': llm_analysis.get('requirements', []),
+                        'constraint_interpretation': llm_analysis.get('constraints', [])
+                    }
+                },
+                'metadata': {
+                    'timestamp': datetime.now().isoformat(),
+                    'model_used': self.model_name
+                },
+                'status': 'success'
+            }
+
+            # Log decision
+            self.log_decision({
+                'action': 'process_user_input',
+                'input': str(input_data),
+                'llm_insights': llm_analysis
+            })
+
+            # Cleanup LLM resources
+            self._cleanup_llm()
+
+            return results
+
+        except Exception as e:
+            error_msg = f"Error executing user input processing: {str(e)}"
+            self.state.errors.append(error_msg)
+            self.logger.error(error_msg)
+            return {'status': 'error', 'error': error_msg}
+        
+
+    def _analyze_intent(self, query: str, constraints: str) -> Dict:
+        """
+        Process user input through LangChain and structure the results
+        
+        Args:
+            query: User's query string
+            constraints: User's constraints string
+            
+        Returns:
+            Dict containing:
+                - purpose: str - Main purpose identified
+                - requirements: List[str] - Key requirements
+                - constraints: List[str] - Interpreted constraints
+        """
+        try:
+            # Get raw LLM analysis
+            raw_analysis = self.intent_chain.run(
+                query=query,
+                constraints=constraints if constraints else "None specified"
+            )
+            
+            # Structure the analysis
+            analysis_dict = {
+                'purpose': '',
+                'requirements': [],
+                'constraints': []
+            }
+            
+            # Parse the raw output
+            current_section = ''
+            for line in raw_analysis.split('\n'):
+                line = line.strip()
+                if 'purpose' in line.lower():
+                    current_section = 'purpose'
+                elif 'requirement' in line.lower():
+                    current_section = 'requirements'
+                elif 'constraint' in line.lower():
+                    current_section = 'constraints'
+                elif line:
+                    if current_section == 'purpose':
+                        analysis_dict['purpose'] = line
+                    elif current_section in ['requirements', 'constraints']:
+                        analysis_dict[current_section].append(line)
+            
+            return analysis_dict
+            
+        except Exception as e:
+            self.logger.error(f"Error in intent analysis: {str(e)}")
+            return {'purpose': '', 'requirements': [], 'constraints': []}
+
+    def _cleanup_llm(self):
+        """
+        Cleanup LLM resources to manage memory
+        """
+        try:
+            # Clear any cached data
+            if hasattr(self, 'intent_chain'):
+                # Clear any stored predictions
+                if hasattr(self.intent_chain, 'clear'):
+                    self.intent_chain.clear()
+            
+            # Force garbage collection
+            import gc
+            gc.collect()
+            
+            # Clear CUDA cache if available
+            import torch
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                
+            self.logger.info("LLM resources cleaned up")
+            
+        except Exception as e:
+            self.logger.error(f"Error in LLM cleanup: {str(e)}")
+
+

app.py

@@ -0,0 +1,54 @@
+# app.py
+import gradio as gr
+import logging
+from interface.app import create_interface
+from utils.resource_manager import ResourceManager
+from config.settings import config
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+def initialize_app():
+    """Initialize application components"""
+    try:
+        # Initialize resource manager
+        resource_manager = ResourceManager()
+        
+        # Check initial resources
+        resources_ok, error_msg = resource_manager.check_resources()
+        if not resources_ok:
+            logger.error(f"Resource check failed: {error_msg}")
+            raise Exception(error_msg)
+        
+        # Create Gradio interface
+        interface = create_interface()
+        
+        return interface
+        
+    except Exception as e:
+        logger.error(f"Initialization error: {str(e)}")
+        raise
+
+if __name__ == "__main__":
+    try:
+        # Initialize app
+        interface = initialize_app()
+        
+        # Launch Gradio interface
+        interface.launch(
+            share=True,
+            server_name="0.0.0.0",
+            server_port=7860,
+            enable_queue=True,
+            max_threads=4,  # Limit concurrent processing
+            auth=None,  # Add authentication if needed
+            ssl_verify=False,  # For HuggingFace spaces
+        )
+        
+    except Exception as e:
+        logger.error(f"Application startup failed: {str(e)}")
+        raise

config/settings.py

@@ -0,0 +1,11 @@
+# config/settings.py
+from dataclasses import dataclass
+
+@dataclass
+class GlobalConfig:
+    top_k: int = 5  # default value
+    max_wiki_results: int = 5  # our existing ContextLearnAgent parameter
+    report_format: str = "detailed"  # for AssemblerAgent
+
+# Can be accessed throughout the application
+config = GlobalConfig()

interface/__init__.py

@@ -0,0 +1,49 @@
+# interface/app.py
+import gradio as gr
+from typing import List, Dict
+from .handlers import process_inputs
+from .display import format_results
+from .utils import validate_inputs
+from config.settings import config
+
+def create_interface():
+    """Create the Gradio interface"""
+    
+    # Create interface components
+    with gr.Blocks(title="Pump Inspection Analysis") as app:
+        gr.Markdown("# Pump Inspection Analysis System")
+        
+        with gr.Row():
+            # Input components
+            with gr.Column():
+                query = gr.Textbox(
+                    label="What would you like to analyze?",
+                    placeholder="e.g., Check safety issues in pump systems",
+                    lines=3
+                )
+                
+                constraints = gr.Textbox(
+                    label="Any specific constraints? (optional)",
+                    placeholder="e.g., Exclude routine maintenance issues",
+                    lines=2
+                )
+                
+                top_k = gr.Slider(
+                    minimum=1,
+                    maximum=10,
+                    value=5,
+                    step=1,
+                    label="Number of top results to show"
+                )
+                
+                report_format = gr.Radio(
+                    choices=["summary", "detailed"],
+                    value="summary",
+                    label="Report Format"
+                )
+                
+                images = gr.File(
+                    file_count="multiple",
+                    label="Upload Images",
+                    file_types=["image"]
+                )

interface/app.py

@@ -0,0 +1,153 @@
+# interface/app.py
+import gradio as gr
+from typing import List, Dict
+from .handlers import process_inputs
+from .display import format_results
+from .utils import validate_inputs
+from config.settings import config
+from .handlers import ProcessingHandler
+from .utils import InputValidator
+
+def create_interface():
+    """Create the Gradio interface"""
+    handler = ProcessingHandler()  
+    
+    # Create interface components
+    with gr.Blocks(title="Image Inspection analysis") as app:
+        gr.Markdown("# Image Inspection analysis System")
+        
+        with gr.Row():
+            # Input components
+            with gr.Column():
+                query = gr.Textbox(
+                    label="What would you like to analyze?",
+                    placeholder="e.g., Check safety issues in pump systems",
+                    lines=3
+                )
+                
+                constraints = gr.Textbox(
+                    label="Any specific constraints? (optional)",
+                    placeholder="e.g., Exclude routine maintenance issues",
+                    lines=2
+                )
+                
+                top_k = gr.Slider(
+                    minimum=1,
+                    maximum=10,
+                    value=5,
+                    step=1,
+                    label="Number of top results to show"
+                )
+                
+                report_format = gr.Radio(
+                    choices=["summary", "detailed"],
+                    value="summary",
+                    label="Report Format"
+                )
+                
+                images = gr.File(
+                    file_count="multiple",
+                    label="Upload Images",
+                    file_types=["image"]
+                )
+
+                submit_btn = gr.Button("Analyze", variant="primary")
+
+            # Right column - Outputs
+            with gr.Column():
+                with gr.Tab("Results"):
+                    analysis_status = gr.Markdown("Ready for analysis...")
+                    
+                    results_box = gr.Markdown(
+                        visible=False,
+                        label="Analysis Results"
+                    )
+                    
+                    selected_images = gr.Gallery(
+                        label="Selected Relevant Images",
+                        visible=False,
+                        columns=2,
+                        height=400
+                    )
+                    
+                    confidence_scores = gr.Json(
+                        label="Confidence Scores",
+                        visible=False
+                    )
+                
+                with gr.Tab("Processing Details"):
+                    processing_status = gr.JSON(
+                        label="Processing Steps",
+                        visible=False
+                    )
+                    
+                with gr.Tab("Errors"):
+                    error_box = gr.Markdown(
+                        visible=False
+                    )
+
+        # Second: Helper functions for UI updates
+        def update_ui_on_error(error_msg):
+            return {
+                results_box: gr.update(visible=True, value=error_msg),
+                selected_images: gr.update(visible=False),
+                confidence_scores: gr.update(visible=False),
+                processing_status: gr.update(visible=True, value={'status': 'error'})
+            }
+        
+        def update_ui_on_success(results):
+            return {
+                results_box: gr.update(visible=True, value=results['content']),
+                selected_images: gr.update(visible=True, value=results['images']),
+                confidence_scores: gr.update(visible=True, value=results['scores']),
+                processing_status: gr.update(visible=True, value={'status': 'success'})
+            }
+        
+        def validate_and_process(query, constraints, images, top_k, report_format):
+            # Validate inputs
+            is_valid, error_message = InputValidator.validate_inputs(
+                query, constraints, images
+            )
+            
+            if not is_valid:
+                return update_ui_on_error(error_message)
+            
+            # If valid, proceed with processing
+            return handler.process_inputs(
+                query, constraints, images, top_k, report_format
+            )
+        
+        # Single submit button with combined functionality
+        submit_btn = gr.Button("Analyze", variant="primary")
+        
+        # Connect submit button to both clear and process
+        submit_btn.click(
+            fn=lambda: [
+                gr.Markdown.update(visible=True, value="Validating inputs......"),
+                gr.Markdown.update(visible=False),
+                gr.Gallery.update(visible=False),
+                gr.Json.update(visible=False),
+                gr.JSON.update(visible=False)
+            ],
+            inputs=None,
+            outputs=[analysis_status, results_box, selected_images, 
+                    confidence_scores, processing_status]
+        ).then(  # Chain the processing after clearing
+            fn=validate_and_process,
+            inputs=[query, constraints, images, top_k, report_format],
+            outputs=[results_box, selected_images, confidence_scores, processing_status],
+            #show_progress=True
+        ).then(
+            # Third: Update UI based on results
+            fn=lambda results: update_ui_on_success(results) if results['status'] == 'success' 
+                            else update_ui_on_error(results['error']),
+            inputs=[processing_status],
+            outputs=[results_box, selected_images, confidence_scores, processing_status]
+        )
+
+    return app
+
+# Launch the interface
+if __name__ == "__main__":
+    interface = create_interface()
+    interface.launch()

interface/display.py

@@ -0,0 +1,50 @@
+# interface/display.py
+from typing import Dict, List
+import json
+
+class DisplayFormatter:
+    @staticmethod
+    def format_error(error_msg: str) -> Dict:
+        """Format error messages for display"""
+        return {
+            'status': 'error',
+            'content': f"""
+            ### ❌ Error
+            {error_msg}
+            
+            Please try again or contact support if the issue persists.
+            """
+        }
+    
+    @staticmethod
+    def format_results(report: Dict, images: List[str], scores: Dict) -> Dict:
+        """Format successful results for display"""
+        try:
+            # Format main results
+            markdown_content = f"""
+            ### 📊 Analysis Results
+            
+            #### Query Analysis
+            {report.get('query_analysis', 'No analysis available')}
+            
+            #### Key Findings
+            {report.get('key_findings', 'No findings available')}
+            
+            #### Recommendations
+            {report.get('recommendations', 'No recommendations available')}
+            """
+            
+            # Format confidence scores
+            confidence_display = {
+                'Model Performance': scores.get('model_decisions', {}),
+                'Overall Confidence': f"{scores.get('average_confidence', 0) * 100:.1f}%"
+            }
+            
+            return {
+                'status': 'success',
+                'content': markdown_content,
+                'scores': confidence_display,
+                'images': images
+            }
+        except Exception as e:
+            return DisplayFormatter.format_error(f"Error formatting results: {str(e)}")

interface/handlers.py

@@ -0,0 +1,129 @@
+# interface/handlers.py
+import gradio as gr
+from typing import Dict, List, Tuple
+from agents.user_input_agent import UserInputAgent
+from agents.context_agent import ContextLearnAgent
+from agents.image_agent import ImageAnalyzerAgent
+from agents.assembler_agent import AssemblerAgent
+from models.data_models import UserInput, ContextInput, ImageAnalysisInput
+from config.settings import config
+from .display import DisplayFormatter
+from utils.resource_manager import ResourceManager
+
+class ProcessingHandler:
+    def __init__(self):
+        self.formatter = DisplayFormatter()
+        self.resource_manager = ResourceManager()
+
+        """Initialize agents"""
+        self.user_agent = UserInputAgent()
+        self.context_agent = ContextLearnAgent()
+        self.image_agent = ImageAnalyzerAgent()
+        self.assembler_agent = AssemblerAgent()
+        
+    async def process_inputs(
+        self,
+        query: str,
+        constraints: str,
+        images: List[str],
+        top_k: int,
+        report_format: str,
+        progress=gr.Progress()
+    ) -> Tuple[str, List[str], Dict, Dict]:
+        """
+        Process inputs through agent pipeline
+        
+        Returns:
+            Tuple containing:
+            - results_markdown: str
+            - selected_images: List[str]
+            - confidence_scores: Dict
+            - processing_details: Dict
+        """
+        try:
+            resources_ok, error_msg = self.resource_manager.check_resources()
+            if not resources_ok:
+                raise Exception(error_msg)
+            processing_details = {'status': 'processing'}
+
+            # Step 1: Process user input
+            progress(0.1, desc="Processing user input...")
+            user_results = self.user_agent.execute(
+                UserInput(query=query, constraints=constraints)
+            )
+
+            if user_results['status'] == 'error':
+                raise Exception(user_results['error'])
+            
+            # Step 2: Gather context
+            progress(0.3, desc="Learning context...")
+            context_results = self.context_agent.execute(
+                ContextInput(
+                    processed_query=user_results['processed_input']['original_query'],
+                    constraints=user_results['processed_input']['constraints'],
+                    domain="oil_and_gas"
+                )
+            )
+            # Step 3: Analyze images
+            progress(0.5, desc="Analyzing images...")
+            image_results = self.image_agent.execute(
+                ImageAnalysisInput(
+                    images=images,
+                    context=context_results['gathered_context'],
+                    constraints=user_results['processed_input']['constraints'],
+                    top_k=top_k
+                )
+            )
+            # Step 4: Assemble final report
+            progress(0.8, desc="Assembling report...")
+            final_report = self.assembler_agent.execute({
+                'user_input_results': user_results,
+                'context_results': context_results,
+                'image_results': image_results,
+                'report_format': report_format
+            })
+            
+            progress(1.0, desc="Complete!")
+            
+            # Format results for Gradio
+            # return (
+            #     final_report['summary'],           # results_markdown
+            #     image_results['selected_images'],  # selected_images
+            #     image_results['model_decisions'],  # confidence_scores
+            #     processing_details                 # processing_details
+            # )
+
+            # Monitor resources during processing
+            self.resource_manager.monitor_and_cleanup()
+
+            # Format results
+            formatted_results = self.formatter.format_results(
+                final_report['summary'],
+                image_results['selected_images'],
+                image_results['model_decisions']
+            )
+
+            # Final cleanup
+            self.resource_manager.cleanup()
+            
+            return (
+                formatted_results['content'],
+                formatted_results['images'],
+                formatted_results['scores'],
+                {'status': 'success', 'details': processing_details}
+            )
+        
+        except Exception as e:
+            self.resource_manager.cleanup()  # Ensure cleanup on error
+            error_format = self.formatter.format_error(str(e))
+            return (
+                error_format['content'],
+                [],
+                {},
+                {'status': 'error', 'error': str(e)}
+            )
+            # error_msg = f"Error during processing: {str(e)}"
+            # processing_details['status'] = 'error'
+            # processing_details['error'] = error_msg
+            # return "", [], {}, processing_details
+        

interface/utils.py

@@ -0,0 +1,89 @@
+# interface/utils.py
+from typing import Tuple, List
+import os
+from PIL import Image
+
+class InputValidator:
+    # Constants for validation
+    MAX_QUERY_LENGTH = 500
+    MIN_IMAGES = 1
+    MAX_IMAGES = 10
+    ALLOWED_IMAGE_TYPES = ['.jpg', '.jpeg', '.png']
+    MAX_IMAGE_SIZE_MB = 5
+    MAX_IMAGE_RESOLUTION = (2048, 2048)
+    ALLOWED_FORMATS = ['summary', 'detailed']
+    
+    @staticmethod
+    def validate_inputs(
+        query: str,
+        constraints: str,
+        images: List[str],
+        top_k: int,
+        report_format: str
+    ) -> Tuple[bool, str]:
+        """
+        Validate all user inputs
+        
+        Args:
+            query: User's query string
+            constraints: User's constraints string
+            images: List of image paths
+            top_k: Number of top results to return
+            report_format: Type of report to generate
+            
+        Returns:
+            Tuple(is_valid: bool, error_message: str)
+        """
+        try:
+            # Query validation
+            if not query or query.isspace():
+                return False, "Query is required"
+            if len(query) > InputValidator.MAX_QUERY_LENGTH:
+                return False, f"Query too long (max {InputValidator.MAX_QUERY_LENGTH} characters)"
+            
+            # Images validation
+            if not images:
+                return False, "At least one image is required"
+            if len(images) > InputValidator.MAX_IMAGES:
+                return False, f"Too many images. Maximum allowed: {InputValidator.MAX_IMAGES}"
+            
+            # Process each image
+            for img_path in images:
+                # File type check
+                file_ext = os.path.splitext(img_path)[1].lower()
+                if file_ext not in InputValidator.ALLOWED_IMAGE_TYPES:
+                    return False, f"Invalid image type: {file_ext}. Allowed types: {', '.join(InputValidator.ALLOWED_IMAGE_TYPES)}"
+                
+                # File size check
+                file_size_mb = os.path.getsize(img_path) / (1024 * 1024)
+                if file_size_mb > InputValidator.MAX_IMAGE_SIZE_MB:
+                    return False, f"Image too large: {file_size_mb:.1f}MB. Maximum size: {InputValidator.MAX_IMAGE_SIZE_MB}MB"
+                
+                # Image integrity and resolution check
+                try:
+                    with Image.open(img_path) as img:
+                        img.verify()  # Verify image integrity
+                        width, height = img.size
+                        if width > InputValidator.MAX_IMAGE_RESOLUTION[0] or height > InputValidator.MAX_IMAGE_RESOLUTION[1]:
+                            return False, f"Image resolution too high. Maximum: {InputValidator.MAX_IMAGE_RESOLUTION[0]}x{InputValidator.MAX_IMAGE_RESOLUTION[1]}"
+                except Exception as e:
+                    return False, f"Invalid or corrupted image: {os.path.basename(img_path)}"
+            
+            # Top-k validation
+            if not isinstance(top_k, int) or top_k < 1:
+                return False, "Top-k must be a positive integer"
+            if top_k > len(images):
+                return False, f"Top-k ({top_k}) cannot be larger than number of images ({len(images)})"
+            
+            # Report format validation
+            if report_format not in InputValidator.ALLOWED_FORMATS:
+                return False, f"Invalid report format. Allowed formats: {', '.join(InputValidator.ALLOWED_FORMATS)}"
+            
+            # Optional constraints validation
+            if constraints and len(constraints) > InputValidator.MAX_QUERY_LENGTH:
+                return False, f"Constraints too long (max {InputValidator.MAX_QUERY_LENGTH} characters)"
+            
+            return True, ""
+            
+        except Exception as e:
+            return False, f"Validation error: {str(e)}"

models/data_models.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass, field
+from typing import List, Dict
+
+# @dataclass
+# class ContextInput:
+#     processed_query: str
+#     constraints: List[str]
+#     search_sources: List[str] = field(default_factory=lambda: ['wikipedia'])
+#     max_results: int = 5
+@dataclass
+class ContextInput:
+    processed_query: str
+    constraints: List[str]
+    domain: str  # e.g., "oil_and_gas", "fine_art"
+    max_results: int = 5
+    min_confidence: float = 0.7
+
+@dataclass
+class UserInput:
+    query: str
+    constraints: str
+    top_k: int
+
+@dataclass
+class ImageAnalysisInput:
+    images: List[str]  # List of image paths
+    context: Dict      # Context from ContextLearnAgent
+    constraints: List[str]
+    top_k: int = 5
+
+@dataclass
+class AgentState:
+    """Tracks the current state of the agent"""
+    intent: str = ""
+    thoughts: List[str] = field(default_factory=list)
+    decisions: List[Dict] = field(default_factory=list)
+    errors: List[str] = field(default_factory=list)
+
+@dataclass
+class AssemblerInput:
+    user_input_results: Dict  # From UserInputAgent
+    context_results: Dict     # From ContextLearnAgent
+    image_results: Dict      # From ImageAnalyzerAgent
+    report_format: str = "detailed"  # or "summary"

models/model_config.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+from typing import Literal
+
+@dataclass
+class ModelConfig:
+    name: str
+    type: Literal['lightweight', 'advanced']
+    task: Literal['caption', 'classification', 'detection', 'summarization']
+    threshold: float = 0.5

requirements.txt

@@ -0,0 +1,31 @@
+# Core dependencies
+gradio>=3.50.2
+torch>=2.0.0
+transformers>=4.30.0
+langchain>=0.0.300
+
+# Image processing
+Pillow>=10.0.0
+numpy>=1.24.0
+
+# Resource management
+psutil>=5.9.0
+
+# Web and API
+requests>=2.31.0
+httpx>=0.24.1
+
+# Utilities
+python-dotenv>=1.0.0
+tqdm>=4.65.0
+pandas>=2.0.0
+
+# Text processing
+nltk>=3.8.1
+beautifulsoup4>=4.12.0
+
+# Wikipedia access
+wikipedia-api>=0.6.0
+
+# Logging and monitoring
+logging>=0.5.1.2

utils/resource_manager.py

@@ -0,0 +1,54 @@
+# utils/resource_manager.py
+import torch
+import gc
+import psutil
+import os
+from typing import Dict, Tuple
+import logging
+
+class ResourceManager:
+    # Constants
+    MAX_MEMORY_USAGE_PCT = 90  # Maximum memory usage percentage
+    CLEANUP_THRESHOLD_PCT = 80  # Threshold to trigger cleanup
+    
+    def __init__(self):
+        self.logger = logging.getLogger(__name__)
+        self._active_models: Dict = {}
+    
+    def check_resources(self) -> Tuple[bool, str]:
+        """Check if sufficient resources are available"""
+        try:
+            memory = psutil.Process(os.getpid()).memory_percent()
+            if memory > self.MAX_MEMORY_USAGE_PCT:
+                return False, f"Memory usage too high: {memory:.1f}%"
+            return True, ""
+        except Exception as e:
+            return False, f"Resource check error: {str(e)}"
+    
+    def cleanup(self):
+        """Force cleanup of resources"""
+        try:
+            # Clear models
+            self._active_models.clear()
+            
+            # Clear CUDA cache
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            
+            # Force garbage collection
+            gc.collect()
+            
+            self.logger.info("Resource cleanup completed")
+            
+        except Exception as e:
+            self.logger.error(f"Error during cleanup: {str(e)}")
+    
+    def monitor_and_cleanup(self):
+        """Monitor resources and cleanup if needed"""
+        try:
+            memory = psutil.Process(os.getpid()).memory_percent()
+            if memory > self.CLEANUP_THRESHOLD_PCT:
+                self.logger.warning(f"High memory usage ({memory:.1f}%), triggering cleanup")
+                self.cleanup()
+        except Exception as e:
+            self.logger.error(f"Monitoring error: {str(e)}")