letta-nightly 0.6.0.dev20241204232133__tar.gz → 0.6.1.dev20241205211219__tar.gz

{letta_nightly-0.6.0.dev20241204232133 → letta_nightly-0.6.1.dev20241205211219}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: letta-nightly
-Version: 0.6.0.dev20241204232133
+Version: 0.6.1.dev20241205211219
 Summary: Create LLM agents with long-term memory and custom tools
 License: Apache License
 Author: Letta Team

{letta_nightly-0.6.0.dev20241204232133 → letta_nightly-0.6.1.dev20241205211219}/letta/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.6.0"
+__version__ = "0.6.1"
 
 # import clients
 from letta.client.client import LocalClient, RESTClient, create_client

{letta_nightly-0.6.0.dev20241204232133 → letta_nightly-0.6.1.dev20241205211219}/letta/agent.py

@@ -1253,9 +1253,23 @@ class Agent(BaseAgent):
         self._messages = new_messages
 
     def rebuild_system_prompt(self, force=False, update_timestamp=True):
-        """Rebuilds the system message with the latest memory object and any shared memory block updates"""
+        """Rebuilds the system message with the latest memory object and any shared memory block updates
+
+        Updates to core memory blocks should trigger a "rebuild", which itself will create a new message object
+
+        Updates to the memory header should *not* trigger a rebuild, since that will simply flood recall storage with excess messages
+        """
+
         curr_system_message = self.messages[0]  # this is the system + memory bank, not just the system prompt
 
+        # note: we only update the system prompt if the core memory is changed
+        # this means that the archival/recall memory statistics may be someout out of date
+        curr_memory_str = self.agent_state.memory.compile()
+        if curr_memory_str in curr_system_message["content"] and not force:
+            # NOTE: could this cause issues if a block is removed? (substring match would still work)
+            printd(f"Memory hasn't changed, skipping system prompt rebuild")
+            return
+
         # If the memory didn't update, we probably don't want to update the timestamp inside
         # For example, if we're doing a system prompt swap, this should probably be False
         if update_timestamp:

{letta_nightly-0.6.0.dev20241204232133 → letta_nightly-0.6.1.dev20241205211219}/letta/client/client.py

@@ -2859,8 +2859,12 @@ class LocalClient(AbstractClient):
         Returns:
             job (Job): Data loading job including job status and metadata
         """
-        metadata_ = {"type": "embedding", "filename": filename, "source_id": source_id}
-        job = self.server.create_job(user_id=self.user_id, metadata=metadata_)
+        job = Job(
+            user_id=self.user_id,
+            status=JobStatus.created,
+            metadata_={"type": "embedding", "filename": filename, "source_id": source_id},
+        )
+        job = self.server.job_manager.create_job(pydantic_job=job, actor=self.user)
 
         # TODO: implement blocking vs. non-blocking
         self.server.load_file_to_source(source_id=source_id, file_path=filename, job_id=job.id)
@@ -2870,16 +2874,16 @@ class LocalClient(AbstractClient):
         self.server.source_manager.delete_file(file_id, actor=self.user)
 
     def get_job(self, job_id: str):
-        return self.server.get_job(job_id=job_id)
+        return self.server.job_manager.get_job_by_id(job_id=job_id, actor=self.user)
 
     def delete_job(self, job_id: str):
-        return self.server.delete_job(job_id)
+        return self.server.job_manager.delete_job(job_id=job_id, actor=self.user)
 
     def list_jobs(self):
-        return self.server.list_jobs(user_id=self.user_id)
+        return self.server.job_manager.list_jobs(actor=self.user)
 
     def list_active_jobs(self):
-        return self.server.list_active_jobs(user_id=self.user_id)
+        return self.server.job_manager.list_jobs(actor=self.user, statuses=[JobStatus.created, JobStatus.running])
 
     def create_source(self, name: str, embedding_config: Optional[EmbeddingConfig] = None) -> Source:
         """

{letta_nightly-0.6.0.dev20241204232133 → letta_nightly-0.6.1.dev20241205211219}/letta/metadata.py

@@ -12,15 +12,14 @@ from letta.orm.base import Base
 from letta.schemas.agent import PersistedAgentState
 from letta.schemas.api_key import APIKey
 from letta.schemas.embedding_config import EmbeddingConfig
