More prompt format

private_gpt/components/llm/llm_component.py

@@ -4,7 +4,6 @@
 from llama_index.llms import MockLLM
 from llama_index.llms.base import LLM
 
-from private_gpt.components.llm.prompt_helper import get_prompt_style
 from private_gpt.paths import models_path
 from private_gpt.settings.settings import Settings
 
@@ -23,9 +22,14 @@ def __init__(self, settings: Settings) -> None:
  case "local":
  from llama_index.llms import LlamaCPP
 
- prompt_style_cls = get_prompt_style(settings.local.prompt_style)
- prompt_style = prompt_style_cls(
- default_system_prompt=settings.local.default_system_prompt
+ from private_gpt.components.llm.prompt.prompt_helper import (
+ get_prompt_style,
+ )
+
+ prompt_style = get_prompt_style(
+ prompt_style=settings.local.prompt_style,
+ template_name=settings.local.template_name,
+ default_system_prompt=settings.local.default_system_prompt,
  )
 
  self.llm = LlamaCPP(
@@ -43,6 +47,7 @@ def __init__(self, settings: Settings) -> None:
  completion_to_prompt=prompt_style.completion_to_prompt,
  verbose=True,
  )
+ # prompt_style.improve_prompt_format(llm=cast(LlamaCPP, self.llm))
 
  case "sagemaker":
  from private_gpt.components.llm.custom.sagemaker import SagemakerLLM