import dataclasses import os import re import time from typing import Any, Dict, List, Tuple from uuid import uuid1 import loguru import requests from tenacity import * from pentestgpt.utils.llm_api import LLMAPI logger = loguru.logger logger.remove() @dataclasses.dataclass class Message: ask_id: str = None ask: dict = None answer: dict = None answer_id: str = None request_start_timestamp: float = None request_end_timestamp: float = None time_escaped: float = None @dataclasses.dataclass class Conversation: conversation_id: str = None message_list: List[Message] = dataclasses.field(default_factory=list) def __hash__(self): return hash(self.conversation_id) def __eq__(self, other): if not isinstance(other, Conversation): return False return self.conversation_id == other.conversation_id class OllamaAPI(LLMAPI): def __init__(self, config_class, use_langfuse_logging=False): self.name = str(config_class.model) self.history_length = 10 # maintain conversation history self.config = config_class self.model = config_class.model self.api_base = config_class.api_base self.log_dir = config_class.log_dir self.conversation_dict: Dict[str, Conversation] = {} self.use_langfuse_logging = use_langfuse_logging if self.log_dir: logger.add(sink=os.path.join(self.log_dir, "ollama.log"), level="WARNING") # Test connection to Ollama self._test_connection() def _test_connection(self): """Test if Ollama server is running and model is available""" try: # Check if Ollama server is running response = requests.get(f"{self.api_base}/api/tags", timeout=5) if response.status_code != 200: raise Exception(f"Ollama server not accessible at {self.api_base}") # Check if model is available models = response.json().get("models", []) model_names = [model["name"] for model in models] if self.model not in model_names: print(f"Warning: Model '{self.model}' not found in Ollama.") print(f"Available models: {', '.join(model_names)}") print(f"You can pull the model using: ollama pull {self.model}") except requests.exceptions.RequestException as e: raise Exception(f"Cannot connect to Ollama server at {self.api_base}: {e}") @retry( wait=wait_random_exponential(min=1, max=10), stop=stop_after_attempt(3), retry=retry_if_exception_type(Exception), ) def _chat_completion(self, messages: List[Dict]) -> str: """ Send chat completion request to Ollama Parameters: ----------- messages: List[Dict] List of message dictionaries with 'role' and 'content' keys Returns: -------- str: Response content from the model """ try: # Convert messages to Ollama format payload = { "model": self.model, "messages": messages, "stream": False } response = requests.post( f"{self.api_base}/api/chat", json=payload, timeout=120 ) if response.status_code != 200: raise Exception(f"Ollama API error: {response.status_code} - {response.text}") result = response.json() return result["message"]["content"] except requests.exceptions.Timeout: raise Exception("Ollama request timed out") except requests.exceptions.RequestException as e: raise Exception(f"Ollama request failed: {e}") except KeyError as e: raise Exception(f"Unexpected Ollama response format: {e}") def send_new_message(self, message: str, conversation_id: str = None) -> Tuple[str, str]: """ Send a new message and start a new conversation Parameters: ----------- message: str The message to send conversation_id: str, optional If provided, use this conversation ID Returns: -------- Tuple[str, str]: (response_content, conversation_id) """ if conversation_id is None: conversation_id = str(uuid1()) # Create new conversation conversation = Conversation(conversation_id=conversation_id) self.conversation_dict[conversation_id] = conversation # Send message messages = [{"role": "user", "content": message}] msg_obj = Message( ask_id=str(uuid1()), ask={"role": "user", "content": message}, request_start_timestamp=time.time(), ) try: response_content = self._chat_completion(messages) msg_obj.request_end_timestamp = time.time() msg_obj.time_escaped = msg_obj.request_end_timestamp - msg_obj.request_start_timestamp msg_obj.answer = {"role": "assistant", "content": response_content} msg_obj.answer_id = str(uuid1()) conversation.message_list.append(msg_obj) return response_content, conversation_id except Exception as e: logger.error(f"Error in send_new_message: {e}") raise e def send_message(self, message: str, conversation_id: str) -> str: """ Send a message in an existing conversation Parameters: ----------- message: str The message to send conversation_id: str The conversation ID to continue Returns: -------- str: Response content from the model """ if conversation_id not in self.conversation_dict: # If conversation doesn't exist, create new one response, _ = self.send_new_message(message, conversation_id) return response conversation = self.conversation_dict[conversation_id] # Build message history messages = [] # Add recent messages from history (limited by history_length) recent_messages = conversation.message_list[-self.history_length:] for msg in recent_messages: if msg.ask: messages.append(msg.ask) if msg.answer: messages.append(msg.answer) # Add current message messages.append({"role": "user", "content": message}) msg_obj = Message( ask_id=str(uuid1()), ask={"role": "user", "content": message}, request_start_timestamp=time.time(), ) try: response_content = self._chat_completion(messages) msg_obj.request_end_timestamp = time.time() msg_obj.time_escaped = msg_obj.request_end_timestamp - msg_obj.request_start_timestamp msg_obj.answer = {"role": "assistant", "content": response_content} msg_obj.answer_id = str(uuid1()) conversation.message_list.append(msg_obj) return response_content except Exception as e: logger.error(f"Error in send_message: {e}") raise e def get_conversation_history(self, conversation_id: str = None) -> List[Dict]: """ Get conversation history Parameters: ----------- conversation_id: str, optional The conversation ID. If None, returns None Returns: -------- List[Dict] or None: List of messages or None if conversation not found """ if conversation_id is None or conversation_id not in self.conversation_dict: return None conversation = self.conversation_dict[conversation_id] history = [] for msg in conversation.message_list: if msg.ask: history.append(msg.ask) if msg.answer: history.append(msg.answer) return history def refresh(self): """ Refresh the API connection - for Ollama, just test connection """ self._test_connection() logger.info("Ollama connection refreshed") def get_model_list(self) -> List[str]: """ Get list of available models from Ollama Returns: -------- List[str]: List of available model names """ try: response = requests.get(f"{self.api_base}/api/tags", timeout=10) if response.status_code == 200: models = response.json().get("models", []) return [model["name"] for model in models] else: return [] except Exception as e: logger.error(f"Error getting model list: {e}") return []