beastbullet-experts / code /task_contracts.py

Complete package: 18 experts + full codebase

ae1e90a 11 months ago

7.1 kB

#!/usr/bin/env python3
"""
Task Contracts: Output Discipline for Execution Engine
Defines strict contracts for each task type to constrain LLM behavior
"""

from dataclasses import dataclass
from typing import Optional, Dict
from enum import Enum

class TaskType(Enum):
    """Supported task types with strict contracts"""
    GRAMMAR_FIX = "grammar_fix"
    MATH_CALC = "math_calc"
    CODE_GEN = "code_gen"
    EXPLANATION = "explanation"
    QUESTION_ANSWER = "qa"
    GENERAL = "general"

@dataclass
class TaskContract:
    """
    Execution contract for a task type
    
    This enforces:
    - Token budgets (prevent rambling)
    - Output style (prevent hallucination)
    - Creativity level (prevent story generation)
    - Format requirements (structured output)
    """
    task_type: TaskType
    max_tokens: int
    temperature: float  # 0.0 = deterministic, 1.0 = creative
    style: str  # "rewrite_only", "factual", "code_only", "concise"
    allow_explanation: bool
    output_format: str  # "single_line", "code_block", "paragraph"
    system_instruction: str
    
    def get_prompt_template(self, user_input: str, tool_result: Optional[str] = None) -> str:
        """Generate disciplined prompt based on contract"""
        
        if self.task_type == TaskType.GRAMMAR_FIX:
            return f"""Task: Grammar correction.
Rules:
- Output ONE corrected sentence only
- Do not add explanation
- Do not add commentary
- Preserve original meaning
Input: "{user_input}"
Corrected:"""

        elif self.task_type == TaskType.MATH_CALC:
            if tool_result:
                # Tool already gave answer - just format it
                return f"""The answer is: {tool_result}"""
            else:
                return f"""Task: Solve this math problem.
Rules:
- Show final answer only
- No explanation unless asked
- Use exact arithmetic
Problem: {user_input}
Answer:"""

        elif self.task_type == TaskType.CODE_GEN:
            return f"""Task: Write code.
Rules:
- Output code only
- Include brief comments
- No explanation outside code
- Follow best practices
Request: {user_input}
Code:
```python"""

        elif self.task_type == TaskType.EXPLANATION:
            return f"""Task: Explain concisely.
Rules:
- Maximum 3 sentences
- Focus on key points
- No unnecessary details
Question: {user_input}
Answer:"""

        elif self.task_type == TaskType.QUESTION_ANSWER:
            return f"""Task: Answer the question directly.
Rules:
- Direct answer first
- Brief explanation if needed
- No rambling
Question: {user_input}
Answer:"""

        else:  # GENERAL
            return user_input


# Task contract definitions
TASK_CONTRACTS = {
    TaskType.GRAMMAR_FIX: TaskContract(
        task_type=TaskType.GRAMMAR_FIX,
        max_tokens=48,
        temperature=0.0,
        style="rewrite_only",
        allow_explanation=False,
        output_format="single_line",
        system_instruction="You are a grammar correction tool. Output only the corrected sentence."
    ),
    
    TaskType.MATH_CALC: TaskContract(
        task_type=TaskType.MATH_CALC,
        max_tokens=32,
        temperature=0.0,
        style="factual",
        allow_explanation=False,
        output_format="single_line",
        system_instruction="You are a math solver. Output only the final answer."
    ),
    
    TaskType.CODE_GEN: TaskContract(
        task_type=TaskType.CODE_GEN,
        max_tokens=200,
        temperature=0.2,
        style="code_only",
        allow_explanation=False,
        output_format="code_block",
        system_instruction="You are a code generator. Output only clean, working code."
    ),
    
    TaskType.EXPLANATION: TaskContract(
        task_type=TaskType.EXPLANATION,
        max_tokens=128,
        temperature=0.3,
        style="concise",
        allow_explanation=True,
        output_format="paragraph",
        system_instruction="You are a concise explainer. Maximum 3 sentences."
    ),
    
    TaskType.QUESTION_ANSWER: TaskContract(
        task_type=TaskType.QUESTION_ANSWER,
        max_tokens=96,
        temperature=0.3,
        style="direct",
        allow_explanation=True,
        output_format="paragraph",
        system_instruction="You are a direct answerer. Answer first, explain briefly if needed."
    ),
    
    TaskType.GENERAL: TaskContract(
        task_type=TaskType.GENERAL,
        max_tokens=256,
        temperature=0.7,
        style="conversational",
        allow_explanation=True,
        output_format="paragraph",
        system_instruction="You are a helpful assistant."
    ),
}


