AI Services

AIService Class

ai_services.AIService

AIService handles interactions with multiple AI services (OpenAI, GCP, and Azure).

Attributes:

Name	Type	Description
technology	str	The technology field from the config.
domain	str	The domain field from the config.
segment	str	The segment field from the config.
anonymous_access	bool	The anonymous access field from the config.
debug_mode	bool	Flag for enabling debug mode.
gcp_project	str	GCP project ID for Gemini API Vertex.
gcp_location	str	GCP location for Gemini API Vertex.
gcp_model_id	str	Model ID for Gemini API Vertex.
azure_endpoint	str	Azure OpenAI endpoint.
azure_location	str	Azure OpenAI location/region.

Source code in src/ai_services.py

class AIService:
    """
    AIService handles interactions with multiple AI services (OpenAI, GCP, and Azure).

    Attributes:
        technology (str): The technology field from the config.
        domain (str): The domain field from the config.
        segment (str): The segment field from the config.
        anonymous_access (bool): The anonymous access field from the config.
        debug_mode (bool): Flag for enabling debug mode.
        gcp_project (str): GCP project ID for Gemini API Vertex.
        gcp_location (str): GCP location for Gemini API Vertex.
        gcp_model_id (str): Model ID for Gemini API Vertex.
        azure_endpoint (str): Azure OpenAI endpoint.
        azure_location (str): Azure OpenAI location/region.
    """

    def __init__(self, api_key=False, gcp_project=None, gcp_location=None, gcp_model_id=None, azure_endpoint=None, azure_location=None, debug_mode=False):
        """
        Initialize AIService with API key for OpenAI, Azure, and GCP project details.

        Args:
            api_key (str): The API key for OpenAI or Azure.
            gcp_project (str): GCP project ID for Google AI.
            gcp_location (str): GCP location for Google AI.
            gcp_model_id (str): Model ID for Google AI.
            azure_endpoint (str): The Azure OpenAI API endpoint.
            azure_location (str): The Azure OpenAI API location/region.
            debug_mode (bool): If True, enables debug logging.
        """
        self.technology = config.get('server', 'technology', fallback='generic')
        self.domain = config.get('server', 'domain', fallback='localhost')
        self.segment = config.get('server', 'segment', fallback='general')
        self.anonymous_access = config.getboolean('server', 'anonymous_access', fallback=False)

        openai.api_key = api_key  # Set the API key directly

        self.gcp_project = gcp_project
        self.gcp_location = gcp_location
        self.gcp_model_id = gcp_model_id
        self.azure_endpoint = azure_endpoint  # New attribute for Azure OpenAI
        self.azure_location = azure_location  # New attribute for Azure location
        self.debug_mode = debug_mode

        if self.debug_mode:
            logging.getLogger('ai_services').setLevel(logging.DEBUG)
            logging.getLogger('urllib3').setLevel(logging.DEBUG)
        else:
            logging.getLogger('ai_services').setLevel(logging.CRITICAL)
            logging.getLogger('urllib3').setLevel(logging.CRITICAL)

    def query_responses(self, prompt, response_type, use_openai=True):
        """
        Query AI services (OpenAI or Google Gemini Vertex) for responses based on the provided prompt and response type.

        Args:
            prompt (str): The prompt to send to the AI service.
            response_type (str): The type of response expected (e.g., "email").
            use_openai (bool): Whether to use OpenAI (True) or Gemini Vertex (False).

        Returns:
            str: The response text from the AI service.
        """
        if use_openai:
            return self._query_openai(prompt, response_type)
        else:
            return self._query_gcp_gemini(prompt, response_type)

    def _query_openai(self, prompt, response_type):
        """
        Query OpenAI's API for a response to the provided prompt.

        Args:
            prompt (str): The prompt to send to OpenAI.
            response_type (str): The type of response expected (e.g., "email").

        Returns:
            str: The response text from OpenAI, or an empty string if there was an error.
        """
        for attempt in range(2):
            try:
                if self.debug_mode:
                    logger.debug(f"Querying OpenAI for {response_type} responses...")
                response = openai.ChatCompletion.create(
                    model="gpt-4",
                    messages=[
                        {"role": "system", "content": "You are a helpful assistant."},
                        {"role": "user", "content": prompt}
                    ],
                    max_tokens=500
                )
                response_text = response.choices[0]['message']['content'].strip()
                self._save_raw_response(response_text, response_type)
                return response_text
            except Exception as e:
                if self.debug_mode:
                    logger.error(f"Error querying OpenAI (attempt {attempt+1}/2): {e}")
                if attempt == 1:
                    logger.critical("Failed to communicate with AI after 2 attempts. Exiting.")
                time.sleep(1)
        return ""

    def _query_gcp_gemini(self, prompt, response_type):
        """
        Query Google's Gemini API Vertex for a response to the provided prompt.

        Args:
            prompt (str): The prompt to send to Google Gemini API.
            response_type (str): The type of response expected (e.g., "email").

        Returns:
            str: The response text from Google Gemini Vertex, or an empty string if there was an error.
        """
        return()
        try:
            if self.debug_mode:
                logger.debug(f"Querying Google Gemini Vertex for {response_type} responses...")

            client = aiplatform.gapic.PredictionServiceClient()
            endpoint = f"projects/{self.gcp_project}/locations/{self.gcp_location}/endpoints/{self.gcp_model_id}"

            instances = [{"content": prompt}]
            parameters = {}
            request = predict.instance.PredictRequest(
                endpoint=endpoint,
                instances=[json_format.ParseDict(instances, predict.instance.Value())],
                parameters=json_format.ParseDict(parameters, predict.instance.Value()),
            )
            response = client.predict(request=request)
            response_text = response.predictions[0].get("content", "").strip()
            self._save_raw_response(response_text, response_type)
            return response_text
        except Exception as e:
            if self.debug_mode:
                logger.error(f"Error querying Google Gemini Vertex: {e}")
            return ""

    def _save_raw_response(self, response_text, response_type):
        """
        Save the raw response text to a file.

        Args:
            response_text (str): The response text from the AI service.
            response_type (str): The type of response (e.g., "email").
        """
        filename = f'files/{response_type}_raw_response.txt'
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(response_text)
        if self.debug_mode:
            logger.debug(f"Raw response saved in {filename}")

    def _store_responses(self, responses, response_type):
        """
        Store the parsed responses in a JSON file.

        Args:
            responses (dict): The parsed responses.
            response_type (str): The type of response (e.g., "email").
        """
        filename = f'files/{response_type}_responses.json'
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(responses, f)
        if self.debug_mode:
            logger.debug(f"Responses stored in {filename}")

    def load_responses(self, response_type):
        """
        Load responses from a saved file.

        Args:
            response_type (str): The type of response (e.g., "smtp").

        Returns:
            dict or str: The loaded response as a JSON dictionary if valid, 
                        otherwise returns the raw text as a string.
        """
        filename = f'files/{response_type}_response.txt'
        if os.path.exists(filename):
            with open(filename, 'r', encoding='utf-8') as f:
                try:
                    # Try to parse the file as JSON
                    return json.load(f)
                except json.JSONDecodeError:
                    # If parsing fails, return the raw text
                    f.seek(0)  # Reset file pointer to the beginning
                    raw_text = f.read()
                    logger.error(f"Failed to parse {filename} as JSON. Returning raw text.")
                    return raw_text
        return {}  # Return empty dictionary if no file is found

    def cleanup_and_parse_json(self, text):
        """
        Clean up and parse a JSON string from text.

        Args:
            text (str): The text containing JSON.

        Returns:
            dict: The parsed JSON object, or an empty dict if parsing fails.
        """
        try:
            start = text.find('{')
            end = text.rfind('}') + 1
            if start == -1 or end == 0:
                if self.debug_mode:
                    logger.error("Invalid JSON structure detected.")
                    logger.debug(f"Raw text: {text}")
                return {}

            json_text = text[start:end]
            return json.loads(json_text)
        except json.JSONDecodeError as e:
            if self.debug_mode:
                logger.error(f"Error parsing JSON response: {e}")
                logger.debug(f"Raw text for cleanup: {text}")
            return {}

    def generate_emails(self, segment, domain, email_num):
        """
        Generate a sample email related to the given segment and domain.

        This method uses the OpenAI API to generate a sample email, including
        the subject, body, and recipient address. The email content is saved
        to a file for later use.

        Args:
            segment (str): The segment or topic of the email.
            domain (str): The domain to use for the email address.
            email_num (int): The identifier number for the email.

        Returns:
            str: The generated email content.
        """
        try:
            prompt = (
                f"Generate an email related to the segment: {segment} for the domain {domain}. "
                f"The email should include a subject, body, and a recipient address at the domain."
            )
            response = openai.ChatCompletion.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": prompt}
                ],
                max_tokens=500
            )
            response_text = response.choices[0]['message']['content'].strip()

            # Save the raw response to a file
            filename = f'files/email{email_num}_raw_response.txt'
            with open(filename, 'w', encoding='utf-8') as f:
                f.write(response_text)
            if self.debug_mode:
                logger.debug(f"Raw response saved in {filename}")

            return response_text
        except Exception as e:
            if self.debug_mode:
                logger.error(f"Error querying OpenAI for email {email_num}: {e}")
            return "No response"

