Fix test pathing

pytest.ini

@@ -1 +1,2 @@
-pythonpath=.,src
+[pytest]
+pythonpath=src

src/app.py

@@ -13,7 +13,7 @@ def init_llm():
     return llm
 
 
-def get_response(prompt):
+def function_agent(prompt):
     try:
         return llm.generate_function_call(
             prompt, config.chat_template, config.num_fewshot, config.max_depth
@@ -22,11 +22,12 @@ def get_response(prompt):
         return f"An error occurred: {str(e)}"
 
 
-def get_output(context, user_input):
+def output_agent(context, user_input):
+    """Takes the output of the RAG and generates a final response."""
     try:
         config.status.update(label=":bulb: Preparing answer..")
         script_dir = os.path.dirname(os.path.abspath(__file__))
-        prompt_path = os.path.join(script_dir, 'prompt_assets', 'output_sys_prompt.yml')
+        prompt_path = os.path.join(script_dir, "prompt_assets", "output_sys_prompt.yml")
         prompt_schema = llm.prompter.read_yaml_file(prompt_path)
         sys_prompt = (
             llm.prompter.format_yaml_prompt(prompt_schema, dict())
@@ -41,6 +42,30 @@ def get_output(context, user_input):
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
+def query_agent(prompt):
+    """Modifies the prompt and runs inference on it."""
+    try:
+        config.status.update(label=":brain: Starting inference..")
+        script_dir = os.path.dirname(os.path.abspath(__file__))
+        prompt_path = os.path.join(script_dir, "prompt_assets", "output_sys_prompt.yml")
+        prompt_schema = llm.prompter.read_yaml_file(prompt_path)
+        sys_prompt = llm.prompter.format_yaml_prompt(prompt_schema, dict())
+        convo = [
+            {"role": "system", "content": sys_prompt},
+            {"role": "user", "content": prompt},
+        ]
+        response = llm.run_inference(convo)
+        return response
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
+
+
+def get_response(input_text: str):
+    """This is the main function that generates the final response."""
+    agent_resp = function_agent(input_text)
+    output = output_agent(agent_resp, input_text)
+    return output
+
 
 def main():
     st.title("LLM-ADE 9B Demo")
@@ -51,26 +76,25 @@ def main():
         if input_text:
             with st.status("Generating response...") as status:
                 config.status = status
-                agent_resp = get_response(input_text)
-                st.write(get_output(agent_resp, input_text))
+                st.write(get_response(input_text))
                 config.status.update(label="Finished!", state="complete", expanded=True)
         else:
             st.warning("Please enter some text to generate a response.")
 
 
-llm = init_llm()
-
-
 def main_headless(prompt: str):
     start = time()
-    agent_resp = get_response(prompt)
-    print("\033[94m" + get_output(agent_resp, prompt) + "\033[0m")
+    print("\033[94m" + get_response(prompt) + "\033[0m")
     print(f"Time taken: {time() - start:.2f}s\n" + "-" * 20)
 
 
+llm = init_llm()
+
+
 if __name__ == "__main__":
     if config.headless:
         import fire
+
         fire.Fire(main_headless)
     else:
         main()

src/functioncall.py

@@ -39,8 +39,7 @@ class ModelInference:
 
     def process_completion_and_validate(self, completion, chat_template):
         if completion:
-            # completion = f"<tool_call>\n{completion}\n</tool_call>"
-            breakpoint()
+            # completion = f"<tool_call>\n{completion}\n</tool_call>"]
             validation, tool_calls, error_message = validate_and_extract_tool_calls(completion)
 
             if validation:
@@ -85,6 +84,7 @@ class ModelInference:
 
             def recursive_loop(prompt, completion, depth):
                 nonlocal max_depth
+                breakpoint()
                 tool_calls, assistant_message, error_message = self.process_completion_and_validate(completion, chat_template)
                 prompt.append({"role": "assistant", "content": assistant_message})
 

tests/test_quality.py

@@ -1,10 +1,11 @@
 import sys
 import json
 from io import StringIO
+from app import get_response
 
 def test_quality():
     """Tests if the expected functions and values are used"""
-    with open("qa_questions.json") as f:
+    with open("tests/qa_questions.json") as f:
         qs = json.load(f)
     
     for q in qs:
@@ -12,11 +13,13 @@ def test_quality():
         stdout = StringIO()
         sys.stdout = stdout
         
-        for include in q['includes']:
+        get_response(q["question"])
+        
+        for include in q["expecteds"]["includes"]:
             assert include in stdout.getvalue(), f"Expected {include} in output"
-        for exclude in q['excludes']:
+        for exclude in q["expecteds"]["excludes"]:
             assert exclude not in stdout.getvalue(), f"Expected {exclude} not in output"
-        for function in q['functions']:
+        for function in q["expecteds"]["functions"]:
             assert f"Invoking function call {function}" in stdout.getvalue(), f"{function} was not invoked"
         
         # Restore stdout