syedkhizarrayaz/BM-AI-Analysis-And-Alert-Prioritization-Agent

1import requests2import json3import asyncio4import aiohttp5import time6import threading7from typing import Optional, Dict, Any8from functools import lru_cache9import concurrent.futures10from requests.adapters import HTTPAdapter11from urllib3.util.retry import Retry12 13 14class LocalLLM:15    """16    A class to interact with Ollama local LLM using qwen2.5:1.5b model.17    Provides methods to generate responses from text prompts.18    """19    20    def __init__(self, model: str = "qwen2.5:1.5b", base_url: str = "http://localhost:11434"):21        """22        Initialize the LocalLLM class with optimized settings.23        24        Args:25            model (str): The Ollama model to use. Defaults to "qwen2.5:1.5b"26            base_url (str): The base URL for Ollama API. Defaults to "http://localhost:11434"27        """28        self.model = model29        self.base_url = base_url30        self.api_url = f"{base_url}/api/generate"31        32        # Performance optimizations33        self.session = self._create_optimized_session()34        self.response_cache = {}35        self.cache_lock = threading.Lock()36        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=10)37        38        # Preload model for faster responses39        self._preload_model()40    41    def _create_optimized_session(self):42        """Create an optimized HTTP session with connection pooling"""43        session = requests.Session()44        45        # Configure retry strategy46        retry_strategy = Retry(47            total=3,48            backoff_factor=0.1,49            status_forcelist=[429, 500, 502, 503, 504],50        )51        52        # Configure adapter with connection pooling53        adapter = HTTPAdapter(54            max_retries=retry_strategy,55            pool_connections=20,56            pool_maxsize=20,57            pool_block=False58        )59        60        session.mount("http://", adapter)61        session.mount("https://", adapter)62        63        return session64    65    def _preload_model(self):66        """Preload the model to keep it in memory for faster responses"""67        try:68            payload = {69                "model": self.model,70                "prompt": "Hello",71                "stream": False,72                "options": {73                    "num_ctx": 2048,74                    "num_predict": 50,75                    "temperature": 0.1,76                    "top_p": 0.9,77                    "repeat_penalty": 1.1,78                    "num_thread": 10  # Use 90% of available cores79                }80            }81            82            # Send a warm-up request83            response = self.session.post(84                self.api_url,85                json=payload,86                timeout=3087            )88            if response.status_code == 200:89                print(f"✅ Model {self.model} preloaded successfully")90            else:91                print(f"⚠️ Model preload failed: {response.status_code}")92                93        except Exception as e:94            print(f"⚠️ Model preload error: {e}")95    96    def _get_cache_key(self, prompt: str, **kwargs) -> str:97        """Generate cache key for prompt"""98        return f"{hash(prompt)}_{hash(str(sorted(kwargs.items())))}"99    100    def generate_response(self, prompt: str, stream: bool = False, **kwargs) -> str:101        """102        Generate a response from the LLM for the given prompt with caching and optimization.103        104        Args:105            prompt (str): The input prompt/question106            stream (bool): Whether to stream the response. Defaults to False107            **kwargs: Additional parameters for the API call (temperature, top_p, etc.)108            109        Returns:110            str: The generated response text111        """112        # Check cache first113        cache_key = self._get_cache_key(prompt, **kwargs)114        with self.cache_lock:115            if cache_key in self.response_cache:116                return self.response_cache[cache_key]117        118        # Optimized payload with performance settings119        payload = {120            "model": self.model,121            "prompt": prompt,122            "stream": stream,123            "options": {124                "num_ctx": 2048,125                "num_predict": 200,126                "temperature": 0.1,127                "top_p": 0.9,128                "repeat_penalty": 1.1,129                "num_thread": 10,  # Use 90% of available cores130                "num_gpu": 0,  # Disable GPU for CPU optimization131                "num_batch": 512,132                "num_keep": 0,133                "seed": -1,134                "tfs_z": 1.0,135                "typical_p": 1.0,136                "repeat_last_n": 64,137                "penalize_newline": True,138                "stop": ["</s>", "[/INST]", "\n\n\n"],139                **kwargs140            }141        }142        143        try:144            start_time = time.time()145            response = self.session.post(146                self.api_url,147                json=payload,148                headers={"Content-Type": "application/json"},149                timeout=30  # Reduced timeout for faster failure detection150            )151            response.raise_for_status()152            153            if stream:154                result = self._handle_streaming_response(response)155            else:156                result = response.json().get("response", "")157            158            # Cache the result159            with self.cache_lock:160                self.response_cache[cache_key] = result161                # Limit cache size to prevent memory issues162                if len(self.response_cache) > 1000:163                    # Remove oldest entries164                    oldest_key = next(iter(self.response_cache))165                    del self.response_cache[oldest_key]166            167            end_time = time.time()168            print(f"⚡ Response generated in {end_time - start_time:.3f} seconds")169            170            return result171                172        except requests.exceptions.RequestException as e:173            raise Exception(f"Failed to generate response: {str(e)}")174    175    def _handle_streaming_response(self, response) -> str:176        """177        Handle streaming response from Ollama API.178        179        Args:180            response: The streaming response object181            182        Returns:183            str: The complete response text184        """185        full_response = ""186        for line in response.iter_lines():187            if line:188                try:189                    data = json.loads(line.decode('utf-8'))190                    if 'response' in data:191                        full_response += data['response']192                    if data.get('done', False):193                        break194                except json.JSONDecodeError:195                    continue196        return full_response197    198    def chat(self, messages: list, stream: bool = False, **kwargs) -> str:199        """200        Chat with the LLM using a conversation format.201        202        Args:203            messages (list): List of message dictionaries with 'role' and 'content' keys204            stream (bool): Whether to stream the response. Defaults to False205            **kwargs: Additional parameters for the API call206            207        Returns:208            str: The generated response text209        """210        # Convert messages to a single prompt for qwen2.5:1.5b211        prompt = self._format_messages_as_prompt(messages)212        return self.generate_response(prompt, stream, **kwargs)213    214    def _format_messages_as_prompt(self, messages: list) -> str:215        """216        Format conversation messages into a single prompt.217        218        Args:219            messages (list): List of message dictionaries220            221        Returns:222            str: Formatted prompt string223        """224        formatted_prompt = ""225        for message in messages:226            role = message.get('role', 'user')227            content = message.get('content', '')228            229            if role == 'system':230                formatted_prompt += f"System: {content}\n\n"231            elif role == 'user':232                formatted_prompt += f"User: {content}\n\n"233            elif role == 'assistant':234                formatted_prompt += f"Assistant: {content}\n\n"235        236        formatted_prompt += "Assistant:"237        return formatted_prompt238    239    def is_model_available(self) -> bool:240        """241        Check if the specified model is available in Ollama.242        243        Returns:244            bool: True if model is available, False otherwise245        """246        try:247            response = requests.get(f"{self.base_url}/api/tags", timeout=5)248            response.raise_for_status()249            250            models = response.json().get('models', [])251            model_names = [model['name'] for model in models]252            253            return self.model in model_names254            255        except requests.exceptions.RequestException:256            return False257    258    def pull_model(self) -> bool:259        """260        Pull the model if it's not available locally.261        262        Returns:263            bool: True if model was pulled successfully, False otherwise264        """265        try:266            payload = {"name": self.model}267            response = requests.post(268                f"{self.base_url}/api/pull",269                json=payload,270                headers={"Content-Type": "application/json"},271                timeout=300  # 5 minutes timeout for model pulling272            )273            response.raise_for_status()274            return True275            276        except requests.exceptions.RequestException:277            return False278 279 280# Example usage and testing281if __name__ == "__main__":282    # Initialize the LLM283    llm = LocalLLM()284    285    # Check if model is available286    if not llm.is_model_available():287        print(f"Model {llm.model} not found. Attempting to pull...")288        if llm.pull_model():289            print("Model pulled successfully!")290        else:291            print("Failed to pull model. Please ensure Ollama is running.")292            exit(1)293    294    # Example 1: Simple prompt295    try:296        response = llm.generate_response("Hello! How are you today?")297        print("Response:", response)298    except Exception as e:299        print(f"Error: {e}")300    301    # Example 2: Chat conversation302    try:303        messages = [304            {"role": "user", "content": "What is the capital of France?"},305            {"role": "assistant", "content": "The capital of France is Paris."},306            {"role": "user", "content": "What is the population of Paris?"}307        ]308        response = llm.chat(messages)309        print("Chat response:", response)310    except Exception as e:311        print(f"Error: {e}")312