Spaces:
Sleeping
Sleeping
| 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) | |