__init__(api_key=False, gcp_project=None, gcp_location=None, gcp_model_id=None, azure_endpoint=None, azure_location=None, debug_mode=False)

Initialize AIService with API key for OpenAI, Azure, and GCP project details.

Parameters:

Name	Type	Description	Default
api_key	str	The API key for OpenAI or Azure.	False
gcp_project	str	GCP project ID for Google AI.	None
gcp_location	str	GCP location for Google AI.	None
gcp_model_id	str	Model ID for Google AI.	None
azure_endpoint	str	The Azure OpenAI API endpoint.	None
azure_location	str	The Azure OpenAI API location/region.	None
debug_mode	bool	If True, enables debug logging.	False

Source code in src/ai_services.py

def __init__(self, api_key=False, gcp_project=None, gcp_location=None, gcp_model_id=None, azure_endpoint=None, azure_location=None, debug_mode=False):
    """
    Initialize AIService with API key for OpenAI, Azure, and GCP project details.

    Args:
        api_key (str): The API key for OpenAI or Azure.
        gcp_project (str): GCP project ID for Google AI.
        gcp_location (str): GCP location for Google AI.
        gcp_model_id (str): Model ID for Google AI.
        azure_endpoint (str): The Azure OpenAI API endpoint.
        azure_location (str): The Azure OpenAI API location/region.
        debug_mode (bool): If True, enables debug logging.
    """
    self.technology = config.get('server', 'technology', fallback='generic')
    self.domain = config.get('server', 'domain', fallback='localhost')
    self.segment = config.get('server', 'segment', fallback='general')
    self.anonymous_access = config.getboolean('server', 'anonymous_access', fallback=False)

    openai.api_key = api_key  # Set the API key directly

    self.gcp_project = gcp_project
    self.gcp_location = gcp_location
    self.gcp_model_id = gcp_model_id
    self.azure_endpoint = azure_endpoint  # New attribute for Azure OpenAI
    self.azure_location = azure_location  # New attribute for Azure location
    self.debug_mode = debug_mode

    if self.debug_mode:
        logging.getLogger('ai_services').setLevel(logging.DEBUG)
        logging.getLogger('urllib3').setLevel(logging.DEBUG)
    else:
        logging.getLogger('ai_services').setLevel(logging.CRITICAL)
        logging.getLogger('urllib3').setLevel(logging.CRITICAL)

