fix: remove manual step logging (#738)
This commit is contained in:
@@ -12,7 +12,7 @@ from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.middleware.cors import CORSMiddleware
|
||||
|
||||
from letta.__init__ import __version__
|
||||
from letta.constants import ADMIN_PREFIX, API_PREFIX, OPENAI_API_PREFIX
|
||||
from letta.constants import ADMIN_PREFIX, API_PREFIX
|
||||
from letta.errors import BedrockPermissionError, LettaAgentNotFoundError, LettaUserNotFoundError
|
||||
from letta.log import get_logger
|
||||
from letta.orm.errors import DatabaseTimeoutError, ForeignKeyConstraintViolationError, NoResultFound, UniqueConstraintViolationError
|
||||
|
||||
@@ -560,9 +560,6 @@ async def process_message_background(
|
||||
)
|
||||
server.job_manager.update_job_by_id(job_id=job_id, job_update=job_update, actor=actor)
|
||||
|
||||
# Add job usage statistics
|
||||
server.job_manager.add_job_usage(job_id=job_id, usage=result.usage, actor=actor)
|
||||
|
||||
except Exception as e:
|
||||
# Update job status to failed
|
||||
job_update = JobUpdate(
|
||||
|
||||
@@ -214,46 +214,6 @@ class JobManager:
|
||||
step_count=len(latest_stats),
|
||||
)
|
||||
|
||||
@enforce_types
|
||||
def add_job_usage(
|
||||
self,
|
||||
job_id: str,
|
||||
usage: LettaUsageStatistics,
|
||||
step_id: Optional[str] = None,
|
||||
actor: PydanticUser = None,
|
||||
) -> None:
|
||||
"""
|
||||
Add usage statistics for a job.
|
||||
|
||||
Args:
|
||||
job_id: The ID of the job
|
||||
usage: Usage statistics for the job
|
||||
step_id: Optional ID of the specific step within the job
|
||||
actor: The user making the request
|
||||
|
||||
Raises:
|
||||
NoResultFound: If the job does not exist or user does not have access
|
||||
"""
|
||||
with self.session_maker() as session:
|
||||
# First verify job exists and user has access
|
||||
self._verify_job_access(session, job_id, actor, access=["write"])
|
||||
|
||||
# Manually log step with usage data
|
||||
# TODO(@caren): log step under the hood and remove this
|
||||
usage_stats = Step(
|
||||
job_id=job_id,
|
||||
completion_tokens=usage.completion_tokens,
|
||||
prompt_tokens=usage.prompt_tokens,
|
||||
total_tokens=usage.total_tokens,
|
||||
step_count=usage.step_count,
|
||||
step_id=step_id,
|
||||
)
|
||||
if actor:
|
||||
usage_stats._set_created_and_updated_by_fields(actor.id)
|
||||
|
||||
session.add(usage_stats)
|
||||
session.commit()
|
||||
|
||||
@enforce_types
|
||||
def get_run_messages_cursor(
|
||||
self,
|
||||
|
||||
Reference in New Issue
Block a user