fix: Invalidate cache on project reset (#3113)

* Invalidate cache on project reset

* Do not reset cache if it doesn't exist

src/phoenix/db/insertion/span.py

@@ -14,6 +14,10 @@ class SpanInsertionResult(NamedTuple):
     project_rowid: int
 
 
+class ClearProjectSpansResult(NamedTuple):
+    project_rowid: int
+
+
 async def insert_span(
     session: AsyncSession,
     span: Span,

src/phoenix/server/api/context.py

@@ -11,6 +11,7 @@
 
 from phoenix.core.model_schema import Model
 from phoenix.server.api.dataloaders import (
+    CacheForDataLoaders,
     DocumentEvaluationsDataLoader,
     DocumentEvaluationSummaryDataLoader,
     DocumentRetrievalMetricsDataLoader,
@@ -49,6 +50,7 @@ class Context:
     response: Optional[Response]
     db: Callable[[], AsyncContextManager[AsyncSession]]
     data_loaders: DataLoaders
+    cache_for_dataloaders: Optional[CacheForDataLoaders]
     model: Model
     export_path: Path
     corpus: Optional[Model] = None

src/phoenix/server/api/dataloaders/__init__.py

@@ -1,12 +1,13 @@
 from dataclasses import dataclass, field
 from functools import singledispatchmethod
+from typing import Union
 
 from phoenix.db.insertion.evaluation import (
     DocumentEvaluationInsertionResult,
     SpanEvaluationInsertionResult,
     TraceEvaluationInsertionResult,
 )
-from phoenix.db.insertion.span import SpanInsertionResult
+from phoenix.db.insertion.span import ClearProjectSpansResult, SpanInsertionResult
 
 from .document_evaluation_summaries import (
     DocumentEvaluationSummaryCache,
@@ -61,7 +62,7 @@ class CacheForDataLoaders:
     )
 
     @singledispatchmethod
-    def invalidate(self, result: SpanInsertionResult) -> None:
+    def invalidate(self, result: Union[SpanInsertionResult, ClearProjectSpansResult]) -> None:
         project_rowid, *_ = result
         self.latency_ms_quantile.invalidate(project_rowid)
         self.token_count.invalidate(project_rowid)

src/phoenix/server/api/schema.py

@@ -12,6 +12,7 @@
 
 from phoenix.config import DEFAULT_PROJECT_NAME
 from phoenix.db import models
+from phoenix.db.insertion.span import ClearProjectSpansResult
 from phoenix.pointcloud.clustering import Hdbscan
 from phoenix.server.api.context import Context
 from phoenix.server.api.helpers import ensure_list
@@ -294,6 +295,8 @@ async def clear_project(self, info: Info[Context, None], id: GlobalID) -> Query:
         delete_statement = delete(models.Trace).where(models.Trace.project_rowid == project_id)
         async with info.context.db() as session:
             await session.execute(delete_statement)
+            if cache := info.context.cache_for_dataloaders:
+                cache.invalidate(ClearProjectSpansResult(project_rowid=project_id))
         return Query()
 
 

src/phoenix/server/app.py

@@ -225,6 +225,7 @@ async def get_context(
                 ),
                 trace_evaluations=TraceEvaluationsDataLoader(self.db),
             ),
+            cache_for_dataloaders=self.cache_for_dataloaders,
         )
 
 

tests/server/api/types/conftest.py

@@ -46,6 +46,7 @@ def create_context(
             export_path=Path(TemporaryDirectory().name),
             db=None,  # TODO(persistence): add mock for db
             data_loaders=None,  # TODO(persistence): add mock for data_loaders
+            cache_for_dataloaders=None,
         )
 
     return create_context