fix: typing in prompt transformation

zylon-ai · Oct 27, 2023 · 880a485 · 880a485
1 parent 3dd4185
commit 880a485
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/private_gpt/components/llm/custom/sagemaker.py b/private_gpt/components/llm/custom/sagemaker.py
@@ -21,8 +21,6 @@
 )
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
-
     from llama_index.callbacks import CallbackManager
     from llama_index.llms import (
         CompletionResponseGen,
@@ -113,10 +111,10 @@ class SagemakerLLM(CustomLLM):
     context_window: int = Field(
         description="The maximum number of context tokens for the model."
     )
-    messages_to_prompt: Callable[..., str] = Field(
+    messages_to_prompt: Any = Field(
         description="The function to convert messages to a prompt.", exclude=True
     )
-    completion_to_prompt: Callable[..., str] = Field(
+    completion_to_prompt: Any = Field(
         description="The function to convert a completion to a prompt.", exclude=True
     )
     generate_kwargs: dict[str, Any] = Field(

diff --git a/private_gpt/components/llm/llm_component.py b/private_gpt/components/llm/llm_component.py
@@ -37,6 +37,8 @@ def __init__(self) -> None:
 
                 self.llm = SagemakerLLM(
                     endpoint_name=settings.sagemaker.endpoint_name,
+                    messages_to_prompt=messages_to_prompt,
+                    completion_to_prompt=completion_to_prompt,
                 )
             case "openai":
                 from llama_index.llms import OpenAI