+ env
This commit is contained in:
@@ -8,6 +8,7 @@ from repos.albums_repo import AlbumsRepo
|
||||
from repos.project_repo import ProjectRepo
|
||||
from repos.idea_repo import IdeaRepo
|
||||
from repos.post_repo import PostRepo
|
||||
from repos.environment_repo import EnvironmentRepo
|
||||
|
||||
|
||||
from typing import Optional
|
||||
@@ -23,3 +24,4 @@ class DAO:
|
||||
self.users = UsersRepo(client, db_name)
|
||||
self.ideas = IdeaRepo(client, db_name)
|
||||
self.posts = PostRepo(client, db_name)
|
||||
self.environments = EnvironmentRepo(client, db_name)
|
||||
|
||||
73
repos/environment_repo.py
Normal file
73
repos/environment_repo.py
Normal file
@@ -0,0 +1,73 @@
|
||||
from typing import List, Optional
|
||||
from datetime import datetime
|
||||
from bson import ObjectId
|
||||
from motor.motor_asyncio import AsyncIOMotorClient
|
||||
from models.Environment import Environment
|
||||
|
||||
|
||||
class EnvironmentRepo:
|
||||
def __init__(self, client: AsyncIOMotorClient, db_name="bot_db"):
|
||||
self.collection = client[db_name]["environments"]
|
||||
|
||||
async def create_env(self, env: Environment) -> Environment:
|
||||
env_dict = env.model_dump(exclude={"id"})
|
||||
res = await self.collection.insert_one(env_dict)
|
||||
env.id = str(res.inserted_id)
|
||||
return env
|
||||
|
||||
async def get_env(self, env_id: str) -> Optional[Environment]:
|
||||
res = await self.collection.find_one({"_id": ObjectId(env_id)})
|
||||
if not res:
|
||||
return None
|
||||
res["id"] = str(res.pop("_id"))
|
||||
return Environment(**res)
|
||||
|
||||
async def get_character_envs(self, character_id: str) -> List[Environment]:
|
||||
cursor = self.collection.find({"character_id": character_id})
|
||||
envs = []
|
||||
async for doc in cursor:
|
||||
doc["id"] = str(doc.pop("_id"))
|
||||
envs.append(Environment(**doc))
|
||||
return envs
|
||||
|
||||
async def update_env(self, env_id: str, update_data: dict) -> bool:
|
||||
update_data["updated_at"] = datetime.utcnow()
|
||||
res = await self.collection.update_one(
|
||||
{"_id": ObjectId(env_id)},
|
||||
{"$set": update_data}
|
||||
)
|
||||
return res.modified_count > 0
|
||||
|
||||
async def delete_env(self, env_id: str) -> bool:
|
||||
res = await self.collection.delete_one({"_id": ObjectId(env_id)})
|
||||
return res.deleted_count > 0
|
||||
|
||||
async def add_asset(self, env_id: str, asset_id: str) -> bool:
|
||||
res = await self.collection.update_one(
|
||||
{"_id": ObjectId(env_id)},
|
||||
{
|
||||
"$addToSet": {"asset_ids": asset_id},
|
||||
"$set": {"updated_at": datetime.utcnow()}
|
||||
}
|
||||
)
|
||||
return res.modified_count > 0
|
||||
|
||||
async def add_assets(self, env_id: str, asset_ids: List[str]) -> bool:
|
||||
res = await self.collection.update_one(
|
||||
{"_id": ObjectId(env_id)},
|
||||
{
|
||||
"$addToSet": {"asset_ids": {"$each": asset_ids}},
|
||||
"$set": {"updated_at": datetime.utcnow()}
|
||||
}
|
||||
)
|
||||
return res.modified_count > 0
|
||||
|
||||
async def remove_asset(self, env_id: str, asset_id: str) -> bool:
|
||||
res = await self.collection.update_one(
|
||||
{"_id": ObjectId(env_id)},
|
||||
{
|
||||
"$pull": {"asset_ids": asset_id},
|
||||
"$set": {"updated_at": datetime.utcnow()}
|
||||
}
|
||||
)
|
||||
return res.modified_count > 0
|
||||
@@ -97,11 +97,12 @@ class GenerationRepo:
|
||||
async def get_usage_stats(self, created_by: Optional[str] = None, project_id: Optional[str] = None) -> dict:
|
||||
"""
|
||||
Calculates usage statistics (runs, tokens, cost) using MongoDB aggregation.
|
||||
Includes even soft-deleted generations to reflect actual expenditure.
|
||||
"""
|
||||
pipeline = []
|
||||
|
||||
# 1. Match active done generations
|
||||
match_stage = {"is_deleted": False, "status": GenerationStatus.DONE}
|
||||
# 1. Match all done generations (including soft-deleted)
|
||||
match_stage = {"status": GenerationStatus.DONE}
|
||||
if created_by:
|
||||
match_stage["created_by"] = created_by
|
||||
if project_id:
|
||||
@@ -156,10 +157,11 @@ class GenerationRepo:
|
||||
async def get_usage_breakdown(self, group_by: str = "created_by", project_id: Optional[str] = None, created_by: Optional[str] = None) -> List[dict]:
|
||||
"""
|
||||
Returns usage statistics grouped by user or project.
|
||||
Includes even soft-deleted generations to reflect actual expenditure.
|
||||
"""
|
||||
pipeline = []
|
||||
|
||||
match_stage = {"is_deleted": False, "status": GenerationStatus.DONE}
|
||||
match_stage = {"status": GenerationStatus.DONE}
|
||||
if project_id:
|
||||
match_stage["project_id"] = project_id
|
||||
if created_by:
|
||||
|
||||
Reference in New Issue
Block a user