99 lines
3.6 KiB
Python
99 lines
3.6 KiB
Python
import dspy
|
|
from .utils import OPENROUTER_API_BASE, OPENROUTER_API_KEY
|
|
from .modules import signature_generator
|
|
from modaic import PrecompiledAgent, PrecompiledConfig
|
|
|
|
class PromptToSignatureConfig(PrecompiledConfig):
|
|
lm: str = "openrouter/anthropic/claude-haiku-4.5"
|
|
refine_lm: str = "openai/gpt-4o"
|
|
max_tokens: int = 16000
|
|
temperature: float = 1.0
|
|
max_attempts_to_refine: int = 5
|
|
|
|
|
|
class PromptToSignatureAgent(PrecompiledAgent):
|
|
config: PromptToSignatureConfig
|
|
|
|
def __init__(self, config: PromptToSignatureConfig, **kwargs):
|
|
super().__init__(config, **kwargs)
|
|
self.signature_generator = signature_generator
|
|
self.signature_refiner = dspy.Refine(
|
|
module=self.signature_generator,
|
|
N=config.max_attempts_to_refine,
|
|
reward_fn=PromptToSignatureAgent.validate_signature_with_feedback,
|
|
threshold=1.0,
|
|
)
|
|
|
|
lm = dspy.LM(
|
|
model=config.lm,
|
|
max_tokens=config.max_tokens,
|
|
temperature=config.temperature,
|
|
api_base=OPENROUTER_API_BASE,
|
|
api_key=OPENROUTER_API_KEY
|
|
)
|
|
refine_lm = dspy.LM(
|
|
model=config.refine_lm,
|
|
max_tokens=config.max_tokens,
|
|
temperature=config.temperature,
|
|
)
|
|
|
|
self.signature_generator.set_lm(lm)
|
|
self.signature_refiner.set_lm(refine_lm)
|
|
|
|
def forward(self, prompt: str, refine: bool = False) -> dspy.Prediction:
|
|
if not prompt:
|
|
raise ValueError("Prompt is required!!")
|
|
|
|
if refine:
|
|
try:
|
|
result = self.signature_refiner(prompt=prompt)
|
|
except Exception as e:
|
|
print(f"Refinement failed: {e}")
|
|
print("💡 Try adjusting your prompt or increasing max attempts")
|
|
return None
|
|
else:
|
|
result = self.signature_generator(prompt)
|
|
|
|
return result
|
|
|
|
def generate_code(self, prediction: dspy.Prediction) -> str:
|
|
return self.signature_generator.generate_code(prediction)
|
|
|
|
@staticmethod # attached metric for refinement
|
|
def validate_signature_with_feedback(args, pred):
|
|
"""Validation function for dspy.Refine that asks user for feedback"""
|
|
|
|
# display the generated signature
|
|
print("\n" + "=" * 60)
|
|
print("🔍 Review Generated Signature")
|
|
print("=" * 60)
|
|
|
|
# show the signature name and description
|
|
print(f"Signature Name: {pred.signature_name}")
|
|
print(f"Description: {pred.task_description}")
|
|
|
|
# show the fields in a simple format
|
|
print(f"\nFields ({len(pred.signature_fields)}):")
|
|
for i, field in enumerate(pred.signature_fields, 1):
|
|
role_emoji = "📥" if field.role.value == "input" else "📤"
|
|
print(
|
|
f" {i}. {role_emoji} {field.name} ({field.type.value}) - {field.description}"
|
|
)
|
|
|
|
# ask for user approval (in an app, this would be a state variable)
|
|
is_satisfied = input("Are you satisfied with this signature? (y/n): ")
|
|
is_satisfied = is_satisfied.lower() == "y"
|
|
|
|
if is_satisfied:
|
|
print("✓ Signature approved!")
|
|
return 1.0
|
|
else:
|
|
# ask for feedback (in an app, this would be a state variable)
|
|
feedback = input("Please provide feedback for improvement: ")
|
|
if not feedback:
|
|
raise ValueError(
|
|
"Feedback is required if you are not satisfied with the signature!"
|
|
)
|
|
|
|
print(f"📝 Feedback recorded: {feedback}")
|
|
return dspy.Prediction(score=0.0, feedback=feedback) |