[FIX] faulty imports when using PydanticProgramMode.LM_FORMAT_ENFORCER (

run-llama#11355)
uptrain-ai · Feb 28, 2024 · 4dc8ac7 · 4dc8ac7
1 parent d38ce1a
commit 4dc8ac7
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/llama-index-core/llama_index/core/llms/utils.py b/llama-index-core/llama_index/core/llms/utils.py
@@ -64,7 +64,7 @@ def resolve_llm(
                 "llm must start with str 'local' or of type LLM or BaseLanguageModel"
             )
         try:
-            from llama_index.llms.llama.utils import (
+            from llama_index.llms.llama_cpp.llama_utils import (
                 completion_to_prompt,
                 messages_to_prompt,
             )  # pants: no-infer-dep

diff --git a/llama-index-core/llama_index/core/program/utils.py b/llama-index-core/llama_index/core/program/utils.py
@@ -90,9 +90,15 @@ def get_program_for_llm(
             **kwargs,
         )
     elif pydantic_program_mode == PydanticProgramMode.LM_FORMAT_ENFORCER:
-        from llama_index.core.program.lmformatenforcer_program import (
-            LMFormatEnforcerPydanticProgram,
-        )
+        try:
+            from llama_index.program.lmformatenforcer import (
+                LMFormatEnforcerPydanticProgram,
+            )  # pants: no-infer-dep
+        except ImportError:
+            raise ImportError(
+                "This mode requires the `llama-index-program-lmformatenforcer package. Please"
+                " install it by running `pip install llama-index-program-lmformatenforcer`."
+            )
 
         return LMFormatEnforcerPydanticProgram.from_defaults(
             output_cls=output_cls,