mirror of
https://github.com/codelion/optillm.git
synced 2025-05-28 09:39:38 +03:00
78 lines
2.6 KiB
Python
78 lines
2.6 KiB
Python
"""
|
|
Utility functions for the System Prompt Learning (SPL) plugin.
|
|
"""
|
|
|
|
import re
|
|
import uuid
|
|
import logging
|
|
from typing import Tuple, Optional, List, Dict, Any
|
|
|
|
from optillm.plugins.spl.prompts import STRATEGY_APPLICATION_PROMPT
|
|
|
|
# Setup logging
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def extract_thinking(response: str) -> Tuple[str, Optional[str]]:
|
|
"""
|
|
Extract thinking content from <think>...</think> tags and the response after.
|
|
|
|
Args:
|
|
response: The model's response
|
|
|
|
Returns:
|
|
Tuple[str, Optional[str]]: The cleaned response and the thinking content (if any)
|
|
"""
|
|
thinking_content = None
|
|
final_response = response
|
|
|
|
# Check if there are thinking tags
|
|
think_pattern = r'<think>(.*?)</think>'
|
|
think_matches = re.findall(think_pattern, response, re.DOTALL)
|
|
|
|
if think_matches:
|
|
# Extract thinking content (concatenate if multiple blocks)
|
|
thinking_content = "\n".join(think_matches)
|
|
|
|
# Extract the response part (everything after the last </think> tag)
|
|
final_parts = response.split('</think>')
|
|
if len(final_parts) > 1:
|
|
final_response = final_parts[-1].strip()
|
|
|
|
return final_response, thinking_content
|
|
|
|
def augment_system_prompt(system_prompt: str, strategies: List[Any]) -> str:
|
|
"""
|
|
Augment the system prompt with selected strategies and reasoning examples.
|
|
Instructs the LLM to apply the strategies in its solution.
|
|
|
|
Args:
|
|
system_prompt: The original system prompt
|
|
strategies: A list of strategies to add to the prompt
|
|
|
|
Returns:
|
|
str: The augmented system prompt
|
|
"""
|
|
if not strategies:
|
|
return system_prompt
|
|
|
|
# Build the individual strategy sections
|
|
strategies_section = ""
|
|
|
|
for i, strategy in enumerate(strategies, 1):
|
|
strategies_section += f"Strategy {i} for {strategy.problem_type} problems:\n{strategy.strategy_text}\n\n"
|
|
|
|
# Add a sample reasoning example if available
|
|
if strategy.reasoning_examples:
|
|
# Use the most recent reasoning example (last one)
|
|
reasoning = strategy.reasoning_examples[-1]
|
|
if reasoning:
|
|
strategies_section += f"Example reasoning process:\n<think>\n{reasoning}\n</think>\n\n"
|
|
|
|
# Format the application prompt with the strategies section
|
|
strategy_prompt = STRATEGY_APPLICATION_PROMPT.format(strategies_section=strategies_section)
|
|
|
|
# Append the formatted strategy section to the system prompt
|
|
augmented_prompt = system_prompt + "\n\n" + strategy_prompt
|
|
|
|
return augmented_prompt
|