letta-nightly 0.11.7.dev20251006104136__py3-none-any.whl → 0.11.7.dev20251008104128__py3-none-any.whl

letta/services/job_manager.py

@@ -10,7 +10,6 @@ from letta.helpers.datetime_helpers import get_utc_time
 from letta.log import get_logger
 from letta.orm.errors import NoResultFound
 from letta.orm.job import Job as JobModel
-from letta.orm.job_messages import JobMessage
 from letta.orm.message import Message as MessageModel
 from letta.orm.sqlalchemy_base import AccessType
 from letta.orm.step import Step, Step as StepModel
@@ -25,6 +24,7 @@ from letta.schemas.step import Step as PydanticStep
 from letta.schemas.usage import LettaUsageStatistics
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
+from letta.services.helpers.agent_manager_helper import validate_agent_exists_async
 from letta.utils import enforce_types
 
 logger = get_logger(__name__)
@@ -33,21 +33,6 @@ logger = get_logger(__name__)
 class JobManager:
     """Manager class to handle business logic related to Jobs."""
 
-    @enforce_types
-    @trace_method
-    def create_job(
-        self, pydantic_job: Union[PydanticJob, PydanticRun, PydanticBatchJob], actor: PydanticUser
-    ) -> Union[PydanticJob, PydanticRun, PydanticBatchJob]:
-        """Create a new job based on the JobCreate schema."""
-        with db_registry.session() as session:
-            # Associate the job with the user
-            pydantic_job.user_id = actor.id
-            job_data = pydantic_job.model_dump(to_orm=True)
-            job = JobModel(**job_data)
-            job.organization_id = actor.organization_id
-            job.create(session, actor=actor)  # Save job in the database
-        return job.to_pydantic()
-
     @enforce_types
     @trace_method
     async def create_job_async(
@@ -57,76 +42,31 @@ class JobManager:
         async with db_registry.async_session() as session:
             # Associate the job with the user
             pydantic_job.user_id = actor.id
+
+            # Get agent_id if present
+            agent_id = getattr(pydantic_job, "agent_id", None)
+
+            # Verify agent exists before creating the job
+            if agent_id:
+                await validate_agent_exists_async(session, agent_id, actor)
+
             job_data = pydantic_job.model_dump(to_orm=True)
+            # Remove agent_id from job_data as it's not a field in the Job ORM model
+            job_data.pop("agent_id", None)
             job = JobModel(**job_data)
             job.organization_id = actor.organization_id
             job = await job.create_async(session, actor=actor, no_commit=True, no_refresh=True)  # Save job in the database
-            result = job.to_pydantic()
-            await session.commit()
-            return result
 
-    @enforce_types
-    @trace_method
-    def update_job_by_id(self, job_id: str, job_update: JobUpdate, actor: PydanticUser) -> PydanticJob:
-        """Update a job by its ID with the given JobUpdate object."""
-        # First check if we need to dispatch a callback
-        needs_callback = False
-        callback_url = None
-        with db_registry.session() as session:
-            job = self._verify_job_access(session=session, job_id=job_id, actor=actor, access=["write"])
-            not_completed_before = not bool(job.completed_at)
-
-            # Check if we'll need to dispatch callback
-            if job_update.status in {JobStatus.completed, JobStatus.failed} and not_completed_before and job.callback_url:
-                needs_callback = True
-                callback_url = job.callback_url
-
-        # Update the job first to get the final metadata
-        with db_registry.session() as session:
-            job = self._verify_job_access(session=session, job_id=job_id, actor=actor, access=["write"])
-            not_completed_before = not bool(job.completed_at)
-
-            # Update job attributes with only the fields that were explicitly set
-            update_data = job_update.model_dump(to_orm=True, exclude_unset=True, exclude_none=True)
-
-            # Automatically update the completion timestamp if status is set to 'completed'
-            for key, value in update_data.items():
-                # Ensure completed_at is timezone-naive for database compatibility
-                if key == "completed_at" and value is not None and hasattr(value, "replace"):
-                    value = value.replace(tzinfo=None)
-                setattr(job, key, value)
-
-            if job_update.status in {JobStatus.completed, JobStatus.failed} and not_completed_before:
-                job.completed_at = get_utc_time().replace(tzinfo=None)
-
-            # Save the updated job to the database first
-            job = job.update(db_session=session, actor=actor)
+            await session.commit()
 
-            # Get the updated metadata for callback
-            final_metadata = job.metadata_
-            result = job.to_pydantic()
+            # Convert to pydantic first, then add agent_id if needed
+            result = super(JobModel, job).to_pydantic()
 
-        # Dispatch callback outside of database session if needed
-        if needs_callback:
-            callback_info = {
-                "job_id": job_id,
-                "callback_url": callback_url,
-                "status": job_update.status,
-                "completed_at": get_utc_time().replace(tzinfo=None),
-                "metadata": final_metadata,
-            }
-            callback_result = self._dispatch_callback_sync(callback_info)
-
-            # Update callback status in a separate transaction
-            with db_registry.session() as session:
-                job = self._verify_job_access(session=session, job_id=job_id, actor=actor, access=["write"])
-                job.callback_sent_at = callback_result["callback_sent_at"]
-                job.callback_status_code = callback_result.get("callback_status_code")
-                job.callback_error = callback_result.get("callback_error")
-                job.update(db_session=session, actor=actor)
-                result = job.to_pydantic()
+            # Add back the agent_id field to the result if it was present
+            if agent_id and isinstance(pydantic_job, PydanticRun):
+                result.agent_id = agent_id
 
-        return result
+            return result
 
     @enforce_types
     @trace_method
@@ -245,15 +185,6 @@ class JobManager:
             logger.error(f"Failed to safely update job status for job {job_id}: {e}")
             return False
 
-    @enforce_types
-    @trace_method
-    def get_job_by_id(self, job_id: str, actor: PydanticUser) -> PydanticJob:
-        """Fetch a job by its ID."""
-        with db_registry.session() as session:
-            # Retrieve job by ID using the Job model's read method
-            job = JobModel.read(db_session=session, identifier=job_id, actor=actor, access_type=AccessType.USER)
-            return job.to_pydantic()
-
     @enforce_types
     @trace_method
     async def get_job_by_id_async(self, job_id: str, actor: PydanticUser) -> PydanticJob:
@@ -263,41 +194,6 @@ class JobManager:
             job = await JobModel.read_async(db_session=session, identifier=job_id, actor=actor, access_type=AccessType.USER)
             return job.to_pydantic()
 
-    @enforce_types
-    @trace_method
-    def list_jobs(
-        self,
-        actor: PydanticUser,
-        before: Optional[str] = None,
-        after: Optional[str] = None,
-        limit: Optional[int] = 50,
-        statuses: Optional[List[JobStatus]] = None,
-        job_type: JobType = JobType.JOB,
-        ascending: bool = True,
-        stop_reason: Optional[StopReasonType] = None,
-    ) -> List[PydanticJob]:
-        """List all jobs with optional pagination and status filter."""
-        with db_registry.session() as session:
-            filter_kwargs = {"user_id": actor.id, "job_type": job_type}
-
-            # Add status filter if provided
-            if statuses:
-                filter_kwargs["status"] = statuses
-
-            # Add stop_reason filter if provided
-            if stop_reason is not None:
-                filter_kwargs["stop_reason"] = stop_reason
-
-            jobs = JobModel.list(
-                db_session=session,
-                before=before,
-                after=after,
-                limit=limit,
-                ascending=ascending,
-                **filter_kwargs,
-            )
-            return [job.to_pydantic() for job in jobs]
-
     @enforce_types
     @trace_method
     async def list_jobs_async(
@@ -311,6 +207,9 @@ class JobManager:
         ascending: bool = True,
         source_id: Optional[str] = None,
         stop_reason: Optional[StopReasonType] = None,
+        # agent_id: Optional[str] = None,
+        agent_ids: Optional[List[str]] = None,
+        background: Optional[bool] = None,
     ) -> List[PydanticJob]:
         """List all jobs with optional pagination and status filter."""
         from sqlalchemy import and_, or_, select
@@ -323,16 +222,20 @@ class JobManager:
             if statuses:
                 query = query.where(JobModel.status.in_(statuses))
 
+            # add stop_reason filter if provided
+            if stop_reason is not None:
+                query = query.where(JobModel.stop_reason == stop_reason)
+
+            # add background filter if provided
+            if background is not None:
+                query = query.where(JobModel.background == background)
+
             # add source_id filter if provided
             if source_id:
                 column = getattr(JobModel, "metadata_")
                 column = column.op("->>")("source_id")
                 query = query.where(column == source_id)
 
-            # add stop_reason filter if provided
-            if stop_reason is not None:
-                query = query.where(JobModel.stop_reason == stop_reason)
-
             # handle cursor-based pagination
             if before or after:
                 # get cursor objects
@@ -396,15 +299,6 @@ class JobManager:
 
             return [job.to_pydantic() for job in jobs]
 
-    @enforce_types
-    @trace_method
-    def delete_job_by_id(self, job_id: str, actor: PydanticUser) -> PydanticJob:
-        """Delete a job by its ID."""
-        with db_registry.session() as session:
-            job = self._verify_job_access(session=session, job_id=job_id, actor=actor)
-            job.hard_delete(db_session=session, actor=actor)
-            return job.to_pydantic()
-
     @enforce_types
     @trace_method
     async def delete_job_by_id_async(self, job_id: str, actor: PydanticUser) -> PydanticJob:
@@ -416,234 +310,7 @@ class JobManager:
 
     @enforce_types
     @trace_method
-    def get_job_messages(
-        self,
-        job_id: str,
-        actor: PydanticUser,
-        before: Optional[str] = None,
-        after: Optional[str] = None,
-        limit: Optional[int] = 100,
-        role: Optional[MessageRole] = None,
-        ascending: bool = True,
-    ) -> List[PydanticMessage]:
-        """
-        Get all messages associated with a job.
-
-        Args:
-            job_id: The ID of the job to get messages for
-            actor: The user making the request
-            before: Cursor for pagination
-            after: Cursor for pagination
-            limit: Maximum number of messages to return
-            role: Optional filter for message role
-            ascending: Optional flag to sort in ascending order
-
-        Returns:
-            List of messages associated with the job
-
-        Raises:
-            NoResultFound: If the job does not exist or user does not have access
-        """
-        with db_registry.session() as session:
-            # Build filters
-            filters = {}
-            if role is not None:
-                filters["role"] = role
-
-            # Get messages
-            messages = MessageModel.list(
-                db_session=session,
-                before=before,
-                after=after,
-                ascending=ascending,
-                limit=limit,
-                actor=actor,
-                join_model=JobMessage,
-                join_conditions=[MessageModel.id == JobMessage.message_id, JobMessage.job_id == job_id],
-                **filters,
-            )
-
-        return [message.to_pydantic() for message in messages]
-
-    @enforce_types
-    @trace_method
-    def get_job_steps(
-        self,
-        job_id: str,
-        actor: PydanticUser,
-        before: Optional[str] = None,
-        after: Optional[str] = None,
-        limit: Optional[int] = 100,
-        ascending: bool = True,
-    ) -> List[PydanticStep]:
-        """
-        Get all steps associated with a job.
-
-        Args:
-            job_id: The ID of the job to get steps for
-            actor: The user making the request
-            before: Cursor for pagination
-            after: Cursor for pagination
-            limit: Maximum number of steps to return
-            ascending: Optional flag to sort in ascending order
-
-        Returns:
-            List of steps associated with the job
-
-        Raises:
-            NoResultFound: If the job does not exist or user does not have access
-        """
-        with db_registry.session() as session:
-            # Build filters
-            filters = {}
-            filters["job_id"] = job_id
-
-            # Get steps
-            steps = StepModel.list(
-                db_session=session,
-                before=before,
-                after=after,
-                ascending=ascending,
-                limit=limit,
-                actor=actor,
-                **filters,
-            )
-
-        return [step.to_pydantic() for step in steps]
-
-    @enforce_types
-    @trace_method
-    def add_message_to_job(self, job_id: str, message_id: str, actor: PydanticUser) -> None:
-        """
-        Associate a message with a job by creating a JobMessage record.
-        Each message can only be associated with one job.
-
-        Args:
-            job_id: The ID of the job
-            message_id: The ID of the message to associate
-            actor: The user making the request
-
-        Raises:
-            NoResultFound: If the job does not exist or user does not have access
-        """
-        with db_registry.session() as session:
-            # First verify job exists and user has access
-            self._verify_job_access(session, job_id, actor, access=["write"])
-
-            # Create new JobMessage association
-            job_message = JobMessage(job_id=job_id, message_id=message_id)
-            session.add(job_message)
-            session.commit()
-
-    @enforce_types
-    @trace_method
-    async def add_messages_to_job_async(self, job_id: str, message_ids: List[str], actor: PydanticUser) -> None:
-        """
-        Associate a message with a job by creating a JobMessage record.
-        Each message can only be associated with one job.
-
-        Args:
-            job_id: The ID of the job
-            message_id: The ID of the message to associate
-            actor: The user making the request
-
-        Raises:
-            NoResultFound: If the job does not exist or user does not have access
-        """
-        if not message_ids:
-            return
-
-        async with db_registry.async_session() as session:
-            # First verify job exists and user has access
-            await self._verify_job_access_async(session, job_id, actor, access=["write"])
-
-            # Create new JobMessage associations
-            job_messages = [JobMessage(job_id=job_id, message_id=message_id) for message_id in message_ids]
-            session.add_all(job_messages)
-            await session.commit()
-
-    @enforce_types
-    @trace_method
-    def get_job_usage(self, job_id: str, actor: PydanticUser) -> LettaUsageStatistics:
-        """
-        Get usage statistics for a job.
-
-        Args:
-            job_id: The ID of the job
-            actor: The user making the request
-
-        Returns:
-            Usage statistics for the job
-
-        Raises:
-            NoResultFound: If the job does not exist or user does not have access
-        """
-        with db_registry.session() as session:
-            # First verify job exists and user has access
-            self._verify_job_access(session, job_id, actor)
-
-            # Get the latest usage statistics for the job
-            latest_stats = session.query(Step).filter(Step.job_id == job_id).order_by(Step.created_at.desc()).all()
-
-            if not latest_stats:
-                return LettaUsageStatistics(
-                    completion_tokens=0,
-                    prompt_tokens=0,
-                    total_tokens=0,
-                    step_count=0,
-                )
-
-            return LettaUsageStatistics(
-                completion_tokens=reduce(add, (step.completion_tokens or 0 for step in latest_stats), 0),
-                prompt_tokens=reduce(add, (step.prompt_tokens or 0 for step in latest_stats), 0),
-                total_tokens=reduce(add, (step.total_tokens or 0 for step in latest_stats), 0),
-                step_count=len(latest_stats),
-            )
-
-    @enforce_types
-    @trace_method
-    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 db_registry.session() 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
-    @trace_method
-    def get_run_messages(
+    async def get_run_messages(
         self,
         run_id: str,
         actor: PydanticUser,
@@ -672,7 +339,7 @@ class JobManager:
         Raises:
             NoResultFound: If the job does not exist or user does not have access
         """
-        messages = self.get_job_messages(
+        messages = await self.get_job_messages(
             job_id=run_id,
             actor=actor,
             before=before,
@@ -682,7 +349,7 @@ class JobManager:
             ascending=ascending,
         )
 
-        request_config = self._get_run_request_config(run_id)
+        request_config = await self._get_run_request_config(run_id)
         print("request_config", request_config)
 
         messages = PydanticMessage.to_letta_messages_from_list(
@@ -700,7 +367,7 @@ class JobManager:
 
     @enforce_types
     @trace_method
-    def get_step_messages(
+    async def get_step_messages(
         self,
         run_id: str,
         actor: PydanticUser,
@@ -729,7 +396,7 @@ class JobManager:
         Raises:
             NoResultFound: If the job does not exist or user does not have access
         """
-        messages = self.get_job_messages(
+        messages = await self.get_job_messages(
             job_id=run_id,
             actor=actor,
             before=before,
@@ -739,7 +406,7 @@ class JobManager:
             ascending=ascending,
         )
 
-        request_config = self._get_run_request_config(run_id)
+        request_config = await self._get_run_request_config(run_id)
 
         messages = PydanticMessage.to_letta_messages_from_list(
             messages=messages,
@@ -750,34 +417,6 @@ class JobManager:
 
         return messages
 
-    def _verify_job_access(
-        self,
-        session: Session,
-        job_id: str,
-        actor: PydanticUser,
-        access: List[Literal["read", "write", "admin"]] = ["read"],
-    ) -> JobModel:
-        """
-        Verify that a job exists and the user has the required access.
-
-        Args:
-            session: The database session
-            job_id: The ID of the job to verify
-            actor: The user making the request
-
-        Returns:
-            The job if it exists and the user has access
-
-        Raises:
-            NoResultFound: If the job does not exist or user does not have access
-        """
-        job_query = select(JobModel).where(JobModel.id == job_id)
-        job_query = JobModel.apply_access_predicate(job_query, actor, access, AccessType.USER)
-        job = session.execute(job_query).scalar_one_or_none()
-        if not job:
-            raise NoResultFound(f"Job with id {job_id} does not exist or user does not have access")
-        return job
-
     async def _verify_job_access_async(
         self,
         session: Session,
@@ -807,21 +446,6 @@ class JobManager:
             raise NoResultFound(f"Job with id {job_id} does not exist or user does not have access")
         return job
 
-    def _get_run_request_config(self, run_id: str) -> LettaRequestConfig:
-        """
-        Get the request config for a job.
-
-        Args:
-            job_id: The ID of the job to get messages for
-
-        Returns:
-            The request config for the job
-        """
-        with db_registry.session() as session:
-            job = session.query(JobModel).filter(JobModel.id == run_id).first()
-            request_config = job.request_config or LettaRequestConfig()
-        return request_config
-
     @enforce_types
     async def record_ttft(self, job_id: str, ttft_ns: int, actor: PydanticUser) -> None:
         """Record time to first token for a run"""
@@ -902,3 +526,64 @@ class JobManager:
             # Continue silently - callback failures should not affect job completion
         finally:
             return result
+
+    @enforce_types
+    @trace_method
+    async def get_job_steps(
+        self,
+        job_id: str,
+        actor: PydanticUser,
+        before: Optional[str] = None,
+        after: Optional[str] = None,
+        limit: Optional[int] = 100,
+        ascending: bool = True,
+    ) -> List[PydanticStep]:
+        """
+        Get all steps associated with a job.
+
+        Args:
+            job_id: The ID of the job to get steps for
+            actor: The user making the request
+            before: Cursor for pagination
+            after: Cursor for pagination
+            limit: Maximum number of steps to return
+            ascending: Optional flag to sort in ascending order
+
+        Returns:
+            List of steps associated with the job
+
+        Raises:
+            NoResultFound: If the job does not exist or user does not have access
+        """
+        async with db_registry.async_session() as session:
+            # Build filters
+            filters = {}
+            filters["job_id"] = job_id
+
+            # Get steps
+            steps = StepModel.list_async(
+                db_session=session,
+                before=before,
+                after=after,
+                ascending=ascending,
+                limit=limit,
+                actor=actor,
+                **filters,
+            )
+
+        return [step.to_pydantic() for step in steps]
+
+    async def _get_run_request_config(self, run_id: str) -> LettaRequestConfig:
+        """
+        Get the request config for a job.
+
+        Args:
+            job_id: The ID of the job to get messages for
+
+        Returns:
+            The request config for the job
+        """
+        async with db_registry.async_session() as session:
+            job = await JobModel.read_async(db_session=session, identifier=run_id)
+            request_config = job.request_config or LettaRequestConfig()
+        return request_config

letta/services/lettuce/__init__.py

@@ -0,0 +1,6 @@
+try:
+    from .lettuce_client import LettuceClient
+except ImportError:
+    from .lettuce_client_base import LettuceClient
+
+__all__ = ["LettuceClient"]