From c7884314167622e0e99560721fdeda1b477f1bf9 Mon Sep 17 00:00:00 2001 From: Adam Wilson Date: Sat, 12 Jul 2025 12:18:19 -0600 Subject: [PATCH] prompt template IDs; fluent text generation service stubs --- docs/to-do.md | 3 + src/text_generation/common/constants.py | 6 +- ...tract_text_generation_completion_result.py | 13 +++ .../domain/text_generation_response.py | 0 ..._semantic_similarity_guidelines_service.py | 31 +++++-- ...ract_text_generation_completion_service.py | 25 +++++- .../nlp/text_generation_completion_service.py | 84 ++++++++++++++++--- tests/integration/test_violation_rate.py | 14 ++-- 8 files changed, 147 insertions(+), 29 deletions(-) create mode 100644 src/text_generation/domain/abstract_text_generation_completion_result.py delete mode 100644 src/text_generation/domain/text_generation_response.py diff --git a/docs/to-do.md b/docs/to-do.md index 114b7c80b..b78b0f967 100644 --- a/docs/to-do.md +++ b/docs/to-do.md @@ -8,7 +8,10 @@ ### Prompt Templates [ X ] Base Phi-3 template +[ ] CoT template [ ] Few Shot template with examples +[ ] Reflextion template + ### Prompt Templates: Supporting Logic diff --git a/src/text_generation/common/constants.py b/src/text_generation/common/constants.py index eb4484619..bed11930b 100644 --- a/src/text_generation/common/constants.py +++ b/src/text_generation/common/constants.py @@ -4,6 +4,8 @@ class Constants: SYSTEM_TOKEN = "<|system|>" USER_TOKEN = "<|user|>" - # prompt template IDs class PromptTemplateIds: - PHI_3_MINI_4K_INSTRUCT_BASIC = "phi-3-mini-4k-instruct-basic" \ No newline at end of file + PHI_3_MINI_4K_INSTRUCT_BASIC = "phi-3-mini-4k-instruct-basic" + PHI_3_MINI_4K_INSTRUCT_CHAIN_OF_THOUGHT = "phi-3-mini-4k-instruct-cot" + PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_EXAMPLES = "phi-3-mini-4k-instruct-few-shot" + PHI_3_MINI_4K_INSTRUCT_REFLEXION = "phi-3-mini-4k-instruct-reflexion" \ No newline at end of file diff --git a/src/text_generation/domain/abstract_text_generation_completion_result.py b/src/text_generation/domain/abstract_text_generation_completion_result.py new file mode 100644 index 000000000..4dbff9ce8 --- /dev/null +++ b/src/text_generation/domain/abstract_text_generation_completion_result.py @@ -0,0 +1,13 @@ +import abc + +class AbstractTextGenerationCompletionResult(abc.ABC): + + @abc.abstractmethod + def get_text(self) -> str: + """Return the generated text.""" + pass + + @abc.abstractmethod + def get_metadata(self) -> dict: + """Return metadata about the generation.""" + pass diff --git a/src/text_generation/domain/text_generation_response.py b/src/text_generation/domain/text_generation_response.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/text_generation/services/guidelines/rag_enhanced_semantic_similarity_guidelines_service.py b/src/text_generation/services/guidelines/rag_enhanced_semantic_similarity_guidelines_service.py index 27cec3f88..40f813e2d 100644 --- a/src/text_generation/services/guidelines/rag_enhanced_semantic_similarity_guidelines_service.py +++ b/src/text_generation/services/guidelines/rag_enhanced_semantic_similarity_guidelines_service.py @@ -1,11 +1,28 @@ -from src.text_generation.services.guidelines.abstract_rag_enhanced_semantic_similarity_guidelines_service import AbstractRagEnhancedSemanticSimilarityGuidelinesService +import abc -class RagEnhancedSemanticSimilarityGuidelinesService(AbstractRagEnhancedSemanticSimilarityGuidelinesService): - def analyze(self, prompt_input_text: str) -> float: +class AbstractChainOfThoughtSecurityGuidelinesService(abc.ABC): + """Abstract service for chain of thought security guidelines.""" + + @abc.abstractmethod + def apply_guidelines(self, context: dict) -> dict: + """Apply chain of thought security guidelines to context.""" + pass - - # TODO - check semantic similarity score - # TODO - retry with summarized prompt? task decomposition - result could contain original score and improved score - raise NotImplementedError \ No newline at end of file +class AbstractRetrievalAugmentedGenerationContextSecurityGuidelinesService(abc.ABC): + """Abstract service for RAG context security guidelines.""" + + @abc.abstractmethod + def apply_guidelines(self, context: dict) -> dict: + """Apply RAG context security guidelines to context.""" + pass + + +class AbstractReflexionSecurityGuidelinesService(abc.ABC): + """Abstract service for reflexion security guidelines.""" + + @abc.abstractmethod + def apply_guidelines(self, context: dict) -> dict: + """Apply reflexion security guidelines to context.""" + pass \ No newline at end of file diff --git a/src/text_generation/services/nlp/abstract_text_generation_completion_service.py b/src/text_generation/services/nlp/abstract_text_generation_completion_service.py index c78e62c72..579402bc0 100644 --- a/src/text_generation/services/nlp/abstract_text_generation_completion_service.py +++ b/src/text_generation/services/nlp/abstract_text_generation_completion_service.py @@ -1,7 +1,30 @@ import abc +from src.text_generation.domain.abstract_text_generation_completion_result import AbstractTextGenerationCompletionResult + class AbstractTextGenerationCompletionService(abc.ABC): + @abc.abstractmethod - def invoke(self, user_prompt: str) -> str: + def without_guidelines(self) -> 'AbstractTextGenerationCompletionService': + """Skip all security guidelines.""" + raise NotImplementedError + + @abc.abstractmethod + def with_chain_of_thought_guidelines(self) -> 'AbstractTextGenerationCompletionService': + """Enable chain of thought security guidelines.""" + raise NotImplementedError + + @abc.abstractmethod + def with_rag_context_guidelines(self) -> 'AbstractTextGenerationCompletionService': + """Enable RAG context security guidelines.""" + raise NotImplementedError + + @abc.abstractmethod + def with_reflexion_guidelines(self) -> 'AbstractTextGenerationCompletionService': + """Enable reflexion security guidelines.""" + raise NotImplementedError + + @abc.abstractmethod + def invoke(self, user_prompt: str) -> AbstractTextGenerationCompletionResult: raise NotImplementedError \ No newline at end of file diff --git a/src/text_generation/services/nlp/text_generation_completion_service.py b/src/text_generation/services/nlp/text_generation_completion_service.py index a25a6a094..0633e1cee 100644 --- a/src/text_generation/services/nlp/text_generation_completion_service.py +++ b/src/text_generation/services/nlp/text_generation_completion_service.py @@ -8,29 +8,61 @@ from src.text_generation.services.nlp.abstract_text_generation_completion_servic from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel -class TextGenerationCompletionService(AbstractTextGenerationCompletionService): - +class TextGenerationCompletionService( + AbstractTextGenerationCompletionService): def __init__( self, foundation_model: AbstractFoundationModel, - prompt_template_service: AbstractPromptTemplateService): + prompt_template_service: AbstractPromptTemplateService, + chain_of_thought_service: AbstractChainOfThoughtSecurityGuidelinesService, + rag_context_service: AbstractRetrievalAugmentedGenerationContextSecurityGuidelinesService, + reflexion_service: AbstractReflexionSecurityGuidelinesService): super().__init__() - self.language_model_pipeline = foundation_model.create_pipeline() - self.prompt_template_service = prompt_template_service self.constants = Constants() + + self._language_model_pipeline = foundation_model.create_pipeline() + self._prompt_template_service = prompt_template_service + self._chain_of_thought_service = chain_of_thought_service + self._rag_context_service = rag_context_service + self._reflexion_service = reflexion_service + + self._use_guidelines = True + self._use_chain_of_thought = True + self._use_rag_context = True + self._use_reflexion = True def _extract_assistant_response(self, text): if self.constants.ASSISTANT_TOKEN in text: return text.split(self.constants.ASSISTANT_TOKEN)[-1].strip() return text + def without_guidelines(self) -> AbstractTextGenerationCompletionService: + """Skip all security guidelines.""" + self._use_guidelines = False + return self + + def with_chain_of_thought_guidelines(self) -> AbstractTextGenerationCompletionService: + """Enable chain-of-thought (CoT) security guidelines.""" + self._use_chain_of_thought = True + return self + + def with_rag_context_guidelines(self) -> AbstractTextGenerationCompletionService: + """Enable RAG-enriched examples context security guidelines.""" + self._use_rag_context = True + return self + + def with_reflexion_guidelines(self) -> AbstractTextGenerationCompletionService: + """Enable reflexion security guidelines.""" + self._use_reflexion = True + return self + def create_chain(self): prompt_template_id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_BASIC - prompt_template = self.prompt_template_service.get(id=prompt_template_id) + prompt_template = self._prompt_template_service.get(id=prompt_template_id) return ( { "question": RunnablePassthrough() } | prompt_template - | self.language_model_pipeline + | self._language_model_pipeline | StrOutputParser() | self._extract_assistant_response ) @@ -38,8 +70,36 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService): def invoke(self, user_prompt: str) -> str: if not user_prompt: raise ValueError(f"Parameter 'user_prompt' cannot be empty or None") - try: - chain = self.create_chain() - return chain.invoke(user_prompt) - except Exception as e: - raise e \ No newline at end of file + + if self._use_guidelines == False: + try: + chain = self.create_chain() + return chain.invoke(user_prompt) + except Exception as e: + raise e + + # security guidelines combinations + if self._use_chain_of_thought and self._use_rag_context and self._use_reflexion: + # All three enabled: CoT + RAG + Reflexion + pass + elif self._use_chain_of_thought and self._use_rag_context and not self._use_reflexion: + # CoT + RAG only + pass + elif self._use_chain_of_thought and not self._use_rag_context and self._use_reflexion: + # CoT + Reflexion only + pass + elif self._use_chain_of_thought and not self._use_rag_context and not self._use_reflexion: + # CoT only + pass + elif not self._use_chain_of_thought and self._use_rag_context and self._use_reflexion: + # RAG + Reflexion only + pass + elif not self._use_chain_of_thought and self._use_rag_context and not self._use_reflexion: + # RAG only + pass + elif not self._use_chain_of_thought and not self._use_rag_context and self._use_reflexion: + # Reflexion only + pass + else: + # None enabled (all False) + pass \ No newline at end of file diff --git a/tests/integration/test_violation_rate.py b/tests/integration/test_violation_rate.py index 657303c71..23a17e485 100644 --- a/tests/integration/test_violation_rate.py +++ b/tests/integration/test_violation_rate.py @@ -195,13 +195,13 @@ def example_with_fluent_service_call( # TODO: should be callable like this actually: - completion_result: TextGenerationCompletionResult = None - completion_result = (text_generation_completion_service - .without_guidelines() - .with_chain_of_thought_guidelines() - .with_rag_example_guidelines() - .with_reflexion_guidelines() - .invoke(user_prompt=prompt)) + completion_result: TextGenerationCompletionResult = ( + text_generation_completion_service + .without_guidelines() + .with_chain_of_thought_guidelines() + .with_rag_example_guidelines() + .with_reflexion_guidelines() + .invoke(user_prompt=prompt)) completion: GuidelinesProcessedCompletion = (generative_ai_security_guidelines_service