-from letta.schemas.enums import JobStatus, ToolRuleType
-from letta.schemas.job import Job
+from letta.schemas.enums import ToolRuleType
 from letta.schemas.llm_config import LLMConfig
 from letta.schemas.openai.chat_completions import ToolCall, ToolCallFunction
 from letta.schemas.tool_rule import ChildToolRule, InitToolRule, TerminalToolRule
 from letta.schemas.user import User
 from letta.services.per_agent_lock_manager import PerAgentLockManager
 from letta.settings import settings
-from letta.utils import enforce_types, get_utc_time, printd
+from letta.utils import enforce_types, printd
 
 
 class LLMConfigColumn(TypeDecorator):
@@ -258,31 +257,6 @@ class AgentSourceMappingModel(Base):
         return f"<AgentSourceMapping(user_id='{self.user_id}', agent_id='{self.agent_id}', source_id='{self.source_id}')>"
 
 
-class JobModel(Base):
-    __tablename__ = "jobs"
-    __table_args__ = {"extend_existing": True}
-
-    id = Column(String, primary_key=True)
-    user_id = Column(String)
-    status = Column(String, default=JobStatus.pending)
-    created_at = Column(DateTime(timezone=True), server_default=func.now())
-    completed_at = Column(DateTime(timezone=True), onupdate=func.now())
-    metadata_ = Column(JSON)
-
-    def __repr__(self) -> str:
-        return f"<Job(id='{self.id}', status='{self.status}')>"
-
-    def to_record(self):
-        return Job(
-            id=self.id,
-            user_id=self.user_id,
-            status=self.status,
-            created_at=self.created_at,
-            completed_at=self.completed_at,
-            metadata_=self.metadata_,
-        )
-
-
 class MetadataStore:
     uri: Optional[str] = None
 
@@ -455,40 +429,3 @@ class MetadataStore:
                 AgentSourceMappingModel.agent_id == agent_id, AgentSourceMappingModel.source_id == source_id
             ).delete()
             session.commit()
-
-    @enforce_types
-    def create_job(self, job: Job):
-        with self.session_maker() as session:
-            session.add(JobModel(**vars(job)))
-            session.commit()
-
-    def delete_job(self, job_id: str):
-        with self.session_maker() as session:
-            session.query(JobModel).filter(JobModel.id == job_id).delete()
-            session.commit()
-
-    def get_job(self, job_id: str) -> Optional[Job]:
-        with self.session_maker() as session:
-            results = session.query(JobModel).filter(JobModel.id == job_id).all()
-            if len(results) == 0:
-                return None
-            assert len(results) == 1, f"Expected 1 result, got {len(results)}"
-            return results[0].to_record()
-
-    def list_jobs(self, user_id: str) -> List[Job]:
-        with self.session_maker() as session:
-            results = session.query(JobModel).filter(JobModel.user_id == user_id).all()
-            return [r.to_record() for r in results]
-
-    def update_job(self, job: Job) -> Job:
-        with self.session_maker() as session:
-            session.query(JobModel).filter(JobModel.id == job.id).update(vars(job))
-            session.commit()
-        return Job
-
-    def update_job_status(self, job_id: str, status: JobStatus):
-        with self.session_maker() as session:
-            session.query(JobModel).filter(JobModel.id == job_id).update({"status": status})
-            if status == JobStatus.COMPLETED:
-                session.query(JobModel).filter(JobModel.id == job_id).update({"completed_at": get_utc_time()})
-            session.commit()

{letta_nightly-0.6.0.dev20241204232133 → letta_nightly-0.6.1.dev20241205211219}/letta/orm/__init__.py

@@ -2,6 +2,7 @@ from letta.orm.base import Base
 from letta.orm.block import Block
 from letta.orm.blocks_agents import BlocksAgents
 from letta.orm.file import FileMetadata
+from letta.orm.job import Job
 from letta.orm.organization import Organization
 from letta.orm.sandbox_config import SandboxConfig, SandboxEnvironmentVariable
 from letta.orm.source import Source

letta_nightly-0.6.1.dev20241205211219/letta/orm/job.py

