91 lines
2.5 KiB
Python
91 lines
2.5 KiB
Python
"""Prompt loader for AI agents.
|
|
|
|
Loads prompts from agent spec files or YAML configuration.
|
|
"""
|
|
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
|
|
import yaml
|
|
|
|
AGENT_SPECS_DIR = Path("docs/specs/agents")
|
|
AGENT_PROMPTS_YAML = Path("docs/agent_prompts.yaml")
|
|
|
|
|
|
def get_agent_roles() -> list[str]:
|
|
"""Get all available agent roles."""
|
|
roles = []
|
|
if AGENT_PROMPTS_YAML.exists():
|
|
data = yaml.safe_load(AGENT_PROMPTS_YAML.read_text(encoding="utf-8"))
|
|
return list(data.keys()) if data else []
|
|
|
|
if AGENT_SPECS_DIR.exists():
|
|
for f in sorted(AGENT_SPECS_DIR.glob("*.md")):
|
|
roles.append(f.stem)
|
|
|
|
return roles
|
|
|
|
|
|
def load_prompt_from_yaml(role: str) -> Optional[dict[str, Any]]:
|
|
"""Load prompt config from YAML file."""
|
|
if not AGENT_PROMPTS_YAML.exists():
|
|
return None
|
|
|
|
data = yaml.safe_load(AGENT_PROMPTS_YAML.read_text(encoding="utf-8"))
|
|
return data.get(role) if data else None
|
|
|
|
|
|
def load_prompt_from_spec(role: str) -> Optional[dict[str, Any]]:
|
|
"""Load prompt from agent spec markdown file."""
|
|
spec_path = AGENT_SPECS_DIR / f"{role}.md"
|
|
if not spec_path.exists():
|
|
return None
|
|
|
|
content = spec_path.read_text(encoding="utf-8")
|
|
|
|
def extract_field(label: str) -> Optional[str]:
|
|
match = re.search(rf"\*\*{label}:\*\*\s*(.+)", content)
|
|
return match.group(1).strip() if match else None
|
|
|
|
def extract_prompt() -> Optional[str]:
|
|
match = re.search(
|
|
r"## Prompt Template\n+```\n(.+?)\n```",
|
|
content,
|
|
re.DOTALL,
|
|
)
|
|
return match.group(1).strip() if match else None
|
|
|
|
provider_str = extract_field("Провайдер") or extract_field("Provider") or "yandex_gpt"
|
|
provider = "gigachat" if "gigachat" in (provider_str or "").lower() else "yandex_gpt"
|
|
|
|
prompt_text = extract_prompt()
|
|
if not prompt_text:
|
|
return None
|
|
|
|
return {
|
|
"system_prompt": prompt_text,
|
|
"provider": provider,
|
|
"temperature": 0.7,
|
|
"max_tokens": 2000,
|
|
}
|
|
|
|
|
|
def get_prompt_config(role: str) -> Optional[dict[str, Any]]:
|
|
"""Get prompt configuration for a role.
|
|
|
|
Tries YAML first, falls back to spec markdown.
|
|
|
|
Args:
|
|
role: Agent role name (e.g. "coordinator", "business_analyst")
|
|
|
|
Returns:
|
|
Dict with keys: system_prompt, provider, temperature, max_tokens
|
|
or None if not found.
|
|
"""
|
|
config = load_prompt_from_yaml(role)
|
|
if config:
|
|
return config
|
|
|
|
return load_prompt_from_spec(role)
|