|
|
from abc import ABC, abstractmethod |
|
|
from typing import Callable, List, Tuple |
|
|
|
|
|
from langchain_core.language_models import BaseLanguageModel |
|
|
from langchain_core.language_models.chat_models import BaseChatModel |
|
|
from langchain_core.language_models.llms import BaseLLM |
|
|
from langchain_core.prompts import BasePromptTemplate |
|
|
from pydantic import BaseModel, Field |
|
|
|
|
|
|
|
|
class BasePromptSelector(BaseModel, ABC): |
|
|
"""Base class for prompt selectors.""" |
|
|
|
|
|
@abstractmethod |
|
|
def get_prompt(self, llm: BaseLanguageModel) -> BasePromptTemplate: |
|
|
"""Get default prompt for a language model.""" |
|
|
|
|
|
|
|
|
class ConditionalPromptSelector(BasePromptSelector): |
|
|
"""Prompt collection that goes through conditionals.""" |
|
|
|
|
|
default_prompt: BasePromptTemplate |
|
|
"""Default prompt to use if no conditionals match.""" |
|
|
conditionals: List[ |
|
|
Tuple[Callable[[BaseLanguageModel], bool], BasePromptTemplate] |
|
|
] = Field(default_factory=list) |
|
|
"""List of conditionals and prompts to use if the conditionals match.""" |
|
|
|
|
|
def get_prompt(self, llm: BaseLanguageModel) -> BasePromptTemplate: |
|
|
"""Get default prompt for a language model. |
|
|
|
|
|
Args: |
|
|
llm: Language model to get prompt for. |
|
|
|
|
|
Returns: |
|
|
Prompt to use for the language model. |
|
|
""" |
|
|
for condition, prompt in self.conditionals: |
|
|
if condition(llm): |
|
|
return prompt |
|
|
return self.default_prompt |
|
|
|
|
|
|
|
|
def is_llm(llm: BaseLanguageModel) -> bool: |
|
|
"""Check if the language model is a LLM. |
|
|
|
|
|
Args: |
|
|
llm: Language model to check. |
|
|
|
|
|
Returns: |
|
|
True if the language model is a BaseLLM model, False otherwise. |
|
|
""" |
|
|
return isinstance(llm, BaseLLM) |
|
|
|
|
|
|
|
|
def is_chat_model(llm: BaseLanguageModel) -> bool: |
|
|
"""Check if the language model is a chat model. |
|
|
|
|
|
Args: |
|
|
llm: Language model to check. |
|
|
|
|
|
Returns: |
|
|
True if the language model is a BaseChatModel model, False otherwise. |
|
|
""" |
|
|
return isinstance(llm, BaseChatModel) |
|
|
|
