File size: 2,860 Bytes
71dd5b6 3b10d4b 71dd5b6 38698de 3b10d4b 71dd5b6 3b10d4b 0c1d317 3b10d4b 71dd5b6 3b10d4b d7ee021 71dd5b6 d7ee021 ef52f9b 3b10d4b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
from __future__ import annotations
from datetime import datetime
from typing import Any, Dict, List, Optional
from pymongo import DESCENDING
from Database.connections import get_img_gen_collection
import streamlit as st
from core.logger import logger
def insert_img_gen_run_result(
*,
prompts: List[str],
image_urls: List[str],
category: str = "general",
platform: Optional[str] = None,
model_name: Optional[str] = None,
offer: Optional[str] = None,
framework: Optional[str] = None,
user_query: Optional[str] = None,
created_by: Optional[str] = None,
) -> str:
col = get_img_gen_collection()
if col is None:
raise ValueError("img_gen collection not available.")
if not prompts or not image_urls:
raise ValueError("prompts and image_urls cannot be empty")
doc: Dict[str, Any] = {
"source": "img_generator",
"category": category or "general",
"prompts": prompts,
"image_urls": image_urls,
"platform": platform,
"model_name": model_name,
"offer": offer if offer else '',
"framework": framework if framework else '',
"user_query": user_query,
"created_by": created_by,
"created_at": datetime.utcnow(),
}
res = col.insert_one(doc)
return str(res.inserted_id)
def find_img_gen_results(
*,
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]]:
try:
col = get_img_gen_collection()
if col is None:
return []
q: Dict[str, Any] = {"source": "img_generator"}
if category:
q["category"] = category
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
cur = (
col.find(q)
.sort("created_at", DESCENDING)
.limit(limit)
)
out: List[Dict[str, Any]] = []
for d in cur:
d["_id"] = str(d.get("_id"))
out.append(d)
return out
except Exception as e:
st.error(f"Operations File: {e}")
def list_img_gen_categories(*, created_by: Optional[str] = None) -> List[str]:
col = get_img_gen_collection()
if col is None:
return []
q: Dict[str, Any] = {"source": "img_generator"}
if created_by:
q["created_by"] = created_by
cats = col.distinct("category", q)
cats = [c for c in cats if c]
cats.sort()
return cats
|