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