fixes
This commit is contained in:
@@ -175,6 +175,8 @@ class AssetsRepo:
|
||||
filter["linked_char_id"] = character_id
|
||||
if created_by:
|
||||
filter["created_by"] = created_by
|
||||
if project_id is None:
|
||||
filter["project_id"] = None
|
||||
if project_id:
|
||||
filter["project_id"] = project_id
|
||||
return await self.collection.count_documents(filter)
|
||||
|
||||
@@ -15,26 +15,24 @@ class CharacterRepo:
|
||||
character.id = str(op.inserted_id)
|
||||
return character
|
||||
|
||||
async def get_character(self, character_id: str, with_image_data: bool = False) -> Character | None:
|
||||
args = {}
|
||||
if not with_image_data:
|
||||
args["character_image_data"] = 0
|
||||
res = await self.collection.find_one({"_id": ObjectId(character_id)}, args)
|
||||
async def get_character(self, character_id: str) -> Character | None:
|
||||
res = await self.collection.find_one({"_id": ObjectId(character_id)})
|
||||
if res is None:
|
||||
return None
|
||||
else:
|
||||
res["id"] = str(res.pop("_id"))
|
||||
return Character(**res)
|
||||
|
||||
async def get_all_characters(self, created_by: Optional[str] = None, project_id: Optional[str] = None) -> List[Character]:
|
||||
async def get_all_characters(self, created_by: Optional[str] = None, project_id: Optional[str] = None, limit: int = 100, offset: int = 0) -> List[Character]:
|
||||
filter = {}
|
||||
if created_by:
|
||||
filter["created_by"] = created_by
|
||||
if project_id is None:
|
||||
filter["project_id"] = None
|
||||
if project_id:
|
||||
filter["project_id"] = project_id
|
||||
|
||||
args = {"character_image_data": 0} # don't return image data for list
|
||||
res = await self.collection.find(filter, args).to_list(None)
|
||||
res = await self.collection.find(filter).skip(offset).limit(limit).to_list(None)
|
||||
chars = []
|
||||
for doc in res:
|
||||
doc["id"] = str(doc.pop("_id"))
|
||||
|
||||
@@ -65,6 +65,8 @@ class GenerationRepo:
|
||||
args["status"] = status
|
||||
if created_by is not None:
|
||||
args["created_by"] = created_by
|
||||
if project_id is None:
|
||||
args["project_id"] = None
|
||||
if project_id is not None:
|
||||
args["project_id"] = project_id
|
||||
if idea_id is not None:
|
||||
@@ -92,6 +94,121 @@ class GenerationRepo:
|
||||
async def update_generation(self, generation: Generation, ):
|
||||
res = await self.collection.update_one({"_id": ObjectId(generation.id)}, {"$set": generation.model_dump()})
|
||||
|
||||
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.
|
||||
"""
|
||||
pipeline = []
|
||||
|
||||
# 1. Match active done generations
|
||||
match_stage = {"is_deleted": False, "status": GenerationStatus.DONE}
|
||||
if created_by:
|
||||
match_stage["created_by"] = created_by
|
||||
if project_id:
|
||||
match_stage["project_id"] = project_id
|
||||
|
||||
pipeline.append({"$match": match_stage})
|
||||
|
||||
# 2. Group by null (total)
|
||||
pipeline.append({
|
||||
"$group": {
|
||||
"_id": None,
|
||||
"total_runs": {"$sum": 1},
|
||||
"total_tokens": {
|
||||
"$sum": {
|
||||
"$cond": [
|
||||
{"$and": [{"$gt": ["$input_token_usage", 0]}, {"$gt": ["$output_token_usage", 0]}]},
|
||||
{"$add": ["$input_token_usage", "$output_token_usage"]},
|
||||
{"$ifNull": ["$token_usage", 0]}
|
||||
]
|
||||
}
|
||||
},
|
||||
"total_input_tokens": {"$sum": {"$ifNull": ["$input_token_usage", 0]}},
|
||||
"total_output_tokens": {"$sum": {"$ifNull": ["$output_token_usage", 0]}},
|
||||
"total_cost": {
|
||||
"$sum": {
|
||||
"$add": [
|
||||
{"$multiply": [{"$ifNull": ["$input_token_usage", 0]}, 0.000002]},
|
||||
{"$multiply": [{"$ifNull": ["$output_token_usage", 0]}, 0.00012]}
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
cursor = self.collection.aggregate(pipeline)
|
||||
res = await cursor.to_list(1)
|
||||
|
||||
if not res:
|
||||
return {
|
||||
"total_runs": 0,
|
||||
"total_tokens": 0,
|
||||
"total_input_tokens": 0,
|
||||
"total_output_tokens": 0,
|
||||
"total_cost": 0.0
|
||||
}
|
||||
|
||||
result = res[0]
|
||||
result.pop("_id")
|
||||
result["total_cost"] = round(result["total_cost"], 4)
|
||||
return result
|
||||
|
||||
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.
|
||||
"""
|
||||
pipeline = []
|
||||
|
||||
match_stage = {"is_deleted": False, "status": GenerationStatus.DONE}
|
||||
if project_id:
|
||||
match_stage["project_id"] = project_id
|
||||
if created_by:
|
||||
match_stage["created_by"] = created_by
|
||||
|
||||
pipeline.append({"$match": match_stage})
|
||||
|
||||
pipeline.append({
|
||||
"$group": {
|
||||
"_id": f"${group_by}",
|
||||
"total_runs": {"$sum": 1},
|
||||
"total_tokens": {
|
||||
"$sum": {
|
||||
"$cond": [
|
||||
{"$and": [{"$gt": ["$input_token_usage", 0]}, {"$gt": ["$output_token_usage", 0]}]},
|
||||
{"$add": ["$input_token_usage", "$output_token_usage"]},
|
||||
{"$ifNull": ["$token_usage", 0]}
|
||||
]
|
||||
}
|
||||
},
|
||||
"total_input_tokens": {"$sum": {"$ifNull": ["$input_token_usage", 0]}},
|
||||
"total_output_tokens": {"$sum": {"$ifNull": ["$output_token_usage", 0]}},
|
||||
"total_cost": {
|
||||
"$sum": {
|
||||
"$add": [
|
||||
{"$multiply": [{"$ifNull": ["$input_token_usage", 0]}, 0.000002]},
|
||||
{"$multiply": [{"$ifNull": ["$output_token_usage", 0]}, 0.00012]}
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
pipeline.append({"$sort": {"total_cost": -1}})
|
||||
|
||||
cursor = self.collection.aggregate(pipeline)
|
||||
res = await cursor.to_list(None)
|
||||
|
||||
results = []
|
||||
for item in res:
|
||||
entity_id = item.pop("_id")
|
||||
item["total_cost"] = round(item["total_cost"], 4)
|
||||
results.append({
|
||||
"entity_id": str(entity_id) if entity_id else "unknown",
|
||||
"stats": item
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
async def get_generations_by_group(self, group_id: str) -> List[Generation]:
|
||||
res = await self.collection.find({"generation_group_id": group_id, "is_deleted": False}).sort("created_at", 1).to_list(None)
|
||||
generations: List[Generation] = []
|
||||
|
||||
@@ -39,8 +39,17 @@ class IdeaRepo:
|
||||
"from": "generations",
|
||||
"let": {"idea_id": "$str_id"},
|
||||
"pipeline": [
|
||||
{"$match": {"$expr": {"$eq": ["$idea_id", "$$idea_id"]}}},
|
||||
{"$sort": {"created_at": -1}}, # Ensure we get the latest
|
||||
{
|
||||
"$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"
|
||||
|
||||
Reference in New Issue
Block a user