Update llm/llm.py

llm/llm.py

@@ -1,3 +1,9 @@
+# {'prompts': [
+#    {'id': 1, 'prompt_template': 'Question: {question}\nAnswer: \n', 'description': 'simple question without a prompt', 'rate': 1}, 
+#    {'id': 2, 'prompt_template': "Question: {question}\nAnswer: Write a concise answer on the question with 
+#               one example if it's possible. CONCISE ANSWER.\n", 'description': 'simple concise prompt', 'rate': 3}]}
+
+
 import os
 
 import yaml
@@ -29,9 +35,8 @@ class LLM_chain:
     def __call__(self, entity: str, id: int = 0):
         try:
             data = self.__read_yaml()
-            print(data)
-            prompts = data["prompts"]
-            template = prompts["prompt_template"][1]
+            prompts = data["prompts"][id] # get second prompt from yaml, need change id parameter to get other prompt
+            template = prompts["prompt_template"]
             prompt = PromptTemplate(template=template, input_variables=["entity"])
             llm_chain = LLMChain(prompt=prompt, llm=self.llm, verbose=True)
             output = llm_chain.invoke(entity)