513 lines
19 KiB
Python
513 lines
19 KiB
Python
"""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)
|