Update app_pages/text_image_generation_library.py

app_pages/text_image_generation_library.py

@@ -1,24 +1,28 @@
 from __future__ import annotations
 
 from datetime import datetime, timedelta, date
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 import streamlit as st
 
 from database.connections import get_results_collection
 
 
-def _distinct_options(col, field: str, source: str) -> List[str]:
+def _distinct_options(col, field: str, source: str, created_by: Optional[str] = None) -> List[str]:
     try:
+        q: Dict[str, Any] = {"source": source}
+        if created_by:
+            q["created_by"] = str(created_by)
+
         if field == "platform":
             vals = (
-                col.distinct("platform", {"source": source})
-                + col.distinct("settings.platform", {"source": source})
-                + col.distinct("provider_settings.platform", {"source": source})
-                + col.distinct("metadata.platform", {"source": source})
+                col.distinct("platform", q)
+                + col.distinct("settings.platform", q)
+                + col.distinct("provider_settings.platform", q)
+                + col.distinct("metadata.platform", q)
             )
         else:
-            vals = col.distinct(field, {"source": source})
+            vals = col.distinct(field, q)
 
         vals = [v for v in vals if v not in (None, "", [])]
         vals = sorted(set(vals))
@@ -34,14 +38,21 @@ def _query_docs(
     end_d: date,
     category: str,
     platform: str,
+    uid: Optional[str] = None,
     limit: int = 200,
 ) -> List[Dict[str, Any]]:
+
     try:
         start_dt = datetime.combine(start_d, datetime.min.time())
         end_dt = datetime.combine(end_d + timedelta(days=1), datetime.min.time())
         q: Dict[str, Any] = {"source": source, "created_at": {"$gte": start_dt, "$lt": end_dt}}
+
+        if uid:
+            q["created_by"] = str(uid)
+
         if category and category != "All":
             q["category"] = category
+
         if platform and platform != "All":
             q["$or"] = [
                 {"platform": platform},
@@ -49,14 +60,15 @@ def _query_docs(
                 {"provider_settings.platform": platform},
                 {"metadata.platform": platform},
             ]
+
         cursor = col.find(q).sort("created_at", -1).limit(limit)
         return list(cursor)
     except Exception:
         return []
 
 
-
 def _render_grid(docs: List[Dict[str, Any]], key_prefix: str) -> None:
+    """Render results in a 4-column grid with metadata captions."""
     seen = set()
     items = []
     for d in docs:
@@ -90,8 +102,8 @@ def _render_grid(docs: List[Dict[str, Any]], key_prefix: str) -> None:
                 st.error("Failed to load image")
 
 
-def render_text_image_library() -> None:
-
+def render_text_image_library(uid: Optional[str] = None) -> None:
+    """Streamlit entrypoint for text-based image library."""
     col = get_results_collection()
     if col is None:
         st.warning("Database not available. Set MONGO_URI and restart.")
@@ -106,8 +118,8 @@ def render_text_image_library() -> None:
     with c2:
         end_date = st.date_input("End date", value=today, key="img_text_end")
 
-    cat_opts = _distinct_options(col, "category", source="text")
-    plat_opts = _distinct_options(col, "platform", source="text")
+    cat_opts = _distinct_options(col, "category", source="text", created_by=uid)
+    plat_opts = _distinct_options(col, "platform", source="text", created_by=uid)
 
     r1, r2 = st.columns(2)
     with r1:
@@ -116,5 +128,6 @@ def render_text_image_library() -> None:
         platform = st.selectbox("Platform", options=plat_opts, key="img_text_plat")
 
     st.markdown("---")
-    docs = _query_docs(col, "text", start_date, end_date, category, platform)
+    docs = _query_docs(col, "text", start_date, end_date, category, platform, uid=uid)
     _render_grid(docs, key_prefix="text")
+