Новый раг

app/modules/rag/persistence/job_repository.py

@@ -0,0 +1,95 @@
+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