Custom Prompt Templates

from dataclasses import dataclass, field
from string import Template
from typing import Dict, List, Union, Any

@dataclass
class PromptTemplateManager:
    role_mapping: Dict[str, str] = field(
        default_factory=lambda: {
            "system": "system", 
            "user": "user", 
            "assistant": "assistant"
        }
    )
    templates: Dict[str, Union[Template, List[Dict[str, Any]]]] = field(
        init=False,
        default_factory=dict,
    )
    
    def __post_init__(self):
        """Load all templates from the templates directory."""
        self.templates_dir = os.path.join(
            os.path.dirname(__file__), "templates"
        )
        self._load_templates()
    
    def _load_templates(self):
        """Load all .py files from templates/ directory."""
        for filename in os.listdir(self.templates_dir):
            if filename.endswith(".py") and filename != "__init__.py":
                script_name = os.path.splitext(filename)[0]
                module = importlib.import_module(f"remem.prompts.templates.{script_name}")
                
                if not hasattr(module, "prompt_template"):
                    raise AttributeError(
                        f"Module '{script_name}' does not define 'prompt_template'"
                    )
                
                prompt_template = module.prompt_template
                
                # Convert to Template if string
                if isinstance(prompt_template, str):
                    self.templates[script_name] = Template(prompt_template)
                # Handle chat history format
                elif isinstance(prompt_template, list):
                    for item in prompt_template:
                        item["role"] = self.role_mapping.get(item["role"], item["role"])
                        item["content"] = Template(item["content"]) \
                            if isinstance(item["content"], str) else item["content"]
                    self.templates[script_name] = prompt_template
    
    def render(self, name: str, **kwargs) -> Union[str, List[Dict[str, Any]]]:
        """Render a template with variables."""
        template = self.get_template(name)
        
        if isinstance(template, Template):
            return template.substitute(**kwargs)
        elif isinstance(template, list):
            return [
                {"role": item["role"], "content": item["content"].substitute(**kwargs)}
                for item in template
            ]
    
    def get_template(self, name: str):
        """Get a template by name."""
        if name not in self.templates:
            raise KeyError(f"Template '{name}' not found.")
        return self.templates[name]

# src/remem/prompts/templates/my_template.py
from string import Template

prompt_template = [
    {
        "role": "system",
        "content": "You are an expert at extracting information."
    },
    {
        "role": "user",
        "content": "Extract entities from: $passage"
    }
]

# src/remem/prompts/templates/simple_template.py
from string import Template

prompt_template = Template("""
Analyze the following text:

$text

Provide a summary.
""")

ner_system = """Extract named entities from the given paragraph.
Respond with a JSON list containing all named entities."""

ner_user = """Extract entities from:

Return JSON: {"named_entities": [...]}"""

one_shot_example = {
    "passage": "Radio City started on 3 July 2001...",
    "output": '{"named_entities": ["Radio City", "India", "3 July 2001"]}'
}

prompt_template = [
    {"role": "system", "content": ner_system},
    {"role": "user", "content": one_shot_example["passage"]},
    {"role": "assistant", "content": one_shot_example["output"]},
    {"role": "user", "content": Template(ner_user)}
]

ner_conditioned_re_system = """Your task is to construct an RDF graph from the given passages and named entity lists.
Respond with a JSON list of triples, with each triple representing a relationship.

Requirements:
- Each triple should contain at least one, but preferably two, named entities
- Clearly resolve pronouns to their specific names
"""

ner_conditioned_re_frame = """Convert the paragraph into a JSON dict with triples.
Paragraph:

{named_entity_json}
"""

prompt_template = [
    {"role": "system", "content": ner_conditioned_re_system},
    {"role": "user", "content": one_shot_input},
    {"role": "assistant", "content": one_shot_output},
    {"role": "user", "content": Template(ner_conditioned_re_frame)}
]

triple_extraction_input_message = self.prompt_template_manager.render(
    name="triple_extraction",
    passage=passage,
    named_entity_json=named_entities_str
)

conversation_qa_system = """As an advanced reading comprehension assistant with multi-step reasoning capabilities, your task is to analyze retrieved information and corresponding questions meticulously.
You have access to information gathered through iterative agent reasoning steps. Your response contains two lines, starting with `Thought: ` and `Answer: `, respectively.
For `Thought: `, you will methodically break down the reasoning process, incorporating insights from the agent's multi-step exploration, illustrating how you arrive at conclusions.
Conclude with "Answer: " to present a concise, definitive response, devoid of additional elaborations.
Output `no information available` after "Answer: " if the answer is indeed not present."""