cleanup_and_parse_json(text)

Clean up and parse a JSON string from text.

Parameters:

Name	Type	Description	Default
text	str	The text containing JSON.	required

Returns:

Name	Type	Description
dict		The parsed JSON object, or an empty dict if parsing fails.

Source code in src/ai_services.py

def cleanup_and_parse_json(self, text):
    """
    Clean up and parse a JSON string from text.

    Args:
        text (str): The text containing JSON.

    Returns:
        dict: The parsed JSON object, or an empty dict if parsing fails.
    """
    try:
        start = text.find('{')
        end = text.rfind('}') + 1
        if start == -1 or end == 0:
            if self.debug_mode:
                logger.error("Invalid JSON structure detected.")
                logger.debug(f"Raw text: {text}")
            return {}

        json_text = text[start:end]
        return json.loads(json_text)
    except json.JSONDecodeError as e:
        if self.debug_mode:
            logger.error(f"Error parsing JSON response: {e}")
            logger.debug(f"Raw text for cleanup: {text}")
        return {}

generate_emails(segment, domain, email_num)

Generate a sample email related to the given segment and domain.

This method uses the OpenAI API to generate a sample email, including the subject, body, and recipient address. The email content is saved to a file for later use.

Parameters:

Name	Type	Description	Default
segment	str	The segment or topic of the email.	required
domain	str	The domain to use for the email address.	required
email_num	int	The identifier number for the email.	required

