mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-27 18:39:00 +08:00
Merge branch 'release-1.10.0' into fix/mask-global-variable-output-exposure
This commit is contained in:
@ -8480,5 +8480,5 @@
|
||||
}
|
||||
]
|
||||
},
|
||||
"generated_at": "2026-04-23T21:12:19Z"
|
||||
"generated_at": "2026-04-24T03:31:30Z"
|
||||
}
|
||||
|
||||
25
RELEASE.md
25
RELEASE.md
@ -39,9 +39,28 @@ This step also usually lasts about a week.
|
||||
After QA and bugfixing are complete for both OSS and Desktop:
|
||||
|
||||
* Final releases are cut from their respective RC branches.
|
||||
* Release timing is coordinated with Langflow’s DevRel team.
|
||||
* Release timing is coordinated with Langflow's DevRel team.
|
||||
* For at least 24 hours after release, Discord, GitHub, and other support channels should be monitored for critical bug reports.
|
||||
|
||||
### 4. Release Artifacts
|
||||
|
||||
The release workflow automatically publishes the following artifacts:
|
||||
|
||||
* **PyPI Packages:**
|
||||
* `langflow` - Main package with all integrations
|
||||
* `langflow-base` - Core framework without integrations
|
||||
* `lfx` - Lightweight executor CLI
|
||||
* `langflow-sdk` - SDK for programmatic access (when updated)
|
||||
|
||||
* **Docker Images:**
|
||||
* `langflowai/langflow` - Full Langflow image
|
||||
* `langflowai/langflow-backend` - Backend-only image (published independently)
|
||||
* `langflowai/langflow-frontend` - Frontend-only image (published independently)
|
||||
* `langflowai/langflow-ep` - Enterprise edition image (published independently)
|
||||
* `langflowai/langflow-base` - Base image without integrations
|
||||
|
||||
**Note:** Backend, frontend, and enterprise images are published separately from the main image and will be built even if the main version already exists on Docker Hub.
|
||||
|
||||
## Branch Model
|
||||
|
||||
| Branch | Purpose | Merge Policy |
|
||||
@ -101,6 +120,10 @@ git merge --ff-only release-X.Y.Z # Fast-forward main to include RC changes
|
||||
|
||||
* Follows [Semantic Versioning](https://semver.org): `MAJOR.MINOR.PATCH`.
|
||||
* RC tags use `-rc.N`, e.g. `v1.8.0-rc.1`.
|
||||
* **All tags MUST start with `v` prefix** (e.g., `v1.9.1`, not `1.9.1`).
|
||||
* The release workflow validates this format and rejects tags without the `v` prefix.
|
||||
* Duplicate tags (e.g., both `1.8.3` and `v1.8.3`) cause GitHub's release notes generation to use the wrong base comparison, resulting in incomplete changelogs.
|
||||
* The workflow automatically checks for and prevents duplicate tags.
|
||||
|
||||
## Roles
|
||||
|
||||
|
||||
@ -131,6 +131,8 @@ langflow:
|
||||
enabled: true
|
||||
driver:
|
||||
value: "postgresql"
|
||||
host:
|
||||
value: "postgresql-svc.langflow.svc.cluster.local"
|
||||
port:
|
||||
value: "5432"
|
||||
user:
|
||||
|
||||
@ -131,6 +131,8 @@ langflow:
|
||||
enabled: true
|
||||
driver:
|
||||
value: "postgresql"
|
||||
host:
|
||||
value: "postgresql-svc.langflow.svc.cluster.local"
|
||||
port:
|
||||
value: "5432"
|
||||
user:
|
||||
|
||||
@ -131,6 +131,8 @@ langflow:
|
||||
enabled: true
|
||||
driver:
|
||||
value: "postgresql"
|
||||
host:
|
||||
value: "postgresql-svc.langflow.svc.cluster.local"
|
||||
port:
|
||||
value: "5432"
|
||||
user:
|
||||
|
||||
@ -0,0 +1,235 @@
|
||||
"""add_memory_base_schema
|
||||
|
||||
Consolidates all Memory Base schema changes into a single migration:
|
||||
- job.dedupe_key (nullable String) + ix_job_dedupe_key
|
||||
- message.run_id (nullable UUID) + ix_message_run_id
|
||||
- message.is_output (bool, default false)
|
||||
- memory_base table + ix_memory_base_flow_id + ix_memory_base_user_id
|
||||
- memory_base_session table + three indexes
|
||||
- message_ingestion_record table + three indexes
|
||||
- memory_base_workflow_run table + two indexes
|
||||
|
||||
Phase: EXPAND
|
||||
|
||||
Revision ID: mb00a1b2c3d4
|
||||
Revises: d306e5c17c41
|
||||
Create Date: 2026-04-14 00:00:00.000000
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from langflow.utils import migration
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "mb00a1b2c3d4" # pragma: allowlist secret
|
||||
down_revision: str | None = "d306e5c17c41" # pragma: allowlist secret
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# job.dedupe_key #
|
||||
# ------------------------------------------------------------------ #
|
||||
inspector = sa.inspect(conn)
|
||||
existing_job_indexes = {idx["name"] for idx in inspector.get_indexes("job")}
|
||||
with op.batch_alter_table("job", schema=None) as batch_op:
|
||||
if not migration.column_exists("job", "dedupe_key", conn):
|
||||
batch_op.add_column(sa.Column("dedupe_key", sa.String(), nullable=True))
|
||||
if "ix_job_dedupe_key" not in existing_job_indexes:
|
||||
batch_op.create_index(batch_op.f("ix_job_dedupe_key"), ["dedupe_key"], unique=False)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# message.run_id + message.is_output #
|
||||
# ------------------------------------------------------------------ #
|
||||
with op.batch_alter_table("message", schema=None) as batch_op:
|
||||
if not migration.column_exists("message", "run_id", conn):
|
||||
batch_op.add_column(sa.Column("run_id", sa.Uuid(), nullable=True))
|
||||
if not migration.column_exists("message", "is_output", conn):
|
||||
batch_op.add_column(sa.Column("is_output", sa.Boolean(), nullable=False, server_default=sa.text("false")))
|
||||
|
||||
existing_message_indexes = {idx["name"] for idx in sa.inspect(conn).get_indexes("message")}
|
||||
if "ix_message_run_id" not in existing_message_indexes:
|
||||
op.create_index("ix_message_run_id", "message", ["run_id"])
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# memory_base #
|
||||
# ------------------------------------------------------------------ #
|
||||
if not migration.table_exists("memory_base", conn):
|
||||
op.create_table(
|
||||
"memory_base",
|
||||
sa.Column("id", sa.Uuid(), nullable=False),
|
||||
sa.Column("name", sa.String(), nullable=False),
|
||||
sa.Column("flow_id", sa.Uuid(), nullable=False),
|
||||
sa.Column("user_id", sa.Uuid(), nullable=False),
|
||||
sa.Column("threshold", sa.Integer(), nullable=False, server_default=sa.text("50")),
|
||||
sa.Column("auto_capture", sa.Boolean(), nullable=False, server_default=sa.text("true")),
|
||||
sa.Column("embedding_model", sa.String(), nullable=False, server_default=sa.text("''")),
|
||||
sa.Column("preprocessing", sa.Boolean(), nullable=False, server_default=sa.text("false")),
|
||||
sa.Column("preproc_model", sa.String(), nullable=True),
|
||||
sa.Column("preproc_instructions", sa.String(), nullable=True),
|
||||
sa.Column("kb_name", sa.String(), nullable=False),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
sa.UniqueConstraint("user_id", "name", name="uq_memory_base_user_name"),
|
||||
sa.Index("ix_memory_base_flow_id", "flow_id"),
|
||||
sa.Index("ix_memory_base_user_id", "user_id"),
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# memory_base_session #
|
||||
# ------------------------------------------------------------------ #
|
||||
if not migration.table_exists("memory_base_session", conn):
|
||||
op.create_table(
|
||||
"memory_base_session",
|
||||
sa.Column("id", sa.Uuid(), nullable=False),
|
||||
sa.Column(
|
||||
"memory_base_id",
|
||||
sa.Uuid(),
|
||||
sa.ForeignKey("memory_base.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("session_id", sa.String(), nullable=False),
|
||||
sa.Column("cursor_id", sa.Uuid(), nullable=True),
|
||||
sa.Column("total_processed", sa.Integer(), nullable=False, server_default=sa.text("0")),
|
||||
sa.Column("last_sync_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
sa.UniqueConstraint("memory_base_id", "session_id", name="uq_memory_base_session"),
|
||||
)
|
||||
op.create_index("ix_memory_base_session_memory_base_id", "memory_base_session", ["memory_base_id"])
|
||||
op.create_index("ix_memory_base_session_session_id", "memory_base_session", ["session_id"])
|
||||
op.create_index(
|
||||
"ix_memory_base_session_lookup",
|
||||
"memory_base_session",
|
||||
["memory_base_id", "session_id"],
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# message_ingestion_record #
|
||||
# ------------------------------------------------------------------ #
|
||||
if not migration.table_exists("message_ingestion_record", conn):
|
||||
op.create_table(
|
||||
"message_ingestion_record",
|
||||
sa.Column("id", sa.Uuid(), nullable=False),
|
||||
sa.Column(
|
||||
"message_id",
|
||||
sa.Uuid(),
|
||||
sa.ForeignKey("message.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"memory_base_id",
|
||||
sa.Uuid(),
|
||||
sa.ForeignKey("memory_base.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"job_id",
|
||||
sa.Uuid(),
|
||||
sa.ForeignKey("job.job_id", ondelete="SET NULL"),
|
||||
nullable=True,
|
||||
),
|
||||
sa.Column("session_id", sa.String(), nullable=False),
|
||||
sa.Column("ingested_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
sa.UniqueConstraint(
|
||||
"message_id",
|
||||
"session_id",
|
||||
"memory_base_id",
|
||||
name="uq_mir_message_session_mb",
|
||||
),
|
||||
)
|
||||
op.create_index("ix_mir_message_id", "message_ingestion_record", ["message_id"])
|
||||
op.create_index("ix_mir_job_id", "message_ingestion_record", ["job_id"])
|
||||
op.create_index(
|
||||
"ix_mir_memory_base_session",
|
||||
"message_ingestion_record",
|
||||
["memory_base_id", "session_id"],
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# memory_base_workflow_run #
|
||||
# ------------------------------------------------------------------ #
|
||||
if not migration.table_exists("memory_base_workflow_run", conn):
|
||||
op.create_table(
|
||||
"memory_base_workflow_run",
|
||||
sa.Column("id", sa.Uuid(), nullable=False),
|
||||
sa.Column(
|
||||
"memory_base_id",
|
||||
sa.Uuid(),
|
||||
sa.ForeignKey("memory_base.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("session_id", sa.String(), nullable=False),
|
||||
sa.Column(
|
||||
"workflow_job_id",
|
||||
sa.Uuid(),
|
||||
sa.ForeignKey("job.job_id", ondelete="SET NULL"),
|
||||
nullable=True,
|
||||
),
|
||||
sa.Column(
|
||||
"ingestion_job_id",
|
||||
sa.Uuid(),
|
||||
sa.ForeignKey("job.job_id", ondelete="SET NULL"),
|
||||
nullable=True,
|
||||
),
|
||||
sa.Column("recorded_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
sa.UniqueConstraint(
|
||||
"memory_base_id",
|
||||
"session_id",
|
||||
"workflow_job_id",
|
||||
name="uq_mbwr_mb_session_wf_job",
|
||||
),
|
||||
)
|
||||
op.create_index("ix_mbwr_mb_session", "memory_base_workflow_run", ["memory_base_id", "session_id"])
|
||||
op.create_index("ix_mbwr_ingestion_job_id", "memory_base_workflow_run", ["ingestion_job_id"])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
|
||||
# Children first (FK dependencies) ----------------------------------- #
|
||||
if migration.table_exists("memory_base_workflow_run", conn):
|
||||
op.drop_index("ix_mbwr_ingestion_job_id", table_name="memory_base_workflow_run")
|
||||
op.drop_index("ix_mbwr_mb_session", table_name="memory_base_workflow_run")
|
||||
op.drop_table("memory_base_workflow_run")
|
||||
|
||||
if migration.table_exists("message_ingestion_record", conn):
|
||||
op.drop_index("ix_mir_memory_base_session", table_name="message_ingestion_record")
|
||||
op.drop_index("ix_mir_job_id", table_name="message_ingestion_record")
|
||||
op.drop_index("ix_mir_message_id", table_name="message_ingestion_record")
|
||||
op.drop_table("message_ingestion_record")
|
||||
|
||||
if migration.table_exists("memory_base_session", conn):
|
||||
op.drop_index("ix_memory_base_session_lookup", table_name="memory_base_session")
|
||||
op.drop_index("ix_memory_base_session_session_id", table_name="memory_base_session")
|
||||
op.drop_index("ix_memory_base_session_memory_base_id", table_name="memory_base_session")
|
||||
op.drop_table("memory_base_session")
|
||||
|
||||
if migration.table_exists("memory_base", conn):
|
||||
op.drop_index("ix_memory_base_user_id", table_name="memory_base")
|
||||
op.drop_index("ix_memory_base_flow_id", table_name="memory_base")
|
||||
op.drop_table("memory_base")
|
||||
|
||||
# Message column/index ----------------------------------------------- #
|
||||
existing_message_indexes = {idx["name"] for idx in sa.inspect(conn).get_indexes("message")}
|
||||
if "ix_message_run_id" in existing_message_indexes:
|
||||
op.drop_index("ix_message_run_id", table_name="message")
|
||||
with op.batch_alter_table("message", schema=None) as batch_op:
|
||||
if migration.column_exists("message", "is_output", conn):
|
||||
batch_op.drop_column("is_output")
|
||||
if migration.column_exists("message", "run_id", conn):
|
||||
batch_op.drop_column("run_id")
|
||||
|
||||
# Job column/index --------------------------------------------------- #
|
||||
with op.batch_alter_table("job", schema=None) as batch_op:
|
||||
existing_job_indexes = {idx["name"] for idx in sa.inspect(conn).get_indexes("job")}
|
||||
if "ix_job_dedupe_key" in existing_job_indexes:
|
||||
batch_op.drop_index(batch_op.f("ix_job_dedupe_key"))
|
||||
if migration.column_exists("job", "dedupe_key", conn):
|
||||
batch_op.drop_column("dedupe_key")
|
||||
@ -29,7 +29,15 @@ from langflow.exceptions.component import ComponentBuildError
|
||||
from langflow.schema.message import ErrorMessage
|
||||
from langflow.schema.schema import OutputValue
|
||||
from langflow.services.database.models.flow.model import Flow
|
||||
from langflow.services.deps import get_chat_service, get_telemetry_service, session_scope
|
||||
from langflow.services.database.models.jobs.model import JobType
|
||||
from langflow.services.deps import (
|
||||
get_chat_service,
|
||||
get_job_service,
|
||||
get_memory_base_service,
|
||||
get_task_service,
|
||||
get_telemetry_service,
|
||||
session_scope,
|
||||
)
|
||||
from langflow.services.job_queue.service import JobQueueNotFoundError, JobQueueService
|
||||
from langflow.services.telemetry.schema import ComponentInputsPayload, ComponentPayload, PlaygroundPayload
|
||||
|
||||
@ -527,35 +535,81 @@ async def generate_flow_events(
|
||||
event_manager.on_error(data=error_message.data)
|
||||
raise
|
||||
|
||||
# Create a WORKFLOW job record so memory-base on_flow_output can track this run.
|
||||
# Best-effort: failures here must never break the build path.
|
||||
_build_job_svc = None
|
||||
_build_run_id: uuid.UUID | None = None
|
||||
try:
|
||||
_build_run_id = uuid.UUID(graph.run_id) if graph.run_id else None
|
||||
if _build_run_id is not None:
|
||||
_build_job_svc = get_job_service()
|
||||
await _build_job_svc.create_job(
|
||||
job_id=_build_run_id,
|
||||
flow_id=flow_id,
|
||||
user_id=current_user.id,
|
||||
job_type=JobType.WORKFLOW,
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
await logger.awarning(
|
||||
"Failed to create workflow job for /build — memory base tracking disabled for flow %s",
|
||||
flow_id,
|
||||
exc_info=True,
|
||||
)
|
||||
_build_job_svc = None
|
||||
|
||||
event_manager.on_vertices_sorted(data={"ids": ids, "to_run": vertices_to_run})
|
||||
|
||||
vertex_timedeltas: list[float] = []
|
||||
event_manager.on_build_start(data={})
|
||||
tasks = []
|
||||
for vertex_id in ids:
|
||||
task = asyncio.create_task(build_vertices(vertex_id, graph, event_manager, vertex_timedeltas))
|
||||
tasks.append(task)
|
||||
try:
|
||||
await asyncio.gather(*tasks)
|
||||
except asyncio.CancelledError:
|
||||
background_tasks.add_task(graph.end_all_traces_in_context())
|
||||
raise
|
||||
except Exception as e:
|
||||
await logger.aerror(f"Error building vertices: {e}")
|
||||
custom_component = graph.get_vertex(vertex_id).custom_component
|
||||
trace_name = getattr(custom_component, "trace_name", None)
|
||||
error_message = ErrorMessage(
|
||||
flow_id=flow_id,
|
||||
exception=e,
|
||||
session_id=graph.session_id,
|
||||
trace_name=trace_name,
|
||||
)
|
||||
event_manager.on_error(data=error_message.data)
|
||||
raise
|
||||
|
||||
async def _run_vertex_build() -> None:
|
||||
tasks = []
|
||||
for vertex_id in ids:
|
||||
task = asyncio.create_task(build_vertices(vertex_id, graph, event_manager, vertex_timedeltas))
|
||||
tasks.append(task)
|
||||
try:
|
||||
await asyncio.gather(*tasks)
|
||||
except asyncio.CancelledError:
|
||||
background_tasks.add_task(graph.end_all_traces_in_context())
|
||||
raise
|
||||
except Exception as e:
|
||||
await logger.aerror(f"Error building vertices: {e}")
|
||||
custom_component = graph.get_vertex(vertex_id).custom_component
|
||||
trace_name = getattr(custom_component, "trace_name", None)
|
||||
error_message = ErrorMessage(
|
||||
flow_id=flow_id,
|
||||
exception=e,
|
||||
session_id=graph.session_id,
|
||||
trace_name=trace_name,
|
||||
)
|
||||
event_manager.on_error(data=error_message.data)
|
||||
raise
|
||||
|
||||
if _build_job_svc and _build_run_id:
|
||||
await _build_job_svc.execute_with_status(_build_run_id, _run_vertex_build)
|
||||
else:
|
||||
await _run_vertex_build()
|
||||
|
||||
build_duration = sum(vertex_timedeltas)
|
||||
event_manager.on_end(data={"build_duration": build_duration})
|
||||
await graph.end_all_traces()
|
||||
|
||||
# Fire memory-base auto-capture hook — non-blocking background effect.
|
||||
# Must use fire_and_forget_task (not background_tasks.add_task) because
|
||||
# generate_flow_events runs as an asyncio task; by the time the flow
|
||||
# finishes, FastAPI has already drained the background_tasks queue and any
|
||||
# tasks added after that point are silently dropped.
|
||||
try:
|
||||
_run_id_uuid = uuid.UUID(graph.run_id) if graph.run_id else None # type-cast only; same run_id set on graph
|
||||
await get_task_service().fire_and_forget_task(
|
||||
get_memory_base_service().on_flow_output,
|
||||
flow_id=flow_id,
|
||||
session_id=graph.session_id or str(flow_id),
|
||||
job_id=_run_id_uuid,
|
||||
)
|
||||
except (RuntimeError, ValueError, OSError):
|
||||
await logger.awarning("Memory base hook scheduling failed for flow %s", flow_id, exc_info=True)
|
||||
|
||||
await event_manager.queue.put((None, None, time.time()))
|
||||
|
||||
|
||||
|
||||
@ -15,6 +15,7 @@ from langflow.api.v1 import (
|
||||
login_router,
|
||||
mcp_projects_router,
|
||||
mcp_router,
|
||||
memories_router,
|
||||
model_options_router,
|
||||
models_router,
|
||||
monitor_router,
|
||||
@ -68,6 +69,7 @@ router_v1.include_router(folders_router)
|
||||
router_v1.include_router(projects_router)
|
||||
router_v1.include_router(starter_projects_router)
|
||||
router_v1.include_router(knowledge_bases_router)
|
||||
router_v1.include_router(memories_router)
|
||||
router_v1.include_router(mcp_router)
|
||||
router_v1.include_router(voice_mode_router)
|
||||
router_v1.include_router(mcp_projects_router)
|
||||
|
||||
@ -445,7 +445,7 @@ class KBIngestionHelper:
|
||||
splitter_kwargs["separators"] = [resolved_separator]
|
||||
text_splitter = RecursiveCharacterTextSplitter(**splitter_kwargs)
|
||||
|
||||
embeddings = await KBIngestionHelper._build_embeddings(embedding_provider, embedding_model, current_user)
|
||||
embeddings = await KBIngestionHelper.build_embeddings(embedding_provider, embedding_model, current_user)
|
||||
|
||||
client = KBStorageHelper.get_fresh_chroma_client(kb_path)
|
||||
chroma = Chroma(
|
||||
@ -462,40 +462,30 @@ class KBIngestionHelper:
|
||||
continue
|
||||
|
||||
chunks = text_splitter.split_text(content)
|
||||
for i in range(0, len(chunks), INGESTION_BATCH_SIZE):
|
||||
if await KBIngestionHelper._is_job_cancelled(job_service, task_job_id):
|
||||
raise IngestionCancelledError
|
||||
docs = [
|
||||
Document(
|
||||
page_content=c,
|
||||
metadata={
|
||||
"source": source_name or file_name,
|
||||
"file_name": file_name,
|
||||
"chunk_index": i,
|
||||
"total_chunks": len(chunks),
|
||||
"ingested_at": datetime.now(timezone.utc).isoformat(),
|
||||
"job_id": job_id_str,
|
||||
},
|
||||
)
|
||||
for i, c in enumerate(chunks)
|
||||
]
|
||||
|
||||
batch = chunks[i : i + INGESTION_BATCH_SIZE]
|
||||
docs = [
|
||||
Document(
|
||||
page_content=c,
|
||||
metadata={
|
||||
"source": source_name or file_name,
|
||||
"file_name": file_name,
|
||||
"chunk_index": i + j,
|
||||
"total_chunks": len(chunks),
|
||||
"ingested_at": datetime.now(timezone.utc).isoformat(),
|
||||
"job_id": job_id_str,
|
||||
},
|
||||
)
|
||||
for j, c in enumerate(batch)
|
||||
]
|
||||
|
||||
for attempt in range(MAX_RETRY_ATTEMPTS):
|
||||
if await KBIngestionHelper._is_job_cancelled(job_service, task_job_id):
|
||||
raise IngestionCancelledError
|
||||
try:
|
||||
await chroma.aadd_documents(docs)
|
||||
break
|
||||
except Exception as e:
|
||||
if attempt == MAX_RETRY_ATTEMPTS - 1:
|
||||
raise
|
||||
wait = (attempt + 1) * EXPONENTIAL_BACKOFF_MULTIPLIER
|
||||
await logger.awarning("Write failed, retrying in %ds: %s", wait, e)
|
||||
await asyncio.sleep(wait)
|
||||
|
||||
await asyncio.sleep(0.01)
|
||||
written = await KBIngestionHelper.write_documents_to_chroma(
|
||||
documents=docs,
|
||||
chroma=chroma,
|
||||
task_job_id=task_job_id,
|
||||
job_service=job_service,
|
||||
)
|
||||
if written < len(docs):
|
||||
# Job was cancelled mid-file
|
||||
raise IngestionCancelledError
|
||||
|
||||
total_chunks_created += len(chunks)
|
||||
processed_files.append(file_name)
|
||||
@ -555,13 +545,69 @@ class KBIngestionHelper:
|
||||
KBStorageHelper.release_chroma_resources(kb_path)
|
||||
|
||||
@staticmethod
|
||||
async def _is_job_cancelled(job_service: JobService, job_id: uuid.UUID) -> bool:
|
||||
async def write_documents_to_chroma(
|
||||
*,
|
||||
documents: list[Document],
|
||||
chroma: Chroma,
|
||||
task_job_id: uuid.UUID,
|
||||
job_service: JobService,
|
||||
) -> int:
|
||||
"""Write pre-built Documents into an open Chroma collection.
|
||||
|
||||
This is the shared primitive used by both file-based KB ingestion
|
||||
(``perform_ingestion``) and message-based Memory Base ingestion.
|
||||
|
||||
Documents must already be chunked and have their metadata populated
|
||||
by the caller — this method only handles the batched write, cancellation
|
||||
checking, and retry logic.
|
||||
|
||||
Args:
|
||||
documents: LangChain Document objects ready for embedding.
|
||||
chroma: An already-constructed ``Chroma`` instance pointing at the
|
||||
target collection.
|
||||
task_job_id: Job ID used to poll for cancellation.
|
||||
job_service: Service for checking job status.
|
||||
|
||||
Returns:
|
||||
Number of documents successfully written. If the job is cancelled
|
||||
mid-batch this will be less than ``len(documents)``.
|
||||
|
||||
Raises:
|
||||
Exception: Re-raises any non-cancellation write failure after the
|
||||
retry budget is exhausted.
|
||||
"""
|
||||
written = 0
|
||||
for i in range(0, len(documents), INGESTION_BATCH_SIZE):
|
||||
if await KBIngestionHelper.is_job_cancelled(job_service, task_job_id):
|
||||
return written
|
||||
|
||||
batch = documents[i : i + INGESTION_BATCH_SIZE]
|
||||
for attempt in range(MAX_RETRY_ATTEMPTS):
|
||||
if await KBIngestionHelper.is_job_cancelled(job_service, task_job_id):
|
||||
return written
|
||||
try:
|
||||
await chroma.aadd_documents(batch)
|
||||
break
|
||||
except Exception as e:
|
||||
if attempt == MAX_RETRY_ATTEMPTS - 1:
|
||||
raise
|
||||
wait = (attempt + 1) * EXPONENTIAL_BACKOFF_MULTIPLIER
|
||||
await logger.awarning("Write failed, retrying in %ds: %s", wait, e)
|
||||
await asyncio.sleep(wait)
|
||||
|
||||
written += len(batch)
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
return written
|
||||
|
||||
@staticmethod
|
||||
async def is_job_cancelled(job_service: JobService, job_id: uuid.UUID) -> bool:
|
||||
"""Internal helper to check if a job has been cancelled."""
|
||||
job = await job_service.get_job_by_job_id(job_id)
|
||||
return job is not None and job.status == JobStatus.CANCELLED
|
||||
|
||||
@staticmethod
|
||||
async def _build_embeddings(provider: str, model: str, current_user):
|
||||
async def build_embeddings(provider: str, model: str, current_user):
|
||||
"""Internal helper to build embeddings object."""
|
||||
options = get_embedding_model_options(user_id=current_user.id)
|
||||
selected_option = next((o for o in options if o["provider"] == provider and o["name"] == model), None)
|
||||
|
||||
@ -10,6 +10,7 @@ from langflow.api.v1.knowledge_bases import router as knowledge_bases_router
|
||||
from langflow.api.v1.login import router as login_router
|
||||
from langflow.api.v1.mcp import router as mcp_router
|
||||
from langflow.api.v1.mcp_projects import router as mcp_projects_router
|
||||
from langflow.api.v1.memories import router as memories_router
|
||||
from langflow.api.v1.model_options import router as model_options_router
|
||||
from langflow.api.v1.models import router as models_router
|
||||
from langflow.api.v1.monitor import router as monitor_router
|
||||
@ -36,6 +37,7 @@ __all__ = [
|
||||
"login_router",
|
||||
"mcp_projects_router",
|
||||
"mcp_router",
|
||||
"memories_router",
|
||||
"model_options_router",
|
||||
"models_router",
|
||||
"monitor_router",
|
||||
|
||||
@ -6,7 +6,7 @@ import time
|
||||
from collections.abc import AsyncGenerator
|
||||
from http import HTTPStatus
|
||||
from typing import TYPE_CHECKING, Annotated
|
||||
from uuid import uuid4
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
import orjson
|
||||
import sqlalchemy as sa
|
||||
@ -62,8 +62,17 @@ from langflow.services.auth.utils import (
|
||||
from langflow.services.cache.utils import save_uploaded_file
|
||||
from langflow.services.database.models.flow.model import Flow, FlowRead
|
||||
from langflow.services.database.models.flow.utils import get_all_webhook_components_in_flow
|
||||
from langflow.services.database.models.jobs.model import JobType
|
||||
from langflow.services.database.models.user.model import User, UserRead
|
||||
from langflow.services.deps import get_auth_service, get_session_service, get_settings_service, get_telemetry_service
|
||||
from langflow.services.deps import (
|
||||
get_auth_service,
|
||||
get_job_service,
|
||||
get_memory_base_service,
|
||||
get_session_service,
|
||||
get_settings_service,
|
||||
get_task_service,
|
||||
get_telemetry_service,
|
||||
)
|
||||
from langflow.services.event_manager import create_webhook_event_manager, webhook_event_manager
|
||||
from langflow.services.telemetry.schema import RunPayload
|
||||
from langflow.utils.compression import compress_response
|
||||
@ -173,8 +182,8 @@ async def simple_run_flow(
|
||||
graph = Graph.from_payload(
|
||||
graph_data, flow_id=flow_id_str, user_id=str(user_id), flow_name=flow.name, context=context
|
||||
)
|
||||
if run_id is None:
|
||||
run_id = str(uuid4())
|
||||
run_id_uuid = uuid4() if run_id is None else UUID(run_id)
|
||||
run_id = str(run_id_uuid)
|
||||
graph.set_run_id(run_id)
|
||||
inputs = None
|
||||
if input_request.input_value is not None:
|
||||
@ -197,15 +206,57 @@ async def simple_run_flow(
|
||||
and (input_request.output_type == "any" or input_request.output_type in vertex.id.lower()) # type: ignore[operator]
|
||||
)
|
||||
]
|
||||
task_result, session_id = await run_graph_internal(
|
||||
graph=graph,
|
||||
flow_id=flow_id_str,
|
||||
session_id=input_request.session_id,
|
||||
inputs=inputs,
|
||||
outputs=outputs,
|
||||
stream=stream,
|
||||
event_manager=event_manager,
|
||||
)
|
||||
|
||||
# Create a WORKFLOW job record so memory-base on_flow_output can track this run.
|
||||
if user_id is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Authentication required to run flows.",
|
||||
)
|
||||
|
||||
try:
|
||||
_job_svc = get_job_service()
|
||||
await _job_svc.create_job(
|
||||
job_id=run_id_uuid,
|
||||
flow_id=flow.id,
|
||||
user_id=user_id,
|
||||
job_type=JobType.WORKFLOW,
|
||||
)
|
||||
task_result, session_id = await _job_svc.execute_with_status(
|
||||
run_id_uuid,
|
||||
run_graph_internal,
|
||||
graph=graph,
|
||||
flow_id=flow_id_str,
|
||||
session_id=input_request.session_id,
|
||||
inputs=inputs,
|
||||
outputs=outputs,
|
||||
stream=stream,
|
||||
event_manager=event_manager,
|
||||
)
|
||||
except Exception as exc:
|
||||
await logger.aerror(
|
||||
"Workflow job execution failed for flow %s: %s",
|
||||
flow.id,
|
||||
str(exc),
|
||||
exc_info=True,
|
||||
)
|
||||
raise APIException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
exception=exc,
|
||||
flow=flow,
|
||||
) from exc
|
||||
|
||||
# Fire memory-base auto-capture hook — non-blocking background effect.
|
||||
try:
|
||||
_run_id_uuid = UUID(graph.run_id) if graph.run_id else None # type-cast only
|
||||
await get_task_service().fire_and_forget_task(
|
||||
get_memory_base_service().on_flow_output,
|
||||
flow_id=flow.id,
|
||||
session_id=session_id,
|
||||
job_id=_run_id_uuid,
|
||||
)
|
||||
except (RuntimeError, ValueError, OSError):
|
||||
await logger.awarning("Memory base hook scheduling failed for flow %s", flow.id, exc_info=True)
|
||||
|
||||
return RunResponse(outputs=task_result, session_id=session_id)
|
||||
|
||||
@ -970,6 +1021,18 @@ async def experimental_run_flow(
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(exc)) from exc
|
||||
|
||||
# Fire memory-base auto-capture hook — non-blocking background effect.
|
||||
try:
|
||||
_run_id_uuid = UUID(graph.run_id) if graph.run_id else None # type-cast only
|
||||
await get_task_service().fire_and_forget_task(
|
||||
get_memory_base_service().on_flow_output,
|
||||
flow_id=flow.id,
|
||||
session_id=session_id,
|
||||
job_id=_run_id_uuid,
|
||||
)
|
||||
except (RuntimeError, ValueError, OSError):
|
||||
await logger.awarning("Memory base hook scheduling failed for flow %s", flow.id, exc_info=True)
|
||||
|
||||
return RunResponse(outputs=task_result, session_id=session_id)
|
||||
|
||||
|
||||
|
||||
@ -74,6 +74,30 @@ def _resolve_kb_path(kb_name: str, current_user: CurrentActiveUser) -> Path:
|
||||
return kb_path
|
||||
|
||||
|
||||
def _is_memory_base_associated(metadata: dict[str, Any]) -> bool:
|
||||
"""Return True if the KB metadata indicates an association with a Memory Base."""
|
||||
source_types = metadata.get("source_types")
|
||||
return isinstance(source_types, list) and "memory" in source_types
|
||||
|
||||
|
||||
def _check_memory_base_association(kb_name: str, current_user: CurrentActiveUser) -> None:
|
||||
"""Raise 403 if the KB is associated with a Memory Base.
|
||||
|
||||
Designed as a FastAPI dependency for per-KB routes — FastAPI injects
|
||||
``kb_name`` from the path parameter and ``current_user`` via its own
|
||||
dependency. The list endpoint filters memory KBs inline using
|
||||
``_is_memory_base_associated`` directly.
|
||||
"""
|
||||
kb_path = _resolve_kb_path(kb_name, current_user)
|
||||
|
||||
metadata = KBAnalysisHelper.get_metadata(kb_path, fast=True)
|
||||
if _is_memory_base_associated(metadata):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"Access denied: knowledge base '{kb_name}' is managed by a Memory Base.",
|
||||
)
|
||||
|
||||
|
||||
@router.post("", status_code=HTTPStatus.CREATED)
|
||||
@router.post("/", status_code=HTTPStatus.CREATED)
|
||||
async def create_knowledge_base(
|
||||
@ -275,7 +299,7 @@ async def preview_chunks(
|
||||
return {"files": file_previews}
|
||||
|
||||
|
||||
@router.post("/{kb_name}/ingest", status_code=HTTPStatus.OK)
|
||||
@router.post("/{kb_name}/ingest", status_code=HTTPStatus.OK, dependencies=[Depends(_check_memory_base_association)])
|
||||
async def ingest_files_to_knowledge_base(
|
||||
kb_name: str,
|
||||
current_user: CurrentActiveUser,
|
||||
@ -427,6 +451,8 @@ async def list_knowledge_bases(
|
||||
try:
|
||||
# Use deep update (fast=False) to ensure legacy KBs are migrated on first view
|
||||
metadata = KBAnalysisHelper.get_metadata(kb_dir, fast=False)
|
||||
if _is_memory_base_associated(metadata):
|
||||
continue # Skip KBs that are associated with a Memory Base
|
||||
|
||||
# Extract KB ID from metadata (stored as string, convert to UUID)
|
||||
kb_id_str = metadata.get("id")
|
||||
@ -503,7 +529,7 @@ async def list_knowledge_bases(
|
||||
return knowledge_bases
|
||||
|
||||
|
||||
@router.get("/{kb_name}", status_code=HTTPStatus.OK)
|
||||
@router.get("/{kb_name}", status_code=HTTPStatus.OK, dependencies=[Depends(_check_memory_base_association)])
|
||||
async def get_knowledge_base(kb_name: str, current_user: CurrentActiveUser) -> KnowledgeBaseInfo:
|
||||
"""Get detailed information about a specific knowledge base."""
|
||||
try:
|
||||
@ -543,7 +569,7 @@ async def get_knowledge_base(kb_name: str, current_user: CurrentActiveUser) -> K
|
||||
raise HTTPException(status_code=500, detail="Error getting knowledge base.") from e
|
||||
|
||||
|
||||
@router.get("/{kb_name}/chunks", status_code=HTTPStatus.OK)
|
||||
@router.get("/{kb_name}/chunks", status_code=HTTPStatus.OK, dependencies=[Depends(_check_memory_base_association)])
|
||||
async def get_knowledge_base_chunks(
|
||||
kb_name: str,
|
||||
current_user: CurrentActiveUser,
|
||||
@ -636,7 +662,7 @@ async def get_knowledge_base_chunks(
|
||||
KBStorageHelper.release_chroma_resources(kb_path)
|
||||
|
||||
|
||||
@router.delete("/{kb_name}", status_code=HTTPStatus.OK)
|
||||
@router.delete("/{kb_name}", status_code=HTTPStatus.OK, dependencies=[Depends(_check_memory_base_association)])
|
||||
async def delete_knowledge_base(kb_name: str, current_user: CurrentActiveUser) -> dict[str, str]:
|
||||
"""Delete a specific knowledge base."""
|
||||
try:
|
||||
@ -703,7 +729,7 @@ async def delete_knowledge_bases_bulk(request: BulkDeleteRequest, current_user:
|
||||
return result
|
||||
|
||||
|
||||
@router.post("/{kb_name}/cancel", status_code=HTTPStatus.OK)
|
||||
@router.post("/{kb_name}/cancel", status_code=HTTPStatus.OK, dependencies=[Depends(_check_memory_base_association)])
|
||||
async def cancel_ingestion(
|
||||
kb_name: str,
|
||||
current_user: CurrentActiveUser,
|
||||
|
||||
345
src/backend/base/langflow/api/v1/memories.py
Normal file
345
src/backend/base/langflow/api/v1/memories.py
Normal file
@ -0,0 +1,345 @@
|
||||
"""REST API for Memory Base management.
|
||||
|
||||
Endpoints:
|
||||
POST /memories - Create
|
||||
GET /memories - List (current user, paginated)
|
||||
GET /memories/{id} - Get one
|
||||
GET /memories/{id}/sessions - List sessions (tracked + untracked from MessageTable)
|
||||
PATCH /memories/{id} - Update (name / threshold / auto_capture / preprocessing)
|
||||
DELETE /memories/{id} - Delete (cancels active tasks + removes KB from disk)
|
||||
POST /memories/{id}/flush - Manual flush / trigger ingestion
|
||||
POST /memories/{id}/regenerate - Regenerate from mismatch
|
||||
|
||||
Edge cases enforced:
|
||||
409 Conflict - name already in use for this user (on create).
|
||||
409 Conflict - active ingestion task already running for same (mb, session).
|
||||
404 Not Found - memory base does not belong to the current user.
|
||||
422 Unprocessable - preprocessing=true but preproc_model missing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from http import HTTPStatus
|
||||
from typing import Annotated, Any
|
||||
|
||||
from fastapi import APIRouter, Body, Depends, HTTPException
|
||||
from fastapi_pagination import Page, Params
|
||||
from fastapi_pagination.ext.sqlmodel import apaginate
|
||||
from pydantic import BaseModel
|
||||
from sqlmodel import col, select
|
||||
|
||||
from langflow.api.utils import CurrentActiveUser
|
||||
from langflow.services.database.models.memory_base.model import (
|
||||
MemoryBase,
|
||||
MemoryBaseCreate,
|
||||
MemoryBaseRead,
|
||||
MemoryBaseSessionRead,
|
||||
MemoryBaseUpdate,
|
||||
)
|
||||
from langflow.services.database.models.message.model import MessageTable
|
||||
from langflow.services.deps import get_memory_base_service, session_scope
|
||||
from langflow.services.jobs import DuplicateJobError
|
||||
|
||||
router = APIRouter(tags=["Memories"], prefix="/memories", include_in_schema=False)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Request / Response schemas #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
|
||||
class MessageReadResponse(BaseModel):
|
||||
"""Slim message projection for Memory Base session message listings.
|
||||
|
||||
Only messages that have been ingested into the requested Memory Base are returned.
|
||||
``job_id`` and ``ingested_at`` are sourced from MessageIngestionRecord.
|
||||
"""
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
id: uuid.UUID
|
||||
timestamp: datetime | None = None
|
||||
sender: str
|
||||
sender_name: str
|
||||
session_id: str
|
||||
text: str
|
||||
content_blocks: list[dict[str, Any]] = []
|
||||
job_id: uuid.UUID | None = None
|
||||
ingested_at: datetime | None = None
|
||||
|
||||
|
||||
class FlushRequest(BaseModel):
|
||||
session_id: str
|
||||
|
||||
|
||||
class MismatchResponse(BaseModel):
|
||||
mismatch_detected: bool
|
||||
|
||||
|
||||
class RegenerateResponse(BaseModel):
|
||||
job_ids: list[str]
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# CRUD #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
|
||||
@router.post("", status_code=HTTPStatus.CREATED)
|
||||
@router.post("/", status_code=HTTPStatus.CREATED)
|
||||
async def create_memory_base(
|
||||
current_user: CurrentActiveUser,
|
||||
payload: Annotated[MemoryBaseCreate, Body(embed=False)] = ...,
|
||||
) -> MemoryBaseRead:
|
||||
"""Create a new Memory Base.
|
||||
|
||||
- kb_name is auto-generated as `{sanitized_name}_{8hex}`.
|
||||
- KB directory and embedding_metadata.json are created on disk immediately.
|
||||
- Returns 409 if a Memory Base with the same name already exists for this user.
|
||||
- Returns 422 if preprocessing=true but preproc_model is missing.
|
||||
"""
|
||||
try:
|
||||
mb = await get_memory_base_service().create(payload, user_id=current_user.id)
|
||||
except PermissionError as exc:
|
||||
# Flow not found or belongs to another user — return 404 to avoid info-leak
|
||||
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
return MemoryBaseRead.model_validate(mb)
|
||||
|
||||
|
||||
@router.get("", status_code=HTTPStatus.OK)
|
||||
@router.get("/", status_code=HTTPStatus.OK)
|
||||
async def list_memory_bases(
|
||||
current_user: CurrentActiveUser,
|
||||
params: Annotated[Params, Depends()],
|
||||
flow_id: uuid.UUID | None = None,
|
||||
) -> Page[MemoryBaseRead]:
|
||||
"""List all Memory Bases owned by the current user (paginated) for a flow_id.
|
||||
|
||||
Query params (from fastapi-pagination):
|
||||
page - 1-based page number (default 1)
|
||||
size - page size (default 50)
|
||||
"""
|
||||
async with session_scope() as db:
|
||||
stmt = get_memory_base_service().list_for_user_stmt(user_id=current_user.id, flow_id=flow_id)
|
||||
return await apaginate(
|
||||
db, stmt, params=params, transformer=lambda items: [MemoryBaseRead.model_validate(m) for m in items]
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{memory_base_id}", status_code=HTTPStatus.OK)
|
||||
async def get_memory_base(
|
||||
memory_base_id: uuid.UUID,
|
||||
current_user: CurrentActiveUser,
|
||||
) -> MemoryBaseRead:
|
||||
"""Get details for a specific Memory Base."""
|
||||
mb = await get_memory_base_service().get(memory_base_id, user_id=current_user.id)
|
||||
if mb is None:
|
||||
raise HTTPException(status_code=404, detail="Memory base not found")
|
||||
return MemoryBaseRead.model_validate(mb)
|
||||
|
||||
|
||||
@router.get("/{memory_base_id}/sessions", status_code=HTTPStatus.OK)
|
||||
async def list_sessions(
|
||||
memory_base_id: uuid.UUID,
|
||||
current_user: CurrentActiveUser,
|
||||
params: Annotated[Params, Depends()],
|
||||
) -> Page[MemoryBaseSessionRead]:
|
||||
"""List persisted sessions for this Memory Base (paginated).
|
||||
|
||||
Only sessions that have been synced at least once (i.e. have a
|
||||
MemoryBaseSession row) are returned. Results are ordered by
|
||||
last_sync_at descending.
|
||||
|
||||
Each item includes ``pending_count``: the number of completed flow runs
|
||||
remaining before the next auto-capture ingestion is triggered.
|
||||
"""
|
||||
from langflow.services.memory_base.ingestion import count_pending_messages
|
||||
|
||||
async with session_scope() as db:
|
||||
try:
|
||||
mb = await get_memory_base_service().get_memory_base_or_404(db, memory_base_id, current_user.id)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
||||
|
||||
stmt = get_memory_base_service().sessions_stmt(memory_base_id, current_user.id)
|
||||
raw_page = await apaginate(db, stmt, params=params)
|
||||
|
||||
items: list[MemoryBaseSessionRead] = []
|
||||
for s in raw_page.items:
|
||||
pending_count = await count_pending_messages(db, mb, s)
|
||||
read = MemoryBaseSessionRead.model_validate(s)
|
||||
read.pending_count = pending_count
|
||||
items.append(read)
|
||||
|
||||
return raw_page.model_copy(update={"items": items})
|
||||
|
||||
|
||||
@router.get("/{memory_base_id}/sessions/{session_id}/messages", status_code=HTTPStatus.OK)
|
||||
async def list_session_messages(
|
||||
memory_base_id: uuid.UUID,
|
||||
session_id: str,
|
||||
current_user: CurrentActiveUser,
|
||||
params: Annotated[Params, Depends()],
|
||||
) -> Page[MessageReadResponse]:
|
||||
"""List messages ingested into this Memory Base session (paginated).
|
||||
|
||||
Only messages that have been successfully ingested into the requested Memory Base
|
||||
are returned. Messages are ordered by timestamp ascending.
|
||||
Each item includes ``job_id`` and ``ingested_at`` from the MessageIngestionRecord.
|
||||
|
||||
Returns 404 if the Memory Base does not belong to the current user.
|
||||
"""
|
||||
from sqlalchemy import and_
|
||||
|
||||
from langflow.services.database.models.memory_base.model import MessageIngestionRecord
|
||||
|
||||
async with session_scope() as db:
|
||||
mb_stmt = select(MemoryBase).where(MemoryBase.id == memory_base_id).where(MemoryBase.user_id == current_user.id)
|
||||
result = await db.exec(mb_stmt)
|
||||
if result.first() is None:
|
||||
raise HTTPException(status_code=404, detail="Memory base not found")
|
||||
|
||||
# INNER JOIN — only messages that were actually ingested into this MB/session pair.
|
||||
# No extra WHERE filters needed:
|
||||
# - mir.session_id == session_id in the JOIN guarantees msg.session_id == session_id
|
||||
# (session_id is denormalized from the message at ingestion time — immutable).
|
||||
# - flow_id is implicitly correct: ingestion only ever touches messages from mb.flow_id,
|
||||
# and MB ownership is already verified above.
|
||||
msg_stmt = (
|
||||
select(MessageTable, MessageIngestionRecord)
|
||||
.join(
|
||||
MessageIngestionRecord,
|
||||
and_(
|
||||
MessageIngestionRecord.message_id == MessageTable.id,
|
||||
MessageIngestionRecord.memory_base_id == memory_base_id,
|
||||
MessageIngestionRecord.session_id == session_id,
|
||||
),
|
||||
)
|
||||
.order_by(col(MessageTable.timestamp).asc())
|
||||
)
|
||||
return await apaginate(
|
||||
db,
|
||||
msg_stmt,
|
||||
params=params,
|
||||
transformer=lambda rows: [
|
||||
MessageReadResponse(
|
||||
id=msg.id,
|
||||
timestamp=msg.timestamp,
|
||||
sender=msg.sender,
|
||||
sender_name=msg.sender_name,
|
||||
session_id=msg.session_id,
|
||||
text=msg.text,
|
||||
content_blocks=msg.content_blocks or [],
|
||||
job_id=mir.job_id,
|
||||
ingested_at=mir.ingested_at,
|
||||
)
|
||||
for msg, mir in rows
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@router.patch("/{memory_base_id}", status_code=HTTPStatus.OK)
|
||||
async def update_memory_base(
|
||||
memory_base_id: uuid.UUID,
|
||||
current_user: CurrentActiveUser,
|
||||
patch: Annotated[MemoryBaseUpdate, Body(embed=False)] = ...,
|
||||
) -> MemoryBaseRead:
|
||||
"""Update mutable parameters (threshold, auto_capture, preprocessing, etc.).
|
||||
|
||||
Threshold changes only take effect at the next auto-capture trigger.
|
||||
Any already-running ingestion task continues with its original arguments.
|
||||
"""
|
||||
mb = await get_memory_base_service().update(memory_base_id, user_id=current_user.id, patch=patch)
|
||||
if mb is None:
|
||||
raise HTTPException(status_code=404, detail="Memory base not found")
|
||||
return MemoryBaseRead.model_validate(mb)
|
||||
|
||||
|
||||
@router.delete("/{memory_base_id}", status_code=HTTPStatus.NO_CONTENT)
|
||||
async def delete_memory_base(
|
||||
memory_base_id: uuid.UUID,
|
||||
current_user: CurrentActiveUser,
|
||||
) -> None:
|
||||
"""Delete a Memory Base.
|
||||
|
||||
Active ingestion tasks are forcefully cancelled before the DB record is
|
||||
removed. The associated KB directory is deleted from disk afterwards
|
||||
(best-effort — a disk failure will not affect the 204 response).
|
||||
"""
|
||||
deleted = await get_memory_base_service().delete(memory_base_id, user_id=current_user.id)
|
||||
if not deleted:
|
||||
raise HTTPException(status_code=404, detail="Memory base not found")
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Ingestion trigger #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
|
||||
@router.post("/{memory_base_id}/flush", status_code=HTTPStatus.ACCEPTED)
|
||||
async def flush_memory_base(
|
||||
memory_base_id: uuid.UUID,
|
||||
current_user: CurrentActiveUser,
|
||||
body: Annotated[FlushRequest, Body(embed=False)] = ...,
|
||||
) -> dict:
|
||||
"""Manually trigger an ingestion / sync job regardless of the threshold.
|
||||
|
||||
Returns 409 Conflict if an ingestion task is already in progress for the
|
||||
given (memory_base_id, session_id) pair to prevent concurrent indexing.
|
||||
"""
|
||||
try:
|
||||
job_id = await get_memory_base_service().trigger_ingestion(
|
||||
memory_base_id=memory_base_id,
|
||||
user_id=current_user.id,
|
||||
session_id=body.session_id,
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
||||
except DuplicateJobError as exc:
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
except RuntimeError as exc:
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
|
||||
return {"job_id": job_id}
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Mismatch detection & regeneration #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
|
||||
@router.get("/{memory_base_id}/mismatch", status_code=HTTPStatus.OK)
|
||||
async def check_mismatch(
|
||||
memory_base_id: uuid.UUID,
|
||||
current_user: CurrentActiveUser,
|
||||
) -> MismatchResponse:
|
||||
"""Detect if the vector store is empty while metadata records processed messages.
|
||||
|
||||
The UI should surface a "Mismatch Detected" warning and offer a Regenerate button.
|
||||
"""
|
||||
try:
|
||||
detected = await get_memory_base_service().check_mismatch(memory_base_id, user_id=current_user.id)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
||||
return MismatchResponse(mismatch_detected=detected)
|
||||
|
||||
|
||||
@router.post("/{memory_base_id}/regenerate", status_code=HTTPStatus.ACCEPTED)
|
||||
async def regenerate_memory_base(
|
||||
memory_base_id: uuid.UUID,
|
||||
current_user: CurrentActiveUser,
|
||||
) -> RegenerateResponse:
|
||||
"""Regenerate the Knowledge Base by resetting all session cursors and re-ingesting.
|
||||
|
||||
Use this to recover from external Chroma directory deletions or vector DB corruption.
|
||||
All MemoryBaseSession.cursor_id values are set to None before re-running ingestion.
|
||||
"""
|
||||
try:
|
||||
job_ids = await get_memory_base_service().regenerate(memory_base_id, user_id=current_user.id)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
||||
return RegenerateResponse(job_ids=job_ids)
|
||||
@ -29,6 +29,7 @@ from uuid import UUID, uuid4
|
||||
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Query, Request, status
|
||||
from fastapi.responses import StreamingResponse
|
||||
from lfx.graph.graph.base import Graph
|
||||
from lfx.log.logger import logger
|
||||
from lfx.schema.workflow import (
|
||||
WORKFLOW_EXECUTION_RESPONSES,
|
||||
WORKFLOW_STATUS_RESPONSES,
|
||||
@ -65,7 +66,7 @@ from langflow.services.auth.utils import api_key_security
|
||||
from langflow.services.database.models.flow.model import FlowRead
|
||||
from langflow.services.database.models.jobs.model import JobType
|
||||
from langflow.services.database.models.user.model import UserRead
|
||||
from langflow.services.deps import get_job_service, get_task_service
|
||||
from langflow.services.deps import get_job_service, get_memory_base_service, get_task_service
|
||||
|
||||
# Configuration constants
|
||||
EXECUTION_TIMEOUT = 300 # 5 minutes default timeout for sync execution
|
||||
@ -391,6 +392,18 @@ async def execute_sync_workflow(
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# Fire memory-base auto-capture hook — non-blocking background effect.
|
||||
try:
|
||||
_run_id_uuid = UUID(graph.run_id) if graph.run_id else None # type-cast only; same run_id set on graph
|
||||
await get_task_service().fire_and_forget_task(
|
||||
get_memory_base_service().on_flow_output,
|
||||
flow_id=flow.id,
|
||||
session_id=execution_session_id,
|
||||
job_id=_run_id_uuid,
|
||||
)
|
||||
except (RuntimeError, ValueError, OSError):
|
||||
await logger.awarning("Memory base hook scheduling failed for flow %s", flow.id, exc_info=True)
|
||||
|
||||
# Build RunResponse
|
||||
run_response = RunResponse(outputs=task_result, session_id=execution_session_id)
|
||||
# Convert to WorkflowExecutionResponse
|
||||
@ -470,10 +483,39 @@ async def execute_workflow_background(
|
||||
user_id=api_key_user.id,
|
||||
)
|
||||
|
||||
# Closure captures flow identity for the memory-base hook.
|
||||
# run_id is the same as job_id — graph.set_run_id(job_id) was called above.
|
||||
_hook_flow_id = flow.id
|
||||
_hook_run_id = job_id
|
||||
|
||||
async def _run_and_notify(**kwargs):
|
||||
"""Thin wrapper: execute graph then fire memory-base hook as a background effect.
|
||||
|
||||
The hook is dispatched non-blocking after graph completion. Any failure in
|
||||
the hook is swallowed so it never affects the job status of the graph run.
|
||||
"""
|
||||
result = await run_graph_internal(**kwargs)
|
||||
_, _effective_session_id = result
|
||||
try:
|
||||
# Direct await — we are already inside a background task; awaiting here
|
||||
# is non-blocking from the client's perspective and avoids the race
|
||||
# condition that arises when dispatching a second fire_and_forget from
|
||||
# within an already-running fire_and_forget task.
|
||||
await get_memory_base_service().on_flow_output(
|
||||
flow_id=_hook_flow_id,
|
||||
session_id=_effective_session_id,
|
||||
job_id=_hook_run_id,
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
await logger.awarning(
|
||||
"Memory base hook failed for flow %s, but workflow succeeded.", _hook_flow_id, exc_info=True
|
||||
)
|
||||
return result
|
||||
|
||||
await task_service.fire_and_forget_task(
|
||||
job_service.execute_with_status,
|
||||
job_id=job_id,
|
||||
run_coro_func=run_graph_internal,
|
||||
run_coro_func=_run_and_notify,
|
||||
graph=graph,
|
||||
flow_id=flow_id_str,
|
||||
session_id=session_id,
|
||||
@ -673,7 +715,7 @@ async def stop_workflow(
|
||||
task_service = get_task_service()
|
||||
|
||||
try:
|
||||
# 1. Fetch Job
|
||||
# 1. Fetch Job and verify ownership
|
||||
job = await job_service.get_job_by_job_id(job_id, user_id=api_key_user.id)
|
||||
except Exception as exc:
|
||||
raise HTTPException(
|
||||
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@ -8,6 +8,7 @@ from .flow_version import FlowVersion
|
||||
from .flow_version_deployment_attachment import FlowVersionDeploymentAttachment
|
||||
from .folder import Folder
|
||||
from .jobs import Job
|
||||
from .memory_base import MemoryBase, MemoryBaseSession, MemoryBaseWorkflowRun, MessageIngestionRecord
|
||||
from .message import MessageTable
|
||||
from .traces.model import SpanTable, TraceTable
|
||||
from .transactions import TransactionTable
|
||||
@ -24,6 +25,10 @@ __all__ = [
|
||||
"FlowVersionDeploymentAttachment",
|
||||
"Folder",
|
||||
"Job",
|
||||
"MemoryBase",
|
||||
"MemoryBaseSession",
|
||||
"MemoryBaseWorkflowRun",
|
||||
"MessageIngestionRecord",
|
||||
"MessageTable",
|
||||
"SSOConfig",
|
||||
"SSOUserProfile",
|
||||
|
||||
@ -63,6 +63,9 @@ class JobBase(SQLModel):
|
||||
asset_type: str | None = Field(
|
||||
index=False, nullable=True
|
||||
) # Polymorphic: records if job is related to an entity like a KB, workflow, etc.
|
||||
dedupe_key: str | None = Field(
|
||||
index=True, nullable=True
|
||||
) # Optional idempotency key to prevent duplicate jobs for the same asset and operation.
|
||||
|
||||
|
||||
class Job(JobBase, table=True): # type: ignore[call-arg]
|
||||
|
||||
@ -0,0 +1,21 @@
|
||||
from langflow.services.database.models.memory_base.model import (
|
||||
MemoryBase,
|
||||
MemoryBaseCreate,
|
||||
MemoryBaseRead,
|
||||
MemoryBaseSession,
|
||||
MemoryBaseSessionRead,
|
||||
MemoryBaseUpdate,
|
||||
MemoryBaseWorkflowRun,
|
||||
MessageIngestionRecord,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"MemoryBase",
|
||||
"MemoryBaseCreate",
|
||||
"MemoryBaseRead",
|
||||
"MemoryBaseSession",
|
||||
"MemoryBaseSessionRead",
|
||||
"MemoryBaseUpdate",
|
||||
"MemoryBaseWorkflowRun",
|
||||
"MessageIngestionRecord",
|
||||
]
|
||||
@ -0,0 +1,199 @@
|
||||
from datetime import datetime, timezone
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
import sqlalchemy as sa
|
||||
from pydantic import model_validator
|
||||
from sqlalchemy import Column, DateTime, ForeignKey, Index, UniqueConstraint
|
||||
from sqlmodel import Field, Relationship, SQLModel
|
||||
|
||||
|
||||
class MemoryBaseBase(SQLModel):
|
||||
name: str = Field(index=False)
|
||||
flow_id: UUID = Field(index=True)
|
||||
user_id: UUID = Field(index=True)
|
||||
threshold: int = Field(default=50)
|
||||
auto_capture: bool = Field(default=True)
|
||||
# Preprocessing config — accepted in payload but logic deferred to future scope
|
||||
embedding_model: str = Field(default="")
|
||||
preprocessing: bool = Field(default=False)
|
||||
preproc_model: str | None = Field(default=None)
|
||||
preproc_instructions: str | None = Field(default=None)
|
||||
|
||||
|
||||
class MemoryBase(MemoryBaseBase, table=True): # type: ignore[call-arg]
|
||||
__tablename__ = "memory_base"
|
||||
__table_args__ = (UniqueConstraint("user_id", "name", name="uq_memory_base_user_name"),)
|
||||
|
||||
id: UUID = Field(default_factory=uuid4, primary_key=True)
|
||||
# kb_name is auto-generated at creation time — not user-supplied
|
||||
kb_name: str = Field(default="")
|
||||
created_at: datetime = Field(
|
||||
default_factory=lambda: datetime.now(timezone.utc),
|
||||
sa_column=Column(DateTime(timezone=True), nullable=False),
|
||||
)
|
||||
|
||||
sessions: list["MemoryBaseSession"] = Relationship(
|
||||
back_populates="memory_base",
|
||||
sa_relationship_kwargs={"cascade": "all, delete-orphan"},
|
||||
)
|
||||
|
||||
|
||||
class MemoryBaseCreate(MemoryBaseBase):
|
||||
user_id: UUID | None = None # Derived from auth token in the endpoint; not required in request body
|
||||
|
||||
@model_validator(mode="after")
|
||||
def preproc_model_required_when_preprocessing(self) -> "MemoryBaseCreate":
|
||||
if self.preprocessing and not self.preproc_model:
|
||||
msg = "preproc_model is required when preprocessing is enabled"
|
||||
raise ValueError(msg)
|
||||
return self
|
||||
|
||||
|
||||
class MemoryBaseUpdate(SQLModel):
|
||||
name: str | None = None
|
||||
threshold: int | None = None
|
||||
auto_capture: bool | None = None
|
||||
preprocessing: bool | None = None
|
||||
preproc_model: str | None = None
|
||||
preproc_instructions: str | None = None
|
||||
|
||||
|
||||
class MemoryBaseRead(MemoryBaseBase):
|
||||
id: UUID
|
||||
kb_name: str
|
||||
created_at: datetime
|
||||
|
||||
|
||||
class MemoryBaseSessionBase(SQLModel):
|
||||
"""Fields shared between the table class and response schemas."""
|
||||
|
||||
session_id: str = Field(index=True)
|
||||
cursor_id: UUID | None = Field(default=None)
|
||||
total_processed: int = Field(default=0)
|
||||
last_sync_at: datetime | None = Field(
|
||||
default=None,
|
||||
sa_column=Column(DateTime(timezone=True), nullable=True),
|
||||
)
|
||||
|
||||
|
||||
class MemoryBaseSession(MemoryBaseSessionBase, table=True): # type: ignore[call-arg]
|
||||
__tablename__ = "memory_base_session"
|
||||
|
||||
__table_args__ = (
|
||||
UniqueConstraint("memory_base_id", "session_id", name="uq_memory_base_session"),
|
||||
Index("ix_memory_base_session_lookup", "memory_base_id", "session_id"),
|
||||
)
|
||||
|
||||
id: UUID = Field(default_factory=uuid4, primary_key=True)
|
||||
|
||||
# FK defined via sa_column so Alembic sees the same shape as the migration:
|
||||
# inline ForeignKey on the column with ondelete="CASCADE".
|
||||
# This matches the pattern used by the File model (ForeignKey on sa_column).
|
||||
memory_base_id: UUID = Field(
|
||||
sa_column=Column(
|
||||
sa.Uuid(),
|
||||
ForeignKey("memory_base.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
index=True,
|
||||
)
|
||||
)
|
||||
|
||||
memory_base: MemoryBase = Relationship(back_populates="sessions")
|
||||
|
||||
|
||||
class MemoryBaseSessionRead(MemoryBaseSessionBase):
|
||||
id: UUID
|
||||
memory_base_id: UUID # Explicit — not in base to keep base free of DB-layer FK
|
||||
pending_count: int = Field(default=0)
|
||||
|
||||
|
||||
class MemoryBaseWorkflowRun(SQLModel, table=True): # type: ignore[call-arg]
|
||||
"""Tracks WORKFLOW job runs per (memory_base, session) for threshold-based ingestion.
|
||||
|
||||
One row per WORKFLOW job, per session, per memory base.
|
||||
- ``workflow_job_id``: the WORKFLOW job that produced this run (SET NULL on job deletion).
|
||||
- ``ingestion_job_id``: set only after the ingestion job that processed this run completes
|
||||
successfully. NULL means the run is still pending (not yet counted toward an ingestion).
|
||||
|
||||
Count pending = COUNT(*) WHERE ingestion_job_id IS NULL for a given (memory_base_id, session_id).
|
||||
"""
|
||||
|
||||
__tablename__ = "memory_base_workflow_run"
|
||||
__table_args__ = (
|
||||
UniqueConstraint("memory_base_id", "session_id", "workflow_job_id", name="uq_mbwr_mb_session_wf_job"),
|
||||
Index("ix_mbwr_mb_session", "memory_base_id", "session_id"),
|
||||
Index("ix_mbwr_ingestion_job_id", "ingestion_job_id"),
|
||||
)
|
||||
|
||||
id: UUID = Field(default_factory=uuid4, primary_key=True)
|
||||
memory_base_id: UUID = Field(
|
||||
sa_column=Column(
|
||||
sa.Uuid(),
|
||||
ForeignKey("memory_base.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
)
|
||||
)
|
||||
session_id: str = Field(sa_column=Column(sa.String(), nullable=False))
|
||||
workflow_job_id: UUID | None = Field(
|
||||
default=None,
|
||||
sa_column=Column(
|
||||
sa.Uuid(),
|
||||
ForeignKey("job.job_id", ondelete="SET NULL"),
|
||||
nullable=True,
|
||||
),
|
||||
)
|
||||
ingestion_job_id: UUID | None = Field(
|
||||
default=None,
|
||||
sa_column=Column(
|
||||
sa.Uuid(),
|
||||
ForeignKey("job.job_id", ondelete="SET NULL"),
|
||||
nullable=True,
|
||||
),
|
||||
)
|
||||
recorded_at: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=False))
|
||||
|
||||
|
||||
class MessageIngestionRecord(SQLModel, table=True): # type: ignore[call-arg]
|
||||
"""M-N join table recording which messages were ingested into which Memory Base by which job.
|
||||
|
||||
One record per (message, session, memory_base) — enforced by the unique constraint.
|
||||
Records are written only after a confirmed successful Chroma write (write-on-success).
|
||||
On regenerate, all records for the memory_base are deleted atomically alongside the
|
||||
cursor reset so that re-ingestion starts clean.
|
||||
"""
|
||||
|
||||
__tablename__ = "message_ingestion_record"
|
||||
__table_args__ = (
|
||||
UniqueConstraint("message_id", "session_id", "memory_base_id", name="uq_mir_message_session_mb"),
|
||||
Index("ix_mir_message_id", "message_id"),
|
||||
Index("ix_mir_job_id", "job_id"),
|
||||
Index("ix_mir_memory_base_session", "memory_base_id", "session_id"),
|
||||
)
|
||||
|
||||
id: UUID = Field(default_factory=uuid4, primary_key=True)
|
||||
|
||||
message_id: UUID = Field(
|
||||
sa_column=Column(
|
||||
sa.Uuid(),
|
||||
ForeignKey("message.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
)
|
||||
)
|
||||
memory_base_id: UUID = Field(
|
||||
sa_column=Column(
|
||||
sa.Uuid(),
|
||||
ForeignKey("memory_base.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
)
|
||||
)
|
||||
job_id: UUID | None = Field(
|
||||
default=None,
|
||||
sa_column=Column(
|
||||
sa.Uuid(),
|
||||
ForeignKey("job.job_id", ondelete="SET NULL"),
|
||||
nullable=True,
|
||||
),
|
||||
)
|
||||
# Denormalized from MessageTable.session_id — immutable, avoids JOIN on the hot query path
|
||||
session_id: str = Field(sa_column=Column(sa.String(), nullable=False))
|
||||
ingested_at: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=False))
|
||||
@ -32,6 +32,7 @@ class MessageBase(SQLModel):
|
||||
properties: Properties = Field(default_factory=Properties)
|
||||
category: str = Field(default="message")
|
||||
content_blocks: list[ContentBlock] = Field(default_factory=list)
|
||||
session_metadata: dict | None = Field(default=None)
|
||||
|
||||
@field_serializer("timestamp")
|
||||
def serialize_timestamp(self, value):
|
||||
@ -59,7 +60,7 @@ class MessageBase(SQLModel):
|
||||
return value
|
||||
|
||||
@classmethod
|
||||
def from_message(cls, message: "Message", flow_id: str | UUID | None = None):
|
||||
def from_message(cls, message: "Message", flow_id: str | UUID | None = None, run_id: str | UUID | None = None):
|
||||
if message.text is None or not message.sender or not message.sender_name:
|
||||
msg = "The message does not have the required fields (text, sender, sender_name)."
|
||||
raise ValueError(msg)
|
||||
@ -114,6 +115,13 @@ class MessageBase(SQLModel):
|
||||
msg = f"Flow ID {flow_id} is not a valid UUID"
|
||||
raise ValueError(msg) from exc
|
||||
|
||||
if isinstance(run_id, str):
|
||||
try:
|
||||
run_id = UUID(run_id)
|
||||
except ValueError as exc:
|
||||
msg = f"Run ID {run_id} is not a valid UUID"
|
||||
raise ValueError(msg) from exc
|
||||
|
||||
return cls(
|
||||
sender=message.sender,
|
||||
sender_name=message.sender_name,
|
||||
@ -123,6 +131,7 @@ class MessageBase(SQLModel):
|
||||
files=message.files or [],
|
||||
timestamp=timestamp,
|
||||
flow_id=flow_id,
|
||||
run_id=run_id,
|
||||
properties=properties,
|
||||
category=message.category,
|
||||
content_blocks=content_blocks,
|
||||
@ -149,6 +158,8 @@ class MessageTable(MessageBase, table=True): # type: ignore[call-arg]
|
||||
|
||||
id: UUID = Field(default_factory=uuid4, primary_key=True)
|
||||
flow_id: UUID | None = Field(default=None)
|
||||
run_id: UUID | None = Field(default=None, index=True)
|
||||
is_output: bool = Field(default=False)
|
||||
|
||||
files: list[str] = Field(sa_column=Column(JSON))
|
||||
properties: dict | Properties = Field( # type: ignore[assignment]
|
||||
@ -207,11 +218,9 @@ class MessageTable(MessageBase, table=True): # type: ignore[call-arg]
|
||||
|
||||
@field_serializer("properties", "content_blocks", "session_metadata")
|
||||
@classmethod
|
||||
def serialize_properties_or_content_blocks(cls, value) -> dict | list[dict] | None:
|
||||
def serialize_properties_or_content_blocks(cls, value) -> dict | list[dict]:
|
||||
# Redundant sanitization here acts as a defensive measure for rows
|
||||
# already in the database that might contain NaN/Infinity values.
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, list):
|
||||
value = [cls.serialize_properties_or_content_blocks(item) for item in value]
|
||||
elif hasattr(value, "model_dump"):
|
||||
@ -224,8 +233,9 @@ class MessageTable(MessageBase, table=True): # type: ignore[call-arg]
|
||||
|
||||
class MessageRead(MessageBase):
|
||||
id: UUID
|
||||
flow_id: UUID | None = Field()
|
||||
flow_id: UUID | None = None
|
||||
session_metadata: dict | None = None
|
||||
run_id: UUID | None = None
|
||||
|
||||
|
||||
class MessageCreate(MessageBase):
|
||||
@ -243,3 +253,5 @@ class MessageUpdate(SQLModel):
|
||||
error: bool | None = None
|
||||
properties: Properties | None = None
|
||||
session_metadata: dict | None = None
|
||||
category: str | None = None
|
||||
content_blocks: list[ContentBlock] | None = None
|
||||
|
||||
@ -269,3 +269,10 @@ def get_flow_events_service():
|
||||
from langflow.services.flow_events.factory import FlowEventsServiceFactory
|
||||
|
||||
return get_service(ServiceType.FLOW_EVENTS_SERVICE, FlowEventsServiceFactory())
|
||||
|
||||
|
||||
def get_memory_base_service():
|
||||
"""Retrieves the MemoryBaseService instance from the service manager."""
|
||||
from langflow.services.memory_base.factory import MemoryBaseServiceFactory
|
||||
|
||||
return get_service(ServiceType.MEMORY_BASE_SERVICE, MemoryBaseServiceFactory())
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
"""Job service package."""
|
||||
|
||||
from langflow.services.jobs.exceptions import DuplicateJobError
|
||||
from langflow.services.jobs.service import JobService
|
||||
|
||||
__all__ = ["JobService"]
|
||||
__all__ = ["DuplicateJobError", "JobService"]
|
||||
|
||||
16
src/backend/base/langflow/services/jobs/exceptions.py
Normal file
16
src/backend/base/langflow/services/jobs/exceptions.py
Normal file
@ -0,0 +1,16 @@
|
||||
"""Domain exceptions for the jobs service."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class JobError(RuntimeError):
|
||||
"""Base exception for job-domain errors."""
|
||||
|
||||
|
||||
class DuplicateJobError(JobError):
|
||||
"""Raised by create_job() when a non-retryable job with the same dedupe_key already exists.
|
||||
|
||||
(QUEUED, IN_PROGRESS, or COMPLETED).
|
||||
FAILED and CANCELLED are retryable and do not trigger this error.
|
||||
Extends RuntimeError so existing except RuntimeError callers keep working.
|
||||
"""
|
||||
@ -10,15 +10,16 @@ if TYPE_CHECKING:
|
||||
from datetime import datetime, timezone
|
||||
from uuid import UUID
|
||||
|
||||
from sqlmodel import col, func, select
|
||||
|
||||
from langflow.services.base import Service
|
||||
from langflow.services.database.models.jobs.crud import (
|
||||
get_job_by_job_id,
|
||||
get_jobs_by_flow_id,
|
||||
get_latest_jobs_by_asset_ids,
|
||||
update_job_status,
|
||||
)
|
||||
from langflow.services.database.models.jobs.model import Job, JobStatus, JobType
|
||||
from langflow.services.deps import session_scope
|
||||
from langflow.services.jobs.exceptions import DuplicateJobError
|
||||
|
||||
|
||||
class JobService(Service):
|
||||
@ -30,11 +31,14 @@ class JobService(Service):
|
||||
"""Initialize the job service."""
|
||||
self.set_ready()
|
||||
|
||||
async def get_jobs_by_flow_id(self, flow_id: UUID | str, page: int = 1, page_size: int = 10) -> list[Job]:
|
||||
"""Get jobs for a specific flow with pagination.
|
||||
async def get_jobs_by_flow_id(
|
||||
self, flow_id: UUID | str, user_id: UUID, page: int = 1, page_size: int = 10
|
||||
) -> list[Job]:
|
||||
"""Get jobs for a specific flow with pagination, filtered by user.
|
||||
|
||||
Args:
|
||||
flow_id: The flow ID to filter jobs by
|
||||
user_id: The user ID to enforce ownership
|
||||
page: Page number (1-indexed)
|
||||
page_size: Number of jobs per page
|
||||
|
||||
@ -45,7 +49,16 @@ class JobService(Service):
|
||||
flow_id = UUID(flow_id)
|
||||
|
||||
async with session_scope() as session:
|
||||
return await get_jobs_by_flow_id(session, flow_id, page=page, size=page_size)
|
||||
stmt = (
|
||||
select(Job)
|
||||
.where(Job.flow_id == flow_id)
|
||||
.where((Job.user_id == user_id) | (Job.user_id.is_(None)))
|
||||
.order_by(col(Job.created_at).desc())
|
||||
.offset((page - 1) * page_size)
|
||||
.limit(page_size)
|
||||
)
|
||||
result = await session.exec(stmt)
|
||||
return list(result.all())
|
||||
|
||||
async def get_job_by_job_id(self, job_id: UUID | str, user_id: UUID | None = None) -> Job | None:
|
||||
"""Get job for a specific job ID.
|
||||
@ -62,7 +75,11 @@ class JobService(Service):
|
||||
job_id = UUID(job_id)
|
||||
|
||||
async with session_scope() as session:
|
||||
return await get_job_by_job_id(session, job_id, user_id=user_id)
|
||||
stmt = select(Job).where(Job.job_id == job_id)
|
||||
if user_id:
|
||||
stmt = stmt.where((Job.user_id == user_id) | (Job.user_id.is_(None)))
|
||||
result = await session.exec(stmt)
|
||||
return result.first()
|
||||
|
||||
async def create_job(
|
||||
self,
|
||||
@ -72,16 +89,19 @@ class JobService(Service):
|
||||
asset_id: UUID | None = None,
|
||||
asset_type: str | None = None,
|
||||
user_id: UUID | None = None,
|
||||
dedupe_key: str | None = None,
|
||||
) -> Job:
|
||||
"""Create a new job record with QUEUED status.
|
||||
|
||||
Args:
|
||||
job_id: The job ID
|
||||
flow_id: The flow ID
|
||||
user_id: The user ID
|
||||
job_type: The job type
|
||||
asset_id: The asset ID
|
||||
asset_type: The asset type
|
||||
user_id: The user ID who owns this job
|
||||
dedupe_key: Optional idempotency key to prevent duplicate jobs for the same batch
|
||||
|
||||
Returns:
|
||||
Created Job object
|
||||
@ -93,6 +113,18 @@ class JobService(Service):
|
||||
flow_id = UUID(flow_id)
|
||||
|
||||
async with session_scope() as session:
|
||||
if dedupe_key is not None:
|
||||
stmt = (
|
||||
select(func.count())
|
||||
.select_from(Job)
|
||||
.where(Job.dedupe_key == dedupe_key)
|
||||
.where(col(Job.status).in_([JobStatus.QUEUED, JobStatus.IN_PROGRESS, JobStatus.COMPLETED]))
|
||||
)
|
||||
result = await session.exec(stmt)
|
||||
if result.one() > 0:
|
||||
msg = f"A non-retryable job with dedupe_key={dedupe_key!r} already exists"
|
||||
raise DuplicateJobError(msg)
|
||||
|
||||
job = Job(
|
||||
job_id=job_id,
|
||||
flow_id=flow_id,
|
||||
@ -101,6 +133,7 @@ class JobService(Service):
|
||||
asset_id=asset_id,
|
||||
asset_type=asset_type,
|
||||
user_id=user_id,
|
||||
dedupe_key=dedupe_key,
|
||||
)
|
||||
session.add(job)
|
||||
await session.flush()
|
||||
@ -210,3 +243,18 @@ class JobService(Service):
|
||||
await logger.ainfo(f"Job {job_id} completed successfully")
|
||||
await self.update_job_status(job_id, JobStatus.COMPLETED, finished_timestamp=True)
|
||||
return result
|
||||
|
||||
async def _validate_ownership(self, job_id: UUID, user_id: UUID) -> Job:
|
||||
"""Verify that a job exists and belongs to the specified user.
|
||||
|
||||
Raises:
|
||||
ValueError: If the job is not found or is NOT owned by the user.
|
||||
"""
|
||||
job = await self.get_job_by_job_id(job_id)
|
||||
if job is None:
|
||||
msg = f"Job {job_id} not found"
|
||||
raise ValueError(msg)
|
||||
if job.user_id is not None and job.user_id != user_id:
|
||||
msg = f"Access denied for job {job_id}"
|
||||
raise ValueError(msg)
|
||||
return job
|
||||
|
||||
@ -0,0 +1,3 @@
|
||||
from langflow.services.memory_base.service import MemoryBaseService
|
||||
|
||||
__all__ = ["MemoryBaseService"]
|
||||
@ -0,0 +1,139 @@
|
||||
"""Document building and KB metadata sync helpers for Memory Base ingestion.
|
||||
|
||||
Extracted from task.py to separate "document shaping" from "ingestion orchestration".
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from langchain_core.documents import Document
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from lfx.log.logger import logger
|
||||
|
||||
from langflow.api.utils.kb_helpers import KBAnalysisHelper, KBStorageHelper
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
from langchain_chroma import Chroma
|
||||
|
||||
from langflow.services.database.models.message.model import MessageTable
|
||||
|
||||
# Chunk size for splitting long messages before embedding
|
||||
MESSAGE_CHUNK_SIZE = 1000
|
||||
MESSAGE_CHUNK_OVERLAP = 100
|
||||
|
||||
|
||||
def extract_content_block_text(content_blocks: list) -> str:
|
||||
"""Extract embeddable text from content blocks of type text, code, and json.
|
||||
|
||||
Blocks of any other type (tool_use, error, media, etc.) are skipped.
|
||||
Each extracted piece is separated by a blank line so chunk boundaries
|
||||
remain readable in the vector store.
|
||||
"""
|
||||
parts: list[str] = []
|
||||
for block in content_blocks:
|
||||
# content_blocks are stored as JSON; each block is a dict at runtime.
|
||||
contents: list = block.get("contents", []) if isinstance(block, dict) else []
|
||||
for entry in contents:
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
entry_type = entry.get("type")
|
||||
if entry_type == "text":
|
||||
fragment = (entry.get("text") or "").strip()
|
||||
elif entry_type == "code":
|
||||
lang = entry.get("language") or ""
|
||||
code = (entry.get("code") or "").strip()
|
||||
fragment = f"```{lang}\n{code}\n```" if code else ""
|
||||
elif entry_type == "json":
|
||||
data = entry.get("data")
|
||||
fragment = json.dumps(data, ensure_ascii=False) if data is not None else ""
|
||||
else:
|
||||
continue
|
||||
if fragment:
|
||||
parts.append(fragment)
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
def build_documents_from_messages(
|
||||
messages: list[MessageTable],
|
||||
*,
|
||||
session_id: str,
|
||||
flow_id: str,
|
||||
job_id: str = "",
|
||||
) -> list[Document]:
|
||||
"""Convert MessageTable rows into LangChain Documents.
|
||||
|
||||
Each message's embeddable text is the concatenation of msg.text and any
|
||||
content-block fragments whose type is text, code, or json. Other block
|
||||
types (tool_use, error, media, ...) are ignored. Long combined texts are
|
||||
split by RecursiveCharacterTextSplitter before embedding.
|
||||
"""
|
||||
splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=MESSAGE_CHUNK_SIZE,
|
||||
chunk_overlap=MESSAGE_CHUNK_OVERLAP,
|
||||
)
|
||||
docs: list[Document] = []
|
||||
for msg in messages:
|
||||
parts: list[str] = []
|
||||
if msg.text and msg.text.strip():
|
||||
parts.append(msg.text.strip())
|
||||
cb_text = extract_content_block_text(msg.content_blocks or [])
|
||||
if cb_text:
|
||||
parts.append(cb_text)
|
||||
|
||||
text = "\n\n".join(parts)
|
||||
if not text:
|
||||
continue
|
||||
chunks = splitter.split_text(text)
|
||||
for i, chunk in enumerate(chunks):
|
||||
docs.append(
|
||||
Document(
|
||||
page_content=chunk,
|
||||
metadata={
|
||||
"message_id": str(msg.id),
|
||||
"session_id": session_id,
|
||||
"flow_id": flow_id,
|
||||
"sender": msg.sender,
|
||||
"sender_name": msg.sender_name,
|
||||
"timestamp": msg.timestamp.isoformat() if msg.timestamp else "",
|
||||
"run_id": str(msg.run_id) if msg.run_id else "",
|
||||
"chunk_index": i,
|
||||
"total_chunks": len(chunks),
|
||||
"source": f"memory_base/{session_id}",
|
||||
"job_id": job_id,
|
||||
},
|
||||
)
|
||||
)
|
||||
return docs
|
||||
|
||||
|
||||
def sync_kb_metadata(*, kb_path: Path, chroma: Chroma) -> None:
|
||||
"""Update embedding_metadata.json after a successful Memory Base ingestion.
|
||||
|
||||
Mirrors the post-write metadata sync in ``KBIngestionHelper.perform_ingestion``:
|
||||
- Refreshes chunk / word / character counts from the live Chroma collection.
|
||||
- Updates on-disk size.
|
||||
- Stamps ``is_memory_base: true`` (required for Knowledge Retrieval filtering).
|
||||
- Sets ``source_types: ["memory"]`` to distinguish from file-based KBs.
|
||||
|
||||
Called while the Chroma client is still open so that ``update_text_metrics``
|
||||
can query the collection directly without opening a second client.
|
||||
"""
|
||||
try:
|
||||
metadata = KBAnalysisHelper.get_metadata(kb_path, fast=True)
|
||||
KBAnalysisHelper.update_text_metrics(kb_path, metadata, chroma=chroma)
|
||||
metadata["size"] = KBStorageHelper.get_directory_size(kb_path)
|
||||
metadata["is_memory_base"] = True
|
||||
# Preserve any existing source_types but always include "memory"
|
||||
existing = set(metadata.get("source_types") or [])
|
||||
existing.add("memory")
|
||||
metadata["source_types"] = sorted(existing)
|
||||
(kb_path / "embedding_metadata.json").write_text(json.dumps(metadata, indent=2))
|
||||
except (OSError, json.JSONDecodeError, ValueError):
|
||||
# Metadata sync is best-effort; a failure here must not block the cursor advance.
|
||||
# Note: this runs inside asyncio.to_thread so we use sync logging here.
|
||||
# The lfx logger's sync .warning() method goes through the same structured pipeline.
|
||||
logger.warning("KB metadata sync failed for kb_path=%s", kb_path, exc_info=True)
|
||||
@ -0,0 +1,26 @@
|
||||
"""Embedding provider inference for MemoryBase.
|
||||
|
||||
Extracted from MemoryBaseService to keep single-responsibility per file.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# Provider inference map — mirrors provider_patterns in KBAnalysisHelper._detect_embedding_provider
|
||||
# so we can derive the provider from a model name string without filesystem access.
|
||||
_MODEL_TO_PROVIDER: list[tuple[list[str], str]] = [
|
||||
(["text-embedding", "ada-", "gpt-"], "OpenAI"),
|
||||
(["embed-english", "embed-multilingual"], "Cohere"),
|
||||
(["sentence-transformers", "bert-", "huggingface"], "HuggingFace"),
|
||||
(["palm", "gecko", "google"], "Google"),
|
||||
(["ollama"], "Ollama"),
|
||||
(["azure"], "Azure OpenAI"),
|
||||
]
|
||||
|
||||
|
||||
def infer_embedding_provider(embedding_model: str) -> str:
|
||||
"""Derive embedding provider name from a model string."""
|
||||
lower = embedding_model.lower()
|
||||
for patterns, provider in _MODEL_TO_PROVIDER:
|
||||
if any(p in lower for p in patterns):
|
||||
return provider
|
||||
return "OpenAI" # Safe default — matches _resolve_embedding fallback
|
||||
14
src/backend/base/langflow/services/memory_base/factory.py
Normal file
14
src/backend/base/langflow/services/memory_base/factory.py
Normal file
@ -0,0 +1,14 @@
|
||||
"""Factory for creating MemoryBaseService instances."""
|
||||
|
||||
from langflow.services.factory import ServiceFactory
|
||||
from langflow.services.memory_base.service import MemoryBaseService
|
||||
|
||||
|
||||
class MemoryBaseServiceFactory(ServiceFactory):
|
||||
"""Factory for creating MemoryBaseService instances."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(MemoryBaseService)
|
||||
|
||||
def create(self):
|
||||
return MemoryBaseService()
|
||||
421
src/backend/base/langflow/services/memory_base/ingestion.py
Normal file
421
src/backend/base/langflow/services/memory_base/ingestion.py
Normal file
@ -0,0 +1,421 @@
|
||||
"""Ingestion orchestration for MemoryBase — auto-capture, regeneration, and mismatch detection.
|
||||
|
||||
Extracted from MemoryBaseService to keep single-responsibility per file.
|
||||
The MemoryBaseService delegates to these functions for all ingestion-related work.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from lfx.log.logger import logger
|
||||
from sqlmodel import col, func, select
|
||||
|
||||
from langflow.api.utils.kb_helpers import KBAnalysisHelper, KBStorageHelper
|
||||
from langflow.services.database.models.jobs.model import Job, JobStatus, JobType
|
||||
from langflow.services.database.models.memory_base.model import (
|
||||
MemoryBase,
|
||||
MemoryBaseSession,
|
||||
MemoryBaseWorkflowRun,
|
||||
)
|
||||
from langflow.services.deps import get_job_service, get_task_service, session_scope
|
||||
from langflow.services.jobs import DuplicateJobError
|
||||
from langflow.services.memory_base.kb_path_helpers import (
|
||||
hash_session_id,
|
||||
resolve_embedding,
|
||||
resolve_kb_username,
|
||||
resolve_kb_username_by_user_id,
|
||||
validate_kb_path,
|
||||
)
|
||||
from langflow.services.memory_base.task import IngestionRequest, ingest_memory_task
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
|
||||
async def trigger_ingestion(
|
||||
memory_base_id: uuid.UUID,
|
||||
user_id: uuid.UUID,
|
||||
session_id: str,
|
||||
*,
|
||||
get_mb_or_raise,
|
||||
get_or_create_session,
|
||||
) -> str:
|
||||
"""Manually trigger (or auto-trigger) an ingestion sync.
|
||||
|
||||
Returns:
|
||||
job_id string for the newly created job.
|
||||
|
||||
Raises:
|
||||
ValueError: If MemoryBase not found.
|
||||
RuntimeError: If a job is already active (caller should return 409).
|
||||
"""
|
||||
async with session_scope() as db:
|
||||
mb = await get_mb_or_raise(db, memory_base_id, user_id)
|
||||
|
||||
# Ensure a session record exists
|
||||
mbs = await get_or_create_session(db, memory_base_id, session_id)
|
||||
|
||||
# Snapshot the cursor NOW (immutable arg for the task)
|
||||
cursor_id_snapshot = mbs.cursor_id
|
||||
|
||||
# Build dedupe_key from the latest uncovered WORKFLOW run for idempotency.
|
||||
latest_job_id = await _get_latest_pending_workflow_job_id(db, mb, mbs)
|
||||
dedupe_key: str | None = None
|
||||
if latest_job_id is not None:
|
||||
dedupe_key = f"ingestion:{memory_base_id}:{session_id}:{latest_job_id}"
|
||||
|
||||
kb_username = await resolve_kb_username(db, mb.user_id)
|
||||
embedding_provider, embedding_model = resolve_embedding(mb.kb_name, kb_username)
|
||||
|
||||
# Create tracking job
|
||||
job_service = get_job_service()
|
||||
job_id = uuid.uuid4()
|
||||
await job_service.create_job(
|
||||
job_id=job_id,
|
||||
flow_id=mb.flow_id,
|
||||
user_id=mb.user_id,
|
||||
job_type=JobType.INGESTION,
|
||||
asset_id=memory_base_id,
|
||||
asset_type="memory_base",
|
||||
dedupe_key=dedupe_key,
|
||||
)
|
||||
|
||||
task_service = get_task_service()
|
||||
await task_service.fire_and_forget_task(
|
||||
job_service.execute_with_status,
|
||||
job_id=job_id,
|
||||
run_coro_func=ingest_memory_task,
|
||||
request=IngestionRequest(
|
||||
memory_base_id=memory_base_id,
|
||||
session_id=session_id,
|
||||
flow_id=mb.flow_id,
|
||||
kb_name=mb.kb_name,
|
||||
kb_username=kb_username,
|
||||
user_id=mb.user_id,
|
||||
embedding_provider=embedding_provider,
|
||||
embedding_model=embedding_model,
|
||||
cursor_id=cursor_id_snapshot,
|
||||
task_job_id=job_id,
|
||||
job_service=job_service,
|
||||
),
|
||||
)
|
||||
|
||||
return str(job_id)
|
||||
|
||||
|
||||
async def on_flow_output(
|
||||
flow_id: uuid.UUID,
|
||||
session_id: str,
|
||||
job_id: uuid.UUID | None,
|
||||
*,
|
||||
get_or_create_session,
|
||||
) -> None:
|
||||
"""Called after a flow run completes.
|
||||
|
||||
For every MemoryBase watching this flow with auto_capture=True:
|
||||
1. Record the workflow run in the tracking table (inside _maybe_trigger).
|
||||
2. Count uncovered WORKFLOW runs for this session.
|
||||
3. If count >= threshold, fire ingestion task.
|
||||
"""
|
||||
async with session_scope() as db:
|
||||
stmt = (
|
||||
select(MemoryBase).where(MemoryBase.flow_id == flow_id).where(MemoryBase.auto_capture == True) # noqa: E712
|
||||
)
|
||||
result = await db.exec(stmt)
|
||||
memory_bases = list(result.all())
|
||||
|
||||
hashed_sid = hash_session_id(session_id)
|
||||
for mb in memory_bases:
|
||||
try:
|
||||
await logger.adebug(
|
||||
"Auto-capture check | memory_base=%s threshold=%s session=%s",
|
||||
mb.id,
|
||||
mb.threshold,
|
||||
hashed_sid,
|
||||
)
|
||||
await _maybe_trigger(
|
||||
mb=mb, session_id=session_id, job_id=job_id, get_or_create_session=get_or_create_session
|
||||
)
|
||||
except (RuntimeError, ValueError, OSError):
|
||||
await logger.aerror("Auto-capture failed for memory_base=%s session=%s", mb.id, hashed_sid, exc_info=True)
|
||||
|
||||
|
||||
async def _maybe_trigger(
|
||||
*,
|
||||
mb: MemoryBase,
|
||||
session_id: str,
|
||||
job_id: uuid.UUID | None,
|
||||
get_or_create_session,
|
||||
) -> None:
|
||||
async with session_scope() as db:
|
||||
mbs = await get_or_create_session(db, mb.id, session_id)
|
||||
|
||||
# Record this workflow run before evaluating the threshold.
|
||||
await _insert_workflow_run(db, mb.id, session_id, job_id)
|
||||
|
||||
pending = await count_pending_messages(db, mb, mbs)
|
||||
|
||||
if pending < mb.threshold:
|
||||
return
|
||||
|
||||
cursor_id_snapshot = mbs.cursor_id
|
||||
|
||||
# Build dedupe_key from the latest pending WORKFLOW run for idempotency.
|
||||
latest_wf_job_id = await _get_latest_pending_workflow_job_id(db, mb, mbs)
|
||||
dedupe_key: str | None = None
|
||||
if latest_wf_job_id is not None:
|
||||
dedupe_key = f"ingestion:{mb.id}:{session_id}:{latest_wf_job_id}"
|
||||
|
||||
kb_username = await resolve_kb_username(db, mb.user_id)
|
||||
|
||||
embedding_provider, embedding_model = resolve_embedding(mb.kb_name, kb_username)
|
||||
|
||||
job_service = get_job_service()
|
||||
job_id = uuid.uuid4()
|
||||
try:
|
||||
await job_service.create_job(
|
||||
job_id=job_id,
|
||||
flow_id=mb.flow_id,
|
||||
user_id=mb.user_id,
|
||||
job_type=JobType.INGESTION,
|
||||
asset_id=mb.id,
|
||||
asset_type="memory_base",
|
||||
dedupe_key=dedupe_key,
|
||||
)
|
||||
except DuplicateJobError:
|
||||
await logger.adebug("Auto-capture: duplicate job for dedupe_key=%s - skipping.", dedupe_key)
|
||||
return
|
||||
|
||||
task_service = get_task_service()
|
||||
await task_service.fire_and_forget_task(
|
||||
job_service.execute_with_status,
|
||||
job_id=job_id,
|
||||
run_coro_func=ingest_memory_task,
|
||||
request=IngestionRequest(
|
||||
memory_base_id=mb.id,
|
||||
session_id=session_id,
|
||||
flow_id=mb.flow_id,
|
||||
kb_name=mb.kb_name,
|
||||
kb_username=kb_username,
|
||||
user_id=mb.user_id,
|
||||
embedding_provider=embedding_provider,
|
||||
embedding_model=embedding_model,
|
||||
cursor_id=cursor_id_snapshot,
|
||||
task_job_id=job_id,
|
||||
job_service=job_service,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def check_mismatch(
|
||||
memory_base_id: uuid.UUID,
|
||||
user_id: uuid.UUID,
|
||||
*,
|
||||
get_mb_or_raise,
|
||||
) -> bool:
|
||||
"""Return True if metadata claims processed rows but vector store is empty."""
|
||||
async with session_scope() as db:
|
||||
mb = await get_mb_or_raise(db, memory_base_id, user_id)
|
||||
stmt = select(func.sum(MemoryBaseSession.total_processed)).where(
|
||||
MemoryBaseSession.memory_base_id == memory_base_id
|
||||
)
|
||||
result = await db.exec(stmt)
|
||||
total_processed: int = result.first() or 0
|
||||
|
||||
if total_processed == 0:
|
||||
return False
|
||||
|
||||
kb_username = await resolve_kb_username_by_user_id(user_id)
|
||||
kb_root = KBStorageHelper.get_root_path()
|
||||
if not kb_root:
|
||||
return False
|
||||
kb_path = kb_root / kb_username / mb.kb_name
|
||||
validate_kb_path(kb_root, kb_path)
|
||||
if not await asyncio.to_thread(kb_path.exists):
|
||||
return True
|
||||
|
||||
metadata = KBAnalysisHelper.get_metadata(kb_path, fast=True)
|
||||
return int(metadata.get("chunks", 0)) == 0
|
||||
|
||||
|
||||
async def regenerate(
|
||||
memory_base_id: uuid.UUID,
|
||||
user_id: uuid.UUID,
|
||||
*,
|
||||
get_mb_or_raise,
|
||||
trigger_ingestion_fn,
|
||||
) -> list[str]:
|
||||
"""Reset all session cursors to None and re-trigger ingestion per session.
|
||||
|
||||
Used to recover from FS / Vector DB mismatch (Chroma dir deleted externally).
|
||||
Returns list of newly created job IDs.
|
||||
Also deletes all MessageIngestionRecord rows for this memory base atomically
|
||||
with the cursor reset so that re-ingestion starts clean without hitting the
|
||||
unique constraint.
|
||||
"""
|
||||
from sqlalchemy import delete as sa_delete
|
||||
|
||||
from langflow.services.database.models.memory_base.model import MessageIngestionRecord
|
||||
|
||||
async with session_scope() as db:
|
||||
await get_mb_or_raise(db, memory_base_id, user_id)
|
||||
|
||||
stmt = select(MemoryBaseSession).where(MemoryBaseSession.memory_base_id == memory_base_id)
|
||||
result = await db.exec(stmt)
|
||||
sessions = list(result.all())
|
||||
|
||||
for s in sessions:
|
||||
s.cursor_id = None
|
||||
db.add(s)
|
||||
|
||||
# Delete existing ingestion records so re-ingestion inserts fresh rows
|
||||
await db.exec( # type: ignore[call-overload]
|
||||
sa_delete(MessageIngestionRecord).where(MessageIngestionRecord.memory_base_id == memory_base_id)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
job_ids: list[str] = []
|
||||
for s in sessions:
|
||||
try:
|
||||
jid = await trigger_ingestion_fn(memory_base_id, user_id, s.session_id)
|
||||
job_ids.append(jid)
|
||||
except DuplicateJobError:
|
||||
await logger.awarning(
|
||||
"Regenerate: duplicate batch already ingested for session %s - skipped.",
|
||||
hash_session_id(s.session_id),
|
||||
)
|
||||
except RuntimeError:
|
||||
await logger.awarning(
|
||||
"Regenerate: active job exists for session %s - reset cursor but skipped trigger.",
|
||||
hash_session_id(s.session_id),
|
||||
)
|
||||
return job_ids
|
||||
|
||||
|
||||
async def cancel_active_jobs(*, memory_base_id: uuid.UUID, db: AsyncSession) -> None:
|
||||
"""Cancel all IN_PROGRESS or QUEUED jobs for this memory base."""
|
||||
stmt = (
|
||||
select(Job)
|
||||
.where(Job.asset_id == memory_base_id)
|
||||
.where(Job.asset_type == "memory_base")
|
||||
.where(col(Job.status).in_([JobStatus.IN_PROGRESS, JobStatus.QUEUED]))
|
||||
)
|
||||
result = await db.exec(stmt)
|
||||
active_jobs = list(result.all())
|
||||
|
||||
task_service = get_task_service()
|
||||
job_service = get_job_service()
|
||||
for job in active_jobs:
|
||||
try:
|
||||
await task_service.revoke_task(job.job_id)
|
||||
await job_service.update_job_status(job.job_id, JobStatus.CANCELLED)
|
||||
await logger.ainfo("Cancelled job %s for memory_base %s", job.job_id, memory_base_id)
|
||||
except (RuntimeError, ValueError, OSError):
|
||||
await logger.awarning(
|
||||
"Could not cancel job %s for memory_base %s", job.job_id, memory_base_id, exc_info=True
|
||||
)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Shared query helpers (public — used by service.py and memories.py) #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
|
||||
async def count_pending_messages(db: AsyncSession, mb: MemoryBase, mbs: MemoryBaseSession) -> int:
|
||||
"""Count WORKFLOW runs for this (memory_base, session) not yet covered by a completed ingestion.
|
||||
|
||||
A row in memory_base_workflow_run with ingestion_job_id IS NULL means the run
|
||||
has not been processed by any ingestion job. Count pending = number of such rows.
|
||||
This is session-scoped and time-independent; job failures leave rows NULL so they
|
||||
are correctly re-counted on the next threshold check.
|
||||
"""
|
||||
stmt = (
|
||||
select(func.count())
|
||||
.select_from(MemoryBaseWorkflowRun)
|
||||
.where(MemoryBaseWorkflowRun.memory_base_id == mb.id)
|
||||
.where(MemoryBaseWorkflowRun.session_id == mbs.session_id)
|
||||
.where(MemoryBaseWorkflowRun.ingestion_job_id == None) # noqa: E711
|
||||
)
|
||||
try:
|
||||
result = await db.exec(stmt)
|
||||
row = result.first()
|
||||
if row is None:
|
||||
return 0
|
||||
return int(row)
|
||||
except (TypeError, ValueError, OSError) as e:
|
||||
await logger.aerror("Error counting pending workflow runs: %s", e)
|
||||
return 0
|
||||
|
||||
|
||||
async def _insert_workflow_run(
|
||||
db: AsyncSession,
|
||||
memory_base_id: uuid.UUID,
|
||||
session_id: str,
|
||||
job_id: uuid.UUID | None,
|
||||
) -> None:
|
||||
"""Record a WORKFLOW job run for (memory_base_id, session_id).
|
||||
|
||||
Verifies that job_id refers to a WORKFLOW type job before inserting.
|
||||
Uses dialect-specific INSERT ... ON CONFLICT DO NOTHING for idempotency —
|
||||
safe to call multiple times with the same arguments.
|
||||
Skips silently if job_id is None or the job is not of WORKFLOW type.
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
hashed_sid = hash_session_id(session_id)
|
||||
if job_id is None:
|
||||
await logger.awarning(
|
||||
"on_flow_output called with no job_id for memory_base=%s session=%s — run not recorded.",
|
||||
memory_base_id,
|
||||
hashed_sid,
|
||||
)
|
||||
return
|
||||
|
||||
job_result = await db.exec(select(Job).where(Job.job_id == job_id).where(Job.type == JobType.WORKFLOW))
|
||||
if job_result.first() is None:
|
||||
await logger.awarning(
|
||||
"job_id=%s is not a WORKFLOW job — skipping workflow run record for memory_base=%s session=%s.",
|
||||
job_id,
|
||||
memory_base_id,
|
||||
hashed_sid,
|
||||
)
|
||||
return
|
||||
|
||||
row = {
|
||||
"id": uuid.uuid4(),
|
||||
"memory_base_id": memory_base_id,
|
||||
"session_id": session_id,
|
||||
"workflow_job_id": job_id,
|
||||
"ingestion_job_id": None,
|
||||
"recorded_at": datetime.now(timezone.utc),
|
||||
}
|
||||
conn = await db.connection()
|
||||
if conn.dialect.name == "postgresql":
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
|
||||
stmt = pg_insert(MemoryBaseWorkflowRun).values([row]).on_conflict_do_nothing()
|
||||
else:
|
||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||
|
||||
stmt = sqlite_insert(MemoryBaseWorkflowRun).values([row]).on_conflict_do_nothing()
|
||||
await db.exec(stmt) # type: ignore[call-overload]
|
||||
await db.commit()
|
||||
|
||||
|
||||
async def _get_latest_pending_workflow_job_id(
|
||||
db: AsyncSession, mb: MemoryBase, mbs: MemoryBaseSession
|
||||
) -> uuid.UUID | None:
|
||||
"""Return the workflow_job_id of the most recent uncovered workflow run for this session."""
|
||||
stmt = (
|
||||
select(MemoryBaseWorkflowRun.workflow_job_id)
|
||||
.where(MemoryBaseWorkflowRun.memory_base_id == mb.id)
|
||||
.where(MemoryBaseWorkflowRun.session_id == mbs.session_id)
|
||||
.where(MemoryBaseWorkflowRun.ingestion_job_id == None) # noqa: E711
|
||||
.order_by(col(MemoryBaseWorkflowRun.recorded_at).desc())
|
||||
.limit(1)
|
||||
)
|
||||
result = await db.exec(stmt)
|
||||
return result.first()
|
||||
@ -0,0 +1,149 @@
|
||||
"""KB path resolution and username helpers for MemoryBase.
|
||||
|
||||
Extracted from MemoryBaseService to keep single-responsibility per file.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from lfx.log.logger import logger
|
||||
from sqlmodel import select
|
||||
|
||||
from langflow.api.utils.kb_helpers import KBAnalysisHelper, KBStorageHelper
|
||||
from langflow.services.deps import session_scope
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
|
||||
def validate_kb_path(kb_root: Path, kb_path: Path) -> None:
|
||||
"""Assert that kb_path is contained within kb_root (path traversal guard).
|
||||
|
||||
Prevents crafted usernames with '..' segments from escaping the KB root directory.
|
||||
Follows the same pattern as services/storage/local.py:save_file.
|
||||
"""
|
||||
kb_root_resolved = kb_root.resolve()
|
||||
kb_path_resolved = kb_path.resolve()
|
||||
if not kb_path_resolved.is_relative_to(kb_root_resolved):
|
||||
msg = "KB path escapes root directory"
|
||||
raise ValueError(msg)
|
||||
|
||||
|
||||
def hash_session_id(session_id: str) -> str:
|
||||
"""Return a truncated SHA-256 hash for safe logging of session IDs."""
|
||||
return hashlib.sha256(session_id.encode()).hexdigest()[:12]
|
||||
|
||||
|
||||
def sanitize_kb_name(name: str) -> str:
|
||||
"""Lowercase, replace spaces/hyphens with underscores, strip non-alphanum."""
|
||||
sanitized = name.strip().lower()
|
||||
sanitized = re.sub(r"[\s\-]+", "_", sanitized)
|
||||
sanitized = re.sub(r"[^\w]", "", sanitized)
|
||||
return sanitized or "memory"
|
||||
|
||||
|
||||
async def resolve_kb_username(db: AsyncSession, user_id: uuid.UUID) -> str:
|
||||
"""Look up the username for a user_id within an existing DB session."""
|
||||
from langflow.services.database.models.user.model import User
|
||||
|
||||
stmt = select(User.username).where(User.id == user_id)
|
||||
result = await db.exec(stmt)
|
||||
username = result.first()
|
||||
if not username:
|
||||
msg = f"User {user_id} not found"
|
||||
raise ValueError(msg)
|
||||
return username
|
||||
|
||||
|
||||
async def resolve_kb_username_by_user_id(user_id: uuid.UUID) -> str:
|
||||
"""Look up the username for a user_id using a fresh DB session."""
|
||||
async with session_scope() as db:
|
||||
return await resolve_kb_username(db, user_id)
|
||||
|
||||
|
||||
def resolve_embedding(kb_name: str, kb_username: str) -> tuple[str, str]:
|
||||
"""Read embedding provider/model from KB metadata.json, with sane defaults."""
|
||||
kb_root = KBStorageHelper.get_root_path()
|
||||
if not kb_root:
|
||||
return "OpenAI", "text-embedding-3-small"
|
||||
kb_path: Path = kb_root / kb_username / kb_name
|
||||
metadata = KBAnalysisHelper.get_metadata(kb_path, fast=True)
|
||||
provider = metadata.get("embedding_provider") or "OpenAI"
|
||||
model = metadata.get("embedding_model") or "text-embedding-3-small"
|
||||
return provider, model
|
||||
|
||||
|
||||
async def initialize_kb(
|
||||
*,
|
||||
kb_name: str,
|
||||
kb_username: str,
|
||||
embedding_provider: str,
|
||||
embedding_model: str,
|
||||
) -> None:
|
||||
"""Create KB directory, initialize Chroma, and write embedding_metadata.json.
|
||||
|
||||
Mirrors the logic in knowledge_bases.py:create_knowledge_base so Memory Base
|
||||
KBs are immediately visible with the correct metadata (including is_memory_base: true).
|
||||
"""
|
||||
import chromadb
|
||||
|
||||
kb_root = KBStorageHelper.get_root_path()
|
||||
if not kb_root:
|
||||
await logger.awarning("KB root path not configured — Memory Base KB will not be initialized on disk.")
|
||||
return
|
||||
|
||||
kb_path: Path = kb_root / kb_username / kb_name
|
||||
validate_kb_path(kb_root, kb_path)
|
||||
await asyncio.to_thread(kb_path.mkdir, parents=True, exist_ok=True)
|
||||
|
||||
# Initialize Chroma collection so the directory is non-empty and readable
|
||||
try:
|
||||
client = KBStorageHelper.get_fresh_chroma_client(kb_path)
|
||||
client.create_collection(name=kb_name)
|
||||
except (OSError, ValueError, chromadb.errors.ChromaError) as exc:
|
||||
await logger.awarning("Initial Chroma setup for %s failed: %s", kb_name, exc)
|
||||
finally:
|
||||
client = None # type: ignore[assignment]
|
||||
KBStorageHelper.release_chroma_resources(kb_path)
|
||||
|
||||
embedding_metadata = {
|
||||
"id": str(uuid.uuid4()),
|
||||
"embedding_provider": embedding_provider,
|
||||
"embedding_model": embedding_model,
|
||||
"is_memory_base": True,
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
"chunks": 0,
|
||||
"words": 0,
|
||||
"characters": 0,
|
||||
"avg_chunk_size": 0.0,
|
||||
"size": 0,
|
||||
"source_types": ["memory"],
|
||||
}
|
||||
await asyncio.to_thread(
|
||||
(kb_path / "embedding_metadata.json").write_text,
|
||||
json.dumps(embedding_metadata, indent=2),
|
||||
)
|
||||
|
||||
|
||||
async def delete_kb(*, kb_name: str, kb_username: str) -> None:
|
||||
"""Remove the KB directory from disk. Logs on failure, does not raise."""
|
||||
if not kb_name:
|
||||
return
|
||||
kb_root = KBStorageHelper.get_root_path()
|
||||
if not kb_root:
|
||||
return
|
||||
kb_path = kb_root / kb_username / kb_name
|
||||
validate_kb_path(kb_root, kb_path)
|
||||
try:
|
||||
await asyncio.to_thread(KBStorageHelper.delete_storage, kb_path, kb_name)
|
||||
except (OSError, ValueError):
|
||||
await logger.awarning("Could not delete KB '%s' from disk after Memory Base deletion.", kb_name, exc_info=True)
|
||||
288
src/backend/base/langflow/services/memory_base/service.py
Normal file
288
src/backend/base/langflow/services/memory_base/service.py
Normal file
@ -0,0 +1,288 @@
|
||||
"""MemoryBase service — CRUD and session state management.
|
||||
|
||||
Ingestion orchestration, KB path helpers, and embedding inference are in
|
||||
separate modules (ingestion.py, kb_path_helpers.py, embedding_helpers.py)
|
||||
to keep this file focused on data access and business-rule enforcement.
|
||||
|
||||
Edge cases handled:
|
||||
- Name uniqueness per user: 409 if a Memory Base with the same name already exists.
|
||||
- Deletion during sync: cancels active tasks before DB deletion.
|
||||
- KB deletion on delete: removes the associated KB directory from disk.
|
||||
- Concurrent task prevention: returns 409 if a job is already IN_PROGRESS.
|
||||
- Threshold updates: deferred; does not re-evaluate pending count immediately.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlmodel import col, select
|
||||
|
||||
from langflow.services.base import Service
|
||||
from langflow.services.database.models.memory_base.model import (
|
||||
MemoryBase,
|
||||
MemoryBaseCreate,
|
||||
MemoryBaseSession,
|
||||
MemoryBaseUpdate,
|
||||
)
|
||||
from langflow.services.deps import session_scope
|
||||
from langflow.services.memory_base.embedding_helpers import infer_embedding_provider
|
||||
from langflow.services.memory_base.ingestion import (
|
||||
cancel_active_jobs,
|
||||
)
|
||||
from langflow.services.memory_base.ingestion import (
|
||||
check_mismatch as _check_mismatch,
|
||||
)
|
||||
from langflow.services.memory_base.ingestion import (
|
||||
on_flow_output as _on_flow_output,
|
||||
)
|
||||
from langflow.services.memory_base.ingestion import (
|
||||
regenerate as _regenerate,
|
||||
)
|
||||
from langflow.services.memory_base.ingestion import (
|
||||
trigger_ingestion as _trigger_ingestion,
|
||||
)
|
||||
from langflow.services.memory_base.kb_path_helpers import (
|
||||
delete_kb,
|
||||
initialize_kb,
|
||||
resolve_kb_username,
|
||||
sanitize_kb_name,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
|
||||
class MemoryBaseService(Service):
|
||||
"""Service layer for MemoryBase CRUD and session state management."""
|
||||
|
||||
name = "memory_base_service"
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# CRUD #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
async def create(self, payload: MemoryBaseCreate, user_id: uuid.UUID) -> MemoryBase:
|
||||
# 1. Verify that the referenced flow belongs to this user.
|
||||
async with session_scope() as db:
|
||||
from langflow.services.database.models.flow.model import Flow
|
||||
|
||||
flow_result = await db.exec(select(Flow).where(Flow.id == payload.flow_id).where(Flow.user_id == user_id))
|
||||
if flow_result.first() is None:
|
||||
msg = f"Flow {payload.flow_id} not found"
|
||||
raise PermissionError(msg)
|
||||
|
||||
# 2. Resolve username — needed for the KB path.
|
||||
async with session_scope() as db:
|
||||
kb_username = await resolve_kb_username(db, user_id)
|
||||
|
||||
# 3. Auto-generate kb_name: sanitized_name_<8hex>
|
||||
kb_name = f"{sanitize_kb_name(payload.name)}_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# 4. Create KB directory and embedding_metadata.json on disk.
|
||||
embedding_provider = infer_embedding_provider(payload.embedding_model)
|
||||
await initialize_kb(
|
||||
kb_name=kb_name,
|
||||
kb_username=kb_username,
|
||||
embedding_provider=embedding_provider,
|
||||
embedding_model=payload.embedding_model,
|
||||
)
|
||||
|
||||
# 5. Uniqueness check + insert.
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
async with session_scope() as db:
|
||||
existing = await db.exec(
|
||||
select(MemoryBase).where(MemoryBase.user_id == user_id).where(MemoryBase.name == payload.name)
|
||||
)
|
||||
if existing.first() is not None:
|
||||
msg = f"A Memory Base named '{payload.name}' already exists for this user"
|
||||
raise ValueError(msg)
|
||||
|
||||
mb = MemoryBase(
|
||||
**payload.model_dump(exclude={"user_id"}),
|
||||
user_id=user_id,
|
||||
kb_name=kb_name,
|
||||
)
|
||||
db.add(mb)
|
||||
try:
|
||||
await db.commit()
|
||||
except IntegrityError:
|
||||
msg = f"A Memory Base named '{payload.name}' already exists for this user"
|
||||
raise ValueError(msg) from None
|
||||
await db.refresh(mb)
|
||||
|
||||
return mb
|
||||
|
||||
async def list_for_user(self, user_id: uuid.UUID) -> list[MemoryBase]:
|
||||
async with session_scope() as db:
|
||||
stmt = select(MemoryBase).where(MemoryBase.user_id == user_id)
|
||||
result = await db.exec(stmt)
|
||||
return list(result.all())
|
||||
|
||||
def list_for_user_stmt(self, user_id: uuid.UUID, flow_id: uuid.UUID | None = None): # type: ignore[return]
|
||||
"""Return the SQLModel select statement for pagination at the API layer."""
|
||||
stmt = select(MemoryBase).where(MemoryBase.user_id == user_id)
|
||||
if flow_id is not None:
|
||||
stmt = stmt.where(MemoryBase.flow_id == flow_id)
|
||||
return stmt
|
||||
|
||||
async def get(self, memory_base_id: uuid.UUID, user_id: uuid.UUID) -> MemoryBase | None:
|
||||
async with session_scope() as db:
|
||||
stmt = select(MemoryBase).where(MemoryBase.id == memory_base_id).where(MemoryBase.user_id == user_id)
|
||||
result = await db.exec(stmt)
|
||||
return result.first()
|
||||
|
||||
async def update(
|
||||
self,
|
||||
memory_base_id: uuid.UUID,
|
||||
user_id: uuid.UUID,
|
||||
patch: MemoryBaseUpdate,
|
||||
) -> MemoryBase | None:
|
||||
"""Update mutable fields.
|
||||
|
||||
Threshold changes take effect on the NEXT auto-capture trigger; any
|
||||
already-running ingestion task ignores the change (immutable args).
|
||||
"""
|
||||
async with session_scope() as db:
|
||||
stmt = select(MemoryBase).where(MemoryBase.id == memory_base_id).where(MemoryBase.user_id == user_id)
|
||||
result = await db.exec(stmt)
|
||||
mb = result.first()
|
||||
if mb is None:
|
||||
return None
|
||||
for field, value in patch.model_dump(exclude_unset=True).items():
|
||||
setattr(mb, field, value)
|
||||
db.add(mb)
|
||||
await db.commit()
|
||||
await db.refresh(mb)
|
||||
return mb
|
||||
|
||||
async def delete(self, memory_base_id: uuid.UUID, user_id: uuid.UUID) -> bool:
|
||||
"""Delete a MemoryBase and its associated KB directory."""
|
||||
async with session_scope() as db:
|
||||
stmt = select(MemoryBase).where(MemoryBase.id == memory_base_id).where(MemoryBase.user_id == user_id)
|
||||
result = await db.exec(stmt)
|
||||
mb = result.first()
|
||||
if mb is None:
|
||||
return False
|
||||
|
||||
kb_name = mb.kb_name
|
||||
kb_username = await resolve_kb_username(db, user_id)
|
||||
|
||||
# Cancel active ingestion jobs before removing the DB record
|
||||
await cancel_active_jobs(memory_base_id=memory_base_id, db=db)
|
||||
|
||||
await db.delete(mb)
|
||||
await db.commit()
|
||||
|
||||
# Delete the corresponding KB from disk (best-effort — DB already committed)
|
||||
await delete_kb(kb_name=kb_name, kb_username=kb_username)
|
||||
|
||||
return True
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Sessions #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
async def verify_ownership(self, memory_base_id: uuid.UUID, user_id: uuid.UUID) -> None:
|
||||
"""Raise ValueError if the Memory Base does not belong to user_id."""
|
||||
async with session_scope() as db:
|
||||
await self.get_memory_base_or_404(db, memory_base_id, user_id)
|
||||
|
||||
def sessions_stmt(self, memory_base_id: uuid.UUID, user_id: uuid.UUID): # type: ignore[return]
|
||||
"""Return the select statement for persisted sessions, for use with apaginate.
|
||||
|
||||
Inline-joins MemoryBase to verify ownership in the SQL itself, so a
|
||||
caller that forgets a pre-check cannot leak other users' sessions.
|
||||
"""
|
||||
return (
|
||||
select(MemoryBaseSession)
|
||||
.join(MemoryBase, MemoryBase.id == MemoryBaseSession.memory_base_id)
|
||||
.where(MemoryBaseSession.memory_base_id == memory_base_id)
|
||||
.where(MemoryBase.user_id == user_id)
|
||||
.order_by(col(MemoryBaseSession.last_sync_at).desc())
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Ingestion delegation #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
async def trigger_ingestion(
|
||||
self,
|
||||
memory_base_id: uuid.UUID,
|
||||
user_id: uuid.UUID,
|
||||
session_id: str,
|
||||
) -> str:
|
||||
return await _trigger_ingestion(
|
||||
memory_base_id,
|
||||
user_id,
|
||||
session_id,
|
||||
get_mb_or_raise=self.get_memory_base_or_404,
|
||||
get_or_create_session=self._get_or_create_session,
|
||||
)
|
||||
|
||||
async def on_flow_output(
|
||||
self,
|
||||
flow_id: uuid.UUID,
|
||||
session_id: str,
|
||||
job_id: uuid.UUID | None,
|
||||
) -> None:
|
||||
await _on_flow_output(
|
||||
flow_id,
|
||||
session_id,
|
||||
job_id,
|
||||
get_or_create_session=self._get_or_create_session,
|
||||
)
|
||||
|
||||
async def check_mismatch(self, memory_base_id: uuid.UUID, user_id: uuid.UUID) -> bool:
|
||||
return await _check_mismatch(
|
||||
memory_base_id,
|
||||
user_id,
|
||||
get_mb_or_raise=self.get_memory_base_or_404,
|
||||
)
|
||||
|
||||
async def regenerate(self, memory_base_id: uuid.UUID, user_id: uuid.UUID) -> list[str]:
|
||||
return await _regenerate(
|
||||
memory_base_id,
|
||||
user_id,
|
||||
get_mb_or_raise=self.get_memory_base_or_404,
|
||||
trigger_ingestion_fn=self.trigger_ingestion,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Public query helpers #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
async def get_memory_base_or_404(
|
||||
self, db: AsyncSession, memory_base_id: uuid.UUID, user_id: uuid.UUID
|
||||
) -> MemoryBase:
|
||||
"""Fetch a MemoryBase or raise ValueError (mapped to 404 at the API layer)."""
|
||||
stmt = select(MemoryBase).where(MemoryBase.id == memory_base_id).where(MemoryBase.user_id == user_id)
|
||||
result = await db.exec(stmt)
|
||||
mb = result.first()
|
||||
if mb is None:
|
||||
msg = f"MemoryBase {memory_base_id} not found"
|
||||
raise ValueError(msg)
|
||||
return mb
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Internal helpers #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
async def _get_or_create_session(
|
||||
self, db: AsyncSession, memory_base_id: uuid.UUID, session_id: str
|
||||
) -> MemoryBaseSession:
|
||||
stmt = (
|
||||
select(MemoryBaseSession)
|
||||
.where(MemoryBaseSession.memory_base_id == memory_base_id)
|
||||
.where(MemoryBaseSession.session_id == session_id)
|
||||
)
|
||||
result = await db.exec(stmt)
|
||||
mbs = result.first()
|
||||
if mbs is None:
|
||||
mbs = MemoryBaseSession(memory_base_id=memory_base_id, session_id=session_id)
|
||||
db.add(mbs)
|
||||
await db.commit()
|
||||
await db.refresh(mbs)
|
||||
return mbs
|
||||
444
src/backend/base/langflow/services/memory_base/task.py
Normal file
444
src/backend/base/langflow/services/memory_base/task.py
Normal file
@ -0,0 +1,444 @@
|
||||
"""Background task for Memory Base ingestion.
|
||||
|
||||
Design principles enforced here:
|
||||
- Cursor atomicity: cursor_id is NEVER updated before ingestion confirms success.
|
||||
- Retry safety: If a job fails, cursor_id remains at the last known good position.
|
||||
- Serialization: A per-(memory_base_id, session_id) distributed lock prevents concurrent
|
||||
jobs from racing to write the same messages into Chroma. Uses PostgreSQL advisory locks
|
||||
for cross-worker safety, with an in-process asyncio.Lock fallback for SQLite (dev/test).
|
||||
The lock is acquired before any DB or Chroma access and released in a finally block.
|
||||
- Live cursor: After acquiring the lock, the current cursor_id is re-read from the DB
|
||||
(not the dispatch-time snapshot) so the pending message fetch always starts from the
|
||||
true latest position, even if a prior job advanced the cursor while this job waited.
|
||||
- Path safety: kb_path is validated against kb_root before any filesystem operation.
|
||||
|
||||
The actual Chroma write logic is shared with KB file ingestion via
|
||||
``KBIngestionHelper.write_documents_to_chroma`` — no duplicate batching/retry code here.
|
||||
|
||||
Document building and KB metadata sync live in ``document_builders.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import types
|
||||
import weakref
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from langchain_chroma import Chroma
|
||||
from lfx.log.logger import logger
|
||||
from sqlalchemy import text
|
||||
from sqlmodel import Session, col, select
|
||||
|
||||
from langflow.api.utils.kb_helpers import KBIngestionHelper, KBStorageHelper
|
||||
from langflow.services.database.models.memory_base.model import MemoryBaseSession, MemoryBaseWorkflowRun
|
||||
from langflow.services.database.models.message.model import MessageTable
|
||||
from langflow.services.deps import get_settings_service, session_scope
|
||||
from langflow.services.memory_base.document_builders import build_documents_from_messages, sync_kb_metadata
|
||||
from langflow.services.memory_base.kb_path_helpers import hash_session_id, validate_kb_path
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
from langflow.services.jobs.service import JobService
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class IngestionRequest:
|
||||
"""Typed parameter bundle for ``ingest_memory_task``.
|
||||
|
||||
All fields needed to run an ingestion job are grouped here so callers
|
||||
construct one object instead of threading 11+ loose kwargs.
|
||||
"""
|
||||
|
||||
memory_base_id: uuid.UUID
|
||||
session_id: str
|
||||
flow_id: uuid.UUID
|
||||
kb_name: str
|
||||
kb_username: str
|
||||
user_id: uuid.UUID
|
||||
embedding_provider: str
|
||||
embedding_model: str
|
||||
cursor_id: uuid.UUID | None
|
||||
task_job_id: uuid.UUID
|
||||
job_service: JobService
|
||||
|
||||
|
||||
# The ingestion lock timeout is read from settings (max_ingestion_timeout_secs).
|
||||
# If the timeout expires before the lock is acquired, an asyncio.TimeoutError is raised.
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Distributed locking: PostgreSQL advisory locks with in-process fallback
|
||||
# ---------------------------------------------------------------------------
|
||||
# In multi-worker deployments, an asyncio.Lock is process-local and cannot
|
||||
# serialize across workers. We use PostgreSQL session-level advisory locks
|
||||
# keyed on a hash of (memory_base_id, session_id). For SQLite (dev/test) we
|
||||
# fall back to the in-process asyncio.Lock which is sufficient for a single worker.
|
||||
|
||||
_session_ingestion_locks: weakref.WeakValueDictionary[tuple, asyncio.Lock] = weakref.WeakValueDictionary()
|
||||
|
||||
|
||||
def _get_or_create_session_lock(key: tuple) -> asyncio.Lock:
|
||||
"""Return the asyncio.Lock for the given key (SQLite fallback only)."""
|
||||
lock = _session_ingestion_locks.get(key)
|
||||
if lock is None:
|
||||
lock = asyncio.Lock()
|
||||
_session_ingestion_locks[key] = lock
|
||||
return lock
|
||||
|
||||
|
||||
def _compute_advisory_key(memory_base_id: uuid.UUID, session_id: str) -> int:
|
||||
"""Compute a stable int64 advisory lock key from (memory_base_id, session_id)."""
|
||||
raw = f"{memory_base_id}:{session_id}".encode()
|
||||
return int(hashlib.sha256(raw).hexdigest()[:16], 16) % (2**63 - 1)
|
||||
|
||||
|
||||
async def _is_postgres() -> bool:
|
||||
"""Return True if the database backend is PostgreSQL."""
|
||||
from langflow.services.deps import get_db_service
|
||||
|
||||
db_service = get_db_service()
|
||||
return db_service.engine.dialect.name == "postgresql"
|
||||
|
||||
|
||||
async def _pg_advisory_lock(db: Session, key: int) -> None:
|
||||
"""Acquire a PostgreSQL session-level advisory lock with retry and timeout.
|
||||
|
||||
The lock is held on the specific connection of the shared 'db' session.
|
||||
"""
|
||||
timeout = get_settings_service().settings.max_ingestion_timeout_secs
|
||||
deadline = asyncio.get_event_loop().time() + timeout
|
||||
backoff = 0.1
|
||||
max_backoff = 5.0
|
||||
|
||||
while True:
|
||||
conn = await db.connection()
|
||||
result = await conn.execute(text(f"SELECT pg_try_advisory_lock({key})"))
|
||||
acquired = result.scalar()
|
||||
|
||||
if acquired:
|
||||
return
|
||||
|
||||
remaining = deadline - asyncio.get_event_loop().time()
|
||||
if remaining <= 0:
|
||||
raise asyncio.TimeoutError
|
||||
|
||||
await asyncio.sleep(min(backoff, remaining))
|
||||
backoff = min(backoff * 2, max_backoff)
|
||||
|
||||
|
||||
async def _pg_advisory_unlock(db: Session, key: int) -> None:
|
||||
"""Release a PostgreSQL session-level advisory lock on the shared session."""
|
||||
conn = await db.connection()
|
||||
await conn.execute(text(f"SELECT pg_advisory_unlock({key})"))
|
||||
|
||||
|
||||
async def _acquire_session_lock(db: Session, memory_base_id: uuid.UUID, session_id: str) -> int | asyncio.Lock:
|
||||
"""Acquire the distributed ingestion lock. Returns the key (PG) or Lock (SQLite)."""
|
||||
timeout = get_settings_service().settings.max_ingestion_timeout_secs
|
||||
if await _is_postgres():
|
||||
key = _compute_advisory_key(memory_base_id, session_id)
|
||||
await _pg_advisory_lock(db, key)
|
||||
return key
|
||||
lock = _get_or_create_session_lock((memory_base_id, session_id))
|
||||
await asyncio.wait_for(lock.acquire(), timeout=timeout)
|
||||
return lock
|
||||
|
||||
|
||||
async def _release_session_lock(db: Session, lock_handle: int | asyncio.Lock) -> None:
|
||||
"""Release the distributed ingestion lock."""
|
||||
if isinstance(lock_handle, int):
|
||||
await _pg_advisory_unlock(db, lock_handle)
|
||||
else:
|
||||
lock_handle.release()
|
||||
|
||||
|
||||
async def _read_live_cursor(db: Session, memory_base_id: uuid.UUID, session_id: str) -> uuid.UUID | None:
|
||||
"""Read current cursor_id from shared 'db' session inside the serialization lock."""
|
||||
stmt = (
|
||||
select(MemoryBaseSession.cursor_id)
|
||||
.where(MemoryBaseSession.memory_base_id == memory_base_id)
|
||||
.where(MemoryBaseSession.session_id == session_id)
|
||||
)
|
||||
result = await db.exec(stmt)
|
||||
return result.first()
|
||||
|
||||
|
||||
async def ingest_memory_task(*, request: IngestionRequest) -> dict:
|
||||
"""Ingest pending output messages from a session into the target Knowledge Base.
|
||||
|
||||
Accepts a single ``IngestionRequest`` dataclass that bundles all required parameters.
|
||||
|
||||
Serialization: acquires a per-(memory_base_id, session_id) distributed lock before
|
||||
any DB or Chroma access. Uses PostgreSQL advisory locks for cross-worker
|
||||
serialization (multi-worker safe) with an in-process asyncio.Lock fallback for
|
||||
SQLite. Concurrent jobs for the same session wait up to max_ingestion_timeout_secs;
|
||||
if the lock cannot be acquired in time, asyncio.TimeoutError is re-raised so
|
||||
execute_with_status records JobStatus.TIMED_OUT.
|
||||
|
||||
Live cursor: after acquiring the lock, the current cursor_id is re-read from the DB.
|
||||
``cursor_id`` on the request is the dispatch-time snapshot kept only for logging.
|
||||
"""
|
||||
# Unpack for readability within the function body
|
||||
memory_base_id = request.memory_base_id
|
||||
session_id = request.session_id
|
||||
flow_id = request.flow_id
|
||||
kb_name = request.kb_name
|
||||
kb_username = request.kb_username
|
||||
user_id = request.user_id
|
||||
embedding_provider = request.embedding_provider
|
||||
embedding_model = request.embedding_model
|
||||
cursor_id = request.cursor_id
|
||||
task_job_id = request.task_job_id
|
||||
job_service = request.job_service
|
||||
|
||||
hashed_sid = hash_session_id(session_id)
|
||||
await logger.adebug(
|
||||
"Ingestion job started | memory_base=%s session=%s dispatch_cursor=%s job=%s",
|
||||
memory_base_id,
|
||||
hashed_sid,
|
||||
cursor_id,
|
||||
task_job_id,
|
||||
)
|
||||
kb_root = KBStorageHelper.get_root_path()
|
||||
if not kb_root:
|
||||
msg = "Knowledge base root path is not configured"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
kb_path: Path = kb_root / kb_username / kb_name
|
||||
|
||||
# ---- Path traversal guard ----
|
||||
validate_kb_path(kb_root, kb_path)
|
||||
|
||||
# ---- 0. Acquire per-session serialization lock ----
|
||||
async with session_scope() as db:
|
||||
try:
|
||||
lock_handle = await _acquire_session_lock(db, memory_base_id, session_id)
|
||||
except asyncio.TimeoutError:
|
||||
await logger.awarning(
|
||||
"Ingestion lock wait timeout | memory_base=%s session=%s job=%s.",
|
||||
memory_base_id,
|
||||
hashed_sid,
|
||||
task_job_id,
|
||||
)
|
||||
raise
|
||||
|
||||
try:
|
||||
# ---- 0b. Re-read live cursor inside the lock ----
|
||||
live_cursor_id = await _read_live_cursor(db, memory_base_id, session_id)
|
||||
await logger.adebug(
|
||||
"Ingestion lock acquired | memory_base=%s session=%s live_cursor=%s job=%s",
|
||||
memory_base_id,
|
||||
hashed_sid,
|
||||
live_cursor_id,
|
||||
task_job_id,
|
||||
)
|
||||
|
||||
# ---- 1. Fetch pending output messages for this session ----
|
||||
messages = await _fetch_pending_messages(
|
||||
db,
|
||||
flow_id=flow_id,
|
||||
session_id=session_id,
|
||||
cursor_id=live_cursor_id,
|
||||
)
|
||||
if not messages:
|
||||
await logger.ainfo(
|
||||
"MemoryBase %s / session %s: no pending messages, skipping.", memory_base_id, hashed_sid
|
||||
)
|
||||
return {"message": "No pending messages", "ingested": 0}
|
||||
|
||||
# ---- 2. Build documents from messages ----
|
||||
job_id_str = str(task_job_id)
|
||||
documents = build_documents_from_messages(
|
||||
messages, session_id=session_id, flow_id=str(flow_id), job_id=job_id_str
|
||||
)
|
||||
|
||||
if not documents:
|
||||
return {"message": "No non-empty messages to ingest", "ingested": 0}
|
||||
|
||||
# ---- 3. Check cancellation before touching the vector store ----
|
||||
if await KBIngestionHelper.is_job_cancelled(job_service, task_job_id):
|
||||
return {"message": "Job cancelled before ingestion", "ingested": 0}
|
||||
|
||||
# ---- 4. Open Chroma, write, then sync KB metadata ----
|
||||
user_stub = types.SimpleNamespace(id=user_id)
|
||||
embeddings = await KBIngestionHelper.build_embeddings(embedding_provider, embedding_model, user_stub)
|
||||
|
||||
client = KBStorageHelper.get_fresh_chroma_client(kb_path)
|
||||
written = 0
|
||||
try:
|
||||
chroma = Chroma(client=client, embedding_function=embeddings, collection_name=kb_name)
|
||||
|
||||
written = await KBIngestionHelper.write_documents_to_chroma(
|
||||
documents=documents,
|
||||
chroma=chroma,
|
||||
task_job_id=task_job_id,
|
||||
job_service=job_service,
|
||||
)
|
||||
|
||||
if written == len(documents):
|
||||
await asyncio.to_thread(sync_kb_metadata, kb_path=kb_path, chroma=chroma)
|
||||
except Exception:
|
||||
await logger.aerror(
|
||||
"Ingestion write failed | memory_base=%s session=%s job=%s. Rolling back partial writes...",
|
||||
memory_base_id,
|
||||
hashed_sid,
|
||||
task_job_id,
|
||||
)
|
||||
await KBIngestionHelper.cleanup_chroma_chunks_by_job(task_job_id, kb_path, kb_name)
|
||||
raise
|
||||
finally:
|
||||
KBStorageHelper.release_chroma_resources(kb_path)
|
||||
|
||||
if written < len(documents):
|
||||
await logger.awarning("Ingestion job %s was cancelled. Cleaning up partial data...", task_job_id)
|
||||
await KBIngestionHelper.cleanup_chroma_chunks_by_job(task_job_id, kb_path, kb_name)
|
||||
return {"message": "Job cancelled during ingestion", "ingested": 0}
|
||||
|
||||
# ---- 5. Bulk-stamp ingestion metadata ----
|
||||
await _mark_messages_ingested(db, messages=messages, job_id=task_job_id, memory_base_id=memory_base_id)
|
||||
|
||||
# ---- 6. Update cursor atomically ONLY after confirmed success ----
|
||||
last_message_id = messages[-1].id
|
||||
ingested_count = len(messages)
|
||||
await _advance_cursor(
|
||||
db,
|
||||
memory_base_id=memory_base_id,
|
||||
session_id=session_id,
|
||||
new_cursor_id=last_message_id,
|
||||
ingested_count=ingested_count,
|
||||
task_job_id=task_job_id,
|
||||
)
|
||||
|
||||
await logger.ainfo(
|
||||
"Ingestion job finished | memory_base=%s session=%s job=%s ingested=%d",
|
||||
memory_base_id,
|
||||
hashed_sid,
|
||||
task_job_id,
|
||||
ingested_count,
|
||||
)
|
||||
return {"message": "Success", "ingested": ingested_count}
|
||||
|
||||
finally:
|
||||
await _release_session_lock(db, lock_handle)
|
||||
|
||||
|
||||
async def _fetch_pending_messages(
|
||||
db: Session,
|
||||
*,
|
||||
flow_id: uuid.UUID,
|
||||
session_id: str,
|
||||
cursor_id: uuid.UUID | None,
|
||||
) -> list[MessageTable]:
|
||||
"""Fetch all messages for this session that come after cursor_id using shared session."""
|
||||
from sqlalchemy import and_, or_
|
||||
|
||||
stmt = (
|
||||
select(MessageTable)
|
||||
.where(MessageTable.flow_id == flow_id)
|
||||
.where(MessageTable.session_id == session_id)
|
||||
.order_by(col(MessageTable.timestamp).asc(), col(MessageTable.id).asc())
|
||||
)
|
||||
if cursor_id is not None:
|
||||
cursor_stmt = select(MessageTable.timestamp, MessageTable.id).where(MessageTable.id == cursor_id)
|
||||
result = await db.exec(cursor_stmt)
|
||||
cursor_row = result.first()
|
||||
if cursor_row:
|
||||
cursor_ts, c_id = cursor_row
|
||||
stmt = stmt.where(
|
||||
or_(
|
||||
col(MessageTable.timestamp) > cursor_ts,
|
||||
and_(
|
||||
col(MessageTable.timestamp) == cursor_ts,
|
||||
col(MessageTable.id) > c_id,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
result = await db.exec(stmt)
|
||||
return list(result.all())
|
||||
|
||||
|
||||
async def _mark_messages_ingested(
|
||||
db: Session,
|
||||
*,
|
||||
messages: list[MessageTable],
|
||||
job_id: uuid.UUID,
|
||||
memory_base_id: uuid.UUID,
|
||||
) -> None:
|
||||
"""Batch-insert ingestion records for all successfully ingested messages using shared session."""
|
||||
from uuid import uuid4 as _uuid4
|
||||
|
||||
from langflow.services.database.models.memory_base.model import MessageIngestionRecord
|
||||
|
||||
ingested_at = datetime.now(timezone.utc)
|
||||
rows = [
|
||||
{
|
||||
"id": _uuid4(),
|
||||
"message_id": msg.id,
|
||||
"memory_base_id": memory_base_id,
|
||||
"job_id": job_id,
|
||||
"session_id": msg.session_id,
|
||||
"ingested_at": ingested_at,
|
||||
}
|
||||
for msg in messages
|
||||
]
|
||||
conn = await db.connection()
|
||||
if conn.dialect.name == "postgresql":
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
|
||||
stmt = pg_insert(MessageIngestionRecord).values(rows).on_conflict_do_nothing()
|
||||
else:
|
||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||
|
||||
stmt = sqlite_insert(MessageIngestionRecord).values(rows).on_conflict_do_nothing()
|
||||
await db.exec(stmt) # type: ignore[call-overload]
|
||||
await db.commit()
|
||||
|
||||
|
||||
async def _advance_cursor(
|
||||
db: Session,
|
||||
*,
|
||||
memory_base_id: uuid.UUID,
|
||||
session_id: str,
|
||||
new_cursor_id: uuid.UUID,
|
||||
ingested_count: int,
|
||||
task_job_id: uuid.UUID,
|
||||
) -> None:
|
||||
"""Atomically advance the cursor using the shared 'db' session."""
|
||||
from sqlalchemy import update as sa_update
|
||||
|
||||
stmt = (
|
||||
select(MemoryBaseSession)
|
||||
.where(MemoryBaseSession.memory_base_id == memory_base_id)
|
||||
.where(MemoryBaseSession.session_id == session_id)
|
||||
)
|
||||
result = await db.exec(stmt)
|
||||
mbs = result.first()
|
||||
if mbs is None:
|
||||
await logger.awarning(
|
||||
"MemoryBaseSession for (%s, %s) vanished before cursor update.",
|
||||
memory_base_id,
|
||||
hash_session_id(session_id),
|
||||
)
|
||||
return
|
||||
|
||||
mbs.cursor_id = new_cursor_id
|
||||
mbs.total_processed += ingested_count
|
||||
mbs.last_sync_at = datetime.now(timezone.utc)
|
||||
db.add(mbs)
|
||||
|
||||
# Stamp all pending workflow run rows for this session.
|
||||
await db.exec( # type: ignore[call-overload]
|
||||
sa_update(MemoryBaseWorkflowRun)
|
||||
.where(MemoryBaseWorkflowRun.memory_base_id == memory_base_id)
|
||||
.where(MemoryBaseWorkflowRun.session_id == session_id)
|
||||
.where(MemoryBaseWorkflowRun.ingestion_job_id == None) # noqa: E711
|
||||
.values(ingestion_job_id=task_job_id)
|
||||
)
|
||||
|
||||
await db.commit()
|
||||
@ -22,3 +22,4 @@ class ServiceType(str, Enum):
|
||||
MCP_COMPOSER_SERVICE = "mcp_composer_service"
|
||||
JOB_SERVICE = "jobs_service"
|
||||
FLOW_EVENTS_SERVICE = "flow_events_service"
|
||||
MEMORY_BASE_SERVICE = "memory_base_service"
|
||||
|
||||
@ -94,6 +94,12 @@ async def setup_superuser(settings_service: SettingsService, session: AsyncSessi
|
||||
)
|
||||
if user is not None:
|
||||
await logger.adebug("Superuser created successfully.")
|
||||
# When the default superuser is recreated (e.g. after a DB reset in
|
||||
# AUTO_LOGIN mode) the per-user MCP servers config file saved under the
|
||||
# previous UUID becomes orphaned on disk. Best-effort recover it so
|
||||
# users don't lose their MCP server configuration across restarts.
|
||||
if is_default and settings_service.auth_settings.AUTO_LOGIN:
|
||||
await migrate_orphaned_mcp_servers_config(session, settings_service, user)
|
||||
except Exception as exc:
|
||||
logger.exception(exc)
|
||||
msg = "Could not create superuser. Please create a superuser manually."
|
||||
@ -103,6 +109,149 @@ async def setup_superuser(settings_service: SettingsService, session: AsyncSessi
|
||||
settings_service.auth_settings.reset_credentials()
|
||||
|
||||
|
||||
async def migrate_orphaned_mcp_servers_config(
|
||||
session: AsyncSession,
|
||||
settings_service: SettingsService,
|
||||
current_user,
|
||||
) -> bool:
|
||||
"""Best-effort recovery of MCP servers config files orphaned by a DB reset.
|
||||
|
||||
The MCP servers config is persisted on disk at
|
||||
``{config_dir}/{user_id}/_mcp_servers_{user_id}.json`` and tracked in the DB
|
||||
via a ``File`` row. When Langflow starts with a fresh database but the same
|
||||
config directory (common in containerized deployments without a persisted
|
||||
DB volume), the default superuser is recreated with a new UUID and the
|
||||
previously saved MCP config files become unreachable.
|
||||
|
||||
Recovery rules (intentionally conservative to avoid importing another user's
|
||||
config — MCP server entries can contain ``env`` and ``headers`` auth material):
|
||||
|
||||
* If the new user already has an MCP config file on disk without a matching
|
||||
``File`` row, re-register the row (self-heal a partial previous migration).
|
||||
* If exactly one orphaned ``_mcp_servers_{uuid}.json`` is found in the config
|
||||
directory, migrate it.
|
||||
* If multiple orphans are found, skip and log — we can't safely identify the
|
||||
previous default superuser's file without extra metadata, so leave manual
|
||||
recovery to the operator.
|
||||
|
||||
Returns True when a file was migrated or a missing DB row was restored.
|
||||
"""
|
||||
from pathlib import Path
|
||||
from uuid import UUID
|
||||
|
||||
import aiofiles
|
||||
import anyio
|
||||
|
||||
from langflow.services.database.models.file.model import File as UserFile
|
||||
|
||||
try:
|
||||
config_dir_value = settings_service.settings.config_dir
|
||||
if not config_dir_value:
|
||||
return False
|
||||
|
||||
config_dir = Path(config_dir_value)
|
||||
if not config_dir.exists() or not config_dir.is_dir():
|
||||
return False
|
||||
|
||||
# The current user's DB record is fresh; nothing to migrate if they
|
||||
# somehow already have an MCP config row (defensive guard).
|
||||
name_without_ext = f"_mcp_servers_{current_user.id}"
|
||||
existing_stmt = (
|
||||
select(UserFile).where(UserFile.user_id == current_user.id).where(UserFile.name == name_without_ext)
|
||||
)
|
||||
if (await session.exec(existing_stmt)).first() is not None:
|
||||
return False
|
||||
|
||||
current_user_dir = str(current_user.id)
|
||||
new_dir = config_dir / current_user_dir
|
||||
new_filename = f"_mcp_servers_{current_user.id}.json"
|
||||
new_file_path = new_dir / new_filename
|
||||
db_path = f"{current_user.id}/{new_filename}"
|
||||
|
||||
# Case 1: a previous migration attempt copied the file but failed before
|
||||
# committing the DB row. Re-register the existing file instead of
|
||||
# returning early and leaving the user with an invisible config.
|
||||
if new_file_path.exists():
|
||||
try:
|
||||
size = new_file_path.stat().st_size
|
||||
except OSError as exc:
|
||||
await logger.awarning(
|
||||
"Cannot stat existing MCP config %s while self-healing DB row: %s",
|
||||
new_file_path,
|
||||
exc,
|
||||
)
|
||||
return False
|
||||
session.add(UserFile(user_id=current_user.id, name=name_without_ext, path=db_path, size=size))
|
||||
await session.commit()
|
||||
await logger.ainfo(
|
||||
"Restored missing MCP servers config DB row for user %s from existing file %s",
|
||||
current_user.id,
|
||||
new_file_path,
|
||||
)
|
||||
return True
|
||||
|
||||
def _find_orphans() -> list[tuple[float, Path]]:
|
||||
orphans: list[tuple[float, Path]] = []
|
||||
for entry in config_dir.iterdir():
|
||||
if not entry.is_dir() or entry.name == current_user_dir:
|
||||
continue
|
||||
try:
|
||||
UUID(entry.name)
|
||||
except ValueError:
|
||||
continue
|
||||
mcp_path = entry / f"_mcp_servers_{entry.name}.json"
|
||||
if mcp_path.is_file():
|
||||
try:
|
||||
mtime = mcp_path.stat().st_mtime
|
||||
except OSError:
|
||||
continue
|
||||
orphans.append((mtime, mcp_path))
|
||||
return orphans
|
||||
|
||||
orphans = await anyio.to_thread.run_sync(_find_orphans)
|
||||
if not orphans:
|
||||
return False
|
||||
|
||||
# Ambiguous: more than one candidate could belong to different users.
|
||||
# Refuse to migrate rather than risk importing unrelated auth material.
|
||||
if len(orphans) > 1:
|
||||
orphan_paths = ", ".join(str(p) for _, p in sorted(orphans, key=lambda i: i[0], reverse=True))
|
||||
await logger.awarning(
|
||||
"Found %d orphaned MCP servers config files in %s; skipping automatic "
|
||||
"migration to avoid restoring the wrong one. Move the intended file to "
|
||||
"%s to recover. Candidates: %s",
|
||||
len(orphans),
|
||||
config_dir,
|
||||
new_file_path,
|
||||
orphan_paths,
|
||||
)
|
||||
return False
|
||||
|
||||
_, orphan_path = orphans[0]
|
||||
|
||||
async with aiofiles.open(str(orphan_path), "rb") as src:
|
||||
data = await src.read()
|
||||
|
||||
await anyio.to_thread.run_sync(lambda: new_dir.mkdir(parents=True, exist_ok=True))
|
||||
async with aiofiles.open(str(new_file_path), "wb") as dst:
|
||||
await dst.write(data)
|
||||
|
||||
session.add(UserFile(user_id=current_user.id, name=name_without_ext, path=db_path, size=len(data)))
|
||||
await session.commit()
|
||||
|
||||
await logger.ainfo(
|
||||
"Migrated orphaned MCP servers config from %s to user %s",
|
||||
orphan_path,
|
||||
current_user.id,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
# Never let migration failure block startup.
|
||||
await logger.awarning("Failed to migrate orphaned MCP servers config: %s", exc)
|
||||
return False
|
||||
else:
|
||||
return True
|
||||
|
||||
|
||||
async def teardown_superuser(settings_service, session: AsyncSession) -> None:
|
||||
"""Teardown the superuser."""
|
||||
# If AUTO_LOGIN is True, we will remove the default superuser
|
||||
|
||||
@ -39,6 +39,7 @@ class TestAgentComponent(ComponentTestBaseWithoutClient):
|
||||
return {
|
||||
"_type": "Agent",
|
||||
"add_current_date_tool": True,
|
||||
"add_calculator_tool": True,
|
||||
"agent_description": "A helpful agent",
|
||||
"model": MockLanguageModel(),
|
||||
"handle_parsing_errors": True,
|
||||
@ -450,6 +451,199 @@ class TestAgentComponent(ComponentTestBaseWithoutClient):
|
||||
# Note: The provider-specific field name mapping happens inside get_llm,
|
||||
# so we just verify max_tokens is passed correctly
|
||||
|
||||
async def test_should_append_calculator_tool_when_add_calculator_toggle_is_true(
|
||||
self, component_class, default_kwargs
|
||||
):
|
||||
"""Calculator tool is appended when the toggle is enabled.
|
||||
|
||||
Given add_calculator_tool=True, When get_agent_requirements runs,
|
||||
Then self.tools contains a StructuredTool derived from CalculatorComponent.
|
||||
"""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from langchain_core.tools import StructuredTool
|
||||
|
||||
default_kwargs["add_calculator_tool"] = True
|
||||
default_kwargs["add_current_date_tool"] = False # isolate: only calculator
|
||||
component = await self.component_setup(component_class, default_kwargs)
|
||||
component.model = [{"name": "gpt-4o", "provider": "OpenAI", "metadata": {}}]
|
||||
component.get_memory_data = AsyncMock(return_value=[])
|
||||
component._get_shared_callbacks = list
|
||||
component.set_tools_callbacks = lambda *_: None
|
||||
|
||||
with patch("lfx.components.models_and_agents.agent.get_llm") as mock_get_llm:
|
||||
mock_get_llm.return_value = MockLanguageModel()
|
||||
_, _, tools = await component.get_agent_requirements()
|
||||
|
||||
assert len(tools) == 1
|
||||
assert isinstance(tools[0], StructuredTool)
|
||||
assert "evaluate" in tools[0].name.lower(), f"Expected a Calculator-derived tool; got name={tools[0].name!r}"
|
||||
|
||||
async def test_should_not_append_calculator_tool_when_add_calculator_toggle_is_false(
|
||||
self, component_class, default_kwargs
|
||||
):
|
||||
"""Calculator tool is skipped when the toggle is disabled.
|
||||
|
||||
Given add_calculator_tool=False, When get_agent_requirements runs,
|
||||
Then no Calculator tool is appended to self.tools.
|
||||
"""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
default_kwargs["add_calculator_tool"] = False
|
||||
default_kwargs["add_current_date_tool"] = False
|
||||
component = await self.component_setup(component_class, default_kwargs)
|
||||
component.model = [{"name": "gpt-4o", "provider": "OpenAI", "metadata": {}}]
|
||||
component.get_memory_data = AsyncMock(return_value=[])
|
||||
component._get_shared_callbacks = list
|
||||
component.set_tools_callbacks = lambda *_: None
|
||||
|
||||
with patch("lfx.components.models_and_agents.agent.get_llm") as mock_get_llm:
|
||||
mock_get_llm.return_value = MockLanguageModel()
|
||||
_, _, tools = await component.get_agent_requirements()
|
||||
|
||||
assert tools == []
|
||||
|
||||
def test_should_replace_current_date_and_model_name_when_both_placeholders_present(self, component_class):
|
||||
"""Unit test: helper replaces both placeholders with concrete values."""
|
||||
component = component_class()
|
||||
component.model = [{"name": "gpt-4o", "provider": "OpenAI", "metadata": {}}]
|
||||
|
||||
prompt = "Today is {current_date}. You are powered by {model_name}."
|
||||
result = component._inject_dynamic_prompt_values(prompt)
|
||||
|
||||
assert "{current_date}" not in result
|
||||
assert "{model_name}" not in result
|
||||
assert "gpt-4o" in result
|
||||
|
||||
def test_should_leave_literal_braces_untouched_when_prompt_has_no_known_placeholders(self, component_class):
|
||||
"""Adversarial: prompts with literal JSON like {"key": 1} must not raise and must stay intact."""
|
||||
component = component_class()
|
||||
component.model = [{"name": "gpt-4o", "provider": "OpenAI", "metadata": {}}]
|
||||
|
||||
prompt = 'Respond with JSON: {"key": 1, "nested": {"a": [1, 2]}}.'
|
||||
result = component._inject_dynamic_prompt_values(prompt)
|
||||
|
||||
assert result == prompt
|
||||
|
||||
def test_should_return_empty_when_prompt_is_empty(self, component_class):
|
||||
"""Edge case: empty/None prompt is returned as-is without raising."""
|
||||
component = component_class()
|
||||
assert component._inject_dynamic_prompt_values("") == ""
|
||||
assert component._inject_dynamic_prompt_values(None) is None
|
||||
|
||||
async def test_should_inject_dynamic_values_into_system_prompt_when_message_response_runs(
|
||||
self, component_class, default_kwargs
|
||||
):
|
||||
"""Integration: message_response must call self.set with the resolved system_prompt."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
default_kwargs["system_prompt"] = "Powered by {model_name}."
|
||||
default_kwargs["add_calculator_tool"] = False
|
||||
default_kwargs["add_current_date_tool"] = False
|
||||
component = await self.component_setup(component_class, default_kwargs)
|
||||
component.model = [{"name": "gpt-4o", "provider": "OpenAI", "metadata": {}}]
|
||||
component.get_memory_data = AsyncMock(return_value=[])
|
||||
component._get_shared_callbacks = list
|
||||
component.set_tools_callbacks = lambda *_: None
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
def fake_set(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return component
|
||||
|
||||
component.set = fake_set
|
||||
component.create_agent_runnable = MagicMock(return_value=MagicMock())
|
||||
component.run_agent = AsyncMock(return_value=MagicMock())
|
||||
|
||||
with patch("lfx.components.models_and_agents.agent.get_llm") as mock_get_llm:
|
||||
mock_get_llm.return_value = MockLanguageModel()
|
||||
await component.message_response()
|
||||
|
||||
assert captured.get("system_prompt") == "Powered by gpt-4o."
|
||||
|
||||
async def test_should_not_mutate_format_instructions_when_json_response_runs(self, component_class, default_kwargs):
|
||||
"""Regression: injection must only touch agent_instructions, not format_instructions.
|
||||
|
||||
Ensures literal {current_date}/{model_name} tokens in user-authored
|
||||
format_instructions survive intact while the main system_prompt is
|
||||
still replaced by the helper.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
default_kwargs["system_prompt"] = "Powered by {model_name}."
|
||||
default_kwargs["format_instructions"] = "Return JSON with fields {current_date} and {model_name} preserved."
|
||||
default_kwargs["add_calculator_tool"] = False
|
||||
default_kwargs["add_current_date_tool"] = False
|
||||
component = await self.component_setup(component_class, default_kwargs)
|
||||
component.model = [{"name": "gpt-4o", "provider": "OpenAI", "metadata": {}}]
|
||||
component.get_memory_data = AsyncMock(return_value=[])
|
||||
component._get_shared_callbacks = list
|
||||
component.set_tools_callbacks = lambda *_: None
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
def fake_set(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return component
|
||||
|
||||
component.set = fake_set
|
||||
component.create_agent_runnable = MagicMock(return_value=MagicMock())
|
||||
component.run_agent = AsyncMock(return_value=MagicMock(content="{}"))
|
||||
|
||||
with patch("lfx.components.models_and_agents.agent.get_llm") as mock_get_llm:
|
||||
mock_get_llm.return_value = MockLanguageModel()
|
||||
await component.json_response()
|
||||
|
||||
prompt = captured.get("system_prompt") or ""
|
||||
assert "Powered by gpt-4o." in prompt, "agent_instructions should have placeholders replaced"
|
||||
assert "{current_date}" in prompt, "format_instructions literal braces must survive"
|
||||
assert "{model_name} preserved" in prompt, "format_instructions literal braces must survive"
|
||||
|
||||
async def test_should_accept_add_calculator_tool_in_default_keys(self, component_class, default_kwargs):
|
||||
"""update_build_config's default_keys validation must include add_calculator_tool."""
|
||||
from lfx.schema.dotdict import dotdict
|
||||
|
||||
with patch("lfx.components.models_and_agents.agent.get_language_model_options") as mock_opts:
|
||||
mock_opts.return_value = [
|
||||
{
|
||||
"name": "gpt-4o",
|
||||
"provider": "OpenAI",
|
||||
"icon": "OpenAI",
|
||||
"metadata": {
|
||||
"model_class": "ChatOpenAI",
|
||||
"model_name_param": "model",
|
||||
"api_key_param": "api_key",
|
||||
},
|
||||
}
|
||||
]
|
||||
component = await self.component_setup(component_class, default_kwargs)
|
||||
frontend_node = component.to_frontend_node()
|
||||
build_config = frontend_node["data"]["node"]["template"]
|
||||
|
||||
# add_calculator_tool must be present in the build_config already; if not,
|
||||
# update_build_config will error listing it as missing.
|
||||
assert "add_calculator_tool" in build_config
|
||||
|
||||
updated_config = await component.update_build_config(
|
||||
dotdict(build_config), mock_opts.return_value, field_name="model"
|
||||
)
|
||||
assert "add_calculator_tool" in updated_config
|
||||
|
||||
def test_should_have_placeholders_in_default_system_prompt(self, component_class):
|
||||
"""Default system_prompt ships with placeholders for the dynamic injection.
|
||||
|
||||
Ensures {current_date} and {model_name} are visible on a fresh agent so
|
||||
that the dynamic injection has an observable effect out-of-the-box.
|
||||
"""
|
||||
prompt_input = next(
|
||||
(inp for inp in component_class.inputs if getattr(inp, "name", None) == "system_prompt"),
|
||||
None,
|
||||
)
|
||||
assert prompt_input is not None
|
||||
assert "{current_date}" in prompt_input.value
|
||||
assert "{model_name}" in prompt_input.value
|
||||
|
||||
|
||||
class TestAgentComponentWithClient(ComponentTestBaseWithClient):
|
||||
@pytest.fixture
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from lfx.components.chroma import ChromaVectorStoreComponent
|
||||
@ -8,6 +9,37 @@ from lfx.schema.data import Data
|
||||
from tests.base import ComponentTestBaseWithoutClient, VersionComponentMapping
|
||||
|
||||
|
||||
def test_remote_chroma_server_uses_http_client() -> None:
|
||||
mock_client = MagicMock()
|
||||
mock_chroma = MagicMock()
|
||||
mock_chroma.get.return_value = {"ids": [], "documents": [], "metadatas": []}
|
||||
|
||||
with (
|
||||
patch("chromadb.HttpClient", return_value=mock_client) as mock_http_client,
|
||||
patch("langchain_chroma.Chroma", return_value=mock_chroma) as mock_chroma_class,
|
||||
):
|
||||
component = ChromaVectorStoreComponent().set(
|
||||
collection_name="remote_collection",
|
||||
persist_directory=None,
|
||||
embedding=None,
|
||||
chroma_server_host="chroma.example.com",
|
||||
chroma_server_http_port=8100,
|
||||
chroma_server_ssl_enabled=True,
|
||||
ingest_data=[],
|
||||
limit=None,
|
||||
)
|
||||
|
||||
assert component.build_vector_store() is mock_chroma
|
||||
|
||||
mock_http_client.assert_called_once_with(host="chroma.example.com", port=8100, ssl=True)
|
||||
mock_chroma_class.assert_called_once_with(
|
||||
persist_directory=None,
|
||||
client=mock_client,
|
||||
embedding_function=None,
|
||||
collection_name="remote_collection",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.api_key_required
|
||||
class TestChromaVectorStoreComponent(ComponentTestBaseWithoutClient):
|
||||
@pytest.fixture
|
||||
|
||||
@ -714,7 +714,7 @@ class TestPerformIngestionTask:
|
||||
|
||||
@patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client")
|
||||
@patch("langflow.api.utils.kb_helpers.Chroma")
|
||||
@patch("langflow.api.utils.kb_helpers.KBIngestionHelper._build_embeddings", new_callable=AsyncMock)
|
||||
@patch("langflow.api.utils.kb_helpers.KBIngestionHelper.build_embeddings", new_callable=AsyncMock)
|
||||
@patch("langflow.api.utils.kb_helpers.KBAnalysisHelper.get_metadata")
|
||||
@patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_directory_size")
|
||||
@patch("langflow.api.utils.kb_helpers.KBAnalysisHelper.update_text_metrics")
|
||||
@ -764,7 +764,7 @@ class TestPerformIngestionTask:
|
||||
|
||||
@patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client")
|
||||
@patch("langflow.api.utils.kb_helpers.Chroma")
|
||||
@patch("langflow.api.utils.kb_helpers.KBIngestionHelper._build_embeddings", new_callable=AsyncMock)
|
||||
@patch("langflow.api.utils.kb_helpers.KBIngestionHelper.build_embeddings", new_callable=AsyncMock)
|
||||
@patch("langflow.api.utils.kb_helpers.KBIngestionHelper.cleanup_chroma_chunks_by_job", new_callable=AsyncMock)
|
||||
async def test_perform_ingestion_rollback(
|
||||
self, mock_cleanup, mock_build, mock_chroma, mock_fresh_client, mock_kb_path
|
||||
|
||||
247
src/backend/tests/unit/test_mcp_servers_orphan_migration.py
Normal file
247
src/backend/tests/unit/test_mcp_servers_orphan_migration.py
Normal file
@ -0,0 +1,247 @@
|
||||
"""Tests for migrate_orphaned_mcp_servers_config in langflow.services.utils.
|
||||
|
||||
Verifies that MCP server config files written under a previous default
|
||||
superuser's UUID are picked up and migrated to the new default superuser
|
||||
when the database is reset but the config directory is preserved
|
||||
(typical of containerized deployments without a persisted DB volume).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import TYPE_CHECKING
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
from langflow.services.database.models.file.model import File as UserFile
|
||||
from langflow.services.deps import get_settings_service, session_scope
|
||||
from langflow.services.utils import migrate_orphaned_mcp_servers_config
|
||||
from sqlmodel import select
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def initialized_services(monkeypatch, tmp_path):
|
||||
"""Initialize DB + services with an isolated config dir."""
|
||||
from langflow.services.utils import initialize_services, teardown_services
|
||||
from lfx.services.manager import get_service_manager
|
||||
|
||||
db_path = tmp_path / "test.db"
|
||||
config_dir = tmp_path / "config"
|
||||
config_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
monkeypatch.setenv("LANGFLOW_DATABASE_URL", f"sqlite:///{db_path}")
|
||||
monkeypatch.setenv("LANGFLOW_CONFIG_DIR", str(config_dir))
|
||||
monkeypatch.setenv("LANGFLOW_AUTO_LOGIN", "true")
|
||||
|
||||
get_service_manager().factories.clear()
|
||||
get_service_manager().services.clear()
|
||||
|
||||
await initialize_services()
|
||||
|
||||
yield config_dir
|
||||
|
||||
await teardown_services()
|
||||
|
||||
|
||||
def _write_orphan(config_dir: Path, payload: dict, *, mtime_offset: float = 0.0) -> Path:
|
||||
"""Create an orphaned _mcp_servers_{uuid}.json file in a UUID-named folder."""
|
||||
orphan_id = uuid4()
|
||||
orphan_dir = config_dir / str(orphan_id)
|
||||
orphan_dir.mkdir(parents=True, exist_ok=True)
|
||||
orphan_path = orphan_dir / f"_mcp_servers_{orphan_id}.json"
|
||||
orphan_path.write_text(json.dumps(payload))
|
||||
if mtime_offset:
|
||||
stat = orphan_path.stat()
|
||||
os.utime(orphan_path, (stat.st_atime + mtime_offset, stat.st_mtime + mtime_offset))
|
||||
return orphan_path
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(30)
|
||||
async def test_migrate_orphaned_mcp_servers_config_recovers_previous_user_config(
|
||||
initialized_services,
|
||||
):
|
||||
"""A single orphaned file is migrated to the current default superuser."""
|
||||
config_dir: Path = initialized_services
|
||||
expected_payload = {"mcpServers": {"my-server": {"command": "uvx", "args": ["mcp-proxy"]}}}
|
||||
orphan_path = _write_orphan(config_dir, expected_payload)
|
||||
|
||||
settings = get_settings_service()
|
||||
|
||||
async with session_scope() as session:
|
||||
from langflow.services.database.models.user.model import User
|
||||
from lfx.services.settings.constants import DEFAULT_SUPERUSER
|
||||
|
||||
user = (await session.exec(select(User).where(User.username == DEFAULT_SUPERUSER))).first()
|
||||
assert user is not None, "default superuser should exist after initialize_services"
|
||||
|
||||
# Simulate the fresh-DB scenario: the user has no MCP config row yet.
|
||||
stmt = select(UserFile).where(UserFile.user_id == user.id).where(UserFile.name == f"_mcp_servers_{user.id}")
|
||||
assert (await session.exec(stmt)).first() is None
|
||||
|
||||
migrated = await migrate_orphaned_mcp_servers_config(session, settings, user)
|
||||
assert migrated is True
|
||||
|
||||
# DB record should exist and point at the new user-specific path.
|
||||
stmt = select(UserFile).where(UserFile.user_id == user.id).where(UserFile.name == f"_mcp_servers_{user.id}")
|
||||
new_record = (await session.exec(stmt)).first()
|
||||
assert new_record is not None
|
||||
assert new_record.path == f"{user.id}/_mcp_servers_{user.id}.json"
|
||||
|
||||
# File should live under the new user's folder with the same contents.
|
||||
target = config_dir / str(user.id) / f"_mcp_servers_{user.id}.json"
|
||||
assert target.exists()
|
||||
assert json.loads(target.read_text()) == expected_payload
|
||||
# Orphan source is left intact (best-effort copy, not destructive move).
|
||||
assert orphan_path.exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(30)
|
||||
async def test_migrate_orphaned_mcp_servers_config_skips_when_multiple_orphans(
|
||||
initialized_services,
|
||||
):
|
||||
"""With multiple orphan candidates we refuse to guess and leave state untouched.
|
||||
|
||||
MCP server entries can contain env/headers auth material, so importing an
|
||||
unrelated user's config would be a security hazard. Operators must resolve
|
||||
the ambiguity manually.
|
||||
"""
|
||||
config_dir: Path = initialized_services
|
||||
|
||||
old_payload = {"mcpServers": {"old": {}}}
|
||||
new_payload = {"mcpServers": {"new": {}}}
|
||||
|
||||
old_path = _write_orphan(config_dir, old_payload, mtime_offset=-3600)
|
||||
new_path = _write_orphan(config_dir, new_payload)
|
||||
|
||||
settings = get_settings_service()
|
||||
|
||||
async with session_scope() as session:
|
||||
from langflow.services.database.models.user.model import User
|
||||
from lfx.services.settings.constants import DEFAULT_SUPERUSER
|
||||
|
||||
user = (await session.exec(select(User).where(User.username == DEFAULT_SUPERUSER))).first()
|
||||
migrated = await migrate_orphaned_mcp_servers_config(session, settings, user)
|
||||
assert migrated is False
|
||||
|
||||
stmt = select(UserFile).where(UserFile.user_id == user.id).where(UserFile.name == f"_mcp_servers_{user.id}")
|
||||
assert (await session.exec(stmt)).first() is None
|
||||
|
||||
target = config_dir / str(user.id) / f"_mcp_servers_{user.id}.json"
|
||||
assert not target.exists()
|
||||
# Both orphans are preserved on disk for manual recovery.
|
||||
assert json.loads(old_path.read_text()) == old_payload
|
||||
assert json.loads(new_path.read_text()) == new_payload
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(30)
|
||||
async def test_migrate_orphaned_mcp_servers_config_no_orphans_is_noop(
|
||||
initialized_services,
|
||||
):
|
||||
"""With no orphaned files the function returns False and does not touch the DB."""
|
||||
config_dir: Path = initialized_services
|
||||
# No UUID-named subdirectories should exist.
|
||||
from uuid import UUID
|
||||
|
||||
def _is_uuid_dir(p: Path) -> bool:
|
||||
if not p.is_dir():
|
||||
return False
|
||||
try:
|
||||
UUID(p.name)
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
assert not [p for p in config_dir.iterdir() if _is_uuid_dir(p)]
|
||||
|
||||
settings = get_settings_service()
|
||||
|
||||
async with session_scope() as session:
|
||||
from langflow.services.database.models.user.model import User
|
||||
from lfx.services.settings.constants import DEFAULT_SUPERUSER
|
||||
|
||||
user = (await session.exec(select(User).where(User.username == DEFAULT_SUPERUSER))).first()
|
||||
migrated = await migrate_orphaned_mcp_servers_config(session, settings, user)
|
||||
assert migrated is False
|
||||
|
||||
stmt = select(UserFile).where(UserFile.user_id == user.id).where(UserFile.name == f"_mcp_servers_{user.id}")
|
||||
assert (await session.exec(stmt)).first() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(30)
|
||||
async def test_migrate_orphaned_mcp_servers_config_self_heals_missing_db_row(
|
||||
initialized_services,
|
||||
):
|
||||
"""Re-register missing DB rows when the on-disk file already exists.
|
||||
|
||||
If the user's file is on disk but the DB row is missing, the migration
|
||||
should recreate the row instead of leaving the config invisible. This
|
||||
covers recovery from a previous migration attempt that wrote the file
|
||||
but crashed before committing the DB row.
|
||||
"""
|
||||
config_dir: Path = initialized_services
|
||||
settings = get_settings_service()
|
||||
|
||||
async with session_scope() as session:
|
||||
from langflow.services.database.models.user.model import User
|
||||
from lfx.services.settings.constants import DEFAULT_SUPERUSER
|
||||
|
||||
user = (await session.exec(select(User).where(User.username == DEFAULT_SUPERUSER))).first()
|
||||
|
||||
# Simulate a file-exists / DB-row-missing state.
|
||||
existing_dir = config_dir / str(user.id)
|
||||
existing_dir.mkdir(parents=True, exist_ok=True)
|
||||
existing_payload = {"mcpServers": {"keep-me": {}}}
|
||||
existing_path = existing_dir / f"_mcp_servers_{user.id}.json"
|
||||
existing_path.write_text(json.dumps(existing_payload))
|
||||
|
||||
migrated = await migrate_orphaned_mcp_servers_config(session, settings, user)
|
||||
assert migrated is True
|
||||
|
||||
stmt = select(UserFile).where(UserFile.user_id == user.id).where(UserFile.name == f"_mcp_servers_{user.id}")
|
||||
row = (await session.exec(stmt)).first()
|
||||
assert row is not None
|
||||
assert row.path == f"{user.id}/_mcp_servers_{user.id}.json"
|
||||
assert row.size == existing_path.stat().st_size
|
||||
|
||||
# File contents are untouched.
|
||||
assert json.loads(existing_path.read_text()) == existing_payload
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(30)
|
||||
async def test_migrate_orphaned_mcp_servers_config_skips_when_row_already_present(
|
||||
initialized_services,
|
||||
):
|
||||
"""When the DB row already exists, do nothing — avoids double-registration."""
|
||||
config_dir: Path = initialized_services
|
||||
_write_orphan(config_dir, {"mcpServers": {"orphan": {}}})
|
||||
|
||||
settings = get_settings_service()
|
||||
|
||||
async with session_scope() as session:
|
||||
from langflow.services.database.models.user.model import User
|
||||
from lfx.services.settings.constants import DEFAULT_SUPERUSER
|
||||
|
||||
user = (await session.exec(select(User).where(User.username == DEFAULT_SUPERUSER))).first()
|
||||
|
||||
# Pre-register an MCP file row to simulate an already-migrated user.
|
||||
session.add(
|
||||
UserFile(
|
||||
user_id=user.id,
|
||||
name=f"_mcp_servers_{user.id}",
|
||||
path=f"{user.id}/_mcp_servers_{user.id}.json",
|
||||
size=0,
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
migrated = await migrate_orphaned_mcp_servers_config(session, settings, user)
|
||||
assert migrated is False
|
||||
1018
src/backend/tests/unit/test_memory_base_task.py
Normal file
1018
src/backend/tests/unit/test_memory_base_task.py
Normal file
File diff suppressed because it is too large
Load Diff
1719
src/backend/tests/unit/test_memory_bases.py
Normal file
1719
src/backend/tests/unit/test_memory_bases.py
Normal file
File diff suppressed because it is too large
Load Diff
@ -322,3 +322,15 @@ async def test_session_metadata_retrieval():
|
||||
messages = await aget_messages(sender="User", session_id=session_id)
|
||||
assert len(messages) == 1
|
||||
assert messages[0].session_metadata == initial_metadata
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("client")
|
||||
async def test_messageupdate_with_session_metadata(sample_session_metadata):
|
||||
"""Test MessageUpdate schema with session_metadata."""
|
||||
from langflow.services.database.models.message.model import MessageUpdate
|
||||
|
||||
message_update = MessageUpdate(
|
||||
text="Updated text",
|
||||
session_metadata=sample_session_metadata,
|
||||
)
|
||||
assert message_update.session_metadata == sample_session_metadata
|
||||
|
||||
@ -19,10 +19,8 @@ export const useDeleteProviderAccount: useMutationFunctionType<
|
||||
);
|
||||
};
|
||||
|
||||
// TODO: Add retries for transient server-side errors (5xx, timeouts).
|
||||
return mutate(["useDeleteProviderAccount"], fn, {
|
||||
...options,
|
||||
retry: false,
|
||||
onSuccess: (...args) => {
|
||||
queryClient.refetchQueries({ queryKey: ["useGetProviderAccounts"] });
|
||||
options?.onSuccess?.(...args);
|
||||
|
||||
@ -17,10 +17,8 @@ export const useDeleteDeployment: useMutationFunctionType<
|
||||
await api.delete(`${getURL("DEPLOYMENTS")}/${deployment_id}`);
|
||||
};
|
||||
|
||||
// TODO: Add retries for transient server-side errors (5xx, timeouts).
|
||||
return mutate(["useDeleteDeployment"], fn, {
|
||||
...options,
|
||||
retry: false,
|
||||
onSuccess: (...args) => {
|
||||
queryClient.refetchQueries({ queryKey: ["useGetDeployments"] });
|
||||
options?.onSuccess?.(...args);
|
||||
|
||||
@ -50,10 +50,8 @@ export const usePatchDeployment: useMutationFunctionType<
|
||||
return res.data;
|
||||
};
|
||||
|
||||
// TODO: Add retries for transient server-side errors (5xx, timeouts).
|
||||
return mutate(["usePatchDeployment"], fn, {
|
||||
...options,
|
||||
retry: false,
|
||||
onSuccess: (...args) => {
|
||||
queryClient.refetchQueries({ queryKey: ["useGetDeployments"] });
|
||||
queryClient.removeQueries({
|
||||
|
||||
@ -26,10 +26,8 @@ export const usePatchSnapshot: useMutationFunctionType<
|
||||
return data;
|
||||
};
|
||||
|
||||
// TODO: Add retries for transient server-side errors (5xx, timeouts).
|
||||
return mutate(["usePatchSnapshot"], fn, {
|
||||
...options,
|
||||
retry: false,
|
||||
onSuccess: (...args) => {
|
||||
queryClient.refetchQueries({ queryKey: ["useGetDeployments"] });
|
||||
options?.onSuccess?.(...args);
|
||||
|
||||
@ -47,6 +47,5 @@ export const usePostDeploymentRun: useMutationFunctionType<
|
||||
return res.data;
|
||||
};
|
||||
|
||||
// TODO: Add retries for transient server-side errors (5xx, timeouts).
|
||||
return mutate(["usePostDeploymentRun"], fn, { ...options, retry: false });
|
||||
return mutate(["usePostDeploymentRun"], fn, options);
|
||||
};
|
||||
|
||||
@ -53,10 +53,8 @@ export const usePostDeployment: useMutationFunctionType<
|
||||
return res.data;
|
||||
};
|
||||
|
||||
// TODO: Add retries for transient server-side errors (5xx, timeouts).
|
||||
return mutate(["usePostDeployment"], fn, {
|
||||
...options,
|
||||
retry: false,
|
||||
onSuccess: () => {
|
||||
return queryClient.refetchQueries({ queryKey: ["useGetDeployments"] });
|
||||
},
|
||||
|
||||
@ -0,0 +1,189 @@
|
||||
// biome-ignore-all lint/suspicious/noExplicitAny: test mocks
|
||||
const mockQueryClient = {
|
||||
invalidateQueries: jest.fn(),
|
||||
};
|
||||
|
||||
let mockCapturedQueryOptions: any = null;
|
||||
let mockCapturedMutationOptions: any = null;
|
||||
|
||||
jest.mock("@tanstack/react-query", () => ({
|
||||
useQueryClient: jest.fn(() => mockQueryClient),
|
||||
useQuery: jest.fn((options: any) => {
|
||||
mockCapturedQueryOptions = options;
|
||||
return { data: undefined, isLoading: false };
|
||||
}),
|
||||
useMutation: jest.fn((options: any) => {
|
||||
mockCapturedMutationOptions = options;
|
||||
return { mutate: jest.fn(), mutateAsync: jest.fn() };
|
||||
}),
|
||||
}));
|
||||
|
||||
import { UseRequestProcessor } from "../request-processor";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// `axios.isAxiosError` checks `error.isAxiosError === true`, so fixtures must
|
||||
// set that flag to be classified as HTTP errors rather than transient.
|
||||
const axiosError = (status: number) => ({
|
||||
isAxiosError: true,
|
||||
response: { status },
|
||||
});
|
||||
const axiosNetworkError = () => ({ isAxiosError: true, response: undefined });
|
||||
const nonAxiosError = () => new Error("Network Error");
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockCapturedQueryOptions = null;
|
||||
mockCapturedMutationOptions = null;
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Queries: retry policy
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe("UseRequestProcessor.query retry policy", () => {
|
||||
const setup = (options: any = {}) => {
|
||||
const { query } = UseRequestProcessor();
|
||||
query(["k"], async () => ({}), options);
|
||||
if (mockCapturedQueryOptions == null) {
|
||||
throw new Error("query was not called by UseRequestProcessor");
|
||||
}
|
||||
return mockCapturedQueryOptions;
|
||||
};
|
||||
|
||||
it("does not retry on any 4xx response", () => {
|
||||
const { retry } = setup();
|
||||
for (const status of [400, 401, 403, 404, 409, 422, 429, 499]) {
|
||||
expect(retry(0, axiosError(status))).toBe(false);
|
||||
}
|
||||
});
|
||||
|
||||
it("retries up to 5 times on 5xx responses", () => {
|
||||
const { retry } = setup();
|
||||
for (const status of [500, 502, 503, 504]) {
|
||||
expect(retry(0, axiosError(status))).toBe(true);
|
||||
expect(retry(4, axiosError(status))).toBe(true);
|
||||
expect(retry(5, axiosError(status))).toBe(false);
|
||||
}
|
||||
});
|
||||
|
||||
it("retries up to 5 times on axios errors with no response", () => {
|
||||
const { retry } = setup();
|
||||
expect(retry(0, axiosNetworkError())).toBe(true);
|
||||
expect(retry(4, axiosNetworkError())).toBe(true);
|
||||
expect(retry(5, axiosNetworkError())).toBe(false);
|
||||
});
|
||||
|
||||
it("treats non-axios / unknown errors as transient (retries)", () => {
|
||||
const { retry } = setup();
|
||||
expect(retry(0, undefined)).toBe(true);
|
||||
expect(retry(0, {})).toBe(true);
|
||||
expect(retry(0, { response: {} })).toBe(true);
|
||||
expect(retry(0, nonAxiosError())).toBe(true);
|
||||
});
|
||||
|
||||
it("allows per-call options.retry to override the default", () => {
|
||||
const { retry } = setup({ retry: false });
|
||||
expect(retry).toBe(false);
|
||||
});
|
||||
|
||||
it("allows per-call options.retry as a number to override the default", () => {
|
||||
const { retry } = setup({ retry: 10 });
|
||||
expect(retry).toBe(10);
|
||||
});
|
||||
|
||||
it("allows per-call options.retryDelay to override the default", () => {
|
||||
const customDelay = jest.fn(() => 42);
|
||||
const { retryDelay } = setup({ retryDelay: customDelay });
|
||||
expect(retryDelay).toBe(customDelay);
|
||||
});
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Mutations: retry policy
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe("UseRequestProcessor.mutate retry policy", () => {
|
||||
const setup = (options: any = {}) => {
|
||||
const { mutate } = UseRequestProcessor();
|
||||
mutate(["k"], async () => ({}), options);
|
||||
if (mockCapturedMutationOptions == null) {
|
||||
throw new Error("mutate was not called by UseRequestProcessor");
|
||||
}
|
||||
return mockCapturedMutationOptions;
|
||||
};
|
||||
|
||||
it("does not retry on any 4xx response", () => {
|
||||
const { retry } = setup();
|
||||
for (const status of [400, 401, 403, 404, 409, 422, 429, 499]) {
|
||||
expect(retry(0, axiosError(status))).toBe(false);
|
||||
}
|
||||
});
|
||||
|
||||
it("retries up to 3 times on 5xx responses", () => {
|
||||
const { retry } = setup();
|
||||
expect(retry(0, axiosError(500))).toBe(true);
|
||||
expect(retry(2, axiosError(500))).toBe(true);
|
||||
expect(retry(3, axiosError(500))).toBe(false);
|
||||
});
|
||||
|
||||
it("retries up to 3 times on axios errors with no response", () => {
|
||||
const { retry } = setup();
|
||||
expect(retry(0, axiosNetworkError())).toBe(true);
|
||||
expect(retry(2, axiosNetworkError())).toBe(true);
|
||||
expect(retry(3, axiosNetworkError())).toBe(false);
|
||||
});
|
||||
|
||||
it("respects options.retry === false", () => {
|
||||
const { retry } = setup({ retry: false });
|
||||
expect(retry).toBe(false);
|
||||
});
|
||||
|
||||
it("respects options.retry === 0 (nullish coalescing, not falsy)", () => {
|
||||
const { retry } = setup({ retry: 0 });
|
||||
expect(retry).toBe(0);
|
||||
});
|
||||
|
||||
it("respects a custom options.retry callback", () => {
|
||||
const customRetry = jest.fn(() => true);
|
||||
const { retry } = setup({ retry: customRetry });
|
||||
expect(retry).toBe(customRetry);
|
||||
});
|
||||
|
||||
it("respects a custom options.retryDelay", () => {
|
||||
const customDelay = jest.fn(() => 42);
|
||||
const { retryDelay } = setup({ retryDelay: customDelay });
|
||||
expect(retryDelay).toBe(customDelay);
|
||||
});
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// retryDelay: exponential backoff capped at 30s
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe("UseRequestProcessor retryDelay", () => {
|
||||
it("uses exponential backoff capped at 30s for queries", () => {
|
||||
const { query } = UseRequestProcessor();
|
||||
query(["k"], async () => ({}));
|
||||
const { retryDelay } = mockCapturedQueryOptions;
|
||||
expect(retryDelay(0)).toBe(1000);
|
||||
expect(retryDelay(1)).toBe(2000);
|
||||
expect(retryDelay(2)).toBe(4000);
|
||||
expect(retryDelay(3)).toBe(8000);
|
||||
expect(retryDelay(4)).toBe(16000);
|
||||
expect(retryDelay(5)).toBe(30000);
|
||||
expect(retryDelay(10)).toBe(30000);
|
||||
});
|
||||
|
||||
it("uses exponential backoff capped at 30s for mutations", () => {
|
||||
const { mutate } = UseRequestProcessor();
|
||||
mutate(["k"], async () => ({}));
|
||||
const { retryDelay } = mockCapturedMutationOptions;
|
||||
expect(retryDelay(0)).toBe(1000);
|
||||
expect(retryDelay(1)).toBe(2000);
|
||||
expect(retryDelay(2)).toBe(4000);
|
||||
expect(retryDelay(10)).toBe(30000);
|
||||
});
|
||||
});
|
||||
@ -6,11 +6,33 @@ import {
|
||||
useQuery,
|
||||
useQueryClient,
|
||||
} from "@tanstack/react-query";
|
||||
import axios from "axios";
|
||||
import type {
|
||||
MutationFunctionType,
|
||||
QueryFunctionType,
|
||||
} from "../../../types/api";
|
||||
|
||||
// 4xx responses are intentional client-side rejections (auth, validation,
|
||||
// deployment guards, etc.) and won't change on retry.
|
||||
function isClientError(error: unknown): boolean {
|
||||
if (!axios.isAxiosError(error)) return false;
|
||||
const status = error.response?.status;
|
||||
return typeof status === "number" && status >= 400 && status < 500;
|
||||
}
|
||||
|
||||
function makeRetry(maxRetries: number) {
|
||||
return (failureCount: number, error: unknown) => {
|
||||
if (isClientError(error)) return false;
|
||||
return failureCount < maxRetries;
|
||||
};
|
||||
}
|
||||
|
||||
const queryRetry = makeRetry(5);
|
||||
const mutationRetry = makeRetry(3);
|
||||
|
||||
const retryDelay = (attemptIndex: number) =>
|
||||
Math.min(1000 * 2 ** attemptIndex, 30000);
|
||||
|
||||
export function UseRequestProcessor(): {
|
||||
query: QueryFunctionType;
|
||||
mutate: MutationFunctionType;
|
||||
@ -26,8 +48,8 @@ export function UseRequestProcessor(): {
|
||||
return useQuery({
|
||||
queryKey,
|
||||
queryFn,
|
||||
retry: 5,
|
||||
retryDelay: (attemptIndex) => Math.min(1000 * 2 ** attemptIndex, 30000),
|
||||
retry: queryRetry,
|
||||
retryDelay,
|
||||
...options,
|
||||
});
|
||||
}
|
||||
@ -45,8 +67,8 @@ export function UseRequestProcessor(): {
|
||||
options.onSettled && options.onSettled(data, error, variables, context);
|
||||
},
|
||||
...options,
|
||||
retry: options.retry ?? 3,
|
||||
retryDelay: (attemptIndex) => Math.min(1000 * 2 ** attemptIndex, 30000),
|
||||
retry: options.retry ?? mutationRetry,
|
||||
retryDelay: options.retryDelay ?? retryDelay,
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@ -94,7 +94,7 @@ const InputWrapper: React.FC<InputWrapperProps> = ({
|
||||
<div className="flex w-full flex-col-reverse">
|
||||
<div
|
||||
data-testid="input-wrapper"
|
||||
className="flex w-full flex-col rounded-md border cursor-text border-input p-4 hover:border-muted-foreground focus:border-[1.75px] has-[:focus]:border-primary"
|
||||
className="flex w-full flex-col rounded-md border cursor-text border-input bg-muted p-4 hover:border-muted-foreground focus:border-[1.75px] has-[:focus]:border-primary"
|
||||
onClick={onClick}
|
||||
onMouseDown={onMouseDown}
|
||||
onKeyDown={onKeyDown}
|
||||
|
||||
@ -34,7 +34,7 @@ const UploadFileButton = ({
|
||||
/>
|
||||
<Button
|
||||
disabled={isBuilding}
|
||||
className={`btn-playground-actions ${
|
||||
className={`h-7 w-7 px-0 flex items-center justify-center ${
|
||||
isBuilding
|
||||
? "cursor-not-allowed"
|
||||
: "text-muted-foreground hover:text-primary"
|
||||
@ -42,7 +42,7 @@ const UploadFileButton = ({
|
||||
onClick={handleClick}
|
||||
unstyled
|
||||
>
|
||||
<ForwardedIconComponent name="File" />
|
||||
<ForwardedIconComponent className="h-[18px] w-[18px]" name="File" />
|
||||
</Button>
|
||||
</div>
|
||||
</ShadTooltip>
|
||||
|
||||
@ -50,7 +50,7 @@ function ConfirmationModal({
|
||||
const [flag, setFlag] = useState(false);
|
||||
|
||||
useEffect(() => {
|
||||
if (open) setModalOpen(open);
|
||||
if (open !== undefined) setModalOpen(open);
|
||||
}, [open]);
|
||||
|
||||
useEffect(() => {
|
||||
@ -81,7 +81,7 @@ function ConfirmationModal({
|
||||
return (
|
||||
<BaseModal
|
||||
{...props}
|
||||
open={open}
|
||||
open={modalOpen}
|
||||
setOpen={setModalOpen}
|
||||
onOpenAutoFocus={onOpenAutoFocus}
|
||||
>
|
||||
|
||||
@ -218,6 +218,42 @@ describe("Create mode — canGoNext validation", () => {
|
||||
|
||||
expect(result.current.canGoNext).toBe(true);
|
||||
});
|
||||
|
||||
it("blocks when trimmed name does not start with a letter", () => {
|
||||
const { result } = renderCreateHook({
|
||||
initialProvider: mockProvider,
|
||||
initialInstance: mockInstance,
|
||||
});
|
||||
|
||||
act(() => result.current.handleNext()); // → step 2
|
||||
|
||||
act(() => {
|
||||
result.current.setDeploymentName(" 1 Agent");
|
||||
result.current.setSelectedLlm("gpt-4");
|
||||
});
|
||||
|
||||
expect(result.current.canGoNext).toBe(false);
|
||||
expect(result.current.isDeploymentNameValid).toBe(false);
|
||||
expect(result.current.hasDeploymentNameFormatError).toBe(true);
|
||||
});
|
||||
|
||||
it("allows when trimmed name starts with a unicode letter", () => {
|
||||
const { result } = renderCreateHook({
|
||||
initialProvider: mockProvider,
|
||||
initialInstance: mockInstance,
|
||||
});
|
||||
|
||||
act(() => result.current.handleNext()); // → step 2
|
||||
|
||||
act(() => {
|
||||
result.current.setDeploymentName(" Ágent");
|
||||
result.current.setSelectedLlm("gpt-4");
|
||||
});
|
||||
|
||||
expect(result.current.canGoNext).toBe(true);
|
||||
expect(result.current.isDeploymentNameValid).toBe(true);
|
||||
expect(result.current.hasDeploymentNameFormatError).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Step 3 (Attach Flows)", () => {
|
||||
|
||||
@ -180,4 +180,18 @@ describe("Custom tool naming", () => {
|
||||
"Tool Beta",
|
||||
);
|
||||
});
|
||||
|
||||
it("rejects create payload when agent name does not start with a letter", () => {
|
||||
const { result } = renderStepperHook();
|
||||
|
||||
act(() => {
|
||||
result.current.setDeploymentName("1 Agent");
|
||||
result.current.setSelectedLlm("test-model");
|
||||
result.current.handleSelectVersion("flow-1", "ver-1", "v1");
|
||||
});
|
||||
|
||||
expect(() => result.current.buildDeploymentPayload("provider-1")).toThrow(
|
||||
"Deployment name must start with a letter",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
@ -0,0 +1,75 @@
|
||||
import { render, screen } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import DeploymentStepperFooter from "../components/deployment-stepper-footer";
|
||||
|
||||
jest.mock(
|
||||
"@/components/common/genericIconComponent",
|
||||
() =>
|
||||
function MockIcon({ name }: { name: string }) {
|
||||
return <span data-testid={`icon-${name}`} />;
|
||||
},
|
||||
);
|
||||
|
||||
describe("DeploymentStepperFooter", () => {
|
||||
it("shows done state controls when deployed", () => {
|
||||
render(
|
||||
<DeploymentStepperFooter
|
||||
canGoNext={true}
|
||||
currentStep={4}
|
||||
isCreatingAccount={false}
|
||||
isDeployed={true}
|
||||
isDeploying={false}
|
||||
isInDeployPhase={true}
|
||||
isFinalStep={true}
|
||||
minStep={1}
|
||||
actionIcon="Rocket"
|
||||
actionLabel="Deploy"
|
||||
progressLabel="Deploying..."
|
||||
onBack={jest.fn()}
|
||||
onCancel={jest.fn()}
|
||||
onClose={jest.fn()}
|
||||
onPrimaryAction={jest.fn()}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.getByRole("button", { name: "Done" })).toBeInTheDocument();
|
||||
expect(
|
||||
screen.queryByRole("button", { name: "Cancel" }),
|
||||
).not.toBeInTheDocument();
|
||||
expect(
|
||||
screen.queryByRole("button", { name: "Back" }),
|
||||
).not.toBeInTheDocument();
|
||||
expect(
|
||||
screen.queryByTestId("deployment-stepper-next"),
|
||||
).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("calls onClose from done button", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onClose = jest.fn();
|
||||
|
||||
render(
|
||||
<DeploymentStepperFooter
|
||||
canGoNext={true}
|
||||
currentStep={4}
|
||||
isCreatingAccount={false}
|
||||
isDeployed={true}
|
||||
isDeploying={false}
|
||||
isInDeployPhase={true}
|
||||
isFinalStep={true}
|
||||
minStep={1}
|
||||
actionIcon="Rocket"
|
||||
actionLabel="Deploy"
|
||||
progressLabel="Deploying..."
|
||||
onBack={jest.fn()}
|
||||
onCancel={jest.fn()}
|
||||
onClose={onClose}
|
||||
onPrimaryAction={jest.fn()}
|
||||
/>,
|
||||
);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Done" }));
|
||||
|
||||
expect(onClose).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@ -345,7 +345,7 @@ describe("Step labels", () => {
|
||||
renderModal();
|
||||
expect(screen.getByText("Provider")).toBeInTheDocument();
|
||||
expect(screen.getByText("Type")).toBeInTheDocument();
|
||||
expect(screen.getByText("Attach Flows")).toBeInTheDocument();
|
||||
expect(screen.getByText("Flows")).toBeInTheDocument();
|
||||
expect(screen.getByText("Review")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
@ -355,7 +355,7 @@ describe("Step labels", () => {
|
||||
renderModal({ editingDeployment: makeDeployment() });
|
||||
expect(screen.queryByText("Provider")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Type")).toBeInTheDocument();
|
||||
expect(screen.getByText("Attach Flows")).toBeInTheDocument();
|
||||
expect(screen.getByText("Flows")).toBeInTheDocument();
|
||||
expect(screen.getByText("Review")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
@ -0,0 +1,86 @@
|
||||
import { render, screen } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import DeploymentSuccessContent from "../components/deployment-success-content";
|
||||
|
||||
jest.mock(
|
||||
"@/components/common/genericIconComponent",
|
||||
() =>
|
||||
function MockIcon({ name }: { name: string }) {
|
||||
return <span data-testid={`icon-${name}`} />;
|
||||
},
|
||||
);
|
||||
|
||||
describe("DeploymentSuccessContent", () => {
|
||||
it("renders success copy and provider link", () => {
|
||||
render(
|
||||
<DeploymentSuccessContent
|
||||
deploymentName="My Agent"
|
||||
providerName="watsonx Orchestrate"
|
||||
providerUrl="https://www.ibm.com/products/watsonx-orchestrate"
|
||||
showTestButton={true}
|
||||
onTest={jest.fn()}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.getByText("Deployment successful")).toBeInTheDocument();
|
||||
expect(
|
||||
screen.getByText("Deployed to watsonx Orchestrate as draft"),
|
||||
).toBeInTheDocument();
|
||||
expect(
|
||||
screen.getByRole("link", { name: /watsonx Orchestrate/i }),
|
||||
).toHaveAttribute(
|
||||
"href",
|
||||
"https://www.ibm.com/products/watsonx-orchestrate",
|
||||
);
|
||||
});
|
||||
|
||||
it("shows test button and calls handler", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onTest = jest.fn();
|
||||
|
||||
render(
|
||||
<DeploymentSuccessContent
|
||||
deploymentName="My Agent"
|
||||
providerName="watsonx Orchestrate"
|
||||
providerUrl="https://www.ibm.com/products/watsonx-orchestrate"
|
||||
showTestButton={true}
|
||||
onTest={onTest}
|
||||
/>,
|
||||
);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Test Deployment" }));
|
||||
|
||||
expect(onTest).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("hides test button without deployment name", () => {
|
||||
render(
|
||||
<DeploymentSuccessContent
|
||||
providerName="watsonx Orchestrate"
|
||||
providerUrl="https://www.ibm.com/products/watsonx-orchestrate"
|
||||
showTestButton={true}
|
||||
onTest={jest.fn()}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(
|
||||
screen.queryByRole("button", { name: "Test Deployment" }),
|
||||
).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("hides test button when showTestButton is false", () => {
|
||||
render(
|
||||
<DeploymentSuccessContent
|
||||
deploymentName="My Agent"
|
||||
providerName="watsonx Orchestrate"
|
||||
providerUrl="https://www.ibm.com/products/watsonx-orchestrate"
|
||||
showTestButton={false}
|
||||
onTest={jest.fn()}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(
|
||||
screen.queryByRole("button", { name: "Test Deployment" }),
|
||||
).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
@ -92,6 +92,31 @@ describe("Edit mode — basic state", () => {
|
||||
expect(result.current.canGoNext).toBe(true);
|
||||
});
|
||||
|
||||
it("blocks update flow when existing name does not start with a letter", () => {
|
||||
const invalidDeployment = { ...mockDeployment, name: "1 Agent" };
|
||||
const wrapper = ({ children }: { children: React.ReactNode }) => (
|
||||
<DeploymentStepperProvider
|
||||
initialState={{
|
||||
editingDeployment: invalidDeployment,
|
||||
selectedVersionByFlow: initialVersions,
|
||||
initialLlm: "test-model",
|
||||
initialToolNameByFlow: initialToolNames,
|
||||
initialConnectionsByFlow: initialConnections,
|
||||
}}
|
||||
>
|
||||
{children}
|
||||
</DeploymentStepperProvider>
|
||||
);
|
||||
const { result } = renderHook(() => useDeploymentStepper(), { wrapper });
|
||||
|
||||
expect(result.current.canGoNext).toBe(false);
|
||||
expect(result.current.isDeploymentNameValid).toBe(false);
|
||||
expect(result.current.hasDeploymentNameFormatError).toBe(true);
|
||||
expect(() => result.current.buildDeploymentUpdatePayload()).toThrow(
|
||||
"Deployment name must start with a letter",
|
||||
);
|
||||
});
|
||||
|
||||
it("canGoNext on step 2 (Attach) allows proceeding in edit mode", () => {
|
||||
const { result } = renderEditHook();
|
||||
act(() => result.current.handleNext()); // step 2
|
||||
|
||||
@ -230,14 +230,14 @@ beforeEach(() => {
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe("Three-panel layout rendering", () => {
|
||||
it("renders the Attach Flows heading", () => {
|
||||
it("renders the Flows heading", () => {
|
||||
render(<StepAttachFlows />);
|
||||
expect(screen.getByText("Attach Flows")).toBeInTheDocument();
|
||||
expect(screen.getByText("Flows")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("renders the Available Flows panel header", () => {
|
||||
it("renders the Available panel header", () => {
|
||||
render(<StepAttachFlows />);
|
||||
expect(screen.getByText("Available Flows")).toBeInTheDocument();
|
||||
expect(screen.getByText("Available")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("renders flow items in the list", () => {
|
||||
|
||||
@ -13,6 +13,8 @@ const mockSetSelectedLlm = jest.fn();
|
||||
let mockIsEditMode = false;
|
||||
let mockDeploymentType = "agent";
|
||||
let mockDeploymentName = "";
|
||||
let mockIsDeploymentNameValid = false;
|
||||
let mockHasDeploymentNameFormatError = false;
|
||||
let mockDeploymentDescription = "";
|
||||
let mockSelectedLlm = "";
|
||||
let mockSelectedInstance: { id: string } | null = { id: "inst-1" };
|
||||
@ -26,6 +28,8 @@ jest.mock("../contexts/deployment-stepper-context", () => ({
|
||||
setDeploymentType: mockSetDeploymentType,
|
||||
deploymentName: mockDeploymentName,
|
||||
setDeploymentName: mockSetDeploymentName,
|
||||
isDeploymentNameValid: mockIsDeploymentNameValid,
|
||||
hasDeploymentNameFormatError: mockHasDeploymentNameFormatError,
|
||||
deploymentDescription: mockDeploymentDescription,
|
||||
setDeploymentDescription: mockSetDeploymentDescription,
|
||||
selectedLlm: mockSelectedLlm,
|
||||
@ -61,6 +65,8 @@ beforeEach(() => {
|
||||
mockIsEditMode = false;
|
||||
mockDeploymentType = "agent";
|
||||
mockDeploymentName = "";
|
||||
mockIsDeploymentNameValid = false;
|
||||
mockHasDeploymentNameFormatError = false;
|
||||
mockDeploymentDescription = "";
|
||||
mockSelectedLlm = "";
|
||||
mockSelectedInstance = { id: "inst-1" };
|
||||
@ -143,6 +149,30 @@ describe("Name input", () => {
|
||||
screen.queryByText("Name cannot be changed after creation."),
|
||||
).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows validation error when name does not start with a letter", () => {
|
||||
mockDeploymentName = "1 Agent";
|
||||
mockHasDeploymentNameFormatError = true;
|
||||
render(<StepType />);
|
||||
expect(
|
||||
screen.getByText("Agent name must start with a letter."),
|
||||
).toBeInTheDocument();
|
||||
expect(screen.getByPlaceholderText("e.g., Sales Bot")).toHaveAttribute(
|
||||
"aria-invalid",
|
||||
"true",
|
||||
);
|
||||
});
|
||||
|
||||
it("does not show validation error for empty name", () => {
|
||||
render(<StepType />);
|
||||
expect(
|
||||
screen.queryByText("Agent name must start with a letter."),
|
||||
).not.toBeInTheDocument();
|
||||
expect(screen.getByPlaceholderText("e.g., Sales Bot")).toHaveAttribute(
|
||||
"aria-invalid",
|
||||
"false",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@ -0,0 +1,89 @@
|
||||
import ForwardedIconComponent from "@/components/common/genericIconComponent";
|
||||
import { Button } from "@/components/ui/button";
|
||||
|
||||
interface DeploymentStepperFooterProps {
|
||||
canGoNext: boolean;
|
||||
currentStep: number;
|
||||
isCreatingAccount: boolean;
|
||||
isDeployed: boolean;
|
||||
isDeploying: boolean;
|
||||
isInDeployPhase: boolean;
|
||||
isFinalStep: boolean;
|
||||
minStep: number;
|
||||
actionIcon: string;
|
||||
actionLabel: string;
|
||||
progressLabel: string;
|
||||
onBack: () => void;
|
||||
onCancel: () => void;
|
||||
onClose: () => void;
|
||||
onPrimaryAction: () => void;
|
||||
}
|
||||
|
||||
export default function DeploymentStepperFooter({
|
||||
canGoNext,
|
||||
currentStep,
|
||||
isCreatingAccount,
|
||||
isDeployed,
|
||||
isDeploying,
|
||||
isInDeployPhase,
|
||||
isFinalStep,
|
||||
minStep,
|
||||
actionIcon,
|
||||
actionLabel,
|
||||
progressLabel,
|
||||
onBack,
|
||||
onCancel,
|
||||
onClose,
|
||||
onPrimaryAction,
|
||||
}: DeploymentStepperFooterProps) {
|
||||
return (
|
||||
<div className="flex items-center justify-between border-t border-border px-6 py-4">
|
||||
{isDeployed ? (
|
||||
<div />
|
||||
) : (
|
||||
<Button variant="ghost" onClick={onCancel}>
|
||||
Cancel
|
||||
</Button>
|
||||
)}
|
||||
<div className="flex items-center gap-3">
|
||||
{!isDeployed && (
|
||||
<Button
|
||||
variant="outline"
|
||||
onClick={onBack}
|
||||
disabled={currentStep === minStep || isDeploying}
|
||||
>
|
||||
Back
|
||||
</Button>
|
||||
)}
|
||||
{!isInDeployPhase && (
|
||||
<Button
|
||||
onClick={onPrimaryAction}
|
||||
disabled={!canGoNext || isCreatingAccount}
|
||||
data-testid="deployment-stepper-next"
|
||||
>
|
||||
{isFinalStep ? (
|
||||
<>
|
||||
<ForwardedIconComponent name={actionIcon} className="h-4 w-4" />
|
||||
{actionLabel}
|
||||
</>
|
||||
) : isCreatingAccount ? (
|
||||
"Connecting..."
|
||||
) : (
|
||||
"Next"
|
||||
)}
|
||||
</Button>
|
||||
)}
|
||||
{isDeploying && (
|
||||
<Button disabled data-testid="deployment-stepper-next">
|
||||
<ForwardedIconComponent
|
||||
name={actionIcon}
|
||||
className="h-4 w-4 animate-pulse"
|
||||
/>
|
||||
{progressLabel}
|
||||
</Button>
|
||||
)}
|
||||
{isDeployed && <Button onClick={onClose}>Done</Button>}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@ -1,7 +1,5 @@
|
||||
import { useMemo, useState } from "react";
|
||||
import { useParams } from "react-router-dom";
|
||||
import ForwardedIconComponent from "@/components/common/genericIconComponent";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
@ -21,7 +19,9 @@ import {
|
||||
} from "../contexts/deployment-stepper-context";
|
||||
import { useErrorAlert } from "../hooks/use-error-alert";
|
||||
import type { Deployment, DeploymentProvider, ProviderAccount } from "../types";
|
||||
import DeploymentStepper from "./deployment-stepper";
|
||||
import DeploymentStepper, { CREATE_DEPLOYED_STEPS } from "./deployment-stepper";
|
||||
import DeploymentStepperFooter from "./deployment-stepper-footer";
|
||||
import DeploymentSuccessContent from "./deployment-success-content";
|
||||
import StepAttachFlows from "./step-attach-flows";
|
||||
import StepDeployStatus from "./step-deploy-status";
|
||||
import StepProvider from "./step-provider";
|
||||
@ -206,6 +206,7 @@ function DeploymentStepperModalContent({
|
||||
canGoNext,
|
||||
handleNext,
|
||||
handleBack,
|
||||
selectedProvider,
|
||||
selectedInstance,
|
||||
setSelectedInstance,
|
||||
needsProviderAccountCreation,
|
||||
@ -225,6 +226,8 @@ function DeploymentStepperModalContent({
|
||||
const isDeploying = deploymentPhase === "deploying";
|
||||
const isDeployed = deploymentPhase === "deployed";
|
||||
const isInDeployPhase = isDeploying || isDeployed;
|
||||
const providerConsoleUrl = "https://www.ibm.com/products/watsonx-orchestrate";
|
||||
const providerDisplayName = selectedProvider?.name ?? "watsonx Orchestrate";
|
||||
|
||||
// In edit mode, steps are shifted: 1=Type, 2=Attach, 3=Review.
|
||||
const logicalStep = isEditMode ? currentStep + 1 : currentStep;
|
||||
@ -316,19 +319,40 @@ function DeploymentStepperModalContent({
|
||||
className="text-center text-2xl font-semibold"
|
||||
data-testid="stepper-modal-title"
|
||||
>
|
||||
{isEditMode ? "Update Deployment" : "Create New Deployment"}
|
||||
{isDeployed && !isEditMode
|
||||
? "Deployed"
|
||||
: isEditMode
|
||||
? "Update Deployment"
|
||||
: "Create New Deployment"}
|
||||
</h2>
|
||||
<DeploymentStepper />
|
||||
<DeploymentStepper
|
||||
steps={!isEditMode && isDeployed ? CREATE_DEPLOYED_STEPS : undefined}
|
||||
currentStepOverride={!isEditMode && isDeployed ? 4 : undefined}
|
||||
/>
|
||||
</div>
|
||||
|
||||
{/* Content box: step content + footer */}
|
||||
<div className="mx-4 mb-4 mt-4 flex flex-1 flex-col overflow-hidden rounded-lg border border-border bg-background">
|
||||
<div className="flex min-h-0 flex-1 flex-col overflow-y-auto px-6 py-2">
|
||||
{isInDeployPhase ? (
|
||||
<StepDeployStatus
|
||||
phase={isDeploying ? "deploying" : "deployed"}
|
||||
deploymentName={createdDeployment?.name}
|
||||
/>
|
||||
isDeployed && !isEditMode ? (
|
||||
<DeploymentSuccessContent
|
||||
deploymentName={createdDeployment?.name}
|
||||
providerName={providerDisplayName}
|
||||
providerUrl={providerConsoleUrl}
|
||||
showTestButton={
|
||||
!!onTestDeployment &&
|
||||
!!createdDeployment &&
|
||||
!!selectedInstance?.id
|
||||
}
|
||||
onTest={handleTest}
|
||||
/>
|
||||
) : (
|
||||
<StepDeployStatus
|
||||
phase={isDeploying ? "deploying" : "deployed"}
|
||||
deploymentName={createdDeployment?.name}
|
||||
/>
|
||||
)
|
||||
) : (
|
||||
<>
|
||||
{logicalStep === 1 && <StepProvider />}
|
||||
@ -339,61 +363,23 @@ function DeploymentStepperModalContent({
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* Footer */}
|
||||
<div className="flex items-center justify-between border-t border-border px-6 py-4">
|
||||
<Button variant="ghost" onClick={() => setOpen(false)}>
|
||||
{isDeployed ? "Close" : "Cancel"}
|
||||
</Button>
|
||||
<div className="flex items-center gap-3">
|
||||
{!isDeployed && (
|
||||
<Button
|
||||
variant="outline"
|
||||
onClick={handleBack}
|
||||
disabled={currentStep === minStep || isDeploying}
|
||||
>
|
||||
Back
|
||||
</Button>
|
||||
)}
|
||||
{!isInDeployPhase && (
|
||||
<Button
|
||||
onClick={isFinalStep ? handleDeploy : handleStepNext}
|
||||
disabled={!canGoNext || isCreatingAccount}
|
||||
data-testid="deployment-stepper-next"
|
||||
>
|
||||
{isFinalStep ? (
|
||||
<>
|
||||
<ForwardedIconComponent
|
||||
name={actionIcon}
|
||||
className="h-4 w-4"
|
||||
/>
|
||||
{actionLabel}
|
||||
</>
|
||||
) : isCreatingAccount ? (
|
||||
"Connecting..."
|
||||
) : (
|
||||
"Next"
|
||||
)}
|
||||
</Button>
|
||||
)}
|
||||
{isDeploying && (
|
||||
<Button disabled data-testid="deployment-stepper-next">
|
||||
<ForwardedIconComponent
|
||||
name={actionIcon}
|
||||
className="h-4 w-4 animate-pulse"
|
||||
/>
|
||||
{progressLabel}
|
||||
</Button>
|
||||
)}
|
||||
{isDeployed && onTestDeployment && !isEditMode && (
|
||||
<Button
|
||||
data-testid="deployment-stepper-test"
|
||||
onClick={handleTest}
|
||||
>
|
||||
Test
|
||||
</Button>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
<DeploymentStepperFooter
|
||||
canGoNext={canGoNext}
|
||||
currentStep={currentStep}
|
||||
isCreatingAccount={isCreatingAccount}
|
||||
isDeployed={isDeployed}
|
||||
isDeploying={isDeploying}
|
||||
isInDeployPhase={isInDeployPhase}
|
||||
isFinalStep={isFinalStep}
|
||||
minStep={minStep}
|
||||
actionIcon={actionIcon}
|
||||
actionLabel={actionLabel}
|
||||
progressLabel={progressLabel}
|
||||
onBack={handleBack}
|
||||
onCancel={() => setOpen(false)}
|
||||
onClose={() => setOpen(false)}
|
||||
onPrimaryAction={isFinalStep ? handleDeploy : handleStepNext}
|
||||
/>
|
||||
</div>
|
||||
</>
|
||||
);
|
||||
|
||||
@ -1,25 +1,41 @@
|
||||
import { cn } from "@/utils/utils";
|
||||
import { useDeploymentStepper } from "../contexts/deployment-stepper-context";
|
||||
|
||||
const CREATE_STEPS = [
|
||||
export const CREATE_STEPS = [
|
||||
{ number: 1, label: "Provider" },
|
||||
{ number: 2, label: "Type" },
|
||||
{ number: 3, label: "Attach Flows" },
|
||||
{ number: 3, label: "Flows" },
|
||||
{ number: 4, label: "Review" },
|
||||
] as const;
|
||||
|
||||
export const CREATE_DEPLOYED_STEPS = [
|
||||
{ number: 1, label: "Provider" },
|
||||
{ number: 2, label: "Type" },
|
||||
{ number: 3, label: "Flows" },
|
||||
{ number: 4, label: "Deployed" },
|
||||
] as const;
|
||||
|
||||
const EDIT_STEPS = [
|
||||
{ number: 1, label: "Type" },
|
||||
{ number: 2, label: "Attach Flows" },
|
||||
{ number: 2, label: "Flows" },
|
||||
{ number: 3, label: "Review" },
|
||||
] as const;
|
||||
|
||||
export const DEPLOYMENT_STEPS = CREATE_STEPS;
|
||||
|
||||
export default function DeploymentStepper() {
|
||||
interface DeploymentStepperProps {
|
||||
steps?: readonly { number: number; label: string }[];
|
||||
currentStepOverride?: number;
|
||||
}
|
||||
|
||||
export default function DeploymentStepper({
|
||||
steps: stepsProp,
|
||||
currentStepOverride,
|
||||
}: DeploymentStepperProps) {
|
||||
const { currentStep, isEditMode } = useDeploymentStepper();
|
||||
const steps = isEditMode ? EDIT_STEPS : CREATE_STEPS;
|
||||
const progressPercent = ((currentStep - 1) / (steps.length - 1)) * 100;
|
||||
const steps = stepsProp ?? (isEditMode ? EDIT_STEPS : CREATE_STEPS);
|
||||
const activeStep = currentStepOverride ?? currentStep;
|
||||
const progressPercent = ((activeStep - 1) / (steps.length - 1)) * 100;
|
||||
|
||||
return (
|
||||
<div className="relative mx-auto h-[52px] w-full max-w-[700px]">
|
||||
@ -35,7 +51,7 @@ export default function DeploymentStepper() {
|
||||
<div
|
||||
className={cn(
|
||||
"flex h-8 w-8 items-center justify-center rounded-full text-sm font-medium transition-colors",
|
||||
currentStep >= step.number
|
||||
activeStep >= step.number
|
||||
? "bg-foreground text-background"
|
||||
: "bg-muted text-muted-foreground",
|
||||
)}
|
||||
@ -45,7 +61,7 @@ export default function DeploymentStepper() {
|
||||
<span
|
||||
className={cn(
|
||||
"whitespace-nowrap text-xs text-foreground",
|
||||
currentStep >= step.number && "font-medium",
|
||||
activeStep >= step.number && "font-medium",
|
||||
)}
|
||||
>
|
||||
{step.label}
|
||||
|
||||
@ -0,0 +1,63 @@
|
||||
import ForwardedIconComponent from "@/components/common/genericIconComponent";
|
||||
import { Button } from "@/components/ui/button";
|
||||
|
||||
interface DeploymentSuccessContentProps {
|
||||
deploymentName?: string;
|
||||
providerName: string;
|
||||
providerUrl: string;
|
||||
showTestButton: boolean;
|
||||
onTest: () => void;
|
||||
}
|
||||
|
||||
export default function DeploymentSuccessContent({
|
||||
deploymentName,
|
||||
providerName,
|
||||
providerUrl,
|
||||
showTestButton,
|
||||
onTest,
|
||||
}: DeploymentSuccessContentProps) {
|
||||
return (
|
||||
<div className="flex flex-1 flex-col items-center justify-center gap-8 py-10 text-center">
|
||||
<div className="flex h-16 w-16 items-center justify-center rounded-full border-2 border-accent-emerald">
|
||||
<ForwardedIconComponent
|
||||
name="Check"
|
||||
className="h-8 w-8 text-accent-emerald-foreground"
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div className="space-y-2">
|
||||
<h3 className="font-sans text-2xl font-semibold tracking-normal text-foreground">
|
||||
Deployment successful
|
||||
</h3>
|
||||
<p className="font-sans text-base font-normal text-muted-foreground">
|
||||
Deployed to {providerName} as draft
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<div className="flex items-center gap-1 font-sans text-sm font-normal text-muted-foreground">
|
||||
<span>Publish from Draft to Live in</span>
|
||||
<a
|
||||
href={providerUrl}
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="inline-flex items-center gap-1 font-semibold text-foreground hover:underline"
|
||||
>
|
||||
{providerName}
|
||||
<ForwardedIconComponent name="ArrowUpRight" className="h-4 w-4" />
|
||||
</a>
|
||||
</div>
|
||||
|
||||
{showTestButton && deploymentName && (
|
||||
<Button
|
||||
variant="outline"
|
||||
className="h-11 w-full max-w-lg gap-2 rounded-lg border-input font-sans text-base font-semibold text-foreground hover:bg-input"
|
||||
data-testid="deployment-stepper-test"
|
||||
onClick={onTest}
|
||||
>
|
||||
<ForwardedIconComponent name="Play" className="h-5 w-5" />
|
||||
Test Deployment
|
||||
</Button>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@ -29,7 +29,7 @@ export const FlowListPanel = memo(function FlowListPanel({
|
||||
return (
|
||||
<div className="flex w-[280px] flex-shrink-0 flex-col border-r border-border">
|
||||
<div className="border-b border-border p-4 text-sm text-muted-foreground">
|
||||
Available Flows
|
||||
Available
|
||||
</div>
|
||||
<div className="flex-1 space-y-1 overflow-y-auto p-2">
|
||||
{flows.map((flow) => {
|
||||
|
||||
@ -259,7 +259,7 @@ export default function StepAttachFlows() {
|
||||
|
||||
return (
|
||||
<div className="flex min-h-0 flex-1 flex-col gap-4 py-3">
|
||||
<h2 className="text-lg font-semibold">Attach Flows</h2>
|
||||
<h2 className="text-lg font-semibold">Flows</h2>
|
||||
|
||||
<div className="flex min-h-0 flex-1 overflow-hidden rounded-xl border border-border">
|
||||
<FlowListPanel
|
||||
|
||||
@ -32,6 +32,7 @@ export default function StepType() {
|
||||
setDeploymentType,
|
||||
deploymentName,
|
||||
setDeploymentName,
|
||||
hasDeploymentNameFormatError,
|
||||
deploymentDescription,
|
||||
setDeploymentDescription,
|
||||
selectedLlm,
|
||||
@ -142,11 +143,21 @@ export default function StepType() {
|
||||
</span>
|
||||
<Input
|
||||
placeholder="e.g., Sales Bot"
|
||||
className="bg-muted"
|
||||
className={cn(
|
||||
"bg-muted",
|
||||
hasDeploymentNameFormatError &&
|
||||
"border-destructive focus-visible:ring-0",
|
||||
)}
|
||||
value={deploymentName}
|
||||
onChange={(e) => setDeploymentName(e.target.value)}
|
||||
disabled={isEditMode}
|
||||
aria-invalid={hasDeploymentNameFormatError}
|
||||
/>
|
||||
{hasDeploymentNameFormatError && (
|
||||
<span className="mt-1 text-xs text-destructive">
|
||||
Agent name must start with a letter.
|
||||
</span>
|
||||
)}
|
||||
{isEditMode && (
|
||||
<span className="mt-1 text-xs text-muted-foreground">
|
||||
Name cannot be changed after creation.
|
||||
|
||||
@ -73,6 +73,8 @@ interface DeploymentStepperContextType {
|
||||
setDeploymentType: (type: DeploymentType) => void;
|
||||
deploymentName: string;
|
||||
setDeploymentName: (name: string) => void;
|
||||
isDeploymentNameValid: boolean;
|
||||
hasDeploymentNameFormatError: boolean;
|
||||
deploymentDescription: string;
|
||||
setDeploymentDescription: (description: string) => void;
|
||||
selectedLlm: string;
|
||||
@ -168,6 +170,11 @@ export function DeploymentStepperProvider({
|
||||
>(initialState?.initialConnectionsByFlow ?? new Map());
|
||||
|
||||
const [hasToolNameErrors, setHasToolNameErrors] = useState(false);
|
||||
const trimmedDeploymentName = deploymentName.trim();
|
||||
const hasDeploymentNameFormatError =
|
||||
trimmedDeploymentName !== "" && !/^\p{L}/u.test(trimmedDeploymentName);
|
||||
const isDeploymentNameValid =
|
||||
trimmedDeploymentName !== "" && !hasDeploymentNameFormatError;
|
||||
|
||||
// Edit mode: track which pre-existing flows the user wants to detach.
|
||||
const [removedFlowIds, setRemovedFlowIds] = useState<Set<string>>(new Set());
|
||||
@ -253,7 +260,7 @@ export function DeploymentStepperProvider({
|
||||
);
|
||||
}
|
||||
if (logical === 2) {
|
||||
return deploymentName.trim() !== "" && selectedLlm.trim() !== "";
|
||||
return isDeploymentNameValid && selectedLlm.trim() !== "";
|
||||
}
|
||||
if (logical === 3) {
|
||||
// In edit mode, user can proceed without new attachments (may just change desc/LLM).
|
||||
@ -270,6 +277,7 @@ export function DeploymentStepperProvider({
|
||||
selectedInstance,
|
||||
hasValidCredentials,
|
||||
deploymentName,
|
||||
isDeploymentNameValid,
|
||||
selectedLlm,
|
||||
selectedVersionByFlow,
|
||||
isEditMode,
|
||||
@ -357,6 +365,9 @@ export function DeploymentStepperProvider({
|
||||
|
||||
const buildDeploymentPayload = useCallback(
|
||||
(providerId: string): DeploymentCreateRequest => {
|
||||
if (!isDeploymentNameValid) {
|
||||
throw new Error("Deployment name must start with a letter");
|
||||
}
|
||||
const allConnectionIds = new Set<string>();
|
||||
Array.from(attachedConnectionByFlow.values()).forEach((ids) => {
|
||||
ids.forEach((id) => allConnectionIds.add(id));
|
||||
@ -381,7 +392,7 @@ export function DeploymentStepperProvider({
|
||||
...(initialState?.projectId
|
||||
? { project_id: initialState.projectId }
|
||||
: {}),
|
||||
name: deploymentName,
|
||||
name: trimmedDeploymentName,
|
||||
description: deploymentDescription,
|
||||
type: deploymentType,
|
||||
provider_data: {
|
||||
@ -396,10 +407,11 @@ export function DeploymentStepperProvider({
|
||||
buildConnectionPayloads,
|
||||
initialState?.projectId,
|
||||
deploymentDescription,
|
||||
deploymentName,
|
||||
deploymentType,
|
||||
isDeploymentNameValid,
|
||||
selectedLlm,
|
||||
selectedVersionByFlow,
|
||||
trimmedDeploymentName,
|
||||
toolNameByFlow,
|
||||
],
|
||||
);
|
||||
@ -411,6 +423,9 @@ export function DeploymentStepperProvider({
|
||||
"buildDeploymentUpdatePayload called outside edit mode",
|
||||
);
|
||||
}
|
||||
if (!isDeploymentNameValid) {
|
||||
throw new Error("Deployment name must start with a letter");
|
||||
}
|
||||
|
||||
const result: DeploymentUpdateRequest = {
|
||||
deployment_id: editingDeployment.id,
|
||||
@ -510,6 +525,7 @@ export function DeploymentStepperProvider({
|
||||
}, [
|
||||
editingDeployment,
|
||||
deploymentDescription,
|
||||
isDeploymentNameValid,
|
||||
selectedLlm,
|
||||
initialVersionByFlow,
|
||||
initialToolNameByFlow,
|
||||
@ -541,6 +557,8 @@ export function DeploymentStepperProvider({
|
||||
setDeploymentType,
|
||||
deploymentName,
|
||||
setDeploymentName,
|
||||
isDeploymentNameValid,
|
||||
hasDeploymentNameFormatError,
|
||||
deploymentDescription,
|
||||
setDeploymentDescription,
|
||||
selectedLlm,
|
||||
@ -581,6 +599,8 @@ export function DeploymentStepperProvider({
|
||||
credentials,
|
||||
deploymentType,
|
||||
deploymentName,
|
||||
isDeploymentNameValid,
|
||||
hasDeploymentNameFormatError,
|
||||
deploymentDescription,
|
||||
selectedLlm,
|
||||
connections,
|
||||
|
||||
@ -150,12 +150,13 @@ async function selectProvider(page: Page) {
|
||||
async function goToStepType(page: Page) {
|
||||
await selectProvider(page);
|
||||
await page.getByTestId("deployment-stepper-next").click();
|
||||
// Wait for the Type step heading
|
||||
await page.waitForSelector("text=Deployment Type");
|
||||
await expect(
|
||||
page.getByRole("heading", { name: /Deployment Type/i }),
|
||||
).toBeVisible();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helper: navigate steps 1 → 2 → 3 (provider → type → attach flows)
|
||||
// Helper: navigate steps 1 → 2 → 3 (provider → type → flows)
|
||||
// ---------------------------------------------------------------------------
|
||||
async function goToStepAttachFlows(page: Page) {
|
||||
await goToStepType(page);
|
||||
@ -165,13 +166,13 @@ async function goToStepAttachFlows(page: Page) {
|
||||
// Select the LLM model
|
||||
await page.getByRole("combobox").click();
|
||||
await page.getByRole("option", { name: "ibm/granite-13b-chat" }).click();
|
||||
// Advance to attach flows
|
||||
// Advance to flows
|
||||
await page.getByTestId("deployment-stepper-next").click();
|
||||
await page.waitForSelector("text=Attach Flows");
|
||||
await expect(page.getByRole("heading", { name: /^Flows$/i })).toBeVisible();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helper: navigate all steps through attach flows and select a flow+version
|
||||
// Helper: navigate all steps through flows and select a flow+version
|
||||
// ---------------------------------------------------------------------------
|
||||
async function goToStepReview(page: Page) {
|
||||
await goToStepAttachFlows(page);
|
||||
@ -189,7 +190,9 @@ async function goToStepReview(page: Page) {
|
||||
}
|
||||
// Advance to review
|
||||
await page.getByTestId("deployment-stepper-next").click();
|
||||
await page.waitForSelector("text=Review & Confirm");
|
||||
await expect(
|
||||
page.getByRole("heading", { name: /Review & Confirm/i }),
|
||||
).toBeVisible();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@ -197,7 +200,9 @@ async function goToStepReview(page: Page) {
|
||||
// ---------------------------------------------------------------------------
|
||||
test(
|
||||
"deployment-create: opens stepper on New Deployment click",
|
||||
{ tag: ["@deployment", "@workspace"] },
|
||||
{
|
||||
tag: ["@deployment", "@workspace"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -220,7 +225,9 @@ test(
|
||||
// ---------------------------------------------------------------------------
|
||||
test(
|
||||
"deployment-create: step 1 provider - Next disabled without selection, enabled after selecting",
|
||||
{ tag: ["@deployment", "@workspace"] },
|
||||
{
|
||||
tag: ["@deployment", "@workspace"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -245,7 +252,9 @@ test(
|
||||
// ---------------------------------------------------------------------------
|
||||
test(
|
||||
"deployment-create: step 2 type - fill name and select type to enable Next",
|
||||
{ tag: ["@deployment", "@workspace"] },
|
||||
{
|
||||
tag: ["@deployment", "@workspace"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -280,11 +289,13 @@ test(
|
||||
);
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 4: Step 3 (Attach Flows) — select a flow and version, Next enables
|
||||
// Test 4: Step 3 (Flows) — select a flow and version, Next enables
|
||||
// ---------------------------------------------------------------------------
|
||||
test(
|
||||
"deployment-create: step 3 attach flows - select flow and version enables Next",
|
||||
{ tag: ["@deployment", "@workspace"] },
|
||||
{
|
||||
tag: ["@deployment", "@workspace"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -331,7 +342,9 @@ test(
|
||||
// ---------------------------------------------------------------------------
|
||||
test(
|
||||
"deployment-create: step 4 review - shows review content and Deploy button text",
|
||||
{ tag: ["@deployment", "@workspace"] },
|
||||
{
|
||||
tag: ["@deployment", "@workspace"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -361,7 +374,9 @@ test(
|
||||
// ---------------------------------------------------------------------------
|
||||
test(
|
||||
"deployment-create: clicking Deploy triggers POST and shows deploy status",
|
||||
{ tag: ["@deployment", "@workspace"] },
|
||||
{
|
||||
tag: ["@deployment", "@workspace"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -414,7 +429,9 @@ test(
|
||||
// ---------------------------------------------------------------------------
|
||||
test(
|
||||
"deployment-create: user can change tool name on review step",
|
||||
{ tag: ["@deployment", "@workspace"] },
|
||||
{
|
||||
tag: ["@deployment", "@workspace"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -451,7 +468,9 @@ test(
|
||||
// ---------------------------------------------------------------------------
|
||||
test(
|
||||
"deployment-create: review step shows error when tool name already exists in provider",
|
||||
{ tag: ["@deployment", "@workspace"] },
|
||||
{
|
||||
tag: ["@deployment", "@workspace"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -476,7 +495,9 @@ test(
|
||||
// ---------------------------------------------------------------------------
|
||||
test(
|
||||
"deployment-create: review step shows no error when tool name is unique",
|
||||
{ tag: ["@deployment", "@workspace"] },
|
||||
{
|
||||
tag: ["@deployment", "@workspace"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -501,7 +522,9 @@ test(
|
||||
// ---------------------------------------------------------------------------
|
||||
test(
|
||||
"deployment-create: editing tool name to unique value clears duplicate error",
|
||||
{ tag: ["@deployment", "@workspace"] },
|
||||
{
|
||||
tag: ["@deployment", "@workspace"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
|
||||
@ -83,9 +83,23 @@ async function openEditDialog(page: Parameters<typeof test>[2]["page"]) {
|
||||
await page.waitForSelector('[data-testid="stepper-modal-title"]');
|
||||
}
|
||||
|
||||
async function expectDeploymentTypeStep(
|
||||
page: Parameters<typeof test>[2]["page"],
|
||||
) {
|
||||
await expect(
|
||||
page.getByRole("heading", { name: /Deployment Type/i }),
|
||||
).toBeVisible();
|
||||
}
|
||||
|
||||
async function expectFlowsStep(page: Parameters<typeof test>[2]["page"]) {
|
||||
await expect(page.getByRole("heading", { name: /^Flows$/i })).toBeVisible();
|
||||
}
|
||||
|
||||
test(
|
||||
"Opens edit stepper from actions menu",
|
||||
{ tag: ["@release", "@workspace", "@api"] },
|
||||
{
|
||||
tag: ["@release", "@workspace", "@api"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -110,7 +124,9 @@ test(
|
||||
|
||||
test(
|
||||
"Edit mode skips provider step — starts at Type",
|
||||
{ tag: ["@release", "@workspace", "@api"] },
|
||||
{
|
||||
tag: ["@release", "@workspace", "@api"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -126,7 +142,7 @@ test(
|
||||
await openEditDialog(page);
|
||||
|
||||
// Wait for stepper body to render (parallel fetches must complete)
|
||||
await page.waitForSelector('h2:has-text("Deployment Type")');
|
||||
await expectDeploymentTypeStep(page);
|
||||
|
||||
await expect(page.getByText("Deployment Type")).toBeVisible();
|
||||
|
||||
@ -141,7 +157,9 @@ test(
|
||||
|
||||
test(
|
||||
"Name field pre-populated in edit mode",
|
||||
{ tag: ["@release", "@workspace", "@api"] },
|
||||
{
|
||||
tag: ["@release", "@workspace", "@api"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -157,7 +175,7 @@ test(
|
||||
await openEditDialog(page);
|
||||
|
||||
// Wait for stepper body to render
|
||||
await page.waitForSelector('h2:has-text("Deployment Type")');
|
||||
await expectDeploymentTypeStep(page);
|
||||
|
||||
const nameInput = page.getByPlaceholder("e.g., Sales Bot");
|
||||
await expect(nameInput).toHaveValue("Test Deployment");
|
||||
@ -166,7 +184,9 @@ test(
|
||||
|
||||
test(
|
||||
"Submitting PATCH closes modal",
|
||||
{ tag: ["@release", "@workspace", "@api"] },
|
||||
{
|
||||
tag: ["@release", "@workspace", "@api"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -182,13 +202,13 @@ test(
|
||||
await openEditDialog(page);
|
||||
|
||||
// Wait for stepper body to render (parallel fetches must complete)
|
||||
await page.waitForSelector('h2:has-text("Deployment Type")');
|
||||
await expectDeploymentTypeStep(page);
|
||||
|
||||
// Navigate through the stepper steps to reach Review
|
||||
// Step: Type → click Next
|
||||
await page.getByTestId("deployment-stepper-next").click();
|
||||
|
||||
// Step: Attach Flows → click Next
|
||||
// Step: Flows → click Next
|
||||
await page.getByTestId("deployment-stepper-next").click();
|
||||
|
||||
// Step: Review → click Update (final step)
|
||||
@ -209,7 +229,9 @@ test(
|
||||
|
||||
test(
|
||||
"Cancel during edit closes modal without calling PATCH",
|
||||
{ tag: ["@release", "@workspace", "@api"] },
|
||||
{
|
||||
tag: ["@release", "@workspace", "@api"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -225,7 +247,7 @@ test(
|
||||
await openEditDialog(page);
|
||||
|
||||
// Wait for stepper body to render
|
||||
await page.waitForSelector('h2:has-text("Deployment Type")');
|
||||
await expectDeploymentTypeStep(page);
|
||||
|
||||
let patchCalled = false;
|
||||
page.on("request", (req) => {
|
||||
@ -356,7 +378,9 @@ async function setupRoutesWithConnections(
|
||||
// ---------------------------------------------------------------------------
|
||||
test(
|
||||
"Edit mode includes new connections in PATCH request",
|
||||
{ tag: ["@release", "@workspace", "@api"] },
|
||||
{
|
||||
tag: ["@release", "@workspace", "@api"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -399,11 +423,11 @@ test(
|
||||
await page.waitForSelector('[data-testid="stepper-modal-title"]');
|
||||
|
||||
// Step 1 (Type) → Next
|
||||
await page.waitForSelector('h2:has-text("Deployment Type")');
|
||||
await expectDeploymentTypeStep(page);
|
||||
await page.getByTestId("deployment-stepper-next").click();
|
||||
|
||||
// Step 2 (Attach Flows) — flow "f1" should already be attached
|
||||
await page.waitForSelector("text=Attach Flows");
|
||||
// Step 2 (Flows) — flow "f1" should already be attached
|
||||
await expectFlowsStep(page);
|
||||
await page.waitForSelector('[data-testid="flow-item-f1"]');
|
||||
|
||||
// Click the pre-attached flow and its version to open the connection panel
|
||||
@ -455,7 +479,9 @@ test(
|
||||
// ---------------------------------------------------------------------------
|
||||
test(
|
||||
"Edit mode includes removed connections in PATCH request",
|
||||
{ tag: ["@release", "@workspace", "@api"] },
|
||||
{
|
||||
tag: ["@release", "@workspace", "@api"],
|
||||
},
|
||||
async ({ page }) => {
|
||||
test.skip(
|
||||
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
|
||||
@ -497,11 +523,11 @@ test(
|
||||
await page.waitForSelector('[data-testid="stepper-modal-title"]');
|
||||
|
||||
// Step 1 (Type) → Next
|
||||
await page.waitForSelector('h2:has-text("Deployment Type")');
|
||||
await expectDeploymentTypeStep(page);
|
||||
await page.getByTestId("deployment-stepper-next").click();
|
||||
|
||||
// Step 2 (Attach Flows)
|
||||
await page.waitForSelector("text=Attach Flows");
|
||||
// Step 2 (Flows)
|
||||
await expectFlowsStep(page);
|
||||
await page.waitForSelector('[data-testid="flow-item-f1"]');
|
||||
await page.getByTestId("flow-item-f1").click();
|
||||
await page.waitForSelector('[data-testid="version-item-fv1"]');
|
||||
|
||||
103
src/frontend/tests/core/unit/agentDefaultToolsComponent.spec.ts
Normal file
103
src/frontend/tests/core/unit/agentDefaultToolsComponent.spec.ts
Normal file
@ -0,0 +1,103 @@
|
||||
import { expect, test } from "../../fixtures";
|
||||
import { adjustScreenView } from "../../utils/adjust-screen-view";
|
||||
import { awaitBootstrapTest } from "../../utils/await-bootstrap-test";
|
||||
import {
|
||||
closeAdvancedOptions,
|
||||
openAdvancedOptions,
|
||||
} from "../../utils/open-advanced-options";
|
||||
|
||||
/**
|
||||
* Covers the user-facing contract of the "Default Agent Tools" feature:
|
||||
* see CZL/MANUAL_TEST_DEFAULT_AGENT_TOOLS.md scenarios S1, S3, S5.
|
||||
*
|
||||
* Backing unit tests (pytest) already cover runtime behaviour; these tests
|
||||
* validate the UI wiring so a regression in the inputs list or default
|
||||
* prompt value is caught in the release gate.
|
||||
*/
|
||||
|
||||
async function dragAgentOntoCanvas(page: import("@playwright/test").Page) {
|
||||
await awaitBootstrapTest(page);
|
||||
|
||||
await page.waitForSelector('[data-testid="blank-flow"]', { timeout: 30000 });
|
||||
await page.getByTestId("blank-flow").click();
|
||||
|
||||
await page.waitForSelector('[data-testid="sidebar-search-input"]', {
|
||||
timeout: 30000,
|
||||
});
|
||||
|
||||
await page.getByTestId("sidebar-search-input").click();
|
||||
await page.getByTestId("sidebar-search-input").fill("agent");
|
||||
|
||||
await page.waitForSelector('[data-testid="models_and_agentsAgent"]', {
|
||||
timeout: 30000,
|
||||
});
|
||||
|
||||
await page
|
||||
.getByTestId("models_and_agentsAgent")
|
||||
.dragTo(page.locator('//*[@id="react-flow-id"]'));
|
||||
|
||||
await adjustScreenView(page);
|
||||
}
|
||||
|
||||
test(
|
||||
"Agent ships with Calculator and Current Date toggles enabled by default (S1/S3)",
|
||||
{ tag: ["@release", "@workspace", "@components"] },
|
||||
async ({ page }) => {
|
||||
await dragAgentOntoCanvas(page);
|
||||
|
||||
// Focus the Agent node so its advanced-field drawer is reachable.
|
||||
await page.getByTestId("div-generic-node").click();
|
||||
|
||||
await openAdvancedOptions(page);
|
||||
|
||||
// Both advanced toggles exist as show-on-canvas checkboxes.
|
||||
// Their default `value=True` is validated by the pytest suite
|
||||
// (`test_should_have_placeholders_in_default_system_prompt` covers the
|
||||
// default contract of the inputs list).
|
||||
await expect(
|
||||
page.locator('//*[@id="showadd_current_date_tool"]'),
|
||||
).toBeVisible({ timeout: 10000 });
|
||||
await expect(
|
||||
page.locator('//*[@id="showadd_calculator_tool"]'),
|
||||
).toBeVisible({ timeout: 10000 });
|
||||
|
||||
// Flip the Calculator field visible on canvas so we can assert the toggle
|
||||
// is active and can be switched off and on (S3).
|
||||
await page.locator('//*[@id="showadd_calculator_tool"]').click();
|
||||
await closeAdvancedOptions(page);
|
||||
|
||||
await adjustScreenView(page);
|
||||
|
||||
const calculatorToggle = page.getByTestId(
|
||||
"toggle_bool_add_calculator_tool",
|
||||
);
|
||||
await expect(calculatorToggle).toBeVisible({ timeout: 10000 });
|
||||
expect(await calculatorToggle.isChecked()).toBeTruthy();
|
||||
|
||||
// S3 — user disables the toggle.
|
||||
await calculatorToggle.click();
|
||||
expect(await calculatorToggle.isChecked()).toBeFalsy();
|
||||
|
||||
// Re-enable to confirm the control is bi-directional.
|
||||
await calculatorToggle.click();
|
||||
expect(await calculatorToggle.isChecked()).toBeTruthy();
|
||||
},
|
||||
);
|
||||
|
||||
test(
|
||||
"Agent default system prompt contains {current_date} and {model_name} placeholders (S5)",
|
||||
{ tag: ["@release", "@workspace", "@components"] },
|
||||
async ({ page }) => {
|
||||
await dragAgentOntoCanvas(page);
|
||||
|
||||
// The placeholders must be present inside the Agent Instructions textarea
|
||||
// so the dynamic injection has a discoverable effect out of the box.
|
||||
const instructionsTextarea = page.getByTestId("textarea_str_system_prompt");
|
||||
await expect(instructionsTextarea).toBeVisible({ timeout: 10000 });
|
||||
|
||||
// <textarea> stores content in its `value`, not textContent.
|
||||
const promptText = await instructionsTextarea.inputValue();
|
||||
expect(promptText).toContain("{current_date}");
|
||||
expect(promptText).toContain("{model_name}");
|
||||
},
|
||||
);
|
||||
@ -41,13 +41,23 @@ test(
|
||||
|
||||
const cleanCode = await extractAndCleanCode(page);
|
||||
|
||||
// Replace the multiline string in the code
|
||||
// Sanity-check: the original Ollama source has many lines. If we read it
|
||||
// back as a single concatenated line, the textarea-value path is the
|
||||
// problem and Ace surgery downstream cannot save us.
|
||||
const cleanCodeNewlines = (cleanCode.match(/\n/g) || []).length;
|
||||
if (cleanCodeNewlines < 50) {
|
||||
throw new Error(
|
||||
`extractAndCleanCode returned code with only ${cleanCodeNewlines} newlines (length ${cleanCode.length}); expected the multi-line Ollama source. The hidden #codeValue textarea may be returning a stripped value on this platform.`,
|
||||
);
|
||||
}
|
||||
|
||||
// Replace the multiline string in the code.
|
||||
// Use a regex so the match is resilient to line-ending differences
|
||||
// (LF on macOS/Linux vs CRLF on Windows after git checkout).
|
||||
const originalSliderBlockRegex =
|
||||
/name="temperature",\s+display_name="Temperature",\s+value=0\.1,\s+range_spec=RangeSpec\(min=0, max=1, step=0\.01\),\s+advanced=True,/;
|
||||
const newCode = cleanCode.replace(
|
||||
`name="temperature",
|
||||
display_name="Temperature",
|
||||
value=0.1,
|
||||
range_spec=RangeSpec(min=0, max=1, step=0.01),
|
||||
advanced=True,`,
|
||||
originalSliderBlockRegex,
|
||||
`name="temperature",
|
||||
display_name="Temperature",
|
||||
value=0.2,
|
||||
@ -63,9 +73,7 @@ test(
|
||||
);
|
||||
// make sure codes are different
|
||||
expect(cleanCode).not.toEqual(newCode);
|
||||
await page.locator("textarea").last().press(`ControlOrMeta+a`);
|
||||
await page.keyboard.press("Backspace");
|
||||
await page.locator("textarea").last().fill(newCode);
|
||||
await setAceEditorValue(page, newCode);
|
||||
await page.locator('//*[@id="checkAndSaveBtn"]').click();
|
||||
await adjustScreenView(page);
|
||||
|
||||
@ -104,13 +112,24 @@ test(
|
||||
);
|
||||
|
||||
async function extractAndCleanCode(page: Page): Promise<string> {
|
||||
// The hidden `#codeValue` mirror is rendered with `<Input value={code} />`,
|
||||
// i.e. a single-line `<input>` rather than a textarea. Browsers strip
|
||||
// newlines from `.value` of single-line inputs by HTML spec, so reading
|
||||
// `(el as HTMLInputElement).value` returns the entire source concatenated
|
||||
// onto one line. The backend then rejects it as
|
||||
// `invalid decimal literal (<unknown>, line 1)`.
|
||||
//
|
||||
// Read the rendered HTML attribute instead — React serializes the prop as
|
||||
// `value="..."` with newlines encoded (` `, ` `, or literal LF
|
||||
// inside attribute text), and the regex below recovers them faithfully.
|
||||
// This is the same strategy queryInputComponent.spec.ts uses.
|
||||
const outerHTML = await page
|
||||
.locator('//*[@id="codeValue"]')
|
||||
.evaluate((el) => el.outerHTML);
|
||||
|
||||
const valueMatch = outerHTML.match(/value="([\s\S]*?)"/);
|
||||
if (!valueMatch) {
|
||||
throw new Error("Could not find value attribute in the HTML");
|
||||
throw new Error("Could not find value attribute in the #codeValue HTML");
|
||||
}
|
||||
|
||||
const codeContent = valueMatch[1]
|
||||
@ -119,9 +138,45 @@ async function extractAndCleanCode(page: Page): Promise<string> {
|
||||
.replace(/</g, "<")
|
||||
.replace(/>/g, ">")
|
||||
.replace(/'/g, "'")
|
||||
.replace(///g, "/");
|
||||
.replace(///g, "/")
|
||||
.replace(/ /g, "\n")
|
||||
.replace(/ /g, "\r")
|
||||
.replace(/
/gi, "\n")
|
||||
.replace(/
/gi, "\r");
|
||||
|
||||
return codeContent;
|
||||
// Normalize line endings so downstream string operations are OS-agnostic
|
||||
// (Windows git checkouts of .py files can end up with CRLF).
|
||||
return codeContent.replace(/\r\n/g, "\n").replace(/\r/g, "\n");
|
||||
}
|
||||
|
||||
// Set the Ace editor's content using the exact same pattern that
|
||||
// queryInputComponent.spec.ts uses successfully on Windows CI:
|
||||
// `textarea.last().fill(newCode)` against Ace's hidden text-input textarea.
|
||||
// Textareas (unlike single-line `<input>`s) preserve `\n` in `.value`, so
|
||||
// fill() reliably round-trips the multi-line source. Ace's text-input
|
||||
// listener picks up the resulting `input` event, applies it to the buffer,
|
||||
// and fires the `change` event that react-ace listens to — which is what
|
||||
// updates the React `code` state that gets POSTed on save.
|
||||
async function setAceEditorValue(page: Page, newCode: string): Promise<void> {
|
||||
const expectedNewlines = (newCode.match(/\n/g) || []).length;
|
||||
if (expectedNewlines < 10) {
|
||||
throw new Error(
|
||||
`setAceEditorValue: newCode has only ${expectedNewlines} newlines (length ${newCode.length}); upstream extractAndCleanCode likely lost newlines.`,
|
||||
);
|
||||
}
|
||||
|
||||
await page.locator("textarea").last().press("ControlOrMeta+a");
|
||||
await page.keyboard.press("Backspace");
|
||||
await page.locator("textarea").last().fill(newCode);
|
||||
|
||||
// Wait for the change to propagate into the controlled React state. The
|
||||
// `#codeValue` mirror is rendered by `<Input>` (single-line), so browsers
|
||||
// strip newlines from its `.value` — matching the substring is enough to
|
||||
// confirm the new code landed.
|
||||
await expect(page.locator("#codeValue")).toHaveValue(
|
||||
/range_spec=RangeSpec\(min=3, max=30, step=1\)/,
|
||||
{ timeout: 10000 },
|
||||
);
|
||||
}
|
||||
|
||||
async function mutualValidation(page: Page) {
|
||||
|
||||
File diff suppressed because one or more lines are too long
@ -1,7 +1,6 @@
|
||||
from copy import deepcopy
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from chromadb.config import Settings
|
||||
from langchain_chroma import Chroma
|
||||
from typing_extensions import override
|
||||
|
||||
@ -92,23 +91,22 @@ class ChromaVectorStoreComponent(LCVectorStoreComponent):
|
||||
def build_vector_store(self) -> Chroma:
|
||||
"""Builds the Chroma object."""
|
||||
try:
|
||||
from chromadb import Client
|
||||
from langchain_chroma import Chroma
|
||||
except ImportError as e:
|
||||
msg = "Could not import Chroma integration package. Please install it with `pip install langchain-chroma`."
|
||||
raise ImportError(msg) from e
|
||||
# Chroma settings
|
||||
chroma_settings = None
|
||||
client = None
|
||||
if self.chroma_server_host:
|
||||
chroma_settings = Settings(
|
||||
chroma_server_cors_allow_origins=self.chroma_server_cors_allow_origins or [],
|
||||
chroma_server_host=self.chroma_server_host,
|
||||
chroma_server_http_port=self.chroma_server_http_port or None,
|
||||
chroma_server_grpc_port=self.chroma_server_grpc_port or None,
|
||||
chroma_server_ssl_enabled=self.chroma_server_ssl_enabled,
|
||||
try:
|
||||
from chromadb import HttpClient
|
||||
except ImportError as e:
|
||||
msg = "Could not import chromadb. Please install it with `pip install chromadb`."
|
||||
raise ImportError(msg) from e
|
||||
client = HttpClient(
|
||||
host=self.chroma_server_host,
|
||||
port=self.chroma_server_http_port or 8000,
|
||||
ssl=bool(self.chroma_server_ssl_enabled),
|
||||
)
|
||||
client = Client(settings=chroma_settings)
|
||||
|
||||
# Check persist_directory and expand it if it is a relative path
|
||||
persist_directory = self.resolve_path(self.persist_directory) if self.persist_directory is not None else None
|
||||
|
||||
@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from pydantic import ValidationError
|
||||
@ -20,7 +21,7 @@ from lfx.base.models.unified_models import (
|
||||
)
|
||||
from lfx.base.models.watsonx_constants import IBM_WATSONX_URLS
|
||||
from lfx.components.agentics.helpers.model_config import validate_model_selection
|
||||
from lfx.components.helpers import CurrentDateComponent
|
||||
from lfx.components.helpers import CalculatorComponent, CurrentDateComponent
|
||||
from lfx.components.langchain_utilities.tool_calling import ToolCallingAgentComponent
|
||||
from lfx.custom.custom_component.component import get_component_toolkit
|
||||
from lfx.field_typing.range_spec import RangeSpec
|
||||
@ -84,8 +85,14 @@ class AgentComponent(ToolCallingAgentComponent):
|
||||
MultilineInput(
|
||||
name="system_prompt",
|
||||
display_name="Agent Instructions",
|
||||
info="System Prompt: Initial instructions and context provided to guide the agent's behavior.",
|
||||
value="You are a helpful assistant that can use tools to answer questions and perform tasks.",
|
||||
info=(
|
||||
"System Prompt: Initial instructions and context provided to guide the agent's behavior. "
|
||||
"Supports dynamic placeholders: {current_date}, {model_name}."
|
||||
),
|
||||
value=(
|
||||
"You are a helpful assistant that can use tools to answer questions and perform tasks. "
|
||||
"Today is {current_date}. You are powered by {model_name}."
|
||||
),
|
||||
advanced=False,
|
||||
),
|
||||
MessageTextInput(
|
||||
@ -181,6 +188,16 @@ class AgentComponent(ToolCallingAgentComponent):
|
||||
info="If true, will add a tool to the agent that returns the current date.",
|
||||
value=True,
|
||||
),
|
||||
BoolInput(
|
||||
name="add_calculator_tool",
|
||||
display_name="Calculator",
|
||||
advanced=True,
|
||||
info=(
|
||||
"If true, adds a zero-config arithmetic calculator tool to the agent "
|
||||
"(safe: only +, -, *, /, ** operators via AST)."
|
||||
),
|
||||
value=True,
|
||||
),
|
||||
]
|
||||
outputs = [
|
||||
Output(name="response", display_name="Response", method="message_response"),
|
||||
@ -274,11 +291,60 @@ class AgentComponent(ToolCallingAgentComponent):
|
||||
raise TypeError(msg)
|
||||
self.tools.append(current_date_tool)
|
||||
|
||||
# Add calculator tool if enabled (zero-config arithmetic)
|
||||
if getattr(self, "add_calculator_tool", False):
|
||||
if not isinstance(self.tools, list): # type: ignore[has-type]
|
||||
self.tools = []
|
||||
calculator_tool = (await CalculatorComponent(**self.get_base_args()).to_toolkit()).pop(0)
|
||||
|
||||
if not isinstance(calculator_tool, StructuredTool):
|
||||
msg = "CalculatorComponent must be converted to a StructuredTool"
|
||||
raise TypeError(msg)
|
||||
self.tools.append(calculator_tool)
|
||||
|
||||
# Set shared callbacks for tracing the tools used by the agent
|
||||
self.set_tools_callbacks(self.tools, self._get_shared_callbacks())
|
||||
|
||||
return llm_model, self.chat_history, self.tools
|
||||
|
||||
def _get_resolved_model_name(self) -> str:
|
||||
"""Best-effort human-readable model name for {model_name} injection."""
|
||||
try:
|
||||
from langchain_core.language_models import BaseLanguageModel
|
||||
|
||||
if isinstance(self.model, BaseLanguageModel):
|
||||
return type(self.model).__name__
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
if isinstance(self.model, list) and self.model:
|
||||
first = self.model[0]
|
||||
if isinstance(first, dict):
|
||||
name = first.get("name")
|
||||
if isinstance(name, str) and name:
|
||||
return name
|
||||
|
||||
legacy_model_name = getattr(self, "model_name", None)
|
||||
if isinstance(legacy_model_name, str) and legacy_model_name:
|
||||
return legacy_model_name
|
||||
return ""
|
||||
|
||||
def _inject_dynamic_prompt_values(self, prompt: str | None) -> str | None:
|
||||
"""Replace {current_date} / {model_name} placeholders in the system prompt.
|
||||
|
||||
Uses str.replace (not str.format) so user prompts containing literal braces
|
||||
such as JSON examples ({"key": 1}) never break the agent.
|
||||
"""
|
||||
if not prompt:
|
||||
return prompt
|
||||
replacements = {
|
||||
"{current_date}": datetime.now(tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC"),
|
||||
"{model_name}": self._get_resolved_model_name(),
|
||||
}
|
||||
for placeholder, value in replacements.items():
|
||||
prompt = prompt.replace(placeholder, value)
|
||||
return prompt
|
||||
|
||||
async def message_response(self) -> Message:
|
||||
try:
|
||||
llm_model, self.chat_history, self.tools = await self.get_agent_requirements()
|
||||
@ -288,7 +354,7 @@ class AgentComponent(ToolCallingAgentComponent):
|
||||
tools=self.tools or [],
|
||||
chat_history=self.chat_history,
|
||||
input_value=self.input_value,
|
||||
system_prompt=self.system_prompt,
|
||||
system_prompt=self._inject_dynamic_prompt_values(self.system_prompt),
|
||||
)
|
||||
agent = self.create_agent_runnable()
|
||||
result = await self.run_agent(agent)
|
||||
@ -392,8 +458,10 @@ class AgentComponent(ToolCallingAgentComponent):
|
||||
try:
|
||||
system_components = []
|
||||
|
||||
# 1. Agent Instructions (system_prompt)
|
||||
agent_instructions = getattr(self, "system_prompt", "") or ""
|
||||
# 1. Agent Instructions (system_prompt).
|
||||
# Inject dynamic placeholders HERE so user-authored format_instructions
|
||||
# and schema descriptions appended later keep their literal {...} tokens.
|
||||
agent_instructions = self._inject_dynamic_prompt_values(getattr(self, "system_prompt", "") or "") or ""
|
||||
if agent_instructions:
|
||||
system_components.append(f"{agent_instructions}")
|
||||
|
||||
@ -423,7 +491,9 @@ class AgentComponent(ToolCallingAgentComponent):
|
||||
except (ValidationError, ValueError, TypeError, KeyError) as e:
|
||||
await logger.aerror(f"Could not build schema for prompt: {e}", exc_info=True)
|
||||
|
||||
# Combine all components
|
||||
# Combine all components. Injection already applied on agent_instructions
|
||||
# above; do NOT re-run it here so literal {...} tokens in format
|
||||
# instructions / schema descriptions stay intact.
|
||||
combined_instructions = "\n\n".join(system_components) if system_components else ""
|
||||
llm_model, self.chat_history, self.tools = await self.get_agent_requirements()
|
||||
self.set(
|
||||
@ -542,6 +612,7 @@ class AgentComponent(ToolCallingAgentComponent):
|
||||
"tools",
|
||||
"input_value",
|
||||
"add_current_date_tool",
|
||||
"add_calculator_tool",
|
||||
"system_prompt",
|
||||
"agent_description",
|
||||
"max_iterations",
|
||||
|
||||
@ -193,14 +193,7 @@ class YouTubeCommentsComponent(Component):
|
||||
else:
|
||||
request = None
|
||||
|
||||
# Convert to DataFrame
|
||||
comments_df = pd.DataFrame(comments_data)
|
||||
|
||||
# Add video metadata
|
||||
comments_df["video_id"] = video_id
|
||||
comments_df["video_url"] = self.video_url
|
||||
|
||||
# Sort columns for better organization
|
||||
# Define column order
|
||||
column_order = [
|
||||
"video_id",
|
||||
"video_url",
|
||||
@ -217,7 +210,20 @@ class YouTubeCommentsComponent(Component):
|
||||
if self.include_metrics:
|
||||
column_order.extend(["like_count", "reply_count"])
|
||||
|
||||
comments_df = comments_df[column_order]
|
||||
# Handle empty comments case
|
||||
if not comments_data:
|
||||
# Create empty DataFrame with proper columns
|
||||
comments_df = pd.DataFrame(columns=column_order)
|
||||
else:
|
||||
# Convert to DataFrame
|
||||
comments_df = pd.DataFrame(comments_data)
|
||||
|
||||
# Add video metadata
|
||||
comments_df["video_id"] = video_id
|
||||
comments_df["video_url"] = self.video_url
|
||||
|
||||
# Reorder columns
|
||||
comments_df = comments_df[column_order]
|
||||
|
||||
return DataFrame(comments_df)
|
||||
|
||||
|
||||
@ -272,8 +272,9 @@ def create_class(code, class_name):
|
||||
module = ast.parse(code)
|
||||
exec_globals = prepare_global_scope(module)
|
||||
|
||||
future_imports = [n for n in module.body if isinstance(n, ast.ImportFrom) and n.module == "__future__"]
|
||||
class_code = extract_class_code(module, class_name)
|
||||
compiled_class = compile_class_code(class_code)
|
||||
compiled_class = compile_class_code(class_code, future_imports)
|
||||
|
||||
return build_class_constructor(compiled_class, exec_globals, class_name)
|
||||
|
||||
@ -394,11 +395,15 @@ def prepare_global_scope(module):
|
||||
exec_globals = globals().copy()
|
||||
imports = []
|
||||
import_froms = []
|
||||
future_imports = []
|
||||
definitions = []
|
||||
|
||||
for node in module.body:
|
||||
if isinstance(node, ast.Import):
|
||||
imports.append(node)
|
||||
elif isinstance(node, ast.ImportFrom) and node.module == "__future__":
|
||||
# __future__ imports are compiler directives — collect separately
|
||||
future_imports.append(node)
|
||||
elif isinstance(node, ast.ImportFrom) and node.module is not None:
|
||||
import_froms.append(node)
|
||||
elif isinstance(node, ast.ClassDef | ast.FunctionDef | ast.Assign | ast.AnnAssign):
|
||||
@ -468,7 +473,8 @@ def prepare_global_scope(module):
|
||||
raise ModuleNotFoundError(msg)
|
||||
|
||||
if definitions:
|
||||
combined_module = ast.Module(body=definitions, type_ignores=[])
|
||||
# Prepend __future__ imports so compiler directives (e.g. PEP 563 annotations) take effect
|
||||
combined_module = ast.Module(body=future_imports + definitions, type_ignores=[])
|
||||
compiled_code = compile(combined_module, "<string>", "exec")
|
||||
exec(compiled_code, exec_globals)
|
||||
|
||||
@ -491,16 +497,18 @@ def extract_class_code(module, class_name):
|
||||
return class_code
|
||||
|
||||
|
||||
def compile_class_code(class_code):
|
||||
def compile_class_code(class_code, future_imports=None):
|
||||
"""Compiles the AST node of a class into a code object.
|
||||
|
||||
Args:
|
||||
class_code: AST node of the class
|
||||
future_imports: Optional list of __future__ ImportFrom nodes to prepend as compiler directives
|
||||
|
||||
Returns:
|
||||
Compiled code object of the class
|
||||
"""
|
||||
return compile(ast.Module(body=[class_code], type_ignores=[]), "<string>", "exec")
|
||||
body = (future_imports or []) + [class_code]
|
||||
return compile(ast.Module(body=body, type_ignores=[]), "<string>", "exec")
|
||||
|
||||
|
||||
def build_class_constructor(compiled_class, exec_globals, class_name):
|
||||
|
||||
@ -41,6 +41,9 @@ class ResultData(BaseModel):
|
||||
if message is None:
|
||||
continue
|
||||
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
|
||||
if "stream_url" in message and "type" in message:
|
||||
stream_url = StreamURL(location=message["stream_url"])
|
||||
values["outputs"].update({key: OutputValue(message=stream_url, type=message["type"])})
|
||||
|
||||
@ -111,7 +111,7 @@ def build_output_logs(vertex, result) -> dict:
|
||||
type_ = get_type(output_result)
|
||||
|
||||
match type_:
|
||||
case LogType.STREAM if "stream_url" in message:
|
||||
case LogType.STREAM if isinstance(message, dict) and "stream_url" in message:
|
||||
message = StreamURL(location=message["stream_url"])
|
||||
|
||||
case LogType.STREAM:
|
||||
@ -124,9 +124,13 @@ def build_output_logs(vertex, result) -> dict:
|
||||
message = ""
|
||||
|
||||
case LogType.ARRAY:
|
||||
if isinstance(message, DataFrame):
|
||||
if message is None:
|
||||
message = []
|
||||
elif isinstance(message, DataFrame):
|
||||
message = message.to_dict(orient="records")
|
||||
message = [serialize(item) for item in message]
|
||||
message = [serialize(item) for item in message]
|
||||
else:
|
||||
message = [serialize(item) for item in message]
|
||||
name = output.get("name", f"output_{index}")
|
||||
outputs |= {name: OutputValue(message=message, type=type_).model_dump()}
|
||||
|
||||
|
||||
@ -319,6 +319,7 @@ class Settings(BaseSettings):
|
||||
max_flow_version_entries_per_flow: int = 50
|
||||
"""Max version history entries per flow. Oldest entries pruned on next snapshot.
|
||||
|
||||
|
||||
If retroactively lowered below the current count for a flow,
|
||||
the oldest entries are deleted only when the next entry is created.
|
||||
"""
|
||||
@ -336,6 +337,7 @@ class Settings(BaseSettings):
|
||||
max_items_length: int = MAX_ITEMS_LENGTH
|
||||
"""Maximum number of items to store and display in the UI. Lists longer than this
|
||||
will be truncated when displayed in the UI. Does not affect data passed between components nor outputs."""
|
||||
max_ingestion_timeout_secs: int = 600
|
||||
|
||||
# MCP Server
|
||||
mcp_server_enabled: bool = True
|
||||
|
||||
@ -39,6 +39,34 @@ class TestLoggingComponent(Component):
|
||||
assert result.__name__ == "TestLoggingComponent"
|
||||
|
||||
|
||||
def test_create_class_future_annotations_with_type_checking():
|
||||
"""Regression test for issue #12776.
|
||||
|
||||
`from __future__ import annotations` must act as a compiler directive so that TYPE_CHECKING-only
|
||||
imports don't raise NameError at classdefinition time.
|
||||
"""
|
||||
code = dedent("""
|
||||
from __future__ import annotations
|
||||
from typing import TYPE_CHECKING
|
||||
from langflow.custom import Component
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import List
|
||||
|
||||
class TypeCheckingComponent(Component):
|
||||
display_name = "Test"
|
||||
|
||||
def build(self, value: List[str]) -> str:
|
||||
return str(value)
|
||||
""")
|
||||
result = create_class(code, "TypeCheckingComponent")
|
||||
assert result.__name__ == "TypeCheckingComponent"
|
||||
# With PEP 563 active, annotations should be stored as strings rather than evaluated
|
||||
hints = result.build.__annotations__
|
||||
assert hints.get("value") == "List[str]"
|
||||
assert hints.get("return") == "str"
|
||||
|
||||
|
||||
def test_execute_function_supports_aliased_dotted_imports():
|
||||
code = dedent("""
|
||||
import urllib.request as request
|
||||
|
||||
13
src/lfx/tests/unit/graph/test_result_data.py
Normal file
13
src/lfx/tests/unit/graph/test_result_data.py
Normal file
@ -0,0 +1,13 @@
|
||||
import threading
|
||||
|
||||
from lfx.graph.schema import ResultData
|
||||
|
||||
|
||||
def test_result_data_ignores_non_dict_artifact_values():
|
||||
"""Vector DB artifacts may include non-iterable objects such as locks."""
|
||||
lock = threading.Lock()
|
||||
|
||||
result = ResultData(artifacts={"vector_db": lock})
|
||||
|
||||
assert result.artifacts == {"vector_db": lock}
|
||||
assert result.outputs == {}
|
||||
Reference in New Issue
Block a user