92 lines
3.4 KiB
Python
92 lines
3.4 KiB
Python
from typing import Optional, List
|
|
from bson import ObjectId
|
|
from motor.motor_asyncio import AsyncIOMotorClient
|
|
from models.Idea import Idea
|
|
|
|
class IdeaRepo:
|
|
def __init__(self, client: AsyncIOMotorClient, db_name="bot_db"):
|
|
self.collection = client[db_name]["ideas"]
|
|
|
|
async def create_idea(self, idea: Idea) -> str:
|
|
res = await self.collection.insert_one(idea.model_dump())
|
|
return str(res.inserted_id)
|
|
|
|
async def get_idea(self, idea_id: str) -> Optional[Idea]:
|
|
if not ObjectId.is_valid(idea_id):
|
|
return None
|
|
res = await self.collection.find_one({"_id": ObjectId(idea_id)})
|
|
if res:
|
|
res["id"] = str(res.pop("_id"))
|
|
return Idea(**res)
|
|
return None
|
|
|
|
async def get_ideas(self, project_id: Optional[str], user_id: str, limit: int = 20, offset: int = 0) -> List[dict]:
|
|
if project_id:
|
|
match_stage = {"project_id": project_id, "is_deleted": False}
|
|
else:
|
|
match_stage = {"created_by": user_id, "project_id": None, "is_deleted": False}
|
|
|
|
pipeline = [
|
|
{"$match": match_stage},
|
|
{"$sort": {"updated_at": -1}},
|
|
{"$skip": offset},
|
|
{"$limit": limit},
|
|
# Add string id field for lookup
|
|
{"$addFields": {"str_id": {"$toString": "$_id"}}},
|
|
# Lookup generations
|
|
{
|
|
"$lookup": {
|
|
"from": "generations",
|
|
"let": {"idea_id": "$str_id"},
|
|
"pipeline": [
|
|
{
|
|
"$match": {
|
|
"$and": [
|
|
{"$expr": {"$eq": ["$idea_id", "$$idea_id"]}},
|
|
{"status": "done"},
|
|
{"result_list": {"$exists": True, "$not": {"$size": 0}}},
|
|
{"is_deleted": False}
|
|
]
|
|
}
|
|
},
|
|
{"$sort": {"created_at": -1}}, # Ensure we get the latest successful
|
|
{"$limit": 1}
|
|
],
|
|
"as": "generations"
|
|
}
|
|
},
|
|
# Unwind generations array (preserve ideas without generations)
|
|
{"$unwind": {"path": "$generations", "preserveNullAndEmptyArrays": True}},
|
|
# Rename for clarity
|
|
{"$addFields": {
|
|
"last_generation": "$generations",
|
|
"id": "$str_id"
|
|
}},
|
|
{"$project": {"generations": 0, "str_id": 0, "_id": 0}}
|
|
]
|
|
|
|
return await self.collection.aggregate(pipeline).to_list(None)
|
|
|
|
async def delete_idea(self, idea_id: str) -> bool:
|
|
if not ObjectId.is_valid(idea_id):
|
|
return False
|
|
res = await self.collection.update_one(
|
|
{"_id": ObjectId(idea_id)},
|
|
{"$set": {"is_deleted": True}}
|
|
)
|
|
return res.modified_count > 0
|
|
|
|
async def update_idea(self, idea: Idea) -> bool:
|
|
if not idea.id or not ObjectId.is_valid(idea.id):
|
|
return False
|
|
|
|
idea_dict = idea.model_dump()
|
|
if "id" in idea_dict:
|
|
del idea_dict["id"]
|
|
|
|
res = await self.collection.update_one(
|
|
{"_id": ObjectId(idea.id)},
|
|
{"$set": idea_dict}
|
|
)
|
|
return res.modified_count > 0
|