def detect_task_type(query: str) -> TaskType:
    """Detect task type from query"""
    query_lower = query.lower()
    
    # Grammar fix
    if any(kw in query_lower for kw in ['fix this', 'correct this', 'grammar']):
        return TaskType.GRAMMAR_FIX
    
    # Math calculation
    if any(kw in query_lower for kw in ['calculate', 'compute', 'solve']) or \
       any(op in query for op in ['+', '-', '*', '/', '=']):
        return TaskType.MATH_CALC
    
    # Code generation
    if any(kw in query_lower for kw in ['write a', 'function', 'code', 'program', 'script']):
        return TaskType.CODE_GEN
    
    # Explanation
    if any(kw in query_lower for kw in ['explain', 'what is', 'how does']):
        return TaskType.EXPLANATION
    
    # Question answering
    if query.strip().endswith('?'):
        return TaskType.QUESTION_ANSWER
    
    return TaskType.GENERAL


def get_contract(task_type: TaskType) -> TaskContract:
    """Get contract for task type"""
    return TASK_CONTRACTS[task_type]


def should_skip_llm(task_type: TaskType, tool_result: Optional[str]) -> bool:
    """
    Determine if we should skip LLM and return tool result directly
    
    Hard-stop rule: If tool gives final answer, don't re-explain
    """
    if not tool_result:
        return False
    
    # Math: calculator gave answer, return it
    if task_type == TaskType.MATH_CALC:
        return True
    
    # Future: Add more hard-stop rules
    
    return False


def test_contracts():
    """Test task contracts"""
    print("="*70)
    print("🧪 Task Contract System Test")
    print("="*70)
    
    test_queries = [
        "Fix this: Me and him goes to store",
        "Calculate 123 * 456",
        "Write a Python function to reverse a string",
        "Explain what Python is",
        "What is 2 + 2?",
    ]
    
    for query in test_queries:
        task_type = detect_task_type(query)
        contract = get_contract(task_type)
        
        print(f"\n📝 Query: {query}")
        print(f"   Task Type: {task_type.value}")
        print(f"   Max Tokens: {contract.max_tokens}")
        print(f"   Temperature: {contract.temperature}")
        print(f"   Style: {contract.style}")
        print(f"   Format: {contract.output_format}")
        
        # Show prompt
        prompt = contract.get_prompt_template(query)
        print(f"\n   Prompt:")
        print("   " + "\n   ".join(prompt.split("\n")[:5]))
        print("   ...")


