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] = [] @staticmethod 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=[])