Returns:

Name	Type	Description
str		The generated email content.

Source code in src/ai_services.py

def generate_emails(self, segment, domain, email_num):
    """
    Generate a sample email related to the given segment and domain.

    This method uses the OpenAI API to generate a sample email, including
    the subject, body, and recipient address. The email content is saved
    to a file for later use.

    Args:
        segment (str): The segment or topic of the email.
        domain (str): The domain to use for the email address.
        email_num (int): The identifier number for the email.

    Returns:
        str: The generated email content.
    """
    try:
        prompt = (
            f"Generate an email related to the segment: {segment} for the domain {domain}. "
            f"The email should include a subject, body, and a recipient address at the domain."
        )
        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=500
        )
        response_text = response.choices[0]['message']['content'].strip()

        # Save the raw response to a file
        filename = f'files/email{email_num}_raw_response.txt'
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(response_text)
        if self.debug_mode:
            logger.debug(f"Raw response saved in {filename}")

        return response_text
    except Exception as e:
        if self.debug_mode:
            logger.error(f"Error querying OpenAI for email {email_num}: {e}")
        return "No response"

load_responses(response_type)

Load responses from a saved file.

Parameters:

Name	Type	Description	Default
response_type	str	The type of response (e.g., "smtp").	required

Returns:

Type	Description
	dict or str: The loaded response as a JSON dictionary if valid, otherwise returns the raw text as a string.

Source code in src/ai_services.py

def load_responses(self, response_type):
    """
    Load responses from a saved file.

    Args:
        response_type (str): The type of response (e.g., "smtp").

    Returns:
        dict or str: The loaded response as a JSON dictionary if valid, 
                    otherwise returns the raw text as a string.
    """
    filename = f'files/{response_type}_response.txt'
    if os.path.exists(filename):
        with open(filename, 'r', encoding='utf-8') as f:
            try:
                # Try to parse the file as JSON
                return json.load(f)
            except json.JSONDecodeError:
                # If parsing fails, return the raw text
                f.seek(0)  # Reset file pointer to the beginning
                raw_text = f.read()
                logger.error(f"Failed to parse {filename} as JSON. Returning raw text.")
                return raw_text
    return {}  # Return empty dictionary if no file is found

query_responses(prompt, response_type, use_openai=True)

Query AI services (OpenAI or Google Gemini Vertex) for responses based on the provided prompt and response type.

Parameters:

Name	Type	Description	Default
prompt	str	The prompt to send to the AI service.	required
response_type	str	The type of response expected (e.g., "email").	required
use_openai	bool	Whether to use OpenAI (True) or Gemini Vertex (False).	True

Returns:

Name	Type	Description
str		The response text from the AI service.

Source code in src/ai_services.py

def query_responses(self, prompt, response_type, use_openai=True):
    """
    Query AI services (OpenAI or Google Gemini Vertex) for responses based on the provided prompt and response type.

    Args:
        prompt (str): The prompt to send to the AI service.
        response_type (str): The type of response expected (e.g., "email").
        use_openai (bool): Whether to use OpenAI (True) or Gemini Vertex (False).

    Returns:
        str: The response text from the AI service.
    """
    if use_openai:
        return self._query_openai(prompt, response_type)
    else:
        return self._query_gcp_gemini(prompt, response_type)