File size: 2,436 Bytes
c124df1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import datetime
from pydantic import BaseModel
from typing import Dict
from schema import FunctionCall
from utils import (
get_fewshot_examples
)
import yaml
import json
import os
class PromptSchema(BaseModel):
Role: str
Objective: str
Tools: str
Examples: str
Schema: str
Instructions: str
class PromptManager:
def __init__(self):
self.script_dir = os.path.dirname(os.path.abspath(__file__))
def format_yaml_prompt(self, prompt_schema: PromptSchema, variables: Dict) -> str:
formatted_prompt = ""
for field, value in prompt_schema.dict().items():
if field == "Examples" and variables.get("examples") is None:
continue
formatted_value = value.format(**variables)
if field == "Instructions":
formatted_prompt += f"{formatted_value}"
else:
formatted_value = formatted_value.replace("\n", " ")
formatted_prompt += f"{formatted_value}"
return formatted_prompt
def read_yaml_file(self, file_path: str) -> PromptSchema:
with open(file_path, 'r') as file:
yaml_content = yaml.safe_load(file)
prompt_schema = PromptSchema(
Role=yaml_content.get('Role', ''),
Objective=yaml_content.get('Objective', ''),
Tools=yaml_content.get('Tools', ''),
Examples=yaml_content.get('Examples', ''),
Schema=yaml_content.get('Schema', ''),
Instructions=yaml_content.get('Instructions', ''),
)
return prompt_schema
def generate_prompt(self, user_prompt, tools, num_fewshot=None):
prompt_path = os.path.join(self.script_dir, 'prompt_assets', 'sys_prompt.yml')
prompt_schema = self.read_yaml_file(prompt_path)
if num_fewshot is not None:
examples = get_fewshot_examples(num_fewshot)
else:
examples = None
schema_json = json.loads(FunctionCall.schema_json())
variables = {
"date": datetime.date.today(),
"tools": tools,
"examples": examples,
"schema": schema_json
}
sys_prompt = self.format_yaml_prompt(prompt_schema, variables)
prompt = [
{'content': sys_prompt, 'role': 'system'}
]
prompt.extend(user_prompt)
return prompt
|