Spaces:
Sleeping
Sleeping
File size: 2,955 Bytes
5db30ac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | from __future__ import annotations
from typing import Any, Dict
from langchain.chat_models import init_chat_model
from langchain.chat_models.base import BaseChatModel
from langchain_core.messages import AIMessage
from elbench.basic_education_runtime.config import CONFIG
from elbench.basic_education_runtime.entity import ModelConfig
from elbench.basic_education_runtime.spend_callback import SpendCallback
from elbench.utils import spend
class MockRuntimeChatModel:
def __init__(self, model_config: ModelConfig) -> None:
self.model_config = model_config
async def ainvoke(self, messages: list[Any]) -> AIMessage:
prefix = ""
if self.model_config.kargs:
prefix = str(self.model_config.kargs.get("prefix", "")).strip()
last_content = ""
for item in reversed(messages):
content = getattr(item, "content", None)
if isinstance(content, str) and content.strip():
last_content = content.strip()
break
if isinstance(item, dict):
dict_content = item.get("content")
if isinstance(dict_content, str) and dict_content.strip():
last_content = dict_content.strip()
break
summary = last_content[:120]
content = "占位回答：我们一步一步来。<end>"
if summary:
content = f"{prefix} 占位回答：已收到问题。{summary}\n我们一步一步来。<end>".strip()
elif prefix:
content = f"{prefix} {content}".strip()
return AIMessage(content=content)
def bind_tools(self, tools: list[Any], tool_choice: Any = None) -> "MockRuntimeChatModel":
return self
class SystemOneRuntimeModel:
def __init__(self, model_config: ModelConfig) -> None:
self.model_config = model_config
def init_chat_model_from_dict(mc: ModelConfig) -> BaseChatModel:
if mc.type == "mock":
return MockRuntimeChatModel(mc) # type: ignore[return-value]
if mc.type == "system_one":
return SystemOneRuntimeModel(mc) # type: ignore[return-value]
if mc.kargs is not None:
llm = init_chat_model(
model=f"{mc.type}:{mc.model}",
api_key=mc.api_key,
base_url=mc.api_base,
**mc.kargs,
)
else:
llm = init_chat_model(
model=f"{mc.type}:{mc.model}",
api_key=mc.api_key,
base_url=mc.api_base,
)
if mc.role and spend.active_ledger() is not None:
llm.callbacks = [*(llm.callbacks or []), SpendCallback(mc.role)]
return llm
def init_model_map_from_dict() -> Dict[str, BaseChatModel]:
cfg = CONFIG.models
result = {}
for k, v in cfg.items():
result[k] = init_chat_model_from_dict(v)
return result
init_model_map = init_model_map_from_dict
if __name__ == "__main__":
a = init_model_map_from_dict()
print(a)
|