if __name__ == "__main__":
    test_contracts()

	#!/usr/bin/env python3
	"""
	Task Contracts: Output Discipline for Execution Engine
	Defines strict contracts for each task type to constrain LLM behavior
	"""

	from dataclasses import dataclass
	from typing import Optional, Dict
	from enum import Enum

	class TaskType(Enum):
	"""Supported task types with strict contracts"""
	GRAMMAR_FIX = "grammar_fix"
	MATH_CALC = "math_calc"
	CODE_GEN = "code_gen"
	EXPLANATION = "explanation"
	QUESTION_ANSWER = "qa"
	GENERAL = "general"

	@dataclass
	class TaskContract:
	"""
	Execution contract for a task type

	This enforces:
	- Token budgets (prevent rambling)
	- Output style (prevent hallucination)
	- Creativity level (prevent story generation)
	- Format requirements (structured output)
	"""
	task_type: TaskType
	max_tokens: int
	temperature: float # 0.0 = deterministic, 1.0 = creative
	style: str # "rewrite_only", "factual", "code_only", "concise"
	allow_explanation: bool
	output_format: str # "single_line", "code_block", "paragraph"
	system_instruction: str

	def get_prompt_template(self, user_input: str, tool_result: Optional[str] = None) -> str:
	"""Generate disciplined prompt based on contract"""

	if self.task_type == TaskType.GRAMMAR_FIX:
	return f"""Task: Grammar correction.
	Rules:
	- Output ONE corrected sentence only
	- Do not add explanation
	- Do not add commentary
	- Preserve original meaning

	Input: "{user_input}"
	Corrected:"""

	elif self.task_type == TaskType.MATH_CALC:
	if tool_result:
	# Tool already gave answer - just format it
	return f"""The answer is: {tool_result}"""
	else:
	return f"""Task: Solve this math problem.
	Rules:
	- Show final answer only
	- No explanation unless asked
	- Use exact arithmetic

	Problem: {user_input}
	Answer:"""

	elif self.task_type == TaskType.CODE_GEN:
	return f"""Task: Write code.
	Rules:
	- Output code only
	- Include brief comments
	- No explanation outside code
	- Follow best practices

	Request: {user_input}
	Code:
	```python"""

	elif self.task_type == TaskType.EXPLANATION:
	return f"""Task: Explain concisely.
	Rules:
	- Maximum 3 sentences
	- Focus on key points
	- No unnecessary details

	Question: {user_input}
	Answer:"""

	elif self.task_type == TaskType.QUESTION_ANSWER:
	return f"""Task: Answer the question directly.
	Rules:
	- Direct answer first
	- Brief explanation if needed
	- No rambling

	Question: {user_input}
	Answer:"""

	else: # GENERAL
	return user_input


	# Task contract definitions
	TASK_CONTRACTS = {
	TaskType.GRAMMAR_FIX: TaskContract(
	task_type=TaskType.GRAMMAR_FIX,
	max_tokens=48,
	temperature=0.0,
	style="rewrite_only",
	allow_explanation=False,
	output_format="single_line",
	system_instruction="You are a grammar correction tool. Output only the corrected sentence."
	),

	TaskType.MATH_CALC: TaskContract(
	task_type=TaskType.MATH_CALC,
	max_tokens=32,
	temperature=0.0,
	style="factual",
	allow_explanation=False,
	output_format="single_line",
	system_instruction="You are a math solver. Output only the final answer."
	),

	TaskType.CODE_GEN: TaskContract(
	task_type=TaskType.CODE_GEN,
	max_tokens=200,
	temperature=0.2,
	style="code_only",
	allow_explanation=False,
	output_format="code_block",
	system_instruction="You are a code generator. Output only clean, working code."
	),

	TaskType.EXPLANATION: TaskContract(
	task_type=TaskType.EXPLANATION,
	max_tokens=128,
	temperature=0.3,
	style="concise",
	allow_explanation=True,
	output_format="paragraph",
	system_instruction="You are a concise explainer. Maximum 3 sentences."
	),

	TaskType.QUESTION_ANSWER: TaskContract(
	task_type=TaskType.QUESTION_ANSWER,
	max_tokens=96,
	temperature=0.3,
	style="direct",
	allow_explanation=True,
	output_format="paragraph",
	system_instruction="You are a direct answerer. Answer first, explain briefly if needed."
	),

	TaskType.GENERAL: TaskContract(
	task_type=TaskType.GENERAL,
	max_tokens=256,
	temperature=0.7,
	style="conversational",
	allow_explanation=True,
	output_format="paragraph",
	system_instruction="You are a helpful assistant."
	),
	}


	def detect_task_type(query: str) -> TaskType:
	"""Detect task type from query"""
	query_lower = query.lower()

	# Grammar fix
	if any(kw in query_lower for kw in ['fix this', 'correct this', 'grammar']):
	return TaskType.GRAMMAR_FIX

	# Math calculation
	if any(kw in query_lower for kw in ['calculate', 'compute', 'solve']) or \
	any(op in query for op in ['+', '-', '*', '/', '=']):
	return TaskType.MATH_CALC

	# Code generation
	if any(kw in query_lower for kw in ['write a', 'function', 'code', 'program', 'script']):
	return TaskType.CODE_GEN

	# Explanation
	if any(kw in query_lower for kw in ['explain', 'what is', 'how does']):
	return TaskType.EXPLANATION

	# Question answering
	if query.strip().endswith('?'):
	return TaskType.QUESTION_ANSWER

	return TaskType.GENERAL


	def get_contract(task_type: TaskType) -> TaskContract:
	"""Get contract for task type"""
	return TASK_CONTRACTS[task_type]


	def should_skip_llm(task_type: TaskType, tool_result: Optional[str]) -> bool:
	"""
	Determine if we should skip LLM and return tool result directly

	Hard-stop rule: If tool gives final answer, don't re-explain
	"""
	if not tool_result:
	return False

	# Math: calculator gave answer, return it
	if task_type == TaskType.MATH_CALC:
	return True

	# Future: Add more hard-stop rules

	return False


	def test_contracts():
	"""Test task contracts"""
	print("="*70)
	print("🧪 Task Contract System Test")
	print("="*70)

	test_queries = [
	"Fix this: Me and him goes to store",
	"Calculate 123 * 456",
	"Write a Python function to reverse a string",
	"Explain what Python is",
	"What is 2 + 2?",
	]

	for query in test_queries:
	task_type = detect_task_type(query)
	contract = get_contract(task_type)

	print(f"\n📝 Query: {query}")
	print(f" Task Type: {task_type.value}")
	print(f" Max Tokens: {contract.max_tokens}")
	print(f" Temperature: {contract.temperature}")
	print(f" Style: {contract.style}")
	print(f" Format: {contract.output_format}")

	# Show prompt
	prompt = contract.get_prompt_template(query)
	print(f"\n Prompt:")
	print(" " + "\n ".join(prompt.split("\n")[:5]))
	print(" ...")


	if __name__ == "__main__":
	test_contracts()