diff --git a/.git-hooks/pre-commit b/.git-hooks/pre-commit old mode 100755 new mode 100644 diff --git a/.gitignore b/.gitignore index 54c488e68..f10a3c556 100644 --- a/.gitignore +++ b/.gitignore @@ -183,6 +183,10 @@ dvclive # MacOS *.DS_Store +.prettierrc + +*.ipynb +*.pkl job_queue.db-shm job_queue.db-wal @@ -202,3 +206,13 @@ docs/node_modules/ !docs/docs/build/** !docs/i18n/**/build/ !docs/i18n/**/build/** + +**/csv-train.arrow +**/dataset_info.json + +/RAG_benchmark +tests/back/dataloaders/iris/** +tests/back/dataloaders/json/** +tests/back/models/json/** +tests/back/tasks/json/** +None/json/** diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000..4100dd641 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,126 @@ +# AGENTS.md + +This file provides guidance to AI coding agents when working with code in this repository. + +## What is DashAI + +DashAI is a desktop/web graphical toolbox for training, evaluating, and deploying ML models. FastAPI backend + React frontend, optionally wrapped in PyWebView for a native desktop window. Python >= 3.10. + +## Commands + +### Backend + +```bash +pip install -e . -r requirements-dev.txt +pre-commit install + +# Run dev server +python -m DashAI --no-browser --logging-level DEBUG + +# Lint / format (must pass both before committing) +ruff check --fix +ruff format + +# Run all tests (in-memory SQLite — no setup needed) +pytest tests/ + +# Run a single test +pytest tests/back/api/test_components_api.py -v +pytest tests/back/api/test_components_api.py::test_function_name -v + +# Database migrations (auto-runs on startup) +alembic upgrade head +``` + +### Frontend + +```bash +cd DashAI/front + +yarn install # requires Node LTS + Yarn 3.5.0 (enforced by packageManager) +yarn start # dev server on http://localhost:3000 (eslint disabled) +yarn build # eslint disabled +yarn test # eslint disabled +yarn test FileName.test.tsx +yarn lint # runs eslint explicitly (uses eslint.config.js) +``` + +**Important:** `start`, `build`, and `test` scripts all set `DISABLE_ESLINT_PLUGIN=true`. ESLint only runs via the dedicated `yarn lint` command. Prettier is configured in `DashAI/front/.prettierrc` and runs via pre-commit. + +## Architecture + +``` +Browser / PyWebView + → React (port 3000 dev / port 8000 prod) + → FastAPI (/api/v1/...) + → Service layer → SQLite (SQLAlchemy + Alembic) + → Huey job queue (async: training, conversion, exploration) +``` + +The frontend polls `/api/v1/jobs/{job_id}` for long-running task status. + +## Key files + +| Path | Purpose | +|------|---------| +| `DashAI/__main__.py` | CLI entry point (typer). Starts uvicorn + Huey consumer subprocess | +| `DashAI/back/app.py` | FastAPI `create_app` factory | +| `DashAI/back/container.py` | DI container (`kink`) — config, DB engine, ComponentRegistry, job queue | +| `DashAI/back/initial_components.py` | **Registers all components on startup.** Add new ML components here in `get_initial_components()` | +| `DashAI/back/dependencies/config_builder.py` | Builds config dict (paths, logging, calls `get_initial_components()`) | +| `DashAI/back/dependencies/registry/component_registry.py` | `ComponentRegistry` — resolves components by name string | +| `DashAI/back/api/api_v1/endpoints/` | REST endpoints — each file is a FastAPI router | +| `DashAI/back/api/api_v1/api.py` | Mounts all endpoint routers on `api_router_v1` | +| `DashAI/back/core/schema_fields/` | Type system driving dynamic frontend forms | +| `DashAI/back/pipeline/` | DAG pipeline nodes | +| `DashAI/back/plugins/` | Plugin system (PyPI packages with `dashai.plugins` entry point) | +| `DashAI/front/src/components/configurableObject/` | Auto-generates forms from backend component schemas | +| `DashAI/front/src/pages/generative/SessionRouter.jsx` | Routes `/app/generative/sessions/:id` to RAG or non-RAG view based on session `task_name` | +| `DashAI/front/src/pages/generative/RAGSession/` | RAG session setup wizard (RAGSessionSetup) + collapsible section components per pipeline stage | +| `DashAI/front/src/pages/generative/RAGSession/sections/` | Per-stage components: ChunkingSection, RetrieverSection, GeneratorSection, PromptSection | +| `DashAI/front/src/pages/generative/RAGSession/advanced/` | Advanced configuration modals: CompositeRetrieverBuilder, ChunkingConfigurationStep, RetrieverConfigurationStep, etc. | +| `DashAI/front/src/pages/generative/RAGSession/components/` | Reusable bodies: GeneratorBody, PromptBody, PresetCard, AdvancedConfigCard | + +## Key patterns + +**ComponentRegistry** — all ML components (models, metrics, converters, dataloaders, explorers, explainers, tasks, optimizers, jobs, pipeline nodes) are registered at startup and resolved by name string. To add a new component: subclass the relevant base, define its schema, and add it to `get_initial_components()` in `DashAI/back/initial_components.py`. + +**Schema / type system** — every component declares parameters using `BaseSchema` + field classes (`IntField`, `FloatField`, `ComponentField`, `UnionType`, etc.) with `MultilingualString` labels. The frontend uses these schemas to auto-generate configuration forms. + +**Dependency injection (`kink`)** — singletons (config, DB engine, ComponentRegistry, job queue) live in the `di` container. Use `@inject` to receive them. Config is accessed as `di["config"]` and includes `INITIAL_COMPONENTS`. + +**Huey job queue** — in dev mode the Huey consumer runs as a subprocess spawned from `__main__.py`; in PyInstaller bundles it runs as a daemon thread due to limitations with frozen executables. + +## RAG Module + +DashAI includes a **Retrieval-Augmented Generation (RAG)** module — a 4-stage pipeline (Document Loading → Chunking → Retrieval → Generation) for chatting with your documents. + +- **Backend:** `DashAI/back/models/RAG/` — pipeline orchestrator (RAGPipeline), abstract factory (RAGModelsFactory), sub-factories for prompts/chunking/retrievers/LLMs, retriever repository (all SQL), document loader, chunk-set caching. +- **Frontend:** `DashAI/front/src/pages/generative/RAGSession/` — session setup wizard (RAGSessionSetup, RAGSessionPage) with stage sections (ChunkingSection, RetrieverSection, GeneratorSection, PromptSection) and advanced config modals. +- **Jobs:** `RAGJob` extends `GenerativeJob` to run the RAG pipeline as a background task. +- **Full docs (backend + frontend):** See [`docs/rag/`](./docs/rag/) for architecture, data models, retrievers, pipeline orchestration, frontend architecture, testing guide, and known constraints. + +## Testing + +- Backend tests use in-memory SQLite — no setup needed. +- Key fixtures in `tests/back/conftest.py`: `test_path`, `test_datasets_path`, `test_job_queue`. +- Job queue tests: `test_job_queue.set_test_mode(immediate=True)` runs tasks synchronously. +- API tests use FastAPI `TestClient` defined in `tests/back/api/conftest.py` (module-scoped). +- Frontend test command passes `--passWithNoTests` in CI. + +## Adding new things + +**New ML model or converter:** subclass `BaseModel` / `BaseConverter`, define schema, register in `DashAI/back/initial_components.py::get_initial_components()`. + +**New API endpoint:** add file in `DashAI/back/api/api_v1/endpoints/`, include its router in `DashAI/back/api/api_v1/api.py`. + +**New async job:** subclass `BaseJob`, dispatch via `job_queue.enqueue(...)`, track status through jobs API. + +## Data at runtime + +Default `~/.DashAI/` (overridable via `DASHAI_LOCAL_PATH` env var): +`datasets/`, `runs/`, `explanations/`, `notebooks/`, `images/`, `documents/`, `rag/`, `sqlite.db` + +## CI + +Only a `publish.yml` workflow exists (on push to `production` branch). It builds the frontend, publishes to PyPI, and builds Windows/macOS installers. No CI test workflow is defined. diff --git a/DashAI/__main__.py b/DashAI/__main__.py index c2bc3717d..5aa48d4c5 100644 --- a/DashAI/__main__.py +++ b/DashAI/__main__.py @@ -65,7 +65,7 @@ def open_browser() -> None: - _wait_for_backend_server(timeout=120) + _wait_for_backend_server(timeout=1200) url = "http://localhost:8000/app/" webbrowser.open(url=url, new=0, autoraise=True) @@ -119,10 +119,11 @@ def _wait_for_backend_server(host="127.0.0.1", port=8000, timeout=15): import socket import time + timeout = 1000 start_time = time.time() while time.time() - start_time < timeout: try: - with socket.create_connection((host, port), timeout=1): + with socket.create_connection((host, port), timeout=timeout): return True except (OSError, ConnectionRefusedError): time.sleep(0.5) diff --git a/DashAI/alembic/versions/8d5416cac8f1_merge_rag_and_production_heads.py b/DashAI/alembic/versions/8d5416cac8f1_merge_rag_and_production_heads.py new file mode 100644 index 000000000..f37cd9196 --- /dev/null +++ b/DashAI/alembic/versions/8d5416cac8f1_merge_rag_and_production_heads.py @@ -0,0 +1,26 @@ +"""Merge RAG and production heads + +Revision ID: 8d5416cac8f1 +Revises: d4e5f6a7b8c9, f1a2b3c4d5e6 +Create Date: 2026-07-06 12:30:00.000000 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = '8d5416cac8f1' +down_revision: Union[str, None] = ('d4e5f6a7b8c9', 'f1a2b3c4d5e6') +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + pass + + +def downgrade() -> None: + pass diff --git a/DashAI/alembic/versions/a1b2c3d4e5f6_add_composite_retriever_support.py b/DashAI/alembic/versions/a1b2c3d4e5f6_add_composite_retriever_support.py new file mode 100644 index 000000000..c114329b0 --- /dev/null +++ b/DashAI/alembic/versions/a1b2c3d4e5f6_add_composite_retriever_support.py @@ -0,0 +1,47 @@ +"""add composite retriever support + +Revision ID: a1b2c3d4e5f6 +Revises: f928e0b5203d +Create Date: 2026-05-22 12:00:00.000000 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'a1b2c3d4e5f6' +down_revision: Union[str, None] = 'f928e0b5203d' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + with op.batch_alter_table('rag_retriever', schema=None) as batch_op: + batch_op.drop_constraint('chk_retriever_type', type_='check') + batch_op.add_column( + sa.Column('composite_child_ids', sa.JSON(), nullable=True) + ) + batch_op.create_check_constraint( + 'chk_retriever_type', + 'NOT (dense_retriever_id IS NOT NULL AND sparse_retriever_id IS NOT NULL)' + ) + + +def downgrade() -> None: + with op.batch_alter_table('rag_retriever', schema=None) as batch_op: + batch_op.drop_constraint('chk_retriever_type', type_='check') + batch_op.drop_column('composite_child_ids') + batch_op.create_check_constraint( + 'chk_retriever_type', + ( + "(class_name = 'SparseRetriever' " + "AND sparse_retriever_id IS NOT NULL " + "AND dense_retriever_id IS NULL) OR " + "(class_name = 'DenseRetriever' " + "AND dense_retriever_id IS NOT NULL " + "AND sparse_retriever_id IS NULL)" + ) + ) diff --git a/DashAI/alembic/versions/b2c3d4e5f6a7_add_uniqueness_constraints_rag.py b/DashAI/alembic/versions/b2c3d4e5f6a7_add_uniqueness_constraints_rag.py new file mode 100644 index 000000000..e105106cd --- /dev/null +++ b/DashAI/alembic/versions/b2c3d4e5f6a7_add_uniqueness_constraints_rag.py @@ -0,0 +1,97 @@ +"""add uniqueness constraints to rag tables and fix storage tracing + +Revision ID: b2c3d4e5f6a7 +Revises: a1b2c3d4e5f6 +Create Date: 2026-05-25 12:00:00.000000 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = 'b2c3d4e5f6a7' +down_revision: Union[str, None] = 'a1b2c3d4e5f6' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ── Chunk: replace useless (id, document_id) with (document_id, chunk_index, chunking_model_id) ── + with op.batch_alter_table('chunk', schema=None) as batch_op: + batch_op.drop_constraint('uix_chunk_document', type_='unique') + batch_op.create_unique_constraint( + 'uix_document_chunk_index_chunking', + ['document_id', 'chunk_index', 'chunking_model_id'], + ) + + # ── RAGPrompt ── + with op.batch_alter_table('rag_prompt', schema=None) as batch_op: + batch_op.create_unique_constraint( + 'uix_rag_prompt_class_params', + ['class_name', 'parameters'], + ) + + # ── RAGGenerationModel ── + with op.batch_alter_table('rag_generation_model', schema=None) as batch_op: + batch_op.create_unique_constraint( + 'uix_rag_gen_model_class_params', + ['class_name', 'parameters'], + ) + + # ── RAGChunkingModel ── + with op.batch_alter_table('rag_chunking_model', schema=None) as batch_op: + batch_op.create_unique_constraint( + 'uix_rag_chunking_model_class_params', + ['class_name', 'parameters'], + ) + + # ── RAGEmbeddingModel ── + with op.batch_alter_table('rag_embedding_model', schema=None) as batch_op: + batch_op.create_unique_constraint( + 'uix_rag_embedding_model_class_params', + ['class_name', 'parameters'], + ) + + # ── RAGSparseRetriever ── + with op.batch_alter_table('rag_sparse_retriever', schema=None) as batch_op: + batch_op.create_unique_constraint( + 'uix_rag_sparse_retriever', + ['class_name', 'parameters', 'documents_ids', 'chunking_model_id'], + ) + + # ── RAGDenseRetriever ── + with op.batch_alter_table('rag_dense_retriever', schema=None) as batch_op: + batch_op.create_unique_constraint( + 'uix_rag_dense_retriever', + ['class_name', 'parameters', 'document_ids', + 'chunking_model_id', 'embedding_model_id'], + ) + + +def downgrade() -> None: + with op.batch_alter_table('rag_dense_retriever', schema=None) as batch_op: + batch_op.drop_constraint('uix_rag_dense_retriever', type_='unique') + + with op.batch_alter_table('rag_sparse_retriever', schema=None) as batch_op: + batch_op.drop_constraint('uix_rag_sparse_retriever', type_='unique') + + with op.batch_alter_table('rag_embedding_model', schema=None) as batch_op: + batch_op.drop_constraint('uix_rag_embedding_model_class_params', type_='unique') + + with op.batch_alter_table('rag_chunking_model', schema=None) as batch_op: + batch_op.drop_constraint('uix_rag_chunking_model_class_params', type_='unique') + + with op.batch_alter_table('rag_generation_model', schema=None) as batch_op: + batch_op.drop_constraint('uix_rag_gen_model_class_params', type_='unique') + + with op.batch_alter_table('rag_prompt', schema=None) as batch_op: + batch_op.drop_constraint('uix_rag_prompt_class_params', type_='unique') + + with op.batch_alter_table('chunk', schema=None) as batch_op: + batch_op.drop_constraint('uix_document_chunk_index_chunking', type_='unique') + batch_op.create_unique_constraint( + 'uix_chunk_document', + ['id', 'document_id'], + ) diff --git a/DashAI/alembic/versions/c3d4e5f6a7b8_invert_retriever_fk_add_child_table.py b/DashAI/alembic/versions/c3d4e5f6a7b8_invert_retriever_fk_add_child_table.py new file mode 100644 index 000000000..93de3de17 --- /dev/null +++ b/DashAI/alembic/versions/c3d4e5f6a7b8_invert_retriever_fk_add_child_table.py @@ -0,0 +1,134 @@ +"""invert FK direction for retrievers and add rag_retriever_child + +Revision ID: c3d4e5f6a7b8 +Revises: b2c3d4e5f6a7 +Create Date: 2026-05-26 12:00:00.000000 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = 'c3d4e5f6a7b8' +down_revision: Union[str, None] = 'b2c3d4e5f6a7' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ── Drop old FK from RAGPipeline to RAGRetriever ── + with op.batch_alter_table('rag_pipeline', schema=None) as batch_op: + batch_op.drop_constraint( + 'fk_rag_pipeline_retriever_model_id_rag_retriever', type_='foreignkey' + ) + batch_op.drop_column('retriever_model_id') + + # ── Rebuild RAGRetriever: drop old columns, add pipeline_id ── + with op.batch_alter_table('rag_retriever', schema=None) as batch_op: + batch_op.drop_constraint('chk_retriever_type', type_='check') + batch_op.drop_constraint( + 'fk_rag_retriever_dense_retriever_id_rag_dense_retriever', type_='foreignkey' + ) + batch_op.drop_constraint( + 'fk_rag_retriever_sparse_retriever_id_rag_sparse_retriever', type_='foreignkey' + ) + batch_op.drop_column('dense_retriever_id') + batch_op.drop_column('sparse_retriever_id') + batch_op.drop_column('composite_child_ids') + batch_op.add_column( + sa.Column('pipeline_id', sa.Integer(), nullable=False) + ) + batch_op.create_foreign_key( + 'fk_rag_retriever_pipeline_id', + 'rag_pipeline', ['pipeline_id'], ['id'], ondelete='CASCADE', + ) + + # ── Add bridge_id to RAGSparseRetriever ── + with op.batch_alter_table('rag_sparse_retriever', schema=None) as batch_op: + batch_op.add_column( + sa.Column('bridge_id', sa.Integer(), nullable=False) + ) + batch_op.create_foreign_key( + 'fk_rag_sparse_retriever_bridge_id', + 'rag_retriever', ['bridge_id'], ['id'], ondelete='CASCADE', + ) + + # ── Add bridge_id to RAGDenseRetriever ── + with op.batch_alter_table('rag_dense_retriever', schema=None) as batch_op: + batch_op.add_column( + sa.Column('bridge_id', sa.Integer(), nullable=False) + ) + batch_op.create_foreign_key( + 'fk_rag_dense_retriever_bridge_id', + 'rag_retriever', ['bridge_id'], ['id'], ondelete='CASCADE', + ) + + # ── Create RAGRetrieverChild ── + op.create_table( + 'rag_retriever_child', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('parent_id', sa.Integer(), nullable=False), + sa.Column('child_id', sa.Integer(), nullable=False), + sa.Column('child_order', sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ['parent_id'], ['rag_retriever.id'], + name='fk_rag_retriever_child_parent', ondelete='CASCADE', + ), + sa.ForeignKeyConstraint( + ['child_id'], ['rag_retriever.id'], + name='fk_rag_retriever_child_child', ondelete='CASCADE', + ), + sa.PrimaryKeyConstraint('id', name='pk_rag_retriever_child'), + sa.UniqueConstraint( + 'parent_id', 'child_order', + name='uix_retriever_child_order', + ), + ) + + +def downgrade() -> None: + op.drop_table('rag_retriever_child') + + with op.batch_alter_table('rag_dense_retriever', schema=None) as batch_op: + batch_op.drop_constraint('fk_rag_dense_retriever_bridge_id', type_='foreignkey') + batch_op.drop_column('bridge_id') + + with op.batch_alter_table('rag_sparse_retriever', schema=None) as batch_op: + batch_op.drop_constraint('fk_rag_sparse_retriever_bridge_id', type_='foreignkey') + batch_op.drop_column('bridge_id') + + with op.batch_alter_table('rag_retriever', schema=None) as batch_op: + batch_op.drop_constraint('fk_rag_retriever_pipeline_id', type_='foreignkey') + batch_op.drop_column('pipeline_id') + batch_op.add_column( + sa.Column('composite_child_ids', sa.JSON(), nullable=True) + ) + batch_op.add_column( + sa.Column('sparse_retriever_id', sa.Integer(), nullable=True) + ) + batch_op.add_column( + sa.Column('dense_retriever_id', sa.Integer(), nullable=True) + ) + batch_op.create_foreign_key( + 'fk_rag_retriever_sparse_retriever_id_rag_sparse_retriever', + 'rag_sparse_retriever', ['sparse_retriever_id'], ['id'], ondelete='CASCADE', + ) + batch_op.create_foreign_key( + 'fk_rag_retriever_dense_retriever_id_rag_dense_retriever', + 'rag_dense_retriever', ['dense_retriever_id'], ['id'], ondelete='CASCADE', + ) + batch_op.create_check_constraint( + 'chk_retriever_type', + 'NOT (dense_retriever_id IS NOT NULL AND sparse_retriever_id IS NOT NULL)', + ) + + with op.batch_alter_table('rag_pipeline', schema=None) as batch_op: + batch_op.add_column( + sa.Column('retriever_model_id', sa.Integer(), nullable=False) + ) + batch_op.create_foreign_key( + 'fk_rag_pipeline_retriever_model_id_rag_retriever', + 'rag_retriever', ['retriever_model_id'], ['id'], ondelete='CASCADE', + ) diff --git a/DashAI/alembic/versions/d4e5f6a7b8c9_add_chunk_set_architecture.py b/DashAI/alembic/versions/d4e5f6a7b8c9_add_chunk_set_architecture.py new file mode 100644 index 000000000..71b937098 --- /dev/null +++ b/DashAI/alembic/versions/d4e5f6a7b8c9_add_chunk_set_architecture.py @@ -0,0 +1,222 @@ +"""add chunk_set architecture: chunk_set, chunk_set_document, retriever_chunk_set + +Revision ID: d4e5f6a7b8c9 +Revises: c3d4e5f6a7b8 +Create Date: 2026-05-26 18:00:00.000000 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = 'd4e5f6a7b8c9' +down_revision: Union[str, None] = 'c3d4e5f6a7b8' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ── RAGChunkSet ── + op.create_table( + 'rag_chunk_set', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('signature', sa.String(), nullable=False), + sa.Column('parameters', sa.JSON(), nullable=True), + sa.PrimaryKeyConstraint('id', name='pk_rag_chunk_set'), + sa.UniqueConstraint('signature', name='uq_rag_chunk_set_signature'), + ) + + # ── RAGChunkSetDocument ── + op.create_table( + 'rag_chunk_set_document', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('chunk_set_id', sa.Integer(), nullable=False), + sa.Column('document_id', sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ['chunk_set_id'], ['rag_chunk_set.id'], + name='fk_chunk_set_doc_set', ondelete='CASCADE', + ), + sa.ForeignKeyConstraint( + ['document_id'], ['document.id'], + name='fk_chunk_set_doc_document', ondelete='CASCADE', + ), + sa.PrimaryKeyConstraint('id', name='pk_rag_chunk_set_document'), + sa.UniqueConstraint( + 'chunk_set_id', 'document_id', name='uix_chunk_set_document', + ), + ) + + # ── Chunk: replace chunking_model_id with chunk_set_id, add metadata ── + with op.batch_alter_table('chunk', schema=None) as batch_op: + batch_op.drop_constraint('uix_document_chunk_index_chunking', type_='unique') + batch_op.drop_constraint( + 'fk_chunk_chunking_model_id_rag_chunking_model', type_='foreignkey' + ) + batch_op.drop_column('chunking_model_id') + batch_op.add_column( + sa.Column('chunk_set_id', sa.Integer(), nullable=False) + ) + batch_op.add_column( + sa.Column('metadata', sa.JSON(), nullable=True) + ) + batch_op.create_foreign_key( + 'fk_chunk_chunk_set_id', + 'rag_chunk_set', ['chunk_set_id'], ['id'], ondelete='CASCADE', + ) + batch_op.create_unique_constraint( + 'uix_chunk_set_doc_index', + ['chunk_set_id', 'document_id', 'chunk_index'], + ) + + # ── RAGSparseRetriever: drop documents_ids + chunking_model_id, add chunk_set_id ── + with op.batch_alter_table('rag_sparse_retriever', schema=None) as batch_op: + batch_op.drop_constraint('uix_rag_sparse_retriever', type_='unique') + batch_op.drop_constraint( + 'fk_rag_sparse_retriever_chunking_model_id_rag_chunking_model', + type_='foreignkey', + ) + batch_op.drop_column('documents_ids') + batch_op.drop_column('chunking_model_id') + batch_op.add_column( + sa.Column('chunk_set_id', sa.Integer(), nullable=False) + ) + batch_op.create_foreign_key( + 'fk_rag_sparse_retriever_chunk_set_id', + 'rag_chunk_set', ['chunk_set_id'], ['id'], ondelete='CASCADE', + ) + batch_op.create_unique_constraint( + 'uix_rag_sparse_retriever', + ['class_name', 'parameters', 'chunk_set_id'], + ) + + # ── RAGDenseRetriever: drop document_ids + chunking_model_id, add chunk_set_id ── + with op.batch_alter_table('rag_dense_retriever', schema=None) as batch_op: + batch_op.drop_constraint('uix_rag_dense_retriever', type_='unique') + batch_op.drop_constraint( + 'fk_rag_dense_retriever_chunking_model_id_rag_chunking_model', + type_='foreignkey', + ) + batch_op.drop_column('document_ids') + batch_op.drop_column('chunking_model_id') + batch_op.add_column( + sa.Column('chunk_set_id', sa.Integer(), nullable=False) + ) + batch_op.create_foreign_key( + 'fk_rag_dense_retriever_chunk_set_id', + 'rag_chunk_set', ['chunk_set_id'], ['id'], ondelete='CASCADE', + ) + batch_op.create_unique_constraint( + 'uix_rag_dense_retriever', + ['class_name', 'parameters', 'chunk_set_id', 'embedding_model_id'], + ) + + # ── RAGEmbeddingMatrix: chunking_model_id → chunk_set_id ── + with op.batch_alter_table('rag_embedding_matrix', schema=None) as batch_op: + batch_op.drop_constraint('uix_document_chunking_embedding', type_='unique') + batch_op.drop_constraint( + 'fk_rag_embedding_matrix_chunking_model_id_rag_chunking_model', + type_='foreignkey', + ) + batch_op.alter_column('chunking_model_id', new_column_name='chunk_set_id') + batch_op.create_foreign_key( + 'fk_rag_embedding_matrix_chunk_set_id', + 'rag_chunk_set', ['chunk_set_id'], ['id'], ondelete='CASCADE', + ) + batch_op.create_unique_constraint( + 'uix_document_chunk_set_embedding', + ['document_id', 'chunk_set_id', 'embedding_model_id'], + ) + + # ── RAGRetrieverChunkSet ── + op.create_table( + 'rag_retriever_chunk_set', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('retriever_id', sa.Integer(), nullable=False), + sa.Column('chunk_set_id', sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ['retriever_id'], ['rag_retriever.id'], + name='fk_retriever_chunk_set_retriever', ondelete='CASCADE', + ), + sa.ForeignKeyConstraint( + ['chunk_set_id'], ['rag_chunk_set.id'], + name='fk_retriever_chunk_set_chunk_set', ondelete='CASCADE', + ), + sa.PrimaryKeyConstraint('id', name='pk_rag_retriever_chunk_set'), + sa.UniqueConstraint( + 'retriever_id', 'chunk_set_id', name='uix_retriever_chunk_set', + ), + ) + + +def downgrade() -> None: + op.drop_table('rag_retriever_chunk_set') + + with op.batch_alter_table('rag_embedding_matrix', schema=None) as batch_op: + batch_op.drop_constraint('uix_document_chunk_set_embedding', type_='unique') + batch_op.drop_constraint( + 'fk_rag_embedding_matrix_chunk_set_id', type_='foreignkey', + ) + batch_op.alter_column('chunk_set_id', new_column_name='chunking_model_id') + batch_op.create_foreign_key( + 'fk_rag_embedding_matrix_chunking_model_id_rag_chunking_model', + 'rag_chunking_model', ['chunking_model_id'], ['id'], ondelete='CASCADE', + ) + batch_op.create_unique_constraint( + 'uix_document_chunking_embedding', + ['document_id', 'chunking_model_id', 'embedding_model_id'], + ) + + with op.batch_alter_table('rag_dense_retriever', schema=None) as batch_op: + batch_op.drop_constraint('uix_rag_dense_retriever', type_='unique') + batch_op.drop_constraint( + 'fk_rag_dense_retriever_chunk_set_id', type_='foreignkey', + ) + batch_op.drop_column('chunk_set_id') + batch_op.add_column(sa.Column('chunking_model_id', sa.Integer(), nullable=False)) + batch_op.add_column(sa.Column('document_ids', sa.JSON(), nullable=False)) + batch_op.create_foreign_key( + 'fk_rag_dense_retriever_chunking_model_id_rag_chunking_model', + 'rag_chunking_model', ['chunking_model_id'], ['id'], ondelete='CASCADE', + ) + batch_op.create_unique_constraint( + 'uix_rag_dense_retriever', + ['class_name', 'parameters', 'document_ids', + 'chunking_model_id', 'embedding_model_id'], + ) + + with op.batch_alter_table('rag_sparse_retriever', schema=None) as batch_op: + batch_op.drop_constraint('uix_rag_sparse_retriever', type_='unique') + batch_op.drop_constraint( + 'fk_rag_sparse_retriever_chunk_set_id', type_='foreignkey', + ) + batch_op.drop_column('chunk_set_id') + batch_op.add_column(sa.Column('chunking_model_id', sa.Integer(), nullable=False)) + batch_op.add_column(sa.Column('documents_ids', sa.JSON(), nullable=False)) + batch_op.create_foreign_key( + 'fk_rag_sparse_retriever_chunking_model_id_rag_chunking_model', + 'rag_chunking_model', ['chunking_model_id'], ['id'], ondelete='CASCADE', + ) + batch_op.create_unique_constraint( + 'uix_rag_sparse_retriever', + ['class_name', 'parameters', 'documents_ids', 'chunking_model_id'], + ) + + with op.batch_alter_table('chunk', schema=None) as batch_op: + batch_op.drop_constraint('uix_chunk_set_doc_index', type_='unique') + batch_op.drop_constraint('fk_chunk_chunk_set_id', type_='foreignkey') + batch_op.drop_column('chunk_set_id') + batch_op.drop_column('metadata') + batch_op.add_column(sa.Column('chunking_model_id', sa.Integer(), nullable=False)) + batch_op.create_foreign_key( + 'fk_chunk_chunking_model_id_rag_chunking_model', + 'rag_chunking_model', ['chunking_model_id'], ['id'], ondelete='CASCADE', + ) + batch_op.create_unique_constraint( + 'uix_document_chunk_index_chunking', + ['document_id', 'chunk_index', 'chunking_model_id'], + ) + + op.drop_table('rag_chunk_set_document') + op.drop_table('rag_chunk_set') diff --git a/DashAI/alembic/versions/f06652057903_add_rag_tables.py b/DashAI/alembic/versions/f06652057903_add_rag_tables.py new file mode 100644 index 000000000..e489b4de5 --- /dev/null +++ b/DashAI/alembic/versions/f06652057903_add_rag_tables.py @@ -0,0 +1,166 @@ +"""add rag tables + +Revision ID: f06652057903 +Revises: b4f9e70098e7 +Create Date: 2026-04-17 02:11:23.522691 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'f06652057903' +down_revision: Union[str, None] = 'b4f9e70098e7' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('document', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('file_name', sa.String(), nullable=False), + sa.Column('file_type', sa.String(), nullable=False), + sa.Column('file_path', sa.String(), nullable=False), + sa.Column('file_hash', sa.String(), nullable=False), + sa.Column('optional_metadata', sa.JSON(), nullable=True), + sa.Column('created', sa.DateTime(), nullable=False), + sa.Column('last_modified', sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint('id', name=op.f('pk_document')), + sa.UniqueConstraint('file_hash', name=op.f('uq_document_file_hash')) + ) + op.create_table('rag_chunking_model', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('class_name', sa.String(), nullable=False), + sa.Column('parameters', sa.JSON(), nullable=True), + sa.PrimaryKeyConstraint('id', name=op.f('pk_rag_chunking_model')) + ) + op.create_table('rag_embedding_model', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('class_name', sa.String(), nullable=False), + sa.Column('parameters', sa.JSON(), nullable=True), + sa.PrimaryKeyConstraint('id', name=op.f('pk_rag_embedding_model')) + ) + op.create_table('rag_generation_model', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('class_name', sa.String(), nullable=False), + sa.Column('parameters', sa.JSON(), nullable=True), + sa.PrimaryKeyConstraint('id', name=op.f('pk_rag_generation_model')) + ) + op.create_table('rag_prompt', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('class_name', sa.String(), nullable=False), + sa.Column('name', sa.String(), nullable=True), + sa.Column('parameters', sa.JSON(), nullable=True), + sa.Column('created', sa.DateTime(), nullable=False), + sa.Column('last_modified', sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint('id', name=op.f('pk_rag_prompt')) + ) + op.create_table('chunk', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('document_id', sa.Integer(), nullable=False), + sa.Column('chunk_index', sa.Integer(), nullable=False), + sa.Column('chunking_model_id', sa.Integer(), nullable=False), + sa.Column('text', sa.Text(), nullable=False), + sa.ForeignKeyConstraint(['chunking_model_id'], ['rag_chunking_model.id'], name=op.f('fk_chunk_chunking_model_id_rag_chunking_model'), ondelete='CASCADE'), + sa.ForeignKeyConstraint(['document_id'], ['document.id'], name=op.f('fk_chunk_document_id_document'), ondelete='CASCADE'), + sa.PrimaryKeyConstraint('id', name=op.f('pk_chunk')), + sa.UniqueConstraint('id', 'document_id', name='uix_chunk_document') + ) + op.create_table('rag_dense_retriever', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('class_name', sa.String(), nullable=False), + sa.Column('parameters', sa.JSON(), nullable=True), + sa.Column('embedding_model_id', sa.Integer(), nullable=False), + sa.Column('document_ids', sa.JSON(), nullable=False), + sa.Column('chunking_model_id', sa.Integer(), nullable=False), + sa.ForeignKeyConstraint(['chunking_model_id'], ['rag_chunking_model.id'], name=op.f('fk_rag_dense_retriever_chunking_model_id_rag_chunking_model'), ondelete='CASCADE'), + sa.ForeignKeyConstraint(['embedding_model_id'], ['rag_embedding_model.id'], name=op.f('fk_rag_dense_retriever_embedding_model_id_rag_embedding_model'), ondelete='CASCADE'), + sa.PrimaryKeyConstraint('id', name=op.f('pk_rag_dense_retriever')) + ) + op.create_table('rag_embedding_matrix', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('document_id', sa.Integer(), nullable=False), + sa.Column('chunking_model_id', sa.Integer(), nullable=False), + sa.Column('embedding_model_id', sa.Integer(), nullable=False), + sa.Column('storage_folder', sa.String(), nullable=False), + sa.Column('matrix_shape', sa.JSON(), nullable=False), + sa.Column('created', sa.DateTime(), nullable=False), + sa.Column('last_modified', sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint(['chunking_model_id'], ['rag_chunking_model.id'], name=op.f('fk_rag_embedding_matrix_chunking_model_id_rag_chunking_model'), ondelete='CASCADE'), + sa.ForeignKeyConstraint(['document_id'], ['document.id'], name=op.f('fk_rag_embedding_matrix_document_id_document'), ondelete='CASCADE'), + sa.ForeignKeyConstraint(['embedding_model_id'], ['rag_embedding_model.id'], name=op.f('fk_rag_embedding_matrix_embedding_model_id_rag_embedding_model'), ondelete='CASCADE'), + sa.PrimaryKeyConstraint('id', name=op.f('pk_rag_embedding_matrix')), + sa.UniqueConstraint('document_id', 'chunking_model_id', 'embedding_model_id', name='uix_document_chunking_embedding') + ) + op.create_table('rag_sparse_retriever', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('class_name', sa.String(), nullable=False), + sa.Column('parameters', sa.JSON(), nullable=True), + sa.Column('storage_folder', sa.String(), nullable=False), + sa.Column('documents_ids', sa.JSON(), nullable=False), + sa.Column('chunking_model_id', sa.Integer(), nullable=False), + sa.ForeignKeyConstraint(['chunking_model_id'], ['rag_chunking_model.id'], name=op.f('fk_rag_sparse_retriever_chunking_model_id_rag_chunking_model'), ondelete='CASCADE'), + sa.PrimaryKeyConstraint('id', name=op.f('pk_rag_sparse_retriever')) + ) + op.create_table('rag_retriever', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('class_name', sa.String(), nullable=False), + sa.Column('dense_retriever_id', sa.Integer(), nullable=True), + sa.Column('sparse_retriever_id', sa.Integer(), nullable=True), + sa.CheckConstraint("(class_name = 'SparseRetriever' AND sparse_retriever_id IS NOT NULL AND dense_retriever_id IS NULL) OR (class_name = 'DenseRetriever' AND dense_retriever_id IS NOT NULL AND sparse_retriever_id IS NULL)", name=op.f('ck_chk_retriever_type')), + sa.ForeignKeyConstraint(['dense_retriever_id'], ['rag_dense_retriever.id'], name=op.f('fk_rag_retriever_dense_retriever_id_rag_dense_retriever'), ondelete='CASCADE'), + sa.ForeignKeyConstraint(['sparse_retriever_id'], ['rag_sparse_retriever.id'], name=op.f('fk_rag_retriever_sparse_retriever_id_rag_sparse_retriever'), ondelete='CASCADE'), + sa.PrimaryKeyConstraint('id', name=op.f('pk_rag_retriever')) + ) + op.create_table('rag_pipeline', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('session_id', sa.Integer(), nullable=False), + sa.Column('name', sa.String(), nullable=False), + sa.Column('description', sa.String(), nullable=True), + sa.Column('parameters', sa.JSON(), nullable=True), + sa.Column('chunking_model_id', sa.Integer(), nullable=False), + sa.Column('retriever_model_id', sa.Integer(), nullable=False), + sa.Column('prompt_id', sa.Integer(), nullable=False), + sa.Column('generation_model_id', sa.Integer(), nullable=False), + sa.ForeignKeyConstraint(['chunking_model_id'], ['rag_chunking_model.id'], name=op.f('fk_rag_pipeline_chunking_model_id_rag_chunking_model'), ondelete='CASCADE'), + sa.ForeignKeyConstraint(['generation_model_id'], ['rag_generation_model.id'], name=op.f('fk_rag_pipeline_generation_model_id_rag_generation_model'), ondelete='CASCADE'), + sa.ForeignKeyConstraint(['prompt_id'], ['rag_prompt.id'], name=op.f('fk_rag_pipeline_prompt_id_rag_prompt'), ondelete='CASCADE'), + sa.ForeignKeyConstraint(['retriever_model_id'], ['rag_retriever.id'], name=op.f('fk_rag_pipeline_retriever_model_id_rag_retriever'), ondelete='CASCADE'), + sa.ForeignKeyConstraint(['session_id'], ['generative_session.id'], name=op.f('fk_rag_pipeline_session_id_generative_session'), ondelete='CASCADE'), + sa.PrimaryKeyConstraint('id', name=op.f('pk_rag_pipeline')) + ) + op.create_table('rag_document_pipeline_session_link', + sa.Column('id', sa.Integer(), autoincrement=True, nullable=False), + sa.Column('document_id', sa.Integer(), nullable=False), + sa.Column('session_id', sa.Integer(), nullable=False), + sa.Column('pipeline_id', sa.Integer(), nullable=False), + sa.ForeignKeyConstraint(['document_id'], ['document.id'], name=op.f('fk_rag_document_pipeline_session_link_document_id_document'), ondelete='CASCADE'), + sa.ForeignKeyConstraint(['pipeline_id'], ['rag_pipeline.id'], name=op.f('fk_rag_document_pipeline_session_link_pipeline_id_rag_pipeline'), ondelete='CASCADE'), + sa.ForeignKeyConstraint(['session_id'], ['generative_session.id'], name=op.f('fk_rag_document_pipeline_session_link_session_id_generative_session'), ondelete='CASCADE'), + sa.PrimaryKeyConstraint('id', name=op.f('pk_rag_document_pipeline_session_link')), + sa.UniqueConstraint('document_id', 'pipeline_id', name='uix_document_pipeline'), + sa.UniqueConstraint('document_id', 'session_id', name='uix_document_session'), + sa.UniqueConstraint('session_id', 'pipeline_id', name='uix_session_pipeline') + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table('rag_document_pipeline_session_link') + op.drop_table('rag_pipeline') + op.drop_table('rag_retriever') + op.drop_table('rag_sparse_retriever') + op.drop_table('rag_embedding_matrix') + op.drop_table('rag_dense_retriever') + op.drop_table('chunk') + op.drop_table('rag_prompt') + op.drop_table('rag_generation_model') + op.drop_table('rag_embedding_model') + op.drop_table('rag_chunking_model') + op.drop_table('document') + # ### end Alembic commands ### diff --git a/DashAI/alembic/versions/f928e0b5203d_rag_tables_added.py b/DashAI/alembic/versions/f928e0b5203d_rag_tables_added.py new file mode 100644 index 000000000..4e7eb013d --- /dev/null +++ b/DashAI/alembic/versions/f928e0b5203d_rag_tables_added.py @@ -0,0 +1,36 @@ +"""RAG tables added + +Revision ID: f928e0b5203d +Revises: f06652057903 +Create Date: 2026-04-24 18:55:45.236676 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'f928e0b5203d' +down_revision: Union[str, None] = 'f06652057903' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('plugin', schema=None) as batch_op: + batch_op.add_column(sa.Column('latest_version', sa.String(), nullable=False)) + batch_op.drop_column('lastest_version') + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('plugin', schema=None) as batch_op: + batch_op.add_column(sa.Column('lastest_version', sa.VARCHAR(), nullable=False)) + batch_op.drop_column('latest_version') + + # ### end Alembic commands ### diff --git a/DashAI/back/api/api_v1/api.py b/DashAI/back/api/api_v1/api.py index 1beb875a5..87b7ee2dc 100644 --- a/DashAI/back/api/api_v1/api.py +++ b/DashAI/back/api/api_v1/api.py @@ -5,6 +5,7 @@ from DashAI.back.api.api_v1.endpoints.datafile import router as datafile_router from DashAI.back.api.api_v1.endpoints.dataset_source import router as dataset_source from DashAI.back.api.api_v1.endpoints.datasets import router as datasets +from DashAI.back.api.api_v1.endpoints.documents import router as documents from DashAI.back.api.api_v1.endpoints.explainers import router as explainers from DashAI.back.api.api_v1.endpoints.explorers import router as explorers from DashAI.back.api.api_v1.endpoints.folders import router as folders @@ -22,6 +23,7 @@ from DashAI.back.api.api_v1.endpoints.pipelines import router as pipelines from DashAI.back.api.api_v1.endpoints.plugins import router as plugins from DashAI.back.api.api_v1.endpoints.predict import router as predict +from DashAI.back.api.api_v1.endpoints.prompts import router as prompts from DashAI.back.api.api_v1.endpoints.runs import router as runs from DashAI.back.api.api_v1.endpoints.scoring import router as scoring @@ -29,6 +31,7 @@ api_router_v1.include_router(converters, prefix="/converter") api_router_v1.include_router(components, prefix="/component") api_router_v1.include_router(datasets, prefix="/dataset") +api_router_v1.include_router(documents, prefix="/document") api_router_v1.include_router(model_sessions, prefix="/model-session") api_router_v1.include_router(explainers, prefix="/explainer") api_router_v1.include_router(explorers, prefix="/explorer") @@ -39,6 +42,7 @@ api_router_v1.include_router(generative_process, prefix="/generative-process") api_router_v1.include_router(pipelines, prefix="/pipelines") api_router_v1.include_router(plugins, prefix="/plugin") +api_router_v1.include_router(prompts, prefix="/prompt") api_router_v1.include_router(notebook, prefix="/notebook") api_router_v1.include_router(metrics, prefix="/metrics") api_router_v1.include_router(hardware, prefix="/hardware") diff --git a/DashAI/back/api/api_v1/endpoints/components.py b/DashAI/back/api/api_v1/endpoints/components.py index 3c2d5cc53..2ab244c61 100644 --- a/DashAI/back/api/api_v1/endpoints/components.py +++ b/DashAI/back/api/api_v1/endpoints/components.py @@ -82,6 +82,21 @@ def _delete_class(component_dict: Dict[str, Any]) -> Dict[str, Any]: return {key: value for key, value in component_dict.items() if key != "class"} +def _enrich_with_flags( + component_dict: Dict[str, Any], include_flags: bool +) -> Dict[str, Any]: + """Append FLAGS from the component class when include_flags is True.""" + if not include_flags: + return component_dict + cls = component_dict.get("class") + if cls is None: + return component_dict + flags = getattr(cls, "FLAGS", []) + if not flags: + return component_dict + return {**component_dict, "flags": flags} + + @router.get("/") @inject async def get_components( @@ -91,6 +106,7 @@ async def get_components( related_component: Union[str, None] = None, component_parent: Union[str, None] = None, has_related_of_type: Union[str, None] = None, + include_flags: bool = Query(default=False), component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), ) -> List[Dict[str, Any]]: """Retrieve components from the register according to the provided parameters. @@ -230,10 +246,13 @@ async def get_components( components_with_related_type, ) - return [ - _filter_by_language(_delete_class(component_dict), accept_language) - for component_dict in selected_components.values() - ] + result = [] + for component_dict in selected_components.values(): + enriched = _enrich_with_flags(component_dict, include_flags) + cleaned = _delete_class(enriched) + localized = _filter_by_language(cleaned, accept_language) + result.append(localized) + return result @router.get("/{id}/") @@ -241,6 +260,7 @@ async def get_components( def get_component_by_id( id: str, accept_language: str | None = Header(default=None), + include_flags: bool = Query(default=False), component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), ) -> Dict[str, Any]: """Return a specific component using its id (the id is the component class name). @@ -271,7 +291,10 @@ def get_component_by_id( status_code=404, detail=f"Component {id} not found in the registry.", ) - return _filter_by_language(_delete_class(component_registry[id]), accept_language) + raw = component_registry[id] + enriched = _enrich_with_flags(raw, include_flags) + cleaned = _delete_class(enriched) + return _filter_by_language(cleaned, accept_language) @router.post("/", status_code=status.HTTP_201_CREATED) @@ -316,6 +339,30 @@ async def update_component() -> None: ) +@router.get("/{component_name}/children/") +async def get_child_components( + component_name: str, + recursive: bool = False, + component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), +): + """Get child components of a specific component. + + Args: + component_name (str): The name of the component to get children for. + recursive (bool): Whether to get child components recursively. + + Returns: + List[Dict[str, Any]]: A list of child component dictionaries. + """ + children_list = component_registry.get_child_components( + component_name, recursive=recursive + ) + # Remove "class" key from each child component + cleaned_children = [_delete_class(child) for child in children_list] + + return cleaned_children + + @router.get("/image/{component_name}/", status_code=status.HTTP_200_OK) async def get_component_image( component_name: str, diff --git a/DashAI/back/api/api_v1/endpoints/documents.py b/DashAI/back/api/api_v1/endpoints/documents.py new file mode 100644 index 000000000..89e16b58b --- /dev/null +++ b/DashAI/back/api/api_v1/endpoints/documents.py @@ -0,0 +1,215 @@ +import json +import logging +import os +from typing import Any, Dict, List +from urllib.parse import quote + +from fastapi import ( + APIRouter, + Depends, + File, + Form, + HTTPException, + Request, + Response, + UploadFile, + status, +) +from kink import di +from sqlalchemy.orm import sessionmaker + +from DashAI.back.api.api_v1.schemas import DocumentResponse +from DashAI.back.models.RAG.documents import DocumentFileType +from DashAI.back.models.RAG.exceptions import RAGDocumentFileTypeError +from DashAI.back.services.RAG.document_service import DocumentService + +router = APIRouter() +log = logging.getLogger(__name__) + + +base_url = "/api/v1/document" + + +@router.get("/", response_model=List[DocumentResponse]) +async def get_all_documents( + request: Request, + session_factory: sessionmaker = Depends(lambda: di["session_factory"]), +): + """Get all documents with file_url included.""" + with session_factory() as db: + base = str(request.base_url).rstrip("/") + return DocumentService(db).get_all(base_url=base) + + +@router.get("/{document_id}", response_model=DocumentResponse) +async def get_document( + document_id: int, + request: Request, + session_factory: sessionmaker = Depends(lambda: di["session_factory"]), +): + """Get metadata of a document by its ID.""" + with session_factory() as db: + try: + return DocumentService(db).get(document_id) + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail=str(e) + ) from e + + +@router.get("/{document_id}/download") +async def download_document( + document_id: int, + session_factory: sessionmaker = Depends(lambda: di["session_factory"]), +): + """Download the actual file content of a document.""" + + with session_factory() as db: + try: + content, media_type, filename = DocumentService(db).download(document_id) + encoded_name = quote(filename, safe="") + return Response( + content=content, + media_type=media_type, + headers={ + "Content-Disposition": ( + f"attachment; filename*=UTF-8''{encoded_name}" + ) + }, + ) + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail=str(e) + ) from e + + +@router.post("/", response_model=DocumentResponse, status_code=status.HTTP_201_CREATED) +async def upload_document( + file: UploadFile = File(...), + metadata: str = Form(...), + config: Dict[str, Any] = Depends(lambda: di["config"]), + session_factory: sessionmaker = Depends(lambda: di["session_factory"]), +): + """Upload a new document to the RAG system with file content and metadata.""" + docs_folder_path = config["DOCUMENTS_PATH"] + if not docs_folder_path.exists(): + raise HTTPException( + status_code=500, + detail=f"Documents folder {docs_folder_path} does not exist.", + ) + try: + metadata_dict = json.loads(metadata) + except json.JSONDecodeError: + raise HTTPException( + status_code=400, detail="Invalid JSON in metadata" + ) from None + file_name = metadata_dict.get("file_name") + if not file_name: + raise HTTPException( + status_code=400, detail="Missing required 'file_name' in metadata" + ) + optional_metadata = metadata_dict.get("optional_metadata", {}) + if not isinstance(optional_metadata, dict): + raise HTTPException( + status_code=400, detail="'optional_metadata' must be a dictionary" + ) + try: + content_bytes = await file.read() + except Exception: + raise HTTPException( + status_code=400, detail="Failed to read file content" + ) from None + ext = os.path.splitext(file_name)[1].lstrip(".") + try: + file_type = DocumentFileType(ext) + except ValueError: + raise HTTPException( + status_code=400, + detail=f"Unsupported file type: {ext}", + ) from None + with session_factory() as db: + return DocumentService(db).upload( + content_bytes, + file_name, + file_type, + str(docs_folder_path), + optional_metadata, + ) + + +@router.get("/session/{session_id}", response_model=List[DocumentResponse]) +async def get_documents_by_session( + session_id: int, + request: Request, + session_factory: sessionmaker = Depends(lambda: di["session_factory"]), +): + """Get all documents associated with a specific RAG session.""" + with session_factory() as db: + try: + base = str(request.base_url).rstrip("/") + return DocumentService(db).get_by_session(session_id, base_url=base) + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail=str(e) + ) from e + + +@router.get("/related-sessions/{document_id}", response_model=List[int]) +async def get_related_sessions( + document_id: int, + session_factory: sessionmaker = Depends(lambda: di["session_factory"]), +): + """Get all generative session IDs related to a specific document.""" + with session_factory() as db: + try: + return DocumentService(db).get_related_sessions(document_id) + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail=str(e) + ) from e + + +@router.delete("/{document_id}", status_code=status.HTTP_204_NO_CONTENT) +async def delete_document( + document_id: int, + session_factory: sessionmaker = Depends(lambda: di["session_factory"]), +): + """Delete a document from the RAG system by its ID.""" + + with session_factory() as db: + try: + DocumentService(db).delete(document_id) + return Response(status_code=status.HTTP_204_NO_CONTENT) + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail=str(e) + ) from e + + +@router.put("/{document_id}", response_model=DocumentResponse) +async def update_document_metadata( + document_id: int, + metadata: str = Form(...), + session_factory: sessionmaker = Depends(lambda: di["session_factory"]), +): + """Update a document's metadata.""" + with session_factory() as db: + try: + metadata_dict = json.loads(metadata) + except json.JSONDecodeError: + raise HTTPException( + status_code=400, detail="Invalid JSON in metadata" + ) from None + file_name = metadata_dict.get("file_name") + optional_metadata = metadata_dict.get("optional_metadata") + + try: + return DocumentService(db).update_metadata( + document_id, + file_name=file_name, + optional_metadata=optional_metadata, + ) + except RAGDocumentFileTypeError as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) + ) from e diff --git a/DashAI/back/api/api_v1/endpoints/explainers.py b/DashAI/back/api/api_v1/endpoints/explainers.py old mode 100755 new mode 100644 diff --git a/DashAI/back/api/api_v1/endpoints/generative_process.py b/DashAI/back/api/api_v1/endpoints/generative_process.py index 1a1791906..58cc3cf1a 100644 --- a/DashAI/back/api/api_v1/endpoints/generative_process.py +++ b/DashAI/back/api/api_v1/endpoints/generative_process.py @@ -121,7 +121,7 @@ async def upload_generative_process( @router.get("/{process_id}", status_code=status.HTTP_200_OK, response_model=None) async def get_generative_process( - process_id: str, + process_id: int, session_factory: "sessionmaker" = Depends(lambda: di["session_factory"]), component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), ): @@ -185,7 +185,7 @@ async def get_generative_process( "/{process_id}", status_code=status.HTTP_204_NO_CONTENT, response_model=None ) async def delete_generative_process( - process_id: str, + process_id: int, session_factory: "sessionmaker" = Depends(lambda: di["session_factory"]), ): """Delete a generative process by its ID. @@ -233,7 +233,7 @@ async def delete_generative_process( "/session/{session_id}", status_code=status.HTTP_200_OK, response_model=None ) async def get_generative_process_by_session_id( - session_id: str, + session_id: int, session_factory: "sessionmaker" = Depends(lambda: di["session_factory"]), component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), ): diff --git a/DashAI/back/api/api_v1/endpoints/generative_session.py b/DashAI/back/api/api_v1/endpoints/generative_session.py index 8211dbcf9..5f4554826 100644 --- a/DashAI/back/api/api_v1/endpoints/generative_session.py +++ b/DashAI/back/api/api_v1/endpoints/generative_session.py @@ -9,12 +9,22 @@ from DashAI.back.api.api_v1.schemas.generative_session_params import ( GenerativeSessionParams, ) +from DashAI.back.core.component_validation import validate_component_refs +from DashAI.back.core.schema_fields.utils import normalize_payload from DashAI.back.dependencies.database.models import ( GenerativeProcess, GenerativeSession, GenerativeSessionParameterHistory, ProcessData, ) +from DashAI.back.models.base_generative_model import BaseGenerativeModel +from DashAI.back.models.RAG.exceptions.base import RAGWorkflowError +from DashAI.back.services.RAG.cleanup_service import CleanupService +from DashAI.back.services.RAG.document_service import DocumentService +from DashAI.back.services.RAG.prompt_service import PromptService +from DashAI.back.services.RAG.RAG_setup_service import RAGSetupService +from DashAI.back.tasks.base_generative_task import BaseGenerativeTask +from DashAI.back.tasks.RAG_task import RAGTask if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker @@ -33,8 +43,6 @@ async def upload_generative_session( component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), ): """Create a new generative session and log the initial parameters in the history.""" - from DashAI.back.models.base_generative_model import BaseGenerativeModel - from DashAI.back.tasks.base_generative_task import BaseGenerativeTask with session_factory() as db: try: @@ -55,24 +63,84 @@ async def upload_generative_session( f"generative model.", ) - # Validate the model parameters + # Check if the task is registered try: - model_class.SCHEMA.model_validate(params.parameters) - except ValueError as e: + task_class = component_registry[params.task_name]["class"] + except KeyError as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail=f"Invalid parameters for model {params.model_name}: {e}", + detail=f"Task {params.task_name} is not registered.", ) from e - # Check if the task is registered + # RAG Task specific handling + # Frontend will send the ids of the documents to be used in the + # RAG session but RAG pipeline expects the documents paths of the + # backend-stored documents + if task_class == RAGTask: + docs = params.parameters.get("documents", []) + if not docs: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="'documents' must be a non-empty list.", + ) + if not all(isinstance(d, int) for d in docs): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="'documents' must be a list of integers.", + ) + try: + DocumentService(db).validate_exist(docs) + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=str(e), + ) from e + + # Normalise frontend properties wrapper + + params.parameters = normalize_payload(params.parameters) + + # RAG: resolve prompt_id BEFORE schema validation + if task_class == RAGTask: + prompt_service = PromptService(db, component_registry) + if "prompt_id" in params.parameters: + prompt_service.validate_prompt_exists( + params.parameters["prompt_id"] + ) + params.parameters["prompt"] = ( + prompt_service.resolve_prompt_id_to_component( + params.parameters["prompt_id"] + ) + ) + del params.parameters["prompt_id"] + + # Validate schema (now with prompt resolved if applicable) try: - task_class = component_registry[params.task_name]["class"] - except KeyError as e: + model_class.SCHEMA.model_validate(params.parameters) + except ValueError as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail=f"Task {params.task_name} is not registered.", + detail=f"Invalid parameters for model {params.model_name}: {e}", ) from e + # RAG: validate component refs and prompt template + if task_class == RAGTask: + component_errors = validate_component_refs( + params.parameters, component_registry + ) + if component_errors: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="; ".join(component_errors), + ) + + # Validate inline prompt template if present + prompt_ref = params.parameters.get("prompt", {}) + if prompt_ref and isinstance(prompt_ref, dict): + PromptService(db, component_registry).validate_component_ref( + prompt_ref + ) + # Check if the task is a subclass of BaseGenerativeTask if not issubclass(task_class, BaseGenerativeTask): raise HTTPException( @@ -257,6 +325,8 @@ async def delete_generative_session( detail=f"Generative session {session_id} does not exist in DB.", ) + old_parameters = dict(session.parameters or {}) + # Delete all the processes associated with the session processes = ( db.query(GenerativeProcess) @@ -286,7 +356,11 @@ async def delete_generative_session( db.delete(entry) # Finally, delete the session itself db.delete(session) + + CleanupService(db).cleanup_orphaned_resources(session_id, old_parameters) db.commit() + except HTTPException: + raise except exc.SQLAlchemyError as e: log.exception(e) raise HTTPException( @@ -300,7 +374,6 @@ async def delete_generative_session( detail="Internal server error", ) from e finally: - db.rollback() db.close() @@ -405,30 +478,8 @@ async def update_generative_session_params( session_id: int, new_params: dict, session_factory: "sessionmaker" = Depends(lambda: di["session_factory"]), + component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), ): - """Update the parameters of a generative session and log the change. - - Parameters - ---------- - session_id : int - The ID of the generative session to update. - new_params : dict - The new parameters to set for the generative session. - session_factory : Callable[..., ContextManager[Session]] - A factory that creates a context manager that handles a SQLAlchemy session. - - Returns - ------- - dict - A dictionary with the updated generative session. - - Raises - ------ - HTTPException - If the generative session does not exist or if there's an internal - database error. - """ - with session_factory() as db: try: session = db.get(GenerativeSession, session_id) @@ -438,8 +489,43 @@ async def update_generative_session_params( detail=f"Generative session {session_id} does not exist in DB.", ) - updated_parameters = {**session.parameters, **new_params} + old_parameters = dict(session.parameters or {}) + try: + task_class = component_registry[session.task_name]["class"] + except KeyError as e: + raise HTTPException( + status_code=404, + detail=( + f"Task '{session.task_name}' is not registered" + " in the component registry." + ), + ) from e + + # ── RAG-specific validation of new_params ── + if task_class is not None and task_class == RAGTask: + try: + config = di["config"] + normalized = RAGSetupService( + db, component_registry, config["RAG_PATH"] + ).validate_update_payload(new_params) + except (ValueError, RAGWorkflowError) as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) + ) from e + + # Merge validated new params into old params + updated_parameters = {**old_parameters, **normalized} + # Cleanup orphaned RAG resources + + CleanupService(db).cleanup_orphaned_resources( + session_id, old_parameters, updated_parameters + ) + else: + # Non-RAG update: simple merge without RAG validation + updated_parameters = {**old_parameters, **new_params} + + # ── Persist ── session_params_entry = GenerativeSessionParameterHistory( session_id=session.id, parameters=updated_parameters, @@ -449,12 +535,14 @@ async def update_generative_session_params( session.parameters = updated_parameters session.last_modified = datetime.now() - db.commit() db.refresh(session) return {"id": session.id, "parameters": session.parameters} + except HTTPException: + raise except exc.SQLAlchemyError as e: + db.rollback() log.exception(e) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -506,16 +594,18 @@ async def get_generative_session_parameters_history( .all() ) - # Convert the objects to dictionaries - return [ - { - "id": entry.id, - "session_id": entry.session_id, - "parameters": entry.parameters, - "modified_at": entry.modified_at, - } - for entry in parameters_history - ] + # Convert the objects to dictionaries (explicit loop for clarity) + history_list = [] + for entry in parameters_history: + history_list.append( + { + "id": entry.id, + "session_id": entry.session_id, + "parameters": entry.parameters, + "modified_at": entry.modified_at, + } + ) + return history_list except exc.SQLAlchemyError as e: log.exception(e) raise HTTPException( @@ -570,13 +660,18 @@ async def get_parameter_history_entry( .all() ) - parameters_history = [p.__dict__ for p in parameters_history] + params_history_dicts = [] + for p in parameters_history: + params_history_dicts.append(dict(p.__dict__)) events = [] - prev_params = parameters_history[0]["parameters"] + if not params_history_dicts: + return [] + + prev_params = params_history_dicts[0]["parameters"] - for i in range(1, len(parameters_history)): - curr = parameters_history[i] + for i in range(1, len(params_history_dicts)): + curr = params_history_dicts[i] curr_params = curr["parameters"] changes = [] diff --git a/DashAI/back/api/api_v1/endpoints/predict.py b/DashAI/back/api/api_v1/endpoints/predict.py index 188a14b02..66bb2bbcf 100644 --- a/DashAI/back/api/api_v1/endpoints/predict.py +++ b/DashAI/back/api/api_v1/endpoints/predict.py @@ -122,7 +122,7 @@ async def get_all_predictions( # Concatenate datasets to predictions dataset_dict = {dataset.id: dataset for dataset in datasets} for prediction in predictions: - prediction.dataset = dataset_dict.get(prediction.dataset_id, None) + prediction.dataset = dataset_dict.get(prediction.dataset_id) return predictions diff --git a/DashAI/back/api/api_v1/endpoints/prompts.py b/DashAI/back/api/api_v1/endpoints/prompts.py new file mode 100644 index 000000000..2cfc3f2bf --- /dev/null +++ b/DashAI/back/api/api_v1/endpoints/prompts.py @@ -0,0 +1,115 @@ +from fastapi import APIRouter, Depends, HTTPException, status +from kink import di +from sqlalchemy.orm import sessionmaker + +from DashAI.back.api.api_v1.schemas.RAG_prompt import ( + RAGPromptSchema, + RAGPromptUpdateSchema, +) +from DashAI.back.dependencies.database.models import GenerativeSession +from DashAI.back.dependencies.registry import ComponentRegistry +from DashAI.back.models.RAG.exceptions import ( + RAGDatabaseError, + RAGPromptError, + RAGPromptValidationError, +) +from DashAI.back.services.RAG.prompt_service import PromptService + +router = APIRouter() + + +@router.post("/", status_code=status.HTTP_201_CREATED) +async def create_RAG_prompt( # noqa: N802 + prompt: RAGPromptSchema, + component_registry: ComponentRegistry = Depends(lambda: di["component_registry"]), + session_factory: sessionmaker = Depends(lambda: di["session_factory"]), +): + """Create a new RAGPrompt entry in the database.""" + + with session_factory() as db: + try: + service = PromptService(db, component_registry) + result = service.create( + prompt.class_name, prompt.name, prompt.parameters or {} + ) + return {"id": result.id} + except RAGPromptValidationError as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) + ) from e + except RAGDatabaseError as e: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="A prompt with these parameters already exists.", + ) from e + except RAGPromptError as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) + ) from e + + +@router.patch("/{prompt_id}", status_code=status.HTTP_200_OK) +async def update_RAG_prompt( # noqa: N802 + prompt_id: int, + prompt: RAGPromptUpdateSchema, + component_registry: ComponentRegistry = Depends(lambda: di["component_registry"]), + session_factory: sessionmaker = Depends(lambda: di["session_factory"]), +): + """Update an existing prompt in place.""" + + with session_factory() as db: + try: + service = PromptService(db, component_registry) + result = service.update( + prompt_id, + name=prompt.name, + parameters=prompt.parameters, + ) + return result + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail=str(e) + ) from e + + +@router.post("/{prompt_id}/sessions/{session_id}", status_code=status.HTTP_201_CREATED) +async def update_RAG_prompt_for_session( # noqa: N802 + prompt_id: int, + session_id: int, + prompt: RAGPromptUpdateSchema, + component_registry: ComponentRegistry = Depends(lambda: di["component_registry"]), + session_factory: sessionmaker = Depends(lambda: di["session_factory"]), +): + """Create a session-scoped copy of a prompt and attach it to the session.""" + + with session_factory() as db: + try: + service = PromptService(db, component_registry) + prompt_result = service.create_session_copy( + prompt_id, + session_id, + parameters=prompt.parameters, + name=prompt.name, + ) + session = db.get(GenerativeSession, session_id) + return { + "prompt": prompt_result, + "session_id": session_id, + "parameters": session.parameters if session else None, + } + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail=str(e) + ) from e + + +@router.get("/", status_code=status.HTTP_200_OK) +async def get_all_prompts( + session_factory: sessionmaker = Depends(lambda: di["session_factory"]), +): + """Get all available RAG prompts.""" + + component_registry = di["component_registry"] + with session_factory() as db: + service = PromptService(db, component_registry) + return service.get_all() diff --git a/DashAI/back/api/api_v1/schemas/RAG_prompt.py b/DashAI/back/api/api_v1/schemas/RAG_prompt.py new file mode 100644 index 000000000..898f3441b --- /dev/null +++ b/DashAI/back/api/api_v1/schemas/RAG_prompt.py @@ -0,0 +1,29 @@ +from typing import Any, Dict, Optional + +from pydantic import BaseModel + + +class RAGPromptSchema(BaseModel): + """Schema for creating a RAG prompt. + + Attributes: + class_name: Registered prompt component class name. + name: Human-readable name for the prompt. + parameters: Optional configuration dict including template(s). + """ + + class_name: str + name: str + parameters: Optional[Dict[str, Any]] = None + + +class RAGPromptUpdateSchema(BaseModel): + """Schema for updating an existing RAG prompt. + + Attributes: + name: Optional new name for the prompt. + parameters: Optional new configuration dict. + """ + + name: Optional[str] = None + parameters: Optional[Dict[str, Any]] = None diff --git a/DashAI/back/api/api_v1/schemas/__init__.py b/DashAI/back/api/api_v1/schemas/__init__.py index e69de29bb..76f42ae53 100644 --- a/DashAI/back/api/api_v1/schemas/__init__.py +++ b/DashAI/back/api/api_v1/schemas/__init__.py @@ -0,0 +1 @@ +from DashAI.back.api.api_v1.schemas.document import DocumentResponse diff --git a/DashAI/back/api/api_v1/schemas/document.py b/DashAI/back/api/api_v1/schemas/document.py new file mode 100644 index 000000000..a19768d02 --- /dev/null +++ b/DashAI/back/api/api_v1/schemas/document.py @@ -0,0 +1,16 @@ +from datetime import datetime +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel + + +class DocumentResponse(BaseModel): + id: int + file_name: str + file_type: str + file_hash: str + created: datetime + last_modified: datetime + optional_metadata: Optional[Dict[str, Any]] + related_sessions: List[int] | None + file_url: str diff --git a/DashAI/back/api/api_v1/schemas/job_params.py b/DashAI/back/api/api_v1/schemas/job_params.py index 167e54dab..816e0f22e 100644 --- a/DashAI/back/api/api_v1/schemas/job_params.py +++ b/DashAI/back/api/api_v1/schemas/job_params.py @@ -15,5 +15,6 @@ class JobParams(BaseModel): "ConverterJob", "GenerativeJob", "PipelineJob", + "RAGJob", ] kwargs: dict diff --git a/DashAI/back/api/api_v1/schemas/rag_prompt.py b/DashAI/back/api/api_v1/schemas/rag_prompt.py new file mode 100644 index 000000000..898f3441b --- /dev/null +++ b/DashAI/back/api/api_v1/schemas/rag_prompt.py @@ -0,0 +1,29 @@ +from typing import Any, Dict, Optional + +from pydantic import BaseModel + + +class RAGPromptSchema(BaseModel): + """Schema for creating a RAG prompt. + + Attributes: + class_name: Registered prompt component class name. + name: Human-readable name for the prompt. + parameters: Optional configuration dict including template(s). + """ + + class_name: str + name: str + parameters: Optional[Dict[str, Any]] = None + + +class RAGPromptUpdateSchema(BaseModel): + """Schema for updating an existing RAG prompt. + + Attributes: + name: Optional new name for the prompt. + parameters: Optional new configuration dict. + """ + + name: Optional[str] = None + parameters: Optional[Dict[str, Any]] = None diff --git a/DashAI/back/app.py b/DashAI/back/app.py index f4bc5088a..a6160f992 100644 --- a/DashAI/back/app.py +++ b/DashAI/back/app.py @@ -80,6 +80,7 @@ def create_app( _create_path_if_not_exists(config["EXPLANATIONS_PATH"]) _create_path_if_not_exists(config["NOTEBOOK_PATH"]) _create_path_if_not_exists(config["RUNS_PATH"]) + _create_path_if_not_exists(config["DOCUMENTS_PATH"]) _create_path_if_not_exists(config["DATAFILE_PATH"]) logger.debug("3. Creating app container and setting up dependency injection.") diff --git a/DashAI/back/config.py b/DashAI/back/config.py index 728c8a760..5259c9dd5 100644 --- a/DashAI/back/config.py +++ b/DashAI/back/config.py @@ -20,5 +20,8 @@ class DefaultSettings(BaseSettings): IMAGES_PATH: str = "images" RUNS_PATH: str = "runs" EXPLANATIONS_PATH: str = "explanations" + EXPLORATIONS_PATH: str = "explorations" + DOCUMENTS_PATH: str = "documents" + RAG_PATH: str = "RAG" NOTEBOOK_PATH: str = "notebook" DATAFILE_PATH: str = "datafiles" diff --git a/DashAI/back/config_object.py b/DashAI/back/config_object.py index b994e8e93..47cd78d0b 100644 --- a/DashAI/back/config_object.py +++ b/DashAI/back/config_object.py @@ -47,5 +47,12 @@ def validate_and_transform(self, raw_data: dict) -> dict: ValidationError If the given data does not follow the schema associated with the model. """ - schema_instance = self.SCHEMA.model_validate(raw_data) + try: + from pydantic import ValidationError + + schema_instance = self.SCHEMA.model_validate(raw_data) + except ValidationError as e: + # print full error message + print(e.json()) + raise e return fill_objects(schema_instance) diff --git a/DashAI/back/converters/hugging_face/embedding.py b/DashAI/back/converters/hugging_face/embedding.py index cbc726d35..d3a348ee3 100644 --- a/DashAI/back/converters/hugging_face/embedding.py +++ b/DashAI/back/converters/hugging_face/embedding.py @@ -1,6 +1,10 @@ """HuggingFace embedding converter with lazy-loaded dependencies.""" -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, List + +import numpy as np +import torch +from transformers import AutoModel, AutoTokenizer from DashAI.back.converters.category.advanced_preprocessing import ( AdvancedPreprocessingConverter, @@ -178,12 +182,41 @@ def get_output_type(self, column_name: str = None) -> DashAIDataType: def _load_model(self): """Load the embedding model and tokenizer.""" - from transformers import AutoModel, AutoTokenizer self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) self.model = AutoModel.from_pretrained(self.model_name).to(self.device) self.model.eval() + def _encode_texts(self, texts: List[str]) -> List[np.ndarray]: + """Encode a list of texts into embeddings.""" + # Tokenize + encoded = self.tokenizer( + texts, + padding=True, + truncation=True, + max_length=self.max_length, + return_tensors="pt", + ) + + # Move to device + encoded = {k: v.to(self.device) for k, v in encoded.items()} + + # Get embeddings + with torch.no_grad(): + outputs = self.model(**encoded) + hidden_states = outputs.last_hidden_state + + # Apply pooling strategy + if self.pooling_strategy == "mean": + embeddings = torch.mean(hidden_states, dim=1) + elif self.pooling_strategy == "cls": + embeddings = hidden_states[:, 0] + else: # max pooling + embeddings = torch.max(hidden_states, dim=1)[0] + + embeddings_np = embeddings.cpu().numpy() + return embeddings_np + def _process_batch(self, batch: "DashAIDataset") -> "DashAIDataset": """Encode a batch of text columns into dense embedding vectors. @@ -203,7 +236,6 @@ def _process_batch(self, batch: "DashAIDataset") -> "DashAIDataset": dense embedding vector column(s). """ import pyarrow as pa - import torch from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset @@ -213,32 +245,7 @@ def _process_batch(self, batch: "DashAIDataset") -> "DashAIDataset": # Get text data from dataset texts = [row[column] if row[column] is not None else "" for row in batch] - # Tokenize - encoded = self.tokenizer( - texts, - padding=True, - truncation=True, - max_length=self.max_length, - return_tensors="pt", - ) - - # Move to device - encoded = {k: v.to(self.device) for k, v in encoded.items()} - - # Get embeddings - with torch.no_grad(): - outputs = self.model(**encoded) - hidden_states = outputs.last_hidden_state - - # Apply pooling strategy - if self.pooling_strategy == "mean": - embeddings = torch.mean(hidden_states, dim=1) - elif self.pooling_strategy == "cls": - embeddings = hidden_states[:, 0] - else: # max pooling - embeddings = torch.max(hidden_states, dim=1)[0] - - embeddings_np = embeddings.cpu().numpy() + embeddings_np = self._encode_texts(texts) # Append one column per embedding dimension for i in range(embeddings_np.shape[1]): diff --git a/DashAI/back/core/component_validation.py b/DashAI/back/core/component_validation.py new file mode 100644 index 000000000..77a895b79 --- /dev/null +++ b/DashAI/back/core/component_validation.py @@ -0,0 +1,78 @@ +"""Generic component reference validation. + +Recursively walks a parameters dict to find all ``{component, params}`` +references and validates them against the component registry. + +Works for any task (RAG, generative, etc.) — any parameter structure +that uses ``component_field`` produces the ``{component, params}`` pattern. +""" + +from typing import Any + +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry + + +def find_component_refs( + data: Any, path: str = "" +) -> list[tuple[str, str, dict[str, Any]]]: + """Recursively find all component references in a parameters structure. + + A component reference is any ``dict`` that has both ``"component"`` (``str``) + and ``"params"`` (``dict``) keys. Returns a list of + ``(json_path, component_name, params)`` tuples. + + Parameters + ---------- + data : Any + The parameters structure to search (typically a ``dict`` or ``list``). + path : str + Dot-separated JSON path for error reporting (used internally during recursion). + + Returns + ------- + list[tuple[str, str, dict[str, Any]]] + Each entry: ``("prompt", "DefaultRAGGenerationPrompt", {...})``, + ``("chunking_model", "CharacterChunkModel", {...})``, etc. + """ + refs: list[tuple[str, str, dict[str, Any]]] = [] + if isinstance(data, dict): + if ( + "component" in data + and "params" in data + and isinstance(data["component"], str) + and isinstance(data["params"], dict) + ): + refs.append((path, data["component"], data["params"])) + for key, value in data.items(): + child_path = f"{path}.{key}" if path else key + refs.extend(find_component_refs(value, child_path)) + elif isinstance(data, list): + for i, item in enumerate(data): + refs.extend(find_component_refs(item, f"{path}[{i}]")) + return refs + + +def validate_component_refs(data: Any, registry: ComponentRegistry) -> list[str]: + """Validate that all component references exist in the registry. + + Walks ``data`` recursively via :func:`find_component_refs` and checks + each component name against ``registry``. + + Parameters + ---------- + data : Any + Parameters structure to validate. + registry : ComponentRegistry + The application's component registry. + + Returns + ------- + list[str] + Error messages (empty list means all references are valid). + Each message like ``"'NonExistent' at 'chunking_model' not registered."`` + """ + errors: list[str] = [] + for path, name, _params in find_component_refs(data): + if name not in registry: + errors.append(f"Component '{name}' at '{path}' is not registered.") + return errors diff --git a/DashAI/back/core/schema_fields/base_schema.py b/DashAI/back/core/schema_fields/base_schema.py index 3371ffd4b..0c31e425f 100644 --- a/DashAI/back/core/schema_fields/base_schema.py +++ b/DashAI/back/core/schema_fields/base_schema.py @@ -45,4 +45,14 @@ def replace_defs_in_schema(schema: dict): class BaseSchema(BaseModel): - pass + """@classmethod + def model_validate(cls, raw_data: dict): + schema_fields = cls.model_fields + for field_name, field in schema_fields.items(): + if field_name not in raw_data: + continue + if field.annotation._name == 'Optional': + if raw_data[field_name] == "": + raw_data[field_name] = None + + return super().model_validate(raw_data)""" diff --git a/DashAI/back/core/schema_fields/utils.py b/DashAI/back/core/schema_fields/utils.py index 17ac63b39..f6ec0a601 100644 --- a/DashAI/back/core/schema_fields/utils.py +++ b/DashAI/back/core/schema_fields/utils.py @@ -7,6 +7,70 @@ from DashAI.back.dependencies.registry import ComponentRegistry +def normalize_payload(data): + """Recursively normalizes the frontend ``properties`` wrapper into the + canonical ``{component, params}`` format. + + This is the factored-out version of the unwrapping logic found in + ``model_factory._process_param`` (approximately line 162). + + The frontend auto-generated forms (``generateInitialValues`` in + ``utils/schema.js``) produce a ``properties`` wrapper structure: + + .. code-block:: python + + { + "properties": { + "component": "ParentModelName", + "params": { + "comp": { + "component": "ActualModelName", + "params": { ... sub‑params ... }, + }, + }, + }, + } + + This function converts it into the canonical form that + ``fill_objects()`` and ``model_class.SCHEMA.model_validate()`` expect: + + .. code-block:: python + + {"component": "ActualModelName", "params": { ... sub‑params ... }} + + Any scalar, list, or already-canonical dict values are returned + unchanged. + + Parameters + ---------- + data : Any + Arbitrary JSON-like data (output of ``model_dump()`` or raw + frontend payload) that may contain ``properties`` wrappers. + + Returns + ------- + Any + The same data with all ``properties`` wrappers replaced by the + canonical ``{component, params}`` form. + """ + if isinstance(data, dict): + if "properties" in data and len(data) == 1: + inner = data["properties"] + if isinstance(inner, dict) and "component" in inner: + raw_params = inner.get("params", {}) + comp = raw_params.get("comp", {}) + if isinstance(comp, dict) and "component" in comp: + return { + "component": comp["component"], + "params": normalize_payload(comp.get("params", {})), + } + return normalize_payload(inner) + return {k: normalize_payload(v) for k, v in data.items()} + if isinstance(data, list): + return [normalize_payload(item) for item in data] + return data + + @inject def fill_objects( schema_instance: "BaseSchema", diff --git a/DashAI/back/dataset_sources/openml_dataset_source.py b/DashAI/back/dataset_sources/openml_dataset_source.py index 202e398c6..af9e536f4 100644 --- a/DashAI/back/dataset_sources/openml_dataset_source.py +++ b/DashAI/back/dataset_sources/openml_dataset_source.py @@ -197,7 +197,7 @@ def _meta(did: str) -> tuple[str, tuple[str, ...]] | None: metas = list(pool.map(_meta, ids)) entries = [] - for row, did, meta in zip(rows, ids, metas): + for row, did, meta in zip(rows, ids, metas, strict=True): description = "" dataset_tags: list[str] = [] if meta is not None: diff --git a/DashAI/back/dependencies/config_builder.py b/DashAI/back/dependencies/config_builder.py index 3323340b9..69e729451 100644 --- a/DashAI/back/dependencies/config_builder.py +++ b/DashAI/back/dependencies/config_builder.py @@ -63,5 +63,7 @@ def build_config_dict( config["BACK_PATH"] = pathlib.Path(config["BACK_PATH"]).absolute() config["LOGGING_LEVEL"] = getattr(logging, logging_level) config["INITIAL_COMPONENTS"] = get_initial_components() + config["DOCUMENTS_PATH"] = local_path / config["DOCUMENTS_PATH"] + config["RAG_PATH"] = local_path / config["RAG_PATH"] return config diff --git a/DashAI/back/dependencies/database/models.py b/DashAI/back/dependencies/database/models.py index cdf57d968..fd0e0817f 100644 --- a/DashAI/back/dependencies/database/models.py +++ b/DashAI/back/dependencies/database/models.py @@ -307,7 +307,7 @@ class Plugin(Base): Boolean, nullable=False, default=False, server_default="0" ) installed_version: Mapped[str] = mapped_column(String, nullable=False) - lastest_version: Mapped[str] = mapped_column(String, nullable=False) + latest_version: Mapped[str] = mapped_column(String, nullable=False) tags: Mapped[List["Tag"]] = relationship( back_populates="plugin", cascade="all, delete", lazy="selectin" ) @@ -547,6 +547,11 @@ class GenerativeSession(Base): "GenerativeProcess", cascade="all, delete-orphan", back_populates="session" ) + # Relationship with RAGDocumentPipelineSessionLink + pipeline_links: Mapped[List["RAGDocumentPipelineSessionLink"]] = relationship( + back_populates="session" + ) + class Pipeline(Base): __tablename__ = "pipeline" @@ -746,6 +751,565 @@ def delete_result(self) -> None: self.end_time = None +""" +RAG tables +""" + + +class Document(Base): + __tablename__ = "document" + """ + Table to store all the information about a document. + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + file_name: Mapped[str] = mapped_column(String, nullable=False) + file_type: Mapped[str] = mapped_column(String, nullable=False) + file_path: Mapped[str] = mapped_column(String, nullable=False) + file_hash: Mapped[str] = mapped_column(String, nullable=False, unique=True) + optional_metadata: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) + created: Mapped[DateTime] = mapped_column(DateTime, default=datetime.now) + last_modified: Mapped[DateTime] = mapped_column( + DateTime, + default=datetime.now, + onupdate=datetime.now, + ) + + # Create a relationship for the related sessions and related chunks + pipeline_links: Mapped[List["RAGDocumentPipelineSessionLink"]] = relationship( + "RAGDocumentPipelineSessionLink", + cascade="all, delete-orphan", + back_populates="document", + ) + + chunks: Mapped[List["Chunk"]] = relationship( + "Chunk", + back_populates="document", + cascade="all, delete-orphan", + ) + + embedding_matrices: Mapped[List["RAGEmbeddingMatrix"]] = relationship( + "RAGEmbeddingMatrix", back_populates="document", cascade="all, delete-orphan" + ) + + @property + def get_related_sessions(self) -> List["GenerativeSession"]: + """Return a list of sessions related to the document.""" + return [link.session for link in self.pipeline_links] + + @property + def get_related_pipelines(self) -> List["RAGPipeline"]: + """Return a list of pipelines related to the document.""" + return [link.pipeline for link in self.pipeline_links] + + def get_embedding_matrix( + self, chunk_set_id: int, embedding_model_id: int + ) -> Optional["RAGEmbeddingMatrix"]: + for matrix in self.embedding_matrices: + if ( + matrix.chunk_set_id == chunk_set_id + and matrix.embedding_model_id == embedding_model_id + ): + return matrix + return None + + +class Chunk(Base): + __tablename__ = "chunk" + """ + Table to store all the information about a chunk. + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + chunk_set_id: Mapped[int] = mapped_column( + ForeignKey("RAG_chunk_set.id", ondelete="CASCADE"), nullable=False + ) + document_id: Mapped[int] = mapped_column( + ForeignKey("document.id", ondelete="CASCADE"), nullable=False + ) + chunk_index: Mapped[int] = mapped_column(Integer, nullable=False) + text: Mapped[str] = mapped_column(Text, nullable=False) + chunk_metadata: Mapped[Optional[dict]] = mapped_column( + "metadata", JSON, nullable=True + ) + + document: Mapped["Document"] = relationship("Document", back_populates="chunks") + chunk_set: Mapped["RAGChunkSet"] = relationship( + "RAGChunkSet", back_populates="chunks" + ) + + __table_args__ = ( + UniqueConstraint( + "chunk_set_id", + "document_id", + "chunk_index", + name="uix_chunk_set_doc_index", + ), + ) + + +class RAGChunkSet(Base): + __tablename__ = "RAG_chunk_set" + """ + Canonical identity for a set of chunks produced by a processing pipeline. + Two sessions with the same documents + same pipeline config share the same + chunk set (same signature → same row). + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + signature: Mapped[str] = mapped_column(String, nullable=False, unique=True) + parameters: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True) + + chunks: Mapped[List["Chunk"]] = relationship( + "Chunk", + back_populates="chunk_set", + cascade="all, delete-orphan", + ) + documents: Mapped[List["RAGChunkSetDocument"]] = relationship( + "RAGChunkSetDocument", + back_populates="chunk_set", + cascade="all, delete-orphan", + ) + embedding_matrices: Mapped[List["RAGEmbeddingMatrix"]] = relationship( + "RAGEmbeddingMatrix", + back_populates="chunk_set", + ) + retriever_links: Mapped[List["RAGRetrieverChunkSet"]] = relationship( + "RAGRetrieverChunkSet", + back_populates="chunk_set", + cascade="all, delete-orphan", + ) + + +class RAGChunkSetDocument(Base): + __tablename__ = "RAG_chunk_set_document" + """Which documents belong to a chunk set.""" + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + chunk_set_id: Mapped[int] = mapped_column( + ForeignKey("RAG_chunk_set.id", ondelete="CASCADE"), nullable=False + ) + document_id: Mapped[int] = mapped_column( + ForeignKey("document.id", ondelete="CASCADE"), nullable=False + ) + + chunk_set: Mapped["RAGChunkSet"] = relationship( + "RAGChunkSet", + back_populates="documents", + ) + document: Mapped["Document"] = relationship("Document") + + __table_args__ = ( + UniqueConstraint( + "chunk_set_id", + "document_id", + name="uix_chunk_set_document", + ), + ) + + +class RAGRetrieverChunkSet(Base): + __tablename__ = "RAG_retriever_chunk_set" + """Links a retriever to the chunk set it operates on.""" + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + retriever_id: Mapped[int] = mapped_column( + ForeignKey("RAG_retriever.id", ondelete="CASCADE"), nullable=False + ) + chunk_set_id: Mapped[int] = mapped_column( + ForeignKey("RAG_chunk_set.id", ondelete="CASCADE"), nullable=False + ) + + chunk_set: Mapped["RAGChunkSet"] = relationship( + "RAGChunkSet", + back_populates="retriever_links", + ) + + __table_args__ = ( + UniqueConstraint( + "retriever_id", + "chunk_set_id", + name="uix_retriever_chunk_set", + ), + ) + + +class RAGPrompt(Base): + __tablename__ = "RAG_prompt" + """ + Table to store all the information about a RAG prompt. + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + class_name: Mapped[str] = mapped_column(String, nullable=False) + name: Mapped[str] = mapped_column(String, nullable=True) + parameters: Mapped[JSON] = mapped_column(JSON, nullable=True) + created: Mapped[DateTime] = mapped_column(DateTime, default=datetime.now) + last_modified: Mapped[DateTime] = mapped_column( + DateTime, + default=datetime.now, + onupdate=datetime.now, + ) + + # Relationship with RAGPipeline + pipelines: Mapped[List["RAGPipeline"]] = relationship( + "RAGPipeline", back_populates="prompt", cascade="all, delete-orphan" + ) + + __table_args__ = ( + UniqueConstraint( + "class_name", + "parameters", + name="uix_RAG_prompt_class_params", + ), + ) + + +class RAGGenerationModel(Base): + __tablename__ = "RAG_generation_model" + """ + Table to store all the information about a generation model. + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + class_name: Mapped[str] = mapped_column(String, nullable=False) + parameters: Mapped[JSON] = mapped_column(JSON, nullable=True) + + # Relationships + pipelines: Mapped[List["RAGPipeline"]] = relationship( + back_populates="generation_model" + ) + + __table_args__ = ( + UniqueConstraint( + "class_name", + "parameters", + name="uix_RAG_gen_model_class_params", + ), + ) + + +class RAGPipeline(Base): + __tablename__ = "RAG_pipeline" + """ + Table to store all the information about a RAG pipeline. + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + session_id: Mapped[int] = mapped_column( + ForeignKey("generative_session.id", ondelete="CASCADE"), nullable=False + ) + name: Mapped[str] = mapped_column(String, nullable=False) + description: Mapped[str] = mapped_column(String, nullable=True) + parameters: Mapped[JSON] = mapped_column(JSON, nullable=True) + + chunking_model_id: Mapped[Optional[int]] = mapped_column( + ForeignKey("RAG_chunking_model.id", ondelete="CASCADE"), nullable=True + ) + prompt_id: Mapped[Optional[int]] = mapped_column( + ForeignKey("RAG_prompt.id", ondelete="CASCADE"), nullable=True + ) + generation_model_id: Mapped[Optional[int]] = mapped_column( + ForeignKey("RAG_generation_model.id", ondelete="CASCADE"), nullable=True + ) + + # Relationships + chunking_model: Mapped["RAGChunkingModel"] = relationship( + "RAGChunkingModel", back_populates="pipelines" + ) + retriever: Mapped["RAGRetriever"] = relationship( + "RAGRetriever", + back_populates="pipeline", + uselist=False, + ) + prompt: Mapped["RAGPrompt"] = relationship("RAGPrompt", back_populates="pipelines") + generation_model: Mapped["RAGGenerationModel"] = relationship( + "RAGGenerationModel", back_populates="pipelines" + ) + pipeline_links: Mapped[List["RAGDocumentPipelineSessionLink"]] = relationship( + "RAGDocumentPipelineSessionLink", back_populates="pipeline" + ) + + +class RAGChunkingModel(Base): + __tablename__ = "RAG_chunking_model" + """ + Table to store all the information about a chunking model. + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + class_name: Mapped[str] = mapped_column(String, nullable=False) + parameters: Mapped[JSON] = mapped_column(JSON, nullable=True) + + pipelines: Mapped[List["RAGPipeline"]] = relationship( + "RAGPipeline", back_populates="chunking_model" + ) + + __table_args__ = ( + UniqueConstraint( + "class_name", + "parameters", + name="uix_RAG_chunking_model_class_params", + ), + ) + + +class RAGRetriever(Base): + __tablename__ = "RAG_retriever" + """ + Canonical identity row for every retriever — unit (sparse/dense) or composite. + + Sub-tables (RAGSparseRetriever, RAGDenseRetriever) reference this row + via bridge_id. Composite children are stored in RAG_retriever_child. + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + class_name: Mapped[str] = mapped_column(String, nullable=False) + pipeline_id: Mapped[int] = mapped_column( + ForeignKey("RAG_pipeline.id", ondelete="CASCADE"), nullable=False + ) + + pipeline: Mapped["RAGPipeline"] = relationship(back_populates="retriever") + children: Mapped[List["RAGRetrieverChild"]] = relationship( + "RAGRetrieverChild", + foreign_keys="RAGRetrieverChild.parent_id", + back_populates="parent", + cascade="all, delete-orphan", + ) + sparse_detail: Mapped[Optional["RAGSparseRetriever"]] = relationship( + "RAGSparseRetriever", + back_populates="bridge", + uselist=False, + ) + dense_detail: Mapped[Optional["RAGDenseRetriever"]] = relationship( + "RAGDenseRetriever", + back_populates="bridge", + uselist=False, + ) + + +class RAGRetrieverChild(Base): + __tablename__ = "RAG_retriever_child" + """ + Link table for composite retrievers: maps a parent composite to its + ordered child retrievers. Both parent and child reference RAG_retriever.id. + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + parent_id: Mapped[int] = mapped_column( + ForeignKey("RAG_retriever.id", ondelete="CASCADE"), + nullable=False, + ) + child_id: Mapped[int] = mapped_column( + ForeignKey("RAG_retriever.id", ondelete="CASCADE"), + nullable=False, + ) + child_order: Mapped[int] = mapped_column(Integer, nullable=False) + + parent: Mapped["RAGRetriever"] = relationship( + "RAGRetriever", + foreign_keys=[parent_id], + back_populates="children", + ) + + __table_args__ = ( + UniqueConstraint( + "parent_id", + "child_order", + name="uix_retriever_child_order", + ), + ) + + +class RAGSparseRetriever(Base): + __tablename__ = "RAG_sparse_retriever" + """ + Table to store all the information about a sparse retriever. + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + bridge_id: Mapped[int] = mapped_column( + ForeignKey("RAG_retriever.id", ondelete="CASCADE"), + nullable=False, + ) + chunk_set_id: Mapped[int] = mapped_column( + ForeignKey("RAG_chunk_set.id", ondelete="CASCADE"), + nullable=False, + ) + class_name: Mapped[str] = mapped_column(String, nullable=False) + parameters: Mapped[JSON] = mapped_column(JSON, nullable=True) + storage_folder: Mapped[str] = mapped_column(String, nullable=False) + + bridge: Mapped["RAGRetriever"] = relationship( + "RAGRetriever", + back_populates="sparse_detail", + ) + chunk_set: Mapped["RAGChunkSet"] = relationship("RAGChunkSet") + + __table_args__ = ( + UniqueConstraint( + "class_name", + "parameters", + "chunk_set_id", + name="uix_RAG_sparse_retriever", + ), + ) + + +class RAGDenseRetriever(Base): + __tablename__ = "RAG_dense_retriever" + """ + Table to store all the information about a dense retriever. + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + bridge_id: Mapped[int] = mapped_column( + ForeignKey("RAG_retriever.id", ondelete="CASCADE"), + nullable=False, + ) + chunk_set_id: Mapped[int] = mapped_column( + ForeignKey("RAG_chunk_set.id", ondelete="CASCADE"), + nullable=False, + ) + class_name: Mapped[str] = mapped_column(String, nullable=False) + parameters: Mapped[JSON] = mapped_column(JSON, nullable=True) + embedding_model_id: Mapped[int] = mapped_column( + ForeignKey("RAG_embedding_model.id", ondelete="CASCADE"), nullable=False + ) + bridge: Mapped["RAGRetriever"] = relationship( + "RAGRetriever", + back_populates="dense_detail", + ) + chunk_set: Mapped["RAGChunkSet"] = relationship("RAGChunkSet") + embedding_model: Mapped["RAGEmbeddingModel"] = relationship( + "RAGEmbeddingModel", back_populates="dense_retrievers" + ) + + __table_args__ = ( + UniqueConstraint( + "class_name", + "parameters", + "chunk_set_id", + "embedding_model_id", + name="uix_RAG_dense_retriever", + ), + ) + + +class RAGEmbeddingModel(Base): + __tablename__ = "RAG_embedding_model" + """ + Table to store embedding model configurations. + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + class_name: Mapped[str] = mapped_column(String, nullable=False) + parameters: Mapped[JSON] = mapped_column(JSON, nullable=True) + + # Relationships + embedding_matrices: Mapped[List["RAGEmbeddingMatrix"]] = relationship( + back_populates="embedding_model", cascade="all, delete-orphan" + ) + dense_retrievers: Mapped[List["RAGDenseRetriever"]] = relationship( + back_populates="embedding_model" + ) + + __table_args__ = ( + UniqueConstraint( + "class_name", + "parameters", + name="uix_RAG_embedding_model_class_params", + ), + ) + + +class RAGEmbeddingMatrix(Base): + __tablename__ = "RAG_embedding_matrix" + """ + Table to store embedding matrices for each unique combination of + document, chunk set, and embedding model. + """ + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + document_id: Mapped[int] = mapped_column( + ForeignKey("document.id", ondelete="CASCADE"), nullable=False + ) + chunk_set_id: Mapped[int] = mapped_column( + ForeignKey("RAG_chunk_set.id", ondelete="CASCADE"), nullable=False + ) + embedding_model_id: Mapped[int] = mapped_column( + ForeignKey("RAG_embedding_model.id", ondelete="CASCADE"), nullable=False + ) + storage_folder: Mapped[str] = mapped_column(String, nullable=False) + matrix_shape: Mapped[List[int]] = mapped_column(JSON, nullable=False) + created: Mapped[DateTime] = mapped_column(DateTime, default=datetime.now) + last_modified: Mapped[DateTime] = mapped_column( + DateTime, + default=datetime.now, + onupdate=datetime.now, + ) + + document: Mapped["Document"] = relationship( + "Document", back_populates="embedding_matrices" + ) + chunk_set: Mapped["RAGChunkSet"] = relationship( + "RAGChunkSet", + back_populates="embedding_matrices", + ) + embedding_model: Mapped["RAGEmbeddingModel"] = relationship( + "RAGEmbeddingModel", back_populates="embedding_matrices" + ) + + __table_args__ = ( + UniqueConstraint( + "document_id", + "chunk_set_id", + "embedding_model_id", + name="uix_document_chunk_set_embedding", + ), + ) + + @property + def num_chunks(self) -> int: + """Return the number of chunks in this embedding matrix.""" + return self.matrix_shape[0] if self.matrix_shape else 0 + + @property + def embedding_dimension(self) -> int: + """Return the embedding dimension.""" + return self.matrix_shape[1] if len(self.matrix_shape) > 1 else 0 + + @classmethod + def get_by_tuple( + cls, session, document_id: int, chunk_set_id: int, embedding_model_id: int + ) -> Optional["RAGEmbeddingMatrix"]: + return ( + session.query(cls) + .filter( + cls.document_id == document_id, + cls.chunk_set_id == chunk_set_id, + cls.embedding_model_id == embedding_model_id, + ) + .first() + ) + + +""" +RAG relationship tables +""" + + +class RAGDocumentPipelineSessionLink(Base): + __tablename__ = "RAG_document_pipeline_session_link" + + id: Mapped[int] = mapped_column(primary_key=True) + document_id: Mapped[int] = mapped_column( + ForeignKey("document.id", ondelete="CASCADE"), nullable=False + ) + session_id: Mapped[int] = mapped_column( + ForeignKey("generative_session.id", ondelete="CASCADE"), nullable=False + ) + pipeline_id: Mapped[int] = mapped_column( + ForeignKey("RAG_pipeline.id", ondelete="CASCADE"), nullable=False + ) + + # Relationships + document = relationship("Document", back_populates="pipeline_links") + pipeline = relationship("RAGPipeline", back_populates="pipeline_links") + session = relationship("GenerativeSession", back_populates="pipeline_links") + + __table_args__ = ( + UniqueConstraint("document_id", "session_id", name="uix_document_session"), + UniqueConstraint("session_id", "pipeline_id", name="uix_session_pipeline"), + UniqueConstraint("document_id", "pipeline_id", name="uix_document_pipeline"), + ) + + class Datafile(Base): __tablename__ = "datafile" diff --git a/DashAI/back/dependencies/database/utils.py b/DashAI/back/dependencies/database/utils.py index 184034def..a5a42ee22 100644 --- a/DashAI/back/dependencies/database/utils.py +++ b/DashAI/back/dependencies/database/utils.py @@ -64,7 +64,7 @@ def add_plugin_to_db( author=raw_plugin.author, verified=raw_plugin.verified, installed_version=raw_plugin.installed_version, - lastest_version=raw_plugin.lastest_version, + latest_version=raw_plugin.lastest_version, summary=raw_plugin.summary, description=raw_plugin.description, description_content_type=raw_plugin.description_content_type, diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 49ea258f1..96c4d6504 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -112,6 +112,7 @@ from DashAI.back.job.model_job import ModelJob from DashAI.back.job.pipeline_job import PipelineJob from DashAI.back.job.predict_job import PredictJob +from DashAI.back.job.RAG_job import RAGJob # Metrics from DashAI.back.metrics.classification.accuracy import Accuracy @@ -172,6 +173,9 @@ from DashAI.back.models.hugging_face.opus_mt_fr_en_transformer import ( OpusMtFrEnTransformer, ) +from DashAI.back.models.hugging_face.phi_4_mini_instruct_model import ( + Phi4MiniInstructModel, +) from DashAI.back.models.hugging_face.pixart_sigma_model import PixArtSigmaModel from DashAI.back.models.hugging_face.qwen_model import QwenModel from DashAI.back.models.hugging_face.roberta_transformer import RobertaTransformer @@ -209,6 +213,56 @@ from DashAI.back.models.hugging_face.xlnet_transformer import XlnetTransformer from DashAI.back.models.lenet5_image_classifier import LeNet5ImageClassifier from DashAI.back.models.mlp_image_classifier import MLPImageClassifier +from DashAI.back.models.RAG import RAGPipeline +from DashAI.back.models.RAG.chunking_models import ( + CharacterChunkModel, + RecursiveCharacterChunkModel, + TokenChunkModel, +) +from DashAI.back.models.RAG.embeddings.dense import ( + BERTEmbedding, + DistilBERTEmbedding, + E5Embedding, + GemmaEmbedding, + InstructorEmbedding, + LaBSEmbedding, + OpenAIEmbedding, + RoBERTaEmbedding, + SentenceTransformerEmbedding, +) +from DashAI.back.models.RAG.prompts import ( + CustomAugmentationPrompt, + CustomRAGGenerationPrompt, + DefaultAugmentationPrompt, + DefaultQARAGGenerationPrompt, + DefaultRAGGenerationPrompt, +) +from DashAI.back.models.RAG.retrievers.composite.mmr_reranker_retriever import ( + MMRRerankerRetriever, +) +from DashAI.back.models.RAG.retrievers.composite.parallel_retriever import ( + ParallelRetriever, +) +from DashAI.back.models.RAG.retrievers.composite.sequential_retriever import ( + SequentialRetriever, +) +from DashAI.back.models.RAG.retrievers.dense.dense_embedding_retriever import ( + DenseEmbeddingRetriever, +) +from DashAI.back.models.RAG.retrievers.sparse.bm25_retriever import ( + BM25Retriever, + BM25VectorizerModel, +) +from DashAI.back.models.RAG.retrievers.sparse.tfidf_retriever import ( + TFIDFRetriever, + TFIDFVectorizerModel, +) +from DashAI.back.models.remote_models.deepseek_text_to_text_generation_model import ( + DeepSeekTextToTextGenerationModel, +) +from DashAI.back.models.remote_models.openai_text_to_text_generation_model import ( + OpenAITextToTextGenerationModel, +) from DashAI.back.models.resnet18_image_classifier import ResNet18ImageClassifier from DashAI.back.models.resnet50_image_classifier import ResNet50ImageClassifier from DashAI.back.models.scikit_learn.adaboost_classifier import AdaBoostClassifier @@ -279,10 +333,11 @@ # Plugins from DashAI.back.plugins.utils import get_available_plugins - -# Tasks from DashAI.back.tasks.controlnet_task import ControlNetTask from DashAI.back.tasks.image_classification_task import ImageClassificationTask + +# Tasks +from DashAI.back.tasks.RAG_task import RAGTask from DashAI.back.tasks.regression_task import RegressionTask from DashAI.back.tasks.tabular_classification_task import TabularClassificationTask from DashAI.back.tasks.text_classification_task import TextClassificationTask @@ -314,6 +369,7 @@ def get_initial_components(): TextToImageGenerationTask, TextToTextGenerationTask, ControlNetTask, + RAGTask, ImageClassificationTask, # Models AdaBoostClassifier, @@ -340,6 +396,8 @@ def get_initial_components(): HistGradientBoostingClassifier, HistGradientBoostingRegression, KNeighborsClassifier, + QwenModel, + RAGPipeline, KNeighborsRegression, LassoRegression, LinearRegression, @@ -363,7 +421,6 @@ def get_initial_components(): OpusMtEsENTransformer, OpusMtFrEnTransformer, PixArtSigmaModel, - QwenModel, RandomForestClassifier, RobertaTransformer, RandomForestRegression, @@ -379,6 +436,9 @@ def get_initial_components(): StableDiffusionV3Model, StableDiffusionXLModel, StableDiffusionXLV1ControlNet, + OpenAITextToTextGenerationModel, + DeepSeekTextToTextGenerationModel, + Phi4MiniInstructModel, SVC, SVR, T5SmallTransformer, @@ -433,6 +493,7 @@ def get_initial_components(): DatasetJob, GenerativeJob, PipelineJob, + RAGJob, # Explainers KernelShap, PartialDependence, @@ -495,6 +556,35 @@ def get_initial_components(): SMOTEConverter, SMOTEENNConverter, RandomUnderSamplerConverter, + # Chunking Models + CharacterChunkModel, + RecursiveCharacterChunkModel, + TokenChunkModel, + # Encodings + OpenAIEmbedding, + SentenceTransformerEmbedding, + BERTEmbedding, + DistilBERTEmbedding, + RoBERTaEmbedding, + E5Embedding, + GemmaEmbedding, + InstructorEmbedding, + LaBSEmbedding, + # Prompts + DefaultRAGGenerationPrompt, + CustomRAGGenerationPrompt, + DefaultQARAGGenerationPrompt, + DefaultAugmentationPrompt, + CustomAugmentationPrompt, + # Retrievers + BM25Retriever, + BM25VectorizerModel, + TFIDFRetriever, + TFIDFVectorizerModel, + DenseEmbeddingRetriever, + MMRRerankerRetriever, + SequentialRetriever, + ParallelRetriever, ] # Obtener plugins instalados diff --git a/DashAI/back/job/RAG_job.py b/DashAI/back/job/RAG_job.py new file mode 100644 index 000000000..45fa0cb26 --- /dev/null +++ b/DashAI/back/job/RAG_job.py @@ -0,0 +1,271 @@ +import gc +import logging +from typing import TYPE_CHECKING, Any + +import torch +from kink import di, inject +from sqlalchemy import exc + +from DashAI.back.core.enums.status import RunStatus +from DashAI.back.dependencies.database.models import ( + GenerativeProcess, + GenerativeSession, + ProcessData, +) +from DashAI.back.job.base_job import BaseJob, JobError +from DashAI.back.models.RAG.RAG_constants import RAG_PARAM_KEYS as _RAG_PARAM_KEYS +from DashAI.back.models.RAG.RAG_pipeline import ( + RAGPipeline, + RAGPipelineConfig, +) +from DashAI.back.services.RAG.RAG_setup_service import RAGSetupService +from DashAI.back.tasks.RAG_task import RAGTask + +if TYPE_CHECKING: + from sqlalchemy.orm import sessionmaker + +log = logging.getLogger(__name__) + + +class RAGJob(BaseJob): + """RAGJob handles the full RAG pipeline lifecycle as a background job.""" + + @inject + def set_status_as_delivered( + self, session_factory: "sessionmaker" = lambda di: di["session_factory"] + ) -> None: + """Set the status of the job as delivered. + + Args: + session_factory: Injected SQLAlchemy session factory. + + Raises: + JobError: If the generative process does not exist or a DB error + occurs. + """ + generative_process_id: int = self.kwargs["generative_process_id"] + + with session_factory() as db: + process: GenerativeProcess = db.get( + GenerativeProcess, generative_process_id + ) + if not process: + raise JobError( + f"Generative process {generative_process_id} does not exist in DB." + ) + try: + process.set_status_as_delivered() + db.commit() + except exc.SQLAlchemyError as e: + log.exception(e) + raise JobError("Internal database error") from e + + @inject + def set_status_as_error( + self, session_factory: "sessionmaker" = lambda di: di["session_factory"] + ) -> None: + """Set the status of the job as error. + + Args: + session_factory: Injected SQLAlchemy session factory. + """ + generative_process_id: int = self.kwargs.get("generative_process_id") + if generative_process_id is None: + log.warning( + "Cannot set error status: generative_process_id is missing from kwargs" + ) + return + + with session_factory() as db: + process: GenerativeProcess = db.get( + GenerativeProcess, generative_process_id + ) + if not process: + return + try: + process.set_status_as_error() + db.commit() + except exc.SQLAlchemyError as e: + log.exception(e) + + @inject + def get_job_name(self) -> str: + """Get a descriptive name for the job. + + Returns: + A human-readable name including the session name if available, + otherwise a fallback like ``"RAG Process #"``. + """ + generative_process_id = self.kwargs.get("generative_process_id") + if not generative_process_id: + return "RAG Process" + + session_factory = di["session_factory"] + + try: + with session_factory() as db: + process: GenerativeProcess = db.get( + GenerativeProcess, generative_process_id + ) + if process: + session: GenerativeSession = db.get( + GenerativeSession, process.session_id + ) + if session and session.name: + return f"RAG: {session.name}" + return f"RAG Process #{generative_process_id}" + except Exception as e: + log.exception(f"Error getting job name: {e}") + + return f"RAG Process #{generative_process_id}" + + @inject + def run(self) -> None: + """Execute the full RAG pipeline lifecycle as a background job. + + Uses two separate DB sessions to avoid holding a connection open + during LLM inference (which may take minutes): + - **Session 1**: load data, build pipeline, prepare input, mark as started. + - **(no session)**: LLM inference via ``model.generate()``. + - **Session 2**: save output, mark as finished. + + Raises: + JobError: If any stage of execution fails. + """ + component_registry = di["component_registry"] + session_factory = di["session_factory"] + config = di["config"] + + if "generative_process_id" not in self.kwargs: + raise JobError("RAGJob requires 'generative_process_id' in kwargs.") + + generative_process_id: int = self.kwargs["generative_process_id"] + model = None + + # ── Session 1: Load, build, prepare ──────────────────────────── + try: + with session_factory() as db: + generative_process = db.get(GenerativeProcess, generative_process_id) + if not generative_process: + raise JobError( + f"Generative process {generative_process_id} not found in DB." + ) + generative_session = db.get( + GenerativeSession, generative_process.session_id + ) + if not generative_session: + raise JobError( + f"Session {generative_process.session_id} not found in DB." + ) + + # Whitelist-only: accept only known RAG pipeline keys. + raw_params = dict(generative_session.parameters) + extra_keys: set[str] = set(raw_params) - _RAG_PARAM_KEYS + if extra_keys: + log.debug( + "Filtered extra session parameter keys: %s", + sorted(extra_keys), + ) + clean_params = { + k: v for k, v in raw_params.items() if k in _RAG_PARAM_KEYS + } + pipeline_config = RAGPipelineConfig.from_kwargs( + db=db, + component_registry=component_registry, + session_id=generative_session.id, + env_RAG_path=config["RAG_PATH"], + **clean_params, + ) + setup_service = RAGSetupService( + db, + component_registry, + config["RAG_PATH"], + ) + model: RAGPipeline = setup_service.build_pipeline(pipeline_config) + + # Build conversation history from prior finished processes + input_data = generative_process.input + task = RAGTask() + finished_processes: list[GenerativeProcess] = ( + db.query(GenerativeProcess) + .filter( + GenerativeProcess.session_id == generative_session.id, + GenerativeProcess.status == RunStatus.FINISHED, + ) + .all() + ) + history = [ + (p.input[0].data, p.output[0].data) for p in finished_processes + ] + input_data = task.prepare_for_task(input_data, history=history) + + generative_process.set_status_as_started() + db.commit() + + process_id = generative_process.id + log.debug( + "Pipeline built and ready for process %d", generative_process_id + ) + except JobError: + self.set_status_as_error() + raise + except Exception as e: + log.exception(e) + self.set_status_as_error() + raise JobError("Error during RAG pipeline setup.") from e + + # ── Generation (no DB session) ────────────────────────────────── + try: + output: Any = model.generate(input_data) + log.debug("Generation completed for process %d", generative_process_id) + except Exception as e: + log.exception(e) + self.set_status_as_error() + raise JobError("Error during RAG model generation.") from e + + # ── Session 2: Save output ────────────────────────────────────── + try: + with session_factory() as db: + generative_process = db.get(GenerativeProcess, process_id) + if not generative_process: + raise JobError( + f"Generative process {process_id} not found when saving output." + ) + + output_data = task.process_output( + output, images_path=config["IMAGES_PATH"] + ) + outputs_for_database = [] + for o in output_data: + if not isinstance(o, tuple) or len(o) != 2: + raise JobError( + "Output from task must be a list of tuples (data, type)." + ) + data, data_type = o + outputs_for_database.append( + ProcessData( + data=data, + data_type=data_type, + process_id=generative_process.id, + is_input=False, + ) + ) + + db.add_all(outputs_for_database) + db.commit() + + db.refresh(generative_process) + generative_process.set_status_as_finished() + db.commit() + log.debug("Output saved for process %d", generative_process_id) + except Exception as e: + log.exception(e) + self.set_status_as_error() + raise JobError("Error processing and saving RAG generation output.") from e + + finally: + if model: + del model + if torch.cuda.is_available(): + torch.cuda.empty_cache() + gc.collect() diff --git a/DashAI/back/job/generative_job.py b/DashAI/back/job/generative_job.py index e52c58a72..3972b1c4c 100644 --- a/DashAI/back/job/generative_job.py +++ b/DashAI/back/job/generative_job.py @@ -4,19 +4,21 @@ from kink import inject from sqlalchemy import exc +from DashAI.back.core.enums.status import RunStatus from DashAI.back.dependencies.database.models import ( GenerativeProcess, GenerativeSession, ProcessData, ) from DashAI.back.job.base_job import BaseJob, JobError +from DashAI.back.job.RAG_job import RAGJob from DashAI.back.models.base_generative_model import BaseGenerativeModel from DashAI.back.tasks.base_generative_task import BaseGenerativeTask +from DashAI.back.tasks.RAG_task import RAGTask if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker -logging.basicConfig(level=logging.DEBUG) log = logging.getLogger(__name__) @@ -111,7 +113,20 @@ def run( component_registry = di["component_registry"] session_factory = di["session_factory"] config = di["config"] - # (Lazy imports removed to avoid duplicate and unused imports warnings) + + with session_factory() as db: + process = db.get( + GenerativeProcess, self.kwargs.get("generative_process_id") + ) + if process is not None: + session = db.get(GenerativeSession, process.session_id) + if session is not None: + task_class = component_registry[session.task_name]["class"] + if issubclass(task_class, RAGTask): + RAG_job = RAGJob(**self.kwargs) + RAG_job.run() + return + model = None generative_process = None with session_factory() as db: @@ -190,7 +205,7 @@ def run( .filter( GenerativeProcess.session_id == generative_session.id ) - .filter(GenerativeProcess.status == "FINISHED") + .filter(GenerativeProcess.status == RunStatus.FINISHED) .all() ] input_data = task.prepare_for_task( diff --git a/DashAI/back/job/rag_job.py b/DashAI/back/job/rag_job.py new file mode 100644 index 000000000..45fa0cb26 --- /dev/null +++ b/DashAI/back/job/rag_job.py @@ -0,0 +1,271 @@ +import gc +import logging +from typing import TYPE_CHECKING, Any + +import torch +from kink import di, inject +from sqlalchemy import exc + +from DashAI.back.core.enums.status import RunStatus +from DashAI.back.dependencies.database.models import ( + GenerativeProcess, + GenerativeSession, + ProcessData, +) +from DashAI.back.job.base_job import BaseJob, JobError +from DashAI.back.models.RAG.RAG_constants import RAG_PARAM_KEYS as _RAG_PARAM_KEYS +from DashAI.back.models.RAG.RAG_pipeline import ( + RAGPipeline, + RAGPipelineConfig, +) +from DashAI.back.services.RAG.RAG_setup_service import RAGSetupService +from DashAI.back.tasks.RAG_task import RAGTask + +if TYPE_CHECKING: + from sqlalchemy.orm import sessionmaker + +log = logging.getLogger(__name__) + + +class RAGJob(BaseJob): + """RAGJob handles the full RAG pipeline lifecycle as a background job.""" + + @inject + def set_status_as_delivered( + self, session_factory: "sessionmaker" = lambda di: di["session_factory"] + ) -> None: + """Set the status of the job as delivered. + + Args: + session_factory: Injected SQLAlchemy session factory. + + Raises: + JobError: If the generative process does not exist or a DB error + occurs. + """ + generative_process_id: int = self.kwargs["generative_process_id"] + + with session_factory() as db: + process: GenerativeProcess = db.get( + GenerativeProcess, generative_process_id + ) + if not process: + raise JobError( + f"Generative process {generative_process_id} does not exist in DB." + ) + try: + process.set_status_as_delivered() + db.commit() + except exc.SQLAlchemyError as e: + log.exception(e) + raise JobError("Internal database error") from e + + @inject + def set_status_as_error( + self, session_factory: "sessionmaker" = lambda di: di["session_factory"] + ) -> None: + """Set the status of the job as error. + + Args: + session_factory: Injected SQLAlchemy session factory. + """ + generative_process_id: int = self.kwargs.get("generative_process_id") + if generative_process_id is None: + log.warning( + "Cannot set error status: generative_process_id is missing from kwargs" + ) + return + + with session_factory() as db: + process: GenerativeProcess = db.get( + GenerativeProcess, generative_process_id + ) + if not process: + return + try: + process.set_status_as_error() + db.commit() + except exc.SQLAlchemyError as e: + log.exception(e) + + @inject + def get_job_name(self) -> str: + """Get a descriptive name for the job. + + Returns: + A human-readable name including the session name if available, + otherwise a fallback like ``"RAG Process #"``. + """ + generative_process_id = self.kwargs.get("generative_process_id") + if not generative_process_id: + return "RAG Process" + + session_factory = di["session_factory"] + + try: + with session_factory() as db: + process: GenerativeProcess = db.get( + GenerativeProcess, generative_process_id + ) + if process: + session: GenerativeSession = db.get( + GenerativeSession, process.session_id + ) + if session and session.name: + return f"RAG: {session.name}" + return f"RAG Process #{generative_process_id}" + except Exception as e: + log.exception(f"Error getting job name: {e}") + + return f"RAG Process #{generative_process_id}" + + @inject + def run(self) -> None: + """Execute the full RAG pipeline lifecycle as a background job. + + Uses two separate DB sessions to avoid holding a connection open + during LLM inference (which may take minutes): + - **Session 1**: load data, build pipeline, prepare input, mark as started. + - **(no session)**: LLM inference via ``model.generate()``. + - **Session 2**: save output, mark as finished. + + Raises: + JobError: If any stage of execution fails. + """ + component_registry = di["component_registry"] + session_factory = di["session_factory"] + config = di["config"] + + if "generative_process_id" not in self.kwargs: + raise JobError("RAGJob requires 'generative_process_id' in kwargs.") + + generative_process_id: int = self.kwargs["generative_process_id"] + model = None + + # ── Session 1: Load, build, prepare ──────────────────────────── + try: + with session_factory() as db: + generative_process = db.get(GenerativeProcess, generative_process_id) + if not generative_process: + raise JobError( + f"Generative process {generative_process_id} not found in DB." + ) + generative_session = db.get( + GenerativeSession, generative_process.session_id + ) + if not generative_session: + raise JobError( + f"Session {generative_process.session_id} not found in DB." + ) + + # Whitelist-only: accept only known RAG pipeline keys. + raw_params = dict(generative_session.parameters) + extra_keys: set[str] = set(raw_params) - _RAG_PARAM_KEYS + if extra_keys: + log.debug( + "Filtered extra session parameter keys: %s", + sorted(extra_keys), + ) + clean_params = { + k: v for k, v in raw_params.items() if k in _RAG_PARAM_KEYS + } + pipeline_config = RAGPipelineConfig.from_kwargs( + db=db, + component_registry=component_registry, + session_id=generative_session.id, + env_RAG_path=config["RAG_PATH"], + **clean_params, + ) + setup_service = RAGSetupService( + db, + component_registry, + config["RAG_PATH"], + ) + model: RAGPipeline = setup_service.build_pipeline(pipeline_config) + + # Build conversation history from prior finished processes + input_data = generative_process.input + task = RAGTask() + finished_processes: list[GenerativeProcess] = ( + db.query(GenerativeProcess) + .filter( + GenerativeProcess.session_id == generative_session.id, + GenerativeProcess.status == RunStatus.FINISHED, + ) + .all() + ) + history = [ + (p.input[0].data, p.output[0].data) for p in finished_processes + ] + input_data = task.prepare_for_task(input_data, history=history) + + generative_process.set_status_as_started() + db.commit() + + process_id = generative_process.id + log.debug( + "Pipeline built and ready for process %d", generative_process_id + ) + except JobError: + self.set_status_as_error() + raise + except Exception as e: + log.exception(e) + self.set_status_as_error() + raise JobError("Error during RAG pipeline setup.") from e + + # ── Generation (no DB session) ────────────────────────────────── + try: + output: Any = model.generate(input_data) + log.debug("Generation completed for process %d", generative_process_id) + except Exception as e: + log.exception(e) + self.set_status_as_error() + raise JobError("Error during RAG model generation.") from e + + # ── Session 2: Save output ────────────────────────────────────── + try: + with session_factory() as db: + generative_process = db.get(GenerativeProcess, process_id) + if not generative_process: + raise JobError( + f"Generative process {process_id} not found when saving output." + ) + + output_data = task.process_output( + output, images_path=config["IMAGES_PATH"] + ) + outputs_for_database = [] + for o in output_data: + if not isinstance(o, tuple) or len(o) != 2: + raise JobError( + "Output from task must be a list of tuples (data, type)." + ) + data, data_type = o + outputs_for_database.append( + ProcessData( + data=data, + data_type=data_type, + process_id=generative_process.id, + is_input=False, + ) + ) + + db.add_all(outputs_for_database) + db.commit() + + db.refresh(generative_process) + generative_process.set_status_as_finished() + db.commit() + log.debug("Output saved for process %d", generative_process_id) + except Exception as e: + log.exception(e) + self.set_status_as_error() + raise JobError("Error processing and saving RAG generation output.") from e + + finally: + if model: + del model + if torch.cuda.is_available(): + torch.cuda.empty_cache() + gc.collect() diff --git a/DashAI/back/models/RAG/RAG_constants.py b/DashAI/back/models/RAG/RAG_constants.py new file mode 100644 index 000000000..dffe01ba3 --- /dev/null +++ b/DashAI/back/models/RAG/RAG_constants.py @@ -0,0 +1,36 @@ +"""Canonical RAG pipeline parameter keys. +Single source of truth shared across models, services, and API layers. +""" + +RAG_PARAM_DOCUMENTS = "documents" +RAG_PARAM_PROMPT = "prompt" +RAG_PARAM_CHUNKING_MODEL = "chunking_model" +RAG_PARAM_RETRIEVER_MODEL = "retriever_model" +RAG_PARAM_GENERATION_MODEL = "generation_model" + +# All known RAG pipeline parameter keys +RAG_PARAM_KEYS = frozenset( + { + RAG_PARAM_DOCUMENTS, + RAG_PARAM_PROMPT, + RAG_PARAM_CHUNKING_MODEL, + RAG_PARAM_RETRIEVER_MODEL, + RAG_PARAM_GENERATION_MODEL, + } +) + +# Model component keys (each has {component, params} structure) +RAG_MODEL_KEYS = ( + RAG_PARAM_PROMPT, + RAG_PARAM_CHUNKING_MODEL, + RAG_PARAM_RETRIEVER_MODEL, + RAG_PARAM_GENERATION_MODEL, +) + +# Infrastructure / internal keys (not user-visible components) +ENV_RAG_PATH = "env_RAG_path" + +RAG_INFRA_KEYS = frozenset({"session_id", "db", "component_registry", ENV_RAG_PATH}) + +# All known pipeline keys combined (for the "unknown keys" guard) +RAG_PARAM_KEYS_ALL = frozenset(RAG_INFRA_KEYS | RAG_PARAM_KEYS | set(RAG_MODEL_KEYS)) diff --git a/DashAI/back/models/RAG/RAG_pipeline.py b/DashAI/back/models/RAG/RAG_pipeline.py new file mode 100644 index 000000000..b42077dfe --- /dev/null +++ b/DashAI/back/models/RAG/RAG_pipeline.py @@ -0,0 +1,445 @@ +"""RAG pipeline orchestration. + +RAGPipeline receives its dependencies injected (config, factories, +repository, loader) rather than constructing them from raw kwargs. +Orchestrates: document loading → chunk-set creation → chunking → +retrieval → prompt formatting → LLM generation. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Dict, List, Tuple + +from sqlalchemy.orm import Session + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + int_field, + list_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.base_generative_model import BaseGenerativeModel +from DashAI.back.models.RAG.documents import BaseDocument, Chunk +from DashAI.back.models.RAG.exceptions import ( + RAGPipelineConfigError, + RAGPipelineInputError, + RAGPipelineRuntimeError, +) +from DashAI.back.models.RAG.RAG_constants import ( + RAG_INFRA_KEYS, + RAG_MODEL_KEYS, + RAG_PARAM_KEYS, + RAG_PARAM_KEYS_ALL, +) +from DashAI.back.models.text_to_text_generation_model import ( + TextToTextGenerationTaskModel, +) + +log = logging.getLogger(__name__) + +if TYPE_CHECKING: + from DashAI.back.models.RAG.chunking_models.base_chunking_model import ( + BaseChunkingModel, + ) + from DashAI.back.models.RAG.prompts import Prompt + from DashAI.back.models.RAG.retrievers.retriever_model import RetrieverModel + + +@dataclass(frozen=True) +class ModelRef: + """Parsed reference to a component model. + + Represents a ``{component: str, params: dict}`` entry from the + pipeline parameter payload. + """ + + component: str + params: Dict[str, Any] + + +@dataclass(frozen=True) +class ChunkReference: + """A single chunk with its document metadata.""" + + document_id: int + document_name: str + document_position: int + text: str + + def to_dict(self) -> Dict[str, Any]: + """Return a JSON-serializable dict representation.""" + return { + "document_id": self.document_id, + "document_name": self.document_name, + "document_position": self.document_position, + "text": self.text, + } + + +@dataclass(frozen=True) +class RAGGenerationOutput: + """Typed output from RAGPipeline.generate().""" + + message: str + chunks: Dict[str, "ChunkReference"] + + +@dataclass(frozen=True) +class RAGPipelineConfig: + """Structured, validated representation of pipeline initialisation kwargs. + + Parses the raw kwargs dict once and provides typed field access, + eliminating magic strings in the pipeline's __init__. + """ + + session_id: int + db: Session + component_registry: ComponentRegistry + env_RAG_path: str # noqa: N815 + documents: List[int] + prompt: ModelRef + chunking_model: ModelRef + retriever_model: ModelRef + generation_model: ModelRef + + @classmethod + def from_kwargs(cls, **kwargs: Any) -> "RAGPipelineConfig": + """Parse and validate raw kwargs into a RAGPipelineConfig. + + Args: + **kwargs: Raw pipeline parameter dict containing infrastructure + keys (session_id, db, component_registry, env_RAG_path), + model keys (prompt, chunking_model, retriever_model, + generation_model), and the documents list. + + Returns: + A validated RAGPipelineConfig instance. + + Raises: + RAGPipelineConfigError: If required keys are missing, model + refs are not dicts, model refs lack component/params, or + unknown keys are present. + """ + missing: List[str] = [] + for key in RAG_INFRA_KEYS: + if key not in kwargs: + missing.append(key) + for key in RAG_PARAM_KEYS: + if key not in kwargs: + missing.append(key) + if missing: + raise RAGPipelineConfigError( + f"Missing required parameters: {sorted(missing)}" + ) + + model_refs: Dict[str, ModelRef] = {} + for key in RAG_MODEL_KEYS: + raw = kwargs[key] + if not isinstance(raw, dict): + raise RAGPipelineConfigError( + f"'{key}' must be a dict, got {type(raw).__name__}" + ) + if "component" not in raw: + raise RAGPipelineConfigError(f"Missing 'component' in '{key}'") + if "params" not in raw: + raise RAGPipelineConfigError(f"Missing 'params' in '{key}'") + model_refs[key] = ModelRef( + component=raw["component"], + params=raw["params"], + ) + + extra: set[str] = set(kwargs) - RAG_PARAM_KEYS_ALL + if extra: + log.warning( + "Unknown parameters passed to RAGPipelineConfig.from_kwargs(): %s. " + "These keys are not recognized and will cause the pipeline to fail. " + "Check if session parameters contain extra metadata keys.", + sorted(extra), + ) + raise RAGPipelineConfigError( + f"Unknown parameters: {sorted(extra)}. " + "These keys are not recognized pipeline configuration parameters." + ) + + return cls( + session_id=kwargs["session_id"], + db=kwargs["db"], + component_registry=kwargs["component_registry"], + env_RAG_path=kwargs["env_RAG_path"], + documents=kwargs["documents"], + prompt=model_refs["prompt"], + chunking_model=model_refs["chunking_model"], + retriever_model=model_refs["retriever_model"], + generation_model=model_refs["generation_model"], + ) + + +class RAGPipelineSchema(BaseSchema): + """Schema for RAG pipeline configuration parameters. + + Defines the five configurable fields: documents, prompt, + chunking_model, retriever_model, and generation_model. Each + model field accepts a ``{"component": str, "params": dict}`` + structure resolved by the ComponentRegistry. + """ + + documents: schema_field( + list_field(int_field(gt=0)), + placeholder=None, + description=MultilingualString( + en="List of document IDs to use in the RAG pipeline.", + es="Lista de IDs de documentos a usar en el pipeline RAG.", + ), + ) # type: ignore + + prompt: schema_field( + component_field(parent="Prompt"), + placeholder={"component": "DefaultRAGGenerationPrompt", "params": {}}, + description=MultilingualString( + en="Prompt template used in the RAG pipeline.", + es="Plantilla de prompt usada en el pipeline RAG.", + ), + ) # type: ignore + + chunking_model: schema_field( + component_field(parent="BaseChunkingModel"), + description=MultilingualString( + en="Chunking model used to split documents into smaller pieces.", + es="Modelo de fragmentación para dividir documentos en piezas.", + ), + placeholder={"component": "CharacterChunkModel", "params": {}}, + ) # type: ignore + + retriever_model: schema_field( + component_field(parent="RetrieverModel"), + placeholder={"component": "TFIDFRetriever", "params": {}}, + description=MultilingualString( + en="Retriever component used in the RAG pipeline.", + es="Componente recuperador usado en el pipeline RAG.", + ), + ) # type: ignore + + generation_model: schema_field( + component_field(parent="TextToTextGenerationTaskModel"), + placeholder={"component": "", "params": {}}, + description=MultilingualString( + en="Text generation model used in the RAG pipeline.", + es="Modelo de generación de texto usado en el pipeline RAG.", + ), + ) # type: ignore + + +class RAGPipeline(BaseGenerativeModel): + """Retrieval-Augmented Generation pipeline. + + Receives dependencies injected — does not construct factories, + repositories, or loaders. The caller (RAGJob) builds them from + the current DB session and passes them in. + + Orchestrates: document loading → chunk-set creation → chunking → + retrieval → prompt formatting → LLM generation. + """ + + COMPATIBLE_COMPONENTS: List[str] = ["RAGTask"] + SCHEMA: type[BaseSchema] = RAGPipelineSchema + + session_id: int + pipeline_id: int + chunking_model_id: int + documents_ids: List[int] + documents: Dict[int, BaseDocument] + prompt_model: Prompt + chunking_model: BaseChunkingModel + chunks: Dict[int, Dict[int, Chunk]] + retriever: RetrieverModel + llm_model: TextToTextGenerationTaskModel + + DISPLAY_NAME: str = MultilingualString( + en="RAG Pipeline", + es="Flujo de RAG", + pt="Pipeline RAG", + ) + DESCRIPTION: str = MultilingualString( + en=( + "Pipeline for Retrieval-Augmented Generation (RAG) tasks," + " orchestrating document loading, chunking, retrieval," + " prompt formatting, and LLM generation." + ), + es=( + "Pipeline para tareas de Generación Aumentada por" + " Recuperación (RAG), orquestando la carga de documentos," + " chunking, recuperación, formateo de prompts y generación" + " con LLM." + ), + pt=( + "Pipeline para tarefas de Geração Aumentada por Recuperação" + " (RAG), orquestrando carregamento de documentos, chunking," + " recuperação, formatação de prompts e geração com LLM." + ), + ) + COLOR: str = "#e12885" + ICON: str = "Grading" + MODEL_NAME: str = "RAGPipeline" + + def __init__( + self, + config: RAGPipelineConfig, + pipeline_id: int, + chunking_model_id: int, + documents: Dict[int, BaseDocument], + chunks: Dict[int, Dict[int, Chunk]], + prompt_model: Prompt, + chunking_model: BaseChunkingModel, + retriever: RetrieverModel, + llm_model: TextToTextGenerationTaskModel, + ) -> None: + """Initialise the RAG pipeline with fully constructed dependencies. + + Args: + config: Validated pipeline configuration. + pipeline_id: Database ID for the pipeline run. + chunking_model_id: Database ID for the chunking model. + documents: Mapping of document IDs to BaseDocument instances. + chunks: Mapping of document IDs to their chunk dicts. + prompt_model: The prompt template model. + chunking_model: The chunking model instance. + retriever: The retriever model instance. + llm_model: The text generation (LLM) model instance. + """ + self.session_id = config.session_id + self.pipeline_id = pipeline_id + self.chunking_model_id = chunking_model_id + self.documents_ids = config.documents + self.documents = documents + self.prompt_model = prompt_model + self.chunking_model = chunking_model + self.chunks = chunks + self.retriever = retriever + self.llm_model = llm_model + + def single_interaction( + self, + query: str, + ) -> List[Chunk]: + """Retrieve the top-K chunks for a single query. + + Args: + query: The search query string. + + Returns: + Ranked list of retrieved chunks. + + Raises: + RAGPipelineRuntimeError: If the retriever fails. + """ + try: + return self.retriever.retrieve(query) + except Exception as e: + raise RAGPipelineRuntimeError(f"Document retrieval failed: {str(e)}") from e + + def _build_chunk_references( + self, + chunks: List[Chunk], + ) -> Tuple[str, Dict[str, ChunkReference]]: + """Build a text block and reference map from retrieved chunks. + + Args: + chunks: Chunks returned by the retriever. + + Returns: + Joined chunk texts and a mapping from chunk key to + ChunkReference. + """ + chunks_texts: List[str] = [] + chunk_dict: Dict[str, ChunkReference] = {} + for retrieved_chunk in chunks: + document_id: int = retrieved_chunk.document_id + document: BaseDocument | None = self.documents.get(document_id) + if document is None: + raise RAGPipelineRuntimeError( + f"Document with ID {document_id} not found in pipeline documents." + ) + chunk_position: int = retrieved_chunk.document_position + chunk_text: str = retrieved_chunk.text + chunk_ref: str = f"{document_id}_{chunk_position}" + chunk_dict[chunk_ref] = ChunkReference( + document_id=document_id, + document_name=document.file_name, + document_position=chunk_position, + text=chunk_text, + ) + chunks_texts.append( + f"Document {document.file_name}, " + f"chunk nº {chunk_position}, text:\n {chunk_text}" + ) + return "\n\n".join(chunks_texts), chunk_dict + + def generate( + self, + input_data: Tuple[Dict[str, str], ...], + ) -> RAGGenerationOutput: + """Run the full RAG pipeline: retrieve, format, and generate. + + Args: + input_data: Chat-format input with history as earlier entries + and the current user message as the last entry. + + Returns: + The generated message and the chunks used for retrieval. + + Raises: + RAGPipelineInputError: If input_data is empty or malformed. + RAGPipelineRuntimeError: If any pipeline stage fails. + """ + if not input_data: + raise RAGPipelineInputError("input_data must not be empty.") + try: + input_dict: Dict[str, str] = input_data[-1] + input_message: str = input_dict["content"] + except (KeyError, IndexError, TypeError) as e: + raise RAGPipelineInputError( + f"Malformed input_data: expected Tuple[Dict[str, str], ...] " + f"with a 'content' key in the last entry, got {type(input_data)}: {e}" + ) from e + history: Tuple[Dict[str, str], ...] = input_data[:-1] + try: + chunks: List[Chunk] = self.single_interaction(input_message) + except Exception as e: + raise RAGPipelineRuntimeError(f"Failed during retrieval: {str(e)}") from e + try: + chunks_text: str + chunk_dict: Dict[str, ChunkReference] + chunks_text, chunk_dict = self._build_chunk_references(chunks) + prompt: str = self.prompt_model.format( + input=input_message, + chunks=chunks_text, + ) + except Exception as e: + raise RAGPipelineRuntimeError( + f"Failed during prompt formatting: {str(e)}" + ) from e + try: + model_input: List[Dict[str, str]] = list(history) + [ + {"role": "user", "content": prompt} + ] + # NOTE: Output is not streamed — the user waits for the full response. + output: List[Any] = self.llm_model.generate(model_input) + if not output: + raise RAGPipelineRuntimeError("LLM model returned empty output list.") + raw_message = output[0] + if not isinstance(raw_message, str): + raise RAGPipelineRuntimeError( + f"LLM output is not a string, got {type(raw_message).__name__}" + ) + return RAGGenerationOutput(message=raw_message, chunks=chunk_dict) + except RAGPipelineRuntimeError: + raise + except Exception as e: + raise RAGPipelineRuntimeError( + f"Failed during LLM generation: {str(e)}" + ) from e diff --git a/DashAI/back/models/RAG/__init__.py b/DashAI/back/models/RAG/__init__.py new file mode 100644 index 000000000..42ff340a4 --- /dev/null +++ b/DashAI/back/models/RAG/__init__.py @@ -0,0 +1,3 @@ +"""RAG model package, exposes the RAGPipeline class.""" + +from DashAI.back.models.RAG.RAG_pipeline import RAGPipeline # noqa: F401 diff --git a/DashAI/back/models/RAG/chunking_models/__init__.py b/DashAI/back/models/RAG/chunking_models/__init__.py new file mode 100644 index 000000000..bb5d2b3dd --- /dev/null +++ b/DashAI/back/models/RAG/chunking_models/__init__.py @@ -0,0 +1,12 @@ +"""Chunking model layer for the RAG module. + +Provides the abstract base chunking model and concrete implementations +for character-based, recursive character-based, and token-based chunking. +""" + +from DashAI.back.models.RAG.chunking_models.base_chunking_model import BaseChunkingModel +from DashAI.back.models.RAG.chunking_models.character_chunk_model import CharacterChunkModel +from DashAI.back.models.RAG.chunking_models.recursive_character_chunk_model import ( + RecursiveCharacterChunkModel, +) +from DashAI.back.models.RAG.chunking_models.token_chunk_model import TokenChunkModel diff --git a/DashAI/back/models/RAG/chunking_models/base_chunking_model.py b/DashAI/back/models/RAG/chunking_models/base_chunking_model.py new file mode 100644 index 000000000..a18e924d0 --- /dev/null +++ b/DashAI/back/models/RAG/chunking_models/base_chunking_model.py @@ -0,0 +1,137 @@ +from abc import ABCMeta, abstractmethod +from typing import Dict, Final, List + +from DashAI.back.config_object import ConfigObject +from DashAI.back.models.RAG.documents.base_document import BaseDocument +from DashAI.back.models.RAG.documents.chunk import Chunk +from DashAI.back.models.RAG.exceptions import RAGChunkingError + + +class BaseChunkingModel(ConfigObject, metaclass=ABCMeta): + """Abstract base class for all chunking models. + + Defines the interface and common workflow for chunking strategies. + Subclasses must implement :meth:`chunk_text`. + """ + + TYPE: Final[str] = "ChunkingModel" + REQUIRED_EXTRA_KWARGS: Final[List[str]] = ["documents"] + + def __init__(self, **kwargs): + """Initialize the chunking model. + + Validates that all required extra keyword arguments are present, + stores the document mapping, and validates/transforms model parameters. + + Args: + **kwargs: Must include a ``documents`` key mapping document IDs + to :class:`BaseDocument` instances. Remaining keys are + validated and transformed by :meth:`validate_and_transform`. + + Raises: + RAGChunkingError: If a required extra keyword argument is missing. + """ + for required_key in self.REQUIRED_EXTRA_KWARGS: + if required_key not in kwargs: + raise RAGChunkingError( + f"Missing required extra keyword argument: '{required_key}'" + ) + self.documents: Dict[int, BaseDocument] = kwargs.pop("documents") + self.parameters = self.validate_and_transform(kwargs) + self.chunks = None + self.id = None + + def compute_chunks(self) -> None: + """Compute chunks if not already done.""" + if self.chunks is None: + self.chunks = self.chunk_documents() + + @abstractmethod + def chunk_text(self, text: str, **kwargs) -> List[str]: + """ + Chunk the input text into smaller pieces. + + Args: + text (str): The input text to be chunked. + **kwargs: Additional parameters for chunking. + + Returns: + List[str]: A list of text chunks. + """ + raise NotImplementedError("Subclasses must implement the chunk_text method.") + + def chunk_document(self, document: BaseDocument) -> Dict[int, Chunk]: + """ + Chunk a single document using the chunk_text method. + + Args: + document (BaseDocument): The input document to be chunked. + **kwargs: Additional parameters for chunking, passed to `chunk_text`. + + Returns: + Dict[int, Chunk]: A dictionary mapping chunk indices to Chunk objects. + """ + if not isinstance(document, BaseDocument): + raise TypeError("Input must be an instance of BaseDocument.") + text = document.get_text() + chunks_text = self.chunk_text(text) + chunks = {} + for idx, chunk_text in enumerate(chunks_text): + chunk = Chunk( + id=None, + document_position=idx, + document_id=document.id, + text=chunk_text, + ) + chunks[idx] = chunk + return chunks + + def chunk_documents(self) -> Dict[int, Dict[int, Chunk]]: + """ + Chunk documents using the chunk_document method. + + Returns: + Dict[int, Dict[int, Chunk]]: A dictionary mapping document IDs + to their respective dictionaries of chunk indices and Chunk + objects. + """ + chunked_documents = {} + for document_id, document in self.documents.items(): + chunks = self.chunk_document(document) + chunked_documents[document_id] = chunks + return chunked_documents + + def get_chunks(self) -> Dict[int, Dict[int, Chunk]]: + """ + Get the chunks generated by the chunking model. + + Returns: + Dict[int, Dict[int, Chunk]]: A dictionary mapping document IDs + to their respective dictionaries of chunk indices and Chunk + objects. + """ + self.compute_chunks() + return self.chunks + + def get_id(self) -> int: + """ + Get the ID of the chunking model. + + Returns: + int: The ID of the chunking model. + """ + return self.id + + def set_id(self, id: int) -> None: + """ + Set the ID of the chunking model. + + Args: + id (int): The ID to set for the chunking model. + """ + if self.id is None: + self.id = id + else: + raise RAGChunkingError( + "Chunking model ID is already set and cannot be modified." + ) diff --git a/DashAI/back/models/RAG/chunking_models/character_chunk_model.py b/DashAI/back/models/RAG/chunking_models/character_chunk_model.py new file mode 100644 index 000000000..0a17dcac8 --- /dev/null +++ b/DashAI/back/models/RAG/chunking_models/character_chunk_model.py @@ -0,0 +1,95 @@ +from typing import List + +from pydantic import field_validator + +from DashAI.back.core.schema_fields import ( + BaseSchema, + int_field, + schema_field, +) +from DashAI.back.models.RAG.chunking_models.base_chunking_model import ( + BaseChunkingModel, +) +from DashAI.back.models.RAG.exceptions import RAGChunkingOverlapError + + +class CharacterChunkModelSchema(BaseSchema): + """Schema for character-based chunking model.""" + + chunk_size: schema_field( + int_field(gt=1), + placeholder=400, + description="Size of each chunk in characters.", + ) # type: ignore + + chunk_overlap: schema_field( + int_field(gt=0), + placeholder=40, + description=( + "Number of characters to overlap between chunks. " + "Must be less than chunk_size." + ), + ) # type: ignore + + @field_validator("chunk_overlap", mode="after") + @classmethod + def validate_chunk_overlap(cls, v, info): + """Validate that chunk_overlap is less than chunk_size.""" + chunk_size = info.data.get("chunk_size") + if chunk_size is not None and v >= chunk_size: + raise ValueError( + f"chunk_overlap must be less than chunk_size. " + f"Got chunk_overlap={v} and chunk_size={chunk_size}" + ) + return v + + +class CharacterChunkModel(BaseChunkingModel): + """ + Chunking model that splits documents into chunks based on character count. + This model is useful for processing large text documents into manageable pieces. + """ + + SCHEMA = CharacterChunkModelSchema + + def __init__(self, **kwargs): + """ + Initialize the character chunking model with the specified chunk size and + overlap. + + Args: + chunk_size (int): Size of each chunk in characters. + chunk_overlap (int): Number of characters to overlap between chunks. + """ + self.chunk_size = kwargs["chunk_size"] + self.chunk_overlap = kwargs["chunk_overlap"] + super().__init__(**kwargs) + + def chunk_text(self, text: str) -> List[str]: + """ + Chunk the input text into smaller segments based on character count. + + Args: + text (str): The input text to be chunked. + + Returns: + List[str]: A list of text chunks. + """ + if self.chunk_overlap >= self.chunk_size: + raise RAGChunkingOverlapError( + f"chunk_overlap ({self.chunk_overlap}) must be less than " + f"chunk_size ({self.chunk_size})" + ) + chunks = [] + start = 0 + text_length = len(text) + + while start < text_length: + end = min(start + self.chunk_size, text_length) + chunk = text[start:end] + chunks.append(chunk) + if end == text_length: + break + start += self.chunk_size - self.chunk_overlap + + return chunks diff --git a/DashAI/back/models/RAG/chunking_models/chunking_model_factory.py b/DashAI/back/models/RAG/chunking_models/chunking_model_factory.py new file mode 100644 index 000000000..c2ec7f305 --- /dev/null +++ b/DashAI/back/models/RAG/chunking_models/chunking_model_factory.py @@ -0,0 +1,80 @@ +"""Pure factory for chunking models. + +Resolves a component name + parameters into a BaseChunkingModel instance +and chunks the provided documents. Phase 1 only — no DB access. +""" + +from dataclasses import dataclass +from typing import Any + +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.RAG.chunking_models.base_chunking_model import ( + BaseChunkingModel, +) +from DashAI.back.models.RAG.documents.base_document import BaseDocument +from DashAI.back.models.RAG.documents.chunk import Chunk +from DashAI.back.models.RAG.exceptions import RAGComponentNotFoundError + + +@dataclass(frozen=True) +class ChunkingFactoryResult: + """Result of a chunking model creation via :class:`ChunkingModelFactory`. + + Attributes: + model: The instantiated chunking model. + chunks: A dictionary mapping document IDs to their chunk dictionaries. + """ + + model: BaseChunkingModel + chunks: dict[int, dict[int, Chunk]] + + +class ChunkingModelFactory: + """Creates a chunking model and chunks documents. Phase 1 only — no DB.""" + + def __init__( + self, + registry: ComponentRegistry, + documents: dict[int, BaseDocument], + ): + """Initialize the factory with a component registry and document mapping. + + Args: + registry: The component registry used to resolve chunking model + classes by name. + documents: A dictionary mapping document IDs to BaseDocument + instances to be chunked. + """ + self._registry = registry + self._documents = documents + + def create( + self, + component_name: str, + params: dict[str, Any], + ) -> ChunkingFactoryResult: + """Build a chunking model and chunk the documents. + + Parameters + ---------- + component_name : str + Registered component name. + params : dict + Parameters passed to the model constructor. + + Returns + ------- + ChunkingFactoryResult + The instantiated model and its chunk dictionary. + """ + try: + model_class = self._registry[component_name]["class"] + except KeyError as err: + raise RAGComponentNotFoundError( + f"Component '{component_name}' not found in registry" + ) from err + + model_params = {**params, "documents": self._documents} + model = model_class(**model_params) + model.compute_chunks() + return ChunkingFactoryResult(model=model, chunks=model.get_chunks()) diff --git a/DashAI/back/models/RAG/chunking_models/recursive_character_chunk_model.py b/DashAI/back/models/RAG/chunking_models/recursive_character_chunk_model.py new file mode 100644 index 000000000..9365a3b8c --- /dev/null +++ b/DashAI/back/models/RAG/chunking_models/recursive_character_chunk_model.py @@ -0,0 +1,202 @@ +from typing import List + +from pydantic import field_validator + +from DashAI.back.core.schema_fields import ( + BaseSchema, + int_field, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.chunking_models.base_chunking_model import ( + BaseChunkingModel, +) +from DashAI.back.models.RAG.exceptions import RAGChunkingOverlapError + + +class RecursiveCharacterChunkModelSchema(BaseSchema): + """Schema for the recursive character chunking model configuration.""" + + chunk_size: schema_field( + int_field(gt=1), + placeholder=1000, + description=MultilingualString( + en="Size of each chunk in characters.", + es="Tamaño de cada fragmento en caracteres.", + ), + ) # type: ignore + + chunk_overlap: schema_field( + int_field(ge=0), + placeholder=100, + description=MultilingualString( + en=( + "Number of characters to overlap between chunks. " + "Must be less than chunk_size." + ), + es=( + "Número de caracteres a solapar entre fragmentos. " + "Debe ser menor que chunk_size." + ), + ), + ) # type: ignore + + separators: schema_field( + list_field(string_field(), min_items=1), + placeholder=["\n\n", "\n", ".", " ", ""], + description=MultilingualString( + en="Ordered list of separators to use for splitting.", + es="Lista ordenada de separadores a usar para dividir.", + ), + ) # type: ignore + + @field_validator("chunk_overlap", mode="after") + @classmethod + def validate_chunk_overlap(cls, v, info): + """Validate that chunk_overlap is less than chunk_size. + + Args: + v: The value of chunk_overlap. + info: Validation info containing the chunk_size field. + + Returns: + int: The validated chunk_overlap value. + + Raises: + ValueError: If chunk_overlap is greater than or equal to chunk_size. + """ + chunk_size = info.data.get("chunk_size") + if chunk_size is not None and v >= chunk_size: + raise ValueError( + f"chunk_overlap must be less than chunk_size. " + f"Got chunk_overlap={v} and chunk_size={chunk_size}" + ) + return v + + +class RecursiveCharacterChunkModel(BaseChunkingModel): + """Chunking model that recursively splits text using a prioritized list + of separators. + + Falls back to character-level splitting when no separator fits within the + chunk size. + """ + + SCHEMA = RecursiveCharacterChunkModelSchema + + DISPLAY_NAME: str = MultilingualString( + en="Recursive Character Chunk Model", + es="Modelo de Fragmentación Recursiva por Caracteres", + ) + + FLAGS: list[str] = [] + + def __init__(self, **kwargs): + """Initialize the recursive character chunking model. + + Args: + **kwargs: Must include ``chunk_size``, ``chunk_overlap``, and + ``separators``, plus a ``documents`` mapping passed to the + parent class. + """ + self.chunk_size = kwargs["chunk_size"] + self.chunk_overlap = kwargs["chunk_overlap"] + self.separators = kwargs["separators"] + super().__init__(**kwargs) + + def chunk_text(self, text: str) -> List[str]: + """Split text into chunks using recursive separator-based splitting. + + Args: + text: The input text to be chunked. + + Returns: + List[str]: A list of text chunks. + + Raises: + RAGChunkingOverlapError: If chunk_overlap is not less than chunk_size. + """ + if self.chunk_overlap >= self.chunk_size: + raise RAGChunkingOverlapError( + f"chunk_overlap ({self.chunk_overlap}) must be less than " + f"chunk_size ({self.chunk_size})" + ) + if len(text) <= self.chunk_size: + return [text] + return self._split_text(text, 0) + + def _split_text(self, text: str, separator_idx: int) -> List[str]: + """Recursively split text using the separator at the given index, + falling back to the next separator if needed. + + Args: + text: The text to split. + separator_idx: Index into the ordered separators list. + + Returns: + List[str]: A list of text chunks. + """ + if separator_idx >= len(self.separators): + return self._force_split(text) + + separator = self.separators[separator_idx] + + if separator == "" or separator not in text: + return self._split_text(text, separator_idx + 1) + + splits = text.split(separator) + return self._recursive_split(splits, separator_idx + 1) + + def _recursive_split(self, splits: List[str], next_sep_idx: int) -> List[str]: + """Iterate over already-split segments and further split any that exceed + chunk_size. + + Args: + splits: Segments from a previous split operation. + next_sep_idx: Index of the next separator to try for oversize segments. + + Returns: + List[str]: A flat list of chunks all within chunk_size. + """ + result: List[str] = [] + for split in splits: + if len(split) <= self.chunk_size: + result.append(split) + continue + + resolved = False + for sep_idx in range(next_sep_idx, len(self.separators)): + separator = self.separators[sep_idx] + if separator == "" or separator in split: + result.extend(self._split_text(split, sep_idx)) + resolved = True + break + + if not resolved: + result.extend(self._force_split(split)) + + return result + + def _force_split(self, text: str) -> List[str]: + """Split text into fixed-size chunks with overlap when all separators fail. + + Args: + text: The text to split. + + Returns: + List[str]: A list of fixed-size character chunks. + """ + chunks: List[str] = [] + start = 0 + text_length = len(text) + + while start < text_length: + end = min(start + self.chunk_size, text_length) + chunks.append(text[start:end]) + if end == text_length: + break + start += self.chunk_size - self.chunk_overlap + + return chunks diff --git a/DashAI/back/models/RAG/chunking_models/token_chunk_model.py b/DashAI/back/models/RAG/chunking_models/token_chunk_model.py new file mode 100644 index 000000000..6134beb29 --- /dev/null +++ b/DashAI/back/models/RAG/chunking_models/token_chunk_model.py @@ -0,0 +1,128 @@ +from typing import TYPE_CHECKING, List, Optional + +from pydantic import field_validator + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + int_field, + schema_field, +) +from DashAI.back.models.RAG.chunking_models.base_chunking_model import ( + BaseChunkingModel, +) +from DashAI.back.models.RAG.exceptions import RAGChunkingOverlapError + + +class TokenChunkModelSchema(BaseSchema): + """Schema for the token-based chunking model configuration.""" + + tokenizer_name: schema_field( + enum_field( + enum=[ + "intfloat/e5-mistral-7b-instruct", + "dccuchile/bert-base-spanish-wwm-uncased", + ], + ), + placeholder="intfloat/e5-mistral-7b-instruct", + description="The tokenizer model to use for tokenizing the text.", + ) # type: ignore + + chunk_size: schema_field( + int_field(gt=1), + placeholder=400, + description="The size of each chunk in tokens.", + ) # type: ignore + + chunk_overlap: schema_field( + int_field(ge=0), + placeholder=40, + description=( + "The number of overlapping tokens between chunks. " + "Must be less than chunk_size." + ), + ) # type: ignore + + @field_validator("chunk_overlap") + @classmethod + def validate_chunk_overlap(cls, v, info): + """Validate that chunk_overlap is less than chunk_size.""" + chunk_size = info.data.get("chunk_size") + if chunk_size is not None and v >= chunk_size: + raise ValueError( + f"chunk_overlap must be less than chunk_size. " + f"Got chunk_overlap={v} and chunk_size={chunk_size}" + ) + return v + + +if TYPE_CHECKING: + from transformers import AutoTokenizer + + +class TokenChunkModel(BaseChunkingModel): + """Chunking model that splits text into chunks based on token count. + + Uses a HuggingFace ``AutoTokenizer`` to tokenize the text and splits at + token boundaries with configurable overlap. + """ + + SCHEMA = TokenChunkModelSchema + + def __init__(self, **kwargs): + """Initialize the token chunking model. + + Args: + **kwargs: Must include ``chunk_size``, ``chunk_overlap``, and + ``tokenizer_name``, plus a ``documents`` mapping passed to + the parent class. + """ + # Extract parameters before super().__init__() pops "documents" + self.chunk_size = kwargs["chunk_size"] + self.chunk_overlap = kwargs["chunk_overlap"] + self.tokenizer_name = kwargs["tokenizer_name"] + self._tokenizer: Optional["AutoTokenizer"] = None + # self.parameters is set by BaseChunkingModel.validate_and_transform() + super().__init__(**kwargs) + + @property + def tokenizer(self) -> "AutoTokenizer": + """Lazily load and return the HuggingFace tokenizer. + + Returns: + AutoTokenizer: The tokenizer loaded from the configured + ``tokenizer_name``. + """ + if self._tokenizer is None: + from transformers import AutoTokenizer + + self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name) + return self._tokenizer + + def chunk_text(self, text: str) -> List[str]: + """Split text into chunks based on token count. + + Args: + text: The input text to be chunked. + + Returns: + List[str]: A list of text chunks, each converted back to a string + from its constituent tokens. + + Raises: + RAGChunkingOverlapError: If chunk_overlap is not less than chunk_size. + """ + if self.chunk_overlap >= self.chunk_size: + raise RAGChunkingOverlapError( + f"chunk_overlap ({self.chunk_overlap}) must be less than " + f"chunk_size ({self.chunk_size})" + ) + tokens = self.tokenizer.tokenize(text) + + token_chunks = [] + while len(tokens) > 0: + chunk = tokens[: self.chunk_size] + token_chunks.append(self.tokenizer.convert_tokens_to_string(chunk)) + tokens = tokens[self.chunk_size - self.chunk_overlap :] + + return token_chunks diff --git a/DashAI/back/models/RAG/documents/__init__.py b/DashAI/back/models/RAG/documents/__init__.py new file mode 100644 index 000000000..42e5f69fa --- /dev/null +++ b/DashAI/back/models/RAG/documents/__init__.py @@ -0,0 +1,11 @@ +"""Document representation layer for the RAG module. + +Provides abstract and concrete document types (BaseDocument, PDFDocument, +TxtDocument), a Chunk data class, and a DocumentFileType enumeration. +""" + +from .base_document import BaseDocument +from .chunk import Chunk +from .file_type import DocumentFileType +from .pdf_document import PDFDocument +from .txt_document import TxtDocument diff --git a/DashAI/back/models/RAG/documents/base_document.py b/DashAI/back/models/RAG/documents/base_document.py new file mode 100644 index 000000000..61d18a994 --- /dev/null +++ b/DashAI/back/models/RAG/documents/base_document.py @@ -0,0 +1,138 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict, Optional + +from DashAI.back.models.RAG.documents.file_type import DocumentFileType + + +class BaseDocument(ABC): + """Abstract base class for all document types. + + Defines the interface that concrete document classes must implement, + including text extraction, metadata retrieval, and file type detection. + """ + + SUPPORTED_FILE_TYPES = [ + DocumentFileType.PDF, + DocumentFileType.TXT, + # MD, RST, TEX, CSV are handled as TXT via TxtDocument + ] + + def __init__( + self, + id: int, + file_name: str, + file_path: str, + file_hash: str, + created: Optional[str] = None, + optional_metadata: Optional[Dict[str, Any]] = None, + ): + """Initialize a BaseDocument instance. + + Args: + id: The unique identifier of the document. + file_name: The name of the file. + file_path: The path to the file. + file_hash: A hash of the file content. + created: The creation date of the document. + optional_metadata: Additional metadata for the document. + """ + self.id = id + self.file_name = file_name + self.file_path = file_path + self.file_hash = file_hash + self.created = created + self.optional_metadata = ( + optional_metadata if optional_metadata is not None else {} + ) + + @abstractmethod + def get_text(self) -> str: + """ + Get the text content of the document. + + Returns: + str: The text content of the document. + """ + + def get_text_length(self) -> int: + """ + Get the length of the text content of the document. + + Returns: + int: The length of the text content of the document. + """ + return len(self.get_text()) + + @abstractmethod + def get_metadata(self) -> Dict[str, Any]: + """ + Get the metadata of the document. + + Returns: + Dict[str, Any]: The metadata of the document. + """ + + def get_id(self) -> int: + """ + Get the unique identifier of the document. + + Returns: + int: The unique identifier of the document. + """ + return self.id + + def get_file_name(self) -> Optional[str]: + """ + Get the filename of the document. + + Returns: + Optional[str]: The filename of the document, or None if not applicable. + """ + return self.file_name + + def get_file_path(self) -> Optional[str]: + """ + Get the file path of the document. + + Returns: + Optional[str]: The file path of the document, or None if not applicable. + """ + return self.file_path + + def get_file_hash(self) -> str: + """ + Get the hash of the document content. + + Returns: + str: A hash string representing the document content. + """ + return self.file_hash + + def get_file_type(self) -> Optional[DocumentFileType]: + """ + Get the filetype of the document. + + Returns: + Optional[DocumentFileType]: The filetype of the document, + or None if not applicable. + """ + if not self.file_path: + return None + ext = self.file_path.split(".")[-1].lower() + try: + return DocumentFileType(ext) + except ValueError: + return None + + def __repr__(self) -> str: + """Return a string representation of the BaseDocument. + + Returns: + str: A string containing the document id, filename, first 50 + characters of text, and metadata. + """ + return ( + f"BaseDocument(id={self.id}, filename='{self.get_file_name()}', " + f"content='{self.get_text()[:50]}...', " + f"metadata={self.get_metadata()})" + ) diff --git a/DashAI/back/models/RAG/documents/chunk.py b/DashAI/back/models/RAG/documents/chunk.py new file mode 100644 index 000000000..2bd32791a --- /dev/null +++ b/DashAI/back/models/RAG/documents/chunk.py @@ -0,0 +1,36 @@ +from typing import Optional + +from DashAI.back.models.RAG.utils import hash_function + + +class Chunk: + """A chunk of text extracted from a document. + + Note: + ``id`` starts as ``None`` when the chunk is first created in memory + (see :meth:`BaseChunkingModel.chunk_document`). It is later set to the + real database primary key by :meth:`ChunkingService._update_chunk_ids` + after the chunk is persisted. + """ + + def __init__( + self, + id: Optional[int], + document_id: int, + document_position: int, + text: str, + ): + """Initialize a Chunk instance. + + Args: + id: The database primary key, or None if not yet persisted. + document_id: The ID of the parent document. + document_position: The 0-based position of this chunk within + the document. + text: The text content of the chunk. + """ + self.id: Optional[int] = id + self.document_id = document_id + self.document_position = document_position + self.text = text + self.hash = hash_function(text) diff --git a/DashAI/back/models/RAG/documents/file_type.py b/DashAI/back/models/RAG/documents/file_type.py new file mode 100644 index 000000000..6856da94d --- /dev/null +++ b/DashAI/back/models/RAG/documents/file_type.py @@ -0,0 +1,34 @@ +"""Document file type enumeration. + +Single source of truth for file-type strings used across all +layers (models, services, API). Each member's value is the +canonical extension string (without the leading dot). +""" + +from __future__ import annotations + +from enum import Enum + + +class DocumentFileType(str, Enum): + """Supported document file extensions. + + Usage:: + + from DashAI.back.models.RAG.documents.file_type import DocumentFileType + + # As a string (via ``str`` mixin): + ft: str = DocumentFileType.PDF # "pdf" + + # As dict key: + classes: dict[DocumentFileType, type] = { + DocumentFileType.PDF: PDFDocument, + } + """ + + TXT = "txt" + PDF = "pdf" + MD = "md" + RST = "rst" + TEX = "tex" + CSV = "csv" diff --git a/DashAI/back/models/RAG/documents/pdf_document.py b/DashAI/back/models/RAG/documents/pdf_document.py new file mode 100644 index 000000000..05567ea40 --- /dev/null +++ b/DashAI/back/models/RAG/documents/pdf_document.py @@ -0,0 +1,117 @@ +from typing import Any, Dict, Optional + +from DashAI.back.models.RAG.documents.base_document import BaseDocument +from DashAI.back.models.RAG.exceptions import RAGDocumentParsingError + + +def _clean_textract_output(text: str) -> str: + """Clean textract output by removing control characters and normalizing whitespace. + + Args: + text: Raw text extracted by textract. + + Returns: + str: Cleaned text with control characters replaced by spaces + and runs of whitespace collapsed. + """ + import re + + text = re.sub(r"[\x00-\x1f\x7f]", " ", text) + text = re.sub(r"\s+", " ", text) + return text.strip() + + +class PDFDocument(BaseDocument): + """Document implementation for PDF files. + + Supports two text-extraction backends: PyPDF2 and textract. + """ + + def __init__( + self, + id: int, + file_name: str, + file_path: str, + file_hash: str, + created: Optional[str] = None, + optional_metadata: Optional[Dict[str, Any]] = None, + parser: str = "textract", + ): + """Initialize a PDFDocument instance. + + Args: + id: The unique identifier of the document. + file_name: The name of the file. + file_path: The path to the file. + file_hash: A hash of the file content (computed upstream). + created: The creation date of the document. + optional_metadata: Additional metadata for the document. + parser: PDF parser to use ("PyPDF2" or "textract"). + + Raises: + ValueError: If the parser is not "PyPDF2" or "textract". + """ + if parser not in ("PyPDF2", "textract"): + raise ValueError( + f"Unsupported parser '{parser}'. Must be 'PyPDF2' or 'textract'." + ) + self.PARSER = parser + super().__init__( + id=id, + file_name=file_name, + file_path=file_path, + file_hash=file_hash, + created=created, + optional_metadata=optional_metadata, + ) + + def get_text(self) -> str: + """Extract and return the text content of the PDF document. + + Uses either PyPDF2 or textract depending on the chosen parser. + + Returns: + str: The extracted text with leading/trailing whitespace removed. + + Raises: + ValueError: If the PDF file is empty or the parser is unsupported. + RAGDocumentParsingError: If text extraction with textract fails. + """ + if self.PARSER == "PyPDF2": + from PyPDF2 import PdfReader + + reader = PdfReader(self.file_path) + if not reader.pages: + raise ValueError( + f"The PDF file {self.file_path} is empty or not valid." + ) + + text = "" + for page in reader.pages: + text += page.extract_text() or "" + + return text.strip() + elif self.PARSER == "textract": + import textract + + try: + text = textract.process(self.file_path, output_encoding="utf-8").decode( + "utf-8" + ) + text = _clean_textract_output(text) + return text.strip() + except Exception as e: + raise RAGDocumentParsingError( + f"Error extracting text from PDF file {self.file_path}: {str(e)}" + ) from e + else: + raise ValueError(f"Unsupported parser: {self.PARSER}") + + def get_metadata(self) -> Dict[str, Any]: + """Get the optional metadata associated with the document. + + Returns: + Dict[str, Any]: The metadata dictionary, or an empty dict + if no metadata was provided. + """ + return self.optional_metadata if self.optional_metadata else {} diff --git a/DashAI/back/models/RAG/documents/txt_document.py b/DashAI/back/models/RAG/documents/txt_document.py new file mode 100644 index 000000000..4ac18a768 --- /dev/null +++ b/DashAI/back/models/RAG/documents/txt_document.py @@ -0,0 +1,74 @@ +from typing import Any, Dict, Optional + +from DashAI.back.models.RAG.documents.base_document import BaseDocument +from DashAI.back.models.RAG.exceptions import RAGDocumentParsingError + + +class TxtDocument(BaseDocument): + """Document implementation for plain text (.txt) files. + + Reads text content directly from the file system with UTF-8 encoding. + """ + + def __init__( + self, + id: int, + file_name: str, + file_path: str, + file_hash: str, + created: Optional[str] = None, + optional_metadata: Optional[Dict[str, Any]] = None, + ): + """Initialize a TxtDocument instance. + + Args: + id: The unique identifier of the document. + file_name: The name of the file. + file_path: The path to the file. + file_hash: A hash of the file content (computed upstream). + created: The creation date of the document. + optional_metadata: Additional metadata for the document. + """ + super().__init__( + id=id, + file_name=file_name, + file_path=file_path, + file_hash=file_hash, + created=created, + optional_metadata=optional_metadata, + ) + + def get_text(self) -> str: + """Read and return the text content of the document. + + Returns: + str: The text content of the document with leading/trailing + whitespace removed. + + Raises: + RAGDocumentParsingError: If the file is not found, cannot be + read, or contains invalid UTF-8 encoding. + """ + try: + with open(self.file_path, "r", encoding="utf-8") as file: + text = file.read() + except FileNotFoundError as e: + raise RAGDocumentParsingError(f"File not found: {self.file_path}") from e + except IOError as e: + raise RAGDocumentParsingError( + f"Error reading file {self.file_path}: {str(e)}" + ) from e + except UnicodeDecodeError as e: + raise RAGDocumentParsingError( + f"Encoding error reading file {self.file_path}: {str(e)}" + ) from e + return text.strip() + + def get_metadata(self) -> Dict[str, Any]: + """Get the optional metadata associated with the document. + + Returns: + Dict[str, Any]: The metadata dictionary, or an empty dict + if no metadata was provided. + """ + return self.optional_metadata if self.optional_metadata else {} diff --git a/DashAI/back/models/RAG/embeddings/__init__.py b/DashAI/back/models/RAG/embeddings/__init__.py new file mode 100644 index 000000000..763ea8f8f --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/__init__.py @@ -0,0 +1,39 @@ +"""Embedding model exports. + +Provides public aliases for all dense embedding implementations +registered with the DashAI component system. +""" + +from DashAI.back.models.RAG.embeddings.dense.bert_embedding import BERTEmbedding +from DashAI.back.models.RAG.embeddings.dense.distilbert_embedding import ( + DistilBERTEmbedding, +) +from DashAI.back.models.RAG.embeddings.dense.e5_embedding import E5Embedding +from DashAI.back.models.RAG.embeddings.dense.gemma_embedding import GemmaEmbedding +from DashAI.back.models.RAG.embeddings.dense.huggingface_embedding import ( + HuggingFaceEmbedding, +) +from DashAI.back.models.RAG.embeddings.dense.instructor_embedding import ( + InstructorEmbedding, +) +from DashAI.back.models.RAG.embeddings.dense.labse_embedding import LaBSEmbedding +from DashAI.back.models.RAG.embeddings.dense.openai_embedding import OpenAIEmbedding +from DashAI.back.models.RAG.embeddings.dense.roberta_embedding import RoBERTaEmbedding +from DashAI.back.models.RAG.embeddings.dense.sentence_transformer_embedding import ( + SentenceTransformerEmbedding, +) +from DashAI.back.models.RAG.embeddings.dense_embedding import DenseEmbedding + +__all__ = [ + "DenseEmbedding", + "BERTEmbedding", + "DistilBERTEmbedding", + "E5Embedding", + "GemmaEmbedding", + "HuggingFaceEmbedding", + "InstructorEmbedding", + "LaBSEmbedding", + "OpenAIEmbedding", + "RoBERTaEmbedding", + "SentenceTransformerEmbedding", +] diff --git a/DashAI/back/models/RAG/embeddings/dense/__init__.py b/DashAI/back/models/RAG/embeddings/dense/__init__.py new file mode 100644 index 000000000..02e3e3b19 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/__init__.py @@ -0,0 +1,30 @@ +"""Dense embedding model implementations for the RAG module. + +Exposes all concrete dense embedding classes (BERT, DistilBERT, E5, Gemma, +HuggingFace, Instructor, LaBSE, OpenAI, RoBERTa, SentenceTransformer) for +registration in the DashAI component system. +""" + +from DashAI.back.models.RAG.embeddings.dense.bert_embedding import BERTEmbedding +from DashAI.back.models.RAG.embeddings.dense.distilbert_embedding import DistilBERTEmbedding +from DashAI.back.models.RAG.embeddings.dense.e5_embedding import E5Embedding +from DashAI.back.models.RAG.embeddings.dense.gemma_embedding import GemmaEmbedding +from DashAI.back.models.RAG.embeddings.dense.huggingface_embedding import HuggingFaceEmbedding +from DashAI.back.models.RAG.embeddings.dense.instructor_embedding import InstructorEmbedding +from DashAI.back.models.RAG.embeddings.dense.labse_embedding import LaBSEmbedding +from DashAI.back.models.RAG.embeddings.dense.openai_embedding import OpenAIEmbedding +from DashAI.back.models.RAG.embeddings.dense.roberta_embedding import RoBERTaEmbedding +from DashAI.back.models.RAG.embeddings.dense.sentence_transformer_embedding import SentenceTransformerEmbedding + +__all__ = [ + "BERTEmbedding", + "DistilBERTEmbedding", + "E5Embedding", + "GemmaEmbedding", + "HuggingFaceEmbedding", + "InstructorEmbedding", + "LaBSEmbedding", + "OpenAIEmbedding", + "RoBERTaEmbedding", + "SentenceTransformerEmbedding", +] diff --git a/DashAI/back/models/RAG/embeddings/dense/_bert_embedding.py b/DashAI/back/models/RAG/embeddings/dense/_bert_embedding.py new file mode 100644 index 000000000..d6e5edff1 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/_bert_embedding.py @@ -0,0 +1,106 @@ +from typing import Dict + +import torch +from transformers import AutoModel, AutoTokenizer + +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + OverflowHandler, +) + +CLS = "cls" +MEAN = "mean" +MAX = "max" +CONCAT_2 = "concat_2" +CONCAT_3 = "concat_3" +CONCAT_4 = "concat_4" + +POOLING_STRATEGIES: Dict[str, str] = { + CLS: "CLS token", + MEAN: "Mean pooling", + MAX: "Max pooling", + CONCAT_2: "Concat last 2 layers", + CONCAT_3: "Concat last 3 layers", + CONCAT_4: "Concat last 4 layers", +} + + +class _BERTEmbedding(OverflowHandler): + """Internal wrapper for BERT-family embedding models. + + Supports multiple pooling strategies: CLS token, mean, max, and + concatenation of the last N hidden layers. + """ + + POOLING_STRATEGIES: Dict[str, str] = POOLING_STRATEGIES + + def __init__( + self, + model_name: str, + device: str, + model_max_length: int, + overflow_strategy: str, + pooling_strategy: str, + ): + """Initialise the BERT embedding wrapper. + + Args: + model_name: HuggingFace model identifier. + device: Target device (``"cpu"`` or ``"cuda"``). + model_max_length: Maximum token length. + overflow_strategy: ``"truncate"`` or ``"aggregate"``. + pooling_strategy: One of ``"cls"``, ``"mean"``, ``"max"``, + ``"concat_2"``, ``"concat_3"``, ``"concat_4"``. + """ + super().__init__( + model_name=model_name, + device=device, + model_max_length=model_max_length, + overflow_strategy=overflow_strategy, + ) + self.pooling_strategy = pooling_strategy + self.params["pooling_strategy"] = pooling_strategy + + def load(self): + """Download tokenizer and model, enabling hidden states for concat strategies.""" # noqa: E501 + self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) + self.model = AutoModel.from_pretrained( + self.model_name, + output_hidden_states=self.pooling_strategy + in (CONCAT_2, CONCAT_3, CONCAT_4), + ).to(self.device) + + def _pool(self, model_output, attention_mask): + """Pool token embeddings according to the configured strategy. + + Args: + model_output: Output from ``AutoModel`` (may contain ``hidden_states``). + attention_mask: Attention mask tensor. + + Returns: + Pooled embedding tensor of shape ``(batch, embedding_dim)``. + + Raises: + ValueError: If the pooling strategy is unknown. + """ + if self.pooling_strategy == CLS: + return model_output[0][:, 0] + if self.pooling_strategy == MAX: + token_embeddings = model_output[0] + input_mask_expanded = ( + attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() + ) + token_embeddings[input_mask_expanded == 0] = -1e9 + return torch.max(token_embeddings, 1)[0] + if self.pooling_strategy in (CONCAT_2, CONCAT_3, CONCAT_4): + hidden_states = model_output.hidden_states + layer_count = int(self.pooling_strategy.split("_")[1]) + selected = [hidden_states[-(i + 1)][:, 0, :] for i in range(layer_count)] + return torch.cat(selected, dim=1) + token_embeddings = model_output[0] + input_mask_expanded = ( + attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() + ) + pooled = torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp( + input_mask_expanded.sum(1), min=1e-9 + ) + return pooled diff --git a/DashAI/back/models/RAG/embeddings/dense/_e5_embedding.py b/DashAI/back/models/RAG/embeddings/dense/_e5_embedding.py new file mode 100644 index 000000000..ecfadb165 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/_e5_embedding.py @@ -0,0 +1,81 @@ +from typing import List + +import numpy as np +import torch.nn.functional as functional + +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + OverflowHandler, +) + +QUERY_PREFIX = "query: " +PASSAGE_PREFIX = "passage: " + + +class _E5Embedding(OverflowHandler): + """Internal wrapper for E5 embedding models. + + Automatically prepends ``"passage: "`` to documents and ``"query: "`` + to queries, then applies mean pooling with L2 normalisation. + """ + + def __init__( + self, + model_name: str, + device: str, + model_max_length: int, + overflow_strategy: str, + ): + """Initialise the E5 embedding wrapper. + + Args: + model_name: HuggingFace model identifier. + device: Target device (``"cpu"`` or ``"cuda"``). + model_max_length: Maximum token length. + overflow_strategy: ``"truncate"`` or ``"aggregate"``. + """ + super().__init__( + model_name=model_name, + device=device, + model_max_length=model_max_length, + overflow_strategy=overflow_strategy, + ) + + def _pool(self, model_output, attention_mask): + """Mean-pool token embeddings with L2 normalisation. + + Args: + model_output: Output from ``AutoModel``. + attention_mask: Attention mask tensor. + + Returns: + Pooled and normalised tensor of shape ``(batch, embedding_dim)``. + """ + last_hidden = model_output[0] + mask = attention_mask[..., None].bool() + last_hidden = last_hidden.masked_fill(~mask, 0.0) + pooled = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None] + return functional.normalize(pooled, p=2, dim=1) + + def batch_encode(self, texts: List[str]) -> np.ndarray: + """Encode a batch of documents (prepends ``"passage: "``). + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + prefixed = [f"{PASSAGE_PREFIX}{t}" for t in texts] + return self._batch_encode_impl(prefixed) + + def encode(self, text: str) -> np.ndarray: + """Encode a single query (prepends ``"query: "``). + + Args: + text: Input string. + + Returns: + A 1-D float32 NumPy array of shape ``(embedding_dim,)``. + """ + result = self._batch_encode_impl([f"{QUERY_PREFIX}{text}"]) + return result.squeeze() diff --git a/DashAI/back/models/RAG/embeddings/dense/_gemma_embedding.py b/DashAI/back/models/RAG/embeddings/dense/_gemma_embedding.py new file mode 100644 index 000000000..b42ba29f3 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/_gemma_embedding.py @@ -0,0 +1,104 @@ +from typing import List + +import numpy as np +from sentence_transformers import SentenceTransformer + +from DashAI.back.models.RAG.embeddings.dense.huggingface_embedding import ( + HuggingFaceEmbedding, +) +from DashAI.back.models.RAG.exceptions import RAGEmbeddingError + +TASK_PROMPTS = { + "search_result": "task: search result | query: ", + "question_answering": "task: question answering | query: ", + "fact_checking": "task: fact checking | query: ", + "classification": "task: classification | query: ", + "clustering": "task: clustering | query: ", + "sentence_similarity": "task: sentence similarity | query: ", + "code_retrieval": "task: code retrieval | query: ", +} + +DOCUMENT_PROMPT = "title: none | text: " + + +class _GemmaEmbedding(HuggingFaceEmbedding): + """Internal wrapper around Gemma embedding models via SentenceTransformers. + + Uses task-specific query prompts (from :data:`TASK_PROMPTS`) for + encoding queries and a fixed document prompt for passages. + """ + + def __init__( + self, + model_name: str, + device: str, + model_max_length: int, + overflow_strategy: str, + task_type: str, + ): + """Initialise the Gemma embedding wrapper. + + Args: + model_name: HuggingFace model identifier. + device: Target device (``"cpu"`` or ``"cuda"``). + model_max_length: Maximum token length. + overflow_strategy: ``"truncate"`` or ``"aggregate"``. + task_type: Prompt template key from :data:`TASK_PROMPTS`. + """ + super().__init__(model_name=model_name, device=device) + self.model_max_length = model_max_length + self.overflow_strategy = overflow_strategy + self.task_type = task_type + self.params["model_max_length"] = model_max_length + self.params["overflow_strategy"] = overflow_strategy + self.params["task_type"] = task_type + + def _pool(self, model_output, attention_mask): + """Not implemented — Gemma uses SentenceTransformer API directly.""" + raise NotImplementedError( + "Gemma uses SentenceTransformer API, _pool is unused." + ) + + def load(self): + """Load the Gemma model via ``SentenceTransformer``.""" + self.model = SentenceTransformer(self.model_name, device=self.device) + + def batch_encode(self, texts: List[str]) -> np.ndarray: + """Encode a batch of documents with the fixed document prompt. + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + + Raises: + RAGEmbeddingError: If the model has not been loaded yet. + """ + if self.model is None: + raise RAGEmbeddingError( + "Model not loaded. Call load() before batch_encode()." + ) + prompted = [DOCUMENT_PROMPT + t for t in texts] + return self.model.encode( + prompted, normalize_embeddings=True, show_progress_bar=False + ) + + def encode(self, text: str) -> np.ndarray: + """Encode a single query with the task-specific prompt. + + Args: + text: Input string. + + Returns: + A 1-D float32 NumPy array of shape ``(embedding_dim,)``. + + Raises: + RAGEmbeddingError: If the model has not been loaded yet. + """ + if self.model is None: + raise RAGEmbeddingError("Model not loaded. Call load() before encode().") + query_prompt = TASK_PROMPTS[self.task_type] + return self.model.encode( + [query_prompt + text], normalize_embeddings=True, show_progress_bar=False + ).squeeze() diff --git a/DashAI/back/models/RAG/embeddings/dense/_instructor_embedding.py b/DashAI/back/models/RAG/embeddings/dense/_instructor_embedding.py new file mode 100644 index 000000000..8d6427757 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/_instructor_embedding.py @@ -0,0 +1,75 @@ +from typing import List + +import numpy as np + +try: + from InstructorEmbedding import INSTRUCTOR +except ImportError as err: + raise ImportError( + "InstructorEmbedding package is not installed. " + "Install it with: pip install InstructorEmbedding" + ) from err + +from DashAI.back.models.RAG.embeddings.dense.huggingface_embedding import ( + HuggingFaceEmbedding, +) + + +class _InstructorEmbedding(HuggingFaceEmbedding): + """Internal wrapper around the ``InstructorEmbedding`` package. + + Prepends a fixed instruction string to every input text and delegates + encoding to the ``INSTRUCTOR`` model (which handles its own pooling). + """ + + def __init__( + self, + model_name: str, + device: str, + instruction: str, + ): + """Initialise the INSTRUCTOR embedding wrapper. + + Args: + model_name: HuggingFace model identifier for the INSTRUCTOR model. + device: Target device (ignored — INSTRUCTOR manages its own device). + instruction: Instruction text prepended to every query / document. + """ + super().__init__(model_name=model_name, device=device) + self.instruction = instruction + self.params["instruction"] = instruction + + def _pool(self, model_output, attention_mask): + """Not implemented — INSTRUCTOR uses its own encoding API.""" + raise NotImplementedError( + "INSTRUCTOR uses custom encoding API, _pool is unused." + ) + + def load(self): + """Instantiate the ``INSTRUCTOR`` model from the ``InstructorEmbedding`` package.""" # noqa: E501 + self.model = INSTRUCTOR(self.model_name) + self.model._text_length = self.model._input_length + + def batch_encode(self, texts: List[str]) -> np.ndarray: + """Encode a batch of texts with the fixed instruction prefix. + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + pairs = [[self.instruction, text] for text in texts] + return self.model.encode(pairs, show_progress_bar=False) + + def encode(self, text: str) -> np.ndarray: + """Encode a single text with the fixed instruction prefix. + + Args: + text: Input string. + + Returns: + A 1-D float32 NumPy array of shape ``(embedding_dim,)``. + """ + result = self.model.encode([[self.instruction, text]], show_progress_bar=False) + return result.squeeze() diff --git a/DashAI/back/models/RAG/embeddings/dense/_overflow_handler.py b/DashAI/back/models/RAG/embeddings/dense/_overflow_handler.py new file mode 100644 index 000000000..092e0b583 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/_overflow_handler.py @@ -0,0 +1,107 @@ +from abc import abstractmethod +from typing import List + +import numpy as np +import torch + +from DashAI.back.models.RAG.embeddings.dense.huggingface_embedding import ( + HuggingFaceEmbedding, +) + +TRUNCATE = "truncate" +AGGREGATE = "aggregate" + + +class OverflowHandler(HuggingFaceEmbedding): + """Extends :class:`HuggingFaceEmbedding` with overflow handling strategies. + + When text exceeds ``model_max_length``, the tokeniser either truncates + (``"truncate"``) or splits into segments and pools their embeddings + together (``"aggregate"``). + """ + + def __init__( + self, + model_name: str, + device: str, + model_max_length: int, + overflow_strategy: str, + ): + """Initialise the overflow handler. + + Args: + model_name: HuggingFace model identifier. + device: Target device (``"cpu"`` or ``"cuda"``). + model_max_length: Maximum token length before overflow logic kicks in. + overflow_strategy: ``"truncate"`` or ``"aggregate"``. + """ + super().__init__(model_name=model_name, device=device) + self.model_max_length = model_max_length + self.overflow_strategy = overflow_strategy + self.params["model_max_length"] = model_max_length + self.params["overflow_strategy"] = overflow_strategy + + @abstractmethod + def _pool(self, model_output, attention_mask): + """Aggregate token-level hidden states into a single embedding per item. + + Args: + model_output: Output tuple from ``AutoModel``. + attention_mask: Attention mask tensor with shape ``(batch, seq_len)``. + + Returns: + A torch tensor of shape ``(batch, embedding_dim)``. + """ + raise NotImplementedError + + def _batch_encode_impl(self, texts: List[str]) -> np.ndarray: + """Tokenise and encode, applying the configured overflow strategy. + + Args: + texts: Batch of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + + Raises: + ValueError: If ``overflow_strategy`` is not supported. + """ + if self.overflow_strategy not in [TRUNCATE, AGGREGATE]: + raise ValueError( + f"Invalid overflow strategy: {self.overflow_strategy}. " + f"Supported strategies are: {TRUNCATE}, {AGGREGATE}." + ) + encoded = self.tokenizer( + texts, + padding=True, + truncation=True, + max_length=self.model_max_length, + return_tensors="pt", + return_overflowing_tokens=self.overflow_strategy == AGGREGATE, + stride=0, + ) + num_texts = len(texts) + if ( + self.overflow_strategy == AGGREGATE + and "overflow_to_sample_mapping" in encoded + ): + mapping = encoded.pop("overflow_to_sample_mapping") + encoded = {k: v.to(self.device) for k, v in encoded.items()} + with torch.no_grad(): + outputs = self.model(**encoded) + segment_embeddings = self._pool(outputs, encoded["attention_mask"]).cpu() + result = [] + for i in range(num_texts): + segment_indices = (mapping == i).nonzero(as_tuple=True)[0] + if len(segment_indices) == 1: + result.append(segment_embeddings[segment_indices[0]]) + else: + result.append( + torch.mean(segment_embeddings[segment_indices], dim=0) + ) + return torch.stack(result).numpy() + encoded = {k: v.to(self.device) for k, v in encoded.items()} + with torch.no_grad(): + outputs = self.model(**encoded) + embeddings = self._pool(outputs, encoded["attention_mask"]) + return embeddings.cpu().numpy() diff --git a/DashAI/back/models/RAG/embeddings/dense/_sentence_transformer_embedding.py b/DashAI/back/models/RAG/embeddings/dense/_sentence_transformer_embedding.py new file mode 100644 index 000000000..f4ecc8b66 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/_sentence_transformer_embedding.py @@ -0,0 +1,97 @@ +import torch +import torch.nn.functional as functional + +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + OverflowHandler, +) + +TRUNCATE = "truncate" +AGGREGATE = "aggregate" + + +class _SentenceTransformerEmbedding(OverflowHandler): + """Internal wrapper that adds mean / last-token pooling and L2 normalisation + on top of :class:`OverflowHandler`. + + Used by :class:`SentenceTransformerEmbedding` and :class:`LaBSEmbedding`. + """ + + def __init__( + self, + model_name: str, + device: str, + model_max_length: int, + overflow_strategy: str, + normalize: bool, + pooling: str = "mean", + ): + """Initialise the embedding with pooling and normalisation options. + + Args: + model_name: HuggingFace model identifier. + device: Target device (``"cpu"`` or ``"cuda"``). + model_max_length: Maximum number of tokens the model accepts. + overflow_strategy: ``"truncate"`` or ``"aggregate"``. + normalize: Whether to L2-normalise the output embeddings. + pooling: Pooling method — ``"mean"`` or ``"last_token"``. + """ + super().__init__( + model_name=model_name, + device=device, + model_max_length=model_max_length, + overflow_strategy=overflow_strategy, + ) + self.normalize = normalize + self.pooling = pooling + self.params["normalize"] = normalize + self.params["pooling"] = pooling + + def _pool(self, model_output, attention_mask): + """Mean-pool token embeddings, optionally normalised. + + Args: + model_output: Output from ``AutoModel``. + attention_mask: Attention mask tensor. + + Returns: + Pooled embedding tensor of shape ``(batch, embedding_dim)``. + """ + if self.pooling == "last_token": + return self._last_token_pool(model_output, attention_mask) + token_embeddings = model_output[0] + input_mask_expanded = ( + attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() + ) + pooled = torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp( + input_mask_expanded.sum(1), min=1e-9 + ) + if self.normalize: + pooled = functional.normalize(pooled, p=2, dim=1) + return pooled + + def _last_token_pool(self, model_output, attention_mask): + """Extract the last non-padding token's hidden state. + + Handles both left-padded and right-padded sequences. + + Args: + model_output: Output from ``AutoModel``. + attention_mask: Attention mask tensor. + + Returns: + Pooled embedding tensor of shape ``(batch, embedding_dim)``. + """ + last_hidden_states = model_output[0] + left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0] + if left_padding: + pooled = last_hidden_states[:, -1] + else: + sequence_lengths = attention_mask.sum(dim=1) - 1 + batch_size = last_hidden_states.shape[0] + pooled = last_hidden_states[ + torch.arange(batch_size, device=last_hidden_states.device), + sequence_lengths, + ] + if self.normalize: + pooled = functional.normalize(pooled, p=2, dim=1) + return pooled diff --git a/DashAI/back/models/RAG/embeddings/dense/bert_embedding.py b/DashAI/back/models/RAG/embeddings/dense/bert_embedding.py new file mode 100644 index 000000000..8e71f89fb --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/bert_embedding.py @@ -0,0 +1,269 @@ +from typing import Dict, List + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.embeddings.dense._bert_embedding import ( + CLS, + CONCAT_2, + CONCAT_3, + CONCAT_4, + MAX, + MEAN, + _BERTEmbedding, +) +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + AGGREGATE, + TRUNCATE, +) +from DashAI.back.models.RAG.embeddings.dense_embedding import DenseEmbedding + +BERT_POOLING_STRATEGIES = [CLS, MEAN, MAX, CONCAT_2, CONCAT_3, CONCAT_4] + +BERT_MODELS: Dict[str, dict] = { + "google-bert/bert-base-cased": { + "languages": ["en"], + "max_seq_length": 512, + }, + "google-bert/bert-base-uncased": { + "languages": ["en"], + "max_seq_length": 512, + }, + "google-bert/bert-large-cased": { + "languages": ["en"], + "max_seq_length": 512, + }, + "google-bert/bert-large-uncased": { + "languages": ["en"], + "max_seq_length": 512, + }, + "google-bert/bert-base-multilingual-cased": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "pt", + "nl", + "pl", + "ca", + "fi", + "ar", + "zh", + "ja", + "ko", + "ru", + "tr", + "hi", + "sv", + "da", + "no", + "cs", + "ro", + "el", + "he", + "hu", + "th", + "vi", + "id", + "ms", + "bg", + "hr", + "sk", + "sl", + "sr", + "uk", + "et", + "lv", + "lt", + "fa", + "ur", + "mk", + "af", + "bn", + "multi", + ], + "max_seq_length": 512, + }, + "google-bert/bert-base-multilingual-uncased": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "pt", + "nl", + "pl", + "ca", + "fi", + "ar", + "zh", + "ja", + "ko", + "ru", + "tr", + "hi", + "sv", + "da", + "no", + "cs", + "ro", + "el", + "he", + "hu", + "th", + "vi", + "id", + "ms", + "bg", + "hr", + "sk", + "sl", + "sr", + "uk", + "et", + "lv", + "lt", + "fa", + "ur", + "mk", + "af", + "bn", + "multi", + ], + "max_seq_length": 512, + }, +} + +BERT_MODEL_NAMES = list(BERT_MODELS.keys()) + + +class BERTEmbeddingSchema(BaseSchema): + """Configuration schema for :class:`BERTEmbedding`. + + Attributes: + model_name: BERT model for embedding generation. + overflow_strategy: Strategy for chunks exceeding model max sequence length. + device: Device to run the model on. + pooling_strategy: Pooling strategy to aggregate token embeddings. + """ + + model_name: schema_field( + enum_field(BERT_MODEL_NAMES), + placeholder="google-bert/bert-base-cased", + description=MultilingualString( + en="BERT model for embedding generation.", + es="Modelo BERT para generación de embeddings.", + ), + ) # type: ignore + + overflow_strategy: schema_field( + enum_field([TRUNCATE, AGGREGATE]), + placeholder=TRUNCATE, + description=MultilingualString( + en="Strategy for chunks exceeding model max sequence length.", + es="Estrategia para fragmentos que exceden la longitud máxima del modelo.", + ), + ) # type: ignore + + device: schema_field( + enum_field(["cpu", "cuda"]), + placeholder="cpu", + description=MultilingualString( + en="Device to run the model on.", + es="Dispositivo para ejecutar el modelo.", + ), + ) # type: ignore + + pooling_strategy: schema_field( + enum_field(BERT_POOLING_STRATEGIES), + placeholder=MEAN, + description=MultilingualString( + en="Pooling strategy to aggregate token embeddings.", + es="Estrategia de pooling para agregar embeddings de tokens.", + ), + ) # type: ignore + + +class BERTEmbedding(DenseEmbedding): + """Dense embeddings using BERT models with configurable pooling. + + Wraps :class:`_BERTEmbedding` and exposes it as a DashAI component with + a configurable schema (:class:`BERTEmbeddingSchema`). + + Supports CLS, mean, max and concat-layer pooling strategies. + + FLAGS: + FAMILY:bert: Groups this model under the BERT family. + huggingface: Marks the model family as HuggingFace-based. + """ + + SCHEMA = BERTEmbeddingSchema + FLAGS: list[str] = ["FAMILY:bert", "huggingface"] + DISPLAY_NAME: str = MultilingualString( + en="BERT Embedding", + es="Embedding BERT", + ) + DESCRIPTION: str = MultilingualString( + en="Dense embeddings using BERT models with configurable pooling" + " (CLS, mean, max, concat layers).", + es="Embeddings densos usando modelos BERT con pooling configurable" + " (CLS, mean, max, concat layers).", + ) + + def __init__(self, **kwargs): + """Initialise the embedding by validating parameters and creating the internal model. + + Args: + **kwargs: Configuration matching :class:`BERTEmbeddingSchema`. + """ # noqa: E501 + self.params = self.validate_and_transform(kwargs) + model_name = self.params["model_name"] + device = self.params["device"] + overflow_strategy = self.params.get("overflow_strategy", "truncate") + pooling_strategy = self.params["pooling_strategy"] + model_info = BERT_MODELS[model_name] + self._embedding = _BERTEmbedding( + model_name=model_name, + device=device, + model_max_length=model_info["max_seq_length"], + overflow_strategy=overflow_strategy, + pooling_strategy=pooling_strategy, + ) + + def load(self): + """Load the BERT model and tokenizer.""" + self._embedding.load() + + def encode(self, text: str): + """Encode a single text into a dense embedding. + + Args: + text: Input string. + + Returns: + A 1-D NumPy array of shape ``(embedding_dim,)``. + """ + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + """Encode a batch of texts into dense embeddings. + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + return self._embedding.batch_encode(texts) + + def save(self): + """No-op. Persistence is handled externally.""" + + def train(self, **kwargs): + """No-op. Pre-trained models are used as-is.""" diff --git a/DashAI/back/models/RAG/embeddings/dense/distilbert_embedding.py b/DashAI/back/models/RAG/embeddings/dense/distilbert_embedding.py new file mode 100644 index 000000000..142a799e0 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/distilbert_embedding.py @@ -0,0 +1,213 @@ +from typing import Dict, List + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.embeddings.dense._bert_embedding import ( + CLS, + CONCAT_2, + CONCAT_3, + CONCAT_4, + MAX, + MEAN, + _BERTEmbedding, +) +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + AGGREGATE, + TRUNCATE, +) +from DashAI.back.models.RAG.embeddings.dense_embedding import DenseEmbedding + +DISTILBERT_POOLING_STRATEGIES = [CLS, MEAN, MAX, CONCAT_2, CONCAT_3, CONCAT_4] + +DISTILBERT_MODELS: Dict[str, dict] = { + "distilbert/distilbert-base-cased": { + "languages": ["en"], + "max_seq_length": 512, + }, + "distilbert/distilbert-base-uncased": { + "languages": ["en"], + "max_seq_length": 512, + }, + "distilbert/distilbert-base-multilingual-cased": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "pt", + "nl", + "pl", + "ca", + "fi", + "ar", + "zh", + "ja", + "ko", + "ru", + "tr", + "hi", + "sv", + "da", + "no", + "cs", + "ro", + "el", + "he", + "hu", + "th", + "vi", + "id", + "ms", + "bg", + "hr", + "sk", + "sl", + "sr", + "uk", + "et", + "lv", + "lt", + "fa", + "ur", + "mk", + "af", + "bn", + "multi", + ], + "max_seq_length": 512, + }, + "distilbert/distilbert-roberta-base": { + "languages": ["en"], + "max_seq_length": 512, + }, +} + +DISTILBERT_MODEL_NAMES = list(DISTILBERT_MODELS.keys()) + + +class DistilBERTEmbeddingSchema(BaseSchema): + """Configuration schema for :class:`DistilBERTEmbedding`. + + Attributes: + model_name: DistilBERT model for embedding generation. + overflow_strategy: Strategy for chunks exceeding model max sequence length. + device: Device to run the model on. + pooling_strategy: Pooling strategy to aggregate token embeddings. + """ + + model_name: schema_field( + enum_field(DISTILBERT_MODEL_NAMES), + placeholder="distilbert/distilbert-base-cased", + description=MultilingualString( + en="DistilBERT model for embedding generation.", + es="Modelo DistilBERT para generación de embeddings.", + ), + ) # type: ignore + + overflow_strategy: schema_field( + enum_field([TRUNCATE, AGGREGATE]), + placeholder=TRUNCATE, + description=MultilingualString( + en="Strategy for chunks exceeding model max sequence length.", + es="Estrategia para fragmentos que exceden la longitud máxima del modelo.", + ), + ) # type: ignore + + device: schema_field( + enum_field(["cpu", "cuda"]), + placeholder="cpu", + description=MultilingualString( + en="Device to run the model on.", + es="Dispositivo para ejecutar el modelo.", + ), + ) # type: ignore + + pooling_strategy: schema_field( + enum_field(DISTILBERT_POOLING_STRATEGIES), + placeholder=MEAN, + description=MultilingualString( + en="Pooling strategy to aggregate token embeddings.", + es="Estrategia de pooling para agregar embeddings de tokens.", + ), + ) # type: ignore + + +class DistilBERTEmbedding(DenseEmbedding): + """Dense embeddings using DistilBERT models with configurable pooling. + + Wraps :class:`_BERTEmbedding` (reusing BERT pooling logic) and exposes + it as a DashAI component with a configurable schema + (:class:`DistilBERTEmbeddingSchema`). + + FLAGS: + FAMILY:distilbert: Groups this model under the DistilBERT family. + huggingface: Marks the model family as HuggingFace-based. + """ + + SCHEMA = DistilBERTEmbeddingSchema + FLAGS: list[str] = ["FAMILY:distilbert", "huggingface"] + DISPLAY_NAME: str = MultilingualString( + en="DistilBERT Embedding", + es="Embedding DistilBERT", + ) + DESCRIPTION: str = MultilingualString( + en="Dense embeddings using DistilBERT models with configurable pooling.", + es="Embeddings densos usando modelos DistilBERT con pooling configurable.", + ) + + def __init__(self, **kwargs): + """Initialise the embedding by validating parameters and creating the internal model. + + Args: + **kwargs: Configuration matching :class:`DistilBERTEmbeddingSchema`. + """ # noqa: E501 + self.params = self.validate_and_transform(kwargs) + model_name = self.params["model_name"] + device = self.params["device"] + overflow_strategy = self.params.get("overflow_strategy", "truncate") + pooling_strategy = self.params["pooling_strategy"] + model_info = DISTILBERT_MODELS[model_name] + self._embedding = _BERTEmbedding( + model_name=model_name, + device=device, + model_max_length=model_info["max_seq_length"], + overflow_strategy=overflow_strategy, + pooling_strategy=pooling_strategy, + ) + + def load(self): + """Load the DistilBERT model and tokenizer.""" + self._embedding.load() + + def encode(self, text: str): + """Encode a single text into a dense embedding. + + Args: + text: Input string. + + Returns: + A 1-D NumPy array of shape ``(embedding_dim,)``. + """ + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + """Encode a batch of texts into dense embeddings. + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + return self._embedding.batch_encode(texts) + + def save(self): + """No-op. Persistence is handled externally.""" + + def train(self, **kwargs): + """No-op. Pre-trained models are used as-is.""" diff --git a/DashAI/back/models/RAG/embeddings/dense/e5_embedding.py b/DashAI/back/models/RAG/embeddings/dense/e5_embedding.py new file mode 100644 index 000000000..fd9aef89e --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/e5_embedding.py @@ -0,0 +1,293 @@ +from typing import Dict, List + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.embeddings.dense._e5_embedding import _E5Embedding +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + AGGREGATE, + TRUNCATE, +) +from DashAI.back.models.RAG.embeddings.dense_embedding import DenseEmbedding + +E5_MODELS: Dict[str, dict] = { + "intfloat/e5-small-v2": { + "languages": ["en"], + "max_seq_length": 512, + }, + "intfloat/e5-large-v2": { + "languages": ["en"], + "max_seq_length": 512, + }, + "intfloat/multilingual-e5-large": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "pt", + "nl", + "pl", + "ca", + "fi", + "ar", + "zh", + "ja", + "ko", + "ru", + "tr", + "hi", + "sv", + "da", + "no", + "cs", + "ro", + "el", + "he", + "hu", + "th", + "vi", + "id", + "ms", + "bg", + "hr", + "sk", + "sl", + "sr", + "uk", + "et", + "lv", + "lt", + "fa", + "ur", + "mk", + "af", + "bn", + "multi", + ], + "max_seq_length": 512, + }, + "intfloat/multilingual-e5-base": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "pt", + "nl", + "pl", + "ca", + "fi", + "ar", + "zh", + "ja", + "ko", + "ru", + "tr", + "hi", + "sv", + "da", + "no", + "cs", + "ro", + "el", + "he", + "hu", + "th", + "vi", + "id", + "ms", + "bg", + "hr", + "sk", + "sl", + "sr", + "uk", + "et", + "lv", + "lt", + "fa", + "ur", + "mk", + "af", + "bn", + "multi", + ], + "max_seq_length": 512, + }, + "intfloat/multilingual-e5-small": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "pt", + "nl", + "pl", + "ca", + "fi", + "ar", + "zh", + "ja", + "ko", + "ru", + "tr", + "hi", + "sv", + "da", + "no", + "cs", + "ro", + "el", + "he", + "hu", + "th", + "vi", + "id", + "ms", + "bg", + "hr", + "sk", + "sl", + "sr", + "uk", + "et", + "lv", + "lt", + "fa", + "ur", + "mk", + "af", + "bn", + "multi", + ], + "max_seq_length": 512, + }, + "intfloat/e5-mistral-7b-instruct": { + "languages": ["en"], + "max_seq_length": 4096, + }, +} + +E5_MODEL_NAMES = list(E5_MODELS.keys()) + + +class E5EmbeddingSchema(BaseSchema): + """Configuration schema for :class:`E5Embedding`. + + Attributes: + model_name: E5 model for embedding generation (uses query/passage prefixes). + overflow_strategy: Strategy for chunks exceeding model max sequence length. + device: Device to run the model on. + """ + + model_name: schema_field( + enum_field(E5_MODEL_NAMES), + placeholder="intfloat/e5-small-v2", + description=MultilingualString( + en="E5 model for embedding generation (uses query/passage prefixes).", + es="Modelo E5 para generación de embeddings (usa prefijos query/passage).", + ), + ) # type: ignore + + overflow_strategy: schema_field( + enum_field([TRUNCATE, AGGREGATE]), + placeholder=TRUNCATE, + description=MultilingualString( + en="Strategy for chunks exceeding model max sequence length.", + es="Estrategia para fragmentos que exceden la longitud máxima del modelo.", + ), + ) # type: ignore + + device: schema_field( + enum_field(["cpu", "cuda"]), + placeholder="cpu", + description=MultilingualString( + en="Device to run the model on.", + es="Dispositivo para ejecutar el modelo.", + ), + ) # type: ignore + + +class E5Embedding(DenseEmbedding): + """Dense embeddings using E5 models with average pooling + L2 normalization. + + Automatically prepends ``"query: "`` or ``"passage: "`` prefixes to + input text (see :class:`_E5Embedding`). + + Wraps :class:`_E5Embedding` and exposes it as a DashAI component with + a configurable schema (:class:`E5EmbeddingSchema`). + + FLAGS: + FAMILY:e5: Groups this model under the E5 family. + huggingface: Marks the model family as HuggingFace-based. + """ + + SCHEMA = E5EmbeddingSchema + FLAGS: list[str] = ["FAMILY:e5", "huggingface"] + DISPLAY_NAME: str = MultilingualString( + en="E5 Embedding", + es="Embedding E5", + ) + DESCRIPTION: str = MultilingualString( + en="Dense embeddings using E5 models with average pooling + L2" + " normalization + query/passage prefixes.", + es="Embeddings densos usando modelos E5 con average pooling +" + " normalización L2 + prefijos query/passage.", + ) + + def __init__(self, **kwargs): + """Initialise the embedding by validating parameters and creating the internal model. + + Args: + **kwargs: Configuration matching :class:`E5EmbeddingSchema`. + """ # noqa: E501 + self.params = self.validate_and_transform(kwargs) + model_name = self.params["model_name"] + device = self.params["device"] + overflow_strategy = self.params.get("overflow_strategy", "truncate") + model_info = E5_MODELS[model_name] + self._embedding = _E5Embedding( + model_name=model_name, + device=device, + model_max_length=model_info["max_seq_length"], + overflow_strategy=overflow_strategy, + ) + + def load(self): + """Load the E5 model and tokenizer.""" + self._embedding.load() + + def encode(self, text: str): + """Encode a single text into a dense embedding (prepends ``"query: "``). + + Args: + text: Input string. + + Returns: + A 1-D NumPy array of shape ``(embedding_dim,)``. + """ + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + """Encode a batch of texts into dense embeddings (prepends ``"passage: "``). + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + return self._embedding.batch_encode(texts) + + def save(self): + """No-op. Persistence is handled externally.""" + + def train(self, **kwargs): + """No-op. Pre-trained models are used as-is.""" diff --git a/DashAI/back/models/RAG/embeddings/dense/fasttext_embedding.py b/DashAI/back/models/RAG/embeddings/dense/fasttext_embedding.py new file mode 100644 index 000000000..ac5cbb08a --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/fasttext_embedding.py @@ -0,0 +1,113 @@ +from typing import List + +import fasttext +import numpy as np +from huggingface_hub import hf_hub_download + +from DashAI.back.core.schema_fields import enum_field +from DashAI.back.core.schema_fields.base_schema import BaseSchema +from DashAI.back.core.schema_fields.schema_field import schema_field +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.embeddings.dense_embedding import DenseEmbedding +from DashAI.back.models.RAG.exceptions import RAGEmbeddingEmptyInputError + + +class FastTextEmbeddingSchema(BaseSchema): + """Configuration schema for :class:`FastTextEmbedding`. + + Attributes: + model_name: Name of the pre-trained FastText model to use. + pooling_strategy: Pooling strategy (``"mean"`` or ``"max"``). + """ + + model_name: schema_field( + enum_field(["facebook/fasttext-es-vectors", "facebook/fasttext-en-vectors"]), + "facebook/fasttext-en-vectors", + "Name of the pre-trained model to use", + ) # type: ignore + + pooling_strategy: schema_field( + enum_field(["mean", "max"]), + "mean", + "Pooling strategy to use", + ) # type: ignore + + +class FastTextEmbedding(DenseEmbedding): + """Dense embeddings using FastText word vectors with mean or max pooling. + + Downloads the model binary from the HuggingFace Hub and aggregates + word-level vectors into a single sentence-level embedding. + + FLAGS: + FAMILY:fasttext: Groups this model under the FastText family. + """ + + FLAGS: list[str] = ["FAMILY:fasttext"] + SCHEMA = FastTextEmbeddingSchema + DISPLAY_NAME = MultilingualString( + en="FastText Embedding", + es="Embedding FastText", + ) + DESCRIPTION = MultilingualString( + en="Convert text to embeddings using FastText.", + es="Convierte texto a embeddings usando FastText.", + ) + + def __init__(self, **kwargs): + """Initialise the embedding by validating parameters and setting up pooling. + + Args: + **kwargs: Configuration matching :class:`FastTextEmbeddingSchema`. + """ + self.params = self.validate_and_transform(kwargs) + self.model_name = self.params["model_name"] + self.pooling_strategy = self.params["pooling_strategy"] + pooling_functions = {"mean": np.mean, "max": np.max} + self.pooling_function = pooling_functions[self.pooling_strategy] + self.model = None + + def load(self): + """Download and load the FastText binary model from the HuggingFace Hub.""" + if self.model is not None: + return + model_path = hf_hub_download(repo_id=self.model_name, filename="model.bin") + self.model = fasttext.load_model(model_path) + + def save(self): + """No-op. The model is loaded from HF Hub on demand.""" + + def train(self, **kwargs): + """No-op. Pre-trained FastText vectors are used as-is.""" + + def encode(self, text: str) -> np.ndarray: + """Encode a single text by averaging/max-pooling its word vectors. + + Args: + text: Input string (non-empty). + + Returns: + A 1-D float32 NumPy array of shape ``(embedding_dim,)``. + + Raises: + RAGEmbeddingEmptyInputError: If the input text is empty or blank. + """ + if not text or not text.strip(): + raise RAGEmbeddingEmptyInputError("Cannot encode empty text") + token_embeddings = [self.model.get_word_vector(word) for word in text.split()] + return self.pooling_function(np.array(token_embeddings), axis=0) + + def batch_encode(self, texts: List[str]) -> np.ndarray: + """Encode a batch of texts by pooling their word vectors. + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + all_embeddings = [] + for text in texts: + embedding = self.encode(text) + all_embeddings.append(embedding) + return np.array(all_embeddings) diff --git a/DashAI/back/models/RAG/embeddings/dense/gemma_embedding.py b/DashAI/back/models/RAG/embeddings/dense/gemma_embedding.py new file mode 100644 index 000000000..4f60cc3c8 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/gemma_embedding.py @@ -0,0 +1,153 @@ +from typing import Dict, List + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.embeddings.dense._gemma_embedding import ( + TASK_PROMPTS, + _GemmaEmbedding, +) +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + AGGREGATE, + TRUNCATE, +) +from DashAI.back.models.RAG.embeddings.dense_embedding import DenseEmbedding + +TASK_TYPES = list(TASK_PROMPTS.keys()) + +GEMMA_MODELS: Dict[str, dict] = { + "google/embeddinggemma-300m": { + "languages": ["en"], + "max_seq_length": 8192, + }, +} + +GEMMA_MODEL_NAMES = list(GEMMA_MODELS.keys()) + + +class GemmaEmbeddingSchema(BaseSchema): + """Configuration schema for :class:`GemmaEmbedding`. + + Attributes: + model_name: Gemma embedding model (uses SentenceTransformers API). + task_type: Task type that optimises the query embedding for a specific use case. + overflow_strategy: Strategy for chunks exceeding model max sequence length. + device: Device to run the model on. + """ + + model_name: schema_field( + enum_field(GEMMA_MODEL_NAMES), + placeholder="google/embeddinggemma-300m", + description=MultilingualString( + en="Gemma embedding model (uses SentenceTransformers API).", + es="Modelo de embedding Gemma (usa API SentenceTransformers).", + ), + ) # type: ignore + + task_type: schema_field( + enum_field(TASK_TYPES), + placeholder="search_result", + description=MultilingualString( + en="Task type that optimizes the query embedding for a specific use case.", + es="Tipo de tarea que optimiza el embedding de consulta para" + " un caso de uso específico.", + ), + ) # type: ignore + + overflow_strategy: schema_field( + enum_field([TRUNCATE, AGGREGATE]), + placeholder=TRUNCATE, + description=MultilingualString( + en="Strategy for chunks exceeding model max sequence length.", + es="Estrategia para fragmentos que exceden la longitud máxima del modelo.", + ), + ) # type: ignore + + device: schema_field( + enum_field(["cpu", "cuda"]), + placeholder="cpu", + description=MultilingualString( + en="Device to run the model on.", + es="Dispositivo para ejecutar el modelo.", + ), + ) # type: ignore + + +class GemmaEmbedding(DenseEmbedding): + """Dense embeddings using Gemma models via the SentenceTransformers API. + + Wraps :class:`_GemmaEmbedding` and exposes it as a DashAI component with + a configurable schema (:class:`GemmaEmbeddingSchema`). + + Supports task-aware query prompts (search, QA, classification, etc.). + + FLAGS: + FAMILY:gemma: Groups this model under the Gemma family. + huggingface: Marks the model family as HuggingFace-based. + """ + + SCHEMA = GemmaEmbeddingSchema + FLAGS: list[str] = ["FAMILY:gemma", "huggingface"] + DISPLAY_NAME: str = MultilingualString( + en="Gemma Embedding", + es="Embedding Gemma", + ) + DESCRIPTION: str = MultilingualString( + en="Dense embeddings using Gemma models via SentenceTransformers API.", + es="Embeddings densos usando modelos Gemma mediante API SentenceTransformers.", + ) + + def __init__(self, **kwargs): + """Initialise the embedding by validating parameters and creating the internal model. + + Args: + **kwargs: Configuration matching :class:`GemmaEmbeddingSchema`. + """ # noqa: E501 + self.params = self.validate_and_transform(kwargs) + model_name = self.params["model_name"] + device = self.params["device"] + overflow_strategy = self.params.get("overflow_strategy", "truncate") + task_type = self.params["task_type"] + model_info = GEMMA_MODELS[model_name] + self._embedding = _GemmaEmbedding( + model_name=model_name, + device=device, + model_max_length=model_info["max_seq_length"], + overflow_strategy=overflow_strategy, + task_type=task_type, + ) + + def load(self): + """Load the Gemma model via SentenceTransformers.""" + self._embedding.load() + + def encode(self, text: str): + """Encode a single text into a dense embedding with the configured task prompt. + + Args: + text: Input string. + + Returns: + A 1-D NumPy array of shape ``(embedding_dim,)``. + """ + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + """Encode a batch of texts into dense embeddings with document prompts. + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + return self._embedding.batch_encode(texts) + + def save(self): + """No-op. Persistence is handled externally.""" + + def train(self, **kwargs): + """No-op. Pre-trained models are used as-is.""" diff --git a/DashAI/back/models/RAG/embeddings/dense/huggingface_embedding.py b/DashAI/back/models/RAG/embeddings/dense/huggingface_embedding.py new file mode 100644 index 000000000..c52772fab --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/huggingface_embedding.py @@ -0,0 +1,107 @@ +from abc import abstractmethod +from typing import List + +import numpy as np +import torch +from transformers import AutoModel, AutoTokenizer + +from DashAI.back.models.RAG.embeddings.dense_embedding import DenseEmbedding + + +class HuggingFaceEmbedding(DenseEmbedding): + """Abstract base for dense embeddings powered by a HuggingFace ``AutoModel``. + + Handles tokenisation, device placement and inference dispatch; subclasses + only need to implement :meth:`_pool` to convert model outputs into a single + vector per text. + + FLAGS: + abstract: Marker indicating this class is not meant for direct use. + huggingface: Marks the model family as HuggingFace-based. + """ + + FLAGS: list[str] = ["abstract", "huggingface"] + + def __init__(self, model_name: str, device: str): + """Initialise the HuggingFace embedding wrapper. + + Args: + model_name: Name or path of a HuggingFace ``AutoModel``. + device: Target device (``"cpu"`` or ``"cuda"``). + """ + self.model_name = model_name + self.device = device + self.params = { + "model_name": model_name, + "device": device, + } + self.model = None + self.tokenizer = None + + def save(self): + """No-op. Persistence is handled externally.""" + + def train(self, **kwargs): + """No-op. Pre-trained models are used as-is.""" + + @abstractmethod + def _pool(self, model_output, attention_mask): + """Aggregate token-level hidden states into a single embedding per item. + + Args: + model_output: Output tuple from ``AutoModel``. + attention_mask: Attention mask tensor with shape ``(batch, seq_len)``. + + Returns: + A torch tensor of shape ``(batch, embedding_dim)``. + """ + raise NotImplementedError + + def load(self): + """Download the tokenizer and model from HuggingFace Hub and move to device.""" + self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) + self.model = AutoModel.from_pretrained(self.model_name).to(self.device) + + def _batch_encode_impl(self, texts: List[str]) -> np.ndarray: + """Tokenise, forward, pool and return a NumPy array of embeddings. + + Args: + texts: Batch of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + encoded = self.tokenizer( + texts, + padding=True, + truncation=True, + return_tensors="pt", + ) + encoded = {k: v.to(self.device) for k, v in encoded.items()} + with torch.no_grad(): + outputs = self.model(**encoded) + embeddings = self._pool(outputs, encoded["attention_mask"]) + return embeddings.cpu().numpy() + + def batch_encode(self, texts: List[str]) -> np.ndarray: + """Encode a batch of texts into a dense NumPy array. + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + return self._batch_encode_impl(texts) + + def encode(self, text: str) -> np.ndarray: + """Encode a single text into a 1-D embedding vector. + + Args: + text: Input string. + + Returns: + A 1-D float32 NumPy array of shape ``(embedding_dim,)``. + """ + embeddings = self.batch_encode([text]) + return embeddings.squeeze() diff --git a/DashAI/back/models/RAG/embeddings/dense/instructor_embedding.py b/DashAI/back/models/RAG/embeddings/dense/instructor_embedding.py new file mode 100644 index 000000000..819517d05 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/instructor_embedding.py @@ -0,0 +1,140 @@ +from typing import Dict, List + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.embeddings.dense._instructor_embedding import ( + _InstructorEmbedding, +) +from DashAI.back.models.RAG.embeddings.dense_embedding import DenseEmbedding + +INSTRUCTOR_MODELS: Dict[str, dict] = { + "hkunlp/instructor-base": { + "languages": ["en"], + }, + "hkunlp/instructor-large": { + "languages": ["en"], + }, + "hkunlp/instructor-xl": { + "languages": ["en"], + }, +} + +INSTRUCTOR_MODEL_NAMES = list(INSTRUCTOR_MODELS.keys()) + + +class InstructorEmbeddingSchema(BaseSchema): + """Configuration schema for :class:`InstructorEmbedding`. + + Attributes: + model_name: INSTRUCTOR model for instruction-tuned embedding generation. + instruction: Instruction text that guides the embedding model. + device: Device to run the model on. + """ + + model_name: schema_field( + enum_field(INSTRUCTOR_MODEL_NAMES), + placeholder="hkunlp/instructor-base", + description=MultilingualString( + en="INSTRUCTOR model for instruction-tuned embedding generation.", + es="Modelo INSTRUCTOR para generación de embeddings ajustados" + " por instrucción.", + ), + ) # type: ignore + + instruction: schema_field( + string_field(), + placeholder="Represent the document for retrieval:", + description=MultilingualString( + en="Instruction text that guides the embedding model.", + es="Texto de instrucción que guía al modelo de embedding.", + ), + ) # type: ignore + + device: schema_field( + enum_field(["cpu", "cuda"]), + placeholder="cpu", + description=MultilingualString( + en="Device to run the model on.", + es="Dispositivo para ejecutar el modelo.", + ), + ) # type: ignore + + +class InstructorEmbedding(DenseEmbedding): + """Dense embeddings using INSTRUCTOR instruction-tuned models. + + Wraps :class:`_InstructorEmbedding` and exposes it as a DashAI component + with a configurable schema (:class:`InstructorEmbeddingSchema`). + + Prepends a user-defined instruction string to every input text. + + FLAGS: + FAMILY:instructor: Groups this model under the INSTRUCTOR family. + huggingface: Marks the model family as HuggingFace-based. + """ + + SCHEMA = InstructorEmbeddingSchema + FLAGS: list[str] = ["FAMILY:instructor", "huggingface"] + DISPLAY_NAME: str = MultilingualString( + en="INSTRUCTOR Embedding", + es="Embedding INSTRUCTOR", + ) + DESCRIPTION: str = MultilingualString( + en="Dense embeddings using INSTRUCTOR instruction-tuned models.", + es="Embeddings densos usando modelos INSTRUCTOR ajustados por instrucción.", + ) + + def __init__(self, **kwargs): + """Initialise the embedding by validating parameters and creating the internal + model. + + Args: + **kwargs: Configuration matching :class:`InstructorEmbeddingSchema`. + """ + self.params = self.validate_and_transform(kwargs) + model_name = self.params["model_name"] + device = self.params["device"] + instruction = self.params["instruction"] + self._embedding = _InstructorEmbedding( + model_name=model_name, + device=device, + instruction=instruction, + ) + + def load(self): + """Load the INSTRUCTOR model.""" + self._embedding.load() + + def encode(self, text: str): + """Encode a single text into a dense embedding with the configured instruction. + + Args: + text: Input string. + + Returns: + A 1-D NumPy array of shape ``(embedding_dim,)``. + """ + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + """Encode a batch of texts into dense embeddings with the configured + instruction. + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + return self._embedding.batch_encode(texts) + + def save(self): + """No-op. Persistence is handled externally.""" + + def train(self, **kwargs): + """No-op. Pre-trained models are used as-is.""" diff --git a/DashAI/back/models/RAG/embeddings/dense/labse_embedding.py b/DashAI/back/models/RAG/embeddings/dense/labse_embedding.py new file mode 100644 index 000000000..12cc9cd96 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/labse_embedding.py @@ -0,0 +1,187 @@ +from typing import Dict, List + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + AGGREGATE, + TRUNCATE, +) +from DashAI.back.models.RAG.embeddings.dense._sentence_transformer_embedding import ( + _SentenceTransformerEmbedding, +) +from DashAI.back.models.RAG.embeddings.dense_embedding import DenseEmbedding + +LABSE_MODELS: Dict[str, dict] = { + "sentence-transformers/LaBSE": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "pt", + "nl", + "pl", + "ca", + "fi", + "ar", + "zh", + "ja", + "ko", + "ru", + "tr", + "hi", + "sv", + "da", + "no", + "cs", + "ro", + "el", + "he", + "hu", + "th", + "vi", + "id", + "ms", + "bg", + "hr", + "sk", + "sl", + "sr", + "uk", + "et", + "lv", + "lt", + "fa", + "ur", + "mk", + "af", + "bn", + "gu", + "ka", + "ku", + "my", + "sq", + "multi", + ], + "max_seq_length": 512, + }, +} + +LABSE_MODEL_NAMES = list(LABSE_MODELS.keys()) + + +class LaBSEmbeddingSchema(BaseSchema): + """Configuration schema for :class:`LaBSEmbedding`. + + Attributes: + model_name: LaBSE model for multilingual embedding generation (109 languages). + overflow_strategy: Strategy for chunks exceeding model max sequence length. + device: Device to run the model on. + """ + + model_name: schema_field( + enum_field(LABSE_MODEL_NAMES), + placeholder="sentence-transformers/LaBSE", + description=MultilingualString( + en="LaBSE model for multilingual embedding generation (109 languages).", + es="Modelo LaBSE para generación de embeddings multilingües (109 idiomas).", + ), + ) # type: ignore + + overflow_strategy: schema_field( + enum_field([TRUNCATE, AGGREGATE]), + placeholder=TRUNCATE, + description=MultilingualString( + en="Strategy for chunks exceeding model max sequence length.", + es="Estrategia para fragmentos que exceden la longitud máxima del modelo.", + ), + ) # type: ignore + + device: schema_field( + enum_field(["cpu", "cuda"]), + placeholder="cpu", + description=MultilingualString( + en="Device to run the model on.", + es="Dispositivo para ejecutar el modelo.", + ), + ) # type: ignore + + +class LaBSEmbedding(DenseEmbedding): + """Dense embeddings using the LaBSE multilingual model (109 languages). + + Wraps :class:`_SentenceTransformerEmbedding` (mean pooling, L2 + normalisation forced on) and exposes it as a DashAI component with + a configurable schema (:class:`LaBSEmbeddingSchema`). + + FLAGS: + FAMILY:labse: Groups this model under the LaBSE family. + huggingface: Marks the model family as HuggingFace-based. + """ + + SCHEMA = LaBSEmbeddingSchema + FLAGS: list[str] = ["FAMILY:labse", "huggingface"] + DISPLAY_NAME: str = MultilingualString( + en="LaBSE Embedding", + es="Embedding LaBSE", + ) + DESCRIPTION: str = MultilingualString( + en="Dense embeddings using LaBSE multilingual model (109 languages).", + es="Embeddings densos usando el modelo multilingüe LaBSE (109 idiomas).", + ) + + def __init__(self, **kwargs): + """Initialise the embedding by validating parameters and creating the internal model. + + Args: + **kwargs: Configuration matching :class:`LaBSEmbeddingSchema`. + """ # noqa: E501 + self.params = self.validate_and_transform(kwargs) + model_name = self.params["model_name"] + device = self.params["device"] + overflow_strategy = self.params.get("overflow_strategy", "truncate") + model_info = LABSE_MODELS[model_name] + self._embedding = _SentenceTransformerEmbedding( + model_name=model_name, + device=device, + model_max_length=model_info["max_seq_length"], + overflow_strategy=overflow_strategy, + normalize=True, + ) + + def load(self): + """Load the LaBSE model and tokenizer.""" + self._embedding.load() + + def encode(self, text: str): + """Encode a single text into a dense embedding. + + Args: + text: Input string. + + Returns: + A 1-D NumPy array of shape ``(embedding_dim,)``. + """ + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + """Encode a batch of texts into dense embeddings. + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + return self._embedding.batch_encode(texts) + + def save(self): + """No-op. Persistence is handled externally.""" + + def train(self, **kwargs): + """No-op. Pre-trained models are used as-is.""" diff --git a/DashAI/back/models/RAG/embeddings/dense/openai_embedding.py b/DashAI/back/models/RAG/embeddings/dense/openai_embedding.py new file mode 100644 index 000000000..42b70fd62 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/openai_embedding.py @@ -0,0 +1,122 @@ +from typing import List + +import numpy as np +from openai import OpenAI + +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.core.schema_fields.enum_field import enum_field +from DashAI.back.core.schema_fields.schema_field import schema_field +from DashAI.back.core.schema_fields.string_field import string_field +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.embeddings.dense_embedding import DenseEmbedding + + +class OpenAIEmbeddingSchema(BaseSchema): + """Configuration schema for :class:`OpenAIEmbedding`. + + Attributes: + model_name: OpenAI embedding model to use. + api_key: OpenAI API key. + """ + + model_name: schema_field( + enum_field( + [ + "text-embedding-ada-002", + "text-embedding-3-small", + "text-embedding-3-large", + ] + ), + placeholder="text-embedding-3-small", + description=MultilingualString( + en="OpenAI embedding model to use.", + es="Modelo de embedding de OpenAI a utilizar.", + ), + ) # type: ignore + + api_key: schema_field( + string_field(), + placeholder="", + description=MultilingualString( + en="OpenAI API key.", + es="Clave API de OpenAI.", + ), + ) # type: ignore + + +class OpenAIEmbedding(DenseEmbedding): + """Dense embeddings via the OpenAI Embeddings API. + + Supports ``text-embedding-ada-002``, ``text-embedding-3-small``, and + ``text-embedding-3-large`` models. + + FLAGS: + FAMILY:openai: Groups this model under the OpenAI family. + remote: Marks this model as calling a remote API. + """ + + FLAGS: list[str] = ["FAMILY:openai", "remote"] + DISPLAY_NAME = MultilingualString( + en="OpenAI Embedding", + es="Embedding OpenAI", + ) + DESCRIPTION = MultilingualString( + en="OpenAI text embeddings", + es="Embeddings de texto de OpenAI", + ) + SCHEMA = OpenAIEmbeddingSchema + + def __init__(self, **kwargs): + """Initialise the embedding by validating parameters and creating the OpenAI client. + + Args: + **kwargs: Configuration matching :class:`OpenAIEmbeddingSchema`. + + Raises: + ValueError: If the API key is missing or empty. + """ # noqa: E501 + self.params = self.validate_and_transform(kwargs) + self.model_name = self.params["model_name"] + self.api_key = self.params["api_key"] + if not self.api_key or not self.api_key.strip(): + raise ValueError("OpenAI API key is required but was not provided") + self.client = OpenAI(api_key=self.api_key) + + def load(self): + """No-op. The OpenAI client does not require model loading.""" + + def save(self): + """No-op. No local state to persist.""" + + def train(self, **kwargs): + """No-op. OpenAI models cannot be fine-tuned through this interface.""" + + def encode(self, text: str) -> np.ndarray: + """Encode a single text via the OpenAI API. + + Args: + text: Input string. + + Returns: + A 1-D float32 NumPy array of shape ``(embedding_dim,)``. + """ + response = self.client.embeddings.create( + model=self.model_name, + input=text, + ) + return np.array(response.data[0].embedding) + + def batch_encode(self, texts: List[str]) -> np.ndarray: + """Encode a batch of texts via the OpenAI API. + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + response = self.client.embeddings.create( + model=self.model_name, + input=texts, + ) + return np.array([d.embedding for d in response.data]) diff --git a/DashAI/back/models/RAG/embeddings/dense/roberta_embedding.py b/DashAI/back/models/RAG/embeddings/dense/roberta_embedding.py new file mode 100644 index 000000000..ba66bf605 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/roberta_embedding.py @@ -0,0 +1,261 @@ +from typing import Dict, List + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.embeddings.dense._bert_embedding import ( + MAX, + MEAN, + _BERTEmbedding, +) +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + AGGREGATE, + TRUNCATE, +) +from DashAI.back.models.RAG.embeddings.dense_embedding import DenseEmbedding + +ROBERTA_POOLING_STRATEGIES = [MEAN, MAX] + +ROBERTA_MODELS: Dict[str, dict] = { + "FacebookAI/roberta-base": { + "languages": ["en"], + "max_seq_length": 512, + }, + "FacebookAI/roberta-large": { + "languages": ["en"], + "max_seq_length": 512, + }, + "FacebookAI/xlm-roberta-base": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "pt", + "nl", + "pl", + "ca", + "fi", + "ar", + "zh", + "ja", + "ko", + "ru", + "tr", + "hi", + "sv", + "da", + "no", + "cs", + "ro", + "el", + "he", + "hu", + "th", + "vi", + "id", + "ms", + "bg", + "hr", + "sk", + "sl", + "sr", + "uk", + "et", + "lv", + "lt", + "fa", + "ur", + "mk", + "af", + "bn", + "multi", + ], + "max_seq_length": 512, + }, + "FacebookAI/xlm-roberta-large": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "pt", + "nl", + "pl", + "ca", + "fi", + "ar", + "zh", + "ja", + "ko", + "ru", + "tr", + "hi", + "sv", + "da", + "no", + "cs", + "ro", + "el", + "he", + "hu", + "th", + "vi", + "id", + "ms", + "bg", + "hr", + "sk", + "sl", + "sr", + "uk", + "et", + "lv", + "lt", + "fa", + "ur", + "mk", + "af", + "bn", + "multi", + ], + "max_seq_length": 512, + }, +} + +ROBERTA_MODEL_NAMES = list(ROBERTA_MODELS.keys()) + + +class RoBERTaEmbeddingSchema(BaseSchema): + """Configuration schema for :class:`RoBERTaEmbedding`. + + Attributes: + model_name: RoBERTa / XLM-RoBERTa model for embedding generation. + overflow_strategy: Strategy for chunks exceeding model max sequence length. + device: Device to run the model on. + pooling_strategy: Pooling strategy (mean or max; CLS not + recommended for RoBERTa). + """ + + model_name: schema_field( + enum_field(ROBERTA_MODEL_NAMES), + placeholder="FacebookAI/roberta-base", + description=MultilingualString( + en="RoBERTa / XLM-RoBERTa model for embedding generation.", + es="Modelo RoBERTa / XLM-RoBERTa para generación de embeddings.", + ), + ) # type: ignore + + overflow_strategy: schema_field( + enum_field([TRUNCATE, AGGREGATE]), + placeholder=TRUNCATE, + description=MultilingualString( + en="Strategy for chunks exceeding model max sequence length.", + es="Estrategia para fragmentos que exceden la longitud máxima del modelo.", + ), + ) # type: ignore + + device: schema_field( + enum_field(["cpu", "cuda"]), + placeholder="cpu", + description=MultilingualString( + en="Device to run the model on.", + es="Dispositivo para ejecutar el modelo.", + ), + ) # type: ignore + + pooling_strategy: schema_field( + enum_field(ROBERTA_POOLING_STRATEGIES), + placeholder=MEAN, + description=MultilingualString( + en="Pooling strategy to aggregate token embeddings. RoBERTa" + " CLS token is not trained for similarity.", + es="Estrategia de pooling para agregar embeddings de tokens." + " El token CLS de RoBERTa no está entrenado para similitud.", + ), + ) # type: ignore + + +class RoBERTaEmbedding(DenseEmbedding): + """Dense embeddings using RoBERTa / XLM-RoBERTa models with mean/max pooling. + + Wraps :class:`_BERTEmbedding` (reusing BERT pooling logic) and exposes + it as a DashAI component with a configurable schema + (:class:`RoBERTaEmbeddingSchema`). + + Only mean and max pooling are exposed because the RoBERTa CLS token is + not trained for similarity tasks. + + FLAGS: + FAMILY:roberta: Groups this model under the RoBERTa family. + huggingface: Marks the model family as HuggingFace-based. + """ + + SCHEMA = RoBERTaEmbeddingSchema + FLAGS: list[str] = ["FAMILY:roberta", "huggingface"] + DISPLAY_NAME: str = MultilingualString( + en="RoBERTa Embedding", + es="Embedding RoBERTa", + ) + DESCRIPTION: str = MultilingualString( + en="Dense embeddings using RoBERTa / XLM-RoBERTa models with mean/max pooling.", + es="Embeddings densos usando modelos RoBERTa / XLM-RoBERTa con" + " pooling mean/max.", + ) + + def __init__(self, **kwargs): + """Initialise the embedding by validating parameters and creating the internal model. + + Args: + **kwargs: Configuration matching :class:`RoBERTaEmbeddingSchema`. + """ # noqa: E501 + self.params = self.validate_and_transform(kwargs) + model_name = self.params["model_name"] + device = self.params["device"] + overflow_strategy = self.params.get("overflow_strategy", "truncate") + pooling_strategy = self.params["pooling_strategy"] + model_info = ROBERTA_MODELS[model_name] + self._embedding = _BERTEmbedding( + model_name=model_name, + device=device, + model_max_length=model_info["max_seq_length"], + overflow_strategy=overflow_strategy, + pooling_strategy=pooling_strategy, + ) + + def load(self): + """Load the RoBERTa model and tokenizer.""" + self._embedding.load() + + def encode(self, text: str): + """Encode a single text into a dense embedding. + + Args: + text: Input string. + + Returns: + A 1-D NumPy array of shape ``(embedding_dim,)``. + """ + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + """Encode a batch of texts into dense embeddings. + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + return self._embedding.batch_encode(texts) + + def save(self): + """No-op. Persistence is handled externally.""" + + def train(self, **kwargs): + """No-op. Pre-trained models are used as-is.""" diff --git a/DashAI/back/models/RAG/embeddings/dense/sentence_transformer_embedding.py b/DashAI/back/models/RAG/embeddings/dense/sentence_transformer_embedding.py new file mode 100644 index 000000000..183cfbaa9 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/sentence_transformer_embedding.py @@ -0,0 +1,446 @@ +from typing import Dict, List + +from DashAI.back.core.schema_fields import ( + BaseSchema, + bool_field, + enum_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + AGGREGATE, + TRUNCATE, +) +from DashAI.back.models.RAG.embeddings.dense._sentence_transformer_embedding import ( + _SentenceTransformerEmbedding, +) +from DashAI.back.models.RAG.embeddings.dense_embedding import DenseEmbedding + +ST_MODELS: Dict[str, dict] = { + "microsoft/harrier-oss-v1-270m": { + "languages": ["multi"], + "max_seq_length": 32768, + "pooling": "last_token", + }, + "microsoft/harrier-oss-v1-0.6b": { + "languages": ["multi"], + "max_seq_length": 32768, + "pooling": "last_token", + }, + "microsoft/harrier-oss-v1-27b": { + "languages": ["multi"], + "max_seq_length": 32768, + "pooling": "last_token", + }, + "Qwen/Qwen3-Embedding-0.6B": { + "languages": ["multi"], + "max_seq_length": 32768, + "pooling": "mean", + }, + "Qwen/Qwen3-Embedding-4B": { + "languages": ["multi"], + "max_seq_length": 32768, + "pooling": "mean", + }, + "Qwen/Qwen3-Embedding-8B": { + "languages": ["multi"], + "max_seq_length": 32768, + "pooling": "mean", + }, + "sentence-transformers/all-MiniLM-L6-v2": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/all-MiniLM-L12-v2": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/all-mpnet-base-v2": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/all-distilroberta-v1": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/multi-qa-mpnet-base-dot-v1": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + "normalize_default": False, + }, + "sentence-transformers/multi-qa-mpnet-base-cos-v1": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/multi-qa-distilbert-dot-v1": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + "normalize_default": False, + }, + "sentence-transformers/multi-qa-distilbert-cos-v1": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/multi-qa-MiniLM-L6-dot-v1": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + "normalize_default": False, + }, + "sentence-transformers/multi-qa-MiniLM-L6-cos-v1": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/msmarco-bert-base-dot-v5": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + "normalize_default": False, + }, + "sentence-transformers/msmarco-distilbert-dot-v5": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + "normalize_default": False, + }, + "sentence-transformers/msmarco-distilbert-base-tas-b": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/msmarco-distilbert-cos-v5": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/msmarco-MiniLM-L12-cos-v5": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/msmarco-MiniLM-L6-cos-v5": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "pt", + "nl", + "pl", + "ca", + "fi", + "ar", + "zh", + "ja", + "ko", + "ru", + "tr", + "hi", + "sv", + "da", + "no", + "cs", + "ro", + "el", + "he", + "hu", + "th", + "vi", + "id", + "ms", + "bg", + "hr", + "sk", + "sl", + "sr", + "uk", + "et", + "lv", + "lt", + "fa", + "ur", + "mk", + "af", + "bn", + "multi", + ], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/paraphrase-multilingual-mpnet-base-v2": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "pt", + "nl", + "pl", + "ca", + "fi", + "ar", + "zh", + "ja", + "ko", + "ru", + "tr", + "hi", + "sv", + "da", + "no", + "cs", + "ro", + "el", + "he", + "hu", + "th", + "vi", + "id", + "ms", + "bg", + "hr", + "sk", + "sl", + "sr", + "uk", + "et", + "lv", + "lt", + "fa", + "ur", + "mk", + "af", + "bn", + "multi", + ], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/distiluse-base-multilingual-cased-v2": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "pt", + "nl", + "pl", + "ca", + "fi", + "ar", + "zh", + "ja", + "ko", + "ru", + "tr", + "hi", + "sv", + "da", + "no", + "cs", + "ro", + "el", + "he", + "hu", + "th", + "vi", + "id", + "ms", + "bg", + "hr", + "sk", + "sl", + "sr", + "uk", + "et", + "lv", + "lt", + "fa", + "ur", + "mk", + "af", + "bn", + "multi", + ], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/distiluse-base-multilingual-cased-v1": { + "languages": [ + "en", + "es", + "fr", + "de", + "it", + "nl", + "pt", + "ar", + "zh", + "ja", + "ko", + "pl", + "ru", + "tr", + "multi", + ], + "max_seq_length": 512, + "pooling": "mean", + }, + "sentence-transformers/allenai-specter": { + "languages": ["en"], + "max_seq_length": 512, + "pooling": "mean", + }, +} + +ST_MODEL_NAMES = list(ST_MODELS.keys()) + + +class SentenceTransformerEmbeddingSchema(BaseSchema): + """Configuration schema for :class:`SentenceTransformerEmbedding`. + + Attributes: + model_name: Sentence Transformer model for embedding generation. + overflow_strategy: Strategy for chunks exceeding model max sequence length. + normalize: Whether to L2-normalize the output embeddings. + device: Device to run the model on. + """ + + model_name: schema_field( + enum_field(ST_MODEL_NAMES), + placeholder="microsoft/harrier-oss-v1-0.6b", + description=MultilingualString( + en="Sentence Transformer model for embedding generation.", + es="Modelo Sentence Transformer para generación de embeddings.", + ), + ) # type: ignore + + overflow_strategy: schema_field( + enum_field([TRUNCATE, AGGREGATE]), + placeholder=TRUNCATE, + description=MultilingualString( + en="Strategy for chunks exceeding model max sequence length.", + es="Estrategia para fragmentos que exceden la longitud máxima del modelo.", + ), + ) # type: ignore + + normalize: schema_field( + bool_field(), + placeholder=True, + description=MultilingualString( + en="Whether to L2-normalize the output embeddings.", + es="Si normalizar con L2 los embeddings de salida.", + ), + ) # type: ignore + + device: schema_field( + enum_field(["cpu", "cuda"]), + placeholder="cpu", + description=MultilingualString( + en="Device to run the model on.", + es="Dispositivo para ejecutar el modelo.", + ), + ) # type: ignore + + +class SentenceTransformerEmbedding(DenseEmbedding): + """Dense embeddings using Sentence Transformer models. + + Wraps :class:`_SentenceTransformerEmbedding` and exposes it as a + DashAI component with a configurable schema + (:class:`SentenceTransformerEmbeddingSchema`). + + Supports mean / last-token pooling, L2 normalisation, and overflow + strategies (truncate / aggregate). + + FLAGS: + FAMILY:sentence_transformer: Groups this model under the + sentence-transformer family. + huggingface: Marks the model family as HuggingFace-based. + """ + + SCHEMA = SentenceTransformerEmbeddingSchema + FLAGS: list[str] = ["FAMILY:sentence_transformer", "huggingface"] + DISPLAY_NAME: str = MultilingualString( + en="Sentence Transformer Embedding", + es="Embedding Sentence Transformer", + ) + DESCRIPTION: str = MultilingualString( + en="Dense embeddings using Sentence Transformer models with mean" + " pooling and L2 normalization.", + es="Embeddings densos usando modelos Sentence Transformer con mean" + " pooling y normalización L2.", + ) + + def __init__(self, **kwargs): + """Initialise the embedding by validating parameters and creating the internal model. + + Args: + **kwargs: Configuration matching :class:`SentenceTransformerEmbeddingSchema`. # noqa: E501 + """ # noqa: E501 + self.params = self.validate_and_transform(kwargs) + model_name = self.params["model_name"] + device = self.params["device"] + normalize = self.params["normalize"] + overflow_strategy = self.params.get("overflow_strategy", "truncate") + model_info = ST_MODELS[model_name] + pooling = model_info.get("pooling", "mean") + self._embedding = _SentenceTransformerEmbedding( + model_name=model_name, + device=device, + model_max_length=model_info["max_seq_length"], + overflow_strategy=overflow_strategy, + normalize=normalize + if model_info.get("normalize_default", True) + else normalize, + pooling=pooling, + ) + + def load(self): + """Load the Sentence Transformer model and tokenizer.""" + self._embedding.load() + + def encode(self, text: str): + """Encode a single text into a dense embedding. + + Args: + text: Input string. + + Returns: + A 1-D NumPy array of shape ``(embedding_dim,)``. + """ + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + """Encode a batch of texts into dense embeddings. + + Args: + texts: List of input strings. + + Returns: + A ``(batch, embedding_dim)`` float32 NumPy array. + """ + return self._embedding.batch_encode(texts) + + def save(self): + """No-op. Persistence is handled externally.""" + + def train(self, **kwargs): + """No-op. Pre-trained models are used as-is.""" diff --git a/DashAI/back/models/RAG/embeddings/dense_embedding.py b/DashAI/back/models/RAG/embeddings/dense_embedding.py new file mode 100644 index 000000000..f7da3bb82 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense_embedding.py @@ -0,0 +1,61 @@ +from abc import abstractmethod +from typing import Any, List + +import numpy as np + +from DashAI.back.models.base_model import BaseModel + + +class DenseEmbedding(BaseModel): + """Base class for all dense encoding (embedding) models. + + Subclasses must override :meth:`encode`, :meth:`batch_encode`, and + :meth:`train`, and set :attr:`embedding_dim` during initialisation. + """ + + embedding_dim: int + + def __init__(self, **kwargs: Any): + """Initialise the embedding model with arbitrary keyword parameters. + + Args: + **kwargs: Model configuration parameters forwarded to subclasses. + """ + self.params = kwargs + + def encode(self, text: str) -> List[float] | np.ndarray: + """Generate an embedding vector for a single text string. + + Args: + text: The input text to encode. + + Returns: + A 1-D vector (list or ndarray) representing the embedding. + + Raises: + NotImplementedError: Subclasses must implement this method. + """ + raise NotImplementedError("Subclasses must implement this method.") + + def batch_encode(self, texts: List[str]) -> List[List[float]] | np.ndarray: + """Generate embedding vectors for a batch of texts. + + Args: + texts: A list of input texts to encode. + + Returns: + A 2-D array (list of lists or ndarray) of embeddings, one per input. + + Raises: + NotImplementedError: Subclasses must implement this method. + """ + raise NotImplementedError("Subclasses must implement this method.") + + @abstractmethod + def train(self, **kwargs): + """Train the embedding model on the provided data. + + Args: + **kwargs: Training configuration (data, hyperparameters, etc.). + """ + return diff --git a/DashAI/back/models/RAG/embeddings/sparse/__init__.py b/DashAI/back/models/RAG/embeddings/sparse/__init__.py new file mode 100644 index 000000000..37452c250 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/sparse/__init__.py @@ -0,0 +1,5 @@ +"""Sparse embedding models for the RAG module. + +Currently a placeholder package for future sparse embedding implementations +(e.g. SPLADE, BM25-based embeddings). +""" diff --git a/DashAI/back/models/RAG/exceptions/__init__.py b/DashAI/back/models/RAG/exceptions/__init__.py new file mode 100644 index 000000000..500cfcca0 --- /dev/null +++ b/DashAI/back/models/RAG/exceptions/__init__.py @@ -0,0 +1,138 @@ +"""Unified RAG exception hierarchy. + +All RAG-related exceptions live here. Every layer — models, services, +factories, API endpoints, jobs — must import from this package when +raising or catching RAG exceptions so that callers can rely on a single +source of truth. + +Exception tree (simplified):: + + RAGWorkflowError + ├── RAGDocumentError + │ ├── RAGDocumentParsingError + │ ├── RAGDocumentNotFoundError + │ └── RAGDocumentFileTypeError + ├── RAGChunkingError + │ └── RAGChunkingOverlapError + ├── RAGRetrieverError + │ ├── RAGRetrieverMissingParameterError + │ ├── RAGRetrieverCompositeValidationError + │ └── RAGRetrieverEmptyChildrenError + ├── RAGPromptError + │ ├── RAGPromptValidationError + │ └── RAGPromptTemplateError + ├── RAGGenerationError + │ └── RAGGenerationModelError + ├── RAGPipelineError + │ ├── RAGPipelineConfigError + │ ├── RAGPipelineInitializationError + │ ├── RAGPipelineRuntimeError + │ ├── RAGPipelineInputError + │ └── RAGDatabaseError + ├── RAGEmbeddingError + │ ├── RAGEmbeddingLoadError + │ └── RAGEmbeddingEmptyInputError + ├── RAGFactoryError + │ └── RAGComponentNotFoundError + └── RAGTaskError + └── RAGTaskInputError +""" + +from DashAI.back.models.RAG.exceptions.base import RAGWorkflowError + +from DashAI.back.models.RAG.exceptions.document import ( + RAGDocumentError, + RAGDocumentFileTypeError, + RAGDocumentNotFoundError, + RAGDocumentParsingError, +) + +from DashAI.back.models.RAG.exceptions.chunking import ( + RAGChunkingError, + RAGChunkingOverlapError, +) + +from DashAI.back.models.RAG.exceptions.retriever import ( + RAGRetrieverCompositeValidationError, + RAGRetrieverEmptyChildrenError, + RAGRetrieverError, + RAGRetrieverMissingParameterError, +) + +from DashAI.back.models.RAG.exceptions.prompt import ( + RAGPromptError, + RAGPromptTemplateError, + RAGPromptValidationError, +) + +from DashAI.back.models.RAG.exceptions.generation import ( + RAGGenerationError, + RAGGenerationModelError, +) + +from DashAI.back.models.RAG.exceptions.pipeline import ( + RAGDatabaseError, + RAGPipelineConfigError, + RAGPipelineError, + RAGPipelineInitializationError, + RAGPipelineInputError, + RAGPipelineRuntimeError, +) + +from DashAI.back.models.RAG.exceptions.embedding import ( + RAGEmbeddingEmptyInputError, + RAGEmbeddingError, + RAGEmbeddingLoadError, +) + +from DashAI.back.models.RAG.exceptions.factory import ( + RAGComponentNotFoundError, + RAGFactoryError, +) + +from DashAI.back.models.RAG.exceptions.task import ( + RAGTaskError, + RAGTaskInputError, +) + +__all__ = [ + # Base + "RAGWorkflowError", + # Document + "RAGDocumentError", + "RAGDocumentParsingError", + "RAGDocumentNotFoundError", + "RAGDocumentFileTypeError", + # Chunking + "RAGChunkingError", + "RAGChunkingOverlapError", + # Retriever + "RAGRetrieverError", + "RAGRetrieverMissingParameterError", + "RAGRetrieverCompositeValidationError", + "RAGRetrieverEmptyChildrenError", + # Prompt + "RAGPromptError", + "RAGPromptValidationError", + "RAGPromptTemplateError", + # Generation + "RAGGenerationError", + "RAGGenerationModelError", + # Pipeline + "RAGPipelineError", + "RAGPipelineConfigError", + "RAGPipelineInitializationError", + "RAGPipelineRuntimeError", + "RAGPipelineInputError", + "RAGDatabaseError", + # Embedding + "RAGEmbeddingError", + "RAGEmbeddingLoadError", + "RAGEmbeddingEmptyInputError", + # Factory + "RAGFactoryError", + "RAGComponentNotFoundError", + # Task + "RAGTaskError", + "RAGTaskInputError", +] diff --git a/DashAI/back/models/RAG/exceptions/base.py b/DashAI/back/models/RAG/exceptions/base.py new file mode 100644 index 000000000..6576d70fd --- /dev/null +++ b/DashAI/back/models/RAG/exceptions/base.py @@ -0,0 +1,5 @@ +"""Root RAG exception class.""" + + +class RAGWorkflowError(Exception): + """Base exception for all RAG workflow errors.""" diff --git a/DashAI/back/models/RAG/exceptions/chunking.py b/DashAI/back/models/RAG/exceptions/chunking.py new file mode 100644 index 000000000..dcf85d56b --- /dev/null +++ b/DashAI/back/models/RAG/exceptions/chunking.py @@ -0,0 +1,11 @@ +"""Chunking-related RAG exceptions.""" + +from DashAI.back.models.RAG.exceptions.base import RAGWorkflowError + + +class RAGChunkingError(RAGWorkflowError): + """Error during document chunking.""" + + +class RAGChunkingOverlapError(RAGChunkingError): + """chunk_overlap is >= chunk_size, which would cause an infinite loop.""" diff --git a/DashAI/back/models/RAG/exceptions/document.py b/DashAI/back/models/RAG/exceptions/document.py new file mode 100644 index 000000000..1659c401b --- /dev/null +++ b/DashAI/back/models/RAG/exceptions/document.py @@ -0,0 +1,19 @@ +"""Document-related RAG exceptions.""" + +from DashAI.back.models.RAG.exceptions.base import RAGWorkflowError + + +class RAGDocumentError(RAGWorkflowError): + """Error in document loading or processing.""" + + +class RAGDocumentParsingError(RAGDocumentError): + """Error parsing a document file (e.g. PDF extraction failure).""" + + +class RAGDocumentNotFoundError(RAGDocumentError): + """Referenced document does not exist.""" + + +class RAGDocumentFileTypeError(RAGDocumentError): + """Unsupported or unrecognized document file type.""" diff --git a/DashAI/back/models/RAG/exceptions/embedding.py b/DashAI/back/models/RAG/exceptions/embedding.py new file mode 100644 index 000000000..833f7af7a --- /dev/null +++ b/DashAI/back/models/RAG/exceptions/embedding.py @@ -0,0 +1,15 @@ +"""Embedding-related RAG exceptions.""" + +from DashAI.back.models.RAG.exceptions.base import RAGWorkflowError + + +class RAGEmbeddingError(RAGWorkflowError): + """Error during embedding computation or loading.""" + + +class RAGEmbeddingLoadError(RAGEmbeddingError): + """Failed to load embeddings from disk.""" + + +class RAGEmbeddingEmptyInputError(RAGEmbeddingError): + """Empty input text provided for embedding.""" diff --git a/DashAI/back/models/RAG/exceptions/factory.py b/DashAI/back/models/RAG/exceptions/factory.py new file mode 100644 index 000000000..35cc8b0cd --- /dev/null +++ b/DashAI/back/models/RAG/exceptions/factory.py @@ -0,0 +1,11 @@ +"""Factory-related RAG exceptions.""" + +from DashAI.back.models.RAG.exceptions.base import RAGWorkflowError + + +class RAGFactoryError(RAGWorkflowError): + """Error in a RAG factory (component creation).""" + + +class RAGComponentNotFoundError(RAGFactoryError): + """The requested component is not registered in the component registry.""" diff --git a/DashAI/back/models/RAG/exceptions/generation.py b/DashAI/back/models/RAG/exceptions/generation.py new file mode 100644 index 000000000..312db6d83 --- /dev/null +++ b/DashAI/back/models/RAG/exceptions/generation.py @@ -0,0 +1,11 @@ +"""Generation-related RAG exceptions.""" + +from DashAI.back.models.RAG.exceptions.base import RAGWorkflowError + + +class RAGGenerationError(RAGWorkflowError): + """Error during LLM generation.""" + + +class RAGGenerationModelError(RAGGenerationError): + """The generation model is not a valid TextToTextGenerationTaskModel.""" diff --git a/DashAI/back/models/RAG/exceptions/pipeline.py b/DashAI/back/models/RAG/exceptions/pipeline.py new file mode 100644 index 000000000..0c7f74fba --- /dev/null +++ b/DashAI/back/models/RAG/exceptions/pipeline.py @@ -0,0 +1,31 @@ +"""Pipeline orchestration RAG exceptions. + +These were historically defined inside RAG_pipeline.py and are moved +here so the whole hierarchy lives in one package. +""" + +from DashAI.back.models.RAG.exceptions.base import RAGWorkflowError + + +class RAGPipelineError(RAGWorkflowError): + """Base exception for RAG pipeline errors.""" + + +class RAGPipelineConfigError(RAGPipelineError): + """Invalid or missing parameters in pipeline configuration.""" + + +class RAGPipelineInitializationError(RAGPipelineError): + """Error during RAG pipeline initialization.""" + + +class RAGPipelineRuntimeError(RAGPipelineError): + """Error during RAG pipeline execution.""" + + +class RAGDatabaseError(RAGPipelineError): + """Database-related error in RAG pipeline.""" + + +class RAGPipelineInputError(RAGPipelineError): + """Invalid or malformed input data to the RAG pipeline.""" diff --git a/DashAI/back/models/RAG/exceptions/prompt.py b/DashAI/back/models/RAG/exceptions/prompt.py new file mode 100644 index 000000000..75503d406 --- /dev/null +++ b/DashAI/back/models/RAG/exceptions/prompt.py @@ -0,0 +1,15 @@ +"""Prompt-related RAG exceptions.""" + +from DashAI.back.models.RAG.exceptions.base import RAGWorkflowError + + +class RAGPromptError(RAGWorkflowError): + """Error in prompt template handling.""" + + +class RAGPromptValidationError(RAGPromptError): + """Prompt template validation failed (missing placeholders, etc.).""" + + +class RAGPromptTemplateError(RAGPromptError): + """The prompt template is missing required placeholders or is malformed.""" diff --git a/DashAI/back/models/RAG/exceptions/retriever.py b/DashAI/back/models/RAG/exceptions/retriever.py new file mode 100644 index 000000000..44be81e81 --- /dev/null +++ b/DashAI/back/models/RAG/exceptions/retriever.py @@ -0,0 +1,19 @@ +"""Retriever-related RAG exceptions.""" + +from DashAI.back.models.RAG.exceptions.base import RAGWorkflowError + + +class RAGRetrieverError(RAGWorkflowError): + """Error during retrieval.""" + + +class RAGRetrieverMissingParameterError(RAGRetrieverError): + """A required parameter was not provided to the retriever.""" + + +class RAGRetrieverCompositeValidationError(RAGRetrieverError): + """A composite retriever has an invalid configuration.""" + + +class RAGRetrieverEmptyChildrenError(RAGRetrieverError): + """A composite retriever was configured with no children.""" diff --git a/DashAI/back/models/RAG/exceptions/task.py b/DashAI/back/models/RAG/exceptions/task.py new file mode 100644 index 000000000..8aa0b8cd1 --- /dev/null +++ b/DashAI/back/models/RAG/exceptions/task.py @@ -0,0 +1,11 @@ +"""Task-related RAG exceptions.""" + +from DashAI.back.models.RAG.exceptions.base import RAGWorkflowError + + +class RAGTaskError(RAGWorkflowError): + """Error during RAG task execution.""" + + +class RAGTaskInputError(RAGTaskError): + """Invalid or malformed input to a RAG task.""" diff --git a/DashAI/back/models/RAG/llm_factory.py b/DashAI/back/models/RAG/llm_factory.py new file mode 100644 index 000000000..39b56aadd --- /dev/null +++ b/DashAI/back/models/RAG/llm_factory.py @@ -0,0 +1,62 @@ +"""Pure factory for text-to-text generation models (LLMs). + +Resolves a component name + parameters into an instantiated model +via the component registry, without any database access. +DB concerns are handled by LLMService. +""" + +from dataclasses import dataclass +from typing import Any + +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.RAG.exceptions import RAGComponentNotFoundError +from DashAI.back.models.text_to_text_generation_model import ( + TextToTextGenerationTaskModel, +) + + +@dataclass(frozen=True) +class LLMFactoryResult: + """Result of LLM instantiation via LLMFactory.""" + + model: TextToTextGenerationTaskModel + + +class LLMFactory: + """Creates LLM instances from the component registry.""" + + def __init__(self, component_registry: ComponentRegistry): + """Initialize the factory with a component registry. + + Args: + component_registry: The component registry to resolve LLM + components. + """ + self._registry = component_registry + + def create(self, component_name: str, params: dict[str, Any]) -> LLMFactoryResult: + """Build an LLM instance from a registered component. + + Args: + component_name: Name of the registered LLM component. + params: Parameters to pass to the model constructor. + + Returns: + Contains only the instantiated model. + + Raises: + RAGComponentNotFoundError: If the component_name is not found + in the registry. + """ + try: + model_class = self._registry[component_name]["class"] + except KeyError as err: + raise RAGComponentNotFoundError( + f"Component '{component_name}' not found in registry" + ) from err + + if hasattr(model_class, "SCHEMA") and model_class.SCHEMA is not None: + model_class.SCHEMA(**params) + + model = model_class(**params) + return LLMFactoryResult(model=model) diff --git a/DashAI/back/models/RAG/prompts/__init__.py b/DashAI/back/models/RAG/prompts/__init__.py new file mode 100644 index 000000000..a2fdf1136 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/__init__.py @@ -0,0 +1,12 @@ +from DashAI.back.models.RAG.prompts.prompt import Prompt +from DashAI.back.models.RAG.prompts.augmentation import ( + AugmentationPrompt, + DefaultAugmentationPrompt, + CustomAugmentationPrompt +) +from DashAI.back.models.RAG.prompts.generation import ( + RAGGenerationPrompt, + CustomRAGGenerationPrompt, + DefaultRAGGenerationPrompt, + DefaultQARAGGenerationPrompt +) diff --git a/DashAI/back/models/RAG/prompts/augmentation/__init__.py b/DashAI/back/models/RAG/prompts/augmentation/__init__.py new file mode 100644 index 000000000..e5a807072 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/augmentation/__init__.py @@ -0,0 +1,3 @@ +from DashAI.back.models.RAG.prompts.augmentation.augmentation_prompt import AugmentationPrompt +from DashAI.back.models.RAG.prompts.augmentation.custom_augmentation_prompt import CustomAugmentationPrompt +from DashAI.back.models.RAG.prompts.augmentation.default_augmentation_prompt import DefaultAugmentationPrompt diff --git a/DashAI/back/models/RAG/prompts/augmentation/augmentation_prompt.py b/DashAI/back/models/RAG/prompts/augmentation/augmentation_prompt.py new file mode 100644 index 000000000..56b3183f8 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/augmentation/augmentation_prompt.py @@ -0,0 +1,35 @@ +from typing import Any + +from DashAI.back.models.RAG.prompts.prompt import Prompt + + +class AugmentationPrompt(Prompt): + """ + AugmentationPrompt class for generating augmented retrieval prompts, + it uses the language model to generate keywords or phrases that can be used + to augment the input. + """ + + required_placeholders = ["{input}", "{n_search_terms}"] + optional_placeholders = [] + + def __init__(self, **kwargs: Any): + """Initialize the augmentation prompt with a template. + + Args: + template: The prompt template string. + """ + self.template = kwargs.pop("template") + super().__init__(**kwargs) + + def format(self, input: str, n_search_terms: int, **kwargs: Any) -> str: + """ + Instantiate and format the prompt for augmentation. + Args: + input (str): The input to be formatted. + n_search_terms (int): The number of search terms to generate. + **kwargs: Additional keyword arguments for formatting. + Returns: + str: The formatted prompt. + """ + raise NotImplementedError("Subclasses must implement this method.") diff --git a/DashAI/back/models/RAG/prompts/augmentation/custom_augmentation_prompt.py b/DashAI/back/models/RAG/prompts/augmentation/custom_augmentation_prompt.py new file mode 100644 index 000000000..20dbe4c57 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/augmentation/custom_augmentation_prompt.py @@ -0,0 +1,63 @@ +from typing import Any + +from DashAI.back.models.RAG.exceptions import RAGPromptTemplateError +from DashAI.back.models.RAG.prompts.augmentation import AugmentationPrompt + + +class CustomAugmentationPrompt(AugmentationPrompt): + """ + User-defined augmentation prompt template for generating augmented retrieval + prompts. + It uses the language model to generate keywords or phrases that can be used + to augment the input. + """ + + metadata = { + "name": "Custom Augmentation Prompt", + "description": "User-defined prompt template for generating augmented " + "retrieval prompts.", + "type": "augmentation", + "required_placeholders": AugmentationPrompt.required_placeholders, + "optional_placeholders": AugmentationPrompt.optional_placeholders, + "placeholder_descriptions": { + "{input}": "The user input message.", + "{n_search_terms}": "The number of search terms to generate.", + }, + } + + required_placeholders = ["{input}", "{n_search_terms}"] + optional_placeholders = [] + + def __init__(self, **kwargs): + """Initialize the custom augmentation prompt with a user-defined template. + + Args: + template: The user-defined prompt template string. + """ + self.template = kwargs.pop("template") + super().__init__(**kwargs) + + def format( + self, + input: str, + n_search_terms: int, + **kwargs: Any, + ) -> str: + """ + Instantiate and format the prompt for augmentation. + Args: + input (str): The input to be formatted. + n_search_terms (int): The number of search terms to generate. + **kwargs: Additional keyword arguments for formatting. + Returns: + str: The formatted prompt. + """ + if not self.validate_template(self.template): + raise RAGPromptTemplateError( + "Template is missing required placeholders:" + f" {self.required_placeholders}" + ) + buffer = self.template + buffer = buffer.replace("{input}", input) + buffer = buffer.replace("{n_search_terms}", str(n_search_terms)) + return buffer diff --git a/DashAI/back/models/RAG/prompts/augmentation/default_augmentation_prompt.py b/DashAI/back/models/RAG/prompts/augmentation/default_augmentation_prompt.py new file mode 100644 index 000000000..60b54acce --- /dev/null +++ b/DashAI/back/models/RAG/prompts/augmentation/default_augmentation_prompt.py @@ -0,0 +1,170 @@ +from typing import Any + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.exceptions import RAGPromptTemplateError +from DashAI.back.models.RAG.prompts.augmentation.augmentation_prompt import ( + AugmentationPrompt, +) + +TEMPLATES = { + "en": ( + "You are an intelligent and insightful assistant. Your task is to generate " + "keywords or phrases to search for relevant information based on the input " + "provided.\n" + "The keywords or phrases should be relevant to the input and should help in " + "retrieving useful information to improve the precision of the response.\n" + "\n" + "User input:\n" + "{input}\n" + "\n" + "Generate {n_search_terms} keywords or phrases that can be used to search for " + "relevant information. The keywords or phrases should be concise and to the " + "point.\n" + "\n" + "You must respond with a JSON object following this exact format:\n" + "{'keywords': ['keyword_1', 'keyword_2', ..., 'keyword_n']}" + ), + "es": ( + "Eres un asistente inteligente y perspicaz. Tu tarea es generar palabras clave " + "o frases para buscar información relevante basada en la entrada " + "proporcionada.\n" + "Las palabras clave o frases deben ser relevantes para la entrada y ayudar a " + "recuperar información útil para mejorar la precisión de la respuesta.\n" + "\n" + "Entrada del usuario:\n" + "{input}\n" + "\n" + "Genera {n_search_terms} palabras clave o frases que puedan usarse para buscar " + "información relevante. Las palabras clave o frases deben ser concisas y " + "directas.\n" + "\n" + "Debes responder con un objeto JSON siguiendo este formato exacto:\n" + "{'keywords': ['palabra_clave_1', 'palabra_clave_2', ..., 'palabra_clave_n']}" + ), + "pt": ( + "Você é um assistente inteligente e perspicaz. Sua tarefa é gerar " + "palavras-chave ou frases para buscar informações relevantes com base na " + "entrada fornecida.\n" + "As palavras-chave ou frases devem ser relevantes para a entrada e ajudar a " + "recuperar informações úteis para melhorar a precisão da resposta.\n" + "\n" + "Entrada do usuário:\n" + "{input}\n" + "\n" + "Gere {n_search_terms} palavras-chave ou frases que possam ser usadas para " + "buscar informações relevantes. As palavras-chave ou frases devem ser concisas " + "e diretas.\n" + "\n" + "Você deve responder com um objeto JSON seguindo este formato exato:\n" + "{'keywords': ['palavra_chave_1', 'palavra_chave_2', ..., 'palavra_chave_n']}" + ), +} + + +class DefaultAugmentationPromptSchema(BaseSchema): + """Schema for the default augmentation prompt. + + Attributes: + language: Language code for the response (en, es, pt). + template: The prompt template string with placeholders. + """ + + language: schema_field( + enum_field(enum=["en", "es", "pt"]), + placeholder="en", + description=MultilingualString( + en="Language for the generated response.", + es="Idioma de la respuesta generada.", + pt="Idioma da resposta gerada.", + ), + ) + template: schema_field( + string_field(), + placeholder="", + description="The prompt template with placeholders.", + ) + + +class DefaultAugmentationPrompt(AugmentationPrompt): + """Default prompt template for generating augmented retrieval queries.""" + + SCHEMA = DefaultAugmentationPromptSchema + DESCRIPTION: str = MultilingualString( + en="Default prompt template for generating augmented retrieval queries.", + es="Plantilla de prompt predeterminada para generar consultas de " + "recuperación aumentadas.", + pt="Modelo de prompt padrão para gerar consultas de recuperação aumentadas.", + ) + DISPLAY_NAME: str = MultilingualString( + en="Default Augmentation Prompt", + es="Prompt de Aumento Predeterminado", + pt="Prompt de Aumento Padrão", + ) + + required_placeholders = ["{input}", "{n_search_terms}"] + optional_placeholders: list[str] = [] + + metadata = { + "name": "Default Augmentation Prompt", + "description": "Default prompt template for generating augmented retrieval " + "prompts.", + "type": "augmentation", + "required_placeholders": required_placeholders, + "optional_placeholders": optional_placeholders, + "placeholder_descriptions": { + "{input}": "The user input message.", + "{n_search_terms}": "The number of search terms to generate.", + }, + "templates": TEMPLATES, + } + + def __init__(self, **kwargs: Any) -> None: + """Initialize the augmentation prompt with a language and resolved template. + + Args: + language: Language code to select the appropriate template + (one of "en", "es", "pt"). + template: Override template string. If not provided, the + default template for the selected language is used. + """ + self.language = kwargs.pop("language") + if "template" in kwargs: + self.template = kwargs.pop("template") + else: + self.template = TEMPLATES.get(self.language, "") + if not self.template: + raise RAGPromptTemplateError( + f"No template available for language: {self.language}" + ) + super().__init__(**kwargs) + + def format( + self, + input: str, + n_search_terms: int, + **kwargs: Any, + ) -> str: + """Render the augmentation prompt by replacing all placeholders. + + Args: + input: The user's input message. + n_search_terms: Number of search terms the LLM should generate. + + Returns: + Fully formatted prompt string ready for LLM input. + """ + if not self.validate_template(self.template): + raise RAGPromptTemplateError( + "Template is missing required placeholders:" + f" {self.required_placeholders}" + ) + buffer = self.template + buffer = buffer.replace("{input}", input) + buffer = buffer.replace("{n_search_terms}", str(n_search_terms)) + return buffer diff --git a/DashAI/back/models/RAG/prompts/generation/RAG_generation_prompt.py b/DashAI/back/models/RAG/prompts/generation/RAG_generation_prompt.py new file mode 100644 index 000000000..f231c0325 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/RAG_generation_prompt.py @@ -0,0 +1,31 @@ +from typing import Any + +from DashAI.back.models.RAG.prompts.prompt import Prompt + + +class RAGGenerationPrompt(Prompt): + """ + RAGGenerationPrompt class for formatting prompts used in the language + generation step of RAG. + """ + + required_placeholders = ["{input}", "{chunks}"] + optional_placeholders = [] + + def format( + self, + input: str, + chunks: str, + **kwargs: Any, + ) -> str: + """Format the generation prompt with user input and context chunks. + + Args: + input: The user's input message. + chunks: Retrieved document chunks to include as context. + **kwargs: Additional formatting parameters. + + Returns: + The formatted prompt string. + """ + raise NotImplementedError("Subclasses must implement this method.") diff --git a/DashAI/back/models/RAG/prompts/generation/__init__.py b/DashAI/back/models/RAG/prompts/generation/__init__.py new file mode 100644 index 000000000..39804cab6 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/__init__.py @@ -0,0 +1,4 @@ +from DashAI.back.models.RAG.prompts.generation.RAG_generation_prompt import RAGGenerationPrompt +from DashAI.back.models.RAG.prompts.generation.custom_RAG_generation_prompt import CustomRAGGenerationPrompt +from DashAI.back.models.RAG.prompts.generation.default_RAG_generation_prompt import DefaultRAGGenerationPrompt +from DashAI.back.models.RAG.prompts.generation.default_QA_RAG_generation_prompt import DefaultQARAGGenerationPrompt diff --git a/DashAI/back/models/RAG/prompts/generation/custom_RAG_generation_prompt.py b/DashAI/back/models/RAG/prompts/generation/custom_RAG_generation_prompt.py new file mode 100644 index 000000000..2c0af7197 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/custom_RAG_generation_prompt.py @@ -0,0 +1,57 @@ +from typing import Any + +from DashAI.back.models.RAG.exceptions import RAGPromptTemplateError +from DashAI.back.models.RAG.prompts.generation.RAG_generation_prompt import ( + RAGGenerationPrompt, +) + + +class CustomRAGGenerationPrompt(RAGGenerationPrompt): + """ + CustomRAGGenerationPrompt class for user-defined prompt templates + used in the language generation step of RAG. + """ + + metadata = { + "name": "Custom RAG Generation Prompt", + "description": "User-defined prompt template used in the language" + " generation step of RAG.", + "type": "generation", + "required_placeholders": RAGGenerationPrompt.required_placeholders, + "optional_placeholders": RAGGenerationPrompt.optional_placeholders, + "placeholder_descriptions": { + "{input}": "The user input message.", + "{chunks}": "The document chunks to be included in the context.", + }, + } + + def __init__(self, **kwargs: Any): + """Initialize the custom RAG generation prompt. + + Args: + template: The user-defined prompt template string. + """ + self.template = kwargs.pop("template") + + def format(self, input: str, chunks: str, **kwargs: Any) -> str: + """ + Format the prompt using the provided template. + + Args: + input (str): The user input message. + chunks (List[str]): The document chunks to be included in the + context. + **kwargs: Additional keyword arguments for formatting. + + Returns: + str: The formatted prompt. + """ + if not self.validate_template(self.template): + raise RAGPromptTemplateError( + "Template is missing required placeholders:" + f" {self.required_placeholders}" + ) + buffer = self.template + buffer = buffer.replace("{input}", input) + buffer = buffer.replace("{chunks}", chunks) + return buffer diff --git a/DashAI/back/models/RAG/prompts/generation/custom_rag_generation_prompt.py b/DashAI/back/models/RAG/prompts/generation/custom_rag_generation_prompt.py new file mode 100644 index 000000000..2c0af7197 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/custom_rag_generation_prompt.py @@ -0,0 +1,57 @@ +from typing import Any + +from DashAI.back.models.RAG.exceptions import RAGPromptTemplateError +from DashAI.back.models.RAG.prompts.generation.RAG_generation_prompt import ( + RAGGenerationPrompt, +) + + +class CustomRAGGenerationPrompt(RAGGenerationPrompt): + """ + CustomRAGGenerationPrompt class for user-defined prompt templates + used in the language generation step of RAG. + """ + + metadata = { + "name": "Custom RAG Generation Prompt", + "description": "User-defined prompt template used in the language" + " generation step of RAG.", + "type": "generation", + "required_placeholders": RAGGenerationPrompt.required_placeholders, + "optional_placeholders": RAGGenerationPrompt.optional_placeholders, + "placeholder_descriptions": { + "{input}": "The user input message.", + "{chunks}": "The document chunks to be included in the context.", + }, + } + + def __init__(self, **kwargs: Any): + """Initialize the custom RAG generation prompt. + + Args: + template: The user-defined prompt template string. + """ + self.template = kwargs.pop("template") + + def format(self, input: str, chunks: str, **kwargs: Any) -> str: + """ + Format the prompt using the provided template. + + Args: + input (str): The user input message. + chunks (List[str]): The document chunks to be included in the + context. + **kwargs: Additional keyword arguments for formatting. + + Returns: + str: The formatted prompt. + """ + if not self.validate_template(self.template): + raise RAGPromptTemplateError( + "Template is missing required placeholders:" + f" {self.required_placeholders}" + ) + buffer = self.template + buffer = buffer.replace("{input}", input) + buffer = buffer.replace("{chunks}", chunks) + return buffer diff --git a/DashAI/back/models/RAG/prompts/generation/default_QA_RAG_generation_prompt.py b/DashAI/back/models/RAG/prompts/generation/default_QA_RAG_generation_prompt.py new file mode 100644 index 000000000..a0a96a7ac --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/default_QA_RAG_generation_prompt.py @@ -0,0 +1,180 @@ +from typing import Any + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.exceptions import RAGPromptTemplateError +from DashAI.back.models.RAG.prompts.generation.RAG_generation_prompt import ( + RAGGenerationPrompt, +) + +TEMPLATES = { + "en": ( + "You are a question-answering assistant with access to a set of " + "retrieved documents.\n" + "Your task is to answer the user's question using ONLY the " + "information found in the documents below. If the documents do " + "not contain enough information to answer the question, respond " + 'with "I don\'t know" and explain what information is missing.\n' + "Cite relevant passages from the documents to support your answer.\n" + "\n" + "User question:\n" + "{input}\n" + "\n" + "Retrieved documents:\n" + "{chunks}\n" + "\n" + "Answer:" + ), + "es": ( + "Eres un asistente de preguntas y respuestas con acceso a un " + "conjunto de documentos recuperados.\n" + "Tu tarea es responder la pregunta del usuario usando SOLAMENTE " + "la información encontrada en los documentos a continuación. Si " + "los documentos no contienen suficiente información para responder " + 'la pregunta, responde con "No lo sé" y explica qué información ' + "falta.\n" + "Cita pasajes relevantes de los documentos para respaldar tu " + "respuesta.\n" + "\n" + "Pregunta del usuario:\n" + "{input}\n" + "\n" + "Documentos recuperados:\n" + "{chunks}\n" + "\n" + "Respuesta:" + ), + "pt": ( + "Você é um assistente de perguntas e respostas com acesso a um " + "conjunto de documentos recuperados.\n" + "Sua tarefa é responder à pergunta do usuário usando SOMENTE as " + "informações encontradas nos documentos abaixo. Se os documentos " + "não contiverem informações suficientes para responder à pergunta, " + 'responda com "Não sei" e explique quais informações estão ' + "faltando.\n" + "Cite passagens relevantes dos documentos para apoiar sua " + "resposta.\n" + "\n" + "Pergunta do usuário:\n" + "{input}\n" + "\n" + "Documentos recuperados:\n" + "{chunks}\n" + "\n" + "Resposta:" + ), +} + + +class DefaultQARAGGenerationPromptSchema(BaseSchema): + """Schema for the default QA RAG generation prompt. + + Attributes: + language: Language code for the response (en, es, pt). + template: The prompt template string with placeholders. + """ + + language: schema_field( + enum_field(enum=["en", "es", "pt"]), + placeholder="en", + description=MultilingualString( + en="Language for the generated response.", + es="Idioma de la respuesta generada.", + pt="Idioma da resposta gerada.", + ), + ) + template: schema_field( + string_field(), + placeholder="", + description="The prompt template with placeholders.", + ) + + +class DefaultQARAGGenerationPrompt(RAGGenerationPrompt): + """ + Default generation prompt for Question Answering tasks. + This prompt is designed to guide the language model in generating + answers based on provided context chunks. + """ + + SCHEMA = DefaultQARAGGenerationPromptSchema + DESCRIPTION: str = MultilingualString( + en="Default prompt template used in the language generation step" + " of RAG for Question Answering tasks.", + es="Plantilla de prompt predeterminada utilizada en el paso de" + " generación de lenguaje de RAG para tareas de preguntas y" + " respuestas.", + pt="Modelo de prompt padrão usado na etapa de geração de linguagem" + " do RAG para tarefas de perguntas e respostas.", + ) + DISPLAY_NAME: str = MultilingualString( + en="Default QA RAG Generation Prompt", + es="Prompt de Generación RAG de Preguntas y Respuestas Predeterminado", + pt="Prompt de Geração RAG de Perguntas e Respostas Padrão", + ) + + metadata = { + "name": "Default QA RAG Generation Prompt", + "description": "Default prompt template used in the language" + " generation step of RAG for Question Answering tasks.", + "type": "generation", + "required_placeholders": RAGGenerationPrompt.required_placeholders, + "optional_placeholders": RAGGenerationPrompt.optional_placeholders, + "placeholder_descriptions": { + "{input}": "The user input message.", + "{chunks}": "The document chunks to be included in the context.", + }, + "templates": TEMPLATES, + } + + def __init__(self, **kwargs): + """Initialize the default QA RAG generation prompt. + + Args: + language: Language code (one of "en", "es", "pt"). + template: Optional override template string. Falls back to the + default template for the selected language if not provided. + + Raises: + RAGPromptTemplateError: If no template is available for the + selected language. + """ + self.language = kwargs.pop("language") + if "template" in kwargs: + self.template = kwargs.pop("template") + else: + self.template = TEMPLATES.get(self.language, "") + if not self.template: + raise RAGPromptTemplateError( + f"No template available for language: {self.language}" + ) + + def format(self, input: str, chunks: str, **kwargs: Any) -> str: + """Render the QA prompt by replacing placeholders with actual values. + + Args: + input: The user's question. + chunks: Retrieved document chunks to include as context. + **kwargs: Additional formatting parameters. + + Returns: + The fully formatted QA prompt string. + + Raises: + RAGPromptTemplateError: If the template is missing required + placeholders. + """ + if not self.validate_template(self.template): + raise RAGPromptTemplateError( + "Template is missing required placeholders:" + f" {self.required_placeholders}" + ) + buffer = self.template + buffer = buffer.replace("{input}", input) + buffer = buffer.replace("{chunks}", chunks) + return buffer diff --git a/DashAI/back/models/RAG/prompts/generation/default_QA_rag_generation_prompt.py b/DashAI/back/models/RAG/prompts/generation/default_QA_rag_generation_prompt.py new file mode 100644 index 000000000..a0a96a7ac --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/default_QA_rag_generation_prompt.py @@ -0,0 +1,180 @@ +from typing import Any + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.exceptions import RAGPromptTemplateError +from DashAI.back.models.RAG.prompts.generation.RAG_generation_prompt import ( + RAGGenerationPrompt, +) + +TEMPLATES = { + "en": ( + "You are a question-answering assistant with access to a set of " + "retrieved documents.\n" + "Your task is to answer the user's question using ONLY the " + "information found in the documents below. If the documents do " + "not contain enough information to answer the question, respond " + 'with "I don\'t know" and explain what information is missing.\n' + "Cite relevant passages from the documents to support your answer.\n" + "\n" + "User question:\n" + "{input}\n" + "\n" + "Retrieved documents:\n" + "{chunks}\n" + "\n" + "Answer:" + ), + "es": ( + "Eres un asistente de preguntas y respuestas con acceso a un " + "conjunto de documentos recuperados.\n" + "Tu tarea es responder la pregunta del usuario usando SOLAMENTE " + "la información encontrada en los documentos a continuación. Si " + "los documentos no contienen suficiente información para responder " + 'la pregunta, responde con "No lo sé" y explica qué información ' + "falta.\n" + "Cita pasajes relevantes de los documentos para respaldar tu " + "respuesta.\n" + "\n" + "Pregunta del usuario:\n" + "{input}\n" + "\n" + "Documentos recuperados:\n" + "{chunks}\n" + "\n" + "Respuesta:" + ), + "pt": ( + "Você é um assistente de perguntas e respostas com acesso a um " + "conjunto de documentos recuperados.\n" + "Sua tarefa é responder à pergunta do usuário usando SOMENTE as " + "informações encontradas nos documentos abaixo. Se os documentos " + "não contiverem informações suficientes para responder à pergunta, " + 'responda com "Não sei" e explique quais informações estão ' + "faltando.\n" + "Cite passagens relevantes dos documentos para apoiar sua " + "resposta.\n" + "\n" + "Pergunta do usuário:\n" + "{input}\n" + "\n" + "Documentos recuperados:\n" + "{chunks}\n" + "\n" + "Resposta:" + ), +} + + +class DefaultQARAGGenerationPromptSchema(BaseSchema): + """Schema for the default QA RAG generation prompt. + + Attributes: + language: Language code for the response (en, es, pt). + template: The prompt template string with placeholders. + """ + + language: schema_field( + enum_field(enum=["en", "es", "pt"]), + placeholder="en", + description=MultilingualString( + en="Language for the generated response.", + es="Idioma de la respuesta generada.", + pt="Idioma da resposta gerada.", + ), + ) + template: schema_field( + string_field(), + placeholder="", + description="The prompt template with placeholders.", + ) + + +class DefaultQARAGGenerationPrompt(RAGGenerationPrompt): + """ + Default generation prompt for Question Answering tasks. + This prompt is designed to guide the language model in generating + answers based on provided context chunks. + """ + + SCHEMA = DefaultQARAGGenerationPromptSchema + DESCRIPTION: str = MultilingualString( + en="Default prompt template used in the language generation step" + " of RAG for Question Answering tasks.", + es="Plantilla de prompt predeterminada utilizada en el paso de" + " generación de lenguaje de RAG para tareas de preguntas y" + " respuestas.", + pt="Modelo de prompt padrão usado na etapa de geração de linguagem" + " do RAG para tarefas de perguntas e respostas.", + ) + DISPLAY_NAME: str = MultilingualString( + en="Default QA RAG Generation Prompt", + es="Prompt de Generación RAG de Preguntas y Respuestas Predeterminado", + pt="Prompt de Geração RAG de Perguntas e Respostas Padrão", + ) + + metadata = { + "name": "Default QA RAG Generation Prompt", + "description": "Default prompt template used in the language" + " generation step of RAG for Question Answering tasks.", + "type": "generation", + "required_placeholders": RAGGenerationPrompt.required_placeholders, + "optional_placeholders": RAGGenerationPrompt.optional_placeholders, + "placeholder_descriptions": { + "{input}": "The user input message.", + "{chunks}": "The document chunks to be included in the context.", + }, + "templates": TEMPLATES, + } + + def __init__(self, **kwargs): + """Initialize the default QA RAG generation prompt. + + Args: + language: Language code (one of "en", "es", "pt"). + template: Optional override template string. Falls back to the + default template for the selected language if not provided. + + Raises: + RAGPromptTemplateError: If no template is available for the + selected language. + """ + self.language = kwargs.pop("language") + if "template" in kwargs: + self.template = kwargs.pop("template") + else: + self.template = TEMPLATES.get(self.language, "") + if not self.template: + raise RAGPromptTemplateError( + f"No template available for language: {self.language}" + ) + + def format(self, input: str, chunks: str, **kwargs: Any) -> str: + """Render the QA prompt by replacing placeholders with actual values. + + Args: + input: The user's question. + chunks: Retrieved document chunks to include as context. + **kwargs: Additional formatting parameters. + + Returns: + The fully formatted QA prompt string. + + Raises: + RAGPromptTemplateError: If the template is missing required + placeholders. + """ + if not self.validate_template(self.template): + raise RAGPromptTemplateError( + "Template is missing required placeholders:" + f" {self.required_placeholders}" + ) + buffer = self.template + buffer = buffer.replace("{input}", input) + buffer = buffer.replace("{chunks}", chunks) + return buffer diff --git a/DashAI/back/models/RAG/prompts/generation/default_RAG_generation_prompt.py b/DashAI/back/models/RAG/prompts/generation/default_RAG_generation_prompt.py new file mode 100644 index 000000000..e0e0718ab --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/default_RAG_generation_prompt.py @@ -0,0 +1,160 @@ +from typing import Any + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.exceptions import RAGPromptTemplateError +from DashAI.back.models.RAG.prompts.generation.RAG_generation_prompt import ( + RAGGenerationPrompt, +) + +TEMPLATES = { + "en": ( + "You are a helpful AI assistant with access to a set of retrieved " + "documents.\n" + "Use the documents below to support your answer. Prioritize " + "information from the documents over your own knowledge. " + "If the documents do not contain relevant information, say so " + "clearly rather than guessing.\n" + "\n" + "User message:\n" + "{input}\n" + "\n" + "Retrieved documents:\n" + "{chunks}\n" + "\n" + "Answer the user's message based on the documents above." + ), + "es": ( + "Eres un asistente de IA con acceso a un conjunto de documentos " + "recuperados.\n" + "Usa los documentos a continuación para fundamentar tu respuesta. " + "Prioriza la información de los documentos sobre tu propio " + "conocimiento. Si los documentos no contienen información relevante, " + "indícalo claramente en lugar de adivinar.\n" + "\n" + "Mensaje del usuario:\n" + "{input}\n" + "\n" + "Documentos recuperados:\n" + "{chunks}\n" + "\n" + "Responde al mensaje del usuario basándote en los documentos " + "anteriores." + ), + "pt": ( + "Você é um assistente de IA com acesso a um conjunto de documentos " + "recuperados.\n" + "Use os documentos abaixo para fundamentar sua resposta. Priorize " + "as informações dos documentos sobre seu próprio conhecimento. Se " + "os documentos não contiverem informações relevantes, indique isso " + "claramente em vez de adivinhar.\n" + "\n" + "Mensagem do usuário:\n" + "{input}\n" + "\n" + "Documentos recuperados:\n" + "{chunks}\n" + "\n" + "Responda à mensagem do usuário com base nos documentos acima." + ), +} + + +class DefaultRAGGenerationPromptSchema(BaseSchema): + """Schema for the default RAG generation prompt. + + Attributes: + language: Language code for the response (en, es, pt). + template: The prompt template string with placeholders. + """ + + language: schema_field( + enum_field(enum=["en", "es", "pt"]), + placeholder="en", + description=MultilingualString( + en="Language for the generated response.", + es="Idioma de la respuesta generada.", + pt="Idioma da resposta gerada.", + ), + ) + template: schema_field( + string_field(), + placeholder="", + description="The prompt template with placeholders.", + ) + + +class DefaultRAGGenerationPrompt(RAGGenerationPrompt): + """ + Default prompt template used in the language generation step of RAG. + """ + + SCHEMA = DefaultRAGGenerationPromptSchema + DESCRIPTION: str = MultilingualString( + en="Default prompt template used in the language generation step of RAG.", + es="Plantilla de prompt predeterminada utilizada en el paso de" + " generación de lenguaje de RAG.", + pt="Modelo de prompt padrão usado na etapa de geração de linguagem do RAG.", + ) + DISPLAY_NAME: str = MultilingualString( + en="Default RAG Generation Prompt", + es="Prompt de Generación RAG Predeterminado", + pt="Prompt de Geração RAG Padrão", + ) + + metadata = { + "name": "Default RAG Generation Prompt", + "description": "Default prompt template used in the language" + " generation step of RAG.", + "type": "generation", + "required_placeholders": RAGGenerationPrompt.required_placeholders, + "optional_placeholders": RAGGenerationPrompt.optional_placeholders, + "placeholder_descriptions": { + "{input}": "The user input message.", + "{chunks}": "The document chunks to be included in the context.", + }, + "templates": TEMPLATES, + } + + required_placeholders = ["{input}", "{chunks}"] + optional_placeholders = [] + + def __init__(self, **kwargs): + """Initialize the default RAG generation prompt. + + Args: + language: Language code (one of "en", "es", "pt"). + template: The prompt template string. + """ + self.language = kwargs.pop("language") + self.template = kwargs.pop("template") + + def format(self, input: str, chunks: str, **kwargs: Any) -> str: + """Render the prompt by replacing placeholders with actual values. + + Args: + input: The user's input message. + chunks: Retrieved document chunks to include as context. + **kwargs: Additional formatting parameters. + + Returns: + The fully formatted prompt string. + + Raises: + RAGPromptTemplateError: If the template is missing required + placeholders. + """ + if not self.validate_template(self.template): + raise RAGPromptTemplateError( + "Template is missing required placeholders:" + f" {self.required_placeholders}" + ) + buffer = self.template + buffer = buffer.replace("{input}", input) + buffer = buffer.replace("{chunks}", chunks) + return buffer diff --git a/DashAI/back/models/RAG/prompts/generation/default_rag_generation_prompt.py b/DashAI/back/models/RAG/prompts/generation/default_rag_generation_prompt.py new file mode 100644 index 000000000..e0e0718ab --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/default_rag_generation_prompt.py @@ -0,0 +1,160 @@ +from typing import Any + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.exceptions import RAGPromptTemplateError +from DashAI.back.models.RAG.prompts.generation.RAG_generation_prompt import ( + RAGGenerationPrompt, +) + +TEMPLATES = { + "en": ( + "You are a helpful AI assistant with access to a set of retrieved " + "documents.\n" + "Use the documents below to support your answer. Prioritize " + "information from the documents over your own knowledge. " + "If the documents do not contain relevant information, say so " + "clearly rather than guessing.\n" + "\n" + "User message:\n" + "{input}\n" + "\n" + "Retrieved documents:\n" + "{chunks}\n" + "\n" + "Answer the user's message based on the documents above." + ), + "es": ( + "Eres un asistente de IA con acceso a un conjunto de documentos " + "recuperados.\n" + "Usa los documentos a continuación para fundamentar tu respuesta. " + "Prioriza la información de los documentos sobre tu propio " + "conocimiento. Si los documentos no contienen información relevante, " + "indícalo claramente en lugar de adivinar.\n" + "\n" + "Mensaje del usuario:\n" + "{input}\n" + "\n" + "Documentos recuperados:\n" + "{chunks}\n" + "\n" + "Responde al mensaje del usuario basándote en los documentos " + "anteriores." + ), + "pt": ( + "Você é um assistente de IA com acesso a um conjunto de documentos " + "recuperados.\n" + "Use os documentos abaixo para fundamentar sua resposta. Priorize " + "as informações dos documentos sobre seu próprio conhecimento. Se " + "os documentos não contiverem informações relevantes, indique isso " + "claramente em vez de adivinhar.\n" + "\n" + "Mensagem do usuário:\n" + "{input}\n" + "\n" + "Documentos recuperados:\n" + "{chunks}\n" + "\n" + "Responda à mensagem do usuário com base nos documentos acima." + ), +} + + +class DefaultRAGGenerationPromptSchema(BaseSchema): + """Schema for the default RAG generation prompt. + + Attributes: + language: Language code for the response (en, es, pt). + template: The prompt template string with placeholders. + """ + + language: schema_field( + enum_field(enum=["en", "es", "pt"]), + placeholder="en", + description=MultilingualString( + en="Language for the generated response.", + es="Idioma de la respuesta generada.", + pt="Idioma da resposta gerada.", + ), + ) + template: schema_field( + string_field(), + placeholder="", + description="The prompt template with placeholders.", + ) + + +class DefaultRAGGenerationPrompt(RAGGenerationPrompt): + """ + Default prompt template used in the language generation step of RAG. + """ + + SCHEMA = DefaultRAGGenerationPromptSchema + DESCRIPTION: str = MultilingualString( + en="Default prompt template used in the language generation step of RAG.", + es="Plantilla de prompt predeterminada utilizada en el paso de" + " generación de lenguaje de RAG.", + pt="Modelo de prompt padrão usado na etapa de geração de linguagem do RAG.", + ) + DISPLAY_NAME: str = MultilingualString( + en="Default RAG Generation Prompt", + es="Prompt de Generación RAG Predeterminado", + pt="Prompt de Geração RAG Padrão", + ) + + metadata = { + "name": "Default RAG Generation Prompt", + "description": "Default prompt template used in the language" + " generation step of RAG.", + "type": "generation", + "required_placeholders": RAGGenerationPrompt.required_placeholders, + "optional_placeholders": RAGGenerationPrompt.optional_placeholders, + "placeholder_descriptions": { + "{input}": "The user input message.", + "{chunks}": "The document chunks to be included in the context.", + }, + "templates": TEMPLATES, + } + + required_placeholders = ["{input}", "{chunks}"] + optional_placeholders = [] + + def __init__(self, **kwargs): + """Initialize the default RAG generation prompt. + + Args: + language: Language code (one of "en", "es", "pt"). + template: The prompt template string. + """ + self.language = kwargs.pop("language") + self.template = kwargs.pop("template") + + def format(self, input: str, chunks: str, **kwargs: Any) -> str: + """Render the prompt by replacing placeholders with actual values. + + Args: + input: The user's input message. + chunks: Retrieved document chunks to include as context. + **kwargs: Additional formatting parameters. + + Returns: + The fully formatted prompt string. + + Raises: + RAGPromptTemplateError: If the template is missing required + placeholders. + """ + if not self.validate_template(self.template): + raise RAGPromptTemplateError( + "Template is missing required placeholders:" + f" {self.required_placeholders}" + ) + buffer = self.template + buffer = buffer.replace("{input}", input) + buffer = buffer.replace("{chunks}", chunks) + return buffer diff --git a/DashAI/back/models/RAG/prompts/generation/rag_generation_prompt.py b/DashAI/back/models/RAG/prompts/generation/rag_generation_prompt.py new file mode 100644 index 000000000..f231c0325 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/rag_generation_prompt.py @@ -0,0 +1,31 @@ +from typing import Any + +from DashAI.back.models.RAG.prompts.prompt import Prompt + + +class RAGGenerationPrompt(Prompt): + """ + RAGGenerationPrompt class for formatting prompts used in the language + generation step of RAG. + """ + + required_placeholders = ["{input}", "{chunks}"] + optional_placeholders = [] + + def format( + self, + input: str, + chunks: str, + **kwargs: Any, + ) -> str: + """Format the generation prompt with user input and context chunks. + + Args: + input: The user's input message. + chunks: Retrieved document chunks to include as context. + **kwargs: Additional formatting parameters. + + Returns: + The formatted prompt string. + """ + raise NotImplementedError("Subclasses must implement this method.") diff --git a/DashAI/back/models/RAG/prompts/prompt.py b/DashAI/back/models/RAG/prompts/prompt.py new file mode 100644 index 000000000..d18e0f988 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/prompt.py @@ -0,0 +1,116 @@ +from typing import Any, Dict, List + +from DashAI.back.core.schema_fields import BaseSchema, schema_field, string_field +from DashAI.back.models.base_model import BaseModel + + +class PromptSchema(BaseSchema): + """Schema for prompt templates. + + Attributes: + template: The prompt template string with placeholders. + """ + + template: str = schema_field( + string_field(), + placeholder="", + description="The prompt template with placeholders.", + ) + + +class Prompt(BaseModel): + """ + Base class for all RAG prompt templates. + This class defines the interface for creating and formatting prompts. + """ + + SCHEMA = PromptSchema + DESCRIPTION: str = "Base class for RAG prompts." + DISPLAY_NAME: str = "Base RAG Prompt" + REQUIRED_EXTRA_KWARGS = [] + + def load(self, filename: str = "") -> None: + """Load a prompt from a file. + + Args: + filename: Path to the file to load from. If empty, uses the + default. + """ + + def save(self, filename: str = "") -> None: + """Save the prompt to a file. + + Args: + filename: Path to save to. If empty, uses the default. + """ + + def train(self, **kwargs: Any) -> None: + """No-op training method for compatibility with the model interface. + + Args: + **kwargs: Ignored. + """ + + @classmethod + def get_metadata(cls) -> Dict[str, Any]: + """Retrieve class metadata. + + Returns: + Dictionary of metadata attributes if defined, otherwise an empty + dict. + """ + metadata = cls.metadata if hasattr(cls, "metadata") else {} + return metadata + + @classmethod + def get_required_placeholders(cls) -> List[str]: + """ + Get the list of required placeholders for the prompt template. + Returns: + List[str]: List of required placeholders. + Raises: + AttributeError: If the subclass does not define 'required_placeholders'. + """ + if not hasattr(cls, "required_placeholders"): + raise AttributeError( + f"Prompt subclass {cls.__name__} must define " + "'required_placeholders' class attribute." + ) + return cls.required_placeholders + + def get_optional_placeholders(self) -> List[str]: + """ + Get the list of optional placeholders for the prompt template. + Returns: + List[str]: List of optional placeholders. + Raises: + AttributeError: If the subclass does not define 'optional_placeholders'. + """ + if not hasattr(self, "optional_placeholders"): + raise AttributeError( + f"Prompt subclass {type(self).__name__} must define " + "'optional_placeholders' class attribute." + ) + return self.optional_placeholders + + @classmethod + def validate_template(cls, template: str) -> bool: + """ + Validate that the template contains all required placeholders. + Args: + template (str): The prompt template to be validated. + Returns: + bool: True if the template is valid, False otherwise. + """ + return all(placeholder in template for placeholder in cls.required_placeholders) + + def format(self, input: str, **kwargs: Any) -> str: + """ + Instantiate and format the prompt. + Args: + input (str): The input to be formatted. + **kwargs: Additional keyword arguments for formatting. + Returns: + str: The formatted prompt. + """ + raise NotImplementedError("Subclasses must implement this method.") diff --git a/DashAI/back/models/RAG/prompts/prompt_factory.py b/DashAI/back/models/RAG/prompts/prompt_factory.py new file mode 100644 index 000000000..c64e073f4 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/prompt_factory.py @@ -0,0 +1,56 @@ +"""Pure factory for prompt models. Phase 1 only — no DB.""" + +from dataclasses import dataclass +from typing import Any + +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.RAG.exceptions import RAGComponentNotFoundError +from DashAI.back.models.RAG.prompts import Prompt + + +@dataclass(frozen=True) +class PromptFactoryResult: + """Result of prompt instantiation via PromptFactory.""" + + model: Prompt + + +class PromptFactory: + """Creates Prompt instances from the component registry.""" + + def __init__(self, registry: ComponentRegistry): + """Initialize the factory with a component registry. + + Args: + registry: The component registry to resolve prompt components. + """ + self._registry = registry + + def create( + self, component_name: str, params: dict[str, Any] + ) -> PromptFactoryResult: + """Build a Prompt instance from a registered component. + + Args: + component_name: Name of the registered prompt component. + params: Parameters to pass to the prompt constructor. + + Returns: + Contains only the instantiated Prompt model. + + Raises: + RAGComponentNotFoundError: If the component_name is not found + in the registry. + """ + try: + prompt_class = self._registry[component_name]["class"] + except KeyError as err: + raise RAGComponentNotFoundError( + f"Prompt component '{component_name}' not found in registry" + ) from err + + if hasattr(prompt_class, "SCHEMA") and prompt_class.SCHEMA is not None: + prompt_class.SCHEMA(**params) + + model = prompt_class(**params) + return PromptFactoryResult(model=model) diff --git a/DashAI/back/models/RAG/rag_models_factory.py b/DashAI/back/models/RAG/rag_models_factory.py new file mode 100644 index 000000000..b183d4ceb --- /dev/null +++ b/DashAI/back/models/RAG/rag_models_factory.py @@ -0,0 +1,75 @@ +"""Unified factory that delegates to type-specific sub-factories. + +Phase 1 of the 3-phase lifecycle (Construction -> Initialization -> Persistence). +Builds models in memory -- no I/O, no DB access. +""" + +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.RAG.chunking_models.chunking_model_factory import ( + ChunkingFactoryResult, + ChunkingModelFactory, +) +from DashAI.back.models.RAG.documents.base_document import BaseDocument +from DashAI.back.models.RAG.documents.chunk import Chunk +from DashAI.back.models.RAG.llm_factory import LLMFactory, LLMFactoryResult +from DashAI.back.models.RAG.prompts.prompt_factory import ( + PromptFactory, + PromptFactoryResult, +) +from DashAI.back.models.RAG.retrievers.persistence import ( + DensePersistence, + SparsePersistence, +) +from DashAI.back.models.RAG.retrievers.retriever_factory import ( + RetrieverFactory, + RetrieverFactoryResult, +) + + +class RAGModelsFactory: + """Delegates model construction to type-specific sub-factories. + + Every ``create_*`` method: + - Accepts a component name + params dict + - Delegates to the corresponding sub-factory + - Returns the model -- no I/O, no DB + """ + + def __init__(self, registry: ComponentRegistry): + """Initialise the factory with a component registry. + + Args: + registry: The ComponentRegistry used to resolve component + names to their implementations. + """ + self._registry = registry + + def create_chunking_model( + self, + component: str, + params: dict, + documents: dict[int, BaseDocument], + ) -> ChunkingFactoryResult: + """Build a chunking model via ChunkingModelFactory. Phase 1 only.""" + factory = ChunkingModelFactory(self._registry, documents) + return factory.create(component, params.copy()) + + def create_retriever( + self, + component: str, + params: dict, + RAG_path: str, # noqa: N803 + chunks: dict[int, dict[int, Chunk]], + persistence: DensePersistence | SparsePersistence | None = None, + ) -> RetrieverFactoryResult: + """Build a retriever via RetrieverFactory. Phase 1 only.""" + factory = RetrieverFactory(self._registry, RAG_path, chunks) + return factory.create(component, params, persistence) + + def create_prompt(self, component: str, params: dict) -> PromptFactoryResult: + """Build a prompt via PromptFactory. Phase 1 only.""" + return PromptFactory(self._registry).create(component, params) + + def create_llm(self, component: str, params: dict) -> LLMFactoryResult: + """Build an LLM via LLMFactory. Phase 1 only.""" + return LLMFactory(self._registry).create(component, params) diff --git a/DashAI/back/models/RAG/retrievers/__init__.py b/DashAI/back/models/RAG/retrievers/__init__.py new file mode 100644 index 000000000..2be403d2c --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/__init__.py @@ -0,0 +1,32 @@ +"""Convenience re-exports for all retriever types under a single namespace.""" + +from DashAI.back.models.RAG.retrievers.composite import ( + CompositeRetriever, + MMRRerankerRetriever, + ParallelRetriever, + SequentialRetriever, +) +from DashAI.back.models.RAG.retrievers.dense import ( + DenseEmbeddingRetriever, + DenseRetriever, + HuggingFaceDenseRetriever, +) +from DashAI.back.models.RAG.retrievers.enums import MergeStrategy +from DashAI.back.models.RAG.retrievers.exceptions import ( + CompositeValidationError, + MissingParameterError, + RetrieverError, +) +from DashAI.back.models.RAG.retrievers.retriever_factory import ( + RetrieverFactory, + RetrieverFactoryResult, +) +from DashAI.back.models.RAG.retrievers.retriever_model import RetrieverModel +from DashAI.back.models.RAG.retrievers.sparse import ( + BM25Retriever, + BM25VectorizerModel, + SparseRetriever, + TFIDFRetriever, + TFIDFVectorizerModel, +) +from DashAI.back.models.RAG.retrievers.unit_retriever import UnitRetriever diff --git a/DashAI/back/models/RAG/retrievers/composite/__init__.py b/DashAI/back/models/RAG/retrievers/composite/__init__.py new file mode 100644 index 000000000..17e2d4396 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/composite/__init__.py @@ -0,0 +1,14 @@ +"""Convenience re-exports for all composite retriever types.""" + +from DashAI.back.models.RAG.retrievers.composite.composite_retriever import ( + CompositeRetriever, +) +from DashAI.back.models.RAG.retrievers.composite.mmr_reranker_retriever import ( + MMRRerankerRetriever, +) +from DashAI.back.models.RAG.retrievers.composite.parallel_retriever import ( + ParallelRetriever, +) +from DashAI.back.models.RAG.retrievers.composite.sequential_retriever import ( + SequentialRetriever, +) diff --git a/DashAI/back/models/RAG/retrievers/composite/composite_retriever.py b/DashAI/back/models/RAG/retrievers/composite/composite_retriever.py new file mode 100644 index 000000000..104fae83e --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/composite/composite_retriever.py @@ -0,0 +1,108 @@ +from abc import ABC, abstractmethod +from typing import Final, List + +from DashAI.back.models.RAG.documents import Chunk +from DashAI.back.models.RAG.retrievers.retriever_model import RetrieverModel + + +class CompositeRetriever(RetrieverModel, ABC): + """ + Composite: abstract base for retrievers that contain child retrievers. + + Implements the Composite role in the Composite design pattern (GoF). + """ + + TYPE: Final[str] = "RetrieverModel" + FLAGS: list[str] = ["abstract"] + REQUIRED_EXTRA_KWARGS: list = [] + + def __init__(self, **kwargs): + """Initialize the composite retriever with its children. + + Pops the ``children`` key from *kwargs* and validates that it is + a list of :class:`RetrieverModel` instances. + + Args: + **kwargs: Must contain a ``children`` key mapping to a list + of :class:`RetrieverModel` instances. + + Raises: + TypeError: If ``children`` is not a list or contains + non-retriever elements. + """ + children_data = kwargs.pop("children") + if not isinstance(children_data, list): + raise TypeError( + f"'children' must be a list, got {type(children_data).__name__}" + ) + if children_data and not all( + isinstance(c, RetrieverModel) for c in children_data + ): + raise TypeError( + "All elements in 'children' must be RetrieverModel instances" + ) + self._children: List[RetrieverModel] = children_data + super().__init__(**kwargs) + + def add(self, child: RetrieverModel) -> None: + """Add a child retriever. + + Args: + child: The :class:`RetrieverModel` instance to add. + """ + self._children.append(child) + + def remove(self, child: RetrieverModel) -> None: + """Remove a child retriever. + + Args: + child: The :class:`RetrieverModel` instance to remove. + """ + self._children.remove(child) + + def get_children(self) -> List[RetrieverModel]: + """Return a copy of the children list. + + Returns: + A new list containing all child :class:`RetrieverModel` + instances. + """ + return list(self._children) + + @abstractmethod + def retrieve(self, query, **kwargs) -> List[Chunk]: + """Retrieve chunks by delegating to child retrievers. + + Args: + query: The search query string. + **kwargs: Additional retrieval parameters. + + Returns: + A list of :class:`Chunk` instances. + """ + raise NotImplementedError + + def score_chunks(self, chunk_ids: List[int], query: str) -> list: + """Score chunks against a query. + + .. note:: + The base ``CompositeRetriever`` does not implement chunk scoring — + each concrete composite subclass that supports scoring must + override this method. The ``SequentialRetriever``, for instance, + delegates ``score_chunks`` to its leaf children. + + Args: + chunk_ids: List of chunk IDs to score. + query: The search query string. + + Returns: + A list of ``(chunk_id, distance)`` tuples. + + Raises: + NotImplementedError: Always — subclasses must override. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement score_chunks. " + "Concrete composite retrievers that support re-ranking " + "must override this method." + ) diff --git a/DashAI/back/models/RAG/retrievers/composite/mmr_reranker_retriever.py b/DashAI/back/models/RAG/retrievers/composite/mmr_reranker_retriever.py new file mode 100644 index 000000000..c43ac886c --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/composite/mmr_reranker_retriever.py @@ -0,0 +1,208 @@ +import logging +from typing import List, Tuple + +import numpy as np +from sklearn.metrics.pairwise import cosine_similarity + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + float_field, + int_field, + list_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.documents import Chunk +from DashAI.back.models.RAG.retrievers.composite.composite_retriever import ( + CompositeRetriever, +) + +log = logging.getLogger(__name__) + + +class MMRRerankerRetrieverSchema(BaseSchema): + """Schema for :class:`MMRRerankerRetriever`. + + Attributes: + mmr_lambda: Trade-off between relevance and diversity. + retrieval_factor: Multiplier for initial retrieval size. + top_k: Final number of chunks to select. + children: Exactly one child retriever whose results are re-ranked. + """ + + mmr_lambda: schema_field( + float_field(ge=0.0, le=1.0), + placeholder=0.5, + description=MultilingualString( + en="Trade-off between relevance (1.0) and diversity (0.0).", + es="Compromiso entre relevancia (1.0) y diversidad (0.0).", + ), + ) # type: ignore + + retrieval_factor: schema_field( + int_field(gt=1), + placeholder=3, + description=MultilingualString( + en="Multiplier for initial retrieval size.", + es="Multiplicador para el tamaño de recuperación inicial.", + ), + ) # type: ignore + + top_k: schema_field( + int_field(gt=0), + placeholder=5, + description=MultilingualString( + en="Final number of chunks to select.", + es="Número final de fragmentos a seleccionar.", + ), + ) # type: ignore + + children: schema_field( + list_field(component_field(parent="RetrieverModel"), min_items=1, max_items=1), + placeholder=[], + description=MultilingualString( + en="The child retriever whose results will be re-ranked.", + es="El recuperador hijo cuyos resultados serán reordenados.", + ), + ) # type: ignore + + +class MMRRerankerRetriever(CompositeRetriever): + """Re-ranks retrieval results using Maximum Marginal Relevance for diversity. + + Retrieves an expanded candidate set from the child, then selects a + diverse subset by balancing relevance and similarity among selected + chunks. + """ + + FLAGS: list[str] = ["FAMILY:mmr", "composite", "reranker"] + SCHEMA = MMRRerankerRetrieverSchema + DISPLAY_NAME: str = MultilingualString( + en="MMR Reranker", + es="Reordenador MMR", + ) + DESCRIPTION: str = MultilingualString( + en="Re-ranks retrieval results using Maximum Marginal Relevance for diversity.", + es="Reordena resultados de recuperación usando Maximum Marginal Relevance " + "para diversidad.", + ) + + def __init__(self, **kwargs): + """Initialize the MMR reranker. + + Args: + **kwargs: Must contain ``mmr_lambda``, ``retrieval_factor``, + ``top_k``, and ``children`` (exactly one child). + """ + super().__init__(**kwargs) + self.mmr_lambda = self.params.pop("mmr_lambda") + self.retrieval_factor = self.params.pop("retrieval_factor") + self._top_k = self.params.pop("top_k") + + @property + def retrieval_top_k(self) -> int: + """Return the final number of chunks selected. + + Returns: + The value of ``top_k``. + """ + return self._top_k + + def retrieve(self, query: str, **kwargs) -> List[Chunk]: + """Retrieve and re-rank chunks using MMR diversity. + + Fetches an expanded candidate set from the child, computes + pairwise cosine similarity among the candidates, then selects + a diverse subset via the MMR algorithm. + + Args: + query: The search query string. + **kwargs: Additional retrieval parameters. + + Returns: + A list of :class:`Chunk` instances selected for relevance + and diversity. + """ + child = self._children[0] + expanded_k = self._top_k * self.retrieval_factor + candidates = child.retrieve(query, top_k=expanded_k) + if len(candidates) <= self._top_k: + return candidates + + chunk_ids: List[int] = [c.id for c in candidates] + scored: List[Tuple[int, float]] = child.score_chunks(chunk_ids, query) + id_to_dist = dict(scored) + + ordered = [c for c in candidates if c.id in id_to_dist] + if len(ordered) <= self._top_k: + return ordered[: self._top_k] + + relevance = np.array([1.0 - id_to_dist[c.id] for c in ordered]) + ordered_ids: List[int] = [c.id for c in ordered] + try: + vectors = child.get_chunk_vectors(ordered_ids) + except ValueError: + return ordered[: self._top_k] + + if len(vectors) != len(ordered): + log.warning( + "MMRRerankerRetriever: vector count (%d) != ordered chunk count (%d). " + "Falling back to top-k without MMR.", + len(vectors), + len(ordered), + ) + return ordered[: self._top_k] + + pairwise = cosine_similarity(vectors) + selected = self._mmr_select(relevance, pairwise) + return [ordered[i] for i in selected] + + def score_chunks(self, chunk_ids, query): + """Score chunks by delegating to the child retriever. + + Args: + chunk_ids: List of chunk IDs to score. + query: The search query string. + + Returns: + A list of ``(chunk_id, distance)`` tuples from the child. + """ + child = self._children[0] + return child.score_chunks(chunk_ids, query) + + def _mmr_select( + self, + relevance: np.ndarray, + pairwise: np.ndarray, + ) -> List[int]: + """Select indices using Maximum Marginal Relevance. + + Greedily picks the next candidate that maximises: + ``lambda * relevance(c) - (1 - lambda) * max_similarity(c, selected)`` + + Args: + relevance: Relevance scores for all candidates (shape ``(n,)``). + pairwise: Pairwise cosine similarity matrix (shape ``(n, n)``). + + Returns: + List of selected indices (up to ``self._top_k``). + """ + n = len(relevance) + best = int(np.argmax(relevance)) + selected = [best] + remaining = [i for i in range(n) if i != best] + + while len(selected) < self._top_k and remaining: + scores = np.empty(len(remaining)) + for j, cand in enumerate(remaining): + max_pair = max(pairwise[cand][s] for s in selected) + scores[j] = ( + self.mmr_lambda * relevance[cand] + - (1.0 - self.mmr_lambda) * max_pair + ) + best_idx = int(np.argmax(scores)) + selected.append(remaining[best_idx]) + remaining.pop(best_idx) + + return selected diff --git a/DashAI/back/models/RAG/retrievers/composite/parallel_retriever.py b/DashAI/back/models/RAG/retrievers/composite/parallel_retriever.py new file mode 100644 index 000000000..84d3b3dc7 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/composite/parallel_retriever.py @@ -0,0 +1,214 @@ +from typing import List, Tuple + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + enum_field, + list_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.documents import Chunk +from DashAI.back.models.RAG.retrievers.composite.composite_retriever import ( + CompositeRetriever, +) +from DashAI.back.models.RAG.retrievers.enums import MergeStrategy + + +class ParallelRetrieverSchema(BaseSchema): + """Schema for :class:`ParallelRetriever`. + + Attributes: + children: List of at least 2 child retrievers queried in parallel. + merge_strategy: Strategy for merging results + (:attr:`MergeStrategy` values). + """ + + children: schema_field( + list_field(component_field(parent="RetrieverModel"), min_items=2), + placeholder=[], + description=MultilingualString( + en="List of child retrievers queried in parallel.", + es="Lista de recuperadores hijos consultados en paralelo.", + ), + ) # type: ignore + + merge_strategy: schema_field( + enum_field(enum=[s.value for s in MergeStrategy]), + placeholder=MergeStrategy.ROUND_ROBIN.value, + description=MultilingualString( + en=( + f"'{MergeStrategy.ROUND_ROBIN.value}': alternates results" + f" from each retriever. " + f"'{MergeStrategy.INTERLEAVE.value}': merges preserving" + f" internal order." + ), + es=( + f"'{MergeStrategy.ROUND_ROBIN.value}': alterna resultados" + f" de cada recuperador. " + f"'{MergeStrategy.INTERLEAVE.value}': fusiona preservando" + f" el orden interno." + ), + ), + ) # type: ignore + + +class ParallelRetriever(CompositeRetriever): + """Queries multiple retrievers in parallel and merges their results. + + Results are deduplicated by ``(document_id, document_position)``. + """ + + FLAGS: list[str] = ["composite", "parallel"] + SCHEMA = ParallelRetrieverSchema + DISPLAY_NAME: str = MultilingualString( + en="Parallel Retriever", + es="Recuperador Paralelo", + ) + DESCRIPTION: str = MultilingualString( + en="Queries multiple retrievers in parallel and merges their results.", + es="Consulta múltiples recuperadores en paralelo y fusiona sus resultados.", + ) + + def __init__(self, **kwargs): + """Initialize the parallel retriever with a merge strategy. + + Args: + **kwargs: Must contain ``children`` and ``merge_strategy`` + keys. + """ + super().__init__(**kwargs) + self.merge_strategy = MergeStrategy(self.params.pop("merge_strategy")) + + @property + def retrieval_top_k(self) -> int: + """Return the sum of all children's top_k. + + Returns: + Total number of chunks the parallel retriever may return. + """ + return sum(c.retrieval_top_k for c in self._children) + + def retrieve(self, query, **kwargs) -> List[Chunk]: + """Retrieve chunks from all children in parallel and merge. + + Args: + query: The search query string. + **kwargs: Additional retrieval parameters forwarded to each + child. + + Returns: + A deduplicated, merged list of :class:`Chunk` instances. + """ + all_child_results = [] + for child in self._children: + all_child_results.append(child.retrieve(query, **kwargs)) + + total_k = sum(c.retrieval_top_k for c in self._children) + + if self.merge_strategy == MergeStrategy.ROUND_ROBIN: + return self._merge_round_robin(all_child_results, total_k) + return self._merge_interleave(all_child_results, total_k) + + def score_chunks(self, chunk_ids: List[int], query: str) -> List[Tuple[int, float]]: + """Aggregate chunk scores across all children. + + Each child scores the chunks independently, then all scores + for the same chunk are averaged. Results sorted by distance + (ascending = more relevant). Chunks that no child scored + are omitted. + + Args: + chunk_ids: List of chunk IDs to score. + query: The search query string. + + Returns: + A list of ``(chunk_id, average_distance)`` tuples sorted by + distance. + """ + if not chunk_ids or not self._children: + return [] + + accumulator = {cid: [] for cid in chunk_ids} + for child in self._children: + for cid, dist in child.score_chunks(chunk_ids, query): + accumulator[cid].append(dist) + + merged = [] + for cid, dists in accumulator.items(): + if dists: + merged.append((cid, sum(dists) / len(dists))) + merged.sort(key=lambda pair: pair[1]) + return merged + + @staticmethod + def _chunk_key(chunk: Chunk) -> Tuple[int, int]: + """Build a hashable dedup key from document_id and position. + + Both values are ints — see Chunk model definition. + + Args: + chunk: The chunk to generate a key for. + + Returns: + A ``(document_id, document_position)`` tuple. + """ + return (chunk.document_id, chunk.document_position) + + def _merge_round_robin( + self, child_results_list: List[List[Chunk]], total_k: int + ) -> List[Chunk]: + """Merge results using round-robin alternation. + + Takes turns picking from each child's result list, skipping + duplicates. + + Args: + child_results_list: Results from each child retriever. + total_k: Maximum number of chunks to return. + + Returns: + A deduplicated list of up to *total_k* chunks. + """ + results = [] + seen_keys: set[Tuple[int, int]] = set() + max_len = max(len(r) for r in child_results_list) if child_results_list else 0 + + for i in range(max_len): + for child_results in child_results_list: + if i < len(child_results): + chunk = child_results[i] + key = self._chunk_key(chunk) + if key not in seen_keys: + seen_keys.add(key) + results.append(chunk) + if len(results) >= total_k: + return results + return results + + def _merge_interleave( + self, child_results_list: List[List[Chunk]], total_k: int + ) -> List[Chunk]: + """Merge results by interleaving in child order. + + Concatenates each child's list in order, skipping duplicates. + + Args: + child_results_list: Results from each child retriever. + total_k: Maximum number of chunks to return. + + Returns: + A deduplicated list of up to *total_k* chunks. + """ + results = [] + seen_keys: set[Tuple[int, int]] = set() + + for child_results in child_results_list: + for chunk in child_results: + key = self._chunk_key(chunk) + if key not in seen_keys: + seen_keys.add(key) + results.append(chunk) + if len(results) >= total_k: + return results + return results diff --git a/DashAI/back/models/RAG/retrievers/composite/sequential_retriever.py b/DashAI/back/models/RAG/retrievers/composite/sequential_retriever.py new file mode 100644 index 000000000..a1985081e --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/composite/sequential_retriever.py @@ -0,0 +1,159 @@ +from typing import List, Tuple + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + list_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.documents import Chunk +from DashAI.back.models.RAG.exceptions import RAGRetrieverError +from DashAI.back.models.RAG.retrievers.composite.composite_retriever import ( + CompositeRetriever, +) +from DashAI.back.models.RAG.retrievers.exceptions import CompositeValidationError + + +class SequentialRetrieverSchema(BaseSchema): + """Schema for :class:`SequentialRetriever`. + + Attributes: + children: Ordered list of at least 2 child retrievers. + """ + + children: schema_field( + list_field(component_field(parent="RetrieverModel"), min_items=2), + placeholder=[], + description=MultilingualString( + en="Ordered list of child retrievers. The first child retrieves" + " broadly; each subsequent child re-ranks and tightens the" + " results.", + es="Lista ordenada de recuperadores hijos. El primero recupera" + " ampliamente; cada hijo subsiguiente reordena y ajusta los" + " resultados.", + ), + ) # type: ignore + + +class SequentialRetriever(CompositeRetriever): + """Queries multiple retrievers in sequence, re-ranking at each step. + + Each stage narrows the result set by requiring strictly decreasing + ``top_k`` values. The last child is the authoritative scorer. + """ + + FLAGS: list[str] = ["composite", "sequential"] + SCHEMA = SequentialRetrieverSchema + DISPLAY_NAME: str = MultilingualString( + en="Sequential Retriever", + es="Recuperador Secuencial", + ) + DESCRIPTION: str = MultilingualString( + en="Queries multiple retrievers in sequence, re-ranking at each step.", + es="Consulta múltiples recuperadores en secuencia, reordenando en cada paso.", + ) + + def __init__(self, **kwargs): + """Initialize and validate the sequential cascade. + + Args: + **kwargs: Must contain a ``children`` key with at least 2 + :class:`RetrieverModel` instances. + + Raises: + CompositeValidationError: If ``top_k`` values are not strictly + decreasing. + """ + super().__init__(**kwargs) + self._validate() + + def _validate(self) -> None: + """Validate cascade constraints on children. + + Checks that ``top_k`` values are strictly decreasing so each + stage narrows the result set. No type restrictions — any + ``RetrieverModel`` subclass is allowed as a child (Composite + pattern). + + Raises: + CompositeValidationError: If child ``top_k`` values are not + strictly decreasing. + """ + children = getattr(self, "_children", []) + if len(children) < 2: + return + + children_k = [c.retrieval_top_k for c in children] + for i in range(1, len(children_k)): + if children_k[i] >= children_k[i - 1]: + raise CompositeValidationError( + f"Cascade requires strictly decreasing top_k. " + f"Child {i} top_k={children_k[i]} >= " + f"child {i - 1} top_k={children_k[i - 1]}" + ) + + def retrieve(self, query, **kwargs) -> List[Chunk]: + """Retrieve chunks through the sequential cascade. + + The first child performs a broad retrieval; each subsequent child + re-ranks the results and narrows to its ``top_k``. + + Args: + query: The search query string. + **kwargs: Additional retrieval parameters forwarded to the + first child. + + Returns: + A list of :class:`Chunk` instances after all stages. + + Raises: + RAGRetrieverError: If any chunk has a ``None`` ID. + """ + first = self._children[0] + results = first.retrieve(query, **kwargs) + + for c in results: + if c.id is None: + raise RAGRetrieverError( + "Chunk with None id encountered in SequentialRetriever. " + "Chunks must be persisted before sequential retrieval." + ) + + for child in self._children[1:]: + chunk_ids = [c.id for c in results] + scored = child.score_chunks(chunk_ids, query) + scored = scored[: child.retrieval_top_k] + id_to_chunk = {c.id: c for c in results} + results = [id_to_chunk[cid] for cid, _ in scored] + + return results + + def score_chunks(self, chunk_ids: List[int], query: str) -> List[Tuple[int, float]]: + """Score chunks by delegating to the last child in the cascade. + + The last child produces the final ranking, so it is the + authoritative scorer. + + Args: + chunk_ids: List of chunk IDs to score. + query: The search query string. + + Returns: + A list of ``(chunk_id, distance)`` tuples sorted by distance. + """ + if not self._children: + return [] + return self._children[-1].score_chunks(chunk_ids, query) + + @property + def retrieval_top_k(self) -> int: + """Return the ``top_k`` of the last (most restrictive) child. + + Returns: + The ``top_k`` of the last child, or ``1`` if there are no + children. + """ + if self._children: + return self._children[-1].retrieval_top_k + return 1 diff --git a/DashAI/back/models/RAG/retrievers/dense/__init__.py b/DashAI/back/models/RAG/retrievers/dense/__init__.py new file mode 100644 index 000000000..e1c3b7e69 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/dense/__init__.py @@ -0,0 +1,9 @@ +"""Convenience re-exports for all dense retriever types.""" + +from DashAI.back.models.RAG.retrievers.dense.dense_embedding_retriever import ( + DenseEmbeddingRetriever, +) +from DashAI.back.models.RAG.retrievers.dense.dense_retriever import DenseRetriever +from DashAI.back.models.RAG.retrievers.dense.huggingface_dense_retriever import ( + HuggingFaceDenseRetriever, +) diff --git a/DashAI/back/models/RAG/retrievers/dense/_hf_language_utils.py b/DashAI/back/models/RAG/retrievers/dense/_hf_language_utils.py new file mode 100644 index 000000000..ed6710a76 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/dense/_hf_language_utils.py @@ -0,0 +1,123 @@ +from typing import Dict, List, Tuple + +LANGUAGE_LABELS: Dict[str, str] = { + "af": "Afr", + "ar": "Ara", + "bg": "Bul", + "bn": "Ben", + "ca": "Cat", + "cs": "Cze", + "da": "Dan", + "de": "Deu", + "el": "Gre", + "en": "Eng", + "es": "Esp", + "et": "Est", + "fa": "Per", + "fi": "Fin", + "fr": "Fra", + "he": "Heb", + "hi": "Hin", + "hr": "Cro", + "hu": "Hun", + "id": "Ind", + "it": "Ita", + "ja": "Jpn", + "ko": "Kor", + "lt": "Lit", + "lv": "Lav", + "mk": "Mac", + "ms": "May", + "multi": "Multi", + "nl": "Dut", + "no": "Nor", + "pl": "Pol", + "pt": "Por", + "ro": "Rom", + "ru": "Rus", + "sk": "Slo", + "sl": "Svn", + "sr": "Srp", + "sv": "Swe", + "th": "Tha", + "tr": "Tur", + "uk": "Ukr", + "ur": "Urd", + "vi": "Vie", + "zh": "Chi", +} + +MAX_DISPLAY_LANGUAGES: int = 3 + + +def compute_language_summary(languages: List[str]) -> Tuple[str, int]: + """Build a human-readable summary and total count for a list of languages. + + Shows up to ``MAX_DISPLAY_LANGUAGES`` (3) labels; overflow is + indicated with a ``+N`` suffix. + + Args: + languages: List of ISO 639-1 language codes. + + Returns: + A ``(summary_string, total_count)`` tuple. + """ + if not languages: + return "", 0 + labels = [LANGUAGE_LABELS.get(lang, lang.title()) for lang in languages] + if len(labels) <= MAX_DISPLAY_LANGUAGES: + return ", ".join(labels), len(languages) + shown = labels[:MAX_DISPLAY_LANGUAGES] + overflow = len(languages) - MAX_DISPLAY_LANGUAGES + return f"{', '.join(shown)} +{overflow}", len(languages) + + +def build_model_language_summaries( + models: Dict[str, dict], +) -> Dict[str, Dict[str, object]]: + """Build per-model language summaries. + + Args: + models: Mapping from model name to model info dicts (each + containing a ``"languages"`` key). + + Returns: + A dict mapping model name to ``{"summary": str, "count": int, + "labels": list[str]}``. + """ + result: Dict[str, Dict[str, object]] = {} + for model_name, info in models.items(): + summary, count = compute_language_summary(info.get("languages", [])) + result[model_name] = { + "summary": summary, + "count": count, + "labels": [ + LANGUAGE_LABELS.get(lang, lang.title()) + for lang in info.get("languages", []) + ], + } + return result + + +def build_family_language_summary( + models: Dict[str, dict], +) -> Tuple[str, int]: + """Build a language summary across all models in a family. + + Collects unique languages from all models and summarises them. + + Args: + models: Mapping from model name to model info dicts (each + containing a ``"languages"`` key). + + Returns: + A ``(summary_string, total_unique_count)`` tuple. + """ + all_languages: List[str] = [] + seen: set[str] = set() + for info in models.values(): + for lang in info.get("languages", []): + if lang not in seen: + seen.add(lang) + all_languages.append(lang) + return compute_language_summary(all_languages) diff --git a/DashAI/back/models/RAG/retrievers/dense/_hf_metadata_utils.py b/DashAI/back/models/RAG/retrievers/dense/_hf_metadata_utils.py new file mode 100644 index 000000000..160048d52 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/dense/_hf_metadata_utils.py @@ -0,0 +1,34 @@ +from typing import Dict + +from DashAI.back.models.RAG.retrievers.dense._hf_language_utils import ( + build_family_language_summary, + build_model_language_summaries, +) + + +def build_retriever_metadata( + models: Dict[str, dict], + family_name: str, + model_count: int, +) -> Dict[str, object]: + """Build a metadata dictionary describing a retriever family. + + Args: + models: Mapping from model name to model info dicts (each + containing a ``"languages"`` key). + family_name: Human-readable name for the model family. + model_count: Total number of models in the family. + + Returns: + A dictionary with keys ``family``, ``language_summary``, + ``language_count``, ``model_count``, and ``model_languages``. + """ + family_summary, family_count = build_family_language_summary(models) + model_languages = build_model_language_summaries(models) + return { + "family": family_name, + "language_summary": family_summary, + "language_count": family_count, + "model_count": model_count, + "model_languages": model_languages, + } diff --git a/DashAI/back/models/RAG/retrievers/dense/dense_embedding_retriever.py b/DashAI/back/models/RAG/retrievers/dense/dense_embedding_retriever.py new file mode 100644 index 000000000..251cde18f --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/dense/dense_embedding_retriever.py @@ -0,0 +1,96 @@ +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + enum_field, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.retrievers.dense.dense_retriever import DenseRetriever + +METRICS = ["cityblock", "cosine", "euclidean", "l1", "l2", "manhattan"] + + +class DenseEmbeddingRetrieverSchema(BaseSchema): + """Schema for :class:`DenseEmbeddingRetriever`. + + Attributes: + embedding_model: The embedding model component to use. + similarity_metric: Distance metric for vector comparison. + top_k: Number of chunks to select. + """ + + embedding_model: schema_field( + component_field(parent="DenseEmbedding"), + placeholder={"component": "SentenceTransformerEmbedding", "params": {}}, + description=MultilingualString( + en="Embedding model to use for encoding chunks.", + es="Modelo de embedding a usar para codificar fragmentos.", + ), + ) # type: ignore + + similarity_metric: schema_field( + enum_field(enum=METRICS), + placeholder="cosine", + description=MultilingualString( + en="Distance metric for comparing dense vectors.", + es="Métrica de distancia para comparar vectores densos.", + ), + ) # type: ignore + + top_k: schema_field( + int_field(gt=0), + placeholder=5, + description=MultilingualString( + en="Number of chunks to select.", + es="Número de fragmentos a seleccionar.", + ), + ) # type: ignore + + +class DenseEmbeddingRetriever(DenseRetriever): + """Concrete dense retriever that accepts any :class:`DenseEmbedding` component. + + The embedding component is specified in the schema and instantiated + by the factory via ``fill_objects``. + """ + + FLAGS: list[str] = ["dense", "dense_embedding"] + SCHEMA = DenseEmbeddingRetrieverSchema + + DISPLAY_NAME: str = MultilingualString( + en="Dense Embedding Retriever", + es="Recuperador por Embeddings Densos", + ) + DESCRIPTION: str = MultilingualString( + en="Dense retriever using any registered DenseEmbedding for similarity search.", + es="Recuperador denso que usa cualquier DenseEmbedding registrado" + " para búsqueda por similitud.", + ) + + def __init__(self, **kwargs): + """Initialize the dense embedding retriever. + + Pops the ``embedding_model`` instance from kwargs and stores it. + + Args: + **kwargs: Must contain ``embedding_model``, + ``similarity_metric``, and ``top_k``. + """ + embedding_instance = kwargs.pop("embedding_model") + super().__init__(**kwargs) + self.params["embedding_model"] = { + "component": embedding_instance.__class__.__name__, + "params": dict(sorted(embedding_instance.params.items())), + } + self._embedding_instance = embedding_instance + + def init_model(self) -> None: + """Load the embedding model, then initialise the similarity matrix. + + The embedding instance is created by :meth:`fill_objects` during + factory construction but its heavy resources (tokenizer, weights) + are only acquired on the explicit ``load()`` call. + """ + self._embedding_instance.load() + self._init_embedding(self._embedding_instance) diff --git a/DashAI/back/models/RAG/retrievers/dense/dense_retriever.py b/DashAI/back/models/RAG/retrievers/dense/dense_retriever.py new file mode 100644 index 000000000..89fa10bb4 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/dense/dense_retriever.py @@ -0,0 +1,255 @@ +import os +from typing import Dict, List, Tuple + +import numpy as np +from sklearn.metrics.pairwise import pairwise_distances + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.documents import Chunk +from DashAI.back.models.RAG.embeddings import DenseEmbedding +from DashAI.back.models.RAG.exceptions import RAGRetrieverError +from DashAI.back.models.RAG.retrievers.unit_retriever import UnitRetriever + +# NOTE: The entire chunk index (similarity_matrix) is loaded into memory, +# which is fine for typical use with tens to low hundreds of documents but +# may be a bottleneck for very large collections. + + +class DenseRetrieverSchema(BaseSchema): + """Schema for :class:`DenseRetriever`. + + Attributes: + similarity_metric: Distance metric for vector comparison. + top_k: Number of chunks to select. + """ + + similarity_metric: schema_field( + enum_field(enum=["cityblock", "cosine", "euclidean", "l1", "l2", "manhattan"]), + placeholder="cosine", + description=MultilingualString( + en="Distance metric for comparing dense vectors.", + es="Métrica de distancia para comparar vectores densos.", + ), + ) # type: ignore + + top_k: schema_field( + int_field(gt=0), + placeholder=5, + description=MultilingualString( + en="Number of chunks to select.", + es="Número de fragmentos a seleccionar.", + ), + ) # type: ignore + + +class DenseRetriever(UnitRetriever): + """Abstract base for dense (embedding-based) retrievers. + + Maintains a similarity matrix of chunk embeddings loaded from + on-disk ``.npy`` files and retrieves via pairwise distance. + """ + + FLAGS: list[str] = ["abstract"] + DISPLAY_NAME: str = MultilingualString( + en="Embedding Retriever", + es="Recuperador por Embeddings", + ) + DESCRIPTION: str = MultilingualString( + en="Embedding retriever using vector embeddings for similarity search.", + es="Recuperador por embeddings que usa representaciones" + " vectoriales para búsqueda por similitud.", + ) + + SCHEMA = DenseRetrieverSchema + + def __init__(self, **kwargs): + """Initialize the dense retriever. + + Args: + **kwargs: Must contain ``similarity_metric`` and ``top_k``. + """ + super().__init__(**kwargs) + + self.similarity_metric = self.params.pop("similarity_metric") + self._top_k = self.params.pop("top_k") + + def _init_embedding(self, embedding_model): + """Initialise the embedding model and build the similarity matrix. + + Args: + embedding_model: A :class:`DenseEmbedding` instance used to + encode chunks. + + Raises: + TypeError: If *embedding_model* is not a ``DenseEmbedding``. + """ + if not isinstance(embedding_model, DenseEmbedding): + raise TypeError( + f"Expected DenseEmbedding instance, " + f"got {type(embedding_model).__name__}" + ) + self.embedding_model = embedding_model + self.compute_missing_embeddings() + self.init_similarity_matrix() + + def compute_missing_embeddings(self): + """Compute and persist embeddings for chunks that lack them. + + Iterates over all documents; if an ``embeddings.npy`` file does + not yet exist at the expected path, the embedding model is used + to encode the chunk texts and the result is saved. + """ + for doc_id, doc_chunks in self.chunks.items(): + matrix_dir = self._persistence.matrix_dirs.get(doc_id) + if matrix_dir is None: + continue + matrix_path = os.path.join(matrix_dir, "embeddings.npy") + if os.path.exists(matrix_path): + continue + chunk_texts = [chunk.text for chunk in doc_chunks.values()] + if not chunk_texts: + raise RAGRetrieverError(f"No chunks found for document ID {doc_id}.") + embeddings = self.embedding_model.batch_encode(chunk_texts) + os.makedirs(matrix_dir, exist_ok=True) + np.save(matrix_path, embeddings) + + def init_similarity_matrix(self): + """Load all persisted embedding matrices into a single similarity matrix. + + Builds ``similarity_matrix`` (a vertical stack of all per-document + embedding arrays), ``matrix_row_to_chunk_id``, and + ``chunk_id_to_doc_id`` lookup mappings. + """ + self.similarity_matrix = None + self.matrix_row_to_chunk_id = {} + self.chunk_id_to_doc_id = {} + + all_embeddings = [] + row_index = 0 + for doc_id, chunks in self.chunks.items(): + matrix_dir = self._persistence.matrix_dirs.get(doc_id) + if matrix_dir is None: + continue + matrix_path = os.path.join(matrix_dir, "embeddings.npy") + if not os.path.exists(matrix_path): + continue + for chunk_id in chunks: + self.matrix_row_to_chunk_id[row_index] = chunk_id + self.chunk_id_to_doc_id[chunk_id] = doc_id + row_index += 1 + all_embeddings.append(np.load(matrix_path)) + + if all_embeddings: + self.similarity_matrix = np.vstack(all_embeddings) + + @property + def retrieval_top_k(self) -> int: + """Return the configured top-k value. + + Returns: + Number of chunks to retrieve. + """ + return self._top_k + + def retrieve(self, query: str, top_k: int | None = None) -> List[Chunk]: + """Retrieve the top-k chunks by embedding similarity. + + Args: + query: The search query string. + top_k: Override for the default ``top_k``. Uses the + configured value if ``None``. + + Returns: + A list of :class:`Chunk` instances ordered by distance. + + Raises: + ValueError: If the similarity matrix has not been + initialised. + """ + self._check_infra() + if self.similarity_matrix is None: + raise ValueError("Similarity matrix not initialized.") + k = top_k if top_k is not None else self._top_k + query_embedding = self.embedding_model.encode(query) + distances = pairwise_distances( + query_embedding.reshape(1, -1), + self.similarity_matrix, + metric=self.similarity_metric, + )[0] + top_indices = np.argsort(distances)[:k] + results = [] + for idx in top_indices: + chunk_id = self.matrix_row_to_chunk_id[idx] + doc_id = self.chunk_id_to_doc_id[chunk_id] + results.append(self.chunks[doc_id][chunk_id]) + return results + + def score_chunks(self, chunk_ids: List[int], query: str) -> List[Tuple[int, float]]: + """Score a set of chunk IDs against the query. + + Args: + chunk_ids: List of chunk IDs to score. + query: The search query string. + + Returns: + A list of ``(chunk_id, distance)`` tuples sorted by distance. + + Raises: + ValueError: If the similarity matrix has not been + initialised. + """ + self._check_infra() + if self.similarity_matrix is None: + raise ValueError("Similarity matrix not initialized.") + query_embedding = self.embedding_model.encode(query) + chunk_id_to_row = { + self.matrix_row_to_chunk_id[r]: r for r in self.matrix_row_to_chunk_id + } + valid_ids = [cid for cid in chunk_ids if cid in chunk_id_to_row] + if not valid_ids: + return [] + rows = [chunk_id_to_row[cid] for cid in valid_ids] + distances = pairwise_distances( + query_embedding.reshape(1, -1), + self.similarity_matrix[rows], + metric=self.similarity_metric, + )[0] + scored = list(zip(valid_ids, distances.tolist(), strict=True)) + scored.sort(key=lambda x: x[1]) + return scored + + def get_chunk_vectors(self, chunk_ids: List[int]) -> np.ndarray: + """Return embedding vectors for the given chunk IDs. + + Args: + chunk_ids: List of chunk IDs whose vectors are needed. + + Returns: + A 2D numpy array of embedding vectors. + + Raises: + ValueError: If the similarity matrix has not been + initialised, or none of the chunk IDs are found. + """ + if self.similarity_matrix is None: + raise ValueError("Similarity matrix not initialized.") + chunk_id_to_row: Dict[int, int] = { + self.matrix_row_to_chunk_id[r]: r for r in self.matrix_row_to_chunk_id + } + rows = [] + for cid in chunk_ids: + row = chunk_id_to_row.get(cid) + if row is not None: + rows.append(row) + if not rows: + raise ValueError( + f"None of the provided chunk_ids {chunk_ids} were found " + "in the similarity matrix." + ) + return self.similarity_matrix[rows] diff --git a/DashAI/back/models/RAG/retrievers/dense/huggingface_dense_retriever.py b/DashAI/back/models/RAG/retrievers/dense/huggingface_dense_retriever.py new file mode 100644 index 000000000..9876d18a7 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/dense/huggingface_dense_retriever.py @@ -0,0 +1,56 @@ +from abc import abstractmethod + +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.embeddings.dense.huggingface_embedding import ( + HuggingFaceEmbedding, +) +from DashAI.back.models.RAG.retrievers.dense.dense_retriever import DenseRetriever + +METRICS = ["cityblock", "cosine", "euclidean", "l1", "l2", "manhattan"] + + +class HuggingFaceDenseRetriever(DenseRetriever): + """Abstract dense retriever that creates a HuggingFace embedding model. + + Subclasses must implement :meth:`_create_embedding` to return a + :class:`HuggingFaceEmbedding` instance with the desired model + configuration. + """ + + FLAGS: list[str] = ["abstract", "huggingface"] + DISPLAY_NAME: str = MultilingualString( + en="HuggingFace Embedding Retriever", + es="Recuperador por Embeddings HuggingFace", + ) + DESCRIPTION: str = MultilingualString( + en="Dense retriever using HuggingFace embeddings for similarity search.", + es="Recuperador denso que usa embeddings HuggingFace para" + " búsqueda por similitud.", + ) + + def __init__(self, **kwargs): + """Initialize the HuggingFace dense retriever. + + Args: + **kwargs: Forwarded to :class:`DenseRetriever`. + """ + super().__init__(**kwargs) + + @abstractmethod + def _create_embedding(self) -> HuggingFaceEmbedding: + """Create and return a HuggingFace embedding model. + + Returns: + A :class:`HuggingFaceEmbedding` instance. + """ + raise NotImplementedError + + def init_model(self) -> None: + """Create, load, and initialise the HuggingFace embedding. + + Calls ``_create_embedding()``, loads its resources, then + passes it to ``_init_embedding()``. + """ + embedding = self._create_embedding() + embedding.load() + self._init_embedding(embedding) diff --git a/DashAI/back/models/RAG/retrievers/enums.py b/DashAI/back/models/RAG/retrievers/enums.py new file mode 100644 index 000000000..77322e247 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/enums.py @@ -0,0 +1,15 @@ +"""Enum definitions for retriever configuration.""" + +from enum import Enum + + +class MergeStrategy(str, Enum): + """Strategies for merging results from parallel retrievers. + + Attributes: + ROUND_ROBIN: Alternates results from each child retriever. + INTERLEAVE: Concatenates child results preserving internal order. + """ + + ROUND_ROBIN = "round_robin" + INTERLEAVE = "interleave" diff --git a/DashAI/back/models/RAG/retrievers/exceptions.py b/DashAI/back/models/RAG/retrievers/exceptions.py new file mode 100644 index 000000000..535a6203d --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/exceptions.py @@ -0,0 +1,23 @@ +"""Retriever exception classes — re-exported from the unified RAG hierarchy. + +All RAG exception types are defined in +:mod:`DashAI.back.models.RAG.exceptions` to avoid circular imports +and provide a single source of truth. This module re-exports the +retriever-specific subset for convenience. +""" + +from DashAI.back.models.RAG.exceptions import ( + RAGRetrieverCompositeValidationError as CompositeValidationError, +) +from DashAI.back.models.RAG.exceptions import ( + RAGRetrieverError as RetrieverError, +) +from DashAI.back.models.RAG.exceptions import ( + RAGRetrieverMissingParameterError as MissingParameterError, +) + +__all__ = [ + "CompositeValidationError", + "MissingParameterError", + "RetrieverError", +] diff --git a/DashAI/back/models/RAG/retrievers/persistence.py b/DashAI/back/models/RAG/retrievers/persistence.py new file mode 100644 index 000000000..e50585569 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/persistence.py @@ -0,0 +1,29 @@ +from dataclasses import dataclass + + +@dataclass +class SparsePersistence: + """Reference to a sparse retriever's on-disk dump. + + Attributes: + model_dir: Absolute path to the directory containing ``.pkl`` + files. If ``None``, no previous dump exists and the retriever + must train from scratch. + """ + + model_dir: str | None + + +@dataclass +class DensePersistence: + """References to embedding matrices for a dense retriever. + + Attributes: + matrix_dirs: Maps ``document_id`` to the absolute path of the + directory containing ``embeddings.npy`` for that document. + embedding_model_id: Identifier of the embedding model that + produced the vectors. + """ + + matrix_dirs: dict[int, str] + embedding_model_id: int | None diff --git a/DashAI/back/models/RAG/retrievers/retriever_factory.py b/DashAI/back/models/RAG/retrievers/retriever_factory.py new file mode 100644 index 000000000..f4e81b5b4 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/retriever_factory.py @@ -0,0 +1,154 @@ +"""Pure factory for retriever instances. + +Builds retriever models from the component registry. +Phase 1 only -- no I/O (embeddings, similarity matrices). +""" + +from dataclasses import dataclass +from typing import Any + +from DashAI.back.core.schema_fields.utils import fill_objects, normalize_payload +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.RAG.documents.chunk import Chunk +from DashAI.back.models.RAG.exceptions import ( + RAGComponentNotFoundError, + RAGRetrieverError, + RAGRetrieverMissingParameterError, +) +from DashAI.back.models.RAG.retrievers.composite.composite_retriever import ( + CompositeRetriever, +) +from DashAI.back.models.RAG.retrievers.dense.dense_retriever import DenseRetriever +from DashAI.back.models.RAG.retrievers.persistence import ( + DensePersistence, + SparsePersistence, +) +from DashAI.back.models.RAG.retrievers.retriever_model import RetrieverModel +from DashAI.back.models.RAG.retrievers.sparse.sparse_retriever import SparseRetriever + + +@dataclass(frozen=True) +class RetrieverFactoryResult: + """Result of building a retriever via :class:`RetrieverFactory`. + + Attributes: + model: The fully constructed :class:`RetrieverModel` instance. + """ + + model: RetrieverModel + + +class RetrieverFactory: + """Creates retriever instances. Phase 1 only — no I/O. + + Builds retriever models from the component registry, handling both + unit (leaf) and composite retrievers. Schema validation and + infrastructure injection are performed during construction. + + Args: + registry: Application component registry for resolving + component name strings. + RAG_path: Root path for RAG data storage. + chunks: Nested mapping of document IDs to chunk IDs to + :class:`Chunk` instances. + """ + + def __init__( + self, + registry: ComponentRegistry, + RAG_path: str, # noqa: N803 + chunks: dict[int, dict[int, Chunk]], + ): + self._registry = registry + self._RAG_path = RAG_path + self._chunks = chunks + + def create( + self, + component_name: str, + params: dict[str, Any], + persistence: DensePersistence | SparsePersistence | None = None, + ) -> RetrieverFactoryResult: + params = normalize_payload(params) + try: + model_class = self._registry[component_name]["class"] + except KeyError as err: + raise RAGComponentNotFoundError( + f"Component '{component_name}' not found in registry" + ) from err + + if issubclass(model_class, CompositeRetriever): + return self._create_composite(model_class, params) + + return self._create_unit(model_class, params, persistence) + + def _create_composite(self, model_class, params) -> RetrieverFactoryResult: + """Build a composite retriever by recursively creating children. + + Args: + model_class: The :class:`CompositeRetriever` subclass to + instantiate. + params: Configuration parameters; the ``children`` key is + consumed and replaced with built :class:`RetrieverModel` + instances. + + Returns: + A :class:`RetrieverFactoryResult` wrapping the new composite. + """ + children_configs = params.pop("children", []) + children_instances = [ + self.create(c["component"], c["params"]).model for c in children_configs + ] + params["children"] = children_instances + model = model_class(**params) + return RetrieverFactoryResult(model=model) + + def _create_unit( + self, + model_class, + params, + persistence: DensePersistence | SparsePersistence | None = None, + ) -> RetrieverFactoryResult: + """Build a unit (leaf) retriever with schema validation and injection. + + Args: + model_class: The :class:`UnitRetriever` subclass to instantiate. + params: Configuration parameters validated against the + model's schema. + persistence: A :class:`DensePersistence` or + :class:`SparsePersistence` instance matching the + retriever type. + + Returns: + A :class:`RetrieverFactoryResult` wrapping the new unit retriever. + + Raises: + RAGRetrieverMissingParameterError: If the required persistence + type is not provided for the retriever class. + RAGRetrieverError: If the model class is neither a dense nor + sparse retriever. + """ + validated = model_class.SCHEMA.model_validate(params) + resolved = fill_objects(validated, self._registry) + model = model_class(**resolved) + + if issubclass(model_class, DenseRetriever): + if not isinstance(persistence, DensePersistence): + raise RAGRetrieverMissingParameterError( + "A DensePersistence is required for dense retrievers " + "but none was provided." + ) + model.inject_infra(self._RAG_path, self._chunks, persistence) + elif issubclass(model_class, SparseRetriever): + if not isinstance(persistence, SparsePersistence): + raise RAGRetrieverMissingParameterError( + "A SparsePersistence is required for sparse retrievers " + "but none was provided." + ) + model.inject_infra(self._RAG_path, self._chunks, persistence) + else: + raise RAGRetrieverError( + f"Unsupported unit retriever: {model_class.__name__}" + ) + + return RetrieverFactoryResult(model=model) diff --git a/DashAI/back/models/RAG/retrievers/retriever_model.py b/DashAI/back/models/RAG/retrievers/retriever_model.py new file mode 100644 index 000000000..e58d7a905 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/retriever_model.py @@ -0,0 +1,248 @@ +import os +from abc import ABC, abstractmethod +from typing import Any, Dict, Final, List, Tuple + +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.base_model import BaseModel +from DashAI.back.models.RAG.documents import Chunk +from DashAI.back.models.RAG.exceptions import RAGWorkflowError + + +class RetrieverModel(BaseModel, ABC): + """ + Component: abstract base class for all retriever models. + + Implements the Component role in the Composite design pattern (GoF). + """ + + TYPE: Final[str] = "RetrieverModel" + FLAGS: list[str] = ["abstract"] + DISPLAY_NAME: str = MultilingualString( + en="Retriever", + es="Recuperador", + ) + DESCRIPTION: str = MultilingualString( + en="Document retrieval component.", + es="Componente de recuperación de documentos.", + ) + COLOR: str = "#9C27B0" + ICON: str = "Search" + + env_RAG_path: str | os.PathLike | None # noqa: N815 + chunks: Dict[int, Dict[int, Chunk]] + params: Dict[str, Any] + + def __init__(self, **kwargs): + """Initialize the retriever model. + + Stores keyword arguments as ``self.params`` and initialises the + database ID to ``None``. + + Args: + **kwargs: Configuration parameters for the retriever. + """ + self._db_id: int | None = None + self.params = kwargs + + def get_id(self) -> int | None: + """Return the database ID of this retriever, or ``None``. + + Returns: + The database ID if already persisted, otherwise ``None``. + """ + return self._db_id + + def set_id(self, id: int) -> None: + """Assign a database ID to this retriever. + + The ID can only be set once; subsequent calls raise an error. + + Args: + id: The database ID to assign. + + Raises: + RAGWorkflowError: If an ID has already been assigned. + """ + if self._db_id is not None: + raise RAGWorkflowError( + f"ID is already set to {self._db_id}, cannot reassign to {id}." + ) + self._db_id = id + + def _validate_chunks_dict(self) -> None: + """Validate the structure of the ``chunks`` attribute. + + Ensures it is a nested dictionary of the form + ``{doc_id: {chunk_id: Chunk}}`` and that each chunk's + ``document_id`` matches its parent key. + + Raises: + ValueError: If the structure is invalid. + """ + if not isinstance(self.chunks, dict): + raise ValueError("Chunks must be a dictionary.") + for doc_id, doc_chunks in self.chunks.items(): + if not isinstance(doc_id, int): + raise ValueError(f"Document ID {doc_id} must be an integer.") + if not isinstance(doc_chunks, dict): + raise ValueError( + f"Chunks for document ID {doc_id} must be a dictionary." + ) + for chunk_id, chunk in doc_chunks.items(): + if not isinstance(chunk_id, int): + raise ValueError( + f"Chunk ID {chunk_id} in document ID {doc_id}" + f" must be an integer." + ) + if not isinstance(chunk, Chunk): + raise ValueError( + f"Chunk {chunk_id} in document ID {doc_id}" + f" must be an instance of Chunk." + ) + if chunk.document_id != doc_id: + raise ValueError( + f"Chunk {chunk_id} document_id {chunk.document_id}" + f" != doc ID {doc_id}." + ) + + def inject_infra( + self, + env_RAG_path: str | os.PathLike, # noqa: N803 + chunks: Dict[int, Dict[int, Chunk]], + persistence: Any, + ) -> None: + """Inject runtime infrastructure *after* schema validation. + + Subclasses **must** store the three arguments as instance + attributes. Raises ``TypeError`` if any argument has an + unexpected type. + + Subclasses define required params in their ``__init__`` and + pop them from ``self.params``. + + Args: + env_RAG_path: Root directory path for RAG data. + chunks: Nested dictionary mapping document IDs to chunk IDs + to :class:`Chunk` instances. + persistence: Persistence object for saving/loading state. + """ + self.env_RAG_path = env_RAG_path + self.chunks = chunks + self._persistence = persistence + self._validate_chunks_dict() + + @property + def persistence(self): + """Get the persistence object.""" + return self._persistence + + @persistence.setter + def persistence(self, value): + """Set the persistence object. + + Args: + value: The new persistence object. + """ + self._persistence = value + + def init_model(self) -> None: + """Called by the factory **after** ``inject_infra()``. + + Subclasses restore saved state (``load()``) or compute initial + state (``_fit()``, ``_init_embedding()``). Default is a no-op. + """ + + @abstractmethod + def retrieve(self, query, **kwargs) -> List[Chunk]: + """Retrieve the top-k chunks most relevant to the query. + + Args: + query: The search query string. + **kwargs: Additional retrieval parameters (e.g. ``top_k``). + + Returns: + A list of :class:`Chunk` instances ordered by relevance. + """ + raise NotImplementedError + + @abstractmethod + def score_chunks(self, chunk_ids: List[int], query: str) -> List[Tuple[int, float]]: + """Score a set of chunks against a query. + + Args: + chunk_ids: List of chunk IDs to score. + query: The search query string. + + Returns: + A list of ``(chunk_id, distance)`` tuples sorted by distance + (ascending — lower is more relevant). + """ + raise NotImplementedError + + @property + def retrieval_top_k(self) -> int: + """Return the maximum number of chunks this retriever returns. + + Returns: + The top-k value. + """ + raise NotImplementedError + + # ── Child management (Composite pattern) ──────────────────────── + + def add(self, child: "RetrieverModel") -> None: + """Add a child retriever (Composite pattern). + + Args: + child: The child retriever to add. + + Raises: + NotImplementedError: If the subclass does not support children. + """ + raise NotImplementedError + + def remove(self, child: "RetrieverModel") -> None: + """Remove a child retriever (Composite pattern). + + Args: + child: The child retriever to remove. + + Raises: + NotImplementedError: If the subclass does not support children. + """ + raise NotImplementedError + + def get_children(self) -> List["RetrieverModel"]: + """Return the list of child retrievers (Composite pattern). + + Returns: + A list of :class:`RetrieverModel` children. + + Raises: + NotImplementedError: If the subclass does not support children. + """ + raise NotImplementedError + + def save(self, filename: str = "") -> None: + """Persist the retriever's state to disk. + + Args: + filename: Optional filename override. Defaults to an empty + string (subclasses determine their own default path). + """ + + def load(self, filename: str = "") -> None: + """Restore the retriever's state from disk. + + Args: + filename: Optional filename override. Defaults to an empty + string (subclasses determine their own default path). + """ + + def train(self, **kwargs): + """Train the retriever on the injected chunks. + + Args: + **kwargs: Training parameters. Default is a no-op. + """ + return diff --git a/DashAI/back/models/RAG/retrievers/sparse/__init__.py b/DashAI/back/models/RAG/retrievers/sparse/__init__.py new file mode 100644 index 000000000..76796f1aa --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/sparse/__init__.py @@ -0,0 +1,5 @@ +"""Convenience re-exports for all sparse retriever types.""" + +from DashAI.back.models.RAG.retrievers.sparse.sparse_retriever import SparseRetriever +from DashAI.back.models.RAG.retrievers.sparse.tfidf_retriever import TFIDFRetriever, TFIDFVectorizerModel +from DashAI.back.models.RAG.retrievers.sparse.bm25_retriever import BM25Retriever, BM25VectorizerModel diff --git a/DashAI/back/models/RAG/retrievers/sparse/bm25_retriever.py b/DashAI/back/models/RAG/retrievers/sparse/bm25_retriever.py new file mode 100644 index 000000000..0a16c6f26 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/sparse/bm25_retriever.py @@ -0,0 +1,429 @@ +import logging +import os +import pickle +from typing import Dict, List, Tuple + +import numpy as np +from scipy.sparse import lil_matrix +from sklearn.feature_extraction.text import CountVectorizer +from sklearn.metrics.pairwise import pairwise_distances + +from DashAI.back.core.schema_fields import ( + BaseSchema, + bool_field, + component_field, + enum_field, + float_field, + int_field, + list_field, + none_type, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.base_model import BaseModel +from DashAI.back.models.RAG.documents import Chunk +from DashAI.back.models.RAG.retrievers.sparse.sparse_retriever import SparseRetriever + +log = logging.getLogger(__name__) + + +class BM25VectorizerSchema(BaseSchema): + """Schema for the BM25 vectorizer parameters. + + Attributes: + strip_accents: Whether to remove accents during preprocessing. + lowercase: Whether to convert all characters to lowercase. + stop_words: List of stop words (or ``None``). + max_df: Document frequency upper threshold. + min_df: Document frequency lower threshold. + max_features: Maximum number of features (or ``None``). + """ + + strip_accents: schema_field( + none_type(enum_field(enum=["ascii", "unicode"])), + placeholder=None, + description=MultilingualString( + en="Remove accents during preprocessing.", + es="Eliminar acentos durante el preprocesamiento.", + ), + ) # type: ignore + + lowercase: schema_field( + bool_field(), + placeholder=True, + description=MultilingualString( + en="Convert all characters to lowercase.", + es="Convertir todos los caracteres a minúsculas.", + ), + ) # type: ignore + + stop_words: schema_field( + none_type(list_field(string_field(), min_items=1)), + placeholder=None, + description=MultilingualString( + en="List of stop words.", + es="Lista de palabras vacías.", + ), + ) # type: ignore + + max_df: schema_field( + float_field(ge=0.0, le=1.0), + placeholder=1.0, + description=MultilingualString( + en="Ignore terms with document frequency above this threshold.", + es="Ignorar términos con frecuencia de documento superior a este umbral.", + ), + ) # type: ignore + + min_df: schema_field( + float_field(ge=0.0, le=1.0), + placeholder=0.0, + description=MultilingualString( + en="Ignore terms with document frequency below this threshold.", + es="Ignorar términos con frecuencia de documento inferior a este umbral.", + ), + ) # type: ignore + + max_features: schema_field( + none_type(int_field(ge=1)), + placeholder=None, + description=MultilingualString( + en="Maximum number of features.", + es="Número máximo de características.", + ), + ) # type: ignore + + +class BM25VectorizerModel(BaseModel): + """Model component that encapsulates a :class:`CountVectorizer` for BM25. + + The vectorizer provides term-frequency counts; the BM25 weighting + is applied by the parent :class:`BM25Retriever`. + """ + + DISPLAY_NAME: str = MultilingualString( + en="BM25 Vectorizer", + es="Vectorizador BM25", + ) + DESCRIPTION: str = MultilingualString( + en="Vectorizer for BM25Retriever using CountVectorizer with" + " BM25-specific parameters.", + es="Vectorizador para BM25Retriever usando CountVectorizer con" + " parámetros específicos de BM25.", + ) + + SCHEMA = BM25VectorizerSchema + + def __init__(self, **kwargs): + """Initialize and build the underlying ``CountVectorizer``. + + Args: + **kwargs: Parameters matching :class:`BM25VectorizerSchema`. + """ + validated = self.SCHEMA.model_validate(kwargs) + self.params = dict(validated) + self.vectorizer = CountVectorizer( + strip_accents=self.params.pop("strip_accents"), + lowercase=self.params.pop("lowercase"), + stop_words=self.params.pop("stop_words"), + max_df=self.params.pop("max_df"), + min_df=self.params.pop("min_df"), + max_features=self.params.pop("max_features"), + ) + + def load(self): + """No-op load (state managed by the parent retriever).""" + + def save(self): + """No-op save (state managed by the parent retriever).""" + + def train(self): + """No-op train (fitting is done by the parent retriever).""" + + +class BM25RetrieverSchema(BaseSchema): + """Schema for :class:`BM25Retriever`. + + Attributes: + BM25Vectorizer: Parameters for the CountVectorizer component. + k1: BM25 term frequency saturation parameter. + b: BM25 length normalisation parameter. + delta: BM25 IDF smoothing parameter. + similarity_function: Distance metric for vector comparison. + top_k: Number of chunks to select. + """ + + BM25Vectorizer: schema_field( + component_field(parent="BM25VectorizerModel"), + placeholder={"component": "BM25VectorizerModel", "params": {}}, + description=MultilingualString( + en="BM25 Vectorizer parameters.", + es="Parámetros del vectorizador BM25.", + ), + ) # type: ignore + + k1: schema_field( + float_field(ge=0.0), + placeholder=1.5, + description=MultilingualString( + en="BM25 k1 parameter: term frequency saturation.", + es="Parámetro k1 de BM25: saturación de frecuencia de término.", + ), + ) # type: ignore + + b: schema_field( + float_field(ge=0.0, le=1.0), + placeholder=0.75, + description=MultilingualString( + en="BM25 b parameter: length normalization.", + es="Parámetro b de BM25: normalización de longitud.", + ), + ) # type: ignore + + delta: schema_field( + float_field(ge=0.0), + placeholder=0.0, + description=MultilingualString( + en="BM25 delta parameter for IDF smoothing.", + es="Parámetro delta de BM25 para suavizado de IDF.", + ), + ) # type: ignore + + similarity_function: schema_field( + enum_field(["cityblock", "cosine", "euclidean", "l1", "l2", "manhattan"]), + placeholder="cosine", + description=MultilingualString( + en="Distance metric for comparing BM25-weighted vectors.", + es="Métrica de distancia para comparar vectores ponderados BM25.", + ), + ) # type: ignore + + top_k: schema_field( + int_field(ge=1), + placeholder=5, + description=MultilingualString( + en="Number of chunks to select.", + es="Número de fragmentos a seleccionar.", + ), + ) # type: ignore + + +class BM25Retriever(SparseRetriever): + """Sparse retriever using BM25 (Okapi) ranking for document retrieval. + + Computes BM25-weighted term-frequency vectors and retrieves via + pairwise distance. + """ + + FLAGS: list[str] = ["keyword", "sparse"] + DISPLAY_NAME: str = MultilingualString( + en="BM25 Retriever", + es="Recuperador BM25", + ) + DESCRIPTION: str = MultilingualString( + en="Sparse retriever using BM25 (Okapi) ranking for document retrieval.", + es="Recuperador disperso que usa ranking BM25 (Okapi) para" + " recuperar documentos.", + ) + + SCHEMA = BM25RetrieverSchema + + def __init__(self, **kwargs): + """Initialize the BM25 retriever. + + Args: + **kwargs: Must contain ``BM25Vectorizer``, ``k1``, ``b``, + ``delta``, ``similarity_function``, and ``top_k``. + """ + super().__init__(**kwargs) + + self.k1 = self.params.pop("k1") + self.b = self.params.pop("b") + self.delta = self.params.pop("delta") + self.similarity_function_name = self.params.pop("similarity_function") + self._top_k = self.params.pop("top_k") + + vectorizer_model = self.params.pop("BM25Vectorizer") + self._vectorizer = vectorizer_model.vectorizer + + def init_model(self) -> None: + """Restore saved state or fit BM25 from scratch.""" + if not self.load(): + self._fit() + + def load(self) -> bool: + """Load a previously saved BM25 state from disk. + + Returns: + ``True`` if state was loaded successfully, ``False`` if + no saved state exists or loading failed. + """ + if self._persistence.model_dir is None: + return False + model_dir = self._persistence.model_dir + try: + with open(os.path.join(model_dir, "bm25_vectorizer.pkl"), "rb") as f: + vectorizer = pickle.load(f) + with open(os.path.join(model_dir, "bm25_tf_matrix.pkl"), "rb") as f: + tf_matrix = pickle.load(f) + with open(os.path.join(model_dir, "bm25_matrix.pkl"), "rb") as f: + bm25_matrix = pickle.load(f) + with open(os.path.join(model_dir, "bm25_row_to_chunk.pkl"), "rb") as f: + matrix_row_to_chunk_map = pickle.load(f) + self._vectorizer = vectorizer + self._tf_matrix = tf_matrix + self._bm25_matrix = bm25_matrix + self.matrix_row_to_chunk_map = matrix_row_to_chunk_map + return True + except Exception as e: + log.error("Error loading BM25 state: %s", e) + return False + + def save(self) -> None: + """Persist the vectorizer, matrices, and chunk map to disk. + + Raises: + ValueError: If ``persistence.model_dir`` is ``None``. + """ + model_dir = self._persistence.model_dir + if model_dir is None: + raise ValueError( + "Cannot save BM25Retriever: persistence.model_dir is None." + ) + os.makedirs(model_dir, exist_ok=True) + to_save = { + "bm25_vectorizer.pkl": self._vectorizer, + "bm25_tf_matrix.pkl": self._tf_matrix, + "bm25_matrix.pkl": self._bm25_matrix, + "bm25_row_to_chunk.pkl": self.matrix_row_to_chunk_map, + } + for fname, obj in to_save.items(): + with open(os.path.join(model_dir, fname), "wb") as f: + pickle.dump(obj, f) + + def _fit(self): + """Fit the CountVectorizer and compute the BM25-weighted matrix.""" + chunk_texts = [] + current_idx = 0 + self.matrix_row_to_chunk_map: Dict[int, Chunk] = {} + + for doc_chunks in self.chunks.values(): + for chunk in doc_chunks.values(): + chunk_texts.append(chunk.text) + self.matrix_row_to_chunk_map[current_idx] = chunk + current_idx += 1 + + self._vectorizer.fit(chunk_texts) + self._tf_matrix = self._vectorizer.transform(chunk_texts) + + tf = self._tf_matrix.copy() + n_docs = tf.shape[0] + doc_lengths = np.array(tf.sum(axis=1)).flatten() + avgdl = np.mean(doc_lengths) + + doc_freq = np.array((tf > 0).sum(axis=0)).flatten() + idf = np.log((n_docs - doc_freq + 0.5) / (doc_freq + 0.5) + 1.0) + idf += self.delta + idf[idf < 0] = 0 + + bm25 = lil_matrix(tf.shape) + for i in range(tf.shape[0]): + row = tf[i] + length_norm = 1 - self.b + self.b * (doc_lengths[i] / avgdl) + for j in range(row.indptr[0], row.indptr[1]): + col = row.indices[j] + tf_val = row.data[j] + numerator = tf_val * (self.k1 + 1) + denominator = tf_val + self.k1 * length_norm + bm25[i, col] = idf[col] * numerator / denominator + + self._bm25_matrix = bm25.tocsr() + + @property + def retrieval_top_k(self) -> int: + """Return the configured top-k value. + + Returns: + Number of chunks to retrieve. + """ + return self._top_k + + def retrieve(self, query: str, top_k: int | None = None) -> List[Chunk]: + """Retrieve the top-k chunks by BM25-weighted similarity. + + Args: + query: The search query string. + top_k: Override for the default ``top_k``. Uses the + configured value if ``None``. + + Returns: + A list of :class:`Chunk` instances ordered by distance. + """ + self._check_infra() + k = top_k if top_k is not None else self._top_k + query_vec = self._vectorizer.transform([query]) + distances = pairwise_distances( + query_vec, + self._bm25_matrix, + metric=self.similarity_function_name, + ).flatten() + top_indices = np.argsort(distances)[:k] + return [self.matrix_row_to_chunk_map[idx] for idx in top_indices] + + def score_chunks(self, chunk_ids: List[int], query: str) -> List[Tuple[int, float]]: + """Score a set of chunk IDs against the query. + + Args: + chunk_ids: List of chunk IDs to score. + query: The search query string. + + Returns: + A list of ``(chunk_id, distance)`` tuples sorted by distance. + """ + self._check_infra() + query_vec = self._vectorizer.transform([query]) + chunk_id_to_row = {c.id: r for r, c in self.matrix_row_to_chunk_map.items()} + rows, valid_ids = [], [] + for cid in chunk_ids: + row = chunk_id_to_row.get(cid) + if row is not None: + rows.append(row) + valid_ids.append(cid) + if not rows: + return [] + chunk_vectors = self._bm25_matrix[rows] + distances = pairwise_distances( + query_vec, + chunk_vectors, + metric=self.similarity_function_name, + ).flatten() + scored = list(zip(valid_ids, distances.tolist(), strict=True)) + scored.sort(key=lambda x: x[1]) + return scored + + def get_chunk_vectors(self, chunk_ids: List[int]) -> np.ndarray: + """Return the BM25-weighted vectors for the given chunk IDs. + + Args: + chunk_ids: List of chunk IDs whose vectors are needed. + + Returns: + A 2D numpy array of BM25-weighted vectors. + + Raises: + ValueError: If none of the provided chunk IDs are found in + the matrix. + """ + chunk_id_to_row = {c.id: r for r, c in self.matrix_row_to_chunk_map.items()} + rows = [] + for cid in chunk_ids: + row = chunk_id_to_row.get(cid) + if row is not None: + rows.append(row) + if not rows: + raise ValueError( + f"None of the provided chunk_ids {chunk_ids} were found " + "in the BM25 matrix." + ) + return self._bm25_matrix[rows].toarray() diff --git a/DashAI/back/models/RAG/retrievers/sparse/sparse_retriever.py b/DashAI/back/models/RAG/retrievers/sparse/sparse_retriever.py new file mode 100644 index 000000000..dfaca24ee --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/sparse/sparse_retriever.py @@ -0,0 +1,21 @@ +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.retrievers.unit_retriever import UnitRetriever + + +class SparseRetriever(UnitRetriever): + """Abstract base for sparse (keyword-based) retrievers. + + Uses term-frequency based representations (e.g. TF-IDF, BM25) for + document retrieval. + """ + + FLAGS: list[str] = ["abstract"] + DISPLAY_NAME: str = MultilingualString( + en="Keyword Retriever", + es="Recuperador por Palabras Clave", + ) + DESCRIPTION: str = MultilingualString( + en="Sparse retriever using term-frequency based representations.", + es="Recuperador disperso que usa representaciones basadas en" + " frecuencia de términos.", + ) diff --git a/DashAI/back/models/RAG/retrievers/sparse/tfidf_retriever.py b/DashAI/back/models/RAG/retrievers/sparse/tfidf_retriever.py new file mode 100644 index 000000000..3e161f715 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/sparse/tfidf_retriever.py @@ -0,0 +1,438 @@ +import logging +import os +import pickle +from typing import Dict, List, Tuple + +import numpy as np +from sklearn.feature_extraction.text import TfidfVectorizer +from sklearn.metrics.pairwise import pairwise_distances + +from DashAI.back.core.schema_fields import ( + BaseSchema, + bool_field, + component_field, + enum_field, + float_field, + int_field, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.base_model import BaseModel +from DashAI.back.models.RAG.documents import Chunk +from DashAI.back.models.RAG.retrievers.sparse.sparse_retriever import SparseRetriever + +log = logging.getLogger(__name__) + + +class TFIDFVectorizerSchema(BaseSchema): + """Schema for the TF-IDF vectorizer parameters. + + Note: + ``"None"`` is used as a string value instead of ``none_type()`` + because the current schema fields system does not support nullable + enum types. The value is converted to Python ``None`` in + :meth:`TFIDFVectorizerModel.__init__`. + """ + + strip_accents: schema_field( + enum_field(enum=["ascii", "unicode", "None"]), + placeholder="None", + description=MultilingualString( + en="Whether to strip accents from the text.", + es="Si se deben eliminar los acentos del texto.", + ), + ) # type: ignore + + lowercase: schema_field( + bool_field(), + placeholder=True, + description=MultilingualString( + en="Whether to convert all characters to lowercase.", + es="Si se deben convertir todos los caracteres a minúsculas.", + ), + ) # type: ignore + + analyzer: schema_field( + enum_field(enum=["word", "char", "char_wb"]), + placeholder="word", + description=MultilingualString( + en="Whether the feature should be made of word or character n-grams.", + es="Si las características deben ser n-gramas de palabras o caracteres.", + ), + ) # type: ignore + + stop_words: schema_field( + list_field(string_field(), min_items=0), + placeholder=[], + description=MultilingualString( + en="List of stop words. Leave empty to use none.", + es="Lista de palabras vacías. Dejar vacío para no usar ninguna.", + ), + ) # type: ignore + + ngram_range: schema_field( + list_field(int_field(), min_items=2, max_items=2), + placeholder=[1, 1], + description=MultilingualString( + en="Lower and upper boundary of the n-gram range.", + es="Límite inferior y superior del rango de n-gramas.", + ), + ) # type: ignore + + max_df: schema_field( + float_field(ge=0.0, le=1.0), + placeholder=1.0, + description=MultilingualString( + en="Ignore terms with document frequency above this threshold.", + es="Ignorar términos con frecuencia de documento superior a este umbral.", + ), + ) # type: ignore + + min_df: schema_field( + float_field(ge=0.0, le=1.0), + placeholder=0.0, + description=MultilingualString( + en="Ignore terms with document frequency below this threshold.", + es="Ignorar términos con frecuencia de documento inferior a este umbral.", + ), + ) # type: ignore + + max_features: schema_field( + int_field(ge=0), + placeholder=1000, + description=MultilingualString( + en="Maximum number of features. 0 means no limit.", + es="Número máximo de características. 0 significa sin límite.", + ), + ) # type: ignore + + norm: schema_field( + enum_field(enum=["l1", "l2", "None"]), + placeholder="l2", + description=MultilingualString( + en="Norm used to normalize term vectors.", + es="Norma utilizada para normalizar los vectores de términos.", + ), + ) # type: ignore + + use_idf: schema_field( + bool_field(), + placeholder=True, + description=MultilingualString( + en="Enable inverse-document-frequency reweighting.", + es="Activar reponderación por frecuencia inversa de documento.", + ), + ) # type: ignore + + smooth_idf: schema_field( + bool_field(), + placeholder=True, + description=MultilingualString( + en="Smooth IDF weights to prevent zero divisions.", + es="Suavizar pesos IDF para prevenir divisiones por cero.", + ), + ) # type: ignore + + sublinear_tf: schema_field( + bool_field(), + placeholder=False, + description=MultilingualString( + en="Apply sublinear TF scaling (1 + log(tf)).", + es="Aplicar escalado sublineal de TF (1 + log(tf)).", + ), + ) # type: ignore + + +class TFIDFVectorizerModel(BaseModel): + """Model component that encapsulates a :class:`TfidfVectorizer`. + + Validates parameters against :class:`TFIDFVectorizerSchema` and + constructs the underlying scikit-learn vectorizer. + """ + + SCHEMA = TFIDFVectorizerSchema + DISPLAY_NAME: str = MultilingualString( + en="TF-IDF Vectorizer Model", + es="Modelo de Vectorización TF-IDF", + ) + DESCRIPTION: str = MultilingualString( + en="Model component that encapsulates a TF-IDF vectorizer.", + es="Componente del modelo que encapsula un vectorizador TF-IDF.", + ) + + def __init__(self, **kwargs) -> None: + """Initialize and build the underlying ``TfidfVectorizer``. + + Args: + **kwargs: Parameters matching :class:`TFIDFVectorizerSchema`. + """ + validated = self.SCHEMA.model_validate(kwargs) + self.params = dict(validated) + if self.params.get("strip_accents") == "None": + self.params["strip_accents"] = None + stop_words = self.params.get("stop_words") or None + ngram_range = tuple(self.params.get("ngram_range")) + self.model = TfidfVectorizer( + strip_accents=self.params.get("strip_accents"), + lowercase=self.params.get("lowercase"), + analyzer=self.params.get("analyzer"), + stop_words=stop_words, + ngram_range=ngram_range, + max_df=self.params.get("max_df"), + min_df=self.params.get("min_df"), + max_features=self.params.get("max_features"), + norm=self.params.get("norm"), + use_idf=self.params.get("use_idf"), + smooth_idf=self.params.get("smooth_idf"), + sublinear_tf=self.params.get("sublinear_tf"), + ) + + def save(self, filename: str = "") -> None: + """No-op save (state managed by the parent retriever). + + Args: + filename: Ignored. + """ + + def load(self, filename: str = "") -> None: + """No-op load (state managed by the parent retriever). + + Args: + filename: Ignored. + """ + + def train(self, **kwargs): + """No-op train (fitting is done by the parent retriever). + + Args: + **kwargs: Ignored. + """ + return + + +class TFIDFRetrieverSchema(BaseSchema): + """Schema for :class:`TFIDFRetriever`. + + Attributes: + TFIDFVectorizer: Parameters for the TF-IDF vectorizer component. + similarity_function: Distance metric for vector comparison. + top_k: Number of chunks to select. + """ + + TFIDFVectorizer: schema_field( + component_field(parent="TFIDFVectorizerModel"), + placeholder={"component": "TFIDFVectorizerModel", "params": {}}, + description=MultilingualString( + en="TF-IDF Vectorizer parameters.", + es="Parámetros del vectorizador TF-IDF.", + ), + ) # type: ignore + + similarity_function: schema_field( + enum_field(["cityblock", "cosine", "euclidean", "l1", "l2", "manhattan"]), + placeholder="cosine", + description=MultilingualString( + en="Distance metric for comparing TF-IDF vectors.", + es="Métrica de distancia para comparar vectores TF-IDF.", + ), + ) # type: ignore + + top_k: schema_field( + int_field(ge=1), + placeholder=5, + description=MultilingualString( + en="Number of chunks to select.", + es="Número de fragmentos a seleccionar.", + ), + ) # type: ignore + + +class TFIDFRetriever(SparseRetriever): + """Sparse retriever using TF-IDF vectorization for document retrieval. + + Fits a :class:`sklearn.feature_extraction.text.TfidfVectorizer` on + the injected chunks and retrieves via pairwise distance. + """ + + FLAGS: list[str] = ["keyword", "sparse"] + DISPLAY_NAME: str = MultilingualString( + en="TF-IDF Retriever", + es="Recuperador TF-IDF", + ) + DESCRIPTION: str = MultilingualString( + en="Sparse retriever using TF-IDF vectorization for document retrieval.", + es="Recuperador disperso que usa vectorización TF-IDF para" + " recuperar documentos.", + ) + + SCHEMA = TFIDFRetrieverSchema + + def __init__(self, **kwargs): + """Initialize the TF-IDF retriever. + + Args: + **kwargs: Must contain ``TFIDFVectorizer``, + ``similarity_function``, and ``top_k``. + """ + super().__init__(**kwargs) + + vectorizer_model = self.params.pop("TFIDFVectorizer") + self._vectorizer = vectorizer_model.model + self.similarity_function_name = self.params.pop("similarity_function") + self._top_k = self.params.pop("top_k") + + def init_model(self) -> None: + """Restore saved state or fit the vectorizer from scratch.""" + if not self.load(): + self._fit() + + def load(self) -> bool: + """Load a previously saved TF-IDF state from disk. + + Returns: + ``True`` if state was loaded successfully, ``False`` if + no saved state exists or loading failed. + """ + if self._persistence.model_dir is None: + return False + model_dir = self._persistence.model_dir + try: + with open(os.path.join(model_dir, "tfidf_vectorizer.pkl"), "rb") as f: + vectorizer = pickle.load(f) + with open(os.path.join(model_dir, "tf_idf_matrix.pkl"), "rb") as f: + tf_idf_matrix = pickle.load(f) + with open( + os.path.join(model_dir, "matrix_row_to_chunk_map.pkl"), "rb" + ) as f: + matrix_row_to_chunk_map = pickle.load(f) + self._vectorizer = vectorizer + self._tf_idf_matrix = tf_idf_matrix + self.matrix_row_to_chunk_map = matrix_row_to_chunk_map + return True + except Exception as e: + log.error("Error loading TFIDF state: %s", e) + return False + + def save(self) -> None: + """Persist the vectorizer, TF-IDF matrix, and chunk map to disk. + + Raises: + ValueError: If ``persistence.model_dir`` is ``None``. + """ + model_dir = self._persistence.model_dir + if model_dir is None: + raise ValueError( + "Cannot save TFIDFRetriever: persistence.model_dir is None." + ) + os.makedirs(model_dir, exist_ok=True) + to_save = { + "tfidf_vectorizer.pkl": self._vectorizer, + "tf_idf_matrix.pkl": self._tf_idf_matrix, + "matrix_row_to_chunk_map.pkl": self.matrix_row_to_chunk_map, + } + for fname, obj in to_save.items(): + with open(os.path.join(model_dir, fname), "wb") as f: + pickle.dump(obj, f) + + def _fit(self): + """Fit the TF-IDF vectorizer on all chunk texts and build the matrix.""" + chunk_texts = [] + current_idx = 0 + self.matrix_row_to_chunk_map: Dict[int, Chunk] = {} + + for doc_chunks in self.chunks.values(): + for chunk in doc_chunks.values(): + chunk_texts.append(chunk.text) + self.matrix_row_to_chunk_map[current_idx] = chunk + current_idx += 1 + + self._tf_idf_matrix = self._vectorizer.fit_transform(chunk_texts) + + @property + def retrieval_top_k(self) -> int: + """Return the configured top-k value. + + Returns: + Number of chunks to retrieve. + """ + return self._top_k + + def retrieve(self, query: str, top_k: int | None = None) -> List[Chunk]: + """Retrieve the top-k chunks by TF-IDF similarity. + + Args: + query: The search query string. + top_k: Override for the default ``top_k``. Uses the + configured value if ``None``. + + Returns: + A list of :class:`Chunk` instances ordered by distance. + """ + self._check_infra() + k = top_k if top_k is not None else self._top_k + query_vector = self._vectorizer.transform([query]) + distances = pairwise_distances( + query_vector, + self._tf_idf_matrix, + metric=self.similarity_function_name, + ).flatten() + top_indices = np.argsort(distances)[:k] + return [self.matrix_row_to_chunk_map[idx] for idx in top_indices] + + def score_chunks(self, chunk_ids: List[int], query: str) -> List[Tuple[int, float]]: + """Score a set of chunk IDs against the query. + + Args: + chunk_ids: List of chunk IDs to score. + query: The search query string. + + Returns: + A list of ``(chunk_id, distance)`` tuples sorted by distance. + """ + self._check_infra() + query_vector = self._vectorizer.transform([query]) + chunk_id_to_row = {c.id: r for r, c in self.matrix_row_to_chunk_map.items()} + rows, valid_ids = [], [] + for cid in chunk_ids: + row = chunk_id_to_row.get(cid) + if row is not None: + rows.append(row) + valid_ids.append(cid) + if not rows: + return [] + distances = pairwise_distances( + query_vector, + self._tf_idf_matrix[rows], + metric=self.similarity_function_name, + ).flatten() + scored = list(zip(valid_ids, distances.tolist(), strict=True)) + scored.sort(key=lambda x: x[1]) + return scored + + def get_chunk_vectors(self, chunk_ids: List[int]) -> np.ndarray: + """Return the TF-IDF vectors for the given chunk IDs. + + Args: + chunk_ids: List of chunk IDs whose vectors are needed. + + Returns: + A 2D numpy array of TF-IDF vectors. + + Raises: + ValueError: If none of the provided chunk IDs are found in + the matrix. + """ + chunk_id_to_row = {c.id: r for r, c in self.matrix_row_to_chunk_map.items()} + rows = [] + for cid in chunk_ids: + row = chunk_id_to_row.get(cid) + if row is not None: + rows.append(row) + if not rows: + raise ValueError( + f"None of the provided chunk_ids {chunk_ids} were found " + "in the TF-IDF matrix." + ) + return self._tf_idf_matrix[rows].toarray() diff --git a/DashAI/back/models/RAG/retrievers/unit_retriever.py b/DashAI/back/models/RAG/retrievers/unit_retriever.py new file mode 100644 index 000000000..78f37fb4b --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/unit_retriever.py @@ -0,0 +1,103 @@ +from abc import ABC +from typing import Any, Dict, Final, List + +from DashAI.back.models.RAG.documents import Chunk +from DashAI.back.models.RAG.retrievers.persistence import ( + DensePersistence, + SparsePersistence, +) +from DashAI.back.models.RAG.retrievers.retriever_model import RetrieverModel + + +class UnitRetriever(RetrieverModel, ABC): + """Leaf: abstract base for unit retrievers. + + A unit retriever cannot contain children — it is the leaf node in the + Composite design pattern. Concrete subclasses must be either sparse + or dense retrievers. + """ + + TYPE: Final[str] = "RetrieverModel" + FLAGS: list[str] = ["abstract"] + + def __init__(self, **kwargs): + """Initialize the unit retriever. + + Args: + **kwargs: Keyword arguments forwarded to the parent + :class:`RetrieverModel`. + """ + super().__init__(**kwargs) + + def inject_infra( + self, + env_RAG_path: str, # noqa: N803 + chunks: Dict[int, Dict[int, Chunk]], + persistence: Any, + ) -> None: + """Inject runtime infrastructure with type-checked persistence. + + Args: + env_RAG_path: Root directory path for RAG data. + chunks: Nested dictionary mapping document IDs to chunk IDs + to :class:`Chunk` instances. + persistence: A :class:`SparsePersistence` or + :class:`DensePersistence` instance. + + Raises: + TypeError: If *persistence* is not a supported persistence type. + """ + if not isinstance(persistence, (SparsePersistence, DensePersistence)): + raise TypeError( + f"Expected SparsePersistence or DensePersistence, " + f"got {type(persistence).__name__}" + ) + super().inject_infra(env_RAG_path, chunks, persistence) + + def _check_infra(self) -> None: + """Raise if infrastructure has not been injected. + + Raises: + RuntimeError: If ``inject_infra()`` has not been called yet. + """ + if self._persistence is None: + raise RuntimeError( + f"{self.__class__.__name__}: infrastructure not injected. " + "Call inject_infra() before retrieve()." + ) + + def add(self, child: RetrieverModel) -> None: + """Add a child retriever (not supported for unit retrievers). + + Args: + child: The child retriever to add. + + Raises: + TypeError: Always — unit retrievers cannot contain children. + """ + raise TypeError( + f"{self.__class__.__name__} is a unit retriever and" + " cannot contain children." + ) + + def remove(self, child: RetrieverModel) -> None: + """Remove a child retriever (not supported for unit retrievers). + + Args: + child: The child retriever to remove. + + Raises: + TypeError: Always — unit retrievers cannot contain children. + """ + raise TypeError( + f"{self.__class__.__name__} is a unit retriever and" + " cannot contain children." + ) + + def get_children(self) -> List[RetrieverModel]: + """Return the (empty) list of child retrievers. + + Returns: + An empty list — unit retrievers have no children. + """ + return [] diff --git a/DashAI/back/models/RAG/utils.py b/DashAI/back/models/RAG/utils.py new file mode 100644 index 000000000..c2b2074d1 --- /dev/null +++ b/DashAI/back/models/RAG/utils.py @@ -0,0 +1,18 @@ +import hashlib + + +def hash_function(content: str | bytes) -> str: + """Generate a SHA-256 hash for the given content. + + Args: + content: The content to hash (str or bytes). + + Returns: + A hex-encoded SHA-256 hash string. + + Raises: + UnicodeEncodeError: If a str content cannot be encoded as UTF-8. + """ + if isinstance(content, str): + content = content.encode("utf-8") + return hashlib.sha256(content).hexdigest() diff --git a/DashAI/back/models/__init__.py b/DashAI/back/models/__init__.py index 9c0fa90a1..4f745c4db 100644 --- a/DashAI/back/models/__init__.py +++ b/DashAI/back/models/__init__.py @@ -1 +1,52 @@ # flake8: noqa +from DashAI.back.models.base_generative_model import BaseGenerativeModel +from DashAI.back.models.base_model import BaseModel +from DashAI.back.models.hugging_face.distilbert_transformer import DistilBertTransformer +from DashAI.back.models.hugging_face.opus_mt_en_es_transformer import ( + OpusMtEnESTransformer, +) +from DashAI.back.models.hugging_face.qwen_model import QwenModel +from DashAI.back.models.hugging_face.stable_diffusion_v1_depth_controlnet import ( + StableDiffusionXLV1ControlNet, +) +from DashAI.back.models.hugging_face.stable_diffusion_v2_model import ( + StableDiffusionV2Model, +) +from DashAI.back.models.hugging_face.stable_diffusion_v3_model import ( + StableDiffusionV3Model, +) +from DashAI.back.models.scikit_learn.bow_text_classification_model import ( + BagOfWordsTextClassificationModel, +) +from DashAI.back.models.scikit_learn.decision_tree_classifier import ( + DecisionTreeClassifier, +) +from DashAI.back.models.scikit_learn.dummy_classifier import DummyClassifier +from DashAI.back.models.scikit_learn.gradient_boosting_regression import ( + GradientBoostingR, +) +from DashAI.back.models.scikit_learn.hist_gradient_boosting_classifier import ( + HistGradientBoostingClassifier, +) +from DashAI.back.models.scikit_learn.k_neighbors_classifier import KNeighborsClassifier +from DashAI.back.models.scikit_learn.linear_regression import LinearRegression +from DashAI.back.models.scikit_learn.linearSVR import LinearSVR +from DashAI.back.models.scikit_learn.logistic_regression import LogisticRegression +from DashAI.back.models.scikit_learn.mlp_regression import MLPRegression +from DashAI.back.models.scikit_learn.random_forest_classifier import ( + RandomForestClassifier, +) +from DashAI.back.models.scikit_learn.random_forest_regression import ( + RandomForestRegression, +) +from DashAI.back.models.scikit_learn.ridge_regression import RidgeRegression +from DashAI.back.models.scikit_learn.sklearn_like_classifier import ( + SklearnLikeClassifier, +) +from DashAI.back.models.scikit_learn.sklearn_like_model import SklearnLikeModel +from DashAI.back.models.scikit_learn.sklearn_like_regressor import SklearnLikeRegressor +from DashAI.back.models.scikit_learn.svc import SVC + +from DashAI.back.models.RAG import RAGPipeline +from DashAI.back.models.remote_models import OpenAITextToTextGenerationModel +from DashAI.back.models.hugging_face.phi_4_mini_instruct_model import Phi4MiniInstructModel diff --git a/DashAI/back/models/base_model.py b/DashAI/back/models/base_model.py index 6cebb7818..71bf97bbc 100644 --- a/DashAI/back/models/base_model.py +++ b/DashAI/back/models/base_model.py @@ -26,6 +26,7 @@ class BaseModel(ConfigObject, metaclass=ABCMeta): """ TYPE: Final[str] = "Model" + FLAGS: list[str] = [] DISPLAY_NAME: str = "" DESCRIPTION: str = "" COLOR: str = "#795548" diff --git a/DashAI/back/models/hugging_face/__init__.py b/DashAI/back/models/hugging_face/__init__.py index e69de29bb..98f1a13e7 100644 --- a/DashAI/back/models/hugging_face/__init__.py +++ b/DashAI/back/models/hugging_face/__init__.py @@ -0,0 +1,3 @@ +from DashAI.back.models.hugging_face.deep_seek_model import DeepSeekModel +from DashAI.back.models.hugging_face.qwen_model import QwenModel +from DashAI.back.models.hugging_face.phi_4_mini_instruct_model import Phi4MiniInstructModel diff --git a/DashAI/back/models/hugging_face/deep_seek_model.py b/DashAI/back/models/hugging_face/deep_seek_model.py new file mode 100644 index 000000000..ae8a45dad --- /dev/null +++ b/DashAI/back/models/hugging_face/deep_seek_model.py @@ -0,0 +1,93 @@ +from typing import List + +from llama_cpp import Llama + +from DashAI.back.core.schema_fields import ( + BaseSchema, + float_field, + int_field, + schema_field, +) +from DashAI.back.models.text_to_text_generation_model import ( + TextToTextGenerationTaskModel, +) + + +class DeepSeekSchema(BaseSchema): + """Schema for DeepSeek model.""" + + max_tokens: schema_field( + int_field(ge=1), + placeholder=100, + description="Maximum number of tokens to generate.", + ) # type: ignore + + temperature: schema_field( + float_field(ge=0.0, le=1.0), + placeholder=0.7, + description=( + "Sampling temperature. Higher values make the output more random, " + "while lower values make it more focused and deterministic." + ), + ) # type: ignore + + frequency_penalty: schema_field( + float_field(ge=0.0, le=2.0), + placeholder=0.1, + description=( + "Penalty for repeated tokens in the output. Higher values reduce the " + "likelihood of repetition, encouraging more diverse text generation." + ), + ) # type: ignore + + n_ctx: schema_field( + int_field(ge=1), + placeholder=4096, + description=( + "Maximum number of tokens the model can process in a single forward pass " + "(context window size)." + ), + ) # type: ignore + + +class DeepSeekModel(TextToTextGenerationTaskModel): + """DeepSeek model for text generation using llama.cpp library.""" + + SCHEMA = DeepSeekSchema + + def __init__(self, **kwargs): + kwargs = self.validate_and_transform(kwargs) + self.max_tokens = kwargs.pop("max_tokens", 100) + self.temperature = kwargs.pop("temperature", 0.7) + self.frequency_penalty = kwargs.pop("frequency_penalty", 0.1) + self.n_ctx = kwargs.pop("n_ctx", 512) + + self.model_id = "TheBloke/deepseek-llm-7B-base-GGUF" + self.filename = "*Q3_K_S.gguf" + + self.model = Llama.from_pretrained( + repo_id=self.model_id, + filename=self.filename, + verbose=True, + n_ctx=self.n_ctx, + n_gpu_layers=-1, + ) + + def generate(self, prompt: str) -> List[str]: + if len(prompt) > self.model.n_ctx(): + prompt = prompt[-self.model.n_ctx() :] + + """Generate text based on prompts.""" + + output = self.model( + prompt, + max_tokens=self.max_tokens, + temperature=self.temperature, + frequency_penalty=self.frequency_penalty, + stop=["Q:"], + echo=False, + ) + + generated_text = output["choices"][0]["text"] + clean_text = generated_text.replace(prompt, "").strip() + return [clean_text] diff --git a/DashAI/back/models/hugging_face/llama_model.py b/DashAI/back/models/hugging_face/llama_model.py index 18f36239b..7a9f29fb2 100644 --- a/DashAI/back/models/hugging_face/llama_model.py +++ b/DashAI/back/models/hugging_face/llama_model.py @@ -21,6 +21,8 @@ "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF": "*Q4_K_M.gguf", "bartowski/Llama-3.2-1B-Instruct-GGUF": "*Q4_K_M.gguf", "bartowski/Llama-3.2-3B-Instruct-GGUF": "*Q4_K_M.gguf", + "TheBloke/Llama-2-7B-Chat-GGUF": "*Q4_K_M.gguf", + "TheBloke/Llama-2-13B-Chat-GGUF": "*Q4_K_M.gguf", } @@ -39,6 +41,8 @@ class LlamaSchema(BaseSchema): "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF", "bartowski/Llama-3.2-1B-Instruct-GGUF", "bartowski/Llama-3.2-3B-Instruct-GGUF", + "TheBloke/Llama-2-7B-Chat-GGUF", + "TheBloke/Llama-2-13B-Chat-GGUF", ] ), placeholder="bartowski/Llama-3.2-3B-Instruct-GGUF", @@ -50,7 +54,9 @@ class LlamaSchema(BaseSchema): "'Llama-3.2-3B' (~3B parameters) offers a good speed/quality " "trade-off. " "'Meta-Llama-3.1-8B' (~8B parameters) delivers the highest quality " - "at the cost of more RAM and slower inference." + "at the cost of more RAM and slower inference. " + "Llama-2-7B-Chat and Llama-2-13B-Chat are also available via " + "TheBloke's community quantizations." ), es=( "El checkpoint Meta Llama 3.x Instruct a cargar en formato GGUF " @@ -59,7 +65,9 @@ class LlamaSchema(BaseSchema): "ideal para sistemas solo con CPU. " "'Llama-3.2-3B' (~3B parámetros) ofrece un buen equilibrio entre " "velocidad y calidad. 'Meta-Llama-3.1-8B' (~8B parámetros) entrega " - "la mayor calidad a costa de más RAM e inferencia más lenta." + "la mayor calidad a costa de más RAM e inferencia más lenta. " + "Llama-2-7B-Chat y Llama-2-13B-Chat también están disponibles " + "mediante las cuantizaciones comunitarias de TheBloke." ), pt=( "O checkpoint Meta Llama 3.x Instruct para carregar em formato GGUF " @@ -69,7 +77,9 @@ class LlamaSchema(BaseSchema): "'Llama-3.2-3B' (~3B parâmetros) oferece um bom equilíbrio entre " "velocidade e qualidade. 'Meta-Llama-3.1-8B' (~8B parâmetros) " "entrega a maior qualidade ao custo de mais RAM e " - "inferência mais lenta." + "inferência mais lenta. " + "Llama-2-7B-Chat e Llama-2-13B-Chat também estão disponíveis " + "via quantizações comunitárias de TheBloke." ), de=( "Der im GGUF-Format zu ladende Meta Llama 3.x Instruct-Checkpoint " @@ -97,6 +107,53 @@ class LlamaSchema(BaseSchema): ), ) # type: ignore + quantization: schema_field( + enum_field( + enum=[ + "Q2_K", + "Q2_K_L", + "Q3_K_S", + "Q3_K_M", + "Q3_K_L", + "Q3_K_XL", + "Q4_K_S", + "Q4_K_0", + "Q4_K_M", + "Q5_K_S", + "Q5_K_M", + "Q6_K", + "Q6_K_L", + "Q8_0", + "F32", + ], + ), + placeholder="Q4_K_M", + description=MultilingualString( + en=( + "Quantization format for the GGUF checkpoint. 'Q4_K_M' is the " + "default and recommended for most users, offering a good balance " + "between speed and quality. Other formats may be faster or smaller " + "but can reduce output quality." + ), + es=( + "Formato de cuantización para el checkpoint GGUF. 'Q4_K_M' es el " + "predeterminado y recomendado para la mayoría de los usuarios, " + "ofreciendo un buen equilibrio entre velocidad y calidad. Otros " + "formatos pueden ser más rápidos o más pequeños pero pueden reducir " + "la calidad de salida." + ), + pt=( + "Formato de quantização para o checkpoint GGUF. 'Q4_K_M' é o " + "padrão e recomendado para a maioria dos usuários, oferecendo um " + "bom equilíbrio entre velocidade e qualidade. Outros formatos podem " + "ser mais rápidos ou menores, mas podem reduzir a qualidade da saída." + ), + ), + alias=MultilingualString( + en="Quantization", es="Cuantización", pt="Quantização" + ), + ) # type: ignore + max_tokens: schema_field( int_field(ge=1), placeholder=100, @@ -448,7 +505,9 @@ def __init__(self, **kwargs): self.frequency_penalty = kwargs.pop("frequency_penalty", 0.1) self.n_ctx = kwargs.pop("context_window", 512) - self.filename = LLAMA_FILENAME_MAP.get(self.model_name, "*Q4_K_M.gguf") + self.quantization = kwargs.get("quantization", "Q4_K_M") + self.filename = f"*{self.quantization}.gguf" + use_gpu = LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) >= 0 self.model = Llama.from_pretrained( diff --git a/DashAI/back/models/hugging_face/phi_4_mini_instruct_model.py b/DashAI/back/models/hugging_face/phi_4_mini_instruct_model.py new file mode 100644 index 000000000..54964c87c --- /dev/null +++ b/DashAI/back/models/hugging_face/phi_4_mini_instruct_model.py @@ -0,0 +1,255 @@ +from typing import List + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + float_field, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.text_to_text_generation_model import ( + TextToTextGenerationTaskModel, +) +from DashAI.back.models.utils import ( + LLAMA_DEVICE_ENUM, + LLAMA_DEVICE_PLACEHOLDER, + LLAMA_DEVICE_TO_IDX, +) + + +class Phi4MiniInstructSchema(BaseSchema): + """Schema for Phi 4 Mini Instruct model hyperparameters""" + + model_name: schema_field( + enum_field( + enum=[ + "unsloth/Phi-4-mini-instruct-GGUF", + ] + ), + placeholder="unsloth/Phi-4-mini-instruct-GGUF", + description=MultilingualString( + en=( + "Phi-4-mini-instruct is a lightweight open model built upon" + " synthetic data and filtered publicly available websites -" + " with a focus on high-quality, reasoning dense data. The" + " model belongs to the Phi-4 model family and supports" + " 128K token context length. The model underwent an" + " enhancement process, incorporating both supervised" + " fine-tuning and direct preference optimization to support" + " precise instruction adherence and robust safety measures." + ), + es=( + "Phi-4-mini-instruct es un modelo abierto ligero construido" + " sobre datos sintéticos y sitios web disponibles" + " públicamente filtrados, con un enfoque en datos de alta" + " calidad y densos en razonamiento. El modelo pertenece a la" + " familia de modelos Phi-4 y soporta un contexto de longitud" + " de 128K tokens. El modelo se sometió a un proceso de" + " mejora, incorporando tanto ajuste fino (fine-tuning)" + " supervisado como optimización directa de preferencias para" + " soportar una adherencia precisa a las instrucciones y" + " medidas de seguridad." + ), + ), + alias=MultilingualString(en="Model Name", es="Nombre del Modelo"), + ) # type: ignore + + quantization: schema_field( + enum_field( + enum=[ + # "unsloth/Phi-4-mini-instruct-GGUF", + "Phi-4-mini-instruct-Q2_K.gguf", + "Phi-4-mini-instruct-Q2_K_L.gguf", + "Phi-4-mini-instruct-Q3_K_M.gguf", + "Phi-4-mini-instruct-Q4_K_M.gguf", + "Phi-4-mini-instruct-Q5_K_M.gguf", + "Phi-4-mini-instruct-Q6_K.gguf", + "Phi-4-mini-instruct.BF16.gguf", + "Phi-4-mini-instruct.Q8_0.gguf", + ] + ), + placeholder="Phi-4-mini-instruct.Q8_0.gguf", + description=MultilingualString( + en=( + "The specific Phi 4 Mini Instruct model quantization to use. Options " + "include various quantization sizes and the BF16 format. The choice of " + "quantization can affect the model's performance and resource usage, " + "with smaller quantizations typically requiring less memory but " + "potentially sacrificing some accuracy." + ), + es=( + "La cuantización específica del modelo Phi 4 Mini Instruct a" + " utilizar. Las opciones incluyen varios tamaños de" + " cuantización y el formato BF16. La elección de la" + " cuantización puede afectar el rendimiento y el uso de" + " recursos del modelo, generalmente con cuantizaciones más" + " pequeñas requieren menos memoria pero potencialmente" + " sacrifican algo de precisión." + ), + ), + alias=MultilingualString(en="Quantization", es="Cuantización"), + ) # type: ignore + + max_tokens: schema_field( + int_field(ge=1), + placeholder=100, + description=MultilingualString( + en=( + "Maximum number of new tokens the model will generate per response. " + "Roughly 1 token ≈ 0.75 English words. Set to 100-200 for short " + "answers, 500-1000 for detailed explanations or code." + ), + es=( + "Número máximo de tokens nuevos que el modelo generará por respuesta. " + "Aproximadamente 1 token ≈ 0.75 palabras en español. Use 100-200 " + "para respuestas cortas, 500-1000 para explicaciones detalladas " + "o código." + ), + ), + alias=MultilingualString(en="Max tokens", es="Tokens máximos"), + ) # type: ignore + + temperature: schema_field( + float_field(ge=0.0, le=1.0), + placeholder=0.7, + description=MultilingualString( + en=( + "Sampling temperature controlling output randomness (range 0.0-1.0). " + "At 0.0 the model picks the most likely token (deterministic). " + "Around 0.7 balances quality and creativity. At 1.0 outputs are " + "maximally varied." + ), + es=( + "Temperatura de muestreo que controla la aleatoriedad (rango 0.0-1.0). " + "En 0.0 el modelo elige el token más probable (determinista). " + "Alrededor de 0.7 equilibra calidad y creatividad. En 1.0 las salidas " + "son máximamente variadas." + ), + ), + alias=MultilingualString(en="Temperature", es="Temperatura"), + ) # type: ignore + + frequency_penalty: schema_field( + float_field(ge=0.0, le=2.0), + placeholder=0.1, + description=MultilingualString( + en=( + "Penalizes tokens that have already appeared in the output based on " + "frequency (range 0.0-2.0). Higher values discourage repetition." + ), + es=( + "Penaliza los tokens que ya aparecieron en la salida según su " + "frecuencia (rango 0.0-2.0). Valores más altos desincentivan " + "la repetición." + ), + ), + alias=MultilingualString( + en="Frequency penalty", es="Penalización de frecuencia" + ), + ) # type: ignore + + context_window: schema_field( + int_field(ge=1, le=131072), + placeholder=512, + description=MultilingualString( + en=( + "Total token budget for a single forward pass, including prompt and " + "response. Mistral-7B supports up to 32K tokens; Mistral-Nemo " + "supports up to 128K tokens." + ), + es=( + "Presupuesto total de tokens por pasada, incluyendo prompt y " + "respuesta. Mistral-7B soporta hasta 32K tokens; Mistral-Nemo " + "soporta hasta 128K tokens." + ), + ), + alias=MultilingualString(en="Context window", es="Ventana de contexto"), + ) # type: ignore + + device: schema_field( + enum_field(enum=LLAMA_DEVICE_ENUM), + placeholder=LLAMA_DEVICE_PLACEHOLDER, + description=MultilingualString( + en=( + "Hardware device for llama.cpp inference. 'CPU' runs the model " + "fully in RAM. A GPU option offloads all layers for faster inference." + ), + es=( + "Dispositivo de hardware para inferencia con llama.cpp. 'CPU' ejecuta " + "el modelo en RAM. Una opción de GPU descarga todas las capas para " + "inferencia más rápida." + ), + ), + alias=MultilingualString(en="Device", es="Dispositivo"), + ) # type: ignore + + +class Phi4MiniInstructModel(TextToTextGenerationTaskModel): + """Phi 4 Mini Instruct model for text generation using llama.cpp library.""" + + SCHEMA = Phi4MiniInstructSchema + COLOR: str = "#FF5733" + DISPLAY_NAME: str = MultilingualString( + en="Phi 4 Mini Instruct", + es="Phi 4 Mini Instruct", + ) + DESCRIPTION: str = MultilingualString( + en=( + "Phi-4-mini-instruct is a lightweight open model built upon" + " synthetic data and filtered publicly available websites -" + " with a focus on high-quality, reasoning dense data. The" + " model belongs to the Phi-4 model family and supports" + " 128K token context length. The model underwent an" + " enhancement process, incorporating both supervised" + " fine-tuning and direct preference optimization to support" + " precise instruction adherence and robust safety measures." + ), + es=( + "Phi-4-mini-instruct es un modelo abierto ligero construido" + " sobre datos sintéticos y sitios web disponibles públicamente" + " filtrados, con un enfoque en datos de alta calidad y densos" + " en razonamiento. El modelo pertenece a la familia de modelos" + " Phi-4 y soporta un contexto de longitud de 128K tokens. El" + " modelo se sometió a un proceso de mejora, incorporando tanto" + " ajuste fino (fine-tuning) supervisado como optimización" + " directa de preferencias para soportar una adherencia precisa" + " a las instrucciones y medidas de seguridad." + ), + ) + + def __init__(self, **kwargs): + try: + from llama_cpp import Llama + except ImportError as e: + raise RuntimeError( + "llama-cpp-python is not installed. Please install it to use QwenModel." + ) from e + + kwargs = self.validate_and_transform(kwargs) + self.model_name = kwargs.pop("model_name") + self.quantization = kwargs.pop("quantization", "Phi-4-mini-instruct.Q8_0.gguf") + self.max_tokens = kwargs.pop("max_tokens", 100) + self.temperature = kwargs.pop("temperature", 0.7) + self.frequency_penalty = kwargs.pop("frequency_penalty", 0.1) + self.n_ctx = kwargs.pop("context_window", 512) + + use_gpu = LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) >= 0 + + self.model = Llama.from_pretrained( + repo_id=self.model_name, + filename=self.quantization, + verbose=True, + n_ctx=self.n_ctx, + n_gpu_layers=-1 if use_gpu else 0, + main_gpu=LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) if use_gpu else 0, + ) + + def generate(self, prompt: list[dict[str, str]]) -> List[str]: + output = self.model.create_chat_completion( + messages=prompt, + max_tokens=self.max_tokens, + temperature=self.temperature, + frequency_penalty=self.frequency_penalty, + ) + return [output["choices"][0]["message"]["content"]] diff --git a/DashAI/back/models/hugging_face/qwen_model.py b/DashAI/back/models/hugging_face/qwen_model.py index 475c536bc..04cf449be 100644 --- a/DashAI/back/models/hugging_face/qwen_model.py +++ b/DashAI/back/models/hugging_face/qwen_model.py @@ -32,6 +32,10 @@ class QwenSchema(BaseSchema): enum=[ "Qwen/Qwen2.5-0.5B-Instruct-GGUF", "Qwen/Qwen2.5-1.5B-Instruct-GGUF", + # "Qwen/Qwen3-4B-GGUF", This one is not working on llama-cpp 0.3.4 + "Qwen/Qwen2.5-3B-Instruct-GGUF", + "Qwen/Qwen2.5-7B-Instruct-GGUF", + "Qwen/Qwen2.5-14B-Instruct-GGUF", ] ), placeholder="Qwen/Qwen2.5-1.5B-Instruct-GGUF", @@ -431,17 +435,34 @@ def __init__(self, **kwargs): self.frequency_penalty = kwargs.pop("frequency_penalty", 0.1) self.n_ctx = kwargs.pop("context_window", 512) - self.filename = "*8_0.gguf" + self.filename = "*q8_0.gguf" use_gpu = LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) >= 0 - self.model = Llama.from_pretrained( - repo_id=self.model_name, - filename=self.filename, - verbose=True, - n_ctx=self.n_ctx, - n_gpu_layers=-1 if use_gpu else 0, - main_gpu=LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) if use_gpu else 0, - ) + try: + self.model = Llama.from_pretrained( + repo_id=self.model_name, + filename=self.filename, + verbose=True, + n_ctx=self.n_ctx, + n_gpu_layers=-1 if use_gpu else 0, + main_gpu=LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) + if use_gpu + else 0, + ) + except ValueError: + # Fall back to Q2_K single-file quantization when Q8_0 is + # not available (e.g. large models whose Q8_0 is multi-part). + self.filename = "*q2_k.gguf" + self.model = Llama.from_pretrained( + repo_id=self.model_name, + filename=self.filename, + verbose=True, + n_ctx=self.n_ctx, + n_gpu_layers=-1 if use_gpu else 0, + main_gpu=LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) + if use_gpu + else 0, + ) def generate(self, prompt: list[dict[str, str]]) -> List[str]: """Generate a reply for the given chat prompt. diff --git a/DashAI/back/models/remote_models/__init__.py b/DashAI/back/models/remote_models/__init__.py new file mode 100644 index 000000000..9dc7a5df7 --- /dev/null +++ b/DashAI/back/models/remote_models/__init__.py @@ -0,0 +1,4 @@ +from DashAI.back.models.remote_models.deepseek_text_to_text_generation_model import ( + DeepSeekTextToTextGenerationModel, +) +from DashAI.back.models.remote_models.openai_text_to_text_generation_model import OpenAITextToTextGenerationModel diff --git a/DashAI/back/models/remote_models/deepseek_text_to_text_generation_model.py b/DashAI/back/models/remote_models/deepseek_text_to_text_generation_model.py new file mode 100644 index 000000000..3024a5e95 --- /dev/null +++ b/DashAI/back/models/remote_models/deepseek_text_to_text_generation_model.py @@ -0,0 +1,146 @@ +from typing import List + +from openai import OpenAI +from pydantic import field_validator + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + float_field, + int_field, + none_type, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.text_to_text_generation_model import ( + TextToTextGenerationTaskModel, +) + +deepseek_available_models = [ + "deepseek-chat", + "deepseek-reasoner", +] + + +class DeepSeekTextToTextGenerationModelSchema(BaseSchema): + API_key: schema_field( + string_field(), + placeholder="", + description="API key for DeepSeek access.", + ) # type: ignore + + model_name: schema_field( + enum_field(enum=deepseek_available_models), + placeholder="deepseek-chat", + description="The specific DeepSeek model version to use.", + ) # type: ignore + + @field_validator( + "frequency_penalty", + "max_completions_tokens", + "presence_penalty", + "temperature", + "top_p", + mode="before", + ) + def validate_optional_fields(cls, v): # noqa: N805 + if v == "": + return None + return v + + frequency_penalty: schema_field( + none_type(float_field(ge=-2.0, le=2.0)), + placeholder=None, + description=( + "Number between -2.0 and 2.0. Positive values penalize new tokens " + "based on their existing frequency in the text so far, decreasing the " + "model's likelihood to repeat the same line verbatim." + ), + ) # type: ignore + + max_completions_tokens: schema_field( + none_type(int_field(ge=1)), + placeholder=None, + description=( + "An upper bound for the number of tokens that can be generated for a " + "completion, including visible output tokens and reasoning tokens." + ), + ) # type: ignore + + presence_penalty: schema_field( + none_type(float_field(ge=-2.0, le=2.0)), + placeholder=None, + description=( + "Number between -2.0 and 2.0. Positive values penalize new tokens " + "based on whether they appear in the text so far, increasing the " + "model's likelihood to talk about new topics." + ), + ) # type: ignore + + temperature: schema_field( + none_type(float_field(ge=0.0, le=2.0)), + placeholder=None, + description=( + "temperature: What sampling temperature to use, between 0 and 2. " + "Higher values like 0.8 will make the output more random, while lower " + "values like 0.2 will make it more focused and deterministic. " + "We generally recommend altering this or `top_p` but not both." + ), + ) # type: ignore + + top_p: schema_field( + none_type(float_field(ge=0.0, le=2.0)), + placeholder=None, + description=( + "An alternative to sampling with temperature, called nucleus sampling, " + "where the model considers the results of the tokens with top_p " + "probability mass. So 0.1 means only the tokens comprising the top 10% " + "probability mass are considered. " + "We generally recommend altering this or `temperature` but not both." + ), + ) # type: ignore + + +class DeepSeekTextToTextGenerationModel(TextToTextGenerationTaskModel): + """Wrapper around DeepSeek's text-to-text generation models via API.""" + + SCHEMA = DeepSeekTextToTextGenerationModelSchema + DISPLAY_NAME: str = MultilingualString( + en="DeepSeek API", + es="API de DeepSeek", + ) + DESCRIPTION: str = MultilingualString( + en=( + "API for DeepSeek's language models, allowing you to select and" + " configure the text-to-text models supported by the DeepSeek" + " API (https://deepseek.com/). Note that it requires a private" + " API key with an associated cost, that the language model runs" + " on DeepSeek's servers, and that your data will be used" + " according to that company's policies." + ), + es=( + "API para los modelos de lenguaje de DeepSeek, permite" + " seleccionar y configurar los modelos de texto a texto" + " soportados por la API de DeepSeek (https://deepseek.com/)." + " Considere que requiere una API key (clave de API) privada" + " con un costo asociado, que el modelo de lenguaje se ejecuta" + " en los servidores de DeepSeek y que sus datos serán" + " utilizados según las políticas de esa empresa." + ), + ) + + def __init__(self, **kwargs): + kwargs = self.validate_and_transform(kwargs) + self.client = OpenAI( + api_key=kwargs.get("API_key"), + base_url="https://api.deepseek.com", + ) + self.model_name = kwargs.get("model_name") + + def generate(self, prompt: list[dict[str, str]]) -> List[str]: + output = self.client.chat.completions.create( + model=self.model_name, + messages=prompt, + ) + return [output.choices[0].message.content] diff --git a/DashAI/back/models/remote_models/openai_text_to_text_generation_model.py b/DashAI/back/models/remote_models/openai_text_to_text_generation_model.py new file mode 100644 index 000000000..7cc7239c8 --- /dev/null +++ b/DashAI/back/models/remote_models/openai_text_to_text_generation_model.py @@ -0,0 +1,176 @@ +from typing import List + +from openai import OpenAI +from pydantic import field_validator + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + float_field, + int_field, + none_type, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.text_to_text_generation_model import ( + TextToTextGenerationTaskModel, +) + +# Hardcoded since the list of available models is only accessible via API call with +# api key :c +gpt_available_models = [ + "gpt-4-0613", + "gpt-4", + "gpt-3.5-turbo", + "gpt-4-0314", + "gpt-3.5-turbo-instruct", + "gpt-3.5-turbo-instruct-0914", + "gpt-4-1106-preview", + "gpt-3.5-turbo-1106", + "gpt-4-0125-preview", + "gpt-4-turbo-preview", + "gpt-3.5-turbo-0125", + "gpt-4-turbo", + "gpt-4-turbo-2024-04-09", + "gpt-4o", + "gpt-4o-2024-05-13", + "gpt-4o-mini-2024-07-18", + "gpt-4o-mini", + "gpt-4o-2024-08-06", + "gpt-4o-2024-11-20", + "gpt-4o-mini-tts", + "gpt-4.1-2025-04-14", + "gpt-4.1", + "gpt-4.1-mini-2025-04-14", + "gpt-4.1-mini", + "gpt-4.1-nano-2025-04-14", + "gpt-4.1-nano", + "gpt-5-chat-latest", + "gpt-5-2025-08-07", + "gpt-5", + "gpt-5-mini-2025-08-07", + "gpt-5-mini", + "gpt-5-nano-2025-08-07", + "gpt-5-nano", + "gpt-5-pro-2025-10-06", + "gpt-5-pro", + "gpt-3.5-turbo-16k", +] + + +class OpenAITextToTextGenerationModelSchema(BaseSchema): + API_key: schema_field( + string_field(), + placeholder="", + description="API key for OpenAI access.", + ) # type: ignore + + model_name: schema_field( + enum_field(enum=gpt_available_models), + placeholder="gpt-3.5-turbo", + description="The specific OpenAI model version to use.", + ) # type: ignore + + @field_validator( + "frequency_penalty", + "max_completions_tokens", + "presence_penalty", + "temperature", + "top_p", + mode="before", + ) + def validate_optional_fields(cls, v): # noqa: N805 + if v == "": + return None + return v + + frequency_penalty: schema_field( + none_type(float_field(ge=-2.0, le=2.0)), + placeholder=None, + description=( + "Number between -2.0 and 2.0. Positive values penalize new tokens " + "based on their existing frequency in the text so far, decreasing the " + "model's likelihood to repeat the same line verbatim." + ), + ) # type: ignore + + max_completions_tokens: schema_field( + none_type(int_field(ge=1)), + placeholder=None, + description=( + "An upper bound for the number of tokens that can be generated for a " + "completion, including visible output tokens and reasoning tokens." + ), + ) # type: ignore + + presence_penalty: schema_field( + none_type(float_field(ge=-2.0, le=2.0)), + placeholder=None, + description=( + "Number between -2.0 and 2.0. Positive values penalize new tokens " + "based on whether they appear in the text so far, increasing the " + "model's likelihood to talk about new topics." + ), + ) # type: ignore + + temperature: schema_field( + none_type(float_field(ge=0.0, le=2.0)), + placeholder=None, + description=( + "temperature: What sampling temperature to use, between 0 and 2. " + "Higher values like 0.8 will make the output more random, while lower " + "values like 0.2 will make it more focused and deterministic. " + "We generally recommend altering this or `top_p` but not both." + ), + ) # type: ignore + + top_p: schema_field( + none_type(float_field(ge=0.0, le=2.0)), + placeholder=None, + description=( + "An alternative to sampling with temperature, called nucleus sampling, " + "where the model considers the results of the tokens with top_p " + "probability mass. So 0.1 means only the tokens comprising the top 10% " + "probability mass are considered. " + "We generally recommend altering this or `temperature` but not both." + ), + ) # type: ignore + + +class OpenAITextToTextGenerationModel(TextToTextGenerationTaskModel): + """Wrapper around OpenAI's text-to-text generation models.""" + + SCHEMA = OpenAITextToTextGenerationModelSchema + DISPLAY_NAME: str = MultilingualString(en="OpenAI API", es="API de OpenAI") + DESCRIPTION: str = MultilingualString( + en=( + "API for OpenAI's language models, allowing you to select and" + " configure the text-to-text models supported by the OpenAI API" + " (https://openai.com/api/). Note that it requires a private" + " API key with an associated cost, that the language model runs" + " on OpenAI's servers, and that your data will be used" + " according to that company's policies." + ), + es=( + "API para los modelos de lenguaje de OpenAI, permite" + " seleccionar y configurar los modelos de texto a texto" + " soportados por la API de OpenAI (https://openai.com/api/)." + " Considere que requiere una API key (clave de API) privada" + " con un costo asociado, que el modelo de lenguaje se ejecuta" + " en los servidores de OpenAI y que sus datos serán" + " utilizados según las políticas de esa empresa." + ), + ) + + def __init__(self, **kwargs): + kwargs = self.validate_and_transform(kwargs) + self.client = OpenAI(api_key=kwargs.get("API_key")) + self.model_name = kwargs.get("model_name") + + def generate(self, prompt: list[dict[str, str]]) -> List[str]: + output = self.client.chat.completions.create( + model=self.model_name, + messages=prompt, + ) + return [output.choices[0].message.content] diff --git a/DashAI/back/models/text_to_text_generation_model.py b/DashAI/back/models/text_to_text_generation_model.py index a546b9569..4ebf9b04e 100644 --- a/DashAI/back/models/text_to_text_generation_model.py +++ b/DashAI/back/models/text_to_text_generation_model.py @@ -8,4 +8,6 @@ class TextToTextGenerationTaskModel(BaseGenerativeModel): ``generate``. Compatible with ``TextToTextGenerationTask``. """ - COMPATIBLE_COMPONENTS = ["TextToTextGenerationTask"] + COMPATIBLE_COMPONENTS = ["TextToTextGenerationTask", "RAGTask"] + + REQUIRED_EXTRA_KWARGS = [] diff --git a/DashAI/back/services/RAG/RAG_setup_service.py b/DashAI/back/services/RAG/RAG_setup_service.py new file mode 100644 index 000000000..bab636669 --- /dev/null +++ b/DashAI/back/services/RAG/RAG_setup_service.py @@ -0,0 +1,351 @@ +import logging +from typing import Any, Dict + +from sqlalchemy.orm import Session + +from DashAI.back.core.component_validation import validate_component_refs +from DashAI.back.core.schema_fields.utils import normalize_payload +from DashAI.back.dependencies.database.models import RAGPipeline as PipelineDBModel +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.RAG.chunking_models.base_chunking_model import ( + BaseChunkingModel, +) +from DashAI.back.models.RAG.documents import BaseDocument, Chunk +from DashAI.back.models.RAG.prompts.prompt import Prompt +from DashAI.back.models.RAG.RAG_constants import RAG_MODEL_KEYS as _RAG_MODEL_KEYS +from DashAI.back.models.RAG.RAG_models_factory import RAGModelsFactory +from DashAI.back.models.RAG.RAG_pipeline import RAGPipeline, RAGPipelineConfig +from DashAI.back.models.RAG.retrievers.retriever_model import RetrieverModel +from DashAI.back.models.text_to_text_generation_model import ( + TextToTextGenerationTaskModel, +) +from DashAI.back.services.RAG.chunking_service import ChunkingService +from DashAI.back.services.RAG.document_service import DocumentService +from DashAI.back.services.RAG.llm_service import LLMService +from DashAI.back.services.RAG.prompt_service import PromptService +from DashAI.back.services.RAG.retriever_setup_service import RetrieverSetupService + +log = logging.getLogger(__name__) + + +class RAGSetupService: + """Assembles RAG pipeline components into a ready-to-use RAGPipeline instance. + + Does NOT execute ``generate()`` — only assembly. Each sub-component + is built via the corresponding service class so that DB persistence + and model instantiation are handled consistently. + """ + + def __init__( + self, + db: Session, + registry: ComponentRegistry, + RAG_path: str, # noqa: N803 + ): + """Initialize the RAG setup service with DB and registry. + + Args: + db: SQLAlchemy session. + registry: Application component registry. + RAG_path: Base RAG data directory path. + """ + self._db = db + self._registry = registry + self._RAG_path = RAG_path + + self._documents = DocumentService(db) + self._chunking = ChunkingService(db, registry) + self._prompts = PromptService(db, registry) + self._llm = LLMService(db, registry) + + def build_pipeline(self, config: RAGPipelineConfig) -> RAGPipeline: + """Assemble a complete RAG pipeline from configuration. + + Sequence + -------- + 1. Ensure pipeline DB record exists (get or create) + 2. Load documents from the database + 3. Get or create the chunk set (identity via SHA-256 signature) + 4. Create the chunking model and persist chunks + 5. Setup the retriever (embedding / sparse / composite) + 6. Get or create the LLM record + 7. Resolve the prompt component and persist it + 8. Update the pipeline DB record with FK component IDs + 9. Build and return the ``RAGPipeline`` instance + + Parameters + ---------- + config : RAGPipelineConfig + Typed pipeline configuration. + + Returns + ------- + RAGPipeline + Fully assembled pipeline ready for ``generate()``. + + Raises + ------ + ValueError + If any referenced document, component or parameter is invalid. + RuntimeError + If a database error occurs. + """ + pipeline_id = self._ensure_db_record(config.session_id) + + documents = self._documents.load(config.documents) + + chunk_set = self._chunking.get_or_create_chunk_set( + config.documents, + { + "chunking_model": { + "component": config.chunking_model.component, + "params": config.chunking_model.params, + } + }, + ) + + chunking_record_id, chunking_result = self._chunking.create( + documents, + chunk_set.id, + config.chunking_model.component, + config.chunking_model.params, + ) + + retriever_service = RetrieverSetupService( + self._db, + self._registry, + self._RAG_path, + chunking_result.chunks, + chunk_set.id, + pipeline_id, + ) + retriever_result = retriever_service.setup( + config.retriever_model.component, + config.retriever_model.params, + ) + + llm_result = self._llm.get_or_create( + config.generation_model.component, + config.generation_model.params, + ) + + models_factory = RAGModelsFactory(self._registry) + prompt_result = models_factory.create_prompt( + config.prompt.component, + config.prompt.params, + ) + prompt_model = prompt_result.model + prompt_response = self._prompts.create( + class_name=config.prompt.component, + name=f"pipeline_{config.session_id}_{config.prompt.component}", + parameters=config.prompt.params, + ) + + self._update_db_record( + config.session_id, + chunking_record_id, + prompt_response.id, + llm_result.db_record_id, + ) + + pipeline = self._assemble_pipeline_instance( + config=config, + pipeline_id=pipeline_id, + documents=documents, + prompt_model=prompt_model, + chunking_model_id=chunking_record_id, + chunking_model=chunking_result.model, + chunks=chunking_result.chunks, + retriever=retriever_result.model, + llm_model=llm_result.model, + ) + return pipeline + + # ── Private helpers ─────────────────────────────────────────────── + + def _ensure_db_record(self, session_id: int) -> int: + """Get or create pipeline DB record, return its ID. + + Creates a placeholder row with nullable FK columns; these are + patched later in ``_update_db_record``. + + Parameters + ---------- + session_id : int + Generative session identifier. + + Returns + ------- + int + Primary key of the pipeline record. + """ + existing = ( + self._db.query(PipelineDBModel).filter_by(session_id=session_id).first() + ) + if existing is not None: + return existing.id + + record = PipelineDBModel( + session_id=session_id, + name="", + description=None, + parameters=None, + chunking_model_id=None, + prompt_id=None, + generation_model_id=None, + ) + self._db.add(record) + self._db.commit() + self._db.refresh(record) + return record.id + + def _update_db_record( + self, + session_id: int, + chunking_model_id: int, + prompt_id: int, + generation_model_id: int, + ) -> None: + """Patch FK columns on the pipeline record. + + Parameters + ---------- + session_id : int + Generative session identifier. + chunking_model_id : int + FK into ``RAG_chunking_model``. + prompt_id : int + FK into ``RAG_prompt``. + generation_model_id : int + FK into ``RAG_generation_model``. + + Raises + ------ + ValueError + If no pipeline record exists for the given session. + """ + record = ( + self._db.query(PipelineDBModel).filter_by(session_id=session_id).first() + ) + if record is None: + raise ValueError(f"No pipeline record for session {session_id}") + record.chunking_model_id = chunking_model_id + record.prompt_id = prompt_id + record.generation_model_id = generation_model_id + self._db.commit() + + def _assemble_pipeline_instance( + self, + config: RAGPipelineConfig, + pipeline_id: int, + documents: Dict[int, BaseDocument], + prompt_model: Prompt, + chunking_model_id: int, + chunking_model: BaseChunkingModel, + chunks: Dict[int, Dict[int, Chunk]], + retriever: RetrieverModel, + llm_model: TextToTextGenerationTaskModel, + ) -> RAGPipeline: + """Build a ``RAGPipeline`` from pre-assembled components. + + Parameters + ---------- + config : RAGPipelineConfig + pipeline_id : int + documents : Dict[int, BaseDocument] + prompt_model : Prompt + chunking_model_id : int + chunking_model : BaseChunkingModel + chunks : Dict[int, Dict[int, Chunk]] + retriever : RetrieverModel + llm_model : TextToTextGenerationTaskModel + + Returns + ------- + RAGPipeline + """ + return RAGPipeline( + config=config, + pipeline_id=pipeline_id, + chunking_model_id=chunking_model_id, + documents=documents, + chunks=chunks, + prompt_model=prompt_model, + chunking_model=chunking_model, + retriever=retriever, + llm_model=llm_model, + ) + + # ------------------------------------------------------------------ + # Parameter update validation + # ------------------------------------------------------------------ + + def validate_update_payload( + self, + new_params: dict[str, Any], + ) -> dict[str, Any]: + """Validate a RAG parameter update payload. + + Uses the already-injected ``self._db`` and ``self._registry`` + from the instance, so callers do not need to pass them again. + + Normalizes the payload, validates structure of each component ref, + checks components exist in the registry, validates documents and + prompts, and resolves ``prompt_id`` to a ``prompt`` component ref + when present. + + Parameters + ---------- + new_params : dict + Raw update payload from the request body. + + Returns + ------- + dict + Normalized and validated params dict, with ``prompt_id`` + resolved to ``prompt`` if applicable. + + Raises + ------ + ValueError + If any validation check fails. + """ + # 1. Normalize + normalized = normalize_payload(dict(new_params)) + + # 2. Validate structure of each component ref sent + for key in _RAG_MODEL_KEYS: + if key not in normalized: + continue + ref = normalized[key] + if not isinstance(ref, dict): + raise ValueError(f"'{key}' must be a dict, got {type(ref).__name__}.") + if "component" not in ref: + raise ValueError(f"Missing 'component' in '{key}'.") + if "params" not in ref: + raise ValueError(f"Missing 'params' in '{key}'.") + + # 3. Validate components exist in registry + component_errors = validate_component_refs(normalized, self._registry) + if component_errors: + raise ValueError("; ".join(component_errors)) + + # 4. Validate documents if provided + if "documents" in normalized: + docs = normalized["documents"] + if not docs: + raise ValueError("Documents list must not be empty.") + DocumentService(self._db).validate_exist(docs) + + # 5. Validate and resolve prompt if provided + prompt_service = PromptService(self._db, self._registry) + if "prompt_id" in normalized: + prompt_service.validate_prompt_exists(normalized["prompt_id"]) + normalized["prompt"] = prompt_service.resolve_prompt_id_to_component( + normalized["prompt_id"] + ) + del normalized["prompt_id"] + elif "prompt" in normalized: + prompt_service.validate_component_ref(normalized["prompt"]) + + return normalized diff --git a/DashAI/back/services/RAG/__init__.py b/DashAI/back/services/RAG/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/DashAI/back/services/RAG/chunking_service.py b/DashAI/back/services/RAG/chunking_service.py new file mode 100644 index 000000000..812bf5d6e --- /dev/null +++ b/DashAI/back/services/RAG/chunking_service.py @@ -0,0 +1,231 @@ +import hashlib +import json +import logging +from typing import Any + +from sqlalchemy import exc +from sqlalchemy.orm import Session + +from DashAI.back.dependencies.database.models import ( + Chunk as ChunkDBModel, +) +from DashAI.back.dependencies.database.models import ( + RAGChunkingModel as ChunkingModelDBModel, +) +from DashAI.back.dependencies.database.models import ( + RAGChunkSet, + RAGChunkSetDocument, +) +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.RAG.chunking_models.chunking_model_factory import ( + ChunkingFactoryResult, +) +from DashAI.back.models.RAG.documents.base_document import BaseDocument +from DashAI.back.models.RAG.documents.chunk import Chunk +from DashAI.back.models.RAG.RAG_models_factory import RAGModelsFactory + +log = logging.getLogger(__name__) + + +class ChunkingService: + """Service for chunking lifecycle: chunk set identity, model persistence.""" + + def __init__(self, db: Session, registry: ComponentRegistry): + self._db = db + self._registry = registry + + # ── Chunk set identity ── + + def _build_chunk_set_signature( + self, document_ids: list[int], pipeline_config: dict[str, Any] + ) -> str: + """Build SHA-256 signature for deterministic chunk set identity.""" + payload = json.dumps( + { + "doc_ids": sorted(document_ids), + "config": dict(sorted(pipeline_config.items())), + }, + sort_keys=True, + ) + return hashlib.sha256(payload.encode()).hexdigest() + + def get_or_create_chunk_set( + self, + document_ids: list[int], + pipeline_config: dict[str, Any], + ) -> RAGChunkSet: + """Deterministic chunk-set identity via SHA-256. + + 1. Build signature + 2. Query RAGChunkSet by signature -> return if exists (CACHE HIT) + 3. Create RAGChunkSet + RAGChunkSetDocument rows + 4. Return new chunk set + + Raises RuntimeError on DB error. + """ + signature = self._build_chunk_set_signature(document_ids, pipeline_config) + try: + existing = ( + self._db.query(RAGChunkSet).filter_by(signature=signature).first() + ) + if existing: + return existing + + chunk_set = RAGChunkSet(signature=signature, parameters=pipeline_config) + self._db.add(chunk_set) + self._db.commit() + self._db.refresh(chunk_set) + + for doc_id in sorted(document_ids): + self._db.add( + RAGChunkSetDocument( + chunk_set_id=chunk_set.id, + document_id=doc_id, + ) + ) + self._db.commit() + return chunk_set + except exc.SQLAlchemyError as e: + self._db.rollback() + log.exception(e) + raise RuntimeError("Database error during chunk set creation.") from e + + # ── Chunking model lifecycle ── + + def create( + self, + documents: dict[int, BaseDocument], + chunk_set_id: int, + component_name: str, + params: dict[str, Any], + ) -> tuple[int, ChunkingFactoryResult]: + """Lookup-or-create a chunking model and its chunks. + + 1. Sort params + 2. Query RAGChunkingModel -> if found, load chunks from DB, return + 3. If not: instantiate via pure factory, persist model + chunks, return + + Returns (db_record_id, ChunkingFactoryResult(model, chunks)). + """ + sorted_params = dict(sorted(params.items())) + + try: + existing = ( + self._db.query(ChunkingModelDBModel) + .filter_by(class_name=component_name, parameters=sorted_params) + .first() + ) + except exc.SQLAlchemyError as e: + log.exception(e) + raise RuntimeError("Database error during chunking model lookup.") from e + + if existing is not None: + chunks = self._fetch_chunks_from_db(chunk_set_id) + model_params = {**params, "documents": documents} + model_class = self._registry[component_name]["class"] + model = model_class(**model_params) + model.chunks = chunks + return existing.id, ChunkingFactoryResult(model=model, chunks=chunks) + + factory = RAGModelsFactory(self._registry) + result = factory.create_chunking_model(component_name, params, documents) + + try: + record_id = self._save_db_record(result.model) + self._persist_chunks(result.chunks, chunk_set_id) + self._update_chunk_ids(result.model, chunk_set_id) + return record_id, result + except exc.SQLAlchemyError as e: + self._db.rollback() + log.exception(e) + raise RuntimeError( + "Database error during chunking model persistence." + ) from e + + # ── Private DB helpers ── + + def _save_db_record(self, model) -> int: + """Persist a chunking model record to the database. + + Args: + model: Chunking model instance with a ``parameters`` attribute. + + Returns: + Primary key of the newly created DB record. + """ + sorted_params = dict(sorted(model.parameters.items())) + record = ChunkingModelDBModel( + class_name=model.__class__.__name__, + parameters=sorted_params, + ) + self._db.add(record) + self._db.commit() + self._db.refresh(record) + return record.id + + def _fetch_chunks_from_db(self, chunk_set_id: int) -> dict[int, dict[int, Chunk]]: + """Load all chunks for a chunk set from the database. + + Args: + chunk_set_id: Chunk set identifier. + + Returns: + Nested dict keyed by ``document_id`` then ``chunk_index``. + """ + chunks: dict[int, dict[int, Chunk]] = {} + db_chunks = ( + self._db.query(ChunkDBModel).filter_by(chunk_set_id=chunk_set_id).all() + ) + for db_chunk in db_chunks: + if db_chunk.document_id not in chunks: + chunks[db_chunk.document_id] = {} + chunk = Chunk( + id=db_chunk.id, + document_id=db_chunk.document_id, + document_position=db_chunk.chunk_index, + text=db_chunk.text, + ) + chunks[db_chunk.document_id][db_chunk.chunk_index] = chunk + return chunks + + def _persist_chunks( + self, chunks: dict[int, dict[int, Chunk]], chunk_set_id: int + ) -> None: + """Write chunk rows to the database. + + Args: + chunks: Nested dict of chunks keyed by document_id and index. + chunk_set_id: FK to the owning chunk set. + """ + for document_id, document_chunks in chunks.items(): + for idx, chunk in document_chunks.items(): + self._db.add( + ChunkDBModel( + document_id=document_id, + chunk_index=idx, + chunk_set_id=chunk_set_id, + text=chunk.text, + ) + ) + self._db.commit() + + def _update_chunk_ids(self, model, chunk_set_id: int) -> None: + """Patch in-memory chunk ids with their DB-assigned primary keys. + + Args: + model: Chunking model whose ``chunks`` dict will be mutated. + chunk_set_id: Chunk set identifier for the DB lookup. + """ + for document_id, document_chunks in model.get_chunks().items(): + for idx, _chunk in document_chunks.items(): + db_chunk = ( + self._db.query(ChunkDBModel) + .filter_by( + document_id=document_id, + chunk_index=idx, + chunk_set_id=chunk_set_id, + ) + .first() + ) + if db_chunk is not None: + model.chunks[document_id][idx].id = db_chunk.id diff --git a/DashAI/back/services/RAG/cleanup_service.py b/DashAI/back/services/RAG/cleanup_service.py new file mode 100644 index 000000000..e854327ce --- /dev/null +++ b/DashAI/back/services/RAG/cleanup_service.py @@ -0,0 +1,334 @@ +import logging +import shutil +from pathlib import Path + +from sqlalchemy.orm import Session + +from DashAI.back.dependencies.database.models import ( + GenerativeSession, + RAGChunkingModel, + RAGEmbeddingMatrix, + RAGEmbeddingModel, + RAGPipeline, + RAGRetriever, + RAGRetrieverChild, +) +from DashAI.back.services.RAG.retriever_db_service import RetrieverDBService + +log = logging.getLogger(__name__) + +COMPOSITE_RETRIEVER_NAMES = frozenset({"SequentialRetriever", "ParallelRetriever"}) + + +class CleanupService: + """Cascade deletion of RAG resources when a session is deleted or updated. + + Responsible for removing orphaned retriever models, chunking models, + embedding matrices, and storage folders when a generative session is + deleted or its parameters are updated. Order of operations is critical: + retrievers must be cleaned up BEFORE chunking models. + """ + + def __init__(self, db: Session): + self.db = db + self._retriever_db = RetrieverDBService(db) + + def cleanup_orphaned_resources( + self, + session_id: int, + old_parameters: dict, + new_parameters: dict | None = None, + ) -> None: + """Main entry point. Called on session DELETE or parameter UPDATE. + + Order MUST be: retriever cleanup BEFORE chunking cleanup. + (Retriever cleanup queries chunking_model_id from DB — if chunking + models are deleted first, the query returns None.) + + Args: + session_id: The generative session being deleted or updated. + old_parameters: The session's parameters before the change. + new_parameters: The session's parameters after the change, or + ``None`` on full session deletion. + """ + try: + if not old_parameters: + return + + def _component_changed(key: str) -> bool: + if new_parameters is None: + return True + return old_parameters.get(key) != new_parameters.get(key) + + documents_ids = sorted(old_parameters.get("documents") or []) + + # ── Retriever cleanup (MUST run BEFORE chunking) ── + retriever_model_params = old_parameters.get("retriever_model") or {} + retriever_component_name = retriever_model_params.get("component", "") + + should_cleanup_retriever = ( + bool(retriever_model_params) + and _component_changed("retriever_model") + and not self._other_sessions_with_same_config( + session_id, + old_parameters, + keys=("documents", "chunking_model", "retriever_model"), + ) + ) + + if should_cleanup_retriever: + if retriever_component_name in COMPOSITE_RETRIEVER_NAMES: + self._cleanup_composite_retriever(old_parameters, session_id) + else: + self._cleanup_unit_retriever( + old_parameters, + documents_ids, + retriever_model_params, + session_id, + ) + + # ── Chunking model cleanup (AFTER retriever) ── + chunking_model_params = old_parameters.get("chunking_model") or {} + + should_cleanup_chunking = ( + bool(chunking_model_params) + and _component_changed("chunking_model") + and not self._other_sessions_with_same_config( + session_id, + old_parameters, + keys=("documents", "chunking_model"), + ) + ) + + if should_cleanup_chunking: + chunking_models = ( + self.db.query(RAGChunkingModel) + .filter( + RAGChunkingModel.class_name + == chunking_model_params.get("component"), + RAGChunkingModel.parameters + == chunking_model_params.get("params"), + ) + .all() + ) + for chunking_model in chunking_models: + self.db.delete(chunking_model) + + self.db.commit() + except Exception: + self.db.rollback() + raise + + # ── Private helpers ── + + @staticmethod + def _delete_path(path_value: str | None) -> None: + """Delete a filesystem path recursively if it exists. + + Logs a warning if deletion fails (e.g. permission error, file in use). + + Args: + path_value: Absolute path to delete. Silently skipped if + ``None`` or the path does not exist. + """ + if not path_value: + return + path = Path(path_value) + if path.exists(): + try: + shutil.rmtree(path) + except OSError as exc: + log.warning("Failed to remove %s: %s", path_value, exc) + + def _other_sessions_with_same_config( + self, + session_id: int, + expected_parameters: dict, + *, + keys: tuple[str, ...], + ) -> bool: + """Return True if any other session matches all specified keys. + + Used to avoid deleting shared resources that another session still + depends on. + + Args: + session_id: Current session id (excluded from the check). + expected_parameters: Parameter dict to compare against. + keys: Subset of keys to compare for equality. + + Returns: + True if at least one other session shares the same config values + for all specified keys. + """ + other_sessions = ( + self.db.query(GenerativeSession) + .filter(GenerativeSession.id != session_id) + .all() + ) + for other_session in other_sessions: + other_params = other_session.parameters or {} + if all(other_params.get(k) == expected_parameters.get(k) for k in keys): + return True + return False + + def _find_pipeline_id(self, session_id: int) -> int | None: + """Find pipeline DB record ID for a session. + + Args: + session_id: Generative session identifier. + + Returns: + The pipeline record id, or ``None`` if no pipeline exists. + """ + pipeline = self.db.query(RAGPipeline).filter_by(session_id=session_id).first() + return pipeline.id if pipeline else None + + def _cleanup_composite_retriever( + self, + old_parameters: dict, + session_id: int, + ) -> None: + """Cascade-delete a composite retriever and its children. + + Finds the composite bridge record, recursively deletes child + retrievers (sparse detail storage folders, dense detail records), + then removes the parent bridge record. + + Args: + old_parameters: Session parameters containing the retriever + component name. + session_id: Generative session identifier. + """ + pipeline_id = self._find_pipeline_id(session_id) + if pipeline_id is None: + return + + retriever_component_name = old_parameters.get("retriever_model", {}).get( + "component" + ) + + composite_bridges = ( + self.db.query(RAGRetriever) + .filter( + RAGRetriever.pipeline_id == pipeline_id, + RAGRetriever.class_name == retriever_component_name, + ) + .all() + ) + + for bridge in composite_bridges: + child_links = ( + self.db.query(RAGRetrieverChild) + .filter_by(parent_id=bridge.id) + .order_by(RAGRetrieverChild.child_order) + .all() + ) + for link in child_links: + child_bridge = self.db.query(RAGRetriever).get(link.child_id) + if child_bridge is None: + continue + if child_bridge.sparse_detail: + self._delete_path(child_bridge.sparse_detail.storage_folder) + self.db.delete(child_bridge.sparse_detail) + elif child_bridge.dense_detail: + self.db.delete(child_bridge.dense_detail) + self.db.delete(bridge) + + def _cleanup_unit_retriever( + self, + old_parameters: dict, + documents_ids: list, + retriever_model_params: dict, + session_id: int, + ) -> None: + """Cascade-delete a unit retriever (dense or sparse). + + Determines retriever type by inspecting the bridge relationship + (``dense_detail`` vs ``sparse_detail``) rather than parameter heuristics. + + Args: + old_parameters: Session parameters dict. + documents_ids: Document ids associated with the session. + retriever_model_params: The ``retriever_model`` parameter dict. + session_id: Generative session identifier. + """ + pipeline_id = self._find_pipeline_id(session_id) + if pipeline_id is None: + return + + bridge = ( + self.db.query(RAGRetriever) + .filter( + RAGRetriever.pipeline_id == pipeline_id, + RAGRetriever.class_name == retriever_model_params.get("component"), + ) + .first() + ) + if bridge is None: + return + + if bridge.dense_detail is not None: + self._cleanup_dense_retriever(bridge, documents_ids) + elif bridge.sparse_detail is not None: + self._cleanup_sparse_retriever(bridge) + self.db.delete(bridge) + + def _cleanup_dense_retriever( + self, + bridge: RAGRetriever, + documents_ids: list, + ) -> None: + """Cascade-delete a dense retriever and its embedding resources. + + Removes embedding matrix storage folders, deletes matrix DB records, + the embedding model record, and the dense detail record. + + Args: + bridge: The bridge record for the dense retriever. + documents_ids: Document ids whose embeddings to delete. + """ + dense_retriever = bridge.dense_detail + chunk_set_id = dense_retriever.chunk_set_id + embedding_model_id = dense_retriever.embedding_model_id + + embedding_matrices = ( + self.db.query(RAGEmbeddingMatrix) + .filter( + RAGEmbeddingMatrix.chunk_set_id == chunk_set_id, + RAGEmbeddingMatrix.embedding_model_id == embedding_model_id, + RAGEmbeddingMatrix.document_id.in_(documents_ids), + ) + .all() + ) + + matrix_ids = [] + for matrix in embedding_matrices: + self._delete_path(matrix.storage_folder) + matrix_ids.append(matrix.id) + + if matrix_ids: + self.db.query(RAGEmbeddingMatrix).filter( + RAGEmbeddingMatrix.id.in_(matrix_ids) + ).delete(synchronize_session="fetch") + + embedding_model = self.db.query(RAGEmbeddingModel).get(embedding_model_id) + if embedding_model is not None: + self.db.delete(embedding_model) + + self.db.delete(dense_retriever) + + def _cleanup_sparse_retriever( + self, + bridge: RAGRetriever, + ) -> None: + """Cascade-delete a sparse retriever and its storage folder. + + Removes the on-disk storage folder and the sparse detail DB record. + + Args: + bridge: The bridge record for the sparse retriever. + """ + sparse_retriever = bridge.sparse_detail + self._delete_path(sparse_retriever.storage_folder) + self.db.delete(sparse_retriever) diff --git a/DashAI/back/services/RAG/document_service.py b/DashAI/back/services/RAG/document_service.py new file mode 100644 index 000000000..e6955785c --- /dev/null +++ b/DashAI/back/services/RAG/document_service.py @@ -0,0 +1,439 @@ +import logging +import mimetypes +import os +from datetime import datetime +from typing import Dict, List, Tuple + +from sqlalchemy import exc +from sqlalchemy.orm import Session + +from DashAI.back.api.api_v1.schemas import DocumentResponse +from DashAI.back.dependencies.database.models import ( + Document as DocumentDBModel, +) +from DashAI.back.dependencies.database.models import ( + GenerativeSession, +) +from DashAI.back.models.RAG.documents import ( + BaseDocument, + DocumentFileType, + PDFDocument, + TxtDocument, +) +from DashAI.back.models.RAG.exceptions import RAGDocumentFileTypeError +from DashAI.back.models.RAG.utils import hash_function + +log = logging.getLogger(__name__) + +_DOCUMENT_CLASSES: dict[DocumentFileType, type[BaseDocument]] = { + DocumentFileType.TXT: TxtDocument, + DocumentFileType.PDF: PDFDocument, + DocumentFileType.MD: TxtDocument, + DocumentFileType.RST: TxtDocument, + DocumentFileType.TEX: TxtDocument, + # CSV, MD, RST, TEX are parsed as plain text via TxtDocument. + # This is a limitation: CSV files are structured data, not free text. + # A future improvement would add a dedicated CsvDocument parser. + DocumentFileType.CSV: TxtDocument, +} + + +class DocumentService: + """Service layer for document CRUD, file storage, and hydration.""" + + def __init__(self, db: Session): + self.db = db + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _to_response( + self, doc: DocumentDBModel, base_url: str = "" + ) -> DocumentResponse: + return DocumentResponse( + id=doc.id, + file_name=doc.file_name, + file_type=doc.file_type, + file_hash=doc.file_hash, + created=doc.created, + last_modified=doc.last_modified, + optional_metadata=doc.optional_metadata, + related_sessions=[s.id for s in doc.get_related_sessions] + if doc.get_related_sessions + else None, + file_url=f"{base_url}/api/v1/document/{doc.id}/download", + ) + + def _get_document_or_raise(self, document_id: int) -> DocumentDBModel: + doc = self.db.get(DocumentDBModel, document_id) + if doc is None: + raise ValueError(f"Document with ID {document_id} does not exist.") + return doc + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def upload( + self, + file_content: bytes, + file_name: str, + file_type: str | DocumentFileType, + docs_path: str, + optional_metadata: dict = None, + ) -> DocumentResponse: + """Upload a document. + + Handles hash deduplication, file storage, and DB record creation. + + Parameters + ---------- + file_content : bytes + Raw file bytes. + file_name : str + Original file name. + file_type : str | DocumentFileType + File extension / type, e.g. ``DocumentFileType.PDF`` or ``"pdf"``. + docs_path : str + Directory on disk where the file will be written. + optional_metadata : dict, optional + Arbitrary metadata attached to the document. + + Returns + ------- + DocumentResponse + The created or updated document representation. + + Raises + ------ + ValueError + If ``docs_path`` does not exist or a database error occurs. + """ + if isinstance(file_type, DocumentFileType): + file_type = file_type.value + if not os.path.isdir(docs_path): + raise ValueError(f"Documents folder does not exist: {docs_path}") + + optional_metadata = optional_metadata or {} + file_content_hash = hash_function(file_content) + file_path = os.path.join(docs_path, file_name) + + try: + existing = ( + self.db.query(DocumentDBModel) + .filter_by(file_hash=file_content_hash) + .first() + ) + if existing: + existing.file_name = file_name + existing.file_path = file_path + existing.optional_metadata = optional_metadata + self.db.commit() + return self._to_response(existing) + + with open(file_path, "wb") as f: + f.write(file_content) + + doc = DocumentDBModel( + file_name=file_name, + file_type=file_type, + file_path=file_path, + file_hash=file_content_hash, + optional_metadata=optional_metadata or None, + ) + self.db.add(doc) + self.db.commit() + self.db.refresh(doc) + return self._to_response(doc) + + except exc.SQLAlchemyError as e: + log.exception(e) + raise ValueError("Database error during document upload.") from e + + def get(self, document_id: int) -> DocumentResponse: + """Get document metadata by ID. + + Parameters + ---------- + document_id : int + + Returns + ------- + DocumentResponse + + Raises + ------ + ValueError + If the document does not exist. + """ + try: + doc = self._get_document_or_raise(document_id) + return self._to_response(doc) + except exc.SQLAlchemyError as e: + log.exception(e) + raise ValueError("Database error retrieving document.") from e + + def get_all(self, base_url: str = "") -> List[DocumentResponse]: + """Get all documents with ``file_url`` included. + + Parameters + ---------- + base_url : str + Base URL prefix for download links. + + Returns + ------- + list[DocumentResponse] + """ + try: + docs: List[DocumentDBModel] = self.db.query(DocumentDBModel).all() + return [self._to_response(d, base_url) for d in docs] + except exc.SQLAlchemyError as e: + log.exception(e) + raise ValueError("Database error listing documents.") from e + + def get_by_session( + self, session_id: int, base_url: str = "" + ) -> List[DocumentResponse]: + """Get documents linked to a generative session. + + Documents are identified from ``session.parameters["documents"]``. + + Parameters + ---------- + session_id : int + base_url : str + + Returns + ------- + list[DocumentResponse] + """ + try: + session = self.db.get(GenerativeSession, session_id) + if session is None: + raise ValueError(f"GenerativeSession with ID {session_id} not found.") + + document_ids: List[int] = session.parameters.get("documents", []) + if not document_ids: + return [] + + docs = ( + self.db.query(DocumentDBModel) + .filter(DocumentDBModel.id.in_(document_ids)) + .all() + ) + return [self._to_response(d, base_url) for d in docs] + + except exc.SQLAlchemyError as e: + log.exception(e) + raise ValueError("Database error retrieving session documents.") from e + + def delete(self, document_id: int) -> None: + """Delete a document file from disk and its DB record. + + Parameters + ---------- + document_id : int + + Raises + ------ + ValueError + If the document does not exist. + """ + try: + doc = self._get_document_or_raise(document_id) + + if os.path.exists(doc.file_path): + os.remove(doc.file_path) + + self.db.delete(doc) + self.db.commit() + + except exc.SQLAlchemyError as e: + log.exception(e) + raise ValueError("Database error deleting document.") from e + except OSError as e: + log.exception(e) + raise ValueError("Error deleting physical file.") from e + + def update_metadata( + self, + document_id: int, + file_name: str = None, + optional_metadata: dict = None, + ) -> DocumentResponse: + """Update document metadata (``file_name``, ``optional_metadata``). + + Parameters + ---------- + document_id : int + file_name : str, optional + New file name. Also updates ``file_type`` from the extension. + optional_metadata : dict, optional + + Returns + ------- + DocumentResponse + + Raises + ------ + ValueError + If the document does not exist. + """ + try: + doc = self._get_document_or_raise(document_id) + + if file_name is not None: + doc.file_name = file_name + ext = os.path.splitext(file_name)[1].lstrip(".") + try: + doc.file_type = DocumentFileType(ext).value + except ValueError as err: + raise RAGDocumentFileTypeError( + f"Unsupported file type: {ext}" + ) from err + + if optional_metadata is not None: + doc.optional_metadata = optional_metadata + + doc.last_modified = datetime.now() + self.db.commit() + self.db.refresh(doc) + return self._to_response(doc) + + except exc.SQLAlchemyError as e: + log.exception(e) + raise ValueError("Database error updating document metadata.") from e + + def download(self, document_id: int) -> Tuple[bytes, str, str]: + """Return file content, media type, and filename for download. + + Parameters + ---------- + document_id : int + + Returns + ------- + tuple[bytes, str, str] + ``(file_content, media_type, filename)``. + + Raises + ------ + ValueError + If the document or its physical file is not found. + """ + try: + doc = self._get_document_or_raise(document_id) + + if not os.path.exists(doc.file_path): + raise ValueError(f"File not found on disk: {doc.file_path}") + + ext = os.path.splitext(doc.file_name)[1].lower() + media_type, _ = mimetypes.guess_type(doc.file_name) + if media_type is None: + media_type = { + ".txt": "text/plain", + ".pdf": "application/pdf", + }.get(ext, "application/octet-stream") + + with open(doc.file_path, "rb") as f: + content = f.read() + + return content, media_type, doc.file_name + + except exc.SQLAlchemyError as e: + log.exception(e) + raise ValueError("Database error during document download.") from e + + def load(self, document_ids: List[int]) -> Dict[int, BaseDocument]: + """Load and hydrate DB document rows into ``BaseDocument`` instances. + + Parameters + ---------- + document_ids : list[int] + + Returns + ------- + dict[int, BaseDocument] + Mapping from document ID to hydrated document object. + + Raises + ------ + ValueError + If any document ID is not found or the file type is unsupported. + """ + documents: Dict[int, BaseDocument] = {} + for doc_id in document_ids: + db_doc: DocumentDBModel = ( + self.db.query(DocumentDBModel) + .filter(DocumentDBModel.id == doc_id) + .first() + ) + if db_doc is None: + raise ValueError(f"Document with ID {doc_id} not found in database.") + try: + doc_class = _DOCUMENT_CLASSES[DocumentFileType(db_doc.file_type)] + except (KeyError, ValueError) as err: + supported = ", ".join(e.value for e in DocumentFileType) + raise ValueError( + f"Unsupported file type '{db_doc.file_type}'. " + f"Supported types: {supported}." + ) from err + documents[doc_id] = doc_class( + id=db_doc.id, + file_name=db_doc.file_name, + file_path=db_doc.file_path, + file_hash=db_doc.file_hash, + created=db_doc.created, + optional_metadata=db_doc.optional_metadata, + ) + return documents + + def validate_exist(self, document_ids: List[int]) -> None: + """Raise ``ValueError`` if any document ID does not exist in the DB. + + Parameters + ---------- + document_ids : list[int] + + Raises + ------ + ValueError + If one or more document IDs are not found. + """ + existing = ( + self.db.query(DocumentDBModel.id) + .filter(DocumentDBModel.id.in_(document_ids)) + .all() + ) + existing_ids = {row.id for row in existing} + missing = [str(i) for i in document_ids if i not in existing_ids] + if missing: + raise ValueError(f"Documents with IDs {', '.join(missing)} not found.") + + def get_related_sessions(self, document_id: int) -> List[int]: + """Get session IDs linked to a document. + + Parameters + ---------- + document_id : int + + Returns + ------- + list[int] + Session IDs related to the document. + + Raises + ------ + ValueError + If the document is not found. + """ + try: + doc = self._get_document_or_raise(document_id) + if not doc.get_related_sessions: + return [] + return [s.id for s in doc.get_related_sessions] + except exc.SQLAlchemyError as e: + log.exception(e) + raise ValueError("Database error retrieving related sessions.") from e diff --git a/DashAI/back/services/RAG/embedding_storage_service.py b/DashAI/back/services/RAG/embedding_storage_service.py new file mode 100644 index 000000000..838f73179 --- /dev/null +++ b/DashAI/back/services/RAG/embedding_storage_service.py @@ -0,0 +1,227 @@ +"""Service for managing embedding matrix persistence on disk and in DB.""" + +import logging +import os +import shutil +from typing import List + +import numpy as np + +from DashAI.back.dependencies.database.models import RAGEmbeddingMatrix +from DashAI.back.models.RAG.exceptions import RAGEmbeddingLoadError +from DashAI.back.services.RAG.retriever_db_service import RetrieverDBService + +logger = logging.getLogger(__name__) + +EMBEDDINGS_DIRNAME = "embeddings" +EMBEDDINGS_FILENAME = "embeddings.npy" + + +class EmbeddingStorageService: + """Manages embedding matrix persistence: numpy .npy files on disk + DB records. + + Embedding matrices are stored as .npy files in directories like: + ``{env_RAG_path}/embeddings/doc_id-{doc_id}__chunk_set_id-{chunk_set_id}__ + embedding_model_id-{embedding_model_id}/embeddings.npy`` + """ + + def __init__(self, env_RAG_path: str, db_service: RetrieverDBService): # noqa: N803 + """Initialise with the base RAG directory path and DB service. + + Parameters + ---------- + env_RAG_path : str + Base RAG directory path from config. + db_service : RetrieverDBService + Service for database operations on retriever models. + """ + self._env_RAG_path = env_RAG_path + self._db_service = db_service + + # ------------------------------------------------------------------ + # Path helpers + # ------------------------------------------------------------------ + + @staticmethod + def build_matrix_dir_name( + doc_id: int, chunk_set_id: int, embedding_model_id: int + ) -> str: + return ( + f"doc_id-{doc_id}__chunk_set_id-{chunk_set_id}__" + f"embedding_model_id-{embedding_model_id}" + ) + + def _matrix_dir( + self, doc_id: int, chunk_set_id: int, embedding_model_id: int + ) -> str: + """Build the directory path for a specific embedding matrix. + + Format: ``{env_RAG_path}/embeddings/doc_id-{doc_id}__ + chunk_set_id-{chunk_set_id}__embedding_model_id-{embedding_model_id}`` + """ + return os.path.join( + self._env_RAG_path, + EMBEDDINGS_DIRNAME, + self.build_matrix_dir_name(doc_id, chunk_set_id, embedding_model_id), + ) + + def _matrix_path( + self, doc_id: int, chunk_set_id: int, embedding_model_id: int + ) -> str: + """Full path to the ``embeddings.npy`` file. + + Delegates to :meth:`_matrix_dir` and appends the filename. + """ + return os.path.join( + self._matrix_dir(doc_id, chunk_set_id, embedding_model_id), + EMBEDDINGS_FILENAME, + ) + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def save_embeddings( + self, + doc_id: int, + chunk_set_id: int, + embedding_model_id: int, + embeddings: np.ndarray, + ) -> RAGEmbeddingMatrix: + """Save embeddings to disk and create / return a DB record. + + Steps + ----- + 1. Create the directory via :meth:`_matrix_dir`. + 2. Write the array to ``embeddings.npy`` with :func:`numpy.save`. + 3. Persist a matching ``RAGEmbeddingMatrix`` record through + :meth:`RetrieverDBService.save_embedding_matrix`. + + Parameters + ---------- + doc_id : int + Owning document id. + chunk_set_id : int + Chunk set identifier. + embedding_model_id : int + Embedding model identifier. + embeddings : np.ndarray + 2-D embedding matrix (num_chunks x embedding_dimension). + + Returns + ------- + RAGEmbeddingMatrix + The newly persisted database record. + """ + matrix_dir = self._matrix_dir(doc_id, chunk_set_id, embedding_model_id) + os.makedirs(matrix_dir, exist_ok=True) + + matrix_path = self._matrix_path(doc_id, chunk_set_id, embedding_model_id) + np.save(matrix_path, embeddings) + + record = self._db_service.save_embedding_matrix( + document_id=doc_id, + chunk_set_id=chunk_set_id, + embedding_model_id=embedding_model_id, + storage_folder=matrix_dir, + matrix_shape=list(embeddings.shape), + ) + return record + + def load_embeddings( + self, doc_id: int, chunk_set_id: int, embedding_model_id: int + ) -> np.ndarray | None: + """Load an embedding matrix from disk as a numpy array. + + Returns ``None`` when the file does not exist (not an error). + + Parameters + ---------- + doc_id : int + Owning document id. + chunk_set_id : int + Chunk set identifier. + embedding_model_id : int + Embedding model identifier. + + Returns + ------- + np.ndarray or None + The loaded embedding matrix, or ``None`` if not found. + """ + matrix_path = self._matrix_path(doc_id, chunk_set_id, embedding_model_id) + if not os.path.isfile(matrix_path): + return None + try: + return np.load(matrix_path) + except (IOError, OSError, ValueError) as exc: + logger.warning("Failed to load embeddings from %s: %s", matrix_path, exc) + raise RAGEmbeddingLoadError( + f"Failed to load embeddings from {matrix_path}: {exc}" + ) from exc + + def exists(self, doc_id: int, chunk_set_id: int, embedding_model_id: int) -> bool: + """Check whether an ``embeddings.npy`` file exists on disk. + + Parameters + ---------- + doc_id : int + Owning document id. + chunk_set_id : int + Chunk set identifier. + embedding_model_id : int + Embedding model identifier. + + Returns + ------- + bool + ``True`` if the file exists, ``False`` otherwise. + """ + return os.path.isfile( + self._matrix_path(doc_id, chunk_set_id, embedding_model_id) + ) + + def delete_all_for_document( + self, doc_id: int, matrices: List[RAGEmbeddingMatrix] + ) -> None: + """Delete embedding directories and DB records for a document. + + For each provided matrix record, the corresponding on-disk directory + is removed recursively. DB records are bulk-deleted via + :meth:`RetrieverDBService.delete_embedding_matrices_by_ids`. + + Parameters + ---------- + doc_id : int + Document whose embeddings are being deleted. + matrices : list[RAGEmbeddingMatrix] + Embedding matrix records to delete. + """ + matrix_ids = [] + for m in matrices: + if m.document_id == doc_id: + matrix_ids.append(m.id) + self.delete_storage(m.storage_folder) + if matrix_ids: + self._db_service.delete_embedding_matrices_by_ids(matrix_ids) + + @staticmethod + def delete_storage(storage_folder: str) -> None: + """Recursively delete a storage directory if it exists. + + Uses :func:`shutil.rmtree` and silently ignores missing folders. + + Parameters + ---------- + storage_folder : str + Absolute path to the directory to remove. + """ + if os.path.isdir(storage_folder): + try: + shutil.rmtree(storage_folder) + except OSError as exc: + logger.warning( + "Failed to remove storage folder %s: %s", + storage_folder, + exc, + ) diff --git a/DashAI/back/services/RAG/exceptions.py b/DashAI/back/services/RAG/exceptions.py new file mode 100644 index 000000000..c5a341d40 --- /dev/null +++ b/DashAI/back/services/RAG/exceptions.py @@ -0,0 +1,20 @@ +"""Service-layer RAG exceptions. + +Re-exports exception types from the unified RAG exception hierarchy +(:mod:`DashAI.back.models.RAG.exceptions`) so that service-layer code +does not import directly from the models layer. + +All RAG exception types are defined in a single hierarchy under +:mod:`DashAI.back.models.RAG.exceptions`. This module re-exports +the subset needed by services. +""" + +from DashAI.back.models.RAG.exceptions import ( # noqa: F401 + RAGDatabaseError, + RAGPromptValidationError, +) + +__all__ = [ + "RAGDatabaseError", + "RAGPromptValidationError", +] diff --git a/DashAI/back/services/RAG/llm_service.py b/DashAI/back/services/RAG/llm_service.py new file mode 100644 index 000000000..c5d8995f4 --- /dev/null +++ b/DashAI/back/services/RAG/llm_service.py @@ -0,0 +1,125 @@ +"""Service for LLM lookup-or-create DB operations. + +Delegates model instantiation to the component registry, +keeping DB persistence concerns separate from model construction. +""" + +import logging +from dataclasses import dataclass +from typing import Any + +from sqlalchemy import exc +from sqlalchemy.orm import Session + +from DashAI.back.dependencies.database.models import ( + RAGGenerationModel as GenerationDBModel, +) +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.RAG.exceptions import RAGGenerationModelError +from DashAI.back.models.RAG.RAG_models_factory import RAGModelsFactory +from DashAI.back.models.text_to_text_generation_model import ( + TextToTextGenerationTaskModel, +) + +log = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class LLMServiceResult: + """Result of LLM lookup-or-create from LLMService. + + Includes the DB record id alongside the instantiated model. + """ + + db_record_id: int + model: "TextToTextGenerationTaskModel" + + +class LLMService: + """Handles lookup-or-create of LLM records in the RAGGenerationModel table.""" + + def __init__(self, db: Session, registry: ComponentRegistry): + self._db = db + self._registry = registry + + def get_or_create( + self, component_name: str, params: dict[str, Any] + ) -> LLMServiceResult: + """Lookup-or-create an LLM record. + + Steps: + 1. Sort params for deterministic DB query. + 2. Query RAGGenerationModel by (class_name, parameters). + 3. If found — instantiate model from registry with existing params. + 4. If not found — create DB record, instantiate model from registry. + + Args: + component_name: Registered LLM component name. + params: LLM configuration parameters. + + Returns: + An LLMServiceResult with the DB record id and the model instance. + + Raises: + RuntimeError: If a database error occurs during lookup or creation. + RAGGenerationModelError: If the registered component is not a + TextToTextGenerationTaskModel subclass. + """ + sorted_params = dict(sorted(params.items())) + + try: + existing = ( + self._db.query(GenerationDBModel) + .filter_by(class_name=component_name, parameters=sorted_params) + .first() + ) + except exc.SQLAlchemyError as e: + log.exception(e) + raise RuntimeError("Database error during LLM lookup.") from e + + factory = RAGModelsFactory(self._registry) + + if existing is not None: + model_class = self._registry[existing.class_name]["class"] + if not issubclass(model_class, TextToTextGenerationTaskModel): + raise RAGGenerationModelError( + f"Component {existing.class_name} is not a " + f"TextToTextGenerationTaskModel subclass." + ) + result = factory.create_llm(existing.class_name, existing.parameters) + return LLMServiceResult( + db_record_id=existing.id, + model=result.model, + ) + + try: + db_record = GenerationDBModel( + class_name=component_name, + parameters=sorted_params, + ) + self._db.add(db_record) + except exc.SQLAlchemyError as e: + self._db.rollback() + log.exception(e) + raise RuntimeError("Database error during LLM creation.") from e + + result = factory.create_llm(component_name, sorted_params) + model = result.model + model_class = model.__class__ + if not issubclass(model_class, TextToTextGenerationTaskModel): + raise RAGGenerationModelError( + f"Component {component_name} is not a " + f"TextToTextGenerationTaskModel subclass." + ) + + try: + self._db.commit() + self._db.refresh(db_record) + except exc.SQLAlchemyError as e: + self._db.rollback() + log.exception(e) + raise RuntimeError("Database error during LLM creation.") from e + return LLMServiceResult( + db_record_id=db_record.id, + model=model, + ) diff --git a/DashAI/back/services/RAG/prompt_service.py b/DashAI/back/services/RAG/prompt_service.py new file mode 100644 index 000000000..639787e83 --- /dev/null +++ b/DashAI/back/services/RAG/prompt_service.py @@ -0,0 +1,499 @@ +import logging +from dataclasses import dataclass +from datetime import datetime +from typing import Any + +from sqlalchemy import exc, select +from sqlalchemy.orm import Session + +from DashAI.back.dependencies.database.models import ( + GenerativeSession, + GenerativeSessionParameterHistory, + RAGPrompt, +) +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.RAG.prompts.generation.default_QA_RAG_generation_prompt import ( + DefaultQARAGGenerationPrompt, +) +from DashAI.back.models.RAG.prompts.generation.default_RAG_generation_prompt import ( + DefaultRAGGenerationPrompt, +) +from DashAI.back.models.RAG.prompts.prompt import Prompt +from DashAI.back.services.RAG.exceptions import ( + RAGDatabaseError, + RAGPromptValidationError, +) + +log = logging.getLogger(__name__) + + +@dataclass +class PromptResponse: + id: int + class_name: str + name: str | None + parameters: dict[str, Any] | None + created: datetime + last_modified: datetime + + +class PromptService: + """Service layer for RAG prompt CRUD, validation, and session-scoped copies.""" + + def __init__(self, db: Session, registry: ComponentRegistry): + self.db = db + self._registry = registry + + def create( + self, class_name: str, name: str, parameters: dict[str, Any] + ) -> PromptResponse: + """Create a RAGPrompt record. + + Validates templates via ``_validate_prompt_template()``, inserts the + DB record, and returns the serialized prompt response. + + Args: + class_name: Registered prompt component name. + name: Human-readable name for the prompt. + parameters: Prompt configuration including ``template`` or + ``templates``. + + Returns: + The newly created prompt response. + + Raises: + RAGPromptValidationError: If validation fails. + RAGDatabaseError: If a database error occurs. + """ + if class_name not in self._registry: + raise RAGPromptValidationError( + f"Component {class_name} is not registered in the registry." + ) + prompt_class = self._registry[class_name]["class"] + if not issubclass(prompt_class, Prompt): + raise RAGPromptValidationError( + f"Component {class_name} is not a valid Prompt subclass." + ) + + if "templates" in parameters: + for _lang, tmpl in parameters["templates"].items(): + self._validate_prompt_template(class_name, tmpl) + elif "template" in parameters: + self._validate_prompt_template(class_name, parameters["template"]) + else: + raise RAGPromptValidationError( + "Prompt parameters must include 'template' or 'templates'." + ) + + try: + prompt = RAGPrompt( + class_name=class_name, + name=name, + parameters=parameters, + ) + self.db.add(prompt) + self.db.commit() + self.db.refresh(prompt) + return self._serialize_prompt(prompt) + except exc.SQLAlchemyError as e: + self.db.rollback() + log.exception(e) + raise RAGDatabaseError("Error creating prompt in database.") from e + + def update( + self, + prompt_id: int, + name: str | None = None, + parameters: dict[str, Any] | None = None, + ) -> PromptResponse: + """Update an existing prompt in place. + + Validates the template if parameters change. + + Args: + prompt_id: Primary key of the prompt to update. + name: New name (optional). + parameters: New parameters including ``template`` or + ``templates`` (optional). + + Returns: + The updated prompt response. + + Raises: + RAGPromptValidationError: If not found or validation fails. + RAGDatabaseError: If a database error occurs. + """ + try: + prompt = self.db.get(RAGPrompt, prompt_id) + except exc.SQLAlchemyError as e: + log.exception(e) + raise RAGDatabaseError("Error retrieving prompt from database.") from e + + if prompt is None: + raise RAGPromptValidationError( + f"Prompt with ID {prompt_id} does not exist." + ) + + changed = False + + if name is not None: + name = name.strip() + if not name: + raise RAGPromptValidationError("Prompt name cannot be empty.") + if name != prompt.name: + prompt.name = name + changed = True + + if parameters is not None: + if prompt.class_name not in self._registry: + raise RAGPromptValidationError( + f"Component {prompt.class_name} is not registered in the registry." + ) + prompt_class = self._registry[prompt.class_name]["class"] + if not issubclass(prompt_class, Prompt): + raise RAGPromptValidationError( + f"Component {prompt.class_name} is not a valid Prompt subclass." + ) + + if "templates" in parameters: + for _lang, tmpl in parameters["templates"].items(): + self._validate_prompt_template(prompt.class_name, tmpl) + elif "template" in parameters: + self._validate_prompt_template( + prompt.class_name, parameters["template"] + ) + else: + raise RAGPromptValidationError( + "Prompt parameters must include 'template' or 'templates'." + ) + prompt.parameters = parameters + changed = True + + if not changed: + return self._serialize_prompt(prompt) + + try: + self.db.commit() + self.db.refresh(prompt) + return self._serialize_prompt(prompt) + except exc.SQLAlchemyError as e: + self.db.rollback() + log.exception(e) + raise RAGDatabaseError("Error updating prompt in database.") from e + + def get_all(self) -> list[PromptResponse]: + """Get all prompts. + + Seeds default prompts (DefaultRAGGenerationPrompt and + DefaultQARAGGenerationPrompt) if the DB is empty. + + Returns: + A list of serialized prompt responses. + + Raises: + RAGDatabaseError: If a database error occurs. + """ + try: + prompts = self.db.query(RAGPrompt).all() + + if len(prompts) == 0: + default_generation_prompt = RAGPrompt( + class_name=DefaultRAGGenerationPrompt.__name__, + name="Default RAG Generation Prompt", + parameters={ + "templates": DefaultRAGGenerationPrompt.metadata["templates"], + "language": "en", + }, + ) + default_qa_prompt = RAGPrompt( + class_name=DefaultQARAGGenerationPrompt.__name__, + name="Default QA RAG Generation Prompt", + parameters={ + "templates": DefaultQARAGGenerationPrompt.metadata["templates"], + "language": "en", + }, + ) + self.db.add(default_generation_prompt) + self.db.add(default_qa_prompt) + self.db.commit() + prompts = self.db.query(RAGPrompt).all() + + return [self._serialize_prompt(p) for p in prompts] + + except exc.SQLAlchemyError as e: + log.exception(e) + raise RAGDatabaseError("Error listing prompts in database.") from e + + def create_session_copy( + self, + prompt_id: int, + session_id: int, + parameters: dict[str, Any] | None = None, + name: str | None = None, + ) -> PromptResponse: + """Create a session-scoped copy of a prompt. + + Adds ``cloned_for_session`` to the parameters dict to avoid UNIQUE + constraint collisions. Generates a unique name and updates the + session's parameters dict with the new prompt_id. + + Args: + prompt_id: Primary key of the prompt to copy. + session_id: Target session id. + parameters: Override parameters for the copy (optional). + name: Override name for the copy (optional). + + Returns: + The newly created prompt response. + + Raises: + RAGPromptValidationError: If the prompt or session does not exist. + RAGDatabaseError: If a database error occurs. + """ + try: + existing_prompt = self.db.get(RAGPrompt, prompt_id) + except exc.SQLAlchemyError as e: + log.exception(e) + raise RAGDatabaseError("Error retrieving prompt from database.") from e + + if existing_prompt is None: + raise RAGPromptValidationError( + f"Prompt with ID {prompt_id} does not exist." + ) + + try: + session = self.db.get(GenerativeSession, session_id) + except exc.SQLAlchemyError as e: + log.exception(e) + raise RAGDatabaseError("Error retrieving session from database.") from e + + if session is None: + raise RAGPromptValidationError( + f"GenerativeSession with ID {session_id} not found." + ) + + if parameters is not None: + if "templates" in parameters: + for _lang, tmpl in parameters["templates"].items(): + self._validate_prompt_template(existing_prompt.class_name, tmpl) + elif "template" in parameters: + self._validate_prompt_template( + existing_prompt.class_name, parameters["template"] + ) + else: + raise RAGPromptValidationError( + "Prompt parameters must include 'template' or 'templates'." + ) + new_parameters = parameters + else: + new_parameters = existing_prompt.parameters + + new_parameters = dict(new_parameters or {}) + new_parameters["cloned_for_session"] = session_id + + base_name = (name or existing_prompt.name or existing_prompt.class_name).strip() + new_name = self._build_session_prompt_name(base_name, session_id) + + try: + new_prompt = RAGPrompt( + class_name=existing_prompt.class_name, + name=new_name, + parameters=new_parameters, + ) + self.db.add(new_prompt) + self.db.commit() + self.db.refresh(new_prompt) + + session_parameters = dict(session.parameters or {}) + session_parameters["prompt_id"] = new_prompt.id + session.parameters = session_parameters + session.last_modified = datetime.now() + self.db.add( + GenerativeSessionParameterHistory( + session_id=session.id, + parameters=session_parameters, + modified_at=datetime.now(), + ) + ) + self.db.commit() + self.db.refresh(session) + + return self._serialize_prompt(new_prompt) + + except exc.SQLAlchemyError as e: + self.db.rollback() + log.exception(e) + raise RAGDatabaseError("Error creating session copy in database.") from e + + def validate_template(self, class_name: str, template: str) -> None: + """Validate a template against the prompt class's required placeholders. + + Parameters + ---------- + class_name : str + Registered prompt component name. + template : str + The template string to validate. + + Raises + ------ + ValueError + If the component is not registered, is not a Prompt subclass, + or the template is missing required placeholders. + """ + self._validate_prompt_template(class_name, template) + + def resolve_prompt_id_to_component(self, prompt_id: int) -> dict[str, Any]: + """Resolve a prompt_id to a component configuration dict. + + Args: + prompt_id: Primary key of the prompt. + + Returns: + ``{"component": class_name, "params": {"template": ..., "language": ...}}`` + suitable for pipeline configuration. + + Raises: + RAGPromptValidationError: If the prompt does not exist or has no + usable template. + RAGDatabaseError: If a database error occurs. + """ + try: + prompt = self.db.get(RAGPrompt, prompt_id) + except exc.SQLAlchemyError as e: + log.exception(e) + raise RAGDatabaseError("Error retrieving prompt from database.") from e + + if prompt is None: + raise RAGPromptValidationError( + f"Prompt with ID {prompt_id} does not exist." + ) + + params = prompt.parameters or {} + template = params.get("template") + if not template: + if "templates" not in params: + raise RAGPromptValidationError( + f"Prompt {prompt_id} has no 'template' or 'templates'" + ) + language = params.get("language", "en") + templates = params["templates"] + template = templates.get(language) + if not template: + template = next(iter(templates.values()), None) + if not template: + raise RAGPromptValidationError(f"Prompt {prompt_id} has no usable template") + + language = params.get("language", "en") + + return { + "component": prompt.class_name, + "params": { + "template": template, + "language": language, + }, + } + + def validate_component_ref(self, prompt_ref: dict[str, Any]) -> None: + """Validate a prompt component reference. + + Validates that: + - The prompt reference contains a registered ``component`` name. + - If a ``template`` is provided, it contains all required placeholders. + - If no template is provided, the component has a built-in default. + + Parameters + ---------- + prompt_ref : dict + ``{"component": "...", "params": {"template": "...", ...}}`` + + Raises + ------ + RAGPromptValidationError + If validation fails. + """ + if "component" not in prompt_ref: + raise RAGPromptValidationError("Prompt config must contain 'component'.") + + component = prompt_ref.get("component", "") + params = prompt_ref.get("params", {}) + template = params.get("template", "") + if not template and "templates" in params: + template = next(iter(params["templates"].values()), "") + # If no explicit template, the component has a built-in default. + # Only validate if the caller provided a template to override. + if template: + self._validate_prompt_template(component, template) + + def validate_prompt_exists(self, prompt_id: int) -> None: + """Validate that a prompt ID exists in the database. + + Parameters + ---------- + prompt_id : int + + Raises + ------ + RAGPromptValidationError + If the prompt does not exist. + """ + prompt = self.db.get(RAGPrompt, prompt_id) + if prompt is None: + raise RAGPromptValidationError( + f"Prompt with ID {prompt_id} does not exist." + ) + + # ------------------------------------------------------------------ + # Private helpers + # ------------------------------------------------------------------ + + def _validate_prompt_template(self, class_name: str, template: str) -> None: + """Check class is registered, is a Prompt subclass, and template validates. + + Args: + class_name: Registered prompt component name. + template: The template string to validate. + + Raises: + RAGPromptValidationError: On any validation failure. + """ + if class_name not in self._registry: + raise RAGPromptValidationError( + f"Component {class_name} is not registered in the registry." + ) + + prompt_component = self._registry[class_name] + prompt_class = prompt_component["class"] + + if not issubclass(prompt_class, Prompt): + raise RAGPromptValidationError( + f"Component {class_name} is not a valid Prompt subclass." + ) + + if not prompt_class.validate_template(template): + raise RAGPromptValidationError( + f"Invalid template for prompt {class_name}. " + f"Required tokens are: {prompt_class.get_required_placeholders()}" + ) + + def _serialize_prompt(self, prompt: RAGPrompt) -> PromptResponse: + """Convert a RAGPrompt DB model to a PromptResponse.""" + return PromptResponse( + id=prompt.id, + class_name=prompt.class_name, + name=prompt.name, + parameters=prompt.parameters, + created=prompt.created, + last_modified=prompt.last_modified, + ) + + def _build_session_prompt_name(self, base_name: str, session_id: int) -> str: + """Generate a unique name for a session-scoped prompt copy.""" + candidate = f"{base_name} - session {session_id}" + suffix = 2 + while self.db.execute( + select(RAGPrompt.id).where(RAGPrompt.name == candidate) + ).scalar(): + candidate = f"{base_name} - session {session_id} ({suffix})" + suffix += 1 + return candidate diff --git a/DashAI/back/services/RAG/retriever_db_service.py b/DashAI/back/services/RAG/retriever_db_service.py new file mode 100644 index 000000000..33b804383 --- /dev/null +++ b/DashAI/back/services/RAG/retriever_db_service.py @@ -0,0 +1,589 @@ +"""Service layer for retriever-related database operations. + +Encapsulates all SQLAlchemy persistence logic for retriever models, +bridge records, embedding models, and embedding matrices. This service +is consumed by retriever factories and orchestration code; retriever +instances themselves never touch the database. +""" + +from typing import Dict, List, Optional + +from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.orm import Session + +from DashAI.back.dependencies.database.models import ( + RAGDenseRetriever as DenseRetrieverDBModel, +) +from DashAI.back.dependencies.database.models import ( + RAGEmbeddingMatrix as EmbeddingMatrixDBModel, +) +from DashAI.back.dependencies.database.models import ( + RAGEmbeddingModel as EmbeddingDBModel, +) +from DashAI.back.dependencies.database.models import ( + RAGRetriever as RetrieverDBModel, +) +from DashAI.back.dependencies.database.models import ( + RAGRetrieverChild, +) +from DashAI.back.dependencies.database.models import ( + RAGSparseRetriever as SparseRetrieverDBModel, +) + + +class RetrieverDBService: + """All database operations for retriever models. + + This class owns every SQL query, INSERT, and DB-model construction + related to retrievers. Factories and higher-level orchestration code + use it as a collaborator; retriever instances never touch it. + + Every public method is expected to handle SQLAlchemy exceptions so + that callers can react to persistence failures uniformly. + """ + + def __init__(self, db: Session): + self.db = db + + # ── Bridge / identity helpers ────────────────────────────────────── + + def create_bridge( + self, class_name: str, pipeline_id: int, commit: bool = True + ) -> RetrieverDBModel: + """Insert a new canonical identity record in RAG_retriever. + + Every retriever (unit or composite) gets one bridge record. + + Args: + class_name: Retriever component class name. + pipeline_id: FK to the owning pipeline. + commit: When ``True``, commit and refresh; otherwise only flush. + + Returns: + The persisted bridge record with its auto-generated id. + """ + bridge_record = RetrieverDBModel( + class_name=class_name, + pipeline_id=pipeline_id, + ) + try: + self.db.add(bridge_record) + if commit: + self.db.commit() + self.db.refresh(bridge_record) + else: + self.db.flush() + except SQLAlchemyError: + self.db.rollback() + raise + return bridge_record + + def find_bridge_for_sub_table( + self, + sub_db_model: SparseRetrieverDBModel | DenseRetrieverDBModel, + class_name: str, + ) -> Optional[RetrieverDBModel]: + """Given a sub-table detail record, find its parent bridge record. + + Args: + sub_db_model: A dense or sparse sub-table record. + class_name: Expected class name for validation. + + Returns: + The matching bridge record, or ``None`` if not found or mismatch. + """ + try: + bridge = self.db.query(RetrieverDBModel).get(sub_db_model.bridge_id) + except SQLAlchemyError: + raise + if bridge is not None and bridge.class_name == class_name: + return bridge + return None + + # ── Sparse retriever ─────────────────────────────────────────────── + + def find_sparse( + self, + class_name: str, + parameters: Dict[str, object], + chunk_set_id: int, + ) -> Optional[SparseRetrieverDBModel]: + """Look up an existing sparse retriever record by its natural key. + + Parameters + ---------- + class_name : str + Concrete retriever class name (e.g. "TFIDFRetriever"). + parameters : dict + Schema parameters, sorted for deterministic JSON serialisation. + chunk_set_id : int + Which chunk set this retriever was trained on. + + Returns + ------- + SparseRetrieverDBModel or None + """ + parameters = dict(sorted(parameters.items())) + try: + return ( + self.db.query(SparseRetrieverDBModel) + .filter_by( + class_name=class_name, + parameters=parameters, + chunk_set_id=chunk_set_id, + ) + .first() + ) + except SQLAlchemyError: + raise + + def save_sparse( + self, + class_name: str, + parameters: Dict[str, object], + storage_folder: str, + bridge_id: int, + chunk_set_id: int, + commit: bool = True, + ) -> SparseRetrieverDBModel: + """Persist a new sparse retriever record and link it to its bridge. + + Args: + class_name: Retriever component class name. + parameters: Schema parameters. + storage_folder: On-disk folder for the sparse model. + bridge_id: FK to the bridge record. + chunk_set_id: FK to the chunk set. + commit: When ``True``, commit and refresh; otherwise only flush. + + Returns: + The persisted sparse retriever record. + """ + parameters = dict(sorted(parameters.items())) + record = SparseRetrieverDBModel( + bridge_id=bridge_id, + chunk_set_id=chunk_set_id, + class_name=class_name, + parameters=parameters, + storage_folder=storage_folder, + ) + try: + self.db.add(record) + if commit: + self.db.commit() + self.db.refresh(record) + else: + self.db.flush() + except SQLAlchemyError: + self.db.rollback() + raise + return record + + def find_sparse_by_bridge_id( + self, bridge_id: int + ) -> Optional[SparseRetrieverDBModel]: + """Look up a sparse retriever sub-table record by its bridge_id. + + Args: + bridge_id: FK to the bridge record. + + Returns: + The sparse detail record, or ``None`` if not found. + """ + try: + return ( + self.db.query(SparseRetrieverDBModel) + .filter_by(bridge_id=bridge_id) + .first() + ) + except SQLAlchemyError: + raise + + # ── Dense retriever ──────────────────────────────────────────────── + + def find_dense( + self, + class_name: str, + parameters: Dict[str, object], + chunk_set_id: int, + ) -> Optional[DenseRetrieverDBModel]: + """Look up an existing dense retriever record by its natural key. + + Args: + class_name: Retriever component class name. + parameters: Schema parameters. + chunk_set_id: FK to the chunk set. + + Returns: + The dense detail record, or ``None`` if not found. + """ + parameters = dict(sorted(parameters.items())) + try: + return ( + self.db.query(DenseRetrieverDBModel) + .filter_by( + class_name=class_name, + parameters=parameters, + chunk_set_id=chunk_set_id, + ) + .first() + ) + except SQLAlchemyError: + raise + + def save_dense( + self, + class_name: str, + parameters: Dict[str, object], + bridge_id: int, + chunk_set_id: int, + embedding_model_id: int, + commit: bool = True, + ) -> DenseRetrieverDBModel: + """Persist a new dense retriever record and link it to its bridge. + + Args: + class_name: Retriever component class name. + parameters: Schema parameters. + bridge_id: FK to the bridge record. + chunk_set_id: FK to the chunk set. + embedding_model_id: FK to the embedding model. + commit: When ``True``, commit and refresh; otherwise only flush. + + Returns: + The persisted dense retriever record. + """ + parameters = dict(sorted(parameters.items())) + record = DenseRetrieverDBModel( + bridge_id=bridge_id, + chunk_set_id=chunk_set_id, + class_name=class_name, + parameters=parameters, + embedding_model_id=embedding_model_id, + ) + try: + self.db.add(record) + if commit: + self.db.commit() + self.db.refresh(record) + else: + self.db.flush() + except SQLAlchemyError: + self.db.rollback() + raise + return record + + def find_dense_by_bridge_id( + self, bridge_id: int + ) -> Optional[DenseRetrieverDBModel]: + """Look up a dense retriever sub-table record by its bridge_id. + + Args: + bridge_id: FK to the bridge record. + + Returns: + The dense detail record, or ``None`` if not found. + """ + try: + return ( + self.db.query(DenseRetrieverDBModel) + .filter_by(bridge_id=bridge_id) + .first() + ) + except SQLAlchemyError: + raise + + # ── Composite retriever ──────────────────────────────────────────── + + def save_composite( + self, + class_name: str, + pipeline_id: int, + child_bridge_ids: List[int], + commit: bool = True, + ) -> RetrieverDBModel: + """Persist a composite bridge record and its child links. + + Parameters + ---------- + class_name : str + "SequentialRetriever" or "ParallelRetriever". + pipeline_id : int + Owning pipeline. + child_bridge_ids : list[int] + Ordered list of child bridge ids (RAG_retriever.id). + commit : bool + When ``True``, commit and refresh; otherwise skip final commit. + + Returns + ------- + RetrieverDBModel + The persisted bridge record. + """ + bridge_record = RetrieverDBModel( + class_name=class_name, + pipeline_id=pipeline_id, + ) + try: + self.db.add(bridge_record) + self.db.flush() + for order, child_id in enumerate(child_bridge_ids): + self.db.add( + RAGRetrieverChild( + parent_id=bridge_record.id, + child_id=child_id, + child_order=order, + ) + ) + if commit: + self.db.commit() + self.db.refresh(bridge_record) + except SQLAlchemyError: + self.db.rollback() + raise + return bridge_record + + # ── Embedding model / matrix ─────────────────────────────────────── + + def find_or_create_embedding_model( + self, + class_name: str, + parameters: Dict[str, object], + commit: bool = True, + ) -> EmbeddingDBModel: + """Return an existing EmbeddingDBModel record or create one. + + This is idempotent: the same (class_name, parameters) pair + always resolves to the same record. + + Args: + class_name: Embedding model component class name. + parameters: Embedding model parameters. + + Returns: + The existing or newly created embedding model record. + """ + parameters = dict(sorted(parameters.items())) + try: + existing = ( + self.db.query(EmbeddingDBModel) + .filter_by( + class_name=class_name, + parameters=parameters, + ) + .first() + ) + except SQLAlchemyError: + raise + if existing is not None: + return existing + record = EmbeddingDBModel( + class_name=class_name, + parameters=parameters, + ) + try: + self.db.add(record) + if commit: + self.db.commit() + self.db.refresh(record) + else: + self.db.flush() + except SQLAlchemyError: + self.db.rollback() + raise + return record + + # NOTE: All embedding matrices are loaded into memory via np.load() + # for simplicity -- not suitable for very large document collections. + def find_embedding_matrices( + self, + doc_ids: List[int], + chunk_set_id: int, + embedding_model_id: int, + ) -> Dict[int, EmbeddingMatrixDBModel]: + """Return existing embedding matrices keyed by document_id. + + Only matrices that actually exist in the database are returned. + + Args: + doc_ids: List of document ids to look up. + chunk_set_id: FK to the chunk set. + embedding_model_id: FK to the embedding model. + + Returns: + Dict mapping document_id to existing matrix records. + """ + result: Dict[int, EmbeddingMatrixDBModel] = {} + for doc_id in doc_ids: + try: + matrix = ( + self.db.query(EmbeddingMatrixDBModel) + .filter_by( + document_id=doc_id, + chunk_set_id=chunk_set_id, + embedding_model_id=embedding_model_id, + ) + .first() + ) + except SQLAlchemyError: + raise + if matrix is not None: + result[doc_id] = matrix + return result + + def find_embedding_matrix( + self, + document_id: int, + chunk_set_id: int, + embedding_model_id: int, + ) -> Optional[EmbeddingMatrixDBModel]: + """Look up a single embedding matrix record. + + Args: + document_id: FK to the document. + chunk_set_id: FK to the chunk set. + embedding_model_id: FK to the embedding model. + + Returns: + The matrix record, or ``None`` if not found. + """ + try: + return ( + self.db.query(EmbeddingMatrixDBModel) + .filter_by( + document_id=document_id, + chunk_set_id=chunk_set_id, + embedding_model_id=embedding_model_id, + ) + .first() + ) + except SQLAlchemyError: + raise + + def save_embedding_matrix( + self, + document_id: int, + chunk_set_id: int, + embedding_model_id: int, + storage_folder: str, + matrix_shape: List[int], + commit: bool = True, + ) -> EmbeddingMatrixDBModel: + """Persist a new embedding matrix record. + + Args: + document_id: FK to the document. + chunk_set_id: FK to the chunk set. + embedding_model_id: FK to the embedding model. + storage_folder: On-disk path to the matrix directory. + matrix_shape: Shape of the numpy array as ``[rows, cols]``. + commit: When ``True``, commit and refresh; otherwise only flush. + + Returns: + The persisted embedding matrix record. + """ + record = EmbeddingMatrixDBModel( + document_id=document_id, + chunk_set_id=chunk_set_id, + embedding_model_id=embedding_model_id, + storage_folder=storage_folder, + matrix_shape=matrix_shape, + ) + try: + self.db.add(record) + if commit: + self.db.commit() + self.db.refresh(record) + else: + self.db.flush() + except SQLAlchemyError: + self.db.rollback() + raise + return record + + # ── Deletion helpers ─────────────────────────────────────────────── + + def delete_embedding_matrices_by_ids(self, matrix_ids: List[int]) -> None: + """Bulk-delete embedding matrices by their primary keys. + + Args: + matrix_ids: List of primary keys to delete. + """ + try: + self.db.query(EmbeddingMatrixDBModel).filter( + EmbeddingMatrixDBModel.id.in_(matrix_ids) + ).delete(synchronize_session="fetch") + self.db.commit() + except SQLAlchemyError: + self.db.rollback() + raise + + def delete_embedding_model(self, embedding_model_id: int) -> None: + """Delete a single embedding model record by its primary key. + + Args: + embedding_model_id: Primary key of the embedding model to delete. + """ + try: + model = ( + self.db.query(EmbeddingDBModel).filter_by(id=embedding_model_id).first() + ) + if model is not None: + self.db.delete(model) + self.db.commit() + except SQLAlchemyError: + self.db.rollback() + raise + + def delete_dense_detail(self, bridge_id: int) -> None: + """Delete the dense retriever sub-table row linked by bridge_id. + + Args: + bridge_id: FK to the bridge record whose detail should be deleted. + """ + try: + detail = ( + self.db.query(DenseRetrieverDBModel) + .filter_by(bridge_id=bridge_id) + .first() + ) + if detail is not None: + self.db.delete(detail) + self.db.commit() + except SQLAlchemyError: + self.db.rollback() + raise + + def delete_sparse_detail(self, bridge_id: int) -> None: + """Delete the sparse retriever sub-table row linked by bridge_id. + + Args: + bridge_id: FK to the bridge record whose detail should be deleted. + """ + try: + detail = ( + self.db.query(SparseRetrieverDBModel) + .filter_by(bridge_id=bridge_id) + .first() + ) + if detail is not None: + self.db.delete(detail) + self.db.commit() + except SQLAlchemyError: + self.db.rollback() + raise + + def delete_bridge(self, bridge_id: int) -> None: + """Delete a bridge record (RAG_retriever row) by its primary key. + + Args: + bridge_id: Primary key of the bridge record to delete. + """ + try: + bridge = self.db.query(RetrieverDBModel).get(bridge_id) + if bridge is not None: + self.db.delete(bridge) + self.db.commit() + except SQLAlchemyError: + self.db.rollback() + raise diff --git a/DashAI/back/services/RAG/retriever_setup_service.py b/DashAI/back/services/RAG/retriever_setup_service.py new file mode 100644 index 000000000..0058f3c3c --- /dev/null +++ b/DashAI/back/services/RAG/retriever_setup_service.py @@ -0,0 +1,578 @@ +"""Service that orchestrates the complete retriever lifecycle. + +Handles lookup-or-create, factory invocation, DB persistence, and embedding +matrix storage for both unit retrievers (dense/sparse) and composite retrievers +(with recursive children). + +The lifecycle is divided into three explicit phases: + Phase 1 — Construction: build the model in memory (no I/O, no DB). + Phase 2 — Initialization: heavy I/O (embeddings, similarity matrices). + Phase 3 — Persistence: save to DB (bridge records, sub-table rows). +""" + +import json +import logging +import os +from dataclasses import dataclass +from typing import Any + +import numpy as np +from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.orm import Session + +from DashAI.back.core.schema_fields.utils import normalize_payload +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.RAG.documents.chunk import Chunk +from DashAI.back.models.RAG.exceptions import ( + RAGRetrieverEmptyChildrenError, + RAGRetrieverError, + RAGRetrieverMissingParameterError, +) +from DashAI.back.models.RAG.retrievers.composite.composite_retriever import ( + CompositeRetriever, +) +from DashAI.back.models.RAG.retrievers.dense.dense_retriever import DenseRetriever +from DashAI.back.models.RAG.retrievers.persistence import ( + DensePersistence, + SparsePersistence, +) +from DashAI.back.models.RAG.retrievers.retriever_factory import ( + RetrieverFactory, +) +from DashAI.back.models.RAG.retrievers.retriever_model import RetrieverModel +from DashAI.back.models.RAG.retrievers.sparse.sparse_retriever import SparseRetriever +from DashAI.back.models.RAG.utils import hash_function +from DashAI.back.services.RAG.embedding_storage_service import EmbeddingStorageService +from DashAI.back.services.RAG.retriever_db_service import RetrieverDBService + +log = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class RetrieverSetupResult: + db_record_id: int + model: RetrieverModel + + +class RetrieverSetupService: + """Orchestrates the complete retriever lifecycle. + + Combines the pure :class:`RetrieverFactory` with database persistence + and embedding storage so that higher-level pipeline code only needs to + call a single ``setup()`` method per retriever configuration. + + The lifecycle is divided into three explicit phases: + 1. Construction — build the model in memory (no I/O, no DB). + **Note:** Phase 1 includes one exception — an idempotent + ``find_or_create_embedding_model`` call to resolve the embedding + model id used in the :class:`DensePersistence` stub. + 2. Initialization — heavy I/O (embeddings, similarity matrices). + 3. Persistence — save to DB (bridge records, sub-table rows). + + **Composite retrievers** are never cached in the DB. They are always + reconstructed fresh because they are lightweight to build (their + children are individually cached). + """ + + def __init__( + self, + db: Session, + registry: ComponentRegistry, + env_RAG_path: str, # noqa: N803 + chunks: dict[int, dict[int, Chunk]], + chunk_set_id: int, + pipeline_id: int, + ): + """Initialize the retriever setup service. + + Internally creates :attr:`_db_service` (:class:`RetrieverDBService`) + and :attr:`_embedding_service` (:class:`EmbeddingStorageService`). + + Args: + db: SQLAlchemy session. + registry: Application component registry. + env_RAG_path: Base RAG data directory path. + chunks: Nested dict of chunks keyed by document_id and index. + chunk_set_id: FK to the owning chunk set. + pipeline_id: FK to the owning pipeline. + """ + self._db = db + self._registry = registry + self._env_RAG_path = env_RAG_path + self._chunks = chunks + self._chunk_set_id = chunk_set_id + self._pipeline_id = pipeline_id + self._db_service = RetrieverDBService(db) + self._embedding_service = EmbeddingStorageService( + env_RAG_path, self._db_service + ) + + # ── Public API ───────────────────────────────────────────────────── + + def build_model( + self, + component_name: str, + params: dict[str, Any], + persistence: DensePersistence | SparsePersistence | None = None, + ) -> RetrieverModel: + """Phase 1: Build model in memory. No I/O, no DB.""" + factory = RetrieverFactory(self._registry, self._env_RAG_path, self._chunks) + result = factory.create(component_name, params, persistence) + return result.model + + def initialize_model(self, model: RetrieverModel) -> None: + """Phase 2: Initialize model (load embeddings, similarity matrices).""" + model.init_model() + + def persist_model( + self, model: RetrieverModel, component_name: str, sorted_params: dict + ) -> int: + """Phase 3: Persist to DB and return bridge_id.""" + bridge_id = self._save_unit(model, sorted_params) + model.set_id(bridge_id) + return bridge_id + + def setup( + self, component_name: str, params: dict[str, Any] + ) -> RetrieverSetupResult: + """Complete retriever setup: build -> initialize -> persist. + + Parameters + ---------- + component_name : str + Registered component name (e.g. ``"DenseEmbeddingRetriever"``). + params : dict[str, Any] + Raw frontend-style configuration parameters. + + Returns + ------- + RetrieverSetupResult + Tuple of the database record id and the fully-initialised model. + """ + params = normalize_payload(params) + model_class = self._registry[component_name]["class"] + + if issubclass(model_class, CompositeRetriever): + return self._setup_composite(model_class, params) + + sorted_params = dict(sorted(params.items())) + loaded = self._load_unit_from_db( + model_class, model_class.__name__, sorted_params + ) + if loaded is not None: + return RetrieverSetupResult(db_record_id=loaded.get_id(), model=loaded) + + persistence = self._build_persistence_for(model_class, params) + model = self.build_model(component_name, params, persistence) + + self.initialize_model(model) + + bridge_id = self.persist_model(model, component_name, sorted_params) + return RetrieverSetupResult(db_record_id=bridge_id, model=model) + + # ── Private: Composite ───────────────────────────────────────────── + + def _setup_composite( + self, model_class: type, params: dict[str, Any] + ) -> RetrieverSetupResult: + """Set up a composite retriever (Sequential/Parallel) with children. + + Always builds fresh — composites are NOT cached since they are + lightweight to reconstruct (children are individually cached). + + Args: + model_class: The composite retriever class. + params: Configuration parameters including ``children``. + + Returns: + A result with the bridge DB record id and initialized model. + + Raises: + RAGRetrieverEmptyChildrenError: If ``children`` is empty. + """ + children_configs = params.get("children", []) + if not children_configs: + raise RAGRetrieverEmptyChildrenError( + "Composite retriever must have at least one child" + ) + children = [self._setup_child(c) for c in children_configs] + model = model_class( + children=[c.model for c in children], + **{k: v for k, v in params.items() if k != "children"}, + ) + model.inject_infra(self._env_RAG_path, self._chunks, None) + + try: + bridge = self._db_service.save_composite( + model_class.__name__, + self._pipeline_id, + [c.db_record_id for c in children], + commit=True, + ) + except SQLAlchemyError: + log.exception("Failed to save composite retriever.") + raise + + model.set_id(bridge.id) + return RetrieverSetupResult(db_record_id=bridge.id, model=model) + + def _setup_child(self, child_config: dict) -> RetrieverSetupResult: + if "component" not in child_config or "params" not in child_config: + raise RAGRetrieverMissingParameterError( + f"Missing 'component' or 'params' in child config: {child_config}" + ) + return self.setup( + component_name=child_config["component"], + params=child_config["params"], + ) + + # ── Private: Unit ────────────────────────────────────────────────── + + def _load_unit_from_db( + self, model_class: type, class_name: str, sorted_params: dict[str, Any] + ) -> RetrieverModel | None: + """Load a unit retriever (dense or sparse) from the DB cache. + + Args: + model_class: Retriever model class. + class_name: Component class name string. + sorted_params: Deterministically sorted parameters. + + Returns: + Initialized retriever model, or ``None`` if not cached. + """ + if issubclass(model_class, DenseRetriever): + return self._load_dense(class_name, sorted_params) + if issubclass(model_class, SparseRetriever): + return self._load_sparse(class_name, sorted_params) + return None + + def _load_dense( + self, class_name: str, sorted_params: dict[str, Any] + ) -> RetrieverModel | None: + """Load a dense retriever from the DB cache by its natural key. + + Args: + class_name: Component class name string. + sorted_params: Deterministically sorted parameters. + + Returns: + Fully initialized DenseRetriever, or ``None`` if not cached. + + Raises: + SQLAlchemyError: On database lookup failure. + """ + try: + record = self._db_service.find_dense( + class_name, sorted_params, self._chunk_set_id + ) + except SQLAlchemyError: + log.exception("Database error during dense retriever lookup.") + raise + if record is None: + return None + persistence = self._build_dense_persistence(record.embedding_model_id) + factory = RetrieverFactory(self._registry, self._env_RAG_path, self._chunks) + result = factory.create( + class_name, dict(sorted_params), persistence=persistence + ) + model = result.model + model.init_model() + bridge = self._db_service.find_bridge_for_sub_table(record, class_name) + if bridge is not None: + model.set_id(bridge.id) + return model + + def _load_sparse( + self, class_name: str, sorted_params: dict[str, Any] + ) -> RetrieverModel | None: + """Load a sparse retriever from the DB cache by its natural key. + + Args: + class_name: Component class name string. + sorted_params: Deterministically sorted parameters. + + Returns: + Fully initialized SparseRetriever, or ``None`` if not cached. + + Raises: + SQLAlchemyError: On database lookup failure. + """ + try: + record = self._db_service.find_sparse( + class_name, sorted_params, self._chunk_set_id + ) + except SQLAlchemyError: + log.exception("Database error during sparse retriever lookup.") + raise + if record is None: + return None + persistence = SparsePersistence(model_dir=record.storage_folder) + factory = RetrieverFactory(self._registry, self._env_RAG_path, self._chunks) + result = factory.create( + class_name, dict(sorted_params), persistence=persistence + ) + model = result.model + model.init_model() + bridge = self._db_service.find_bridge_for_sub_table(record, class_name) + if bridge is not None: + model.set_id(bridge.id) + return model + + def _build_persistence_for( + self, model_class: type, params: dict + ) -> DensePersistence | SparsePersistence: + """Build the appropriate persistence object for a retriever type. + + Args: + model_class: Retriever model class. + params: Configuration parameters. + + Returns: + A DensePersistence or SparsePersistence instance. + + Raises: + RAGRetrieverError: If the retriever type is unsupported. + """ + if issubclass(model_class, DenseRetriever): + return self._build_dense_persistence_for(params) + elif issubclass(model_class, SparseRetriever): + return self._build_sparse_persistence_for(model_class, params) + else: + raise RAGRetrieverError( + f"Unsupported retriever type: {model_class.__name__}" + ) + + def _build_dense_persistence_for(self, params: dict) -> DensePersistence: + """Build a DensePersistence from the embedding model component ref. + + Reads ``params["embedding_model"]`` which at Phase 1 always has the + schema-component format ``{"component": "", "params": {...}}``, + set by ``DenseEmbeddingRetriever.__init__``. + + When the embedding model key is absent (e.g. for + ``HuggingFaceDenseRetriever`` subclasses that create the embedding + internally during ``init_model()``), an empty persistence is returned + and the caller must handle embedding computation in Phase 2. + + Args: + params: Configuration parameters containing ``embedding_model`` + in component-ref format. + + Returns: + A DensePersistence with matrix directories and embedding model id, + or an empty one when no embedding model reference is present. + """ + emb_ref = params.get("embedding_model") + if not isinstance(emb_ref, dict): + return DensePersistence(matrix_dirs={}, embedding_model_id=0) + + class_name = emb_ref.get("component") + parameters = emb_ref.get("params") + if not class_name or not isinstance(parameters, dict): + return DensePersistence(matrix_dirs={}, embedding_model_id=0) + + try: + emb_record = self._db_service.find_or_create_embedding_model( + class_name, + dict(sorted(parameters.items())), + commit=False, + ) + except SQLAlchemyError: + log.exception("Failed to find or create embedding model record.") + raise + return self._build_dense_persistence(emb_record.id) + + def _build_sparse_persistence_for( + self, model_class: type, params: dict + ) -> SparsePersistence: + """Build a SparsePersistence object with a deterministic storage folder. + + The folder path is derived from a hash of class name, params, and + chunk set id to support idempotent lookups. + + Args: + model_class: Retriever model class. + params: Configuration parameters. + + Returns: + A SparsePersistence instance pointing to the computed folder. + """ + identity = hash_function( + json.dumps( + { + "class_name": model_class.__name__, + "params": dict(sorted(params.items())), + "chunk_set_id": self._chunk_set_id, + }, + sort_keys=True, + ) + )[:16] + storage_folder = os.path.join( + self._env_RAG_path, + "sparse_retrievers", + f"sparse_retriever_id-{identity}", + ) + return SparsePersistence(model_dir=storage_folder) + + def _build_dense_persistence(self, embedding_model_id: int) -> DensePersistence: + """Build a DensePersistence with matrix directories for every document. + + Args: + embedding_model_id: FK to the embedding model record. + + Returns: + A DensePersistence instance with per-document matrix dirs. + """ + matrix_dirs: dict[int, str] = {} + for doc_id in self._chunks: + folder = EmbeddingStorageService.build_matrix_dir_name( + doc_id, + self._chunk_set_id, + embedding_model_id, + ) + matrix_dirs[doc_id] = os.path.join( + self._env_RAG_path, + "embeddings", + folder, + ) + return DensePersistence( + matrix_dirs=matrix_dirs, + embedding_model_id=embedding_model_id, + ) + + # ── Private: Persistence ─────────────────────────────────────────── + + def _save_unit(self, model: RetrieverModel, sorted_params: dict[str, Any]) -> int: + """Persist a unit retriever (dense or sparse) and return its bridge id. + + Args: + model: The initialized retriever model. + sorted_params: Deterministically sorted parameters. + + Returns: + The bridge record primary key. + + Raises: + ValueError: If the retriever type is unsupported. + """ + if isinstance(model, DenseRetriever): + return self._save_dense(model, sorted_params) + if isinstance(model, SparseRetriever): + return self._save_sparse(model, sorted_params) + raise ValueError(f"Unsupported retriever: {type(model).__name__}") + + def _save_dense(self, model: DenseRetriever, sorted_params: dict[str, Any]) -> int: + """Persist a dense retriever: embedding matrices, bridge, and detail record. + + All DB operations share a single transaction; the caller's session + is committed once at the end. + + Args: + model: The initialized DenseRetriever. + sorted_params: Deterministically sorted parameters. + + Returns: + The bridge record primary key. + + Raises: + SQLAlchemyError: On any DB persistence failure. + IOError: If an embedding matrix file cannot be read. + OSError: If an embedding matrix file cannot be read. + ValueError: If an embedding matrix file contains invalid data. + """ + try: + for doc_id, mdir in model.persistence.matrix_dirs.items(): + path = os.path.join(mdir, "embeddings.npy") + if not os.path.exists(path): + continue + try: + exists = self._db_service.find_embedding_matrix( + doc_id, + self._chunk_set_id, + model.persistence.embedding_model_id, + ) + except SQLAlchemyError: + log.exception("Database error during embedding matrix lookup.") + raise + if exists: + continue + try: + shape = list(np.load(path).shape) + except (IOError, OSError, ValueError) as exc: + log.error("Failed to load %s to record its shape: %s", path, exc) + raise + + self._db_service.save_embedding_matrix( + doc_id, + self._chunk_set_id, + model.persistence.embedding_model_id, + mdir, + shape, + commit=False, + ) + + bridge = self._db_service.create_bridge( + model.__class__.__name__, + self._pipeline_id, + commit=False, + ) + self._db_service.save_dense( + model.__class__.__name__, + sorted_params, + bridge.id, + self._chunk_set_id, + model.persistence.embedding_model_id, + commit=False, + ) + + self._db.commit() + except Exception: + self._db.rollback() + log.exception("Failed to persist dense retriever.") + raise + + return bridge.id + + def _save_sparse( + self, model: SparseRetriever, sorted_params: dict[str, Any] + ) -> int: + """Persist a sparse retriever: filesystem save, bridge, and detail record. + + All DB operations share a single transaction; the caller's session + is committed once at the end. + + Args: + model: The initialized SparseRetriever. + sorted_params: Deterministically sorted parameters. + + Returns: + The bridge record primary key. + + Raises: + SQLAlchemyError: On any DB persistence failure. + IOError: If the model cannot be saved to disk. + """ + model.save() + + try: + bridge = self._db_service.create_bridge( + model.__class__.__name__, + self._pipeline_id, + commit=False, + ) + self._db_service.save_sparse( + model.__class__.__name__, + sorted_params, + model.persistence.model_dir, + bridge.id, + self._chunk_set_id, + commit=False, + ) + self._db.commit() + except Exception: + self._db.rollback() + log.exception("Failed to persist sparse retriever.") + raise + + return bridge.id diff --git a/DashAI/back/tasks/RAG_task.py b/DashAI/back/tasks/RAG_task.py new file mode 100644 index 000000000..fb8ba1240 --- /dev/null +++ b/DashAI/back/tasks/RAG_task.py @@ -0,0 +1,174 @@ +import json +from typing import Any, List, Optional, Tuple + +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import ProcessData +from DashAI.back.models.RAG.exceptions import RAGTaskInputError +from DashAI.back.models.RAG.RAG_pipeline import RAGGenerationOutput +from DashAI.back.tasks.base_generative_task import BaseGenerativeTask + + +class RAGTask(BaseGenerativeTask): + """Class for RAG Task. + + Here you can change the methods provided by class Task. + """ + + metadata: dict = { + "inputs": {"str": {"min": 1, "max": 1}}, + "outputs": {"str": {"min": 1, "max": 1}, "Dict": {"min": 1, "max": 1}}, + } + + DISPLAY_NAME: str = MultilingualString( + en="Retrieval-Augmented Generation", + es="Generación Aumentada por Recuperación (RAG)", + ) + DESCRIPTION: str = MultilingualString( + en=""" + This task generates a text response with an LLM model based on + documents provided. + """, + es=""" + Esta tarea genera una respuesta de texto con un modelo LLM basado + en los documentos proporcionados. + """, + ) + + USE_HISTORY: bool = True + + def prepare_for_task( + self, + input: List[ProcessData], + history: Optional[List[Tuple[str, str]]] = None, + ) -> list[dict[str, str]]: + """Prepare the input by including the history. + + Parameters + ---------- + input : str + The current input to be processed. + E.g.: + ["Tell me a joke."] + + history : Optional[List[Tuple[str, str]]], optional + The history of previous inputs and outputs, by default None. E.g.: + [("Hello!", "Hello! How can I assist you today?")] + + Returns + ------- + str + The input prepared with the history to be used by the model. + E.g.: + [{"role": "user", "content": "Hello!"}, + {"role": "assistant", "content": "Hello! How can I assist you today?"}, + {"role": "user", "content": "Tell me a joke."}] + """ + if not input: + raise RAGTaskInputError("Task input list must not be empty") + input = str(input[0].data) + + prepared_input = [{"role": "user", "content": input}] + + if not history: + return prepared_input + + context = [ + ( + {"role": "user", "content": h_input}, + {"role": "assistant", "content": h_output}, + ) + for (h_input, h_output) in history + ] + context = [entry for input_output in context for entry in input_output] + prepared_input = context + prepared_input + return prepared_input + + def prepare_input_for_database( + self, + input: List[str], + **kwargs: Any, + ) -> List[Tuple[str, str]]: + """Prepare the input for the database. + + Parameters + ---------- + input : str + The input to be prepared. + + Returns + ------- + List[Tuple[str, str]] + Input with the new types as a list of tuples containing the data + and its type + + """ + return [(input[0], "str")] + + def process_output( + self, + output: RAGGenerationOutput, + **kwargs: Any, + ) -> List[Tuple[str, str]]: + """Process the output of a generative model. + + Converts the generation output into a list of (data, type) tuples + suitable for database storage. + + Args: + output: The typed output from ``RAGPipeline.generate()``. + + Returns: + A list of ``(data, type)`` tuples: the message as ``"str"`` + and the chunks as ``"Dict"`` (JSON-encoded). + """ + message = output.message + chunks = output.chunks + return [ + (str(message), "str"), + ( + json.dumps( + {k: v.to_dict() for k, v in chunks.items()}, + ensure_ascii=False, + ), + "Dict", + ), + ] + + def process_output_from_database( + self, + output: List[ProcessData], + **kwargs: Any, + ) -> List[ProcessData]: + """Process the output from the database. + + Parameters + ---------- + output : list[str] + The output data to be processed. + + Returns + ------- + list[str] + The processed output data. + """ + + return output + + def process_input_from_database( + self, + input: List[ProcessData], + **kwargs: Any, + ) -> List[ProcessData]: + """Process the input from the database. + + Parameters + ---------- + input : list[str] + The input data to be processed. + + Returns + ------- + list[str] + The processed input data. + """ + return input diff --git a/DashAI/back/tasks/__init__.py b/DashAI/back/tasks/__init__.py index 9c0fa90a1..ca031c6c0 100644 --- a/DashAI/back/tasks/__init__.py +++ b/DashAI/back/tasks/__init__.py @@ -1 +1,12 @@ # flake8: noqa +from DashAI.back.tasks.base_generative_task import BaseGenerativeTask +from DashAI.back.tasks.base_task import BaseTask +from DashAI.back.tasks.classification_task import ClassificationTask +from DashAI.back.tasks.controlnet_task import ControlNetTask +from DashAI.back.tasks.RAG_task import RAGTask +from DashAI.back.tasks.regression_task import RegressionTask +from DashAI.back.tasks.tabular_classification_task import TabularClassificationTask +from DashAI.back.tasks.text_classification_task import TextClassificationTask +from DashAI.back.tasks.text_to_image_generation_task import TextToImageGenerationTask +from DashAI.back.tasks.text_to_text_generation_task import TextToTextGenerationTask +from DashAI.back.tasks.translation_task import TranslationTask diff --git a/DashAI/front/.yarn/releases/yarn-3.5.0.cjs b/DashAI/front/.yarn/releases/yarn-3.5.0.cjs old mode 100755 new mode 100644 diff --git a/DashAI/front/package.json b/DashAI/front/package.json index 915f700fa..04a097a00 100644 --- a/DashAI/front/package.json +++ b/DashAI/front/package.json @@ -11,6 +11,7 @@ "@mui/material": "^7", "@mui/system": "^7", "@mui/utils": "^7", + "@mui/x-data-grid": "^8.27.2", "@mui/x-date-pickers": "^8.27.2", "@tanstack/react-table": "^8.21.3", "axios": "^1.3.4", diff --git a/DashAI/front/src/App.jsx b/DashAI/front/src/App.jsx index 903e514dc..a2decd1ff 100644 --- a/DashAI/front/src/App.jsx +++ b/DashAI/front/src/App.jsx @@ -1,6 +1,12 @@ import React from "react"; -import { BrowserRouter, Outlet, Route, Routes } from "react-router-dom"; +import { + BrowserRouter, + Navigate, + Outlet, + Route, + Routes, +} from "react-router-dom"; import { TourRegistryProvider } from "./contexts/TourRegistryContext"; import ModuleThemeWrapper from "./components/ModuleThemeWrapper"; @@ -13,10 +19,15 @@ import PluginsPage from "./pages/plugins/Plugins"; import PipelinesPage from "./pages/pipelines/Pipelines"; import PluginsDetails from "./pages/plugins/components/PluginsDetails"; import Generative from "./pages/generative/Generative"; +import { GenerativeProvider } from "./components/generative/GenerativeContext"; import NewPipelineWrapper from "./pages/pipelines/newPipelineWrapper"; import HubContent from "./pages/hub/HubContent"; import HubImportPage from "./pages/hub/HubImportPage"; import JobQueueWidget from "./components/jobs/JobQueueWidget"; +import RAGDocumentsPage from "./pages/generative/RAG/RAGDocumentsPage"; +import RAGPromptsPage from "./pages/generative/RAG/RAGPromptsPage"; +import RAGSessionPage from "./pages/generative/RAGSession/RAGSessionPage"; +import SessionRouter from "./pages/generative/SessionRouter"; import { DatasetsAndNotebooksProvider } from "./components/custom/contexts/DatasetsAndNotebooksContext"; function DataSectionLayout() { @@ -62,12 +73,35 @@ function App() { element={} /> } /> + } + /> + + + + } + /> + + + + } + /> } /> } /> - } /> + } + /> } /> } /> => { + const response = await api.get( + `/v1/component/${componentName}/children`, + { + params: { recursive }, + }, + ); + return response.data; +}; export const getComponentById = async (id: string): Promise => { const response = await api.get(`/v1/component/${id}/`); return response.data; diff --git a/DashAI/front/src/api/generativeTask.ts b/DashAI/front/src/api/generativeTask.ts index d64c2c18f..9554483a2 100644 --- a/DashAI/front/src/api/generativeTask.ts +++ b/DashAI/front/src/api/generativeTask.ts @@ -12,15 +12,21 @@ export const getGenerativeTask = async (): Promise => { export const getRelatedComponents = async ( relatedComponent: string, ): Promise => { - const response = await api.get( - `/v1/component/?related_component=${encodeURIComponent(relatedComponent)}`, - ); - return response.data; + try { + const response = await api.get( + `/v1/component/?related_component=${encodeURIComponent(relatedComponent)}`, + ); + console.log("Related components:", response.data); + return response.data; + } catch (error) { + return []; + } }; export const createGenerativeSession = async ( sessionData: Omit, ): Promise => { + console.log("Creating new generative session with data:", sessionData); const response = await api.post( "/v1/generative-session/", sessionData, diff --git a/DashAI/front/src/api/job.ts b/DashAI/front/src/api/job.ts index 684fd434d..d3bd85459 100644 --- a/DashAI/front/src/api/job.ts +++ b/DashAI/front/src/api/job.ts @@ -176,6 +176,24 @@ export const enqueueGenerativeProcessJob = async ( return response.data; }; +export const enqueueRAGProcessJob = async ( + processId: number, +): Promise => { + const data = { + job_type: "RAGJob", + kwargs: { rag_process_id: processId }, + }; + const formData = new FormData(); + formData.append("job_type", data.job_type); + formData.append("kwargs", JSON.stringify(data.kwargs)); + const response = await api.post("/v1/job/", formData, { + headers: { + "Content-Type": "multipart/form-data", + }, + }); + return response.data; +}; + export const enqueueConverterJob = async ( converterId: number, ): Promise => { diff --git a/DashAI/front/src/api/rag.ts b/DashAI/front/src/api/rag.ts new file mode 100644 index 000000000..bbf561acf --- /dev/null +++ b/DashAI/front/src/api/rag.ts @@ -0,0 +1,297 @@ +import api from "./api"; +import { ISession } from "../types/session"; +import { IGenerativeTask } from "../types/generativeTask"; +import { IDocumentResponse } from "../types/documentResponse"; +import { IComponent } from "../types/component"; +import { IRAGPrompt } from "../types/ragPrompt"; +import { getChildComponents } from "./component"; + +/** + * Creates a new RAG prompt via the API. + * @param prompt - The prompt data (class_name, name, optional parameters). + * @returns The created prompt metadata containing the new ID. + */ +export const createRAGPrompt = async (prompt: { + class_name: string; + name: string; + parameters?: Record; +}): Promise<{ id: number }> => { + const response = await api.post("/v1/prompt/", prompt); + if (response.status !== 201) { + throw new Error(`Failed to create RAG prompt: ${response.statusText}`); + } + return response.data; +}; + +/** Fetches all generative sessions filtered to RAGTask. @returns List of RAG sessions. */ +export const getRAGSessions = async (): Promise => { + const response = await api.get("/v1/generative-session/"); + if (response.status !== 200) { + throw new Error(`Failed to fetch RAG sessions: ${response.statusText}`); + } + + const ragSessions = response.data.filter( + (session) => session.task_name === "RAGTask", + ); + return ragSessions; +}; + +/** Fetches a single RAG session by ID. @param sessionId - The session ID. @returns The session object. */ +export const getRAGSession = async (sessionId: number): Promise => { + const response = await api.get( + `/v1/generative-session/${sessionId}`, + ); + if (response.status !== 200) { + throw new Error(`Failed to fetch RAG session: ${response.statusText}`); + } + + return response.data; +}; + +/** + * Creates a new RAG session with the given data, forcing task_name to "RAGTask". + * @param sessionData - Session creation payload (without id/created/last_modified). + * @returns The created session. + */ +export const createRAGSession = async ( + sessionData: Omit, +): Promise => { + const transformedSession: Omit = + { + name: sessionData.name, + description: sessionData.description, + task_name: "RAGTask", + model_name: "RAGPipeline", + display_name: "", + parameters: sessionData.parameters, + }; + + const response = await api.post( + "/v1/generative-session/", + transformedSession, + ); + + if (response.status !== 201) { + throw new Error(`Failed to create RAG session: ${response.statusText}`); + } + + return response.data; +}; + +/** Updates an existing RAG session. @param sessionId - The session ID. @param sessionData - Partial fields to update. @returns The updated session. */ +export const updateRAGSession = async ( + sessionId: number, + sessionData: Partial, +): Promise => { + const response = await api.put( + `/v1/generative-session/${sessionId}`, + sessionData, + ); + if (response.status !== 200) { + throw new Error(`Failed to update RAG session: ${response.statusText}`); + } + + return response.data; +}; + +/** Deletes a RAG session by ID. @param sessionId - The session ID. */ +export const deleteRAGSession = async (sessionId: number): Promise => { + const response = await api.delete(`/v1/generative-session/${sessionId}`); + if (response.status !== 204) { + throw new Error(`Failed to delete RAG session: ${response.statusText}`); + } +}; + +/** Updates only the parameters of an existing RAG session. @param sessionId - The session ID. @param newParams - The new parameters payload. @returns The updated session. */ +export const updateGenerativeSessionParams = async ( + sessionId: number, + newParams: Record, +): Promise => { + const response = await api.put( + `/v1/generative-session/${sessionId}/parameters`, + newParams, + ); + if (response.status !== 200) { + throw new Error( + `Failed to update RAG session parameters: ${response.statusText}`, + ); + } + return response.data; +}; + +/** Fetches all available retriever paradigms (children of RetrieverModel). @returns List of retriever paradigm components. */ +export const getRetrievalParadigm = async (): Promise => { + const response = await getChildComponents("RetrieverModel", true); + if (!response) { + throw new Error(`Failed to fetch retrieval options`); + } + return response; +}; + +/** Fetches child components (specific retrievers) for a given retrieval paradigm. @param retrievalParadigm - Parent paradigm name. @returns List of retriever components. */ +export const getRetrieverComponents = async ( + retrievalParadigm: string, +): Promise => { + const response = await getChildComponents(retrievalParadigm, true); + + if (!response) { + throw new Error(`Failed to fetch retriever components`); + } + + return response; +}; + +/** Fetches generator components related to TextToTextGenerationTask. @returns List of generator components. */ +export const getGeneratorComponents = async (): Promise => { + const response = await api.get( + `/v1/component/?related_component=TextToTextGenerationTask`, + ); + + if (response.status !== 200) { + throw new Error( + `Failed to fetch generator components: ${response.statusText}`, + ); + } + + return response.data; +}; + +/** Fetches chunking model components (children of BaseChunkingModel). @returns List of chunking components. */ +export const getChunkingComponents = async (): Promise => { + const response = await getChildComponents("BaseChunkingModel", false); + if (!response) { + throw new Error(`Failed to fetch chunking components`); + } + return response; +}; + +/** Fetches all uploaded documents. @returns List of document responses. */ +export const loadDocuments = async (): Promise => { + const response = await api.get("/v1/document/"); + if (response.status !== 200) { + throw new Error(`Failed to load documents: ${response.statusText}`); + } + return response.data; +}; + +/** Fetches documents scoped to a specific RAG session. @param sessionId - The session ID. @returns List of document responses. */ +export const getSessionDocuments = async ( + sessionId: number, +): Promise => { + const response = await api.get( + `/v1/document/session/${sessionId}`, + ); + if (response.status !== 200) { + throw new Error(`Failed to load session documents: ${response.statusText}`); + } + return response.data; +}; + +/** Deletes a document by ID. @param documentId - The document ID. */ +export const deleteDocument = async (documentId: number): Promise => { + const response = await api.delete(`/v1/document/${documentId}`); + if (response.status !== 204) { + throw new Error(`Failed to delete document: ${response.statusText}`); + } +}; + +/** + * Uploads a document file with optional metadata via multipart/form-data. + * @param file - The File object to upload. + * @param optional_metadata - Optional metadata (name, source, etc.). + * @returns The saved document response. + */ +export const addDocument = async ({ + file, + optional_metadata, +}: { + file: File; + optional_metadata?: Record; +}): Promise => { + if (optional_metadata) { + optional_metadata.last_modified = file.lastModified; + } + const metadata = { + file_name: file.name, + last_modified: file.lastModified, + optional_metadata, + }; + + const formData = new FormData(); + formData.append("file", file); + formData.append("metadata", JSON.stringify(metadata)); + + const response = await api.post( + "/v1/document/", + formData, + { + headers: { "Content-Type": "multipart/form-data" }, + }, + ); + + if (response.status !== 201) { + throw new Error(`Failed to upload document: ${response.statusText}`); + } + + return response.data; +}; + +/** Class name prefix that identifies non-generation prompt types. */ +const AUGMENTATION_PROMPT_CLASS_PREFIX = "Augmentation"; + +/** + * Checks whether a prompt's class_name corresponds to a generation prompt + * (i.e. NOT an augmentation prompt). + * + * @param className - The prompt component class name to test. + * @returns `true` if the class is a generation prompt, `false` if it is an augmentation prompt. + */ +export function isGenerationPromptClass(className: string): boolean { + return !className.includes(AUGMENTATION_PROMPT_CLASS_PREFIX); +} + +/** Fetches default prompt components (children of RAGGenerationPrompt). @returns List of default prompt components. */ +export const getDefaultPrompts = async (): Promise => { + return getChildComponents("RAGGenerationPrompt", false); +}; + +/** Fetches all saved RAG prompts (user-created). @returns List of RAG prompts. */ +export const getRAGPrompts = async (): Promise => { + const response = await api.get("/v1/prompt/"); + if (response.status !== 200) { + throw new Error(`Failed to fetch RAG prompts: ${response.statusText}`); + } + return response.data; +}; + +/** + * Fetches custom (non-Default) prompt components for the given parent types. + * @param types - Array of parent component type names to fetch children from. + * @returns List of custom prompt components (excluding Default* classes). + */ +export const getCustomPrompts = async ( + types: string[] = ["RAGGenerationPrompt", "AugmentationPrompt"], +): Promise => { + let allChildren: IComponent[] = []; + for (const type of types) { + const response = await api.get( + `/v1/component/${type}/children`, + { params: { recursive: false } }, + ); + if (response.status !== 200) { + throw new Error( + `Failed to fetch ${type} children: ${response.statusText}`, + ); + } + const filtered = response.data.filter( + (child) => + !( + child.name && + typeof child.name === "string" && + child.name.includes("Default") + ), + ); + allChildren.push(...filtered); + } + return allChildren; +}; diff --git a/DashAI/front/src/api/session.ts b/DashAI/front/src/api/session.ts index a08e4d534..2d727a819 100644 --- a/DashAI/front/src/api/session.ts +++ b/DashAI/front/src/api/session.ts @@ -31,7 +31,7 @@ export const updateGenerativeSession = async ({ formData, }: { id: string; - formData: { name?: string; task_name?: string }; + formData: { name?: string; task_name?: string; description?: string }; }): Promise => { const response = await api.patch(`/v1/generative-session/${id}`, null, { params: formData, diff --git a/DashAI/front/src/components/configurableObject/Inputs/ArrayInput.jsx b/DashAI/front/src/components/configurableObject/Inputs/ArrayInput.jsx index 46be6e56e..10b2ecba5 100644 --- a/DashAI/front/src/components/configurableObject/Inputs/ArrayInput.jsx +++ b/DashAI/front/src/components/configurableObject/Inputs/ArrayInput.jsx @@ -14,7 +14,15 @@ function ArrayInput({ itemType, ...props }) { - const [inputValue, setInputValue] = useState(value.join(",")); + // Ensure value is an array before using join + const safeValue = Array.isArray(value) ? value : []; + const [inputValue, setInputValue] = useState(safeValue.join(",")); + + // Update inputValue when value prop changes + useEffect(() => { + const newSafeValue = Array.isArray(value) ? value : []; + setInputValue(newSafeValue.join(",")); + }, [value]); useEffect(() => { if (Array.isArray(value)) { diff --git a/DashAI/front/src/components/configurableObject/Inputs/ClassInput.jsx b/DashAI/front/src/components/configurableObject/Inputs/ClassInput.jsx index 5750881f8..d22be4182 100644 --- a/DashAI/front/src/components/configurableObject/Inputs/ClassInput.jsx +++ b/DashAI/front/src/components/configurableObject/Inputs/ClassInput.jsx @@ -164,6 +164,18 @@ function ClassInput({ open={open} onClose={handleClose} sx={{ display: modal ? "show" : "none" }} + maxWidth="sm" + fullWidth + PaperProps={{ + sx: { + zIndex: 1400, // Ensure it appears above parent Dialog (which has z-index 1300) + }, + }} + BackdropProps={{ + sx: { + zIndex: 1399, // Backdrop should be below the Dialog but above parent + }, + }} > {`${selectedOption} parameters`} diff --git a/DashAI/front/src/components/configurableObject/Inputs/TextInput.jsx b/DashAI/front/src/components/configurableObject/Inputs/TextInput.jsx index fbabeb814..49447ba95 100644 --- a/DashAI/front/src/components/configurableObject/Inputs/TextInput.jsx +++ b/DashAI/front/src/components/configurableObject/Inputs/TextInput.jsx @@ -20,6 +20,9 @@ function TextInput({ description, ...props }) { + const isEmpty = value === undefined || value === ""; + const showError = error || (isEmpty ? `${name} is a required field` : ""); + return ( ); diff --git a/DashAI/front/src/components/custom/TemplateModal.jsx b/DashAI/front/src/components/custom/TemplateModal.jsx new file mode 100644 index 000000000..a06fd16eb --- /dev/null +++ b/DashAI/front/src/components/custom/TemplateModal.jsx @@ -0,0 +1,51 @@ +import React from "react"; +import PropTypes from "prop-types"; +import { + Dialog, + DialogTitle, + DialogContent, + DialogContentText, + DialogActions, + Button, + Typography, +} from "@mui/material"; + +export default function TemplateModal({ + open, + handleClose, + template, + title = "Prompt", + formatText = true, +}) { + return ( + + {title} + + + + {template} + + + + + + + + ); +} + +TemplateModal.propTypes = { + open: PropTypes.bool.isRequired, + handleClose: PropTypes.func.isRequired, + template: PropTypes.string, + title: PropTypes.string, + formatText: PropTypes.bool, +}; diff --git a/DashAI/front/src/components/datasets/ConfigureAndUploadDataset.jsx b/DashAI/front/src/components/datasets/ConfigureAndUploadDataset.jsx index 9a4a27d3e..b6c3a3247 100644 --- a/DashAI/front/src/components/datasets/ConfigureAndUploadDataset.jsx +++ b/DashAI/front/src/components/datasets/ConfigureAndUploadDataset.jsx @@ -1,7 +1,7 @@ import React, { useState, useEffect } from "react"; -import { Grid, Paper } from "@mui/material"; +import { Grid, Paper, Typography } from "@mui/material"; import PropTypes from "prop-types"; -import Upload from "./Upload"; +import Upload from "../shared/Upload"; import { getComponents as getComponentsRequest } from "../../api/component"; import { useSnackbar } from "notistack"; import DataloaderConfiguration from "./DataloaderConfiguration"; diff --git a/DashAI/front/src/components/documents/DocumentTable.jsx b/DashAI/front/src/components/documents/DocumentTable.jsx new file mode 100644 index 000000000..1245fde37 --- /dev/null +++ b/DashAI/front/src/components/documents/DocumentTable.jsx @@ -0,0 +1,233 @@ +import React, { useEffect, useState, useCallback } from "react"; +import { DataGrid, GridToolbar } from "@mui/x-data-grid"; +import { useSnackbar } from "notistack"; +import { + Button, + Grid, + Paper, + Typography, + LinearProgress, + Box, +} from "@mui/material"; +import { + get_documents_json, + delete_document as deleteDocumentRequest, +} from "../../api/documents"; +import { + AddCircleOutline as AddIcon, + Update as UpdateIcon, +} from "@mui/icons-material"; +import DeleteItemModal from "../custom/DeleteItemModal"; +import DocumentSummaryModal from "./DocumentSummaryModal"; + +function DocumentTable({ + handleNewDocument, + updateTableFlag, + setUpdateTableFlag, +}) { + const { enqueueSnackbar } = useSnackbar(); + const [loading, setLoading] = useState(true); + const [documents, setDocuments] = useState([]); + const [selectionModel, setSelectionModel] = useState([]); + + const getDocuments = useCallback(async () => { + setLoading(true); + try { + const docs = await get_documents_json(); + setDocuments(docs); + } catch (error) { + enqueueSnackbar("Error when trying to get the documents", { + variant: "error", + }); + console.error("Error fetching documents:", error); + } finally { + setLoading(false); + } + }, [enqueueSnackbar]); + + const deleteDocument = async (doc_name) => { + try { + await deleteDocumentRequest(doc_name); + setUpdateTableFlag(true); + enqueueSnackbar("Document successfully deleted.", { + variant: "success", + }); + } catch (error) { + enqueueSnackbar("Error when trying to delete the document", { + variant: "error", + }); + console.error("Error deleting document:", error); + } + }; + + const createDeleteHandler = useCallback( + (doc_name) => () => { + deleteDocument(doc_name); + }, + [deleteDocument], + ); + + useEffect(() => { + getDocuments(); + }, [getDocuments]); + + useEffect(() => { + if (updateTableFlag) { + getDocuments(); + setUpdateTableFlag(false); + } + }, [updateTableFlag, setUpdateTableFlag, getDocuments]); + + const columns = React.useMemo( + () => [ + { + field: "id", + headerName: "ID", + minWidth: 30, + editable: false, + }, + { + field: "doc_name", + headerName: "Document Name", + minWidth: 200, + editable: false, + }, + { + field: "created_at", + headerName: "Created At", + minWidth: 200, + editable: false, + }, + { + field: "actions", + type: "actions", + minWidth: 150, + getActions: (params) => [ + , + , + ], + }, + ], + [createDeleteHandler], + ); + + return ( + + + + Current documents + + + + + + + + + + + + + + + {!loading && documents.length === 0 ? ( + + + No documents available + + + ) : ( + String(row.id)} + pageSize={5} + sortModel={[{ field: "id", sort: "asc" }]} + pageSizeOptions={[5, 10]} + checkboxSelection + selectionModel={selectionModel} + onSelectionModelChange={setSelectionModel} + disableRowSelectionOnClick={false} + autoHeight={false} + loading={loading} + slots={{ + toolbar: GridToolbar, + loadingOverlay: LinearProgress, + }} + sx={{ + flex: 1, + minHeight: 300, + "& .MuiDataGrid-cell:focus": { + outline: "none", + }, + }} + /> + )} + + + ); +} + +export default DocumentTable; diff --git a/DashAI/front/src/components/generative/CreateSessionCenter.jsx b/DashAI/front/src/components/generative/CreateSessionCenter.jsx index 2480df2b4..49701a6ca 100644 --- a/DashAI/front/src/components/generative/CreateSessionCenter.jsx +++ b/DashAI/front/src/components/generative/CreateSessionCenter.jsx @@ -7,11 +7,14 @@ import { } from "@mui/material"; import { useCallback, useEffect } from "react"; import { useTranslation } from "react-i18next"; +import { useNavigate } from "react-router-dom"; import ComponentSelector from "../custom/ComponentSelector"; import GenerativeBreadcrumbs from "./GenerativeBreadcrumbs"; import { useCreateSession } from "./CreateSessionContext"; +import { useGenerative } from "./GenerativeContext"; import StepperNavigationFooter from "../shared/StepperNavigationFooter"; import { useTourContext } from "../tour/TourProvider"; +import RAGSessionSetup from "../../pages/generative/RAGSession/RAGSessionSetup"; export default function CreateSessionCenter() { const { t } = useTranslation(["generative", "common"]); @@ -38,6 +41,9 @@ export default function CreateSessionCenter() { handleCreate, } = useCreateSession(); + const navigate = useNavigate(); + const { sessions, setSessions } = useGenerative(); + const handleSelectModelWithTour = useCallback( (model) => { handleSelectModel(model); @@ -56,6 +62,18 @@ export default function CreateSessionCenter() { handleCreate(); }, [handleCreate, tourContext]); + const handleRagSessionCreated = useCallback( + (createdSession) => { + setSessions((prev) => [...prev, createdSession]); + navigate(`/app/generative/sessions/${createdSession.id}`); + }, + [setSessions, navigate], + ); + + const handleRagClose = useCallback(() => { + handleBack(); + }, [handleBack]); + useEffect(() => { if (!tourContext?.run) return; const currentTarget = tourContext.steps?.[tourContext.stepIndex]?.target; @@ -84,18 +102,20 @@ export default function CreateSessionCenter() { }} > - - - {step === 0 - ? t("generative:label.selectModel") - : t("generative:label.configureSession")} - - - {step === 0 - ? t("generative:label.pickAModelGroupedByTask") - : t("generative:label.nameAndDescribeYourSession")} - - + {step === 1 && selectedModel?.task_name === "RAGTask" ? null : ( + + + {step === 0 + ? t("generative:label.selectModel") + : t("generative:label.configureSession")} + + + {step === 0 + ? t("generative:label.pickAModelGroupedByTask") + : t("generative:label.nameAndDescribeYourSession")} + + + )} c.name.toLowerCase().includes("qwen")} /> ) + ) : selectedModel?.task_name === "RAGTask" ? ( + ) : ( - + {step === 1 && selectedModel?.task_name === "RAGTask" ? null : ( + + )} ); } diff --git a/DashAI/front/src/components/generative/CreateSessionLanding.jsx b/DashAI/front/src/components/generative/CreateSessionLanding.jsx index 015d418e6..b47c51d25 100644 --- a/DashAI/front/src/components/generative/CreateSessionLanding.jsx +++ b/DashAI/front/src/components/generative/CreateSessionLanding.jsx @@ -9,7 +9,7 @@ export default function CreateSessionLanding() { const { t } = useTranslation(["generative", "common"]); const tourContext = useTourContext(); - const handleCreateSession = () => { + const handleOption = () => { if (tourContext?.run) { tourContext.nextStep(); } @@ -18,7 +18,7 @@ export default function CreateSessionLanding() { return ( + + + ); + } + return ( {/* Title */} diff --git a/DashAI/front/src/components/generative/DocumentReferencesModal.jsx b/DashAI/front/src/components/generative/DocumentReferencesModal.jsx new file mode 100644 index 000000000..9a14bb352 --- /dev/null +++ b/DashAI/front/src/components/generative/DocumentReferencesModal.jsx @@ -0,0 +1,182 @@ +import React, { useState } from "react"; +import { useTranslation } from "react-i18next"; +import { + Box, + Typography, + IconButton, + List, + ListItem, + Divider, + Dialog, + DialogTitle, + DialogContent, + DialogActions, + Button, +} from "@mui/material"; +import { + Close as CloseIcon, + Article as ArticleIcon, + ContentCopy as CopyIcon, +} from "@mui/icons-material"; + +const DocumentReferencesModal = ({ + open, + onClose, + document, + chunks, + onOpenReference, +}) => { + const { t } = useTranslation("generative"); + if (!document || !chunks) return null; + + const getDocumentTitle = (docId, chunks) => { + const firstChunk = chunks[0]; + if (firstChunk.document_title) return firstChunk.document_title; + if (firstChunk.document_name) return firstChunk.document_name; + if (firstChunk.title) return firstChunk.title; + if (firstChunk.name) return firstChunk.name; + return t("documentReferences.fallbackTitle", { + id: docId, + defaultValue: `Document ${docId}`, + }); + }; + + const handleCopyChunk = async (chunkText) => { + try { + const cleanText = chunkText.replace(/\\n/g, "\n"); + await navigator.clipboard.writeText(cleanText); + // You could add a toast notification here if you have a notification system + } catch (err) { + console.error("Failed to copy text: ", err); + } + }; + + return ( + + + + + + {getDocumentTitle(document.id, chunks)} + + + + + + + + + + {t("documentReferences.chunksCount", { count: chunks.length })} + + + + {chunks.map((chunk, index) => ( + + + + + + + {chunk.document_position + ? t("documentReferences.chunkLabel", { + position: chunk.document_position, + }) + : t("documentReferences.chunkLabel", { + position: index + 1, + })} + + + handleCopyChunk(chunk.text)} + sx={{ + color: "text.secondary", + "&:hover": { + color: "primary.main", + backgroundColor: "action.hover", + }, + }} + > + + + + + + {chunk.text.replace(/\\n/g, "\n")} + + + + {index < chunks.length - 1 && } + + ))} + + + + + + + + ); +}; + +export default DocumentReferencesModal; diff --git a/DashAI/front/src/components/generative/GenerativeChat.jsx b/DashAI/front/src/components/generative/GenerativeChat.jsx index 5f12cd2d3..4e0fb69fe 100644 --- a/DashAI/front/src/components/generative/GenerativeChat.jsx +++ b/DashAI/front/src/components/generative/GenerativeChat.jsx @@ -1,5 +1,4 @@ import { Box, Divider, IconButton, Typography } from "@mui/material"; -import { useTheme } from "@mui/material/styles"; import InfoIcon from "@mui/icons-material/Info"; import ArrowRightAltIcon from "@mui/icons-material/ArrowRightAlt"; import KeyboardArrowDownIcon from "@mui/icons-material/KeyboardArrowDown"; @@ -17,9 +16,15 @@ import { getHistoryBySessionId, getSessionById } from "../../api/session"; import InfoSessionModal from "./InfoSessionModal"; import { useSnackbar } from "notistack"; import { MediaInput } from "./MediaInput"; +import JobQueueWidget from "../jobs/JobQueueWidget"; +import { getRunStatus } from "../../utils/runStatus"; +import TemplateModal from "../custom/TemplateModal"; +import SourcesDisplay from "./SourcesDisplay"; +import RAGBreadcrumbs from "./RAG/RAGBreadcrumbs"; import { Trans, useTranslation } from "react-i18next"; import { useGenerative } from "./GenerativeContext"; import { useTourContext } from "../tour/TourProvider"; +import { useTheme } from "@mui/material/styles"; export default function GenerativeChat() { const theme = useTheme(); @@ -45,9 +50,13 @@ export default function GenerativeChat() { const [showScrollButton, setShowScrollButton] = useState(false); const [sessionInfo, setSessionInfo] = useState(null); const [sessionInfoVisible, setSessionInfoVisible] = useState(false); + const [referenceModalOpen, setReferenceModalOpen] = useState(false); + const [selectedReferenceText, setSelectedReferenceText] = useState(""); + const [referenceModalTitle, setReferenceModalTitle] = useState(""); const { enqueueSnackbar } = useSnackbar(); const { t } = useTranslation(["generative"]); const tourContext = useTourContext(); + const [shouldAutoScroll, setShouldAutoScroll] = useState(true); const scrollToBottom = (force = false) => { const el = chatContainerRef.current; @@ -61,6 +70,20 @@ export default function GenerativeChat() { } }; + const isAtBottom = () => { + if (!chatContainerRef.current) return true; + const { scrollTop, scrollHeight, clientHeight } = chatContainerRef.current; + return Math.abs(scrollHeight - clientHeight - scrollTop) < 5; // 5px threshold + }; + + const handleOpenReference = (ref, key) => { + const title = `Document ${ref.document_id}${ref.document_name ? ` (${ref.document_name})` : ""}${ref.document_position ? ` - Chunk ${ref.document_position}` : ""}`; + setReferenceModalTitle(title); + // Convert escaped newlines to actual newlines + setSelectedReferenceText(ref.text.replace(/\\n/g, "\n")); + setReferenceModalOpen(true); + }; + const handleScroll = () => { const el = chatContainerRef.current; if (!el) return; @@ -79,6 +102,7 @@ export default function GenerativeChat() { const getMessages = () => { getProcessesBySessionId(sessionId).then((response) => { + console.log("Fetched messages:", response); // Add here setIsLoadingMessage(false); setMessages(response); }); @@ -92,6 +116,7 @@ export default function GenerativeChat() { const handleSendMessage = (input) => { setIsLoadingMessage(true); + setShouldAutoScroll(true); // Enable auto-scroll when sending new message postProcess(sessionId, input).then((response) => { // Add the new message to the chat @@ -186,13 +211,92 @@ export default function GenerativeChat() { }, [messages]); useEffect(() => { + console.log("Combining messages and history for display"); // Add here + console.log("Messages:", messages); + console.log("TASK NAME:", taskName); + console.log("session name:", sessionInfo?.name); + console.log("session description:", sessionInfo?.description); let messagesObject = messages.map((process) => { + // Check if there's reference data in the output (only for RAGTask) + let referenceOutput = null; + let mainOutput = process.output; + + if ( + taskName === "RAGTask" && + process.output && + process.output.length > 1 + ) { + // Look for Dict type output that contains reference information + const referenceItem = process.output.find( + (item) => item.data_type === "Dict", + ); + if (referenceItem) { + console.log("Raw reference data:", referenceItem.data); + try { + // The data might be a Python dict string, try to parse it as JSON + let dataStr = referenceItem.data; + + // If it starts with { but isn't valid JSON, it might be a Python dict + // Try to convert Python dict format to JSON format + if (dataStr.startsWith("{") && !dataStr.startsWith('{"')) { + // Replace Python dict format with JSON format + dataStr = dataStr + .replace(/'/g, '"') // Replace single quotes with double quotes + .replace(/True/g, "true") // Replace Python True with JSON true + .replace(/False/g, "false") // Replace Python False with JSON false + .replace(/None/g, "null"); // Replace Python None with JSON null + } + + console.log("Processed data string:", dataStr); + const parsedData = JSON.parse(dataStr); + referenceOutput = parsedData; + // Keep only non-Dict outputs as main output + mainOutput = process.output.filter( + (item) => item.data_type !== "Dict", + ); + } catch (e) { + console.log("Could not parse reference data:", e); + console.log("Original data:", referenceItem.data); + // If parsing fails, try to extract info using regex as fallback for multiple references + try { + // Updated regex to capture all fields: document_id, document_name, document_position, text + const matches = [ + ...referenceItem.data.matchAll( + /(\d+):\s*\{\s*['"]?document_id['"]?\s*:\s*(\d+).*?['"]?document_name['"]?\s*:\s*['"]([^'"]*)['"]\s*.*?['"]?document_position['"]?\s*:\s*(\d+).*?['"]?text['"]?\s*:\s*['"]([^'"]*)['"]/gs, + ), + ]; + console.log("Regex matches for references:", matches); + + if (matches.length > 0) { + referenceOutput = {}; + matches.forEach((match) => { + const refId = match[1]; + referenceOutput[refId] = { + document_id: parseInt(match[2]), + document_name: match[3], + document_position: parseInt(match[4]), + text: match[5], + }; + }); + mainOutput = process.output.filter( + (item) => item.data_type !== "Dict", + ); + console.log("Fallback parsing successful:", referenceOutput); + } + } catch (fallbackError) { + console.log("Fallback parsing also failed:", fallbackError); + } + } + } + } + return { type: "message", timestamp: process.created, id: process.id, input: process.input, - output: process.output, + output: mainOutput, + referenceOutput: referenceOutput, status: process.status, end_time: process.end_time, }; @@ -212,8 +316,9 @@ export default function GenerativeChat() { whiteSpace: "pre-wrap", }} > - {change.parameter}: {change.oldValue}{" "} - {change.newValue}{" "} + {change.parameter}: {formatHistoryValue(change.oldValue)}{" "} + {" "} + {formatHistoryValue(change.newValue)}{" "} )), }; @@ -226,6 +331,34 @@ export default function GenerativeChat() { setMessagesWithHistory(combinedMessages); }, [messages, history]); + const formatHistoryValue = (value) => { + if (value === null || value === undefined) { + return ""; + } + + if (typeof value === "string") { + return value; + } + + if (typeof value === "number" || typeof value === "boolean") { + return String(value); + } + + if (typeof value === "object") { + if (value.component) { + return value.component; + } + + try { + return JSON.stringify(value); + } catch (error) { + return String(value); + } + } + + return String(value); + }; + return ( + {/* RAG Breadcrumbs - only show for RAG tasks */} + {taskName === "RAGTask" && ( + + + + )} + {/* Model display */} @@ -303,10 +443,10 @@ export default function GenerativeChat() { flexDirection="column" justifyContent="flex-start" flexGrow={0} - gap={4} + gap={1} width={"100%"} //height={"100%"} - mt={4} + mt={1} > {message.type === "history" ? ( @@ -325,13 +465,53 @@ export default function GenerativeChat() { isUser={true} /> {message.status === 3 ? ( - + <> + {taskName === "RAGTask" && message.referenceOutput ? ( + // For RAG messages, we need custom layout to insert sources before timestamp + <> + + + {/* Add timestamp after sources with proper alignment */} + + + {new Date( + message.end_time, + ).toLocaleTimeString()} + + + + ) : ( + // For non-RAG messages, use normal ChatBubble with timestamp + + )} + ) : ( )} diff --git a/DashAI/front/src/components/generative/InfoSessionModal.jsx b/DashAI/front/src/components/generative/InfoSessionModal.jsx index c8b9b7451..600830ab5 100644 --- a/DashAI/front/src/components/generative/InfoSessionModal.jsx +++ b/DashAI/front/src/components/generative/InfoSessionModal.jsx @@ -96,18 +96,30 @@ export default function InfoSessionModal({ sessionData, open, onClose }) { {Object.entries(sessionData.parameters || {}).map( - ([key, value]) => ( - - - {key.replace(/_/g, " ")} - - {value} - - ), + ([key, value]) => { + let displayValue = value; + if ( + typeof value === "object" && + value !== null && + !Array.isArray(value) + ) { + displayValue = JSON.stringify(value, null, 2); + } else if (Array.isArray(value)) { + displayValue = JSON.stringify(value); + } + return ( + + + {key.replace(/_/g, " ")} + + {displayValue} + + ); + }, )}
diff --git a/DashAI/front/src/components/generative/MainGenerativeBox.jsx b/DashAI/front/src/components/generative/MainGenerativeBox.jsx new file mode 100644 index 000000000..2cc0250ed --- /dev/null +++ b/DashAI/front/src/components/generative/MainGenerativeBox.jsx @@ -0,0 +1,20 @@ +import React from "react"; +import { Box } from "@mui/material"; + +export default function MainGenerativeBox({ children }) { + return ( + + {children} + + ); +} diff --git a/DashAI/front/src/components/generative/MediaInput.jsx b/DashAI/front/src/components/generative/MediaInput.jsx index fd7f311e9..7bfe3ac68 100644 --- a/DashAI/front/src/components/generative/MediaInput.jsx +++ b/DashAI/front/src/components/generative/MediaInput.jsx @@ -149,7 +149,6 @@ export function MediaInput({ /> ); })} -