73 lines
2.4 KiB
Python
73 lines
2.4 KiB
Python
from dataclasses import dataclass
|
|
from uuid import uuid4
|
|
|
|
from app.modules.rag_session.repository import RagRepository
|
|
from app.schemas.common import ErrorPayload, ModuleName
|
|
from app.schemas.indexing import IndexJobStatus
|
|
|
|
|
|
@dataclass
|
|
class IndexJob:
|
|
index_job_id: str
|
|
rag_session_id: str
|
|
status: IndexJobStatus = IndexJobStatus.QUEUED
|
|
indexed_files: int = 0
|
|
failed_files: int = 0
|
|
cache_hit_files: int = 0
|
|
cache_miss_files: int = 0
|
|
error: ErrorPayload | None = None
|
|
|
|
|
|
class IndexJobStore:
|
|
def __init__(self, repository: RagRepository) -> None:
|
|
self._repo = repository
|
|
|
|
def create(self, rag_session_id: str) -> IndexJob:
|
|
job = IndexJob(index_job_id=str(uuid4()), rag_session_id=rag_session_id)
|
|
self._repo.create_job(job.index_job_id, rag_session_id, job.status.value)
|
|
return job
|
|
|
|
def get(self, index_job_id: str) -> IndexJob | None:
|
|
row = self._repo.get_job(index_job_id)
|
|
if not row:
|
|
return None
|
|
payload = None
|
|
if row.error_code:
|
|
module = ModuleName.RAG
|
|
if row.error_module:
|
|
try:
|
|
module = ModuleName(row.error_module)
|
|
except ValueError:
|
|
module = ModuleName.RAG
|
|
payload = ErrorPayload(
|
|
code=row.error_code,
|
|
desc=row.error_desc or "",
|
|
module=module,
|
|
)
|
|
return IndexJob(
|
|
index_job_id=row.index_job_id,
|
|
rag_session_id=row.rag_session_id,
|
|
status=IndexJobStatus(row.status),
|
|
indexed_files=row.indexed_files,
|
|
failed_files=row.failed_files,
|
|
cache_hit_files=row.cache_hit_files,
|
|
cache_miss_files=row.cache_miss_files,
|
|
error=payload,
|
|
)
|
|
|
|
def save(self, job: IndexJob) -> None:
|
|
error_code = job.error.code if job.error else None
|
|
error_desc = job.error.desc if job.error else None
|
|
error_module = job.error.module.value if job.error else None
|
|
self._repo.update_job(
|
|
job.index_job_id,
|
|
status=job.status.value,
|
|
indexed_files=job.indexed_files,
|
|
failed_files=job.failed_files,
|
|
cache_hit_files=job.cache_hit_files,
|
|
cache_miss_files=job.cache_miss_files,
|
|
error_code=error_code,
|
|
error_desc=error_desc,
|
|
error_module=error_module,
|
|
)
|