CalvinBackup/llmsecops-research
1
0
Fork 0
mirror of https://github.com/lightbroker/llmsecops-research.git synced 2026-07-12 15:46:37 +02:00
Code Issues Packages Projects Releases Wiki Activity

more test and template setup

Browse Source
This commit is contained in:
Adam Wilson
2025-07-25 09:45:03 -06:00
parent 3a62ecfae8
commit d15e9d6794
13 changed files with 112 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/text_generation/common/constants.py
+5 -5
View File
@@ -8,8 +8,8 @@ class Constants:
ALT_COMPLETION_TEXT = "I can't answer that."
class PromptTemplateIds:
PHI_3_MINI_4K_INSTRUCT_BASIC = "phi-3-mini-4k-instruct-basic"
PHI_3_MINI_4K_INSTRUCT_ZERO_SHOT_CHAIN_OF_THOUGHT = "phi-3-mini-4k-instruct-zero-shot-cot"
PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_EXAMPLES = "phi-3-mini-4k-instruct-few-shot"
PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_RAG_PLUS_COT = "phi-3-mini-4k-instruct-few-shot-rag-plus-cot"
PHI_3_MINI_4K_INSTRUCT_REFLEXION = "phi-3-mini-4k-instruct-reflexion"
PHI_3_MINI_4K_INSTRUCT__01_BASIC = "phi-3-mini-4k-instruct-basic"
PHI_3_MINI_4K_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT = "phi-3-mini-4k-instruct-zero-shot-cot"
PHI_3_MINI_4K_INSTRUCT__03_FEW_SHOT_EXAMPLES = "phi-3-mini-4k-instruct-few-shot"
PHI_3_MINI_4K_INSTRUCT__04_FEW_SHOT_RAG_PLUS_COT = "phi-3-mini-4k-instruct-few-shot-rag-plus-cot"
PHI_3_MINI_4K_INSTRUCT__05_REFLEXION = "phi-3-mini-4k-instruct-reflexion"
src/text_generation/services/guidelines/base_security_guidelines_service.py
+6 -2
View File
@@ -28,6 +28,10 @@ class BaseSecurityGuidelinesService(AbstractSecurityGuidelinesService):
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
@@ -36,7 +40,7 @@ class BaseSecurityGuidelinesService(AbstractSecurityGuidelinesService):
| self.response_processing_service.process_text_generation_output
)
def _get_template(self) -> StringPromptTemplate:
def _get_template(self, user_prompt: str) -> StringPromptTemplate:
"""
Get the prompt template for security guidelines.
@@ -50,7 +54,7 @@ class BaseSecurityGuidelinesService(AbstractSecurityGuidelinesService):
raise ValueError(f"Parameter 'user_prompt' cannot be empty or None")
try:
prompt_template = self._get_template()
prompt_template = self._get_template(user_prompt=user_prompt)
chain = self._create_chain(prompt_template)
result = GuidelinesResult(
completion_text=chain.invoke(user_prompt),
src/text_generation/services/guidelines/chain_of_thought_security_guidelines_service.py
+1 -1
View File
@@ -31,5 +31,5 @@ class ChainOfThoughtSecurityGuidelinesService(BaseSecurityGuidelinesService):
StringPromptTemplate: Template configured for CoT processing
"""
return self.prompt_template_service.get(
id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_ZERO_SHOT_CHAIN_OF_THOUGHT
id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT
)
src/text_generation/services/guidelines/rag_context_security_guidelines_configuration_builder.py
+1 -1
View File
@@ -24,7 +24,7 @@ class RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder(
self.embedding_model: EmbeddingModel = embedding_model
self.prompt_template_service = prompt_template_service
self.prompt_injection_example_repository = prompt_injection_example_repository
self.prompt_template_id = self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_EXAMPLES
self.prompt_template_id = self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__03_FEW_SHOT_EXAMPLES
self.vectorstore = self._setup_vectorstore()
def _setup_vectorstore(self):
src/text_generation/services/guidelines/rag_context_security_guidelines_service.py
+1 -1
View File
@@ -29,7 +29,7 @@ class RagContextSecurityGuidelinesService(BaseSecurityGuidelinesService):
Returns:
StringPromptTemplate: Template configured for RAG processing
"""
template_id = self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_EXAMPLES
template_id = self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__03_FEW_SHOT_EXAMPLES
return self.config_builder.get_prompt_template(
template_id=template_id,
user_prompt=user_prompt
src/text_generation/services/guidelines/rag_plus_cot_security_guidelines_service.py
+2 -2
View File
@@ -33,7 +33,7 @@ class RagPlusCotSecurityGuidelinesService(BaseSecurityGuidelinesService):
StringPromptTemplate: Template configured for RAG processing
"""
return self.prompt_template_service.get(
id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_RAG_PLUS_COT
id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__04_FEW_SHOT_RAG_PLUS_COT
)
def _get_template_id(self) -> str:
@@ -43,4 +43,4 @@ class RagPlusCotSecurityGuidelinesService(BaseSecurityGuidelinesService):
Returns:
str: Template ID for RAG + CoT security guidelines
"""
return self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_RAG_PLUS_COT
return self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__04_FEW_SHOT_RAG_PLUS_COT
src/text_generation/services/nlp/text_generation_completion_service.py
+2 -2
View File
@@ -211,11 +211,11 @@ class TextGenerationCompletionService(
def _create_chain_without_guidelines(self):
prompt_template = self.prompt_template_service.get(
id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_BASIC
id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__01_BASIC
)
if prompt_template is None:
raise ValueError(f"Prompt template not found for ID: {self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_BASIC}")
raise ValueError(f"Prompt template not found for ID: {self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__01_BASIC}")
return (
{ "question": RunnablePassthrough() }