"""ExperienceStore - 任务经验存储 提供两种后端实现： - ExperienceStore: 基于 PostgreSQL + pgvector 的语义检索存储 - InMemoryExperienceStore: 基于内存字典的轻量存储（用于测试） 存储任务执行经验（成功路径、失败原因、耗时分布）， 支持按任务类型检索和语义搜索，追踪完成率/耗时/重试率趋势。 """ from __future__ import annotations import logging import math import re import uuid from datetime import datetime, timedelta, timezone from typing import Any from sqlalchemy import text _SAFE_TABLE_NAME_PATTERN = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_]*$') from agentkit.evolution.experience_schema import EvolutionMetrics, TaskExperience from agentkit.memory.embedder import Embedder from agentkit.utils.vector_math import compute_cosine_similarity logger = logging.getLogger(__name__) class ExperienceStore: """任务经验存储 - PostgreSQL + pgvector 混合存储 基于 pgvector 向量索引 + tsvector 全文索引， 支持精确匹配 task_type + 语义相似度排序 + 时效性衰减。 检索策略： 1. pgvector ``<=>`` 算符进行最近邻检索 2. Python 侧 time_decay 重排 3. 混合评分：alpha * cosine + (1 - alpha) * time_decay_score 当 pgvector_enabled=False 或 embedder 不可用时， 回退到客户端 O(N) cosine similarity。 """ def __init__( self, session_factory: Any, experience_model: Any, embedder: Embedder | None = None, decay_rate: float = 0.01, alpha: float = 0.7, retrieve_limit: int = 200, pgvector_enabled: bool = True, table_name: str = "task_experiences", ): """ Args: session_factory: 返回 async context manager 的工厂 experience_model: TaskExperience ORM 模型类 embedder: 嵌入器，用于生成向量 decay_rate: 时间衰减率（越大衰减越快） alpha: 混合评分权重，alpha * cosine + (1-alpha) * time_decay retrieve_limit: 客户端检索时的最大候选行数 pgvector_enabled: 是否使用 pgvector 原生 ``<=>`` 算符检索 table_name: pgvector 查询使用的表名 """ self._session_factory = session_factory self._experience_model = experience_model self._embedder = embedder self._decay_rate = decay_rate self._alpha = alpha self._retrieve_limit = retrieve_limit self._pgvector_enabled = pgvector_enabled self._table_name = table_name if not _SAFE_TABLE_NAME_PATTERN.match(self._table_name): raise ValueError(f"Invalid table_name: {self._table_name}. Must match [a-zA-Z_][a-zA-Z0-9_]*") async def record_experience(self, experience: TaskExperience) -> str: """记录任务经验 如果 experience.embedding 为 None 且 embedder 可用， 自动生成 embedding。 Args: experience: 任务经验数据 Returns: 经验 ID """ if not experience.experience_id: experience.experience_id = str(uuid.uuid4()) # 自动生成 embedding if experience.embedding is None and self._embedder is not None: text = experience.text_for_embedding() try: experience.embedding = await self._embedder.embed(text) except Exception as e: logger.warning(f"Failed to generate embedding for experience {experience.experience_id}: {e}") async with self._session_factory() as db: try: Model = self._experience_model entry = Model( id=experience.experience_id, task_type=experience.task_type, goal=experience.goal, steps_summary=experience.steps_summary, outcome=experience.outcome, duration_seconds=experience.duration_seconds, success_rate=experience.success_rate, failure_reasons=experience.failure_reasons, optimization_tips=experience.optimization_tips, embedding=experience.embedding, created_at=experience.created_at, ) db.add(entry) await db.commit() logger.info( f"Experience recorded: {experience.experience_id} " f"task_type={experience.task_type} outcome={experience.outcome}" ) return experience.experience_id except Exception as e: await db.rollback() logger.error(f"Failed to record experience: {e}") raise async def search( self, query: str, top_k: int = 5, task_type: str | None = None, search_multiplier: int = 5, ) -> list[TaskExperience]: """语义检索相似经验 支持精确匹配 task_type + 语义相似度排序 + 时效性衰减。 Args: query: 搜索查询文本 top_k: 返回的最大结果数 task_type: 可选的任务类型过滤 search_multiplier: 预取行数倍数 """ async with self._session_factory() as db: try: if self._pgvector_enabled and self._embedder: return await self._search_pgvector(db, query, top_k, task_type, search_multiplier) return await self._search_client_side(db, query, top_k, task_type, search_multiplier) except Exception as e: logger.error(f"Failed to search experiences: {e}") return [] async def _search_pgvector( self, db: Any, query: str, top_k: int, task_type: str | None, search_multiplier: int, ) -> list[TaskExperience]: """使用 pgvector ``<=>`` 算符检索，再 Python 侧 time_decay 重排""" query_embedding = await self._embedder.embed(query) fetch_limit = top_k * search_multiplier where_clauses = [] params: dict[str, Any] = {"query_vec": str(query_embedding), "lim": fetch_limit} if task_type: where_clauses.append("task_type = :task_type") params["task_type"] = task_type where_sql = (" WHERE " + " AND ".join(where_clauses)) if where_clauses else "" sql = text( f"SELECT *, embedding <=> :query_vec AS distance " f"FROM {self._table_name}{where_sql} " f"ORDER BY embedding <=> :query_vec " f"LIMIT :lim" ) result = await db.execute(sql, params) rows = result.mappings().all() if not rows: return [] # Re-rank with time_decay in Python items: list[tuple[float, TaskExperience]] = [] for row in rows: row_embedding = row.get("embedding") age_hours = ( (datetime.now(timezone.utc) - row["created_at"]).total_seconds() / 3600 if row.get("created_at") else 0 ) decay = math.exp(-self._decay_rate * age_hours) time_decay_score = (row.get("success_rate") or 0.5) * decay if row_embedding is not None: cosine_sim = compute_cosine_similarity(query_embedding, row_embedding) score = self._alpha * cosine_sim + (1 - self._alpha) * time_decay_score else: score = time_decay_score exp = TaskExperience( experience_id=str(row.get("id", "")), task_type=row.get("task_type", ""), goal=row.get("goal", ""), steps_summary=row.get("steps_summary", ""), outcome=row.get("outcome", "success"), duration_seconds=row.get("duration_seconds", 0.0), success_rate=row.get("success_rate", 1.0), failure_reasons=row.get("failure_reasons") or [], optimization_tips=row.get("optimization_tips") or [], embedding=row_embedding, created_at=row.get("created_at") or datetime.now(timezone.utc), ) items.append((score, exp)) items.sort(key=lambda x: x[0], reverse=True) return [exp for _, exp in items[:top_k]] async def _search_client_side( self, db: Any, query: str, top_k: int, task_type: str | None, search_multiplier: int, ) -> list[TaskExperience]: """客户端 O(N) cosine similarity 检索（回退路径）""" Model = self._experience_model from sqlalchemy import select stmt = select(Model) if task_type: stmt = stmt.where(Model.task_type == task_type) stmt = stmt.order_by(Model.created_at.desc()).limit(top_k * search_multiplier) result = await db.execute(stmt) entries = result.scalars().all() query_embedding = None if self._embedder and entries: query_embedding = await self._embedder.embed(query) items: list[tuple[float, TaskExperience]] = [] for entry in entries: age_hours = ( (datetime.now(timezone.utc) - entry.created_at).total_seconds() / 3600 if entry.created_at else 0 ) decay = math.exp(-self._decay_rate * age_hours) time_decay_score = (entry.success_rate or 0.5) * decay if self._embedder and query_embedding is not None and entry.embedding is not None: cosine_sim = compute_cosine_similarity(query_embedding, entry.embedding) score = self._alpha * cosine_sim + (1 - self._alpha) * time_decay_score else: score = time_decay_score exp = TaskExperience( experience_id=str(entry.id), task_type=entry.task_type, goal=entry.goal, steps_summary=entry.steps_summary, outcome=entry.outcome, duration_seconds=entry.duration_seconds, success_rate=entry.success_rate, failure_reasons=entry.failure_reasons or [], optimization_tips=entry.optimization_tips or [], embedding=entry.embedding, created_at=entry.created_at or datetime.now(timezone.utc), ) items.append((score, exp)) items.sort(key=lambda x: x[0], reverse=True) return [exp for _, exp in items[:top_k]] async def get_metrics( self, task_type: str | None = None, time_window: str = "24h", ) -> list[EvolutionMetrics]: """获取进化指标趋势 按任务类型和时间窗口聚合完成率、平均耗时和重试率。 Args: task_type: 可选的任务类型过滤，None 表示所有类型 time_window: 时间窗口（"1h", "24h", "7d", "30d"） """ window_delta = _parse_time_window(time_window) window_start = datetime.now(timezone.utc) - window_delta window_end = datetime.now(timezone.utc) async with self._session_factory() as db: try: where_clauses = ["created_at >= :window_start"] params: dict[str, Any] = {"window_start": window_start} if task_type: where_clauses.append("task_type = :task_type") params["task_type"] = task_type where_sql = " AND ".join(where_clauses) # 按任务类型聚合 group_by = "task_type" if task_type is None else "" select_clause = "task_type" if task_type: select_clause += f", '{task_type}' as filtered_task_type" sql = text( f"SELECT task_type, " f" COUNT(*) as sample_count, " f" AVG(CASE WHEN outcome = 'success' THEN 1.0 ELSE 0.0 END) as completion_rate, " f" AVG(duration_seconds) as avg_duration, " f" AVG(CASE WHEN success_rate < 1.0 THEN 1.0 ELSE 0.0 END) as retry_rate " f"FROM {self._table_name} " f"WHERE {where_sql} " f"GROUP BY task_type" ) result = await db.execute(sql, params) rows = result.mappings().all() metrics_list = [] for row in rows: metrics_list.append( EvolutionMetrics( task_type=row["task_type"], time_window=time_window, completion_rate=row["completion_rate"] or 0.0, avg_duration=row["avg_duration"] or 0.0, retry_rate=row["retry_rate"] or 0.0, sample_count=row["sample_count"] or 0, window_start=window_start, window_end=window_end, ) ) return metrics_list except Exception as e: logger.error(f"Failed to get metrics: {e}") return [] class InMemoryExperienceStore: """基于内存字典的任务经验存储（用于测试和轻量场景） 无需数据库，纯 dict-based 实现，支持与 ExperienceStore 相同的接口。 """ def __init__( self, embedder: Embedder | None = None, decay_rate: float = 0.01, alpha: float = 0.7, ): self._embedder = embedder self._decay_rate = decay_rate self._alpha = alpha self._experiences: dict[str, TaskExperience] = {} async def record_experience(self, experience: TaskExperience) -> str: """记录任务经验""" if not experience.experience_id: experience.experience_id = str(uuid.uuid4()) # 自动生成 embedding if experience.embedding is None and self._embedder is not None: text = experience.text_for_embedding() try: experience.embedding = await self._embedder.embed(text) except Exception as e: logger.warning(f"Failed to generate embedding for experience {experience.experience_id}: {e}") # 存储副本，避免外部修改影响内部状态 self._experiences[experience.experience_id] = TaskExperience( experience_id=experience.experience_id, task_type=experience.task_type, goal=experience.goal, steps_summary=experience.steps_summary, outcome=experience.outcome, duration_seconds=experience.duration_seconds, success_rate=experience.success_rate, failure_reasons=list(experience.failure_reasons), optimization_tips=list(experience.optimization_tips), embedding=experience.embedding, created_at=experience.created_at, ) logger.info( f"Experience recorded: {experience.experience_id} " f"task_type={experience.task_type} outcome={experience.outcome}" ) return experience.experience_id async def search( self, query: str, top_k: int = 5, task_type: str | None = None, search_multiplier: int = 5, ) -> list[TaskExperience]: """语义检索相似经验""" # 生成 query embedding query_embedding = None if self._embedder: try: query_embedding = await self._embedder.embed(query) except Exception as e: logger.warning(f"Failed to generate query embedding: {e}") # 筛选候选 candidates = list(self._experiences.values()) if task_type: candidates = [e for e in candidates if e.task_type == task_type] # 计算得分 items: list[tuple[float, TaskExperience]] = [] for exp in candidates: age_hours = ( (datetime.now(timezone.utc) - exp.created_at).total_seconds() / 3600 if exp.created_at else 0 ) decay = math.exp(-self._decay_rate * age_hours) time_decay_score = exp.success_rate * decay if query_embedding is not None and exp.embedding is not None: cosine_sim = compute_cosine_similarity(query_embedding, exp.embedding) score = self._alpha * cosine_sim + (1 - self._alpha) * time_decay_score else: score = time_decay_score items.append((score, exp)) items.sort(key=lambda x: x[0], reverse=True) return [exp for _, exp in items[:top_k]] async def get_metrics( self, task_type: str | None = None, time_window: str = "24h", ) -> list[EvolutionMetrics]: """获取进化指标趋势""" window_delta = _parse_time_window(time_window) window_start = datetime.now(timezone.utc) - window_delta window_end = datetime.now(timezone.utc) # 筛选时间窗口内的经验 candidates = [ e for e in self._experiences.values() if e.created_at >= window_start ] if task_type: candidates = [e for e in candidates if e.task_type == task_type] # 按 task_type 分组聚合 groups: dict[str, list[TaskExperience]] = {} for exp in candidates: groups.setdefault(exp.task_type, []).append(exp) metrics_list = [] for tt, exps in groups.items(): n = len(exps) if n == 0: continue completion_rate = sum(1 for e in exps if e.outcome == "success") / n avg_duration = sum(e.duration_seconds for e in exps) / n retry_rate = sum(1 for e in exps if e.success_rate < 1.0) / n metrics_list.append( EvolutionMetrics( task_type=tt, time_window=time_window, completion_rate=completion_rate, avg_duration=avg_duration, retry_rate=retry_rate, sample_count=n, window_start=window_start, window_end=window_end, ) ) return metrics_list # ── 辅助函数 ────────────────────────────────────────────── def _parse_time_window(window: str) -> timedelta: """解析时间窗口字符串为 timedelta 支持格式: "1h", "24h", "7d", "30d" """ unit = window[-1].lower() try: value = int(window[:-1]) except ValueError: return timedelta(hours=24) if unit == "h": return timedelta(hours=value) elif unit == "d": return timedelta(days=value) else: logger.warning(f"Unknown time window unit '{unit}', defaulting to 24h") return timedelta(hours=24)