@@ -0,0 +1,29 @@
+from datetime import datetime
+from typing import TYPE_CHECKING, Optional
+
+from sqlalchemy import JSON, String
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from letta.orm.mixins import UserMixin
+from letta.orm.sqlalchemy_base import SqlalchemyBase
+from letta.schemas.enums import JobStatus
+from letta.schemas.job import Job as PydanticJob
+
+if TYPE_CHECKING:
+    from letta.orm.user import User
+
+
+class Job(SqlalchemyBase, UserMixin):
+    """Jobs run in the background and are owned by a user.
+    Typical jobs involve loading and processing sources etc.
+    """
+
+    __tablename__ = "jobs"
+    __pydantic_model__ = PydanticJob
+
+    status: Mapped[JobStatus] = mapped_column(String, default=JobStatus.created, doc="The current status of the job.")
+    completed_at: Mapped[Optional[datetime]] = mapped_column(nullable=True, doc="The unix timestamp of when the job was completed.")
+    metadata_: Mapped[Optional[dict]] = mapped_column(JSON, default=lambda: {}, doc="The metadata of the job.")
+
+    # relationships
+    user: Mapped["User"] = relationship("User", back_populates="jobs")

{letta_nightly-0.6.0.dev20241204232133 → letta_nightly-0.6.1.dev20241205211219}/letta/orm/sqlalchemy_base.py

@@ -31,24 +31,43 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
     def list(
         cls, *, db_session: "Session", cursor: Optional[str] = None, limit: Optional[int] = 50, **kwargs
     ) -> List[Type["SqlalchemyBase"]]:
-        """List records with optional cursor (for pagination) and limit."""
-        logger.debug(f"Listing {cls.__name__} with kwarg filters {kwargs}")
-        with db_session as session:
-            # Start with the base query filtered by kwargs
-            query = select(cls).filter_by(**kwargs)
+        """
+        List records with optional cursor (for pagination), limit, and automatic filtering.
+
+        Args:
+            db_session: The database session to use.
+            cursor: Optional ID to start pagination from.
+            limit: Maximum number of records to return.
+            **kwargs: Filters passed as equality conditions or iterable for IN filtering.
 
-            # Add a cursor condition if provided
+        Returns:
+            A list of model instances matching the filters.
+        """
+        logger.debug(f"Listing {cls.__name__} with filters {kwargs}")
+        with db_session as session:
+            # Start with a base query
+            query = select(cls)
+
+            # Apply filtering logic
+            for key, value in kwargs.items():
+                column = getattr(cls, key)
+                if isinstance(value, (list, tuple, set)):  # Check for iterables
+                    query = query.where(column.in_(value))
+                else:  # Single value for equality filtering
+                    query = query.where(column == value)
+
+            # Apply cursor for pagination
             if cursor:
                 query = query.where(cls.id > cursor)
 
-            # Add a limit to the query if provided
-            query = query.order_by(cls.id).limit(limit)
-
             # Handle soft deletes if the class has the 'is_deleted' attribute
             if hasattr(cls, "is_deleted"):
                 query = query.where(cls.is_deleted == False)
 
-            # Execute the query and return the results as a list of model instances
+            # Add ordering and limit
+            query = query.order_by(cls.id).limit(limit)
+
+            # Execute the query and return results as model instances
             return list(session.execute(query).scalars())
 
     @classmethod

{letta_nightly-0.6.0.dev20241204232133 → letta_nightly-0.6.1.dev20241205211219}/letta/orm/user.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, List
 
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
@@ -7,7 +7,7 @@ from letta.orm.sqlalchemy_base import SqlalchemyBase
 from letta.schemas.user import User as PydanticUser
 
 if TYPE_CHECKING:
-    from letta.orm.organization import Organization
+    from letta.orm import Job, Organization
 
 
 class User(SqlalchemyBase, OrganizationMixin):
@@ -20,10 +20,10 @@ class User(SqlalchemyBase, OrganizationMixin):
 
     # relationships
     organization: Mapped["Organization"] = relationship("Organization", back_populates="users")
