Update database/operations.py

database/operations.py

@@ -57,31 +57,6 @@ def start_job(
     settings: Dict[str, Any],
     user_prompt: Optional[str] = None,
 ) -> str:
-    """
-    Variation schema (AI Library compatible):
-      {
-        type: "variation",
-        source: "variation",
-        created_by, category, file_name, prompt, lob,
-        status: "in_progress",
-        urls: [],
-        created_at,
-        settings, inputs
-      }
-
-    Generation schema:
-      {
-        type: "generation",
-        source: "text",
-        category, prompt,
-        status: "in_progress",
-        urls: [],
-        settings,
-        inputs,
-        created_by,
-        created_at,
-      }
-    """
     now = datetime.utcnow()
     doc: Dict[str, Any] = {
         "type": type,
@@ -111,9 +86,6 @@ def finish_job(
     outputs_urls: Optional[List[str]] = None,
     provider_update: Optional[Dict[str, Any]] = None,
 ) -> None:
-    """
-    Finalizes a job. Writes top-level `urls` and `status`.
-    """
     set_fields: Dict[str, Any] = {"status": status}
     if outputs_urls is not None:
         set_fields["urls"] = list(dict.fromkeys(outputs_urls or []))
@@ -125,8 +97,7 @@ def finish_job(
     except Exception as e:
         log.error(f"finish_job update failed: {e}")
 
-
-# ---------- Video analyses ----------
+# ------------------- Video Analyses -------------------
 def insert_video_analysis(
     *,
     video_name: str,
@@ -152,30 +123,35 @@ def insert_video_analysis(
     res = col.insert_one(doc)
     return str(res.inserted_id)
 
-
-def list_video_categories() -> List[str]:
+def list_video_categories(created_by: Optional[str] = None) -> List[str]:
     col = get_video_collection()
     if col is None:
         return []
     try:
-        vals = col.distinct("category")
+        q = {}
+        if created_by:
+            q["created_by"] = created_by
+        vals = col.distinct("category", q)
         return sorted({v for v in vals if v not in (None, "", [])})
     except Exception:
         return []
 
-
 def find_video_analyses(
     *,
     category: Optional[str] = None,
     start_date: Optional[datetime] = None,
     end_date: Optional[datetime] = None,
     limit: int = 200,
+    created_by: Optional[str] = None,
 ) -> List[Dict[str, Any]]:
     col = get_video_collection()
     if col is None:
         return []
     q: Dict[str, Any] = {}
-    if category: q["category"] = category
+    if category:
+        q["category"] = category
+    if created_by:
+        q["created_by"] = created_by
     if start_date or end_date:
         rng: Dict[str, Any] = {}
         if start_date: rng["$gte"] = start_date
@@ -188,8 +164,7 @@ def find_video_analyses(
         out.append(d)
     return out
 
-
-
+# ------------------- Image/Text Generations -------------------
 def find_generation_jobs(
     *,
     category: Optional[str] = None,
@@ -197,6 +172,7 @@ def find_generation_jobs(
     end_date: Optional[datetime] = None,
     page: int = 0,
     page_size: int = 20,
+    created_by: Optional[str] = None,
 ) -> Tuple[List[Dict[str, Any]], int]:
     col = get_results_collection()
     if col is None:
@@ -205,10 +181,12 @@ def find_generation_jobs(
     q: Dict[str, Any] = {"type": "generation", "source": "text"}
     if category:
         q["category"] = category
+    if created_by:
+        q["created_by"] = created_by
     if start_date or end_date:
         rng: Dict[str, Any] = {}
         if start_date: rng["$gte"] = start_date
-        if end_date:   rng["$lt"] = end_date
+        if end_date:   rng["$lt"]  = end_date
         q["created_at"] = rng
 
     total = col.count_documents(q)
@@ -225,11 +203,7 @@ def find_generation_jobs(
         out.append(d)
     return out, int(total)
 
-
-script_collection = get_script_collection()
-
-
-
+# ------------------- Script Results -------------------
 def insert_script_result(
     *,
     video_name: str,
@@ -243,10 +217,6 @@ def insert_script_result(
     num_scripts: Optional[int] = None,
     category: Optional[str] = None,
 ) -> None:
-    """
-    Save a script generation run into the results collection.
-    Stored with type/source so the library can query it.
-    """
     col = get_script_collection()
     if col is None:
         raise ValueError("Results collection is not available.")
@@ -286,13 +256,7 @@ def find_script_results(
     if col is None:
         return [], 0
 
-    q: Dict[str, Any] = {
-        "$or": [
-            {"type": "script", "source": "script_generator"},
-            {"type": {"$exists": False}, "response": {"$exists": True}},
-        ]
-    }
-
+    q: Dict[str, Any] = {"type": "script", "source": "script_generator"}
     if start_date or end_date:
         rng: Dict[str, Any] = {}
         if start_date: rng["$gte"] = start_date
@@ -306,10 +270,12 @@ def find_script_results(
         q["video_name"] = {"$regex": video_name_query, "$options": "i"}
 
     total = col.count_documents(q)
-    cur = (col.find(q)
-             .sort("created_at", DESCENDING)
-             .skip(page * page_size)
-             .limit(page_size))
+    cur = (
+        col.find(q)
+        .sort("created_at", DESCENDING)
+        .skip(page * page_size)
+        .limit(page_size)
+    )
 
     out: List[Dict[str, Any]] = []
     for d in cur:
@@ -331,4 +297,4 @@ List[Dict[str, Any]]:
         query["created_at"]["$lt"] = end_date
 
     cursor = col.find(query).sort("created_at", DESCENDING).limit(limit)
-    return list(cursor)
+    return list(cursor)