dify/api/core/chain/llm_chain.py

49 lines
1.8 KiB
Python
Raw Normal View History

from typing import Any, Dict, List, Optional
2023-10-12 03:02:53 +08:00
from core.agent.agent.agent_llm_callback import AgentLLMCallback
from core.entities.application_entities import ModelConfigEntity
from core.entities.message_entities import lc_messages_to_prompt_messages
from core.model_manager import ModelInstance
2023-10-12 03:02:53 +08:00
from core.third_party.langchain.llms.fake import FakeLLM
from langchain import LLMChain as LCLLMChain
from langchain.callbacks.manager import CallbackManagerForChainRun
from langchain.schema import Generation, LLMResult
from langchain.schema.language_model import BaseLanguageModel
2023-10-12 03:02:53 +08:00
class LLMChain(LCLLMChain):
model_config: ModelConfigEntity
2023-10-12 03:02:53 +08:00
"""The language model instance to use."""
llm: BaseLanguageModel = FakeLLM(response="")
parameters: Dict[str, Any] = {}
agent_llm_callback: Optional[AgentLLMCallback] = None
2023-10-12 03:02:53 +08:00
def generate(
self,
input_list: List[Dict[str, Any]],
run_manager: Optional[CallbackManagerForChainRun] = None,
) -> LLMResult:
"""Generate LLM result from inputs."""
prompts, stop = self.prep_prompts(input_list, run_manager=run_manager)
messages = prompts[0].to_messages()
prompt_messages = lc_messages_to_prompt_messages(messages)
model_instance = ModelInstance(
provider_model_bundle=self.model_config.provider_model_bundle,
model=self.model_config.model,
)
result = model_instance.invoke_llm(
prompt_messages=prompt_messages,
stream=False,
stop=stop,
callbacks=[self.agent_llm_callback] if self.agent_llm_callback else None,
model_parameters=self.parameters
2023-10-12 03:02:53 +08:00
)
generations = [
[Generation(text=result.message.content)]
2023-10-12 03:02:53 +08:00
]
return LLMResult(generations=generations)