203 lines
9.8 KiB
Python
203 lines
9.8 KiB
Python
"""User-owned analysis projects and immutable project/thread bindings."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from copy import deepcopy
|
|
from datetime import UTC, datetime
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import ForeignKey, String, Text, delete, select
|
|
from sqlalchemy.orm import Mapped, mapped_column
|
|
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.types import BeijingDateTime, PortableJSON
|
|
|
|
MODES = frozenset({"qa", "strategy", "solution"})
|
|
CONFIG_FIELDS = frozenset({"agent_id", "report_config", "strategy_config"})
|
|
EDITABLE_FIELDS = frozenset({"name", "background", "knowledge_base_ids", "assistant_knowledge_base_ids"}) | CONFIG_FIELDS
|
|
|
|
|
|
class AnalysisProjectConfigRow(Base):
|
|
"""Additive table: upgrades do not alter existing project/thread tables."""
|
|
|
|
__tablename__ = "analysis_project_configs"
|
|
key: Mapped[str] = mapped_column(String(80), primary_key=True)
|
|
config: Mapped[dict] = mapped_column(PortableJSON(), default=dict)
|
|
|
|
|
|
class AnalysisProjectRow(Base):
|
|
__tablename__ = "analysis_projects"
|
|
|
|
id: Mapped[str] = mapped_column(String(36), primary_key=True)
|
|
user_id: Mapped[str] = mapped_column(String(255), index=True)
|
|
name: Mapped[str] = mapped_column(String(160))
|
|
mode: Mapped[str] = mapped_column(String(16))
|
|
background: Mapped[str] = mapped_column(Text, default="")
|
|
knowledge_base_ids: Mapped[list] = mapped_column(PortableJSON(), default=list)
|
|
assistant_knowledge_base_ids: Mapped[list] = mapped_column(PortableJSON(), default=list)
|
|
created_at: Mapped[datetime] = mapped_column(BeijingDateTime())
|
|
updated_at: Mapped[datetime] = mapped_column(BeijingDateTime())
|
|
|
|
|
|
class AnalysisProjectThreadRow(Base):
|
|
__tablename__ = "analysis_project_threads"
|
|
|
|
thread_id: Mapped[str] = mapped_column(String(64), primary_key=True)
|
|
project_id: Mapped[str] = mapped_column(ForeignKey("analysis_projects.id"), index=True)
|
|
created_at: Mapped[datetime] = mapped_column(BeijingDateTime())
|
|
|
|
|
|
def _dict(row):
|
|
data = row.to_dict()
|
|
for key in ("created_at", "updated_at"):
|
|
if isinstance(data.get(key), datetime):
|
|
data[key] = data[key].isoformat()
|
|
return deepcopy(data)
|
|
|
|
|
|
class AnalysisProjectStore:
|
|
"""SQL in persistent deployments; isolated in-memory fallback for memory mode.
|
|
|
|
Every public project operation requires an explicit owner. The unscoped
|
|
thread lookup is only for the Gateway's ownership and run-policy checks.
|
|
Mode has no update path, and bindings are separate from client-editable metadata.
|
|
"""
|
|
|
|
def __init__(self, session_factory=None):
|
|
self._sf = session_factory
|
|
self._projects: dict[str, dict] = {}
|
|
self._threads: dict[str, dict] = {}
|
|
self._configs: dict[str, dict] = {}
|
|
|
|
async def config_for(self, key: str) -> dict:
|
|
if not self._sf:
|
|
return deepcopy(self._configs.get(key, {}))
|
|
async with self._sf() as session:
|
|
row = await session.get(AnalysisProjectConfigRow, key)
|
|
return deepcopy(row.config) if row else {}
|
|
|
|
async def _enrich(self, project: dict | None) -> dict | None:
|
|
if project is None:
|
|
return None
|
|
return {
|
|
**project,
|
|
"agent_id": None,
|
|
"report_config": {},
|
|
"strategy_config": {},
|
|
**await self.config_for(f"project:{project['id']}"),
|
|
}
|
|
|
|
async def create(self, user_id: str, *, mode: str, **fields) -> dict:
|
|
if mode not in MODES or set(fields) - EDITABLE_FIELDS:
|
|
raise ValueError("Invalid project fields")
|
|
now = datetime.now(UTC)
|
|
row = AnalysisProjectRow(
|
|
id=str(uuid4()),
|
|
user_id=user_id,
|
|
mode=mode,
|
|
name=fields["name"],
|
|
background=fields.get("background", ""),
|
|
knowledge_base_ids=list(fields.get("knowledge_base_ids", [])),
|
|
assistant_knowledge_base_ids=list(fields.get("assistant_knowledge_base_ids", [])),
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
if self._sf:
|
|
async with self._sf() as session:
|
|
session.add(row)
|
|
session.add(AnalysisProjectConfigRow(key=f"project:{row.id}", config={key: deepcopy(fields[key]) for key in CONFIG_FIELDS if key in fields}))
|
|
await session.commit()
|
|
else:
|
|
self._projects[row.id] = _dict(row)
|
|
self._configs[f"project:{row.id}"] = {key: deepcopy(fields[key]) for key in CONFIG_FIELDS if key in fields}
|
|
return await self._enrich(_dict(row))
|
|
|
|
async def list(self, user_id: str) -> list[dict]:
|
|
if not self._sf:
|
|
return [await self._enrich(deepcopy(p)) for p in sorted(self._projects.values(), key=lambda p: p["updated_at"], reverse=True) if p["user_id"] == user_id]
|
|
async with self._sf() as session:
|
|
result = await session.execute(select(AnalysisProjectRow).where(AnalysisProjectRow.user_id == user_id).order_by(AnalysisProjectRow.updated_at.desc()))
|
|
rows = [_dict(row) for row in result.scalars()]
|
|
return [await self._enrich(row) for row in rows]
|
|
|
|
async def get(self, project_id: str, user_id: str) -> dict | None:
|
|
if not self._sf:
|
|
row = self._projects.get(project_id)
|
|
return await self._enrich(deepcopy(row)) if row and row["user_id"] == user_id else None
|
|
async with self._sf() as session:
|
|
row = (await session.execute(select(AnalysisProjectRow).where(AnalysisProjectRow.id == project_id, AnalysisProjectRow.user_id == user_id))).scalar_one_or_none()
|
|
data = _dict(row) if row else None
|
|
return await self._enrich(data)
|
|
|
|
async def update(self, project_id: str, user_id: str, **fields) -> dict | None:
|
|
if set(fields) - EDITABLE_FIELDS:
|
|
raise ValueError("项目模式创建后不可更改")
|
|
if not self._sf:
|
|
row = self._projects.get(project_id)
|
|
if not row or row["user_id"] != user_id:
|
|
return None
|
|
row.update({key: deepcopy(value) for key, value in fields.items() if key not in CONFIG_FIELDS}, updated_at=datetime.now(UTC).isoformat())
|
|
self._configs.setdefault(f"project:{project_id}", {}).update({key: deepcopy(value) for key, value in fields.items() if key in CONFIG_FIELDS})
|
|
return await self._enrich(deepcopy(row))
|
|
async with self._sf() as session:
|
|
row = (await session.execute(select(AnalysisProjectRow).where(AnalysisProjectRow.id == project_id, AnalysisProjectRow.user_id == user_id).with_for_update())).scalar_one_or_none()
|
|
if row is None:
|
|
return None
|
|
for key, value in fields.items():
|
|
if key not in CONFIG_FIELDS:
|
|
setattr(row, key, deepcopy(value))
|
|
options = await session.get(AnalysisProjectConfigRow, f"project:{project_id}")
|
|
config = {**(options.config if options else {}), **{key: deepcopy(value) for key, value in fields.items() if key in CONFIG_FIELDS}}
|
|
await session.merge(AnalysisProjectConfigRow(key=f"project:{project_id}", config=config))
|
|
row.updated_at = datetime.now(UTC)
|
|
await session.commit()
|
|
data = _dict(row)
|
|
return await self._enrich(data)
|
|
|
|
async def bind_thread(self, project_id: str, user_id: str, thread_id: str, *, agent_id: str | None = None) -> None:
|
|
if await self.get(project_id, user_id) is None:
|
|
raise ValueError("Project not found")
|
|
row = AnalysisProjectThreadRow(thread_id=thread_id, project_id=project_id, created_at=datetime.now(UTC))
|
|
if not self._sf:
|
|
if thread_id in self._threads:
|
|
raise ValueError("Thread already bound")
|
|
self._threads[thread_id] = _dict(row)
|
|
self._configs[f"thread:{thread_id}"] = {"agent_id": agent_id}
|
|
return
|
|
async with self._sf() as session:
|
|
session.add(row)
|
|
session.add(AnalysisProjectConfigRow(key=f"thread:{thread_id}", config={"agent_id": agent_id}))
|
|
await session.commit()
|
|
|
|
async def project_for_thread(self, thread_id: str) -> dict | None:
|
|
if not self._sf:
|
|
binding = self._threads.get(thread_id)
|
|
return await self._enrich(deepcopy(self._projects.get(binding["project_id"]))) if binding else None
|
|
async with self._sf() as session:
|
|
row = (await session.execute(select(AnalysisProjectRow).join(AnalysisProjectThreadRow).where(AnalysisProjectThreadRow.thread_id == thread_id))).scalar_one_or_none()
|
|
data = _dict(row) if row else None
|
|
return await self._enrich(data)
|
|
|
|
async def thread_ids(self, project_id: str, user_id: str) -> list[str]:
|
|
if await self.get(project_id, user_id) is None:
|
|
return []
|
|
if not self._sf:
|
|
return [key for key, value in reversed(list(self._threads.items())) if value["project_id"] == project_id]
|
|
async with self._sf() as session:
|
|
return list((await session.execute(select(AnalysisProjectThreadRow.thread_id).where(AnalysisProjectThreadRow.project_id == project_id).order_by(AnalysisProjectThreadRow.created_at.desc()))).scalars())
|
|
|
|
async def unbind_failed_creation(self, thread_id: str) -> None:
|
|
if not self._sf:
|
|
self._threads.pop(thread_id, None)
|
|
self._configs.pop(f"thread:{thread_id}", None)
|
|
return
|
|
async with self._sf() as session:
|
|
await session.execute(delete(AnalysisProjectThreadRow).where(AnalysisProjectThreadRow.thread_id == thread_id))
|
|
await session.execute(delete(AnalysisProjectConfigRow).where(AnalysisProjectConfigRow.key == f"thread:{thread_id}"))
|
|
await session.commit()
|
|
|
|
async def unbind_thread(self, thread_id: str) -> None:
|
|
"""Remove a durable project binding when its ordinary thread is deleted."""
|
|
|
|
await self.unbind_failed_creation(thread_id)
|