fix: llm tests failing

This commit is contained in:
Siddhant Rai
2025-03-19 10:25:46 +05:30
parent ab95d90284
commit f00802dd6b
2 changed files with 10 additions and 16 deletions

View File

@@ -1368,22 +1368,22 @@ class GetMessageAnalytics(Resource):
group_format = "%Y-%m-%d"
try:
match_stage = {
"$match": {
"user": user,
}
}
if api_key:
match_stage["$match"]["api_key"] = api_key
pipeline = [
# Initial match for user and API key if provided
{
"$match": {
"user": user,
"api_key": api_key if api_key else {"$exists": False},
}
},
match_stage,
{"$unwind": "$queries"},
# Match queries within the time range
{
"$match": {
"queries.timestamp": {"$gte": start_date, "$lte": end_date}
}
},
# Group by formatted timestamp
{
"$group": {
"_id": {
@@ -1395,7 +1395,6 @@ class GetMessageAnalytics(Resource):
"count": {"$sum": 1},
}
},
# Sort by timestamp
{"$sort": {"_id": 1}},
]
@@ -1546,8 +1545,6 @@ class GetTokenAnalytics(Resource):
}
if api_key:
match_stage["$match"]["api_key"] = api_key
else:
match_stage["$match"]["api_key"] = {"$exists": False}
token_usage_data = token_usage_collection.aggregate(
[
@@ -1683,10 +1680,7 @@ class GetFeedbackAnalytics(Resource):
}
if api_key:
match_stage["$match"]["api_key"] = api_key
else:
match_stage["$match"]["api_key"] = {"$exists": False}
# Unwind the queries array to process each query separately
pipeline = [
match_stage,
{"$unwind": "$queries"},

View File

@@ -5,7 +5,7 @@ from application.usage import gen_token_usage, stream_token_usage
class BaseLLM(ABC):
def __init__(self, decoded_token):
def __init__(self, decoded_token=None):
self.decoded_token = decoded_token
self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}