llmsecops-research/src/text_generation/services/guidelines/base_security_guidelines_service.py

from typing import Optional
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate, StringPromptTemplate
from langchain_core.prompt_values import PromptValue
from langchain_core.runnables import RunnablePassthrough
from langchain.prompts import FewShotPromptTemplate

from src.text_generation.common.constants import Constants
from src.text_generation.domain.abstract_guidelines_processed_completion import AbstractGuidelinesProcessedCompletion
from src.text_generation.domain.guidelines_result import GuidelinesResult
from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel
from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesConfigurationBuilder, AbstractSecurityGuidelinesService
from src.text_generation.services.nlp.abstract_prompt_template_service import AbstractPromptTemplateService
from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService


class BaseSecurityGuidelinesService(AbstractSecurityGuidelinesService):
    """Base service for security guidelines implementations."""
    
    def __init__(
            self,
            foundation_model: AbstractFoundationModel,
            response_processing_service: AbstractResponseProcessingService,
            prompt_template_service: AbstractPromptTemplateService,
            config_builder: Optional[AbstractSecurityGuidelinesConfigurationBuilder] = None):
        super().__init__()
        self.constants = Constants()
        self.foundation_model_pipeline = foundation_model.create_pipeline()
        self.response_processing_service = response_processing_service
        self.prompt_template_service = prompt_template_service
        self.config_builder = config_builder

    def _create_chain(self, prompt_template: PromptTemplate):

        if prompt_template is None:
            raise ValueError("prompt_template cannot be None")

        return (
            { f"{self.constants.INPUT_VARIABLE_TOKEN}": RunnablePassthrough() }
            | prompt_template
            | self.foundation_model_pipeline
            | StrOutputParser()
            | self.response_processing_service.process_text_generation_output
        )

    def _get_template(self, user_prompt: str) -> StringPromptTemplate:
        """
        Get the prompt template for security guidelines.
        
        Returns:
            StringPromptTemplate: Template for processing security guidelines
        """
        raise NotImplementedError("Subclasses must implement _get_template()")

    def _find_llm_step(self, chain):
        if hasattr(chain, 'steps'):
            for i, step in enumerate(chain.steps):
                if step.__class__.__name__ == 'HuggingFacePipeline':
                    return step
        return None

    def _extract_llm_config(self, llm_step):
        if not llm_step:
            return {}
        
        full_config = llm_step.model_dump()
        
        serializable_keys = [
            'batch_size', 
            'device', 
            'do_sample', 
            'temperature', 
            'top_p', 
            'top_k', 
            'max_new_tokens', 
            'max_length',
            'repetition_penalty', 
            'pad_token_id', 
            'eos_token_id',
            'model_id', 
            'task', 
            'return_full_text'
        ]
        
        config = {}
        for key, value in full_config.items():
            if key in serializable_keys and isinstance(value, (str, int, float, bool, type(None))):
                config[key] = value
        return config


    def apply_guidelines(self, user_prompt: str) -> AbstractGuidelinesProcessedCompletion:
        print(f'applying guidelines (if any set)')
        if not user_prompt:
            raise ValueError(f"Parameter 'user_prompt' cannot be empty or None")
        
        try:
            prompt_template: StringPromptTemplate = self._get_template(user_prompt=user_prompt)
            print(f'got prompt template')
            prompt_value: PromptValue = prompt_template.format_prompt(input=user_prompt)

            # Create a comprehensive dict
            prompt_dict = {
                "messages": [
                    {"role": msg.type, "content": msg.content, "additional_kwargs": msg.additional_kwargs}
                    for msg in prompt_value.to_messages()
                ],
                "string_representation": prompt_value.to_string(),
            }

            print(f'creating chain...')
            chain = self._create_chain(prompt_template)
            print(f'Chain type: {type(chain)}')
            print(f'Number of steps: {len(chain.steps) if hasattr(chain, "steps") else "No steps attribute"}')

            # Print each step to see what's at each position
            if hasattr(chain, 'steps'):
                for i, step in enumerate(chain.steps):
                    print(f'Step {i}: {type(step)} - {step.__class__.__name__}')
            print(f'generating completion...')
            completion_text=chain.invoke({"input": user_prompt})
            llm_step = self._find_llm_step(chain)
            llm_config = self._extract_llm_config(llm_step)
            result = GuidelinesResult(
                completion_text=completion_text,
                llm_config=llm_config,
                full_prompt=prompt_dict
            )
            return result
        except Exception as e:
            raise e