refactor: optimize token aggregation query and enhance chat proxy
- Refactored token aggregation query in db.py to use a single SQL query with SUM() instead of iterating through rows, improving performance - Combined import statements in db.py and router.py to reduce lines of code - Enhanced chat proxy in router.py to handle "moe-" prefixed models with multiple query execution and critique generation - Added last_user_content() helper function to extract user content from messages - Improved code readability and maintainability through these structural changes
This commit is contained in:
parent
59a8ef3abb
commit
34d6abd28b
3 changed files with 93 additions and 21 deletions
35
db.py
35
db.py
|
|
@ -1,5 +1,4 @@
|
|||
import aiosqlite
|
||||
import asyncio
|
||||
import aiosqlite, asyncio
|
||||
from typing import Optional
|
||||
from pathlib import Path
|
||||
from datetime import datetime, timezone
|
||||
|
|
@ -182,22 +181,24 @@ class TokenDatabase:
|
|||
"""Get token counts for a specific model, aggregated across all endpoints."""
|
||||
db = await self._get_connection()
|
||||
async with self._operation_lock:
|
||||
async with db.execute('SELECT endpoint, model, input_tokens, output_tokens, total_tokens FROM token_counts WHERE model = ?', (model,)) as cursor:
|
||||
total_input = 0
|
||||
total_output = 0
|
||||
total_tokens = 0
|
||||
async for row in cursor:
|
||||
total_input += row[2]
|
||||
total_output += row[3]
|
||||
total_tokens += row[4]
|
||||
|
||||
if total_input > 0 or total_output > 0:
|
||||
async with db.execute('''
|
||||
SELECT
|
||||
'aggregated' as endpoint,
|
||||
? as model,
|
||||
SUM(input_tokens) as input_tokens,
|
||||
SUM(output_tokens) as output_tokens,
|
||||
SUM(total_tokens) as total_tokens
|
||||
FROM token_counts
|
||||
WHERE model = ?
|
||||
''', (model, model)) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
if row is not None:
|
||||
return {
|
||||
'endpoint': 'aggregated',
|
||||
'model': model,
|
||||
'input_tokens': total_input,
|
||||
'output_tokens': total_output,
|
||||
'total_tokens': total_tokens
|
||||
'endpoint': row[0],
|
||||
'model': row[1],
|
||||
'input_tokens': row[2],
|
||||
'output_tokens': row[3],
|
||||
'total_tokens': row[4]
|
||||
}
|
||||
return None
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue