96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from sqlalchemy import text
|
|
|
|
from app.modules.shared.db import get_engine
|
|
|
|
|
|
@dataclass
|
|
class RagJobRow:
|
|
index_job_id: str
|
|
rag_session_id: str
|
|
status: str
|
|
indexed_files: int
|
|
failed_files: int
|
|
cache_hit_files: int
|
|
cache_miss_files: int
|
|
error_code: str | None
|
|
error_desc: str | None
|
|
error_module: str | None
|
|
|
|
|
|
class RagJobRepository:
|
|
def create_job(self, index_job_id: str, rag_session_id: str, status: str) -> None:
|
|
with get_engine().connect() as conn:
|
|
conn.execute(
|
|
text(
|
|
"""
|
|
INSERT INTO rag_index_jobs (index_job_id, rag_session_id, status)
|
|
VALUES (:jid, :sid, :status)
|
|
"""
|
|
),
|
|
{"jid": index_job_id, "sid": rag_session_id, "status": status},
|
|
)
|
|
conn.commit()
|
|
|
|
def update_job(
|
|
self,
|
|
index_job_id: str,
|
|
*,
|
|
status: str,
|
|
indexed_files: int,
|
|
failed_files: int,
|
|
cache_hit_files: int = 0,
|
|
cache_miss_files: int = 0,
|
|
error_code: str | None = None,
|
|
error_desc: str | None = None,
|
|
error_module: str | None = None,
|
|
) -> None:
|
|
with get_engine().connect() as conn:
|
|
conn.execute(
|
|
text(
|
|
"""
|
|
UPDATE rag_index_jobs
|
|
SET status = :status,
|
|
indexed_files = :indexed,
|
|
failed_files = :failed,
|
|
cache_hit_files = :cache_hit_files,
|
|
cache_miss_files = :cache_miss_files,
|
|
error_code = :ecode,
|
|
error_desc = :edesc,
|
|
error_module = :emodule,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
WHERE index_job_id = :jid
|
|
"""
|
|
),
|
|
{
|
|
"jid": index_job_id,
|
|
"status": status,
|
|
"indexed": indexed_files,
|
|
"failed": failed_files,
|
|
"cache_hit_files": cache_hit_files,
|
|
"cache_miss_files": cache_miss_files,
|
|
"ecode": error_code,
|
|
"edesc": error_desc,
|
|
"emodule": error_module,
|
|
},
|
|
)
|
|
conn.commit()
|
|
|
|
def get_job(self, index_job_id: str) -> RagJobRow | None:
|
|
with get_engine().connect() as conn:
|
|
row = conn.execute(
|
|
text(
|
|
"""
|
|
SELECT index_job_id, rag_session_id, status, indexed_files, failed_files,
|
|
cache_hit_files, cache_miss_files, error_code, error_desc, error_module
|
|
FROM rag_index_jobs
|
|
WHERE index_job_id = :jid
|
|
"""
|
|
),
|
|
{"jid": index_job_id},
|
|
).mappings().fetchone()
|
|
return RagJobRow(**dict(row)) if row else None
|