prompt_template = [
    {"role": "system", "content": conversation_qa_system},
    {"role": "user", "content": "${prompt_user}"}
]

# src/remem/prompts/templates/my_custom_extraction.py
from string import Template

system_prompt = """You are an expert at identifying key concepts and relationships.
Extract structured information from the provided text."""

user_prompt = """Analyze the following passage and extract:
1. Main topics
2. Key entities
3. Important relationships

Passage:
$passage

Return a JSON object with keys: topics, entities, relationships."""

prompt_template = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": Template(user_prompt)}
]

# src/remem/prompts/templates/few_shot_custom.py
from string import Template

system_prompt = """Extract cause-effect relationships from text."""

# Define examples
example_1_input = """Text: Heavy rain caused flooding in the city."""
example_1_output = """{
  "cause": "heavy rain",
  "effect": "flooding in the city"
}"""

example_2_input = """Text: The vaccine prevented the disease outbreak."""
example_2_output = """{
  "cause": "the vaccine",
  "effect": "prevented disease outbreak"
}"""

user_template = """Text: $passage"""

prompt_template = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": example_1_input},
    {"role": "assistant", "content": example_1_output},
    {"role": "user", "content": example_2_input},
    {"role": "assistant", "content": example_2_output},
    {"role": "user", "content": Template(user_template)}
]

# src/remem/prompts/templates/contextual_qa.py
from string import Template

system_prompt = """Answer questions based on provided context."""

user_prompt = """Context:
$context

Previous conversation:
$history

Question: $question

Provide a direct answer based on the context."""

prompt_template = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": Template(user_prompt)}
]

from remem.prompts import PromptTemplateManager

class MyExtractor:
    def __init__(self, llm_model):
        self.prompt_template_manager = PromptTemplateManager(
            role_mapping={"system": "system", "user": "user", "assistant": "assistant"}
        )
        self.llm_model = llm_model
    
    def extract(self, passage: str):
        # Render template
        messages = self.prompt_template_manager.render(
            name="my_custom_extraction",
            passage=passage
        )
        
        # Call LLM
        response, metadata, cache_hit = self.llm_model.infer(
            messages=messages,
            response_format={"type": "json_object"}
        )
        
        return response

# Render with multiple variables
qa_prompt = prompt_manager.render(
    name="contextual_qa",
    context="\n\n".join(retrieved_docs),
    history=conversation_history,
    question=user_question
)

response = llm.infer(messages=qa_prompt)

from remem.prompts import PromptTemplateManager

manager = PromptTemplateManager()
print(manager.list_template_names())
# ['ner', 'triple_extraction', 'rag_qa_unified', 'episodic_gist_extraction_locomo', ...]

manager.print_template("triple_extraction")
# Output:
# Template name: triple_extraction
# Role: system, Content: Your task is to construct an RDF graph...
# Role: user, Content: ...
# Role: assistant, Content: ...
# Role: user, Content: Template('...')

def _extract_chunk(self, chunk_key, passage, template="episodic_fact_extraction"):
    template_name = f"{template}_locomo"
    
    # Check for dataset-specific template
    selected = ["menatqa", "timeqa", "musique", "complex_tr", "2wikimultihopqa"]
    if any(self.global_config.dataset.startswith(prefix) for prefix in selected):
        template_name = f"{template}_wikipedia"
    else:
        for file_name in self.prompt_template_manager.templates:
            if file_name.startswith(f"{template}_"):
                if file_name.split("_")[-1] == self.global_config.dataset.split("_")[0]:
                    template_name = file_name
                    break
    
    extraction_input = self.prompt_template_manager.render(
        name=template_name,
        prompt_user=passage
    )

class AdaptiveExtractor:
    def __init__(self, llm_model, global_config):
        self.prompt_manager = PromptTemplateManager()
        self.global_config = global_config
    
    def select_template(self, task: str) -> str:
        """Select best template for task and dataset."""
        dataset = self.global_config.dataset.split("_")[0]
        
        # Try dataset-specific template
        specific_name = f"{task}_{dataset}"
        if self.prompt_manager.is_template_name_valid(specific_name):
            return specific_name
        
        # Fall back to unified template
        unified_name = f"{task}_unified"
        if self.prompt_manager.is_template_name_valid(unified_name):
            return unified_name
        
        # Default template
        return task
    
    def extract(self, passage: str):
        template_name = self.select_template("episodic_extraction")
        messages = self.prompt_manager.render(
            name=template_name,
            prompt_user=passage
        )
        # ... inference