From 03a76207f572406d8a9e4c01c8a9e1a28b1b7334 Mon Sep 17 00:00:00 2001 From: Dmytro Nikolaiev Date: Thu, 19 Sep 2024 10:59:24 -0400 Subject: [PATCH] Add LLMFunctionWithPrompt to __init__ --- council/llm/__init__.py | 1 + tests/integration/llm/test_llm_function_with_prompt.py | 2 +- tests/unit/llm/test_llm_response_parser_yaml_json.py | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/council/llm/__init__.py b/council/llm/__init__.py index a8b1e56e..c19d55f8 100644 --- a/council/llm/__init__.py +++ b/council/llm/__init__.py @@ -20,6 +20,7 @@ ExecuteLLMRequest, ) from .llm_function import LLMFunction, LLMFunctionError, FunctionOutOfRetryError +from .llm_function_with_prompt import LLMFunctionWithPrompt from .monitored_llm import MonitoredLLM from .chat_gpt_configuration import ChatGPTConfigurationBase diff --git a/tests/integration/llm/test_llm_function_with_prompt.py b/tests/integration/llm/test_llm_function_with_prompt.py index e19118b2..827c2546 100644 --- a/tests/integration/llm/test_llm_function_with_prompt.py +++ b/tests/integration/llm/test_llm_function_with_prompt.py @@ -3,7 +3,7 @@ import dotenv from council import AzureLLM -from council.llm.llm_function_with_prompt import LLMFunctionWithPrompt +from council.llm import LLMFunctionWithPrompt from council.prompt import LLMPromptConfigObject from tests import get_data_filename from tests.integration.llm.test_llm_function import SQLResult diff --git a/tests/unit/llm/test_llm_response_parser_yaml_json.py b/tests/unit/llm/test_llm_response_parser_yaml_json.py index 155c4289..94dd207b 100644 --- a/tests/unit/llm/test_llm_response_parser_yaml_json.py +++ b/tests/unit/llm/test_llm_response_parser_yaml_json.py @@ -152,7 +152,7 @@ def test_no_block(self): assert str(e.exception).strip().endswith("json block is not found") - def test_incorrect_yaml(self): + def test_incorrect_json(self): llm = MockLLM.from_response( """ ```json