69 lines
2.6 KiB
Python
69 lines
2.6 KiB
Python
|
from typing import List, Optional
|
|||
|
|
from th_agenter.llm.base_llm import BaseLLM
|
|||
|
|
from langchain_core.callbacks import CallbackManagerForLLMRun
|
|||
|
|
from langchain_core.messages import BaseMessage, AIMessage, HumanMessage
|
|||
|
|
from langchain_core.outputs import ChatResult, ChatGeneration
|
|||
|
|
|
|||
|
|
|
|||
|
|
class LocalLLM(BaseLLM):
|
|||
|
|
def __init__(self, config):
|
|||
|
|
super().__init__(config)
|
|||
|
|
self.local_config = config
|
|||
|
|
|
|||
|
|
def _validate_config(self):
|
|||
|
|
if not self.local_config.model_path:
|
|||
|
|
raise ValueError("LocalLLM 必须配置 model_path")
|
|||
|
|
|
|||
|
|
def load_model(self):
|
|||
|
|
from langchain_community.llms import LlamaCpp
|
|||
|
|
self.model = LlamaCpp(
|
|||
|
|
model_path=self.local_config.model_path,
|
|||
|
|
temperature=self.local_config.temperature,
|
|||
|
|
max_tokens=self.local_config.max_tokens,
|
|||
|
|
n_ctx=self.local_config.n_ctx,
|
|||
|
|
n_threads=self.local_config.n_threads,
|
|||
|
|
verbose=False
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def _llm_type(self) -> str:
|
|||
|
|
return "llama"
|
|||
|
|
|
|||
|
|
def _generate(
|
|||
|
|
self,
|
|||
|
|
messages: List[BaseMessage],
|
|||
|
|
stop: Optional[List[str]] = None,
|
|||
|
|
run_manager: Optional[CallbackManagerForLLMRun] = None,** kwargs: Any,
|
|||
|
|
) -> ChatResult:
|
|||
|
|
if not self.model:
|
|||
|
|
self.load_model()
|
|||
|
|
# 适配 LlamaCpp（非 Chat 模型）的调用方式
|
|||
|
|
prompt = self._format_messages(messages)
|
|||
|
|
text = self.model.invoke(prompt, stop=stop, **kwargs)
|
|||
|
|
# 构造 ChatResult（LangChain 标准格式）
|
|||
|
|
generation = ChatGeneration(message=AIMessage(content=text))
|
|||
|
|
return ChatResult(generations=[generation])
|
|||
|
|
|
|||
|
|
async def _agenerate(
|
|||
|
|
self,
|
|||
|
|
messages: List[BaseMessage],
|
|||
|
|
stop: Optional[List[str]] = None,
|
|||
|
|
run_manager: Optional[CallbackManagerForLLMRun] = None,** kwargs: Any,
|
|||
|
|
) -> ChatResult:
|
|||
|
|
if not self.model:
|
|||
|
|
self.load_model()
|
|||
|
|
prompt = self._format_messages(messages)
|
|||
|
|
text = await self.model.ainvoke(prompt, stop=stop, **kwargs)
|
|||
|
|
generation = ChatGeneration(message=AIMessage(content=text))
|
|||
|
|
return ChatResult(generations=[generation])
|
|||
|
|
|
|||
|
|
def _format_messages(self, messages: List[BaseMessage]) -> str:
|
|||
|
|
"""将 LangChain 消息列表格式化为本地模型的 Prompt"""
|
|||
|
|
prompt_parts = []
|
|||
|
|
for msg in messages:
|
|||
|
|
if isinstance(msg, HumanMessage):
|
|||
|
|
prompt_parts.append(f"<s>[INST] {msg.content} [/INST]")
|
|||
|
|
elif isinstance(msg, AIMessage):
|
|||
|
|
prompt_parts.append(msg.content)
|
|||
|
|
return "".join(prompt_parts)
|