diff --git a/.gitignore b/.gitignore index 54c488e68..d500df9bc 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,8 @@ docs/node_modules/ !docs/docs/build/** !docs/i18n/**/build/ !docs/i18n/**/build/** + +**/csv-train.arrow +**/dataset_info.json + +/RAG_benchmark 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..f80446c49 100644 --- a/DashAI/back/api/api_v1/endpoints/components.py +++ b/DashAI/back/api/api_v1/endpoints/components.py @@ -235,6 +235,8 @@ async def get_components( for component_dict in selected_components.values() ] + return out + @router.get("/{id}/") @inject @@ -315,6 +317,25 @@ async def update_component() -> None: status_code=status.HTTP_501_NOT_IMPLEMENTED, detail="Method not implemented" ) +@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( 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..41d29b632 --- /dev/null +++ b/DashAI/back/api/api_v1/endpoints/documents.py @@ -0,0 +1,469 @@ +import json +import logging +import os +from datetime import datetime +from typing import Any, Dict, List +from urllib.parse import quote + +from fastapi import ( + APIRouter, + Depends, + File, + Form, + HTTPException, + Request, + Response, + UploadFile, + status, +) +from fastapi.responses import FileResponse +from kink import di +from sqlalchemy import exc +from sqlalchemy.orm import sessionmaker + +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.utils import hash_function + +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: + try: + documents: List[DocumentDBModel] = db.query(DocumentDBModel).all() + + documents_responses = [] + base = str(request.base_url).rstrip("/") + for doc in documents: + documents_responses.append( + 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=[ + session.id for session in doc.get_related_sessions + ] + if getattr(doc, "get_related_sessions", None) + else None, + file_url=f"{base}{base_url}/{doc.id}/download", + ) + ) + + return documents_responses + except exc.SQLAlchemyError as e: + log.exception(e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Internal database error", + ) from e + + +@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: + document = db.get(DocumentDBModel, document_id) + if not document: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Document {document_id} does not exist in DB.", + ) + base = str(request.base_url).rstrip("/") + return DocumentResponse( + id=document.id, + file_name=document.file_name, + file_type=document.file_type, + file_hash=document.file_hash, + created=document.created, + last_modified=document.last_modified, + optional_metadata=document.optional_metadata, + related_sessions=[ + session.id for session in document.get_related_sessions + ] + if getattr(document, "get_related_sessions", None) + else None, + file_url=f"{base}{base_url}/{document.id}/download", + ) + except exc.SQLAlchemyError as e: + log.exception(e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Internal database error", + ) from e + + +@router.get("/{document_id}/download", response_class=FileResponse) +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: + document = db.query(DocumentDBModel).filter_by(id=document_id).first() + if not document: + raise HTTPException( + status_code=404, + detail=f"Document {document_id} not found", + ) + if not os.path.exists(document.file_path): + raise HTTPException( + status_code=500, + detail=f"File not found at {document.file_path}", + ) + # Detect media type by file extension + ext = os.path.splitext(document.file_name)[1].lower() + if ext == ".txt": + media_type = "text/plain" + elif ext == ".pdf": + media_type = "application/pdf" + else: + media_type = "application/octet-stream" + + encoded_name = quote(document.file_name, safe="") + if ext == ".pdf": + return FileResponse( + path=document.file_path, + media_type=media_type, + headers={ + "Content-Disposition": (f"inline; filename*=UTF-8''{encoded_name}"), + }, + ) + else: + return FileResponse( + path=document.file_path, + media_type=media_type, + headers={ + "Content-Disposition": ( + f"attachment; filename*=UTF-8''{encoded_name}" + ), + }, + ) + + +@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"]), + request: Request = None, # type: ignore[assignment] +): + """ + Upload a new document to the RAG system with file content and metadata. + The metadata should be a JSON string with required fields and an 'optional_metadata' + dict. + """ + 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") + last_modified = metadata_dict.get("last_modified") + if not file_name or not last_modified: + raise HTTPException( + status_code=400, + detail="Missing required metadata fields: 'filename' and 'last_modified'", + ) + 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 + + file_content_hash = hash_function(content_bytes) + base = str(request.base_url).rstrip("/") if request else "" + with session_factory() as db: + existing_doc: DocumentDBModel = ( + db.query(DocumentDBModel).filter_by(file_hash=file_content_hash).first() + ) + if existing_doc: + # Update existing document's information + existing_doc.file_name = file_name + existing_doc.optional_metadata = optional_metadata + existing_doc.created = datetime.now() + db.commit() + return DocumentResponse( + id=existing_doc.id, + file_name=existing_doc.file_name, + file_type=existing_doc.file_type, + file_hash=existing_doc.file_hash, + created=existing_doc.created, + last_modified=existing_doc.last_modified, + related_sessions=existing_doc.get_related_sessions, + optional_metadata=existing_doc.optional_metadata, + file_url=f"{base}{base_url}/{existing_doc.id}/download", + ) + else: + # Create a new document entry + try: + file_path = docs_folder_path / file_name + + with open(file_path, "wb") as f: + f.write(content_bytes) + + # Get file type from file extension + file_type = os.path.splitext(file_name)[1].lstrip(".") or "unknown" + + doc = DocumentDBModel( + file_name=file_name, + file_type=file_type, + file_path=str(file_path), + file_hash=file_content_hash, + optional_metadata=optional_metadata or None, + ) + db.add(doc) + db.commit() + db.refresh(doc) + + 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, + related_sessions=doc.get_related_sessions, + optional_metadata=doc.optional_metadata, + file_url=f"{base}{base_url}/{doc.id}/download", + ) + except exc.SQLAlchemyError as e: + log.exception(e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Internal database error", + ) from e + + +@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: + # Get the session + session = db.query(GenerativeSession).filter_by(id=session_id).first() + if not session: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Session {session_id} does not exist.", + ) + + # Extract document IDs from session parameters + document_ids = session.parameters.get("documents", []) + if not document_ids: + return [] + + # Query documents by IDs + documents = ( + db.query(DocumentDBModel) + .filter(DocumentDBModel.id.in_(document_ids)) + .all() + ) + + documents_responses = [] + base = str(request.base_url).rstrip("/") + for doc in documents: + documents_responses.append( + 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=[ + session.id for session in doc.get_related_sessions + ] + if doc.pipeline_links + else None, + file_url=f"{base}{base_url}/{doc.id}/download", + ) + ) + + return documents_responses + except exc.SQLAlchemyError as e: + log.exception(e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Internal database error", + ) 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: + document = db.query(DocumentDBModel).filter_by(id=document_id).first() + if not document: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Document {document_id} does not exist in DB.", + ) + if not document.get_related_sessions: + return [] + related_session_ids = [ + session.id for session in document.get_related_sessions + ] + return related_session_ids + except exc.SQLAlchemyError as e: + log.exception(e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Internal database error", + ) 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: + document = db.query(DocumentDBModel).filter_by(id=document_id).first() + if not document: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Document {document_id} does not exist in DB.", + ) + + # Delete the physical file + if os.path.exists(document.file_path): + os.remove(document.file_path) + + # Delete the database record + db.delete(document) + db.commit() + + return Response(status_code=status.HTTP_204_NO_CONTENT) + except exc.SQLAlchemyError as e: + log.exception(e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Internal database error", + ) from e + except OSError as e: + log.exception(e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Error deleting physical file", + ) 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"]), + request: Request = None, # type: ignore[assignment] +): + """Update a document's metadata.""" + base = str(request.base_url).rstrip("/") if request else "" + with session_factory() as db: + try: + document = db.query(DocumentDBModel).filter_by(id=document_id).first() + if not document: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Document {document_id} does not exist in DB.", + ) + + try: + metadata_dict = json.loads(metadata) + except json.JSONDecodeError: + raise HTTPException( + status_code=400, detail="Invalid JSON in metadata" + ) from None + + # Update fields that are allowed to be modified + if "file_name" in metadata_dict: + document.file_name = metadata_dict["file_name"] + # Update file_type when file_name changes + document.file_type = ( + os.path.splitext(metadata_dict["file_name"])[1].lstrip(".") + or "unknown" + ) + + if "optional_metadata" in metadata_dict: + optional_metadata = metadata_dict["optional_metadata"] + if not isinstance(optional_metadata, dict): + raise HTTPException( + status_code=400, + detail="'optional_metadata' must be a dictionary", + ) + document.optional_metadata = optional_metadata + + document.last_modified = datetime.now() + + db.commit() + db.refresh(document) + + return DocumentResponse( + id=document.id, + file_name=document.file_name, + file_hash=document.file_hash, + created=document.created, + optional_metadata=document.optional_metadata, + related_sessions=[ + session.id for session in document.get_related_sessions + ] + if document.get_related_sessions + else None, + file_url=f"{base}{base_url}/{document.id}/download", + ) + + except exc.SQLAlchemyError as e: + log.exception(e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Internal database error", + ) from e 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..b5a1feb41 100644 --- a/DashAI/back/api/api_v1/endpoints/generative_session.py +++ b/DashAI/back/api/api_v1/endpoints/generative_session.py @@ -1,6 +1,12 @@ import logging +import shutil from datetime import datetime -from typing import TYPE_CHECKING, Union +from pathlib import Path +from typing import TYPE_CHECKING, Final, Union + +COMPOSITE_RETRIEVER_NAMES: Final = frozenset( + {"SequentialRetriever", "ParallelRetriever"} +) from fastapi import APIRouter, Depends, HTTPException, status from kink import di @@ -10,11 +16,22 @@ GenerativeSessionParams, ) from DashAI.back.dependencies.database.models import ( + Document, GenerativeProcess, GenerativeSession, GenerativeSessionParameterHistory, ProcessData, + RAGChunkingModel, + RAGDenseRetriever, + RAGEmbeddingMatrix, + RAGEmbeddingModel, + RAGPipeline, + RAGPrompt, + RAGRetriever, + RAGRetrieverChild, + RAGSparseRetriever, ) +from DashAI.back.tasks.RAG_task import RAGTask if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker @@ -26,6 +43,254 @@ log = logging.getLogger(__name__) +def _delete_path(path_value: str | None) -> None: + if not path_value: + return + path = Path(path_value) + if path.exists(): + shutil.rmtree(path, ignore_errors=True) + + +def _other_sessions_with_same_config( + db, + session_id: int, + expected_parameters: dict, + *, + keys: tuple[str, ...], +) -> bool: + """Return True when any other session matches all keys in `keys`. + + This avoids generator/list comprehensions for clarity. + """ + other_sessions = ( + db.query(GenerativeSession).filter(GenerativeSession.id != session_id).all() + ) + + for other_session in other_sessions: + other_parameters = other_session.parameters or {} + all_keys_match = True + for key in keys: + if other_parameters.get(key) != expected_parameters.get(key): + all_keys_match = False + break + if all_keys_match: + return True + + return False + + +def _cleanup_orphaned_rag_resources( + db, + session_id: int, + old_parameters: dict, + new_parameters: dict | None = None, +) -> None: + 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) + + other_sessions = ( + db.query(GenerativeSession).filter(GenerativeSession.id != session_id).all() + ) + + documents_ids = old_parameters.get("documents") or [] + documents_ids = sorted(documents_ids) + + # ── Retriever cleanup MUST run BEFORE chunking cleanup ────────── + # _cleanup_unit_retriever queries RAGChunkingModel to obtain + # chunking_model_id. If chunking models are deleted first, the + # query returns None and retriever rows are silently skipped. + # ───────────────────────────────────────────────────────────────── + retriever_model_params = old_parameters.get("retriever_model") + if not retriever_model_params: + retriever_model_params = {} + + retriever_component_name = retriever_model_params.get("component", "") + + should_cleanup_retriever = ( + bool(retriever_model_params) + and _component_changed("retriever_model") + and not _other_sessions_with_same_config( + db, + session_id, + old_parameters, + keys=("documents", "chunking_model", "retriever_model"), + ) + ) + + if should_cleanup_retriever: + if retriever_component_name in COMPOSITE_RETRIEVER_NAMES: + _cleanup_composite_retriever(db, old_parameters, documents_ids, session_id) + else: + _cleanup_unit_retriever( + db, old_parameters, documents_ids, retriever_model_params, session_id + ) + + # ── Chunking model cleanup (AFTER retriever) ────────────────────── + chunking_model_params = old_parameters.get("chunking_model") + if not chunking_model_params: + chunking_model_params = {} + + should_cleanup_chunking = ( + bool(chunking_model_params) + and _component_changed("chunking_model") + and not _other_sessions_with_same_config( + db, session_id, old_parameters, keys=("documents", "chunking_model") + ) + ) + + if should_cleanup_chunking: + chunking_models = ( + 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: + db.delete(chunking_model) + + +def _find_pipeline_id(db, session_id: int) -> int | None: + pipeline = db.query(RAGPipeline).filter_by(session_id=session_id).first() + return pipeline.id if pipeline else None + + +def _cleanup_composite_retriever( + db, + old_parameters, + documents_ids, + session_id, +) -> None: + pipeline_id = _find_pipeline_id(db, session_id) + if pipeline_id is None: + return + + retriever_component_name = old_parameters.get("retriever_model", {}).get( + "component" + ) + + composite_bridges = ( + db.query(RAGRetriever) + .filter( + RAGRetriever.pipeline_id == pipeline_id, + RAGRetriever.class_name == retriever_component_name, + ) + .all() + ) + + for bridge in composite_bridges: + child_links = ( + db.query(RAGRetrieverChild) + .filter_by(parent_id=bridge.id) + .order_by(RAGRetrieverChild.child_order) + .all() + ) + for link in child_links: + child_bridge = db.query(RAGRetriever).get(link.child_id) + if child_bridge is None: + continue + if child_bridge.sparse_detail: + _delete_path(child_bridge.sparse_detail.storage_folder) + db.delete(child_bridge.sparse_detail) + elif child_bridge.dense_detail: + db.delete(child_bridge.dense_detail) + db.delete(bridge) + + +def _cleanup_unit_retriever( + db, + old_parameters, + documents_ids, + retriever_model_params, + session_id, +) -> None: + # Previously queried non-existent columns (chunking_model_id, document_ids) + # on RAGDenseRetriever/RAGSparseRetriever. Fixed after schema refactor: + # trace through the pipeline chain to obtain the valid chunk_set_id column. + retriever_params = retriever_model_params.get("params", {}) + + pipeline_id = _find_pipeline_id(db, session_id) + if pipeline_id is None: + return + + bridge = ( + db.query(RAGRetriever) + .filter( + RAGRetriever.pipeline_id == pipeline_id, + RAGRetriever.class_name == retriever_model_params.get("component"), + ) + .first() + ) + if bridge is None: + return + + if "encoding_model" in retriever_params: + dense_retriever = ( + db.query(RAGDenseRetriever) + .filter( + RAGDenseRetriever.bridge_id == bridge.id, + RAGDenseRetriever.class_name == retriever_model_params.get("component"), + RAGDenseRetriever.parameters == retriever_params, + ) + .first() + ) + if dense_retriever is None: + return + + chunk_set_id = dense_retriever.chunk_set_id + embedding_model_id = dense_retriever.embedding_model_id + + embedding_matrices = ( + 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: + _delete_path(matrix.storage_folder) + matrix_ids.append(matrix.id) + + if matrix_ids: + db.query(RAGEmbeddingMatrix).filter( + RAGEmbeddingMatrix.id.in_(matrix_ids) + ).delete(synchronize_session=False) + + embedding_model = db.query(RAGEmbeddingModel).get(embedding_model_id) + if embedding_model is not None: + db.delete(embedding_model) + + db.delete(dense_retriever) + db.delete(bridge) + else: + sparse_retriever = ( + db.query(RAGSparseRetriever) + .filter( + RAGSparseRetriever.bridge_id == bridge.id, + RAGSparseRetriever.class_name + == retriever_model_params.get("component"), + RAGSparseRetriever.parameters == retriever_params, + ) + .first() + ) + if sparse_retriever is None: + return + + _delete_path(sparse_retriever.storage_folder) + db.delete(sparse_retriever) + db.delete(bridge) + + @router.post("/", status_code=status.HTTP_201_CREATED) async def upload_generative_session( params: GenerativeSessionParams, @@ -55,22 +320,53 @@ 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: + try: + assert params.parameters["documents"] + assert isinstance(params.parameters["documents"], list) + assert len(params.parameters["documents"]) > 0 + + except (AssertionError, KeyError) as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="RAG Task requires a non-empty list of document IDs.", + ) from e + + documents_ids = [] + for doc_id in params.parameters["documents"]: + document = db.get(Document, doc_id) + if not document: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Document with ID {doc_id} does not exist.", + ) + documents_ids.append(document.id) + + params.parameters["documents"] = documents_ids + # Continue with the session creation + + # Normalise frontend properties wrapper and validate + from DashAI.back.core.schema_fields.utils import normalize_payload + params.parameters = normalize_payload(params.parameters) 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 # Check if the task is a subclass of BaseGenerativeTask @@ -257,6 +553,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,6 +584,7 @@ async def delete_generative_session( db.delete(entry) # Finally, delete the session itself db.delete(session) + _cleanup_orphaned_rag_resources(db, session_id, old_parameters) db.commit() except exc.SQLAlchemyError as e: log.exception(e) @@ -405,6 +704,7 @@ 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. @@ -438,7 +738,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 {}) + updated_parameters = {**old_parameters, **new_params} + + if "prompt_id" in updated_parameters: + if "prompt" not in updated_parameters: + prompt_db = db.get(RAGPrompt, updated_parameters["prompt_id"]) + if prompt_db: + raw_params = dict(prompt_db.parameters or {}) + prompt_params = { + "template": raw_params.get( + "template", + raw_params.get("templates", {}).get("en", ""), + ), + } + if "language" in raw_params: + prompt_params["language"] = raw_params["language"] + updated_parameters["prompt"] = { + "component": prompt_db.class_name, + "params": prompt_params, + } + del updated_parameters["prompt_id"] + + try: + task_class = component_registry[session.task_name]["class"] + except KeyError: + task_class = None + + if task_class is not None and task_class == RAGTask: + from DashAI.back.models.RAG.RAG_pipeline import RAGPipeline + + try: + RAGPipeline.SCHEMA.model_validate(updated_parameters) + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Invalid parameters for RAG session {session_id}: {e}", + ) from e session_params_entry = GenerativeSessionParameterHistory( session_id=session.id, @@ -450,11 +786,17 @@ async def update_generative_session_params( session.parameters = updated_parameters session.last_modified = datetime.now() + _cleanup_orphaned_rag_resources( + db, session_id, old_parameters, updated_parameters + ) 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 +848,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 +914,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/prompts.py b/DashAI/back/api/api_v1/endpoints/prompts.py new file mode 100644 index 000000000..e3170390a --- /dev/null +++ b/DashAI/back/api/api_v1/endpoints/prompts.py @@ -0,0 +1,365 @@ +import logging +from datetime import datetime +from typing import Any, Dict + +from fastapi import APIRouter, Depends, HTTPException, status +from kink import di +from sqlalchemy import exc, select +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, + GenerativeSessionParameterHistory, + RAGPrompt, +) +from DashAI.back.dependencies.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 + +router = APIRouter() +log = logging.getLogger(__name__) + + +def _validate_prompt_template( + prompt_class_name: str, + template: str, + component_registry: ComponentRegistry, +) -> None: + if prompt_class_name not in component_registry: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Component {prompt_class_name} is not registered.", + ) + + prompt_component = component_registry[prompt_class_name] + prompt_class = prompt_component["class"] + if not issubclass(prompt_class, Prompt): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Component {prompt_class_name} is not a valid Prompt subclass.", + ) + + if not prompt_class.validate_template(template): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=( + f"Invalid template for prompt {prompt_class_name}. " + f"Required tokens are: {prompt_class.get_required_placeholders()}" + ), + ) + + +def _serialize_prompt(prompt: RAGPrompt) -> Dict[str, Any]: + return { + "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(base_name: str, session_id: int, db) -> str: + candidate = f"{base_name} - session {session_id}" + suffix = 2 + while db.execute(select(RAGPrompt.id).where(RAGPrompt.name == candidate)).scalar(): + candidate = f"{base_name} - session {session_id} ({suffix})" + suffix += 1 + return candidate + + +@router.post("/", status_code=status.HTTP_201_CREATED) +async def create_rag_prompt( + 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: + if not prompt.parameters: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail="Prompt parameters are required and must include a template.", + ) + + if "templates" in prompt.parameters: + for _lang, tmpl in prompt.parameters["templates"].items(): + _validate_prompt_template(prompt.class_name, tmpl, component_registry) + elif "template" in prompt.parameters: + _validate_prompt_template( + prompt.class_name, + prompt.parameters["template"], + component_registry, + ) + else: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail="Prompt parameters must include 'template' or 'templates'.", + ) + new_prompt = RAGPrompt( + class_name=prompt.class_name, + name=prompt.name, + parameters=prompt.parameters, + ) + db.add(new_prompt) + db.commit() + db.refresh(new_prompt) + return {"id": new_prompt.id} + + +@router.patch("/{prompt_id}", status_code=status.HTTP_200_OK) +async def update_rag_prompt( + 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: + existing_prompt = db.get(RAGPrompt, prompt_id) + if existing_prompt is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="Prompt not found", + ) + + changed = False + + if prompt.name is not None: + new_name = prompt.name.strip() + if not new_name: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail="Prompt name cannot be empty", + ) + if new_name != existing_prompt.name: + existing_prompt.name = new_name + changed = True + + if prompt.parameters is not None: + if "templates" in prompt.parameters: + for _lang, tmpl in prompt.parameters["templates"].items(): + _validate_prompt_template( + existing_prompt.class_name, + tmpl, + component_registry, + ) + elif "template" in prompt.parameters: + _validate_prompt_template( + existing_prompt.class_name, + prompt.parameters["template"], + component_registry, + ) + else: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail=( + "Prompt parameters must include 'template' or 'templates'." + ), + ) + existing_prompt.parameters = prompt.parameters + changed = True + + if not changed: + raise HTTPException( + status_code=status.HTTP_304_NOT_MODIFIED, + detail="Record not modified", + ) + + existing_prompt.last_modified = datetime.now() + db.commit() + db.refresh(existing_prompt) + return _serialize_prompt(existing_prompt) + except HTTPException: + raise + except exc.SQLAlchemyError as e: + db.rollback() + log.exception(e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Internal database error", + ) from e + + +@router.post("/{prompt_id}/sessions/{session_id}", status_code=status.HTTP_201_CREATED) +async def update_rag_prompt_for_session( + 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: + existing_prompt = db.get(RAGPrompt, prompt_id) + if existing_prompt is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="Prompt not found", + ) + + session = db.get(GenerativeSession, session_id) + if session is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="Generative session not found", + ) + + if prompt.parameters is not None: + if "templates" in prompt.parameters: + for _lang, tmpl in prompt.parameters["templates"].items(): + _validate_prompt_template( + existing_prompt.class_name, + tmpl, + component_registry, + ) + elif "template" in prompt.parameters: + _validate_prompt_template( + existing_prompt.class_name, + prompt.parameters["template"], + component_registry, + ) + else: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail=( + "Prompt parameters must include 'template' or 'templates'." + ), + ) + new_parameters = prompt.parameters + else: + new_parameters = existing_prompt.parameters + + # Make parameters unique per session to avoid UNIQUE constraint + new_parameters = dict(new_parameters or {}) + new_parameters["cloned_for_session"] = session_id + + if prompt.name: + candidate_name = prompt.name + elif existing_prompt.name: + candidate_name = existing_prompt.name + else: + candidate_name = existing_prompt.class_name + + base_name = (candidate_name or "").strip() + new_name = _build_session_prompt_name(base_name, session_id, db) + + new_prompt = RAGPrompt( + class_name=existing_prompt.class_name, + name=new_name, + parameters=new_parameters, + ) + db.add(new_prompt) + db.commit() + 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() + db.add( + GenerativeSessionParameterHistory( + session_id=session.id, + parameters=session_parameters, + modified_at=datetime.now(), + ) + ) + db.commit() + db.refresh(session) + + return { + "prompt": _serialize_prompt(new_prompt), + "session_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, + detail="Internal database error", + ) 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. + + Parameters + ---------- + session_factory : Callable[..., ContextManager[Session]] + A factory that creates a context manager that handles a SQLAlchemy session. + The generated session can be used to access and query the database. + + Returns + ------- + list + A list of dictionaries with all prompts on the database. + + Raises + ------ + HTTPException + If there's an internal database error. + """ + with session_factory() as db: + try: + prompts = 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", + }, + ) + db.add(default_generation_prompt) + db.add(default_qa_prompt) + db.commit() + prompts = db.query(RAGPrompt).all() + + prompt_responses = [] + for prompt in prompts: + prompt_responses.append( + { + "id": prompt.id, + "class_name": prompt.class_name, + "name": prompt.name, + "parameters": prompt.parameters, + "created": prompt.created, + "last_modified": prompt.last_modified, + } + ) + + return prompt_responses + except exc.SQLAlchemyError as e: + log.exception(e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Internal database error", + ) from e diff --git a/DashAI/back/api/api_v1/schemas/__init__.py b/DashAI/back/api/api_v1/schemas/__init__.py index e69de29bb..9d35150b8 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 \ No newline at end of file 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..1bde38a24 --- /dev/null +++ b/DashAI/back/api/api_v1/schemas/document.py @@ -0,0 +1,15 @@ +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 + 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..317c8d303 --- /dev/null +++ b/DashAI/back/api/api_v1/schemas/rag_prompt.py @@ -0,0 +1,14 @@ +from typing import Any, Dict, Optional + +from pydantic import BaseModel + + +class RAGPromptSchema(BaseModel): + class_name: str + name: str + parameters: Optional[Dict[str, Any]] = None + + +class RAGPromptUpdateSchema(BaseModel): + 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..b56a4821e 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..3eb485d61 100644 --- a/DashAI/back/config_object.py +++ b/DashAI/back/config_object.py @@ -47,5 +47,11 @@ 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..fe5f98d14 100644 --- a/DashAI/back/converters/hugging_face/embedding.py +++ b/DashAI/back/converters/hugging_face/embedding.py @@ -1,3 +1,9 @@ +from typing import List + +import numpy as np +import torch +from transformers import AutoModel, AutoTokenizer + """HuggingFace embedding converter with lazy-loaded dependencies.""" from typing import TYPE_CHECKING @@ -178,12 +184,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 +238,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 +247,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/schema_fields/base_schema.py b/DashAI/back/core/schema_fields/base_schema.py index 3371ffd4b..3a4571c4f 100644 --- a/DashAI/back/core/schema_fields/base_schema.py +++ b/DashAI/back/core/schema_fields/base_schema.py @@ -46,3 +46,15 @@ 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/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..938b0c1fd 100644 --- a/DashAI/back/dependencies/database/models.py +++ b/DashAI/back/dependencies/database/models.py @@ -7,6 +7,7 @@ JSON, BigInteger, Boolean, + Column, DateTime, Enum, Float, @@ -307,7 +308,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 +548,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 +752,517 @@ 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[int] = mapped_column( + ForeignKey("rag_chunking_model.id", ondelete="CASCADE"), nullable=False + ) + prompt_id: Mapped[int] = mapped_column( + ForeignKey("rag_prompt.id", ondelete="CASCADE"), nullable=False + ) + generation_model_id: Mapped[int] = mapped_column( + ForeignKey("rag_generation_model.id", ondelete="CASCADE"), nullable=False + ) + + # 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 = Column(Integer, primary_key=True, autoincrement=True) + document_id = Column( + Integer, ForeignKey("document.id", ondelete="CASCADE"), nullable=False + ) + session_id = Column( + Integer, ForeignKey("generative_session.id", ondelete="CASCADE"), nullable=False + ) + pipeline_id = Column( + Integer, 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/dependencies/registry/component_registry.py b/DashAI/back/dependencies/registry/component_registry.py index d99145aaf..e9b1e0209 100644 --- a/DashAI/back/dependencies/registry/component_registry.py +++ b/DashAI/back/dependencies/registry/component_registry.py @@ -217,6 +217,7 @@ def register_component(self, new_component: Type) -> None: "metadata": _metadata, "description": getattr(new_component, "DESCRIPTION", None), "display_name": getattr(new_component, "DISPLAY_NAME", None), + "flags": getattr(new_component, "FLAGS", []), "color": getattr(new_component, "COLOR", None), } diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 49ea258f1..160008c18 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,57 @@ 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, + HuggingFaceEmbedding, + 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 +334,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 +370,7 @@ def get_initial_components(): TextToImageGenerationTask, TextToTextGenerationTask, ControlNetTask, + RAGTask, ImageClassificationTask, # Models AdaBoostClassifier, @@ -340,6 +397,8 @@ def get_initial_components(): HistGradientBoostingClassifier, HistGradientBoostingRegression, KNeighborsClassifier, + QwenModel, + RAGPipeline, KNeighborsRegression, LassoRegression, LinearRegression, @@ -379,6 +438,9 @@ def get_initial_components(): StableDiffusionV3Model, StableDiffusionXLModel, StableDiffusionXLV1ControlNet, + OpenAITextToTextGenerationModel, + DeepSeekTextToTextGenerationModel, + Phi4MiniInstructModel, SVC, SVR, T5SmallTransformer, @@ -433,6 +495,7 @@ def get_initial_components(): DatasetJob, GenerativeJob, PipelineJob, + RAGJob, # Explainers KernelShap, PartialDependence, @@ -495,6 +558,38 @@ def get_initial_components(): SMOTEConverter, SMOTEENNConverter, RandomUnderSamplerConverter, + # RAG + RAGPipeline, + # Chunking Models + CharacterChunkModel, + RecursiveCharacterChunkModel, + TokenChunkModel, + # Encodings + OpenAIEmbedding, + HuggingFaceEmbedding, + 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/generative_job.py b/DashAI/back/job/generative_job.py index e52c58a72..17af4d511 100644 --- a/DashAI/back/job/generative_job.py +++ b/DashAI/back/job/generative_job.py @@ -111,7 +111,23 @@ 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) + + from DashAI.back.job.rag_job import RAGJob + from DashAI.back.tasks.RAG_task import RAGTask + + 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: diff --git a/DashAI/back/job/rag_job.py b/DashAI/back/job/rag_job.py new file mode 100644 index 000000000..01253b582 --- /dev/null +++ b/DashAI/back/job/rag_job.py @@ -0,0 +1,313 @@ +import logging +from typing import TYPE_CHECKING, Any + +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.models.RAG.document_loader import DocumentLoader +from DashAI.back.models.RAG.pipeline_repository import PipelineRepository +from DashAI.back.models.RAG.rag_models_factory import RAGModelsFactory +from DashAI.back.models.RAG.RAG_pipeline import ( + RAGPipeline, + RAGPipelineConfig, +) +from DashAI.back.tasks.RAG_task import RAGTask + +if TYPE_CHECKING: + from sqlalchemy.orm import sessionmaker + +log = logging.getLogger(__name__) + +# Known RAG pipeline parameter keys accepted from generative_session.parameters. +_RAG_PARAM_KEYS: frozenset = frozenset( + {"documents", "prompt", "chunking_model", "retriever_model", "generation_model"} +) + + +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.""" + 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.""" + generative_process_id: int = self.kwargs.get("generative_process_id") + if generative_process_id is None: + 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.""" + generative_process_id = self.kwargs.get("generative_process_id") + if not generative_process_id: + return "RAG Process" + + from kink import di + + 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. + + Loads the generative process and session, filters parameters via + whitelist, instantiates the RAG pipeline, runs inference, and + persists the output. + + Raises: + JobError: If any stage of execution fails. + """ + import gc + + import torch + from kink import di + + component_registry = di["component_registry"] + session_factory = di["session_factory"] + config = di["config"] + + model = None + generative_process = None + # NOTE: The DB session spans the entire job lifecycle including LLM + # inference, which risks connection timeouts for long-running + # generations. No obvious solution yet. + with session_factory() as db: + try: + generative_process_id: int = self.kwargs["generative_process_id"] + + try: + generative_process: GenerativeProcess = db.get( + GenerativeProcess, generative_process_id + ) + if not generative_process: + raise JobError( + f"Generative process {generative_process_id} " + "not found in DB." + ) + log.debug("Loaded generative process %d", generative_process_id) + except Exception as e: + log.exception(e) + if generative_process: + generative_process.set_status_as_error() + db.commit() + raise JobError("Error retrieving generative process.") from e + + try: + generative_session: GenerativeSession = db.get( + GenerativeSession, generative_process.session_id + ) + if not generative_session: + raise JobError( + f"Session {generative_process.session_id} not found in DB." + ) + log.debug( + "Loaded generative session %d", generative_session.id + ) + except Exception as e: + log.exception(e) + generative_process.set_status_as_error() + db.commit() + raise JobError("Error retrieving generative session.") from e + + try: + # Whitelist-only: accept only known RAG pipeline keys. + # generative_session.parameters is a JSON column that may + # contain legacy or metadata keys (e.g. prompt_id). + 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 = {} + for key in raw_params: + if key in _RAG_PARAM_KEYS: + clean_params[key] = raw_params[key] + pipeline_config = RAGPipelineConfig.from_kwargs( + db=db, + component_registry=component_registry, + session_id=generative_session.id, + env_rag_path=config["RAG_PATH"], + **clean_params, + ) + log.debug("Created RAG pipeline config for session %d", generative_session.id) + models = RAGModelsFactory( + db, + component_registry, + config["RAG_PATH"], + ) + repo = PipelineRepository(db) + doc_loader = DocumentLoader(db) + model: RAGPipeline = RAGPipeline( + pipeline_config, + models, + repo, + doc_loader, + ) + log.debug("Created RAG pipeline model for session %d", generative_session.id) + except Exception as e: + log.exception(e) + generative_process.set_status_as_error() + db.commit() + raise JobError( + "Error instantiating RAG pipeline with given parameters." + ) from e + + input_data = generative_process.input + + try: + task = RAGTask() + except Exception as e: + log.exception(e) + generative_process.set_status_as_error() + db.commit() + raise JobError("Error instantiating RAG task.") from e + + try: + history = [ + (proc.input[0].data, proc.output[0].data) + for proc in db.query(GenerativeProcess) + .filter(GenerativeProcess.session_id == generative_session.id) + .filter(GenerativeProcess.status == RunStatus.FINISHED) + .all() + ] + input_data = task.prepare_for_task( + input_data, + history=history, + ) + except Exception as e: + log.exception(e) + generative_process.set_status_as_error() + db.commit() + raise JobError("Error preparing task with history.") from e + + try: + log.debug("Starting generation for process %d", generative_process_id) + generative_process.set_status_as_started() + db.commit() + except exc.SQLAlchemyError as e: + log.exception(e) + generative_process.set_status_as_error() + db.commit() + raise JobError( + "Failed to update process status in database." + ) from e + + try: + output: Any = model.generate(input_data) + log.debug("Generation completed for process %d", generative_process_id) + except Exception as e: + log.exception(e) + generative_process.set_status_as_error() + db.add( + ProcessData( + data=f"Error details: {str(e)}", + data_type="str", + process_id=generative_process.id, + is_input=False, + ) + ) + db.commit() + raise JobError("Error during RAG model generation.") from e + + try: + output = task.process_output( + output, images_path=config["IMAGES_PATH"] + ) + outputs_for_database = [] + for o in output: + if not isinstance(o, tuple) or len(o) != 2: + raise JobError( + "Output from task must be a list of " + "tuples (data, type)." + ) + output_data, output_type = o + process_data = ProcessData( + data=output_data, + data_type=output_type, + process_id=generative_process.id, + is_input=False, + ) + outputs_for_database.append(process_data) + + db.add_all(outputs_for_database) + db.commit() + + db.refresh(generative_process) + generative_process.set_status_as_finished() + db.commit() + log.debug("Output processing completed for process %d", generative_process_id) + except Exception as e: + log.exception(e) + generative_process.set_status_as_error() + db.commit() + 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_pipeline.py b/DashAI/back/models/RAG/RAG_pipeline.py new file mode 100644 index 000000000..a61cd9979 --- /dev/null +++ b/DashAI/back/models/RAG/RAG_pipeline.py @@ -0,0 +1,502 @@ +"""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 + +log = logging.getLogger(__name__) + +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.database.models import ( + RAGChunkSet, +) +from DashAI.back.dependencies.database.models import ( + RAGPipeline as PipelineDBModel, +) +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.base_generative_model import BaseGenerativeModel +from DashAI.back.models.RAG.chunk_set_utils import get_or_create_chunk_set +from DashAI.back.models.RAG.chunking_models.chunking_model_factory import ( + ChunkingFactoryResult, +) +from DashAI.back.models.RAG.document_loader import DocumentLoader +from DashAI.back.models.RAG.documents import BaseDocument, Chunk +from DashAI.back.models.RAG.llm_factory import LLMFactoryResult +from DashAI.back.models.RAG.pipeline_repository import PipelineRepository +from DashAI.back.models.RAG.prompts.prompt_factory import PromptFactoryResult +from DashAI.back.models.RAG.rag_models_factory import RAGModelsFactory +from DashAI.back.models.RAG.retrievers.retriever_factory import ( + RetrieverFactoryResult, +) +from DashAI.back.models.text_to_text_generation_model import ( + TextToTextGenerationTaskModel, +) + +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 + + +class RAGPipelineError(Exception): + """Base exception for RAG pipeline errors.""" + + +class RAGPipelineParametersError(RAGPipelineError): + """Backward-compatible alias for RAGPipelineConfigError.""" + + +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.""" + + +@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: Any + 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 + documents: List[int] + prompt: ModelRef + chunking_model: ModelRef + retriever_model: ModelRef + generation_model: ModelRef + + _MODEL_KEYS: Tuple[str, str, str, str] = ( + "prompt", + "chunking_model", + "retriever_model", + "generation_model", + ) + _INFRA_KEYS: Tuple[str, str, str, str] = ( + "session_id", + "db", + "component_registry", + "env_rag_path", + ) + _PARAM_KEYS: Tuple[str] = ("documents",) + + @classmethod + def validate_model_refs(cls, params: Dict[str, Any]) -> None: + for key in cls._MODEL_KEYS: + if key not in params: + raise RAGPipelineConfigError(f"Missing '{key}'") + raw = params[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}'") + if "documents" not in params: + raise RAGPipelineConfigError("Missing required parameter 'documents'") + + @classmethod + def from_kwargs(cls, **kwargs: Any) -> "RAGPipelineConfig": + missing: List[str] = [] + for key in cls._INFRA_KEYS: + if key not in kwargs: + missing.append(key) + for key in cls._PARAM_KEYS: + if key not in kwargs: + missing.append(key) + for key in cls._MODEL_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 cls._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"], + ) + + all_known: set[str] = ( + set(cls._INFRA_KEYS) | set(cls._PARAM_KEYS) | set(cls._MODEL_KEYS) + ) + extra: set[str] = set(kwargs) - all_known + 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): + 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 + documents_ids: List[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 + + 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, + models: RAGModelsFactory, + repo: PipelineRepository, + doc_loader: DocumentLoader, + ) -> None: + pipeline_record: PipelineDBModel = repo.ensure_db_record( + config.session_id, + ) + + documents = doc_loader.load(config.documents) + + chunk_set: RAGChunkSet = get_or_create_chunk_set( + db=config.db, + document_ids=config.documents, + pipeline_config={ + "chunking_model": { + "component": config.chunking_model.component, + "params": config.chunking_model.params, + } + }, + ) + + prompt_result: PromptFactoryResult = models.create_prompt( + config.prompt.component, + config.prompt.params, + ) + + chunking_result: ChunkingFactoryResult = models.create_chunking_model( + documents, + chunk_set.id, + config.chunking_model.component, + config.chunking_model.params, + ) + + retriever_result: RetrieverFactoryResult = models.create_retriever( + pipeline_record.id, + chunking_result.chunks, + chunk_set.id, + config.retriever_model.component, + config.retriever_model.params, + ) + + llm_result: LLMFactoryResult = models.create_llm( + config.generation_model.component, + config.generation_model.params, + ) + + repo.update_db_record( + session_id=config.session_id, + chunking_model_id=chunking_result.db_record_id, + prompt_id=prompt_result.db_record_id, + generation_model_id=llm_result.db_record_id, + ) + + self.session_id = config.session_id + self.pipeline_id = pipeline_record.id + self.documents_ids = config.documents + self.documents = documents + self.prompt_model = prompt_result.model + self.chunking_model_id = chunking_result.db_record_id + self.chunking_model = chunking_result.model + self.chunks = chunking_result.chunks + self.retriever = retriever_result.model + self.llm_model = llm_result.model + + def single_interaction( + self, + query: str, + ) -> List[Chunk]: + """Retrieve the top-K chunks for a single query. + + Parameters + ---------- + query : str + The search query string. + + Returns + ------- + List[Chunk] + 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. + + Parameters + ---------- + chunks : List[Chunk] + Chunks returned by the retriever. + + Returns + ------- + Tuple[str, Dict[str, ChunkReference]] + 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 = self.documents[document_id] + 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 → generate. + + Parameters + ---------- + input_data : Tuple[Dict[str, str], ...] + Chat-format input with history as earlier entries and the + current user message as the last entry. + + Returns + ------- + RAGGenerationOutput + The generated message and the chunks used for retrieval. + + Raises + ------ + RAGPipelineRuntimeError + If any stage fails (retrieval, formatting, or generation). + """ + if not input_data: + raise RAGPipelineRuntimeError("input_data must not be empty.") + try: + input_dict: Dict[str, str] = input_data[-1] + input_message: str = input_dict["content"] + history: Tuple[Dict[str, str], ...] = input_data[:-1] + 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) + return RAGGenerationOutput(message=output[0], chunks=chunk_dict) + 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..bbfd5ce7a --- /dev/null +++ b/DashAI/back/models/RAG/__init__.py @@ -0,0 +1 @@ +from DashAI.back.models.RAG.RAG_pipeline import RAGPipeline # noqa: F401 diff --git a/DashAI/back/models/RAG/chunk_set_utils.py b/DashAI/back/models/RAG/chunk_set_utils.py new file mode 100644 index 000000000..7daebf40b --- /dev/null +++ b/DashAI/back/models/RAG/chunk_set_utils.py @@ -0,0 +1,57 @@ +import hashlib +import json +from typing import Any, Dict, List + +from sqlalchemy.orm import Session + +from DashAI.back.dependencies.database.models import ( + RAGChunkSet, + RAGChunkSetDocument, +) + + +def _build_chunk_set_signature( + document_ids: List[int], + pipeline_config: Dict[str, Any], +) -> str: + payload = json.dumps( + { + "doc_ids": sorted(document_ids), + "config": dict(sorted(pipeline_config.items())), + }, + sort_keys=True, + ) + return hashlib.sha256(payload.encode()).hexdigest() + + +# NOTE: This does a SELECT then INSERT without a lock, which is safe for +# a single-user platform but could cause IntegrityError under concurrent +# session creation with identical document sets. +def get_or_create_chunk_set( + db: Session, + document_ids: List[int], + pipeline_config: Dict[str, Any], +) -> RAGChunkSet: + signature = _build_chunk_set_signature(document_ids, pipeline_config) + existing = db.query(RAGChunkSet).filter_by(signature=signature).first() + if existing: + return existing + + chunk_set = RAGChunkSet( + signature=signature, + parameters=pipeline_config, + ) + db.add(chunk_set) + db.commit() + db.refresh(chunk_set) + + for doc_id in sorted(document_ids): + db.add( + RAGChunkSetDocument( + chunk_set_id=chunk_set.id, + document_id=doc_id, + ) + ) + + db.commit() + return chunk_set 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..49eda0a1f --- /dev/null +++ b/DashAI/back/models/RAG/chunking_models/__init__.py @@ -0,0 +1,6 @@ +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 \ No newline at end of file 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..45d496f96 --- /dev/null +++ b/DashAI/back/models/RAG/chunking_models/base_chunking_model.py @@ -0,0 +1,113 @@ +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 RAGWorkflowError + + +class BaseChunkingModel(ConfigObject, metaclass=ABCMeta): + """ + Base class for chunking models. + This class should be inherited by any specific chunking model implementation. + """ + + TYPE: Final[str] = "ChunkingModel" + REQUIRED_EXTRA_KWARGS: Final[List[str]] = ["documents"] + + def __init__(self, **kwargs): + """ + Initialize the chunking model with any necessary parameters. + """ + self.documents: Dict[int, BaseDocument] = kwargs.pop("documents") + self.parameters = self.validate_and_transform(kwargs) + self.chunks = self.chunk_documents() + self.id = None + + @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. + """ + 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 RAGWorkflowError( + "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..dd04002a3 --- /dev/null +++ b/DashAI/back/models/RAG/chunking_models/character_chunk_model.py @@ -0,0 +1,89 @@ +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, +) + + +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. + """ + 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..dc6c367c9 --- /dev/null +++ b/DashAI/back/models/RAG/chunking_models/chunking_model_factory.py @@ -0,0 +1,177 @@ +"""Factory for chunking models with full lifecycle encapsulation. + +Single-call interface: create() handles DB-record resolution, +model instantiation, document chunking, and chunk persistence. +""" + +from dataclasses import dataclass +from typing import Any, Dict + +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.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 + + +@dataclass(frozen=True) +class ChunkingFactoryResult: + """Result of chunking model creation via ChunkingModelFactory.""" + + db_record_id: int + model: BaseChunkingModel + chunks: Dict[int, Dict[int, Chunk]] + + +class ChunkingModelFactory: + """Creates a chunking model, chunks documents, and persists everything. + + Encapsulates the full lifecycle: DB-record resolution, model + instantiation, document splitting, and chunk persistence. + The caller calls ``create()`` once — no post-call required. + """ + + def __init__( + self, + db: Session, + registry: ComponentRegistry, + documents: Dict[int, BaseDocument], + chunk_set_id: int, + ): + self._db = db + self._registry = registry + self._documents = documents + self._chunk_set_id = chunk_set_id + + def create( + self, + component_name: str, + params: Dict[str, Any], + ) -> ChunkingFactoryResult: + """Full lifecycle: resolve DB record → instantiate → chunk → persist. + + Returns: + ChunkingFactoryResult with record_id, model, and populated chunks. + """ + sorted_params: Dict[str, Any] = dict(sorted(params.items())) + + existing_record = self._resolve_db_record(component_name, sorted_params) + if existing_record is not None: + params["documents"] = self._documents + model_class = self._registry[component_name]["class"] + model = model_class(**params) + model.set_id(existing_record.id) + self._update_chunk_ids(model) + return ChunkingFactoryResult( + db_record_id=existing_record.id, + model=model, + chunks=model.get_chunks(), + ) + + model_class = self._registry[component_name]["class"] + params["documents"] = self._documents + model = model_class(**params) + record_id = self._save_db_record(model) + + self._persist_chunks(model) + + return ChunkingFactoryResult( + db_record_id=record_id, + model=model, + chunks=model.get_chunks(), + ) + + def _resolve_db_record( + self, class_name: str, sorted_params: Dict[str, Any] + ) -> ChunkingModelDBModel | None: + return ( + self._db.query(ChunkingModelDBModel) + .filter_by(class_name=class_name, parameters=sorted_params) + .first() + ) + + def _save_db_record(self, model: BaseChunkingModel) -> int: + 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) + model.set_id(record.id) + return record.id + + def _fetch_chunks_from_db(self) -> Dict[int, Dict[int, Chunk]]: + chunks: Dict[int, Dict[int, Chunk]] = {} + db_chunks = ( + self._db.query(ChunkDBModel) + .filter_by(chunk_set_id=self._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 _create_chunks_in_db(self, chunks: Dict[int, Dict[int, Chunk]]) -> None: + 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=self._chunk_set_id, + text=chunk.text, + ) + ) + self._db.commit() + self._db.flush() + + def _update_chunk_ids(self, model: BaseChunkingModel) -> None: + 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=self._chunk_set_id, + ) + .first() + ) + if db_chunk is not None: + model.chunks[document_id][idx].id = db_chunk.id + + def _persist_chunks(self, model: BaseChunkingModel) -> None: + existing_chunks = self._fetch_chunks_from_db() + chunks_to_create: Dict[int, Dict[int, Chunk]] = {} + for document_id, document_chunks in model.get_chunks().items(): + if document_id not in existing_chunks: + chunks_to_create[document_id] = document_chunks + else: + for idx, chunk in document_chunks.items(): + if idx not in existing_chunks[document_id]: + if document_id not in chunks_to_create: + chunks_to_create[document_id] = {} + chunks_to_create[document_id][idx] = chunk + + self._create_chunks_in_db(chunks_to_create) + + self._update_chunk_ids(model) 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..40f1571f8 --- /dev/null +++ b/DashAI/back/models/RAG/chunking_models/recursive_character_chunk_model.py @@ -0,0 +1,129 @@ +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, +) + + +class RecursiveCharacterChunkModelSchema(BaseSchema): + 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): + 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): + 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): + 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]: + 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]: + 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]: + 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]: + 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..07b13c642 --- /dev/null +++ b/DashAI/back/models/RAG/chunking_models/token_chunk_model.py @@ -0,0 +1,88 @@ +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, +) + + +class TokenChunkModelSchema(BaseSchema): + 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): + SCHEMA = TokenChunkModelSchema + + def __init__(self, **kwargs): + self.parameters = kwargs + self.chunk_size = self.parameters["chunk_size"] + self.chunk_overlap = self.parameters["chunk_overlap"] + self.tokenizer_name = self.parameters["tokenizer_name"] + self._tokenizer: Optional["AutoTokenizer"] = None + super().__init__(**kwargs) + + @property + def tokenizer(self): + 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]: + 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/document_loader.py b/DashAI/back/models/RAG/document_loader.py new file mode 100644 index 000000000..9ce48ea52 --- /dev/null +++ b/DashAI/back/models/RAG/document_loader.py @@ -0,0 +1,68 @@ +"""Loads Document rows from the database and hydrates in-memory +BaseDocument instances. Keeps the pipeline constructor free of raw +SQL queries. +""" + +from typing import Dict, List + +from sqlalchemy.orm import Session + +from DashAI.back.dependencies.database.models import Document as DBDocument +from DashAI.back.models.RAG.documents import BaseDocument, PDFDocument, TxtDocument + +_DOCUMENT_CLASSES: Dict[str, type[BaseDocument]] = { + "txt": TxtDocument, + "pdf": PDFDocument, + "md": TxtDocument, + "rst": TxtDocument, + "tex": TxtDocument, + "csv": TxtDocument, +} + + +class DocumentLoader: + """Resolves a list of document IDs to in-memory BaseDocument objects.""" + + def __init__(self, db: Session): + self.db = db + + def load(self, document_ids: List[int]) -> Dict[int, BaseDocument]: + """Load and hydrate documents. + + Parameters + ---------- + document_ids : list[int] + IDs present in the ``document`` table. + + Returns + ------- + dict[int, BaseDocument] + Mapping from document ID to the hydrated document object. + + Raises + ------ + ValueError + If any document ID is not found in the database. + """ + documents: Dict[int, BaseDocument] = {} + for doc_id in document_ids: + db_doc: DBDocument | None = ( + self.db.query(DBDocument).filter(DBDocument.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[db_doc.file_type] + except KeyError: + raise ValueError( + f"Unsupported file type '{db_doc.file_type}'. " + f"Supported types: txt, pdf, md, rst, tex, csv." + ) + documents[doc_id] = doc_class( + id=db_doc.id, + file_name=db_doc.file_name, + file_path=db_doc.file_path, + created=db_doc.created, + optional_metadata=db_doc.optional_metadata, + ) + return documents diff --git a/DashAI/back/models/RAG/documents/__init__.py b/DashAI/back/models/RAG/documents/__init__.py new file mode 100644 index 000000000..d88c9bd6b --- /dev/null +++ b/DashAI/back/models/RAG/documents/__init__.py @@ -0,0 +1,4 @@ +from .base_document import BaseDocument +from .pdf_document import PDFDocument +from .txt_document import TxtDocument +from .chunk import Chunk \ No newline at end of file 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..424146d89 --- /dev/null +++ b/DashAI/back/models/RAG/documents/base_document.py @@ -0,0 +1,111 @@ +from abc import ABC +from typing import Any, Dict, Optional + + +class BaseDocument(ABC): + """ + Base class for documents. + """ + + SUPPORTED_FILE_TYPES = ["pdf", "txt"] + + 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 the document. + Args (from database): + id (int): The unique identifier of the document. + file_name (str): The name of the file. + file_path (str): The path to the file. + file_hash (str): A hash of the file content. + created (Optional[str]): The creation date of the document. + optional_metadata (Optional[Dict[str, Any]]): 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 if created else None + self.optional_metadata = optional_metadata if optional_metadata else None + + def get_text(self) -> str: + """ + Get the text content of the document. + + Returns: + str: The text content of the document. + """ + raise NotImplementedError("This method should be implemented by subclasses.") + + 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()) + + def get_metadata(self) -> Dict[str, Any]: + """ + Get the metadata of the document. + + Returns: + Dict[str, Any]: The metadata of the document. + """ + raise NotImplementedError("This method should be implemented by subclasses.") + + 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[str]: + """ + Get the filetype of the document. + + Returns: + Optional[str]: The filetype of the document, or None if not applicable. + """ + return self.file_path.split(".")[-1].lower() if self.file_path else None + + def __repr__(self): + return f"BaseDocument(id={self.id}, filename='{self.get_file_name()}', content='{self.get_text()[:50]}...', 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..4bab49839 --- /dev/null +++ b/DashAI/back/models/RAG/documents/chunk.py @@ -0,0 +1,18 @@ +from DashAI.back.models.RAG.utils import hash_function + + +class Chunk: + """A class representing a chunk of a document.""" + + def __init__( + self, + id: int, + document_id: str, + document_position: int, + text: str, + ): + self.id = 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/pdf_document.py b/DashAI/back/models/RAG/documents/pdf_document.py new file mode 100644 index 000000000..ed62d31d7 --- /dev/null +++ b/DashAI/back/models/RAG/documents/pdf_document.py @@ -0,0 +1,82 @@ +from typing import Any, Dict, Optional + +from DashAI.back.models.RAG.documents.base_document import BaseDocument +from DashAI.back.models.RAG.utils import hash_function + + +def _clean_textract_output(text: str) -> str: + """Clean textract output by removing control characters and normalizing whitespace.""" + import re + + text = re.sub(r"[\x00-\x1f\x7f]", " ", text) + text = re.sub(r"\s+", " ", text) + return text.strip() + + +class PDFDocument(BaseDocument): + """ + Class representing a PDF document. + """ + + def __init__( + self, + id: int, + file_name: str, + file_path: str, + created: Optional[str] = None, + optional_metadata: Optional[Dict[str, Any]] = None, + **kwargs, + ): + """ + Initialize the document. + Args (from database): + id (int): The unique identifier of the document. + file_name (str): The name of the file. + file_path (str): The path to the file. + created (Optional[str]): The creation date of the document. + optional_metadata (Optional[Dict[str, Any]]): Additional metadata for the document. + """ + self.PARSER = kwargs.get("parser", "textract") + file_hash = hash_function(file_path) + 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: + 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 ValueError( + f"Error extracting text from PDF file {self.file_path}: {str(e)}" + ) + else: + raise ValueError(f"Unsupported parser: {self.PARSER}") + + def get_metadata(self): + 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..0b94be276 --- /dev/null +++ b/DashAI/back/models/RAG/documents/txt_document.py @@ -0,0 +1,58 @@ +from typing import Any, Dict, Optional + +from DashAI.back.models.RAG.documents.base_document import BaseDocument +from DashAI.back.models.RAG.utils import hash_function + + +class TxtDocument(BaseDocument): + """ + Class representing a .txt document. + """ + + def __init__( + self, + id: int, + file_name: str, + file_path: str, + created: Optional[str] = None, + optional_metadata: Optional[Dict[str, Any]] = None, + ): + """ + Initialize the document. + Args (from database): + id (int): The unique identifier of the document. + file_name (str): The name of the file. + file_path (str): The path to the file. + file_hash (str): A hash of the file content. + created (Optional[str]): The creation date of the document. + optional_metadata (Optional[Dict[str, Any]]): Additional metadata for the document. + """ + file_hash = hash_function(file_path) + 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: + """ + Get the text content of the document. + + Returns: + str: The text content of the document. + """ + with open(self.file_path, "r", encoding="utf-8") as file: + text = file.read() + return text.strip() + + def get_metadata(self) -> Dict[str, Any]: + """ + Get the metadata of the document. + + Returns: + Dict[str, Any]: The metadata of the document. + """ + 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..f0f81ee2a --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/__init__.py @@ -0,0 +1,35 @@ +"""Embedding model exports.""" + +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..e15b168ba --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/__init__.py @@ -0,0 +1,10 @@ +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 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..0aee9915a --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/_bert_embedding.py @@ -0,0 +1,70 @@ +from typing import Dict + +import torch + +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + OverfloatHandler, +) + +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(OverfloatHandler): + def __init__( + self, + model_name: str, + device: str, + model_max_length: int, + overflow_strategy: str, + pooling_strategy: str, + ): + 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): + from transformers import AutoModel, AutoTokenizer + 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): + 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..ac98a39b0 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/_e5_embedding.py @@ -0,0 +1,42 @@ +from typing import List + +import numpy as np +import torch.nn.functional as F + +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + OverfloatHandler, +) + +QUERY_PREFIX = "query: " +PASSAGE_PREFIX = "passage: " + + +class _E5Embedding(OverfloatHandler): + def __init__( + self, + model_name: str, + device: str, + model_max_length: int, + overflow_strategy: str, + ): + 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): + 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 F.normalize(pooled, p=2, dim=1) + + def batch_encode(self, texts: List[str]) -> np.ndarray: + prefixed = [f"{PASSAGE_PREFIX}{t}" for t in texts] + return self._batch_encode_impl(prefixed) + + def encode(self, text: str) -> np.ndarray: + 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..bb59ff694 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/_gemma_embedding.py @@ -0,0 +1,55 @@ +from typing import List + +import numpy as np + +from DashAI.back.models.RAG.embeddings.dense.huggingface_embedding import ( + HuggingFaceEmbedding, +) + +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): + def __init__( + self, + model_name: str, + device: str, + model_max_length: int, + overflow_strategy: str, + task_type: str, + ): + 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): + raise NotImplementedError("Gemma uses SentenceTransformer API, _pool is unused.") + + def load(self): + from sentence_transformers import SentenceTransformer + + self.model = SentenceTransformer(self.model_name, device=self.device) + + def batch_encode(self, texts: List[str]) -> np.ndarray: + 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: + 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..4fe3651a5 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/_instructor_embedding.py @@ -0,0 +1,35 @@ +from typing import List + +import numpy as np + +from DashAI.back.models.RAG.embeddings.dense.huggingface_embedding import ( + HuggingFaceEmbedding, +) + + +class _InstructorEmbedding(HuggingFaceEmbedding): + def __init__( + self, + model_name: str, + device: str, + instruction: str, + ): + super().__init__(model_name=model_name, device=device) + self.instruction = instruction + self.params["instruction"] = instruction + + def _pool(self, model_output, attention_mask): + raise NotImplementedError("INSTRUCTOR uses custom encoding API, _pool is unused.") + + def load(self): + from InstructorEmbedding import INSTRUCTOR + self.model = INSTRUCTOR(self.model_name) + self.model._text_length = self.model._input_length + + def batch_encode(self, texts: List[str]) -> np.ndarray: + pairs = [[self.instruction, text] for text in texts] + return self.model.encode(pairs, show_progress_bar=False) + + def encode(self, text: str) -> np.ndarray: + 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..2173c6483 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/_overflow_handler.py @@ -0,0 +1,62 @@ +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 OverfloatHandler(HuggingFaceEmbedding): + def __init__( + self, + model_name: str, + device: str, + model_max_length: int, + overflow_strategy: str, + ): + 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): + raise NotImplementedError + + def _batch_encode_impl(self, texts: List[str]) -> np.ndarray: + 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..3d10497a4 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/_sentence_transformer_embedding.py @@ -0,0 +1,56 @@ +import torch +import torch.nn.functional as F + +from DashAI.back.models.RAG.embeddings.dense._overflow_handler import ( + OverfloatHandler, +) + +TRUNCATE = "truncate" +AGGREGATE = "aggregate" + + +class _SentenceTransformerEmbedding(OverfloatHandler): + def __init__( + self, + model_name: str, + device: str, + model_max_length: int, + overflow_strategy: str, + normalize: bool, + pooling: str = "mean", + ): + 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): + 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 = F.normalize(pooled, p=2, dim=1) + return pooled + + def _last_token_pool(self, model_output, attention_mask): + 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 = F.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..12a4c1b27 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/bert_embedding.py @@ -0,0 +1,146 @@ +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): + 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): + 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): + 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): + self._embedding.load() + + def encode(self, text: str): + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + return self._embedding.batch_encode(texts) + + def save(self): + pass + + def train(self, **kwargs): + return 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..ec056e266 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/distilbert_embedding.py @@ -0,0 +1,132 @@ +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): + 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): + 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): + 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): + self._embedding.load() + + def encode(self, text: str): + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + return self._embedding.batch_encode(texts) + + def save(self): + pass + + def train(self, **kwargs): + return 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..21f5678af --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/e5_embedding.py @@ -0,0 +1,131 @@ +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): + 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): + 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): + 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): + self._embedding.load() + + def encode(self, text: str): + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + return self._embedding.batch_encode(texts) + + def save(self): + pass + + def train(self, **kwargs): + return 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..af4db87c1 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/fasttext_embedding.py @@ -0,0 +1,74 @@ +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 + + +class FastTextEmbeddingSchema(BaseSchema): + 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): + """FastText embedding""" + + 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): + 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): + """Load the FastText model.""" + 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): + pass + + def train(self, **kwargs): + pass + + def encode(self, text: str) -> np.ndarray: + """Encode text into an embedding.""" + 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 into embeddings.""" + 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..bd1ec013e --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/gemma_embedding.py @@ -0,0 +1,109 @@ +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): + 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): + 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): + 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): + self._embedding.load() + + def encode(self, text: str): + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + return self._embedding.batch_encode(texts) + + def save(self): + pass + + def train(self, **kwargs): + return 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..e3e6d07e7 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/huggingface_embedding.py @@ -0,0 +1,56 @@ +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): + FLAGS: list[str] = ["abstract", "huggingface"] + + def __init__(self, model_name: str, device: str): + 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): + pass + + def train(self, **kwargs): + return + + @abstractmethod + def _pool(self, model_output, attention_mask): + raise NotImplementedError + + def load(self): + 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: + 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: + return self._batch_encode_impl(texts) + + def encode(self, text: str) -> np.ndarray: + 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..b33f598ee --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/instructor_embedding.py @@ -0,0 +1,95 @@ +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): + 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): + 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): + 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): + self._embedding.load() + + def encode(self, text: str): + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + return self._embedding.batch_encode(texts) + + def save(self): + pass + + def train(self, **kwargs): + return 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..a38a8d33a --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/labse_embedding.py @@ -0,0 +1,102 @@ +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): + 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): + 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): + 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): + self._embedding.load() + + def encode(self, text: str): + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + return self._embedding.batch_encode(texts) + + def save(self): + pass + + def train(self, **kwargs): + return 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..d08d3fb2a --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/openai_embedding.py @@ -0,0 +1,73 @@ +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): + 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): + 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): + self.params = self.validate_and_transform(kwargs) + self.model_name = self.params["model_name"] + self.api_key = self.params["api_key"] + self.client = OpenAI(api_key=self.api_key) + + def load(self): + pass + + def save(self): + pass + + def train(self, **kwargs): + return + + def encode(self, text: str) -> np.ndarray: + 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: + 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..d916ea55d --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/roberta_embedding.py @@ -0,0 +1,134 @@ +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): + 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): + 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): + 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): + self._embedding.load() + + def encode(self, text: str): + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + return self._embedding.batch_encode(texts) + + def save(self): + pass + + def train(self, **kwargs): + return 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..3e0cfa432 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense/sentence_transformer_embedding.py @@ -0,0 +1,266 @@ +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): + 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): + 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): + 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): + self._embedding.load() + + def encode(self, text: str): + return self._embedding.encode(text) + + def batch_encode(self, texts: List[str]): + return self._embedding.batch_encode(texts) + + def save(self): + pass + + def train(self, **kwargs): + return 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..aee70f370 --- /dev/null +++ b/DashAI/back/models/RAG/embeddings/dense_embedding.py @@ -0,0 +1,46 @@ +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 encoding (embedding) models. + This class should be inherited by any specific encoding (embedding) model implementation. + """ + + embedding_dim: int + + def __init__(self, **kwargs: Any): + """ + Initialize the encoding (embedding) model with the given documents. + """ + raise NotImplementedError("Subclasses must implement the __init__ method.") + + def encode(self, text: str) -> List[float] | np.ndarray: + """ + Method to generate encodings (embeddings) for the given text. + This method should be implemented by subclasses. + + :param text: The input text to encode (embed). + :return: A list representing the encodings (embeddings) of the input text. + """ + raise NotImplementedError("Subclasses must implement this method.") + + def batch_encode(self, texts: List[str]) -> List[List[float]] | np.ndarray: + """ + Method to generate encodings (embeddings) for a batch of texts. + This method should be implemented by subclasses. + + :param texts: A list of input texts to encode (embed). + :return: A list of lists representing the encodings (embeddings) of the input texts. + """ + raise NotImplementedError("Subclasses must implement this method.") + + @abstractmethod + def train(self, **kwargs): + """Train the embedding model.""" + 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..e69de29bb diff --git a/DashAI/back/models/RAG/exceptions.py b/DashAI/back/models/RAG/exceptions.py new file mode 100644 index 000000000..abdbb49ee --- /dev/null +++ b/DashAI/back/models/RAG/exceptions.py @@ -0,0 +1,2 @@ +class RAGWorkflowError(Exception): + """Custom exception for RAG workflow errors.""" diff --git a/DashAI/back/models/RAG/llm_factory.py b/DashAI/back/models/RAG/llm_factory.py new file mode 100644 index 000000000..169adb8bd --- /dev/null +++ b/DashAI/back/models/RAG/llm_factory.py @@ -0,0 +1,62 @@ +"""Factory for text-to-text generation models (LLMs). + +Resolves a component name + parameters into an instantiated model +via the component registry, with lookup-or-create DB semantics. +""" + +from dataclasses import dataclass +from typing import Any, Dict + +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.text_to_text_generation_model import ( + TextToTextGenerationTaskModel, +) + + +@dataclass(frozen=True) +class LLMFactoryResult: + """Result of LLM instantiation via LLMFactory.""" + + db_record_id: int + model: TextToTextGenerationTaskModel + + +class LLMFactory: + """Creates LLM instances with lookup-or-create DB semantics.""" + + def __init__(self, db: Session, component_registry: ComponentRegistry): + self._db = db + self._registry = component_registry + + def create(self, component_name: str, params: Dict[str, Any]) -> LLMFactoryResult: + sorted_params = dict(sorted(params.items())) + existing = ( + self._db.query(GenerationDBModel) + .filter_by(class_name=component_name, parameters=sorted_params) + .first() + ) + if existing is not None: + model_class = self._registry[existing.class_name]["class"] + return LLMFactoryResult( + db_record_id=existing.id, + model=model_class(**existing.parameters), + ) + + db_record = GenerationDBModel( + class_name=component_name, + parameters=sorted_params, + ) + self._db.add(db_record) + self._db.commit() + self._db.refresh(db_record) + + model_class = self._registry[component_name]["class"] + return LLMFactoryResult( + db_record_id=db_record.id, + model=model_class(**params), + ) diff --git a/DashAI/back/models/RAG/pipeline_repository.py b/DashAI/back/models/RAG/pipeline_repository.py new file mode 100644 index 000000000..7a2bf26ed --- /dev/null +++ b/DashAI/back/models/RAG/pipeline_repository.py @@ -0,0 +1,62 @@ +"""Persistence layer for the RAGPipeline DB record. + +Keeps DB access out of the pipeline's __init__ so the pipeline can +focus on orchestration. +""" + +from typing import Optional + +from sqlalchemy.orm import Session + +from DashAI.back.dependencies.database.models import RAGPipeline as PipelineDBModel + + +class PipelineRepository: + """Manages the rag_pipeline DB record for a given session.""" + + def __init__(self, db: Session): + self._db = db + + def ensure_db_record(self, session_id: int) -> PipelineDBModel: + """Return the existing pipeline record or create a placeholder. + + The placeholder has filler values for the FK columns; call + ``update_db_record`` once the real component IDs are known. + """ + existing: Optional[PipelineDBModel] = ( + self._db.query(PipelineDBModel).filter_by(session_id=session_id).first() + ) + if existing is not None: + return existing + + record = PipelineDBModel( + session_id=session_id, + name="", + description=None, + parameters=None, + chunking_model_id=0, + prompt_id=0, + generation_model_id=0, + ) + self._db.add(record) + self._db.commit() + self._db.refresh(record) + return record + + def update_db_record( + self, + session_id: int, + chunking_model_id: int, + prompt_id: int, + generation_model_id: int, + ) -> None: + """Patch the pipeline record with real component IDs.""" + pipeline_record: Optional[PipelineDBModel] = ( + self._db.query(PipelineDBModel).filter_by(session_id=session_id).first() + ) + if pipeline_record is None: + raise ValueError(f"No pipeline record for session {session_id}") + pipeline_record.chunking_model_id = chunking_model_id + pipeline_record.prompt_id = prompt_id + pipeline_record.generation_model_id = generation_model_id + self._db.commit() 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..895c4e03d --- /dev/null +++ b/DashAI/back/models/RAG/prompts/augmentation/augmentation_prompt.py @@ -0,0 +1,29 @@ +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): + self.template = kwargs.pop("template") + + 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..4b1d8dfdb --- /dev/null +++ b/DashAI/back/models/RAG/prompts/augmentation/custom_augmentation_prompt.py @@ -0,0 +1,51 @@ +from typing import Any + +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): + self.template = kwargs.pop("template") + + 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. + """ + 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..91b53d8a8 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/augmentation/default_augmentation_prompt.py @@ -0,0 +1,153 @@ +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.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): + 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): + """ + 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. + """ + + 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") + self.template = kwargs.pop("template") or TEMPLATES.get(self.language, "") + + 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. + """ + 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/__init__.py b/DashAI/back/models/RAG/prompts/generation/__init__.py new file mode 100644 index 000000000..2867bb7a7 --- /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..a540f5e84 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/custom_rag_generation_prompt.py @@ -0,0 +1,44 @@ +from typing import Any + +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): + 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. + history (List[Tuple[str, str]]): The chat history to be included in the context. + chunks (List[str]): The document chunks to be included in the context. + **kwargs: Additional keyword arguments for formatting. + + Returns: + str: The formatted prompt. + """ + 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..80f5e4872 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/default_QA_rag_generation_prompt.py @@ -0,0 +1,129 @@ +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.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): + 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): + self.language = kwargs.pop("language") + self.template = kwargs.pop("template") or TEMPLATES.get(self.language, "") + + def format(self, input: str, chunks: str, **kwargs: Any) -> str: + 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..37e5b4c94 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/default_rag_generation_prompt.py @@ -0,0 +1,125 @@ +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.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): + 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): + self.language = kwargs.pop("language") + self.template = kwargs.pop("template") or TEMPLATES.get(self.language, "") + + def format(self, input: str, chunks: str, **kwargs: Any) -> str: + 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..3e1e0aaa2 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/generation/rag_generation_prompt.py @@ -0,0 +1,21 @@ +from typing import Any, Tuple + +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 = [] + + @staticmethod + def format( + input: str, + chunks: str, + history: Tuple[Tuple[str, str], ...] | None = None, + **kwargs: Any, + ) -> str: + 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..f01308d14 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/prompt.py @@ -0,0 +1,86 @@ +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): + 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" + id: int + name: str + REQUIRED_EXTRA_KWARGS = [] + + def load(self, **kwargs: Any) -> None: + pass + + def save(self) -> None: + pass + + def train(self, **kwargs: Any) -> None: + pass + + @classmethod + def get_metadata(cls) -> Dict[str, Any]: + 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. + """ + assert hasattr(cls, "required_placeholders"), ( + "Subclasses 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. + """ + assert hasattr(self, "optional_placeholders"), ( + "Subclasses 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) + + @staticmethod + def format(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..5d292ed78 --- /dev/null +++ b/DashAI/back/models/RAG/prompts/prompt_factory.py @@ -0,0 +1,60 @@ +"""Factory for prompt templates with lookup-or-create semantics. + +Resolves a component name + parameters into a Prompt instance, +creating a RAGPrompt DB record when no matching record exists. +""" + +from dataclasses import dataclass +from typing import Any, Dict + +from sqlalchemy.orm import Session + +from DashAI.back.dependencies.database.models import RAGPrompt as PromptDBModel +from DashAI.back.dependencies.registry.component_registry import ComponentRegistry +from DashAI.back.models.RAG.prompts import Prompt + + +@dataclass(frozen=True) +class PromptFactoryResult: + """Result of prompt instantiation via PromptFactory.""" + + db_record_id: int + model: Prompt + + +class PromptFactory: + """Creates Prompt instances with automatic DB record reuse.""" + + def __init__(self, db: Session, component_registry: ComponentRegistry): + self._db = db + self._registry = component_registry + + def create( + self, component_name: str, params: Dict[str, Any] + ) -> PromptFactoryResult: + sorted_params: Dict[str, Any] = dict(sorted(params.items())) + existing: PromptDBModel | None = ( + self._db.query(PromptDBModel) + .filter_by(class_name=component_name, parameters=sorted_params) + .first() + ) + if existing is not None: + prompt_class = self._registry[existing.class_name]["class"] + return PromptFactoryResult( + db_record_id=existing.id, + model=prompt_class(**existing.parameters), + ) + + db_record = PromptDBModel( + class_name=component_name, + parameters=sorted_params, + ) + self._db.add(db_record) + self._db.commit() + self._db.refresh(db_record) + + prompt_class = self._registry[component_name]["class"] + return PromptFactoryResult( + db_record_id=db_record.id, + model=prompt_class(**params), + ) 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..46470ec49 --- /dev/null +++ b/DashAI/back/models/RAG/rag_models_factory.py @@ -0,0 +1,105 @@ +"""Abstract Factory (GoF) for RAG component creation. + +RAGModelsFactory provides a unified interface for creating the four +RAG component types: prompts, chunking models, retrievers, and LLMs. + +Each ``create_*`` method delegates to a specialised sub-factory that +encapsulates the component's full lifecycle (DB-record resolution, +instantiation, and persistence). +""" + +from typing import Any, Dict + +from sqlalchemy.orm import Session + +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.retriever_factory import ( + RetrieverFactory, + RetrieverFactoryResult, +) + + +class RAGModelsFactory: + """Abstract Factory for the family of RAG components. + + Provides four ``create_*`` methods, each delegating to a + specialised sub-factory. Shared dependencies (db, registry, + env_rag_path) are injected once via the constructor. + """ + + def __init__( + self, + db: Session, + registry: ComponentRegistry, + env_rag_path: str, + ): + self._db = db + self._registry = registry + self._env_rag_path = env_rag_path + + def create_prompt( + self, + component_name: str, + params: Dict[str, Any], + ) -> PromptFactoryResult: + """Create a prompt with lookup-or-create semantics.""" + factory = PromptFactory(self._db, self._registry) + return factory.create(component_name, params) + + def create_chunking_model( + self, + documents: Dict[int, BaseDocument], + chunk_set_id: int, + component_name: str, + params: Dict[str, Any], + ) -> ChunkingFactoryResult: + """Create a chunking model, chunk documents, and persist chunks.""" + factory = ChunkingModelFactory( + self._db, + self._registry, + documents, + chunk_set_id, + ) + return factory.create(component_name, params) + + def create_retriever( + self, + pipeline_id: int, + chunks: Dict[int, Dict[int, Chunk]], + chunk_set_id: int, + component_name: str, + params: Dict[str, Any], + ) -> RetrieverFactoryResult: + """Create a retriever with full persistence lifecycle.""" + factory = RetrieverFactory( + self._db, + pipeline_id, + self._registry, + self._env_rag_path, + chunks, + chunk_set_id, + ) + return factory.create(component_name, params) + + def create_llm( + self, + component_name: str, + params: Dict[str, Any], + ) -> LLMFactoryResult: + """Create an LLM with lookup-or-create DB semantics.""" + factory = LLMFactory(self._db, self._registry) + return factory.create(component_name, 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..929f5addb --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/__init__.py @@ -0,0 +1,32 @@ +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, + ExtraKwargsMissingError, + MissingParameterError, + RetrieverError, + UnitRetrieverChildError, +) +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..347e2a803 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/composite/__init__.py @@ -0,0 +1,12 @@ +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..31fc4dd39 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/composite/composite_retriever.py @@ -0,0 +1,47 @@ +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): + 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 isinstance(children_data[0], RetrieverModel): + raise TypeError( + "'children' must contain RetrieverModel instances, " + f"got {type(children_data[0]).__name__}" + ) + self._children: List[RetrieverModel] = children_data + super().__init__(**kwargs) + + def add(self, child: RetrieverModel) -> None: + self._children.append(child) + + def remove(self, child: RetrieverModel) -> None: + self._children.remove(child) + + def get_children(self) -> List[RetrieverModel]: + return list(self._children) + + @abstractmethod + def retrieve(self, query, **kwargs) -> List[Chunk]: + raise NotImplementedError + + def score_chunks(self, chunk_ids: List[int], query: str) -> list: + return [] 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..750ddd6d6 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/composite/mmr_reranker_retriever.py @@ -0,0 +1,130 @@ +from typing import List + +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, +) + + +class MMRRerankerRetrieverSchema(BaseSchema): + 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): + 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): + 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 self._top_k + + def retrieve(self, query, **kwargs) -> List[Chunk]: + 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 = [c.id for c in candidates] + scored = 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 = [c.id for c in ordered] + try: + vectors = child.get_chunk_vectors(ordered_ids) + except ValueError: + return ordered[: self._top_k] + + pairwise = cosine_similarity(vectors) + selected = self._mmr_select(relevance, pairwise) + return [ordered[i] for i in selected] + + def _mmr_select( + self, + relevance: np.ndarray, + pairwise: np.ndarray, + ) -> List[int]: + 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..e8952b82f --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/composite/parallel_retriever.py @@ -0,0 +1,112 @@ +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): + 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 from each retriever. " + f"'{MergeStrategy.INTERLEAVE.value}': merges preserving internal order." + ), + es=( + f"'{MergeStrategy.ROUND_ROBIN.value}': alterna resultados de cada recuperador. " + f"'{MergeStrategy.INTERLEAVE.value}': fusiona preservando el orden interno." + ), + ), + ) # type: ignore + + +class ParallelRetriever(CompositeRetriever): + 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): + super().__init__(**kwargs) + self.merge_strategy = MergeStrategy(self.params.pop("merge_strategy")) + + @property + def retrieval_top_k(self) -> int: + return sum(c.retrieval_top_k for c in self._children) + + def retrieve(self, query, **kwargs) -> List[Chunk]: + 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) + + @staticmethod + def _chunk_key(chunk: Chunk) -> Tuple[str, int]: + return (chunk.document_id, chunk.document_position) + + def _merge_round_robin( + self, child_results_list: List[List[Chunk]], total_k: int + ) -> List[Chunk]: + results = [] + seen_keys: set[Tuple[str, 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]: + results = [] + seen_keys: set[Tuple[str, 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..6a84b9b9b --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/composite/sequential_retriever.py @@ -0,0 +1,89 @@ +from typing import List + +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.retrievers.composite.composite_retriever import ( + CompositeRetriever, +) +from DashAI.back.models.RAG.retrievers.exceptions import ( + CompositeValidationError, + UnitRetrieverChildError, +) +from DashAI.back.models.RAG.retrievers.unit_retriever import UnitRetriever + + +class SequentialRetrieverSchema(BaseSchema): + 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): + 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): + super().__init__(**kwargs) + self._validate() + + def _validate(self) -> None: + children = getattr(self, "_children", []) + if len(children) < 2: + return + + children_k = [c.retrieval_top_k for c in children] + + for _i, child in enumerate(children): + if not isinstance(child, UnitRetriever): + raise UnitRetrieverChildError( + child_class=child.__class__.__name__, + strategy="cascade", + ) + + 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]: + first = self._children[0] + results = first.retrieve(query, **kwargs) + + 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 + + @property + def retrieval_top_k(self) -> int: + 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..14a5985ab --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/dense/__init__.py @@ -0,0 +1,7 @@ +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..03aa92f16 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/dense/_hf_language_utils.py @@ -0,0 +1,91 @@ +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]: + 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]]: + 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]: + 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..daca8d7a9 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/dense/_hf_metadata_utils.py @@ -0,0 +1,22 @@ +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]: + 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..23e5bb347 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/dense/dense_embedding_retriever.py @@ -0,0 +1,75 @@ +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): + 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 DenseEmbedding component.""" + + 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): + 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..63745336d --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/dense/dense_retriever.py @@ -0,0 +1,179 @@ +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.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): + 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): + 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): + 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): + if not isinstance(embedding_model, DenseEmbedding): + raise TypeError( + f"Expected DenseEmbedding instance, " + f"got {type(embedding_model).__name__}" + ) + self.embedding_model = embedding_model + encoding_class_name = embedding_model.__class__.__name__ + encoding_params = dict(sorted(embedding_model.params.items())) + + self.params["encoding_model"] = { + "class_name": encoding_class_name, + "parameters": encoding_params, + } + + self.compute_missing_embeddings() + self.init_similarity_matrix() + + def compute_missing_embeddings(self): + 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()] + 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): + 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 self._top_k + + def retrieve(self, query: str, top_k: int | None = None) -> List[Chunk]: + 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) + # NOTE: pairwise_distances computes O(n×dim) per query with no + # FAISS/HNSW indexing. Low priority since current use cases have + # small document sets. + 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]]: + 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())) + scored.sort(key=lambda x: x[1]) + return scored + + def get_chunk_vectors(self, chunk_ids: List[int]) -> np.ndarray: + 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..e46526f18 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/dense/huggingface_dense_retriever.py @@ -0,0 +1,33 @@ +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): + 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): + super().__init__(**kwargs) + + @abstractmethod + def _create_embedding(self) -> HuggingFaceEmbedding: + raise NotImplementedError + + def init_model(self) -> None: + 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..33f721321 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/enums.py @@ -0,0 +1,6 @@ +from enum import Enum + + +class MergeStrategy(str, Enum): + 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..b9cba1448 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/exceptions.py @@ -0,0 +1,42 @@ +class RetrieverError(Exception): + """Base exception for all retriever-related errors.""" + + +class MissingParameterError(RetrieverError): + """A required parameter was not provided to the retriever.""" + + def __init__(self, param_name: str, retriever_name: str): + super().__init__( + f"Missing required parameter '{param_name}' " + f"in retriever '{retriever_name}'." + ) + self.param_name = param_name + self.retriever_name = retriever_name + + +class ExtraKwargsMissingError(RetrieverError): + """One or more required infrastructure kwargs are missing.""" + + def __init__(self, missing_keys: set, retriever_name: str): + super().__init__( + f"Missing required extra kwargs {missing_keys} " + f"in retriever '{retriever_name}'." + ) + self.missing_keys = missing_keys + self.retriever_name = retriever_name + + +class CompositeValidationError(RetrieverError): + """A composite retriever has an invalid configuration.""" + + +class UnitRetrieverChildError(RetrieverError): + """A composite retriever received a child that is not a UnitRetriever.""" + + def __init__(self, child_class: str, strategy: str): + super().__init__( + f"Cascade strategy '{strategy}' requires UnitRetriever children, " + f"but child '{child_class}' is a composite." + ) + self.child_class = child_class + self.strategy = strategy diff --git a/DashAI/back/models/RAG/retrievers/persistence.py b/DashAI/back/models/RAG/retrievers/persistence.py new file mode 100644 index 000000000..171e38a7d --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/persistence.py @@ -0,0 +1,25 @@ +from dataclasses import dataclass + + +@dataclass +class SparsePersistence: + """Reference to a sparse retriever's on-disk dump. + + model_dir is the 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. + + matrix_dirs maps document_id to the absolute path of the directory + containing embeddings.npy for that document. + embedding_model_id identifies the embedding model that produced the vectors. + """ + + matrix_dirs: dict[int, str] + embedding_model_id: int 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..a30a60eba --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/retriever_factory.py @@ -0,0 +1,335 @@ +"""Factory for retriever instances with full lifecycle encapsulation. + +Single-call ``create()`` interface encapsulates DB-record resolution, +persistence planning, constructor injection, and composite child recursion. +""" + +import os +from dataclasses import dataclass +from typing import Any, Dict +from uuid import uuid4 + +from sqlalchemy.orm import Session + +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.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.retriever_repository import ( + RetrieverRepository, +) +from DashAI.back.models.RAG.retrievers.sparse.sparse_retriever import SparseRetriever + + +@dataclass(frozen=True) +class RetrieverFactoryResult: + db_record_id: int + model: RetrieverModel + + +_ENCODING_MODEL_KEY = "encoding_model" + + +class RetrieverFactory: + """Creates retriever instances with full persistence lifecycle. + + All DB access delegated to ``RetrieverRepository``. + """ + + def __init__( + self, + db: Session, + pipeline_id: int, + registry: ComponentRegistry, + env_rag_path: str, + chunks: Dict[int, Dict[int, Chunk]], + chunk_set_id: int, + ): + self._db = db + self._pipeline_id = pipeline_id + self._registry = registry + self._env_rag_path = env_rag_path + self._chunks = chunks + self._chunk_set_id = chunk_set_id + self._repo = RetrieverRepository(db) + + def create( + self, + component_name: str, + params: Dict[str, Any], + ) -> RetrieverFactoryResult: + params = normalize_payload(params) + model_class = self._registry[component_name]["class"] + + if issubclass(model_class, CompositeRetriever): + return self._create_composite(model_class, params) + + return self._create_unit(model_class, params) + + def _create_composite(self, model_class, params): + existing = self._load_composite_from_db(model_class, params) + if existing is not None: + return RetrieverFactoryResult( + db_record_id=existing.get_id(), + model=existing, + ) + + children_configs = params.pop("children") + children_instances = [self._create_child_retriever(c) for c in children_configs] + params["children"] = children_instances + model = model_class(**params) + bridge_id = self._save_composite(model) + model.set_id(bridge_id) + return RetrieverFactoryResult(db_record_id=bridge_id, model=model) + + def _create_child_retriever(self, child_config): + return self.create( + component_name=child_config["component"], + params=child_config["params"], + ).model + + def _load_composite_from_db(self, model_class, params): + bridge_record = self._repo.find_composite( + self._pipeline_id, + model_class.__name__, + ) + if bridge_record is None: + return None + child_links = self._repo.find_composite_children(bridge_record.id) + if not child_links: + return None + + children = [] + for link in child_links: + child_bridge = self._repo.get_bridge_by_id(link.child_id) + if child_bridge is None: + return None + child_params = self._load_child_params( + child_bridge.id, child_bridge.class_name + ) + if child_params is None: + return None + children.append( + self.create( + component_name=child_bridge.class_name, + params=child_params, + ).model, + ) + + model = model_class( + children=children, + **{k: v for k, v in params.items() if k != "children"}, + ) + model.set_id(bridge_record.id) + return model + + def _load_child_params( + self, + bridge_id: int, + class_name: str, + ) -> Dict[str, Any] | None: + sparse = self._repo.find_sparse_by_bridge_id(bridge_id) + if sparse is not None and sparse.parameters is not None: + return dict(sorted(sparse.parameters.items())) + dense = self._repo.find_dense_by_bridge_id(bridge_id) + if dense is not None and dense.parameters is not None: + return dict(sorted(dense.parameters.items())) + return None + + def _save_composite(self, model): + child_ids = [] + for child in model.get_children(): + cid = child.get_id() + if cid is None: + raise ValueError(f"Child {child.__class__.__name__} has no bridge ID.") + child_ids.append(cid) + return self._repo.save_composite( + model.__class__.__name__, + self._pipeline_id, + child_ids, + ).id + + def _create_unit(self, 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 RetrieverFactoryResult( + db_record_id=loaded.get_id(), + model=loaded, + ) + + validated = model_class.SCHEMA.model_validate(params) + resolved = fill_objects(validated, self._registry) + model = model_class(**resolved) + + if issubclass(model_class, DenseRetriever): + model.inject_infra( + self._env_rag_path, + self._chunks, + DensePersistence(matrix_dirs={}, embedding_model_id=0), + ) + model.init_model() + self._finalize_dense(model, sorted_params) + elif issubclass(model_class, SparseRetriever): + model.inject_infra( + self._env_rag_path, + self._chunks, + SparsePersistence(model_dir=None), + ) + model.init_model() + else: + raise ValueError(f"Unsupported unit retriever: {model_class.__name__}") + + bridge_id = self._save_unit(model, sorted_params) + model.set_id(bridge_id) + return RetrieverFactoryResult(db_record_id=bridge_id, model=model) + + def _finalize_dense(self, model, sorted_params): + """After init_model() for a new dense model: persist embedding + model, build persistence with the proper ID, compute matrices.""" + enc = model.params.get(_ENCODING_MODEL_KEY) + if not enc: + return + emb_record = self._repo.find_or_create_embedding_model( + enc["class_name"], + dict(sorted(enc["parameters"].items())), + ) + model._persistence = self._dense_persistence(emb_record.id) + model.compute_missing_embeddings() + model.init_similarity_matrix() + + def _dense_persistence(self, embedding_model_id: int) -> DensePersistence: + """Build a DensePersistence with matrix dirs for every known doc.""" + matrix_dirs = {} + for doc_id in self._chunks: + folder = ( + f"doc_id-{doc_id}__" + f"chunk_set_id-{self._chunk_set_id}__" + f"embedding_model_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, + ) + + def _load_unit_from_db(self, model_class, class_name, sorted_params): + 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) + raise ValueError(f"Unsupported unit retriever: {model_class.__name__}") + + def _load_dense(self, class_name, sorted_params): + record = self._repo.find_dense(class_name, sorted_params, self._chunk_set_id) + if record is None: + return None + persistence = self._dense_persistence(record.embedding_model_id) + model_class = self._registry[class_name]["class"] + + validated = model_class.SCHEMA.model_validate(dict(sorted_params)) + resolved = fill_objects(validated, self._registry) + model = model_class(**resolved) + model.inject_infra(self._env_rag_path, self._chunks, persistence) + model.init_model() + bridge = self._repo.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, sorted_params): + record = self._repo.find_sparse(class_name, sorted_params, self._chunk_set_id) + if record is None: + return None + persistence = SparsePersistence(model_dir=record.storage_folder) + model_class = self._registry[class_name]["class"] + + validated = model_class.SCHEMA.model_validate(dict(sorted_params)) + resolved = fill_objects(validated, self._registry) + model = model_class(**resolved) + model.inject_infra(self._env_rag_path, self._chunks, persistence) + model.init_model() + bridge = self._repo.find_bridge_for_sub_table(record, class_name) + if bridge is not None: + model.set_id(bridge.id) + return model + + def _save_unit(self, model, sorted_params): + 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, sorted_params): + import numpy as np + + for doc_id, mdir in model._persistence.matrix_dirs.items(): + path = os.path.join(mdir, "embeddings.npy") + if not os.path.exists(path): + continue + if self._repo.find_embedding_matrix( + doc_id, + self._chunk_set_id, + model._persistence.embedding_model_id, + ): + continue + self._repo.save_embedding_matrix( + doc_id, + self._chunk_set_id, + model._persistence.embedding_model_id, + mdir, + list(np.load(path).shape), + ) + + bridge = self._repo.create_bridge( + model.__class__.__name__, + self._pipeline_id, + ) + self._repo.save_dense( + model.__class__.__name__, + sorted_params, + bridge.id, + self._chunk_set_id, + model._persistence.embedding_model_id, + ) + return bridge.id + + def _save_sparse(self, model, sorted_params): + folder_id = uuid4().hex + storage_folder = os.path.join( + self._env_rag_path, + "sparse_retrievers", + f"sparse_retriever_id-{folder_id}", + ) + model._persistence.model_dir = storage_folder + model.save() + bridge = self._repo.create_bridge( + model.__class__.__name__, + self._pipeline_id, + ) + self._repo.save_sparse( + model.__class__.__name__, + sorted_params, + storage_folder, + bridge.id, + self._chunk_set_id, + ) + return bridge.id 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..f4e1bb935 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/retriever_model.py @@ -0,0 +1,128 @@ +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 + chunks: Dict[int, Dict[int, Chunk]] + chunking_model_id: int | None + params: Dict[str, Any] + + def __init__(self, **kwargs): + self._db_id: int | None = None + self.params = kwargs + + def get_id(self) -> int | None: + return self._db_id + + def set_id(self, id: int) -> None: + 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: + 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} must be an integer." + ) + if not isinstance(chunk, Chunk): + raise ValueError( + f"Chunk {chunk_id} in document ID {doc_id} must be an instance of Chunk." + ) + if chunk.document_id != doc_id: + raise ValueError( + f"Chunk {chunk_id} document_id {chunk.document_id} != doc ID {doc_id}." + ) + + def inject_infra( + self, + env_rag_path: str | os.PathLike, + 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. + """ + self.env_rag_path = env_rag_path + self.chunks = chunks + self._persistence = persistence + self._validate_chunks_dict() + + 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]: + raise NotImplementedError + + @abstractmethod + def score_chunks(self, chunk_ids: List[int], query: str) -> List[Tuple[int, float]]: + raise NotImplementedError + + @property + def retrieval_top_k(self) -> int: + raise NotImplementedError + + # ── Child management (Composite pattern) ──────────────────────── + + def add(self, child: "RetrieverModel") -> None: + raise NotImplementedError + + def remove(self, child: "RetrieverModel") -> None: + raise NotImplementedError + + def get_children(self) -> List["RetrieverModel"]: + raise NotImplementedError + + def save(self, filename: str = "") -> None: + pass + + def load(self, filename: str = "") -> None: + pass + + def train(self, **kwargs): + return diff --git a/DashAI/back/models/RAG/retrievers/retriever_repository.py b/DashAI/back/models/RAG/retrievers/retriever_repository.py new file mode 100644 index 000000000..e3d3c9c0a --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/retriever_repository.py @@ -0,0 +1,337 @@ +"""Repository for retriever-related database operations. + +Separates persistence logic (SQL queries, record creation) from +orchestration logic (component construction, DI, composites). +""" + +from typing import Dict, List, Optional + +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 RetrieverRepository: + """All database operations for retriever models. + + This class owns every SQL query, INSERT, and DB-model construction + related to retrievers. Factories use it as a collaborator; retriever + instances never touch it. + """ + + def __init__(self, db: Session): + self.db = db + + def create_bridge(self, class_name: str, pipeline_id: int) -> RetrieverDBModel: + """Insert a new canonical identity record in rag_retriever. + + Every retriever (unit or composite) gets one bridge record. + Returns the persisted model with its auto-generated id. + """ + bridge_record = RetrieverDBModel( + class_name=class_name, + pipeline_id=pipeline_id, + ) + self.db.add(bridge_record) + self.db.commit() + self.db.refresh(bridge_record) + return bridge_record + + 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())) + return ( + self.db.query(SparseRetrieverDBModel) + .filter_by( + class_name=class_name, + parameters=parameters, + chunk_set_id=chunk_set_id, + ) + .first() + ) + + def save_sparse( + self, + class_name: str, + parameters: Dict[str, object], + storage_folder: str, + bridge_id: int, + chunk_set_id: int, + ) -> SparseRetrieverDBModel: + """Persist a new sparse retriever record and link it to its bridge. + + Returns the persisted model. + """ + 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, + ) + self.db.add(record) + self.db.commit() + self.db.refresh(record) + return record + + 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.""" + parameters = dict(sorted(parameters.items())) + return ( + self.db.query(DenseRetrieverDBModel) + .filter_by( + class_name=class_name, + parameters=parameters, + chunk_set_id=chunk_set_id, + ) + .first() + ) + + def save_dense( + self, + class_name: str, + parameters: Dict[str, object], + bridge_id: int, + chunk_set_id: int, + embedding_model_id: int, + ) -> DenseRetrieverDBModel: + """Persist a new dense retriever record and link it to its bridge.""" + 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, + ) + self.db.add(record) + self.db.commit() + self.db.refresh(record) + return record + + def find_composite( + self, pipeline_id: int, class_name: str + ) -> Optional[RetrieverDBModel]: + """Find the composite bridge record for a given pipeline.""" + return ( + self.db.query(RetrieverDBModel) + .filter_by( + pipeline_id=pipeline_id, + class_name=class_name, + ) + .first() + ) + + def find_composite_children(self, parent_id: int) -> List[RAGRetrieverChild]: + """Return child links ordered by child_order.""" + return ( + self.db.query(RAGRetrieverChild) + .filter_by(parent_id=parent_id) + .order_by(RAGRetrieverChild.child_order) + .all() + ) + + def save_composite( + self, + class_name: str, + pipeline_id: int, + child_bridge_ids: List[int], + ) -> 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). + + Returns + ------- + RetrieverDBModel + The persisted bridge record. + """ + bridge_record = self.create_bridge(class_name, pipeline_id) + 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, + ) + ) + self.db.commit() + return bridge_record + + def find_or_create_embedding_model( + self, + class_name: str, + parameters: Dict[str, object], + ) -> EmbeddingDBModel: + """Return an existing EmbeddingDBModel record or create one. + + This is idempotent: the same (class_name, parameters) pair + always resolves to the same record. + """ + parameters = dict(sorted(parameters.items())) + existing = ( + self.db.query(EmbeddingDBModel) + .filter_by( + class_name=class_name, + parameters=parameters, + ) + .first() + ) + if existing is not None: + return existing + record = EmbeddingDBModel( + class_name=class_name, + parameters=parameters, + ) + self.db.add(record) + self.db.commit() + self.db.refresh(record) + 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. + """ + result: Dict[int, EmbeddingMatrixDBModel] = {} + for doc_id in doc_ids: + matrix = ( + self.db.query(EmbeddingMatrixDBModel) + .filter_by( + document_id=doc_id, + chunk_set_id=chunk_set_id, + embedding_model_id=embedding_model_id, + ) + .first() + ) + if matrix is not None: + result[doc_id] = matrix + return result + + def save_embedding_matrix( + self, + document_id: int, + chunk_set_id: int, + embedding_model_id: int, + storage_folder: str, + matrix_shape: List[int], + ) -> EmbeddingMatrixDBModel: + """Persist a new embedding matrix record. + + Returns the persisted model. + """ + 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, + ) + self.db.add(record) + self.db.commit() + self.db.refresh(record) + return record + + 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.""" + return ( + self.db.query(EmbeddingMatrixDBModel) + .filter_by( + document_id=document_id, + chunk_set_id=chunk_set_id, + embedding_model_id=embedding_model_id, + ) + .first() + ) + + def find_sparse_by_bridge_id(self, bridge_id: int) -> Optional[SparseRetrieverDBModel]: + """Look up a sparse retriever sub-table record by its bridge_id.""" + return ( + self.db.query(SparseRetrieverDBModel) + .filter_by(bridge_id=bridge_id) + .first() + ) + + def find_dense_by_bridge_id(self, bridge_id: int) -> Optional[DenseRetrieverDBModel]: + """Look up a dense retriever sub-table record by its bridge_id.""" + return ( + self.db.query(DenseRetrieverDBModel) + .filter_by(bridge_id=bridge_id) + .first() + ) + + 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.""" + bridge = self.db.query(RetrieverDBModel).get(sub_db_model.bridge_id) + if bridge is not None and bridge.class_name == class_name: + return bridge + return None + + def get_bridge_by_id(self, bridge_id: int) -> Optional[RetrieverDBModel]: + """Fetch a bridge record by its primary key.""" + return self.db.query(RetrieverDBModel).get(bridge_id) 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..07e0f9285 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/sparse/__init__.py @@ -0,0 +1,3 @@ +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..4d5da12f1 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/sparse/bm25_retriever.py @@ -0,0 +1,326 @@ +import logging +import os +import pickle +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, + 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): + 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): + 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): + from sklearn.feature_extraction.text import CountVectorizer + + 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): + pass + + def save(self): + pass + + def train(self): + pass + +class BM25RetrieverSchema(BaseSchema): + 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): + 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): + 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: + if not self.load(): + self._fit() + + def load(self) -> bool: + 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: + self._vectorizer = pickle.load(f) + with open(os.path.join(model_dir, "bm25_tf_matrix.pkl"), "rb") as f: + self._tf_matrix = pickle.load(f) + with open(os.path.join(model_dir, "bm25_matrix.pkl"), "rb") as f: + self._bm25_matrix = pickle.load(f) + with open(os.path.join(model_dir, "bm25_row_to_chunk.pkl"), "rb") as f: + self.matrix_row_to_chunk_map = pickle.load(f) + return True + except Exception as e: + log.info("Error loading BM25 state: %s", e) + return False + + def save(self) -> 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): + 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) + + df = np.array((tf > 0).sum(axis=0)).flatten() + idf = np.log((n_docs - df + 0.5) / (df + 0.5) + 1.0) + idf += self.delta + idf[idf < 0] = 0 + + from scipy.sparse import lil_matrix + + 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 self._top_k + + def retrieve(self, query: str, top_k: int | None = None) -> List[Chunk]: + 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]]: + 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())) + scored.sort(key=lambda x: x[1]) + return scored + + def get_chunk_vectors(self, chunk_ids: List[int]) -> np.ndarray: + 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..2761374eb --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/sparse/sparse_retriever.py @@ -0,0 +1,14 @@ +from DashAI.back.core.utils import MultilingualString +from DashAI.back.models.RAG.retrievers.unit_retriever import UnitRetriever + + +class SparseRetriever(UnitRetriever): + 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..8af4749c9 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/sparse/tfidf_retriever.py @@ -0,0 +1,334 @@ +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): + 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): + 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: + 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: + pass + + def load(self, filename: str = "") -> None: + pass + + def train(self, **kwargs): + return + + +class TFIDFRetrieverSchema(BaseSchema): + 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): + 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): + 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: + if not self.load(): + self._fit() + + def load(self) -> bool: + 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: + self._vectorizer = pickle.load(f) + with open(os.path.join(model_dir, "tf_idf_matrix.pkl"), "rb") as f: + self._tf_idf_matrix = pickle.load(f) + with open( + os.path.join(model_dir, "matrix_row_to_chunk_map.pkl"), "rb" + ) as f: + self.matrix_row_to_chunk_map = pickle.load(f) + return True + except Exception as e: + log.info("Error loading TFIDF state: %s", e) + return False + + def save(self) -> 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): + 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 self._top_k + + + + def retrieve(self, query: str, top_k: int | None = None) -> List[Chunk]: + 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]]: + 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())) + scored.sort(key=lambda x: x[1]) + return scored + + def get_chunk_vectors(self, chunk_ids: List[int]) -> np.ndarray: + 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..c55bf8c84 --- /dev/null +++ b/DashAI/back/models/RAG/retrievers/unit_retriever.py @@ -0,0 +1,53 @@ +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.""" + + TYPE: Final[str] = "RetrieverModel" + FLAGS: list[str] = ["abstract"] + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def inject_infra( + self, + env_rag_path: str, + chunks: Dict[int, Dict[int, Chunk]], + persistence: Any, + ) -> None: + 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.""" + 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: + raise TypeError( + f"{self.__class__.__name__} is a unit retriever and cannot contain children." + ) + + def remove(self, child: RetrieverModel) -> None: + raise TypeError( + f"{self.__class__.__name__} is a unit retriever and cannot contain children." + ) + + def get_children(self) -> List[RetrieverModel]: + return [] diff --git a/DashAI/back/models/RAG/utils.py b/DashAI/back/models/RAG/utils.py new file mode 100644 index 000000000..300162d3d --- /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 bytes. + + 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..73da59ea1 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 \ No newline at end of file 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..425560dca 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,44 @@ 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 +496,10 @@ 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.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..09a2747a8 --- /dev/null +++ b/DashAI/back/models/hugging_face/phi_4_mini_instruct_model.py @@ -0,0 +1,244 @@ +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..4224f238e 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,30 @@ 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..19f1f3f6f --- /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 \ No newline at end of file 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..c92eaa357 --- /dev/null +++ b/DashAI/back/models/remote_models/deepseek_text_to_text_generation_model.py @@ -0,0 +1,143 @@ +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..ac6b9f544 --- /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/tasks/RAG_task.py b/DashAI/back/tasks/RAG_task.py new file mode 100644 index 000000000..bbb281f19 --- /dev/null +++ b/DashAI/back/tasks/RAG_task.py @@ -0,0 +1,164 @@ +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.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="RAG Task", + es="Tarea 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."}] + """ + 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. + + Parameters + ---------- + output : RAGGenerationOutput + The typed output from RAGPipeline.generate(). + """ + 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..999f64b17 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 \ No newline at end of file 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..2f6be13a8 100644 --- a/DashAI/front/src/App.jsx +++ b/DashAI/front/src/App.jsx @@ -1,6 +1,6 @@ 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 +13,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 +67,32 @@ 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..2dac403c7 100644 --- a/DashAI/front/src/api/generativeTask.ts +++ b/DashAI/front/src/api/generativeTask.ts @@ -12,15 +12,22 @@ 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..a75e47a6e 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..45b945c42 --- /dev/null +++ b/DashAI/front/src/api/rag.ts @@ -0,0 +1,295 @@ +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"; + +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; +}; + +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; +}; + +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; +}; + +export const createRAGSession = async ( + sessionData: Omit, +): Promise => { + console.log("Creating RAG session with data:", sessionData); + const params = sessionData.parameters as { + documents: number[]; + chunking_model: { + component: string; + params: Record; + }; + retriever_model: { + component: string; + params: Record; + }; + generation_model: { + component: string; + params: Record; + }; + prompt: { + component: string; + params: Record; + }; + }; + + const transformedSession: Omit = + { + name: sessionData.name, + description: sessionData.description, + task_name: "RAGTask", + model_name: "RAGPipeline", + display_name: "", + parameters: { + documents: params.documents, + chunking_model: { + component: params.chunking_model.component, + params: params.chunking_model.params, + }, + retriever_model: { + component: params.retriever_model.component, + params: params.retriever_model.params, + }, + generation_model: { + component: params.generation_model.component, + params: params.generation_model.params, + }, + prompt: params.prompt, + }, + }; + + 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; +}; + +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; +}; + +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}`); + } +}; + +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; +}; + +export const getRetrievalParadigm = async (): Promise => { + const response = await getChildComponents("RetrieverModel", true); + if (!response) { + throw new Error(`Failed to fetch retrieval options`); + } + return response; +}; + +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; +}; + +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; +}; + +export const getChunkingComponents = async (): Promise => { + const response = await getChildComponents("BaseChunkingModel", false); + if (!response) { + throw new Error(`Failed to fetch chunking components`); + } + return response; +}; + +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; +}; + +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; +}; + +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}`); + } +}; + +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"; + +/** + * Check whether a prompt's class_name corresponds to a generation prompt + * (i.e. NOT an augmentation prompt). + */ +export function isGenerationPromptClass(className: string): boolean { + return !className.includes(AUGMENTATION_PROMPT_CLASS_PREFIX); +} + +export const getDefaultPrompts = async (): Promise => { + return getChildComponents("RAGGenerationPrompt", false); +}; + +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; +}; + +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..d2c6a43f5 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..789127a51 --- /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..1ecd4471e --- /dev/null +++ b/DashAI/front/src/components/documents/DocumentTable.jsx @@ -0,0 +1,205 @@ +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; \ No newline at end of file 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/CreateSessionContext.jsx b/DashAI/front/src/components/generative/CreateSessionContext.jsx index 43a639088..b895bc2fc 100644 --- a/DashAI/front/src/components/generative/CreateSessionContext.jsx +++ b/DashAI/front/src/components/generative/CreateSessionContext.jsx @@ -50,7 +50,8 @@ export function CreateSessionProvider({ children }) { setLoadingModels(true); Promise.all( - tasks.map((task) => + tasks + .map((task) => getRelatedComponents(task.name).then((components) => components.map((c) => ({ ...c, 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..8f3300da6 --- /dev/null +++ b/DashAI/front/src/components/generative/DocumentReferencesModal.jsx @@ -0,0 +1,165 @@ +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; \ No newline at end of file diff --git a/DashAI/front/src/components/generative/GenerativeChat.jsx b/DashAI/front/src/components/generative/GenerativeChat.jsx index 5f12cd2d3..27b9846fc 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,78 @@ 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 +302,8 @@ export default function GenerativeChat() { whiteSpace: "pre-wrap", }} > - {change.parameter}: {change.oldValue}{" "} - {change.newValue}{" "} + {change.parameter}: {formatHistoryValue(change.oldValue)}{" "} + {formatHistoryValue(change.newValue)}{" "} )), }; @@ -226,6 +316,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 +428,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" ? ( @@ -319,19 +444,53 @@ export default function GenerativeChat() { {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({ /> ); })} -