+    jobs: Mapped[List["Job"]] = relationship("Job", back_populates="user", doc="the jobs associated with this user.")
 
     # TODO: Add this back later potentially
     # agents: Mapped[List["Agent"]] = relationship(
     #     "Agent", secondary="users_agents", back_populates="users", doc="the agents associated with this user."
     # )
     # tokens: Mapped[List["Token"]] = relationship("Token", back_populates="user", doc="the tokens associated with this user.")
-    # jobs: Mapped[List["Job"]] = relationship("Job", back_populates="user", doc="the jobs associated with this user.")

{letta_nightly-0.6.0.dev20241204232133 → letta_nightly-0.6.1.dev20241205211219}/letta/schemas/job.py

@@ -4,12 +4,13 @@ from typing import Optional
 from pydantic import Field
 
 from letta.schemas.enums import JobStatus
-from letta.schemas.letta_base import LettaBase
-from letta.utils import get_utc_time
+from letta.schemas.letta_base import OrmMetadataBase
 
 
-class JobBase(LettaBase):
+class JobBase(OrmMetadataBase):
     __id_prefix__ = "job"
+    status: JobStatus = Field(default=JobStatus.created, description="The status of the job.")
+    completed_at: Optional[datetime] = Field(None, description="The unix timestamp of when the job was completed.")
     metadata_: Optional[dict] = Field(None, description="The metadata of the job.")
 
 
@@ -27,12 +28,11 @@ class Job(JobBase):
     """
 
     id: str = JobBase.generate_id_field()
-    status: JobStatus = Field(default=JobStatus.created, description="The status of the job.")
-    created_at: datetime = Field(default_factory=get_utc_time, description="The unix timestamp of when the job was created.")
-    completed_at: Optional[datetime] = Field(None, description="The unix timestamp of when the job was completed.")
-    user_id: str = Field(..., description="The unique identifier of the user associated with the job.")
+    user_id: Optional[str] = Field(None, description="The unique identifier of the user associated with the job.")
 
 
 class JobUpdate(JobBase):
-    id: str = Field(..., description="The unique identifier of the job.")
-    status: Optional[JobStatus] = Field(..., description="The status of the job.")
+    status: Optional[JobStatus] = Field(None, description="The status of the job.")
+
+    class Config:
+        extra = "ignore"  # Ignores extra fields

{letta_nightly-0.6.0.dev20241204232133 → letta_nightly-0.6.1.dev20241205211219}/letta/server/rest_api/routers/v1/jobs.py

@@ -2,6 +2,8 @@ from typing import List, Optional
 
 from fastapi import APIRouter, Depends, Header, HTTPException, Query
 
+from letta.orm.errors import NoResultFound
+from letta.schemas.enums import JobStatus
 from letta.schemas.job import Job
 from letta.server.rest_api.utils import get_letta_server
 from letta.server.server import SyncServer
@@ -21,12 +23,11 @@ def list_jobs(
     actor = server.get_user_or_default(user_id=user_id)
 
     # TODO: add filtering by status
-    jobs = server.list_jobs(user_id=actor.id)
+    jobs = server.job_manager.list_jobs(actor=actor)
 
-    # TODO: eventually use ORM
-    # results = session.query(JobModel).filter(JobModel.user_id == user_id, JobModel.metadata_["source_id"].astext == sourced_id).all()
     if source_id:
         # can't be in the ORM since we have source_id stored in the metadata_
+        # TODO: Probably change this
         jobs = [job for job in jobs if job.metadata_.get("source_id") == source_id]
     return jobs
 
@@ -41,32 +42,39 @@ def list_active_jobs(
     """
     actor = server.get_user_or_default(user_id=user_id)
 
-    return server.list_active_jobs(user_id=actor.id)
+    return server.job_manager.list_jobs(actor=actor, statuses=[JobStatus.created, JobStatus.running])
 
 
 @router.get("/{job_id}", response_model=Job, operation_id="get_job")
 def get_job(
     job_id: str,
+    user_id: Optional[str] = Header(None, alias="user_id"),
     server: "SyncServer" = Depends(get_letta_server),
 ):
     """
     Get the status of a job.
     """
+    actor = server.get_user_or_default(user_id=user_id)
 
