"""ExperienceStore - 任务经验存储 提供两种后端实现： - ExperienceStore: 基于 PostgreSQL + pgvector 的语义检索存储 - InMemoryExperienceStore: 基于内存字典的轻量存储（用于测试） 存储任务执行经验（成功路径、失败原因、耗时分布）， 支持按任务类型检索和语义搜索，追踪完成率/耗时/重试率趋势。 """ from __future__ import annotations import asyncio import logging import math import re import uuid from datetime import datetime, timedelta, timezone from typing import Any from sqlalchemy import text from sqlalchemy.exc import DBAPIError _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 (ConnectionError, RuntimeError, asyncio.TimeoutError, ValueError) 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 (DBAPIError, ValueError, KeyError, RuntimeError) 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 (DBAPIError, ValueError, KeyError, RuntimeError) 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 (DBAPIError, ValueError, KeyError, RuntimeError) 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 (ConnectionError, RuntimeError, asyncio.TimeoutError, ValueError) 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 (ConnectionError, RuntimeError, asyncio.TimeoutError, ValueError) 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)