Spaces:
Sleeping
Sleeping
| import json | |
| import re | |
| from typing import Dict, Any, Literal, Optional, List, Annotated, Tuple, Final | |
| from pydantic import BaseModel, ConfigDict, Field, create_model | |
| from typing_extensions import TypedDict | |
| from langgraph.graph.message import add_messages | |
| from pathlib import Path | |
| from aiosqlite import Connection | |
| from polyfactory.factories.pydantic_factory import ModelFactory | |
| import math | |
| # Common | |
| class State(TypedDict): | |
| messages: Annotated[list, add_messages] | |
| # Memory | |
| class Memory(BaseModel): | |
| path: Path = Path(".") | |
| # RetryConfig | |
| class RetryConfig(BaseModel): | |
| attempt: int = 3 | |
| interval: int = 3 | |
| # Global | |
| class GlobalConfig(BaseModel): | |
| concurrency: int = 8 | |
| recursion_limit: int = 25 | |
| memory: Memory = Memory() | |
| retry: RetryConfig = RetryConfig() | |
| model_call_timeout_seconds: int = 600 | |
| # Model | |
| class ModelConfig(BaseModel): | |
| api_base: Optional[str] | |
| api_key: Optional[str] | |
| kargs: Optional[Dict[str, Any]] = None | |
| model: Optional[str] | |
| type: str = "openai" | |
| # Which part of the run this model plays (target, student, judge); used only for spending limits. | |
| role: Optional[str] = None | |
| # Agent | |
| class Prompt(BaseModel): | |
| role: Optional[str] | |
| content: str | |
| class SwitchConfig(BaseModel): | |
| swap_user_assistant: bool = True | |
| class AgentMemoryConfig(BaseModel): | |
| enable: bool = True | |
| id: Optional[str] = None | |
| keep_turns: int = 3 | |
| # when_switch: SwitchConfig = SwitchConfig() | |
| class AgentConfig(BaseModel): | |
| model: str | |
| prompt: Final[List[Prompt]] | |
| memory: AgentMemoryConfig = AgentMemoryConfig(enable=True) | |
| checkpointer: Optional[Any] = None | |
| # Task | |
| class TaskConfig(BaseModel): | |
| start_prompt: Optional[Prompt] = None | |
| variables: List[Dict[str, str]] = [] | |
| def model_post_init(self, __context): | |
| object.__setattr__(self, "_frozen_start_prompt", True) | |
| def __setattr__(self, name, value): | |
| if getattr(self, f"_frozen_{name}", False) and name == "start_prompt": | |
| raise AttributeError(f"{name} is const and cannot be modified") | |
| super().__setattr__(name, value) | |
| # Runtime Context | |
| class BasicEducationRuntimeContext(BaseModel): | |
| model_config = ConfigDict(arbitrary_types_allowed=True) | |
| conns: List[Connection] = [] | |
| # ExportFormat | |
| class ExportFormat(BaseModel): | |
| task: Dict[str, str] | |
| messages: List[Prompt] = [] | |
| def from_json_file(file_path: Path | str) -> "ExportFormat": | |
| with open(file_path, "r", encoding="utf-8") as f: | |
| data = json.load(f) | |
| return ExportFormat(**data) | |
| def message_function(self, function_call: str) -> str: | |
| if function_call == "as_dialog()": | |
| strings = [] | |
| for message in self.messages: | |
| if message.role is None: | |
| continue | |
| strings.append(f"{message.role}: {message.content}") | |
| return "\n".join(strings) | |
| else: | |
| raise ValueError(f"Unsupported function call: {function_call}") | |
| def replace_template(self, template: str) -> str: | |
| """ | |
| 将输入的模板字符串中的占位符替换为实际值 | |
| 目前支持的占位符有： | |
| task.xxxx 代表task的某个字段 | |
| messages.as_dialog() 代表对话形式的messages | |
| messages.as_dialog() 形式如下： | |
| teacher: xxxxxxx | |
| student: xxxxxxx | |
| """ | |
| # Only runtime placeholders should be substituted. Model outputs may | |
| # legitimately contain braces such as "{1}", which must remain literal. | |
| placeholders = re.findall(r"\{(?:task\.[^{}]+|messages\.[^{}]+)\}", template) | |
| for placeholder in placeholders: | |
| # 获取占位符的名称，即{}中的内容 | |
| placeholder_name = placeholder.strip("{}") | |
| # 获取占位符对应的值，根据占位符的名称从task或messages中获取 | |
| if placeholder_name.startswith("task."): | |
| field_name = placeholder_name.split(".")[1] | |
| value = self.task[field_name] | |
| elif placeholder_name.startswith("messages."): | |
| message_name = placeholder_name.split(".")[1] | |
| # 如果是函数调用形式 | |
| if "(" in message_name: | |
| value = self.message_function(message_name) | |
| else: | |
| raise Exception(f"Invalid message name: {message_name}") | |
| else: | |
| raise Exception(f"Invalid placeholder name: {placeholder_name}") | |
| # 将占位符替换为对应的值 | |
| template = template.replace(placeholder, str(value)) | |
| return template | |
| # FormatField | |
| class FormatField(BaseModel): | |
| field: str | |
| type: Literal["str", "int", "float", "bool", "dict"] | |
| description: str | |
| items: List["FormatField"] = [] | |
| # min: float | int | None = None | |
| max: float | int | None = None | |
| # Evaluation | |
| class EvalConfig(BaseModel): | |
| model: str | |
| name: Optional[str] | |
| prompt: List[Prompt] | |
| format: List[FormatField] | |
| format_mode: Literal["tool", "prompt"] = "tool" | |
| def format_to_json_schema(self) -> str: | |
| model = self.format_to_pydantic() | |
| json_schema = model.model_json_schema() | |
| return json.dumps(json_schema, ensure_ascii=False) | |
| def format_to_json_example(self) -> str: | |
| mmodel: type[BaseModel] = self.format_to_pydantic() | |
| class MMF(ModelFactory): | |
| __model__ = mmodel | |
| return MMF().build().json() | |
| def get_prompts(self) -> Tuple[str, List[Prompt]]: | |
| """获取系统提示和其他提示词""" | |
| system_prompt = "" | |
| other_prompt: List[Prompt] = [] | |
| for p in self.prompt: | |
| if p.role == "system": | |
| if system_prompt != "": | |
| raise ValueError("存在多个系统提示") | |
| system_prompt = p.content | |
| else: | |
| other_prompt.append(p) | |
| return system_prompt, other_prompt | |
| def format_to_pydantic(self) -> type[BaseModel]: | |
| def field_type_from_format(f: FormatField) -> tuple: | |
| """将FormatField转成pydantic字段元组（类型，Field信息）""" | |
| python_type_map = { | |
| "str": str, | |
| "int": int, | |
| "float": float, | |
| "bool": bool, | |
| } | |
| if f.type == "dict": | |
| nested_model = build_model_from_format(f.items, f.field) | |
| return nested_model, Field(..., description=f.description) | |
| if f.type == "int" or f.type == "float": | |
| py_type = python_type_map[f.type] | |
| # if f.min is None: | |
| # if f.max is None: | |
| # return py_type, Field(..., description=f.description) | |
| # else: | |
| # return py_type, Field(..., le=f.max, description=f.description) | |
| # else: | |
| # if f.max is None: | |
| # return py_type, Field(..., ge=f.min, description=f.description) | |
| # else: | |
| # return py_type, Field(..., ge=f.min, le=f.max, description=f.description) | |
| if f.max is None: | |
| return py_type, Field(..., description=f.description) | |
| else: | |
| return py_type, Field(..., le=f.max, description=f.description) | |
| else: | |
| py_type = python_type_map[f.type] | |
| return py_type, Field(..., description=f.description) | |
| def build_model_from_format( | |
| fields: List[FormatField], model_name: str = "DynamicModel" | |
| ) -> type[BaseModel]: | |
| annotations = {} | |
| for f in fields: | |
| annotations[f.field] = field_type_from_format(f) | |
| return create_model(model_name, **annotations) | |
| return build_model_from_format(self.format, "GeneratedModel") | |
| # Runtime | |
| class BasicEducationRuntimeConfig(BaseModel): | |
| globals: GlobalConfig | |
| models: Dict[str, ModelConfig] | |
| agents: Dict[str, AgentConfig] | |
| directions: List[str] | |
| tasks: TaskConfig | |
| evaluation: Optional[EvalConfig] = None | |
| context: BasicEducationRuntimeContext = BasicEducationRuntimeContext(conns=[]) | |