prompt.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. import os
  2. from jinja2 import Template
  3. from openhands.utils.microagent import MicroAgent
  4. class PromptManager:
  5. """
  6. Manages prompt templates and micro-agents for AI interactions.
  7. This class handles loading and rendering of system and user prompt templates,
  8. as well as loading micro-agent specifications. It provides methods to access
  9. rendered system and initial user messages for AI interactions.
  10. Attributes:
  11. prompt_dir (str): Directory containing prompt templates.
  12. agent_skills_docs (str): Documentation of agent skills.
  13. micro_agent (MicroAgent | None): Micro-agent, if specified.
  14. """
  15. def __init__(
  16. self,
  17. prompt_dir: str,
  18. agent_skills_docs: str,
  19. micro_agent: MicroAgent | None = None,
  20. ):
  21. self.prompt_dir: str = prompt_dir
  22. self.agent_skills_docs: str = agent_skills_docs
  23. self.system_template: Template = self._load_template('system_prompt')
  24. self.user_template: Template = self._load_template('user_prompt')
  25. self.micro_agent: MicroAgent | None = micro_agent
  26. def _load_template(self, template_name: str) -> Template:
  27. template_path = os.path.join(self.prompt_dir, f'{template_name}.j2')
  28. if not os.path.exists(template_path):
  29. raise FileNotFoundError(f'Prompt file {template_path} not found')
  30. with open(template_path, 'r') as file:
  31. return Template(file.read())
  32. @property
  33. def system_message(self) -> str:
  34. rendered = self.system_template.render(
  35. agent_skills_docs=self.agent_skills_docs,
  36. ).strip()
  37. return rendered
  38. @property
  39. def initial_user_message(self) -> str:
  40. """This is the initial user message provided to the agent
  41. before *actual* user instructions are provided.
  42. It is used to provide a demonstration of how the agent
  43. should behave in order to solve the user's task. And it may
  44. optionally contain some additional context about the user's task.
  45. These additional context will convert the current generic agent
  46. into a more specialized agent that is tailored to the user's task.
  47. """
  48. rendered = self.user_template.render(
  49. micro_agent=self.micro_agent.content if self.micro_agent else None
  50. )
  51. return rendered.strip()