-    return server.get_job(job_id=job_id)
+    try:
+        return server.job_manager.get_job_by_id(job_id=job_id, actor=actor)
+    except NoResultFound:
+        raise HTTPException(status_code=404, detail="Job not found")
 
 
 @router.delete("/{job_id}", response_model=Job, operation_id="delete_job")
 def delete_job(
     job_id: str,
+    user_id: Optional[str] = Header(None, alias="user_id"),
     server: "SyncServer" = Depends(get_letta_server),
 ):
     """
     Delete a job by its job_id.
     """
-    job = server.get_job(job_id=job_id)
-    if not job:
-        raise HTTPException(status_code=404, detail="Job not found")
+    actor = server.get_user_or_default(user_id=user_id)
 
-    server.delete_job(job_id=job_id)
-    return job
+    try:
+        job = server.job_manager.delete_job_by_id(job_id=job_id, actor=actor)
+        return job
+    except NoResultFound:
+        raise HTTPException(status_code=404, detail="Job not found")

{letta_nightly-0.6.0.dev20241204232133 → letta_nightly-0.6.1.dev20241205211219}/letta/server/rest_api/routers/v1/sources.py

@@ -16,6 +16,7 @@ from letta.schemas.file import FileMetadata
 from letta.schemas.job import Job
 from letta.schemas.passage import Passage
 from letta.schemas.source import Source, SourceCreate, SourceUpdate
+from letta.schemas.user import User
 from letta.server.rest_api.utils import get_letta_server
 from letta.server.server import SyncServer
 from letta.utils import sanitize_filename
@@ -175,13 +176,14 @@ def upload_file_to_source(
         completed_at=None,
     )
     job_id = job.id
-    server.ms.create_job(job)
+    server.job_manager.create_job(job, actor=actor)
 
     # create background task
-    background_tasks.add_task(load_file_to_source_async, server, source_id=source.id, file=file, job_id=job.id, bytes=bytes)
+    background_tasks.add_task(load_file_to_source_async, server, source_id=source.id, file=file, job_id=job.id, bytes=bytes, actor=actor)
 
     # return job information
-    job = server.ms.get_job(job_id=job_id)
+    # Is this necessary? Can we just return the job from create_job?
+    job = server.job_manager.get_job_by_id(job_id=job_id, actor=actor)
     assert job is not None, "Job not found"
     return job
 
