Update database/operations.py

database/operations.py

@@ -4,22 +4,23 @@ from typing import Optional, List, Dict, Any, Tuple
 from datetime import datetime
 from bson import ObjectId
 from pymongo import ASCENDING, DESCENDING
-from pymongo import MongoClient
-
-
-from database.connections import get_video_collection, get_script_collection, get_results_collection
 
+from database.connections import (
+    get_video_collection,
+    get_script_collection,
+    get_results_collection,
+    get_image_collection,
+)
 
 log = logging.getLogger(__name__)
 
-
+# ------------------- Index Setup -------------------
 def _ensure_results_indexes(col):
     try:
         col.create_index([("created_at", DESCENDING)])
         col.create_index([("type", ASCENDING), ("created_at", DESCENDING)])
         col.create_index([("category", ASCENDING), ("created_at", DESCENDING)])
         col.create_index([("created_by", ASCENDING), ("created_at", DESCENDING)])
-
     except Exception:
         pass
 
@@ -31,7 +32,6 @@ def _ensure_video_indexes(col):
     except Exception:
         pass
 
-
 _rc = get_results_collection()
 if _rc is not None:
     _ensure_results_indexes(_rc)
@@ -44,9 +44,11 @@ _sg = get_script_collection()
 if _sg is not None:
     _ensure_video_indexes(_sg)
 
+_ia = get_image_collection()
+if _ia is not None:
+    _ensure_results_indexes(_ia)
 
-
-# ---------- Image jobs (generation / variation) ----------
+# ------------------- Image Jobs -------------------
 def start_job(
     col,
     *,
@@ -70,14 +72,12 @@ def start_job(
         "created_by": created_by,
         "created_at": now,
     }
-
     if "file_name" in (inputs or {}):
         doc["file_name"] = inputs["file_name"]
 
     res = col.insert_one(doc)
     return str(res.inserted_id)
 
-
 def finish_job(
     col,
     job_id: str,
@@ -164,7 +164,85 @@ def find_video_analyses(
         out.append(d)
     return out
 
-# ------------------- Image/Text Generations -------------------
+
+def insert_image_analysis(
+    *,
+    image_name: str,
+    response: Dict[str, Any],
+    category: Optional[str] = None,
+    created_by: Optional[str] = None,
+    analyzer_model: Optional[str] = None,
+    image_meta: Optional[Dict[str, Any]] = None,
+    thumbnail: str = "",
+) -> Optional[str]:
+    col = get_image_collection()
+    if col is None:
+        raise ValueError("Image collection not available")
+
+    doc: Dict[str, Any] = {
+        "type": "image_analysis",
+        "source": "image_analyzer",
+        "category": category or "general",
+        "image": {"name": image_name, **(image_meta or {})},
+        "analyzer_model": analyzer_model or "gpt-4o",
+        "results": response or {},
+        "created_by": created_by,
+        "created_at": datetime.utcnow(),
+        "thumbnail": thumbnail or "",
+    }
+
+    res = col.insert_one(doc)
+    return str(res.inserted_id)
+
+
+
+def list_image_categories(created_by: Optional[str] = None) -> List[str]:
+    col = get_image_collection()
+    if col is None:
+        return []
+    try:
+        q = {"type": "image_analysis"}
+        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_image_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_image_collection()
+    if col is None:
+        return []
+    q: Dict[str, Any] = {"type": "image_analysis"}
+    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
+        q["created_at"] = rng
+
+    cur = col.find(q).sort("created_at", DESCENDING).limit(max(1, int(limit)))
+    out: List[Dict[str, Any]] = []
+    for d in cur:
+        d["_id"] = str(d.get("_id"))
+        out.append(d)
+    return out
+
+
+# ------------------- Generation Jobs -------------------
 def find_generation_jobs(
     *,
     category: Optional[str] = None,
@@ -177,7 +255,6 @@ def find_generation_jobs(
     col = get_results_collection()
     if col is None:
         return [], 0
-
     q: Dict[str, Any] = {"type": "generation", "source": "text"}
     if category:
         q["category"] = category
@@ -188,7 +265,6 @@ def find_generation_jobs(
         if start_date: rng["$gte"] = start_date
         if end_date:   rng["$lt"]  = end_date
         q["created_at"] = rng
-
     total = col.count_documents(q)
     cur = (
         col.find(q)
@@ -196,7 +272,6 @@ def find_generation_jobs(
         .skip(page * page_size)
         .limit(page_size)
     )
-
     out: List[Dict[str, Any]] = []
     for d in cur:
         d["_id"] = str(d.get("_id"))
@@ -219,8 +294,7 @@ def insert_script_result(
 ) -> None:
     col = get_script_collection()
     if col is None:
-        raise ValueError("Results collection is not available.")
-
+        raise ValueError("Script collection not available.")
     doc: Dict[str, Any] = {
         "type": "script",
         "source": "script_generator",
@@ -236,13 +310,11 @@ def insert_script_result(
         "num_scripts": num_scripts,
         "created_at": datetime.utcnow(),
     }
-
     try:
         col.insert_one(doc)
     except Exception as e:
         raise ValueError(f"Failed to insert script result: {e}")
 
-
 def find_script_results(
     *,
     start_date: Optional[datetime] = None,
@@ -255,20 +327,16 @@ def find_script_results(
     col = get_script_collection()
     if col is None:
         return [], 0
-
     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
         if end_date:   rng["$lt"]  = end_date
         q["created_at"] = rng
-
     if created_by:
         q["created_by"] = created_by
-
     if video_name_query:
         q["video_name"] = {"$regex": video_name_query, "$options": "i"}
-
     total = col.count_documents(q)
     cur = (
         col.find(q)
@@ -276,25 +344,9 @@ def find_script_results(
         .skip(page * page_size)
         .limit(page_size)
     )
-
     out: List[Dict[str, Any]] = []
     for d in cur:
         d["_id"] = str(d.get("_id"))
         out.append(d)
     return out, int(total)
 
-
-def get_all_scripts(start_date: Optional[datetime] = None, end_date: Optional[datetime] = None, limit: int = 20) -> \
-List[Dict[str, Any]]:
-    col = get_script_collection()
-    if col is None:
-        return []
-    query = {"type": "generation", "source": "text"}
-
-    if start_date:
-        query["created_at"] = {"$gte": start_date}
-    if end_date:
-        query["created_at"]["$lt"] = end_date
-
-    cursor = col.find(query).sort("created_at", DESCENDING).limit(limit)
-    return list(cursor)