@@ -234,7 +236,7 @@ def delete_file_from_source(
         raise HTTPException(status_code=404, detail=f"File with id={file_id} not found.")
 
 
-def load_file_to_source_async(server: SyncServer, source_id: str, job_id: str, file: UploadFile, bytes: bytes):
+def load_file_to_source_async(server: SyncServer, source_id: str, job_id: str, file: UploadFile, bytes: bytes, actor: User):
     # Create a temporary directory (deleted after the context manager exits)
     with tempfile.TemporaryDirectory() as tmpdirname:
         # Sanitize the filename
@@ -246,4 +248,4 @@ def load_file_to_source_async(server: SyncServer, source_id: str, job_id: str, f
             buffer.write(bytes)
 
         # Pass the file to load_file_to_source
-        server.load_file_to_source(source_id, file_path, job_id)
+        server.load_file_to_source(source_id, file_path, job_id, actor)

{letta_nightly-0.6.0.dev20241204232133 → letta_nightly-0.6.1.dev20241205211219}/letta/server/server.py

@@ -6,7 +6,7 @@ import warnings
 from abc import abstractmethod
 from asyncio import Lock
 from datetime import datetime
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable, List, Optional, Tuple, Union
 
 from composio.client import Composio
 from composio.client.collections import ActionModel, AppModel
@@ -56,7 +56,7 @@ from letta.schemas.embedding_config import EmbeddingConfig
 
 # openai schemas
 from letta.schemas.enums import JobStatus
-from letta.schemas.job import Job
+from letta.schemas.job import Job, JobUpdate
 from letta.schemas.letta_message import FunctionReturn, LettaMessage
 from letta.schemas.llm_config import LLMConfig
 from letta.schemas.memory import (
@@ -75,6 +75,7 @@ from letta.schemas.user import User
 from letta.services.agents_tags_manager import AgentsTagsManager
 from letta.services.block_manager import BlockManager
 from letta.services.blocks_agents_manager import BlocksAgentsManager
+from letta.services.job_manager import JobManager
 from letta.services.organization_manager import OrganizationManager
 from letta.services.per_agent_lock_manager import PerAgentLockManager
 from letta.services.sandbox_config_manager import SandboxConfigManager
@@ -256,6 +257,7 @@ class SyncServer(Server):
         self.agents_tags_manager = AgentsTagsManager()
         self.sandbox_config_manager = SandboxConfigManager(tool_settings)
         self.blocks_agents_manager = BlocksAgentsManager()
+        self.job_manager = JobManager()
 
         # Managers that interface with parallelism
         self.per_agent_lock_manager = PerAgentLockManager()
@@ -393,6 +395,10 @@ class SyncServer(Server):
         agent_lock = self.per_agent_lock_manager.get_lock(agent_id)
         with agent_lock:
             agent_state = self.get_agent(agent_id=agent_id)
+            if agent_state is None:
+                raise ValueError(f"Agent (agent_id={agent_id}) does not exist")
+            elif agent_state.user_id is None:
+                raise ValueError(f"Agent (agent_id={agent_id}) does not have a user_id")
             actor = self.user_manager.get_user_by_id(user_id=agent_state.user_id)
 
             interface = interface or self.default_interface_factory()
@@ -880,7 +886,7 @@ class SyncServer(Server):
         in_memory_agent_state = self.get_agent(agent_state.id)
         return in_memory_agent_state
 
-    def get_agent(self, agent_id: str) -> AgentState:
+    def get_agent(self, agent_id: str) -> Optional[AgentState]:
         """
         Retrieve the full agent state from the DB.
         This gathers data accross multiple tables to provide the full state of an agent, which is passed into the `Agent` object for creation.
@@ -891,6 +897,8 @@ class SyncServer(Server):
         if agent_state is None:
             # agent does not exist
             return None
+        if agent_state.user_id is None:
+            raise ValueError(f"Agent {agent_id} does not have a user_id")
         user = self.user_manager.get_user_by_id(user_id=agent_state.user_id)
 
         # construct the in-memory, full agent state - this gather data stored in different tables but that needs to be passed to `Agent`
@@ -1469,39 +1477,12 @@ class SyncServer(Server):
 
         # TODO: delete data from agent passage stores (?)
 
-    def create_job(self, user_id: str, metadata: Optional[Dict] = None) -> Job:
-        """Create a new job"""
-        job = Job(
-            user_id=user_id,
-            status=JobStatus.created,
-            metadata_=metadata,
-        )
-        self.ms.create_job(job)
-        return job
-
-    def delete_job(self, job_id: str):
-        """Delete a job"""
-        self.ms.delete_job(job_id)
-
-    def get_job(self, job_id: str) -> Job:
-        """Get a job"""
-        return self.ms.get_job(job_id)
-
-    def list_jobs(self, user_id: str) -> List[Job]:
-        """List all jobs for a user"""
-        return self.ms.list_jobs(user_id=user_id)
-
-    def list_active_jobs(self, user_id: str) -> List[Job]:
-        """List all active jobs for a user"""
-        jobs = self.ms.list_jobs(user_id=user_id)
-        return [job for job in jobs if job.status in [JobStatus.created, JobStatus.running]]
-
-    def load_file_to_source(self, source_id: str, file_path: str, job_id: str) -> Job:
+    def load_file_to_source(self, source_id: str, file_path: str, job_id: str, actor: User) -> Job:
 
         # update job
-        job = self.ms.get_job(job_id)
+        job = self.job_manager.get_job_by_id(job_id, actor=actor)
         job.status = JobStatus.running
-        self.ms.update_job(job)
+        self.job_manager.update_job_by_id(job_id=job_id, job_update=JobUpdate(**job.model_dump()), actor=actor)
 
         # try:
         from letta.data_sources.connectors import DirectoryConnector
@@ -1509,23 +1490,12 @@ class SyncServer(Server):
         source = self.source_manager.get_source_by_id(source_id=source_id)
         connector = DirectoryConnector(input_files=[file_path])
         num_passages, num_documents = self.load_data(user_id=source.created_by_id, source_name=source.name, connector=connector)
-        # except Exception as e:
-        #    # job failed with error
-        #    error = str(e)
-        #    print(error)
-        #    job.status = JobStatus.failed
-        #    job.metadata_["error"] = error
-        #    self.ms.update_job(job)
-        #    # TODO: delete any associated passages/files?
-
-        #    # return failed job
-        #    return job
 
         # update job status
         job.status = JobStatus.completed
         job.metadata_["num_passages"] = num_passages
         job.metadata_["num_documents"] = num_documents
-        self.ms.update_job(job)
+        self.job_manager.update_job_by_id(job_id=job_id, job_update=JobUpdate(**job.model_dump()), actor=actor)
 
         return job
 

letta_nightly-0.6.1.dev20241205211219/letta/services/job_manager.py

@@ -0,0 +1,85 @@
+from typing import List, Optional
+
+from letta.orm.job import Job as JobModel
+from letta.schemas.enums import JobStatus
+from letta.schemas.job import Job as PydanticJob
+from letta.schemas.job import JobUpdate
+from letta.schemas.user import User as PydanticUser
+from letta.utils import enforce_types, get_utc_time
+
+
+class JobManager:
+    """Manager class to handle business logic related to Jobs."""
+
+    def __init__(self):
+        # Fetching the db_context similarly as in OrganizationManager
+        from letta.server.server import db_context
+
+        self.session_maker = db_context
+
+    @enforce_types
+    def create_job(self, pydantic_job: PydanticJob, actor: PydanticUser) -> PydanticJob:
+        """Create a new job based on the JobCreate schema."""
+        with self.session_maker() as session:
+            # Associate the job with the user
+            pydantic_job.user_id = actor.id
+            job_data = pydantic_job.model_dump()
+            job = JobModel(**job_data)
+            job.create(session, actor=actor)  # Save job in the database
+        return job.to_pydantic()
+
+    @enforce_types
+    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."""
+        with self.session_maker() as session:
+            # Fetch the job by ID
+            job = JobModel.read(db_session=session, identifier=job_id)  # TODO: Add this later , actor=actor)
+
+            # Update job attributes with only the fields that were explicitly set
+            update_data = job_update.model_dump(exclude_unset=True, exclude_none=True)
+
+            # Automatically update the completion timestamp if status is set to 'completed'
+            if update_data.get("status") == JobStatus.completed and not job.completed_at:
+                job.completed_at = get_utc_time()
+
+            for key, value in update_data.items():
+                setattr(job, key, value)
+
+            # Save the updated job to the database
+            return job.update(db_session=session)  # TODO: Add this later , actor=actor)
+
+    @enforce_types
+    def get_job_by_id(self, job_id: str, actor: PydanticUser) -> PydanticJob:
+        """Fetch a job by its ID."""
+        with self.session_maker() as session:
+            # Retrieve job by ID using the Job model's read method
+            job = JobModel.read(db_session=session, identifier=job_id)  # TODO: Add this later , actor=actor)
+            return job.to_pydantic()
+
+    @enforce_types
+    def list_jobs(
+        self, actor: PydanticUser, cursor: Optional[str] = None, limit: Optional[int] = 50, statuses: Optional[List[JobStatus]] = None
+    ) -> List[PydanticJob]:
+        """List all jobs with optional pagination and status filter."""
+        with self.session_maker() as session:
+            filter_kwargs = {"user_id": actor.id}
+
+            # Add status filter if provided
+            if statuses:
+                filter_kwargs["status"] = statuses
+
+            jobs = JobModel.list(
+                db_session=session,
+                cursor=cursor,
+                limit=limit,
+                **filter_kwargs,
+            )
+            return [job.to_pydantic() for job in jobs]
+
+    @enforce_types
+    def delete_job_by_id(self, job_id: str, actor: PydanticUser) -> PydanticJob:
+        """Delete a job by its ID."""
+        with self.session_maker() as session:
+            job = JobModel.read(db_session=session, identifier=job_id)  # TODO: Add this later , actor=actor)
+            job.hard_delete(db_session=session)  # TODO: Add this later , actor=actor)
+            return job.to_pydantic()