From a3b4cf33607c752d97b55f57097145e9bc457134 Mon Sep 17 00:00:00 2001 From: Ayush8923 <80516839+Ayush8923@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:23:39 +0530 Subject: [PATCH] feat: rag agent flow --- .../090_add_chatbot_session_agent_type.py | 47 +++ .../app/api/docs/chatbot/create_session.md | 4 +- .../api/docs/chatbot/get_latest_session.md | 2 +- .../api/docs/chatbot/send_session_message.md | 4 +- backend/app/api/routes/chatbot/session.py | 21 +- backend/app/crud/chatbot.py | 3 + backend/app/models/__init__.py | 1 + backend/app/models/chatbot.py | 25 ++ .../chatbot/agents/onboarding/graph.py | 44 +-- .../services/chatbot/agents/rag/__init__.py | 0 .../app/services/chatbot/agents/rag/graph.py | 179 ++++++++++ .../services/chatbot/agents/rag/prompts.py | 35 ++ .../services/chatbot/agents/rag/retriever.py | 216 ++++++++++++ .../app/services/chatbot/agents/rag/types.py | 24 ++ .../app/services/chatbot/agents/registry.py | 55 +++ backend/app/services/chatbot/config.py | 63 +++- backend/app/services/chatbot/history.py | 46 ++- .../app/services/chatbot/provider_errors.py | 68 ++++ backend/app/services/chatbot/session.py | 35 +- backend/app/services/chatbot/types.py | 12 +- backend/app/services/chatbot/utils.py | 4 + .../tests/api/routes/chatbot/test_session.py | 206 +++++++++++ .../chatbot/agents/onboarding/test_graph.py | 22 +- .../services/chatbot/agents/rag/__init__.py | 0 .../services/chatbot/agents/rag/test_graph.py | 319 ++++++++++++++++++ .../chatbot/agents/rag/test_retriever.py | 243 +++++++++++++ .../services/chatbot/agents/test_registry.py | 92 +++++ .../app/tests/services/chatbot/test_config.py | 127 ++++++- .../tests/services/chatbot/test_session.py | 170 ++++++++++ backend/app/tests/utils/chatbot.py | 88 ++++- backend/app/tests/utils/collection.py | 5 +- docs/wiki/domain-map.md | 2 +- docs/wiki/modules/chatbot.md | 36 +- 33 files changed, 2084 insertions(+), 114 deletions(-) create mode 100644 backend/app/alembic/versions/090_add_chatbot_session_agent_type.py create mode 100644 backend/app/services/chatbot/agents/rag/__init__.py create mode 100644 backend/app/services/chatbot/agents/rag/graph.py create mode 100644 backend/app/services/chatbot/agents/rag/prompts.py create mode 100644 backend/app/services/chatbot/agents/rag/retriever.py create mode 100644 backend/app/services/chatbot/agents/rag/types.py create mode 100644 backend/app/services/chatbot/agents/registry.py create mode 100644 backend/app/services/chatbot/provider_errors.py create mode 100644 backend/app/tests/services/chatbot/agents/rag/__init__.py create mode 100644 backend/app/tests/services/chatbot/agents/rag/test_graph.py create mode 100644 backend/app/tests/services/chatbot/agents/rag/test_retriever.py create mode 100644 backend/app/tests/services/chatbot/agents/test_registry.py diff --git a/backend/app/alembic/versions/090_add_chatbot_session_agent_type.py b/backend/app/alembic/versions/090_add_chatbot_session_agent_type.py new file mode 100644 index 000000000..6b89a1740 --- /dev/null +++ b/backend/app/alembic/versions/090_add_chatbot_session_agent_type.py @@ -0,0 +1,47 @@ +"""add agent_type to chatbot_session + +Revision ID: 090 +Revises: 089 +Create Date: 2026-10-28 00:00:00.000000 + +Sessions created before this revision all ran the onboarding agent, so the +server default backfills them as 'onboarding' and keeps the column NOT NULL. +""" + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + +revision = "090" +down_revision = "089" +branch_labels = None +depends_on = None + +_AGENT_TYPE_ENUM_NAME = "chatbotagenttypeenum" +_AGENT_TYPE_VALUES = ("onboarding", "rag") +_DEFAULT_AGENT_TYPE = "onboarding" + + +def upgrade(): + agent_type_enum = postgresql.ENUM( + *_AGENT_TYPE_VALUES, + name=_AGENT_TYPE_ENUM_NAME, + create_type=False, + ) + agent_type_enum.create(op.get_bind(), checkfirst=True) + + op.add_column( + "chatbot_session", + sa.Column( + "agent_type", + agent_type_enum, + nullable=False, + server_default=_DEFAULT_AGENT_TYPE, + comment="Agent that runs each turn: onboarding (collects the fields the system_prompt asks for) or rag (answers questions from the config's knowledge bases)", + ), + ) + + +def downgrade(): + op.drop_column("chatbot_session", "agent_type") + sa.Enum(name=_AGENT_TYPE_ENUM_NAME).drop(op.get_bind(), checkfirst=True) diff --git a/backend/app/api/docs/chatbot/create_session.md b/backend/app/api/docs/chatbot/create_session.md index 69346c6ae..0dae93a0f 100644 --- a/backend/app/api/docs/chatbot/create_session.md +++ b/backend/app/api/docs/chatbot/create_session.md @@ -2,8 +2,10 @@ Start a server-side chatbot conversation and return its `session_id` and `chatbo The `system_prompt` and the `config_id` / `config_version` pin are stored on the session, so later turns only send the new user message. `config_id` must be a config in the caller's project whose completion is `provider: "anthropic"`, `type: "text"` with `params.model` set; when `config_version` is omitted, each turn uses the config's latest version at that time. +`agent_type` picks the agent that runs every turn and defaults to `onboarding`, which collects the fields the `system_prompt` asks for. `rag` answers the user's questions from the config's knowledge bases instead, and the `system_prompt` is its goal: it decides whether answers stay within the documents or may use general knowledge. A `rag` session needs a config with non-empty `params.knowledge_base_ids`, each the `knowledge_base_id` of a non-deleted OpenAI collection in the caller's project, plus OpenAI credentials on the project, which are used to search those knowledge bases. `params.instructions` and `params.max_num_results` (default 5, at most 50 passages per search) are also applied. + `name` is the chatbot's display name only; omit it to get `"Untitled Chatbot"`. `contact_ref` is the caller's own identifier for the end user and is stored as-is. `chatbot_ref` identifies the chatbot. Omit it to create a new chatbot: Kaapi generates a ref of the form `LLMResponseWithPrompt-<5 hex chars>` (e.g. `LLMResponseWithPrompt-913f0`), unique within the project, and returns it. To start a new session for the same chatbot, send that returned `chatbot_ref` back; it is used as-is. -A contact has at most one active session per chatbot. Creating a session marks the contact's previous `active` sessions with the same `chatbot_ref` as `inactive`, in the same transaction; those still `in_progress` become `abandoned`, while `complete` ones stay `complete`. Sessions of other chatbots or other contacts are untouched. Returns 404 if the config or version isn't found in the project, 422 if it isn't a usable Anthropic text config, 409 if another session for the same `contact_ref` + `chatbot_ref` was created concurrently (retry the request), and 503 in the unlikely case no unused `chatbot_ref` could be generated (retry). +A contact has at most one active session per chatbot. Creating a session marks the contact's previous `active` sessions with the same `chatbot_ref` as `inactive`, in the same transaction; those still `in_progress` become `abandoned`, while `complete` ones stay `complete`. Sessions of other chatbots or other contacts are untouched. Returns 404 if the config or version isn't found in the project, 422 if it isn't a usable Anthropic text config (or, for `rag`, has no `knowledge_base_ids`, names knowledge bases that aren't this project's OpenAI collections, or the project has no OpenAI credentials), 409 if another session for the same `contact_ref` + `chatbot_ref` was created concurrently (retry the request), and 503 in the unlikely case no unused `chatbot_ref` could be generated (retry). diff --git a/backend/app/api/docs/chatbot/get_latest_session.md b/backend/app/api/docs/chatbot/get_latest_session.md index 3a3aa41a8..e8289f7c8 100644 --- a/backend/app/api/docs/chatbot/get_latest_session.md +++ b/backend/app/api/docs/chatbot/get_latest_session.md @@ -2,4 +2,4 @@ Get the most recently created chatbot session for the given `contact_ref` in the `contact_ref` is required; `chatbot_ref` is optional. Both are matched exactly. `chatbot_ref` is the value returned when the chatbot's first session was created (or the one sent on later creates). Omit it to get the contact's newest session across all chatbots — the response's `chatbot_ref` says which chatbot that is. -The response carries the session's `system_prompt` (as stored at creation), its current `status` (`active`, or `inactive` once a newer session for the same contact and chatbot was created), its `completion_status` (`in_progress` until every field the system prompt asks for is collected, then `complete`; `abandoned` if it was superseded while still `in_progress`), and the flat `collected_fields` map (`{field_name: value}`) gathered so far. Returns 404 when no matching session exists in this project. +The response carries the session's `system_prompt` (as stored at creation), its `agent_type` (`onboarding` or `rag`), its current `status` (`active`, or `inactive` once a newer session for the same contact and chatbot was created), its `completion_status` (`in_progress` until every field the system prompt asks for is collected, then `complete`; `abandoned` if it was superseded while still `in_progress`), and the flat `collected_fields` map (`{field_name: value}`) gathered so far. `rag` sessions don't collect fields, so they stay `in_progress` with an empty map. Returns 404 when no matching session exists in this project. diff --git a/backend/app/api/docs/chatbot/send_session_message.md b/backend/app/api/docs/chatbot/send_session_message.md index c6ff7b84f..67d7c52a9 100644 --- a/backend/app/api/docs/chatbot/send_session_message.md +++ b/backend/app/api/docs/chatbot/send_session_message.md @@ -1,7 +1,7 @@ Queue one turn of a chatbot session; the bot reply is delivered asynchronously to `webhook_url`. -The user `message` is stored immediately and the response is `202` with `status: "queued"`. Once the turn completes, Kaapi POSTs `{"success": true, "data": {"session_id", "chatbot_ref", "reply", "status", "completion_status", "collected_fields"}}` to `webhook_url` (HMAC-signed when the project has a webhook secret). On a model or config failure it POSTs `{"success": false, "error": "...", "data": {"session_id"}}` instead. +The user `message` is stored immediately and the response is `202` with `status: "queued"`. Once the turn completes, Kaapi POSTs `{"success": true, "data": {"session_id", "chatbot_ref", "agent_type", "reply", "status", "completion_status", "collected_fields", "sources"}}` to `webhook_url` (HMAC-signed when the project has a webhook secret). On a model, config, or knowledge-base search failure it POSTs `{"success": false, "error": "...", "data": {"session_id"}}` instead. -Each turn re-reads the conversation against the session's `system_prompt` and saves every field the user has given into `collected_fields`, a flat `{field_name: value}` map with snake_case keys (values are only ever taken from what the user said; a correction overwrites the old value). `completion_status` is `complete` once every required field is collected and `in_progress` otherwise; it is computed from a structured extraction, not from the reply text. A `complete` session still accepts messages. `status` (`active`/`inactive`) is separate and says whether the session is in use: once a newer session is created for the same `contact_ref` + `chatbot_ref`, this one is `inactive` and rejects new messages with 409. A turn that was already queued when that happened is not run; Kaapi POSTs the failure payload with error `[KAAPI] Chatbot session was superseded by a newer session; start a new session.` instead. +For an `onboarding` session, each turn re-reads the conversation against the session's `system_prompt` and saves every field the user has given into `collected_fields`, a flat `{field_name: value}` map with snake_case keys (values are only ever taken from what the user said; a correction overwrites the old value). `completion_status` is `complete` once every required field is collected and `in_progress` otherwise; it is computed from a structured extraction, not from the reply text. A `complete` session still accepts messages. For a `rag` session, the bot searches the config's knowledge bases when the question may be covered by them and answers from the retrieved passages; `collected_fields` is left unchanged and `completion_status` stays `in_progress`. `sources` lists the documents retrieved for the turn as `[{"file_id", "filename"}]`, deduplicated by `file_id` (always `[]` for `onboarding`, and for `rag` turns that didn't search). `status` (`active`/`inactive`) is separate and says whether the session is in use: once a newer session is created for the same `contact_ref` + `chatbot_ref`, this one is `inactive` and rejects new messages with 409. A turn that was already queued when that happened is not run; Kaapi POSTs the failure payload with error `[KAAPI] Chatbot session was superseded by a newer session; start a new session.` instead. On the first call of a session, omit `message` to get the bot's opening greeting; after that `message` is required. `webhook_url` must be a public HTTPS URL. Returns 404 if the session isn't in the caller's project, 409 if the session is inactive, and 422 for an invalid `webhook_url` or a missing `message` on a started session. diff --git a/backend/app/api/routes/chatbot/session.py b/backend/app/api/routes/chatbot/session.py index efa8948e8..de07f1636 100644 --- a/backend/app/api/routes/chatbot/session.py +++ b/backend/app/api/routes/chatbot/session.py @@ -17,6 +17,7 @@ from app.models import ( CHATBOT_REF_MAX_LENGTH, CONTACT_REF_MAX_LENGTH, + ChatbotAgentTypeEnum, ChatbotMessageRoleEnum, ChatbotSessionStatusEnum, ChatbotSessionSummary, @@ -26,7 +27,10 @@ SendChatbotSessionMessageResponse, ) from app.services.chatbot.chatbot_ref import generate_chatbot_ref -from app.services.chatbot.config import resolve_chatbot_model_settings +from app.services.chatbot.config import ( + prepare_rag_search_client, + resolve_chatbot_model_settings, +) from app.services.chatbot.exceptions import ( ChatbotInputError, ChatbotRefGenerationError, @@ -51,9 +55,10 @@ def start_chatbot_session( session: SessionDep, auth_context: AuthContextDep, ) -> APIResponse[CreateChatbotSessionResponse]: + organization_id = auth_context.organization_.id project_id = auth_context.project_.id try: - resolve_chatbot_model_settings( + model_settings = resolve_chatbot_model_settings( session=session, project_id=project_id, config_id=request.config_id, @@ -61,6 +66,14 @@ def start_chatbot_session( max_tokens=None, history_token_limit=None, ) + # Fail at create rather than on the first turn's webhook. + if request.agent_type == ChatbotAgentTypeEnum.RAG: + prepare_rag_search_client( + session=session, + organization_id=organization_id, + project_id=project_id, + model_settings=model_settings, + ) except ChatbotInputError as e: raise HTTPException(status_code=422, detail=str(e)) @@ -74,7 +87,7 @@ def start_chatbot_session( try: chatbot_session = create_chatbot_session( session, - organization_id=auth_context.organization_.id, + organization_id=organization_id, project_id=project_id, name=request.name, contact_ref=request.contact_ref, @@ -82,6 +95,7 @@ def start_chatbot_session( config_id=request.config_id, config_version=request.config_version, system_prompt=request.system_prompt, + agent_type=request.agent_type, ) except ChatbotSessionConflictError as e: raise HTTPException(status_code=409, detail=str(e)) @@ -138,6 +152,7 @@ def read_latest_chatbot_session( contact_ref=chatbot_session.contact_ref, chatbot_ref=chatbot_session.chatbot_ref, system_prompt=chatbot_session.system_prompt, + agent_type=chatbot_session.agent_type, status=chatbot_session.status, completion_status=chatbot_session.completion_status, collected_fields=chatbot_session.collected_fields, diff --git a/backend/app/crud/chatbot.py b/backend/app/crud/chatbot.py index 347bf66d8..f63170c46 100644 --- a/backend/app/crud/chatbot.py +++ b/backend/app/crud/chatbot.py @@ -9,6 +9,7 @@ from app.core.util import is_constraint_violation, now from app.models import ( ACTIVE_CHATBOT_SESSION_UNIQUE_INDEX, + ChatbotAgentTypeEnum, ChatbotCompletionStatusEnum, ChatbotMessageRoleEnum, ChatbotSession, @@ -76,6 +77,7 @@ def create_chatbot_session( config_id: UUID, config_version: int | None, system_prompt: str, + agent_type: ChatbotAgentTypeEnum = ChatbotAgentTypeEnum.ONBOARDING, ) -> ChatbotSession: """Insert a new active session, superseding the contact's active ones for the same chatbot.""" timestamp = now() @@ -96,6 +98,7 @@ def create_chatbot_session( config_id=config_id, config_version=config_version, system_prompt=system_prompt, + agent_type=agent_type, inserted_at=timestamp, updated_at=timestamp, ) diff --git a/backend/app/models/__init__.py b/backend/app/models/__init__.py index 0a112ca9f..2fc8bbe97 100644 --- a/backend/app/models/__init__.py +++ b/backend/app/models/__init__.py @@ -40,6 +40,7 @@ CHATBOT_REF_MAX_LENGTH, CONTACT_REF_MAX_LENGTH, DEFAULT_CHATBOT_NAME, + ChatbotAgentTypeEnum, ChatbotCompletionStatusEnum, ChatbotMessageRoleEnum, ChatbotSession, diff --git a/backend/app/models/chatbot.py b/backend/app/models/chatbot.py index 2310e2ecb..0673dc606 100644 --- a/backend/app/models/chatbot.py +++ b/backend/app/models/chatbot.py @@ -27,6 +27,11 @@ class ChatbotCompletionStatusEnum(StrEnum): ABANDONED = "abandoned" +class ChatbotAgentTypeEnum(StrEnum): + ONBOARDING = "onboarding" + RAG = "rag" + + _SESSION_STATUS_PG_ENUM = ENUM( ChatbotSessionStatusEnum, name="chatbotsessionstatusenum", @@ -39,6 +44,12 @@ class ChatbotCompletionStatusEnum(StrEnum): values_callable=lambda enum_cls: [member.value for member in enum_cls], create_type=False, ) +_AGENT_TYPE_PG_ENUM = ENUM( + ChatbotAgentTypeEnum, + name="chatbotagenttypeenum", + values_callable=lambda enum_cls: [member.value for member in enum_cls], + create_type=False, +) _MESSAGE_ROLE_PG_ENUM = ENUM( ChatbotMessageRoleEnum, name="chatbotmessageroleenum", @@ -118,6 +129,15 @@ class ChatbotSession(SQLModel, table=True): "comment": "Admin-authored flow prompt captured at session creation" }, ) + agent_type: ChatbotAgentTypeEnum = Field( + default=ChatbotAgentTypeEnum.ONBOARDING, + sa_column=Column( + _AGENT_TYPE_PG_ENUM, + nullable=False, + server_default=ChatbotAgentTypeEnum.ONBOARDING.value, + comment="Agent that runs each turn: onboarding (collects the fields the system_prompt asks for) or rag (answers questions from the config's knowledge bases)", + ), + ) status: ChatbotSessionStatusEnum = Field( default=ChatbotSessionStatusEnum.ACTIVE, sa_column=Column( @@ -236,6 +256,10 @@ class CreateChatbotSessionRequest(SQLModel): min_length=1, description="Admin-authored flow prompt (goal, fields to collect, rules); stored for the session's lifetime.", ) + agent_type: ChatbotAgentTypeEnum = Field( + default=ChatbotAgentTypeEnum.ONBOARDING, + description="Agent that runs the session: `onboarding` collects the fields `system_prompt` asks for; `rag` answers questions from the config's `knowledge_base_ids` (the project needs OpenAI credentials).", + ) class CreateChatbotSessionResponse(SQLModel): @@ -249,6 +273,7 @@ class ChatbotSessionSummary(SQLModel): contact_ref: str chatbot_ref: str system_prompt: str + agent_type: ChatbotAgentTypeEnum status: ChatbotSessionStatusEnum completion_status: ChatbotCompletionStatusEnum collected_fields: dict[str, Any] diff --git a/backend/app/services/chatbot/agents/onboarding/graph.py b/backend/app/services/chatbot/agents/onboarding/graph.py index 2c1c69893..4e2efd525 100644 --- a/backend/app/services/chatbot/agents/onboarding/graph.py +++ b/backend/app/services/chatbot/agents/onboarding/graph.py @@ -4,7 +4,6 @@ import json from typing import Any -import anthropic from langchain_core.exceptions import OutputParserException from langchain_core.language_models import LanguageModelInput from langchain_core.language_models.chat_models import BaseChatModel @@ -28,14 +27,16 @@ FieldExtraction, OnboardingState, ) -from app.services.chatbot.exceptions import ChatbotInputError, ChatbotProviderError +from app.services.chatbot.exceptions import ChatbotProviderError from app.services.chatbot.history import ( + extract_final_reply, has_user_turn, insert_after_leading_system, summarize_if_needed, to_langchain_messages, - to_message_dicts, + validate_turn_input, ) +from app.services.chatbot.provider_errors import CHAT_MODEL_ERRORS, to_provider_error from app.services.chatbot.types import ChatbotMessageDict, ChatbotTurnResult from app.services.chatbot.utils import ChatbotModelSettings, get_chat_model @@ -166,28 +167,6 @@ def _get_extraction_model(chat_model: BaseChatModel) -> ExtractionModel: return chat_model.with_structured_output(FieldExtraction, method=EXTRACTION_METHOD) -def _to_provider_error(error: Exception, *, model: str) -> ChatbotProviderError: - if isinstance(error, anthropic.APITimeoutError): - return ChatbotProviderError( - "[KAAPI] Request timed out — retry smaller. If persistent, contact Kaapi. " - "(code: APITimeoutError)" - ) - if isinstance(error, anthropic.APIConnectionError): - return ChatbotProviderError( - "[KAAPI] Network/DNS issue reaching provider — check connectivity. " - "If persistent, contact Kaapi. (code: APIConnectionError)" - ) - if isinstance(error, anthropic.APIStatusError): - return ChatbotProviderError( - f"[ANTHROPIC] Chat model call failed: {error.message} " - f"(code: {error.status_code}, request_id: {error.request_id})" - ) - return ChatbotProviderError( - f"[KAAPI] Unsupported model configuration for '{model}': {error} " - "— update the config's model, thinking, or effort. (code: ValueError)" - ) - - async def run_chatbot_turn( *, system_prompt: str, @@ -197,11 +176,7 @@ async def run_chatbot_turn( collected_fields: dict[str, Any] | None = None, ) -> ChatbotTurnResult: """Run one stateless turn: update the collected fields, then reply.""" - if user_message is None and history: - raise ChatbotInputError( - "[KAAPI] `message` is required once `history` is non-empty; omit it only " - "on the first call to get the opening greeting." - ) + validate_turn_input(history=history, user_message=user_message) messages = to_langchain_messages( system_prompt=system_prompt, history=history, user_message=user_message @@ -222,14 +197,11 @@ async def run_chatbot_turn( ) # ChatAnthropic.invoke is sync; keep it off the event loop. result = await asyncio.to_thread(graph.invoke, initial_state) - except (anthropic.APIConnectionError, anthropic.APIStatusError, ValueError) as e: - raise _to_provider_error(e, model=model_settings.model) from e + except CHAT_MODEL_ERRORS as e: + raise to_provider_error(e, model=model_settings.model) from e - transcript = to_message_dicts( - system_prompt=system_prompt, messages=result["messages"] - ) return ChatbotTurnResult( - transcript=transcript, + reply=extract_final_reply(result["messages"]), collected_fields=result["collected_fields"], is_complete=result.get("is_complete", False), ) diff --git a/backend/app/services/chatbot/agents/rag/__init__.py b/backend/app/services/chatbot/agents/rag/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/backend/app/services/chatbot/agents/rag/graph.py b/backend/app/services/chatbot/agents/rag/graph.py new file mode 100644 index 000000000..2a390af3e --- /dev/null +++ b/backend/app/services/chatbot/agents/rag/graph.py @@ -0,0 +1,179 @@ +"""RAG turn: the model answers from the knowledge base, searching it through a tool loop.""" + +import asyncio +from typing import Any + +from langchain_anthropic import ChatAnthropic +from langchain_core.language_models import LanguageModelInput +from langchain_core.messages import ( + AIMessage, + BaseMessage, + RemoveMessage, + ToolMessage, +) +from langchain_core.runnables import Runnable +from langchain_core.tools import BaseTool +from langgraph.graph import END, START, StateGraph +from langgraph.graph.message import REMOVE_ALL_MESSAGES +from langgraph.graph.state import CompiledStateGraph +from langgraph.prebuilt import ToolNode, tools_condition +from langgraph.prebuilt.tool_node import ToolInvocationError +from openai import OpenAI + +from app.services.chatbot.agents.rag.prompts import ( + CONFIG_INSTRUCTIONS_TEMPLATE, + GOAL_TEMPLATE, + RAG_INSTRUCTION, +) +from app.services.chatbot.agents.rag.retriever import build_search_tool +from app.services.chatbot.agents.rag.types import KnowledgeBasePassage, RagState +from app.services.chatbot.history import ( + extract_final_reply, + summarize_if_needed, + to_langchain_messages, + validate_turn_input, +) +from app.services.chatbot.provider_errors import CHAT_MODEL_ERRORS, to_provider_error +from app.services.chatbot.types import ( + ChatbotMessageDict, + ChatbotSource, + ChatbotTurnResult, +) +from app.services.chatbot.utils import ChatbotModelSettings, get_chat_model + +ToolCallingModel = Runnable[LanguageModelInput, BaseMessage] + +PREPARE_CONTEXT_NODE = "prepare_context_node" +AGENT_NODE = "agent_node" +# tools_condition routes to this exact node name. +TOOLS_NODE = "tools" + +MAX_TOOL_ROUNDS = 3 +# Forced tool use is rejected with thinking on, so the model picks; "none" makes it +# answer. Only ChatAnthropic passes a dict tool_choice through as-is. +TOOL_CHOICE_NONE: dict[str, str] = {"type": "none"} + + +def build_rag_system_prompt(*, goal: str, instructions: str | None) -> str: + system_prompt = RAG_INSTRUCTION + if instructions: + system_prompt += CONFIG_INSTRUCTIONS_TEMPLATE.format(instructions=instructions) + system_prompt += GOAL_TEMPLATE.format(goal=goal) + return system_prompt + + +def _count_tool_rounds(messages: list[BaseMessage]) -> int: + # Stored history is plain text, so every tool call in state belongs to this turn. + rounds = 0 + for message in messages: + if isinstance(message, AIMessage) and message.tool_calls: + rounds += 1 + return rounds + + +def build_rag_graph( + chat_model: ChatAnthropic, + search_tool: BaseTool, + *, + history_token_limit: int, +) -> CompiledStateGraph: + """prepare_context_node -> agent_node <-> tools, ending once the model stops calling tools.""" + tool_model: ToolCallingModel = chat_model.bind_tools([search_tool]) + # Tools stay bound so the prior tool_use blocks in the history remain valid. + answer_model: ToolCallingModel = chat_model.bind_tools( + [search_tool], tool_choice=TOOL_CHOICE_NONE + ) + + def prepare_context_node(state: RagState) -> dict[str, list[BaseMessage]]: + # Summarized once up front so the tool loop never splits a tool_use/tool_result pair. + messages = state["messages"] + context_view = summarize_if_needed( + chat_model, messages, history_token_limit=history_token_limit + ) + if context_view is messages: + return {} + return {"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES), *context_view]} + + def agent_node(state: RagState) -> dict[str, list[BaseMessage]]: + messages = state["messages"] + model = tool_model + if _count_tool_rounds(messages) >= MAX_TOOL_ROUNDS: + model = answer_model + response = model.invoke(messages) + return {"messages": [response]} + + # Retrieval failures must fail the turn (and send a failure webhook), not be fed + # back to the model as a tool result; only bad tool arguments are handed back. + tool_node = ToolNode([search_tool], handle_tool_errors=ToolInvocationError) + + graph = StateGraph(RagState) + graph.add_node(PREPARE_CONTEXT_NODE, prepare_context_node) + graph.add_node(AGENT_NODE, agent_node) + graph.add_node(TOOLS_NODE, tool_node) + graph.add_edge(START, PREPARE_CONTEXT_NODE) + graph.add_edge(PREPARE_CONTEXT_NODE, AGENT_NODE) + graph.add_conditional_edges( + AGENT_NODE, tools_condition, {TOOLS_NODE: TOOLS_NODE, END: END} + ) + graph.add_edge(TOOLS_NODE, AGENT_NODE) + return graph.compile() + + +def _collect_sources(messages: list[BaseMessage]) -> list[ChatbotSource]: + sources: list[ChatbotSource] = [] + seen_file_ids: set[str] = set() + for message in messages: + if not isinstance(message, ToolMessage) or not message.artifact: + continue + passages: list[KnowledgeBasePassage] = message.artifact + for passage in passages: + if passage.file_id in seen_file_ids: + continue + seen_file_ids.add(passage.file_id) + sources.append({"file_id": passage.file_id, "filename": passage.filename}) + return sources + + +async def run_rag_turn( + *, + system_prompt: str, + history: list[ChatbotMessageDict], + user_message: str | None, + model_settings: ChatbotModelSettings, + openai_client: OpenAI, + collected_fields: dict[str, Any] | None = None, +) -> ChatbotTurnResult: + """Run one stateless turn: answer the user from the configured knowledge bases.""" + validate_turn_input(history=history, user_message=user_message) + + rag_system_prompt = build_rag_system_prompt( + goal=system_prompt, instructions=model_settings.instructions + ) + initial_state: RagState = { + "messages": to_langchain_messages( + system_prompt=rag_system_prompt, history=history, user_message=user_message + ), + } + search_tool = build_search_tool( + client=openai_client, + knowledge_base_ids=model_settings.knowledge_base_ids, + max_num_results=model_settings.max_num_results, + ) + try: + graph = build_rag_graph( + get_chat_model(model_settings), + search_tool, + history_token_limit=model_settings.history_token_limit, + ) + # ChatAnthropic.invoke and the OpenAI search are sync; keep them off the event loop. + result = await asyncio.to_thread(graph.invoke, initial_state) + except CHAT_MODEL_ERRORS as e: + raise to_provider_error(e, model=model_settings.model) from e + + final_messages: list[BaseMessage] = result["messages"] + return ChatbotTurnResult( + reply=extract_final_reply(final_messages), + collected_fields=dict(collected_fields) if collected_fields else {}, + is_complete=False, + sources=_collect_sources(final_messages), + ) diff --git a/backend/app/services/chatbot/agents/rag/prompts.py b/backend/app/services/chatbot/agents/rag/prompts.py new file mode 100644 index 000000000..4a484f66e --- /dev/null +++ b/backend/app/services/chatbot/agents/rag/prompts.py @@ -0,0 +1,35 @@ +SEARCH_TOOL_NAME = "search_knowledge_base" +SEARCH_TOOL_DESCRIPTION = ( + "Search the knowledge base documents and return the most relevant passages. " + "Call it whenever the user's question may be answered by the documents." +) +SEARCH_QUERY_DESCRIPTION = ( + "A standalone search query: resolve pronouns and references from the " + "conversation so the query makes sense without it." +) + +NO_PASSAGES_FOUND = ( + "No relevant passages were found in the knowledge base for this query." +) +PASSAGE_TEMPLATE = '\n{text}\n' +PASSAGE_SEPARATOR = "\n\n" + +RAG_INSTRUCTION = ( + "You are an assistant that answers the user's questions using a knowledge base " + f"of documents, searchable with the `{SEARCH_TOOL_NAME}` tool.\n" + f"1. Call `{SEARCH_TOOL_NAME}` when the question may be covered by the documents; " + "skip it for greetings and small talk.\n" + "2. Write each search query as a standalone question, resolving references from " + "the conversation.\n" + "3. Ground your answers in the retrieved passages and never invent facts. If the " + "passages don't cover the question, say so.\n" + "4. The goal decides your scope. If it restricts answers to the knowledge base, " + "politely decline anything the documents don't cover. Otherwise you may use " + "general knowledge, but say that part isn't from the documents.\n" + "5. If the goal and the configuration instructions conflict, follow the goal.\n" + "6. Never mention tools, searches, passages, file names, or ids to the user." +) +CONFIG_INSTRUCTIONS_TEMPLATE = ( + "\n\n\n{instructions}\n" +) +GOAL_TEMPLATE = "\n\n\n{goal}\n" diff --git a/backend/app/services/chatbot/agents/rag/retriever.py b/backend/app/services/chatbot/agents/rag/retriever.py new file mode 100644 index 000000000..6d937bdce --- /dev/null +++ b/backend/app/services/chatbot/agents/rag/retriever.py @@ -0,0 +1,216 @@ +"""Knowledge-base search over OpenAI vector stores; Anthropic has no native file search.""" + +import html +import logging + +import openai +from langchain_core.tools import BaseTool, StructuredTool +from openai import OpenAI +from sqlmodel import Session + +from app.crud import get_provider_credential +from app.crud.rag.open_ai import OPENAI_TIMEOUT_SECONDS +from app.models.llm.constants import Provider +from app.services.chatbot.agents.rag.prompts import ( + NO_PASSAGES_FOUND, + PASSAGE_SEPARATOR, + PASSAGE_TEMPLATE, + SEARCH_TOOL_DESCRIPTION, + SEARCH_TOOL_NAME, +) +from app.services.chatbot.agents.rag.types import ( + KnowledgeBasePassage, + SearchKnowledgeBaseArgs, +) +from app.services.chatbot.exceptions import ChatbotInputError, ChatbotProviderError +from app.services.chatbot.provider_errors import PROVIDER_SERVER_ERROR_STATUS +from app.utils import mask_string + +logger = logging.getLogger(__name__) + +DEFAULT_MAX_NUM_RESULTS = 5 +OPENAI_MAX_NUM_RESULTS = 50 +# A search is an idempotent read, so one SDK retry absorbs a transient blip +# without stretching the turn far past OPENAI_TIMEOUT_SECONDS. +KB_SEARCH_MAX_RETRIES = 1 +CREDENTIAL_API_KEY = "api_key" +PASSAGE_CHUNK_SEPARATOR = "\n" + + +def get_knowledge_base_client( + *, session: Session, organization_id: int, project_id: int +) -> OpenAI: + """OpenAI client from the project's credentials, used only for vector-store search.""" + credentials = get_provider_credential( + session=session, + org_id=organization_id, + project_id=project_id, + provider=Provider.OPENAI.value, + ) + if not isinstance(credentials, dict) or not credentials.get(CREDENTIAL_API_KEY): + error_message = ( + "[KAAPI] The rag agent searches knowledge bases with the project's OpenAI " + "credentials, but none are configured; add OpenAI credentials to this project." + ) + logger.warning( + f"[get_knowledge_base_client] {error_message} | project_id: {project_id}" + ) + raise ChatbotInputError(error_message) + + api_key = credentials[CREDENTIAL_API_KEY] + logger.info( + f"[get_knowledge_base_client] OpenAI client configured | " + f"project_id: {project_id}, api_key: {mask_string(api_key)}" + ) + return OpenAI( + api_key=api_key, + timeout=OPENAI_TIMEOUT_SECONDS, + max_retries=KB_SEARCH_MAX_RETRIES, + ) + + +def _resolve_top_k(max_num_results: int | None) -> int: + if max_num_results is None: + return DEFAULT_MAX_NUM_RESULTS + return min(max_num_results, OPENAI_MAX_NUM_RESULTS) + + +def _to_search_error(error: openai.OpenAIError, *, vector_store_id: str) -> str: + # APITimeoutError subclasses APIConnectionError, so it must be checked first. + if isinstance(error, openai.APITimeoutError): + error_message = ( + "[KAAPI] Knowledge base search timed out — retry the turn. If persistent, " + "contact Kaapi. (code: APITimeoutError)" + ) + logger.error( + f"[search_knowledge_bases] {error_message} | provider=openai, " + f"vector_store_id: {vector_store_id}", + exc_info=error, + ) + return error_message + if isinstance(error, openai.APIConnectionError): + error_message = ( + "[KAAPI] Network/DNS issue reaching OpenAI for knowledge base search — " + "check connectivity. If persistent, contact Kaapi. (code: APIConnectionError)" + ) + logger.error( + f"[search_knowledge_bases] {error_message} | provider=openai, " + f"vector_store_id: {vector_store_id}", + exc_info=error, + ) + return error_message + if isinstance(error, openai.APIStatusError): + error_message = ( + f"[OPENAI] Knowledge base search failed for '{vector_store_id}': " + f"{error.message} (code: {error.status_code}, request_id: {error.request_id})" + ) + # 5xx is provider-side (alert-worthy); 4xx is caller's fault (noise if alerted) + log = ( + logger.error + if error.status_code >= PROVIDER_SERVER_ERROR_STATUS + else logger.warning + ) + log( + f"[search_knowledge_bases] {error_message} | provider=openai, " + f"vector_store_id: {vector_store_id}", + exc_info=error, + ) + return error_message + + error_message = ( + f"[OPENAI] Knowledge base search failed for '{vector_store_id}': {error} " + f"— retry the turn. If persistent, contact Kaapi. (code: {type(error).__name__})" + ) + logger.error( + f"[search_knowledge_bases] {error_message} | provider=openai, " + f"vector_store_id: {vector_store_id}", + exc_info=error, + ) + return error_message + + +def search_knowledge_bases( + *, + client: OpenAI, + knowledge_base_ids: tuple[str, ...], + query: str, + max_num_results: int | None, +) -> list[KnowledgeBasePassage]: + """Search every knowledge base and keep the overall top passages by score.""" + top_k = _resolve_top_k(max_num_results) + passages: list[KnowledgeBasePassage] = [] + for vector_store_id in knowledge_base_ids: + logger.info( + f"[search_knowledge_bases] Searching knowledge base | " + f"vector_store_id: {vector_store_id}, max_num_results: {top_k}" + ) + try: + page = client.vector_stores.search( + vector_store_id, query=query, max_num_results=top_k + ) + except openai.OpenAIError as e: + error_message = _to_search_error(e, vector_store_id=vector_store_id) + raise ChatbotProviderError(error_message) from e + + for result in page.data: + chunk_texts: list[str] = [] + for chunk in result.content: + chunk_texts.append(chunk.text) + passages.append( + KnowledgeBasePassage( + file_id=result.file_id, + filename=result.filename, + score=result.score, + text=PASSAGE_CHUNK_SEPARATOR.join(chunk_texts), + ) + ) + + passages.sort(key=lambda passage: passage.score, reverse=True) + top_passages = passages[:top_k] + logger.info( + f"[search_knowledge_bases] Knowledge base search finished | " + f"knowledge_bases: {len(knowledge_base_ids)}, passages: {len(top_passages)}" + ) + return top_passages + + +def format_passages(passages: list[KnowledgeBasePassage]) -> str: + if not passages: + return NO_PASSAGES_FOUND + + rendered: list[str] = [] + for passage in passages: + rendered.append( + PASSAGE_TEMPLATE.format( + filename=html.escape(passage.filename, quote=True), text=passage.text + ) + ) + return PASSAGE_SEPARATOR.join(rendered) + + +def build_search_tool( + *, + client: OpenAI, + knowledge_base_ids: tuple[str, ...], + max_num_results: int | None, +) -> BaseTool: + """`search_knowledge_base` tool: passages as content, the structured hits as artifact.""" + + def search_knowledge_base( + query: str, + ) -> tuple[str, list[KnowledgeBasePassage]]: + passages = search_knowledge_bases( + client=client, + knowledge_base_ids=knowledge_base_ids, + query=query, + max_num_results=max_num_results, + ) + return format_passages(passages), passages + + return StructuredTool.from_function( + func=search_knowledge_base, + name=SEARCH_TOOL_NAME, + description=SEARCH_TOOL_DESCRIPTION, + args_schema=SearchKnowledgeBaseArgs, + response_format="content_and_artifact", + ) diff --git a/backend/app/services/chatbot/agents/rag/types.py b/backend/app/services/chatbot/agents/rag/types.py new file mode 100644 index 000000000..a1b85ec87 --- /dev/null +++ b/backend/app/services/chatbot/agents/rag/types.py @@ -0,0 +1,24 @@ +from dataclasses import dataclass +from typing import Annotated, TypedDict + +from langchain_core.messages import BaseMessage +from langgraph.graph.message import add_messages +from pydantic import BaseModel, Field + +from app.services.chatbot.agents.rag.prompts import SEARCH_QUERY_DESCRIPTION + + +class RagState(TypedDict): + messages: Annotated[list[BaseMessage], add_messages] + + +@dataclass(frozen=True) +class KnowledgeBasePassage: + file_id: str + filename: str + score: float + text: str + + +class SearchKnowledgeBaseArgs(BaseModel): + query: str = Field(min_length=1, description=SEARCH_QUERY_DESCRIPTION) diff --git a/backend/app/services/chatbot/agents/registry.py b/backend/app/services/chatbot/agents/registry.py new file mode 100644 index 000000000..74d8fa77b --- /dev/null +++ b/backend/app/services/chatbot/agents/registry.py @@ -0,0 +1,55 @@ +import logging +from typing import Any + +from openai import OpenAI + +from app.models.chatbot import ChatbotAgentTypeEnum +from app.services.chatbot.agents.onboarding.graph import run_chatbot_turn +from app.services.chatbot.agents.rag.graph import run_rag_turn +from app.services.chatbot.exceptions import ChatbotInputError +from app.services.chatbot.types import ChatbotMessageDict, ChatbotTurnResult +from app.services.chatbot.utils import ChatbotModelSettings + +logger = logging.getLogger(__name__) + + +async def run_agent_turn( + *, + agent_type: ChatbotAgentTypeEnum, + system_prompt: str, + history: list[ChatbotMessageDict], + user_message: str | None, + model_settings: ChatbotModelSettings, + collected_fields: dict[str, Any] | None = None, + openai_client: OpenAI | None = None, +) -> ChatbotTurnResult: + """Run one turn with the agent the session was created for.""" + if agent_type == ChatbotAgentTypeEnum.ONBOARDING: + return await run_chatbot_turn( + system_prompt=system_prompt, + history=history, + user_message=user_message, + model_settings=model_settings, + collected_fields=collected_fields, + ) + + if agent_type == ChatbotAgentTypeEnum.RAG: + if openai_client is None: + error_message = ( + "[KAAPI] The rag agent needs an OpenAI client to search its " + "knowledge bases." + ) + logger.error(f"[run_agent_turn] {error_message}") + raise ChatbotInputError(error_message) + return await run_rag_turn( + system_prompt=system_prompt, + history=history, + user_message=user_message, + model_settings=model_settings, + openai_client=openai_client, + collected_fields=collected_fields, + ) + + error_message = f"[KAAPI] Unsupported chatbot agent type '{agent_type}'." + logger.error(f"[run_agent_turn] {error_message}") + raise ChatbotInputError(error_message) diff --git a/backend/app/services/chatbot/config.py b/backend/app/services/chatbot/config.py index 3ab7ec419..1014e3be8 100644 --- a/backend/app/services/chatbot/config.py +++ b/backend/app/services/chatbot/config.py @@ -1,14 +1,18 @@ +import logging from typing import Any, get_args from uuid import UUID from fastapi import HTTPException +from openai import OpenAI from pydantic import ValidationError from sqlmodel import Session +from app.crud import CollectionCrud from app.crud.config import ConfigVersionCrud -from app.models import ConfigVersion +from app.models import ConfigVersion, ProviderType from app.models.llm.constants import Provider from app.models.llm.request import ConfigBlob, KaapiTextCompletionConfig +from app.services.chatbot.agents.rag.retriever import get_knowledge_base_client from app.services.chatbot.exceptions import ChatbotInputError from app.services.chatbot.utils import ( DEFAULT_HISTORY_TOKEN_LIMIT, @@ -18,6 +22,8 @@ default_thinking, ) +logger = logging.getLogger(__name__) + _ANTHROPIC_EFFORT_LEVELS: tuple[AnthropicEffort, ...] = get_args(AnthropicEffort) THINKING_TYPE_ENABLED = "enabled" @@ -120,5 +126,60 @@ def resolve_chatbot_model_settings( if history_token_limit is not None else DEFAULT_HISTORY_TOKEN_LIMIT ), + instructions=params.instructions, + knowledge_base_ids=tuple(params.knowledge_base_ids or ()), + max_num_results=params.max_num_results, ) return model_settings + + +def _validate_knowledge_base_ids( + *, session: Session, project_id: int, knowledge_base_ids: tuple[str, ...] +) -> None: + if not knowledge_base_ids: + error_message = ( + "[KAAPI] The rag agent needs a config with non-empty " + "completion.params.knowledge_base_ids; add the collection's knowledge base id." + ) + logger.warning( + f"[_validate_knowledge_base_ids] {error_message} | project_id: {project_id}" + ) + raise ChatbotInputError(error_message) + + known_ids: set[str] = set() + for collection in CollectionCrud(session, project_id).read_all(): + if collection.provider == ProviderType.openai: + known_ids.add(collection.knowledge_base_id) + + unknown_ids: list[str] = [] + for knowledge_base_id in knowledge_base_ids: + if knowledge_base_id not in known_ids: + unknown_ids.append(knowledge_base_id) + if unknown_ids: + error_message = ( + "[KAAPI] Config knowledge_base_ids not found among this project's OpenAI " + f"collections: {', '.join(unknown_ids)}; use the knowledge_base_id of an " + "existing collection." + ) + logger.warning( + f"[_validate_knowledge_base_ids] {error_message} | project_id: {project_id}" + ) + raise ChatbotInputError(error_message) + + +def prepare_rag_search_client( + *, + session: Session, + organization_id: int, + project_id: int, + model_settings: ChatbotModelSettings, +) -> OpenAI: + """Validate the rag agent's knowledge bases and return the client that searches them.""" + _validate_knowledge_base_ids( + session=session, + project_id=project_id, + knowledge_base_ids=model_settings.knowledge_base_ids, + ) + return get_knowledge_base_client( + session=session, organization_id=organization_id, project_id=project_id + ) diff --git a/backend/app/services/chatbot/history.py b/backend/app/services/chatbot/history.py index d3680638c..4967489c3 100644 --- a/backend/app/services/chatbot/history.py +++ b/backend/app/services/chatbot/history.py @@ -1,13 +1,17 @@ """Conversation-history plumbing shared by every chatbot agent graph.""" +import logging + from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage from app.models.chatbot import ChatbotMessageRoleEnum -from app.services.chatbot.exceptions import ChatbotInputError +from app.services.chatbot.exceptions import ChatbotInputError, ChatbotProviderError from app.services.chatbot.types import ChatbotMessageDict from app.services.chatbot.utils import FLOW_KEEP_RECENT_MESSAGES +logger = logging.getLogger(__name__) + SUMMARY_NOTE_PREFIX = "[Earlier conversation summary] " SUMMARIZE_INSTRUCTION = ( "Summarize the conversation so far in 2-3 sentences. Keep every fact the " @@ -112,20 +116,26 @@ def to_langchain_messages( return [SystemMessage(content=system_prompt), *conversation] -def to_message_dicts( - *, system_prompt: str, messages: list[BaseMessage] -) -> list[ChatbotMessageDict]: - transcript: list[ChatbotMessageDict] = [ - {"role": ChatbotMessageRoleEnum.SYSTEM.value, "content": system_prompt} - ] - for message in messages: - if message.id == KICKOFF_MESSAGE_ID: - continue - if isinstance(message, HumanMessage): - role = ChatbotMessageRoleEnum.USER - elif isinstance(message, AIMessage): - role = ChatbotMessageRoleEnum.ASSISTANT - else: - continue - transcript.append({"role": role.value, "content": message.text}) - return transcript +def extract_final_reply(messages: list[BaseMessage]) -> str: + """Text of the turn's last message, which every agent graph ends on.""" + reply = "" + if messages and isinstance(messages[-1], AIMessage): + reply = messages[-1].text + if not reply.strip(): + error_message = ( + "[ANTHROPIC] Chat model returned no reply text — retry the turn. " + "If persistent, contact Kaapi. (code: EmptyReply)" + ) + logger.warning(f"[extract_final_reply] {error_message}") + raise ChatbotProviderError(error_message) + return reply + + +def validate_turn_input( + *, history: list[ChatbotMessageDict], user_message: str | None +) -> None: + if user_message is None and history: + raise ChatbotInputError( + "[KAAPI] `message` is required once `history` is non-empty; omit it only " + "on the first call to get the opening greeting." + ) diff --git a/backend/app/services/chatbot/provider_errors.py b/backend/app/services/chatbot/provider_errors.py new file mode 100644 index 000000000..7430c8145 --- /dev/null +++ b/backend/app/services/chatbot/provider_errors.py @@ -0,0 +1,68 @@ +"""Anthropic chat-model error mapping shared by every chatbot agent.""" + +import logging + +import anthropic + +from app.services.chatbot.exceptions import ChatbotProviderError + +logger = logging.getLogger(__name__) + +PROVIDER_SERVER_ERROR_STATUS = 500 + +# ValueError covers langchain_anthropic's client-side rejection of a model config +# (e.g. thinking/effort the model doesn't support) before any request is sent. +CHAT_MODEL_ERRORS: tuple[type[Exception], ...] = ( + anthropic.APIConnectionError, + anthropic.APIStatusError, + ValueError, +) + + +def to_provider_error(error: Exception, *, model: str) -> ChatbotProviderError: + # APITimeoutError subclasses APIConnectionError, so it must be checked first. + if isinstance(error, anthropic.APITimeoutError): + error_message = ( + "[KAAPI] Request timed out — retry smaller. If persistent, contact Kaapi. " + "(code: APITimeoutError)" + ) + logger.error( + f"[to_provider_error] {error_message} | provider=anthropic, model={model}", + exc_info=error, + ) + return ChatbotProviderError(error_message) + if isinstance(error, anthropic.APIConnectionError): + error_message = ( + "[KAAPI] Network/DNS issue reaching provider — check connectivity. " + "If persistent, contact Kaapi. (code: APIConnectionError)" + ) + logger.error( + f"[to_provider_error] {error_message} | provider=anthropic, model={model}", + exc_info=error, + ) + return ChatbotProviderError(error_message) + if isinstance(error, anthropic.APIStatusError): + error_message = ( + f"[ANTHROPIC] Chat model call failed: {error.message} " + f"(code: {error.status_code}, request_id: {error.request_id})" + ) + # 5xx is provider-side (alert-worthy); 4xx is caller's fault (noise if alerted) + log = ( + logger.error + if error.status_code >= PROVIDER_SERVER_ERROR_STATUS + else logger.warning + ) + log( + f"[to_provider_error] {error_message} | provider=anthropic, model={model}", + exc_info=error, + ) + return ChatbotProviderError(error_message) + + error_message = ( + f"[KAAPI] Unsupported model configuration for '{model}': {error} " + "— update the config's model, thinking, or effort. (code: ValueError)" + ) + logger.warning( + f"[to_provider_error] {error_message} | provider=anthropic, model={model}" + ) + return ChatbotProviderError(error_message) diff --git a/backend/app/services/chatbot/session.py b/backend/app/services/chatbot/session.py index e142ad8d4..af3434f14 100644 --- a/backend/app/services/chatbot/session.py +++ b/backend/app/services/chatbot/session.py @@ -2,6 +2,7 @@ from uuid import UUID from fastapi import HTTPException +from openai import OpenAI from sqlmodel import Session from app.core.db import engine @@ -12,12 +13,16 @@ update_chatbot_session_progress, ) from app.models import ( + ChatbotAgentTypeEnum, ChatbotMessageRoleEnum, ChatbotSessionMessage, ChatbotSessionStatusEnum, ) -from app.services.chatbot.agents.onboarding.graph import run_chatbot_turn -from app.services.chatbot.config import resolve_chatbot_model_settings +from app.services.chatbot.agents.registry import run_agent_turn +from app.services.chatbot.config import ( + prepare_rag_search_client, + resolve_chatbot_model_settings, +) from app.services.chatbot.exceptions import ChatbotInputError, ChatbotProviderError from app.services.chatbot.types import ChatbotMessageDict from app.services.chatbot.utils import ChatbotModelSettings @@ -42,19 +47,12 @@ def _split_pending_turn( if not history: return history, None if history[-1]["role"] != ChatbotMessageRoleEnum.USER: - # No pending user turn; run_chatbot_turn rejects this with ChatbotInputError. + # No pending user turn; the agent runner rejects this with ChatbotInputError. return history, None pending = history.pop() return history, pending["content"] -def _extract_reply(transcript: list[ChatbotMessageDict]) -> str: - for entry in reversed(transcript): - if entry["role"] == ChatbotMessageRoleEnum.ASSISTANT: - return entry["content"] - raise ChatbotInputError("[KAAPI] Chat model returned no assistant reply.") - - def _send_failure_webhook( *, session_id: UUID, @@ -110,7 +108,9 @@ def process_chatbot_session_message( system_prompt = chatbot_session.system_prompt chatbot_ref = chatbot_session.chatbot_ref session_status = chatbot_session.status.value + agent_type = chatbot_session.agent_type collected_fields = dict(chatbot_session.collected_fields) + openai_client: OpenAI | None = None try: model_settings: ChatbotModelSettings = resolve_chatbot_model_settings( @@ -121,6 +121,13 @@ def process_chatbot_session_message( max_tokens=None, history_token_limit=None, ) + if agent_type == ChatbotAgentTypeEnum.RAG: + openai_client = prepare_rag_search_client( + session=session, + organization_id=organization_id, + project_id=project_id, + model_settings=model_settings, + ) except HTTPException as e: _send_failure_webhook( session_id=session_id, @@ -142,15 +149,16 @@ def process_chatbot_session_message( try: turn_result = asyncio.run( - run_chatbot_turn( + run_agent_turn( + agent_type=agent_type, system_prompt=system_prompt, history=history, user_message=user_message, model_settings=model_settings, collected_fields=collected_fields, + openai_client=openai_client, ) ) - reply = _extract_reply(turn_result.transcript) except (ChatbotInputError, ChatbotProviderError) as e: _send_failure_webhook( session_id=session_id, @@ -160,6 +168,7 @@ def process_chatbot_session_message( ) return + reply = turn_result.reply completion_status = turn_result.completion_status with Session(engine) as session: create_chatbot_session_message( @@ -179,10 +188,12 @@ def process_chatbot_session_message( data={ "session_id": str(session_id), "chatbot_ref": chatbot_ref, + "agent_type": agent_type.value, "reply": reply, "status": session_status, "completion_status": completion_status.value, "collected_fields": turn_result.collected_fields, + "sources": turn_result.sources, } ) send_callback(webhook_url, envelope.model_dump(exclude_none=True), webhook_secret) diff --git a/backend/app/services/chatbot/types.py b/backend/app/services/chatbot/types.py index 2333353e2..71df0e814 100644 --- a/backend/app/services/chatbot/types.py +++ b/backend/app/services/chatbot/types.py @@ -1,16 +1,22 @@ -from dataclasses import dataclass -from typing import Any +from dataclasses import dataclass, field +from typing import Any, TypedDict from app.models.chatbot import ChatbotCompletionStatusEnum ChatbotMessageDict = dict[str, str] +class ChatbotSource(TypedDict): + file_id: str + filename: str + + @dataclass(frozen=True) class ChatbotTurnResult: - transcript: list[ChatbotMessageDict] + reply: str collected_fields: dict[str, Any] is_complete: bool + sources: list[ChatbotSource] = field(default_factory=list) @property def completion_status(self) -> ChatbotCompletionStatusEnum: diff --git a/backend/app/services/chatbot/utils.py b/backend/app/services/chatbot/utils.py index d4831b46f..ab6037893 100644 --- a/backend/app/services/chatbot/utils.py +++ b/backend/app/services/chatbot/utils.py @@ -27,6 +27,10 @@ class ChatbotModelSettings: reasoning_effort: AnthropicEffort | None = None thinking: dict[str, Any] = field(default_factory=default_thinking) history_token_limit: int = DEFAULT_HISTORY_TOKEN_LIMIT + instructions: str | None = None + # OpenAI vector store ids (Collection.knowledge_base_id); only the rag agent reads them. + knowledge_base_ids: tuple[str, ...] = () + max_num_results: int | None = None def get_chat_model(model_settings: ChatbotModelSettings) -> ChatAnthropic: diff --git a/backend/app/tests/api/routes/chatbot/test_session.py b/backend/app/tests/api/routes/chatbot/test_session.py index 5501b993a..65c142878 100644 --- a/backend/app/tests/api/routes/chatbot/test_session.py +++ b/backend/app/tests/api/routes/chatbot/test_session.py @@ -9,19 +9,25 @@ from app.api.routes.chatbot import session as session_routes from app.core.config import settings +from app.crud import CollectionCrud from app.crud.chatbot import ( get_chatbot_session, list_chatbot_session_messages, update_chatbot_session_progress, ) +from app.crud.credentials import remove_provider_credential from app.models import ( DEFAULT_CHATBOT_NAME, + ChatbotAgentTypeEnum, ChatbotCompletionStatusEnum, ChatbotMessageRoleEnum, ChatbotSession, ChatbotSessionStatusEnum, Config, + Project, + ProviderType, ) +from app.models.llm.constants import Provider from app.services.chatbot import chatbot_ref as chatbot_ref_module from app.services.chatbot.chatbot_ref import ( CHATBOT_REF_PREFIX, @@ -36,6 +42,7 @@ add_test_chatbot_messages, create_test_chatbot_session, ) +from app.tests.utils.collection import get_vector_store_collection from app.tests.utils.test_data import ( build_anthropic_text_config_blob, create_test_config, @@ -221,6 +228,204 @@ def test_non_anthropic_config_returns_422( assert "only supports configs" in resp.json()["error"] +def _rag_config(db: Session, project_id: int, knowledge_base_ids: list[str]) -> Config: + return create_test_config( + db, + project_id=project_id, + config_blob=build_anthropic_text_config_blob( + knowledge_base_ids=knowledge_base_ids + ), + ) + + +def _project_collection_id(db: Session, auth: TestAuthContext) -> str: + project = db.get(Project, auth.project_id) + assert project is not None + return get_vector_store_collection(db, project).knowledge_base_id + + +def _deleted_collection_id(db: Session, auth: TestAuthContext) -> str: + project = db.get(Project, auth.project_id) + assert project is not None + collection = get_vector_store_collection(db, project) + CollectionCrud(db, project.id).delete_by_id(collection.id) + return collection.knowledge_base_id + + +def _non_openai_collection_id(db: Session, auth: TestAuthContext) -> str: + project = db.get(Project, auth.project_id) + assert project is not None + return get_vector_store_collection( + db, project, provider=ProviderType.google_aistudio + ).knowledge_base_id + + +class TestCreateRagChatbotSession: + def _create( + self, + client: TestClient, + headers: dict[str, str], + config: Config, + contact_ref: str, + **extra: Any, + ) -> Any: + return client.post( + SESSIONS_URL, + headers=headers, + json={ + "config_id": str(config.id), + "contact_ref": contact_ref, + "system_prompt": "Answer questions about fees from the documents.", + **extra, + }, + ) + + def test_rag_session_is_stored_and_reported_by_get( + self, + client: TestClient, + db: Session, + headers: dict[str, str], + user_api_key: TestAuthContext, + ) -> None: + config = _rag_config( + db, user_api_key.project_id, [_project_collection_id(db, user_api_key)] + ) + contact_ref = f"contact-{random_lower_string()}" + + resp = self._create(client, headers, config, contact_ref, agent_type="rag") + + assert resp.status_code == 201 + [stored] = _sessions_for_contact(db, contact_ref) + assert stored.agent_type == ChatbotAgentTypeEnum.RAG + latest = client.get( + SESSIONS_URL, headers=headers, params={"contact_ref": contact_ref} + ) + assert latest.status_code == 200 + assert latest.json()["data"]["agent_type"] == "rag" + + def test_omitted_agent_type_creates_onboarding_session( + self, + client: TestClient, + db: Session, + headers: dict[str, str], + anthropic_config: Config, + ) -> None: + contact_ref = f"contact-{random_lower_string()}" + + resp = self._create(client, headers, anthropic_config, contact_ref) + + assert resp.status_code == 201 + [stored] = _sessions_for_contact(db, contact_ref) + assert stored.agent_type == ChatbotAgentTypeEnum.ONBOARDING + + def test_onboarding_ignores_missing_knowledge_bases( + self, + client: TestClient, + db: Session, + headers: dict[str, str], + user_api_key: TestAuthContext, + ) -> None: + config = _rag_config(db, user_api_key.project_id, ["vs_not_a_collection"]) + + resp = self._create( + client, + headers, + config, + f"contact-{random_lower_string()}", + agent_type="onboarding", + ) + + assert resp.status_code == 201 + + def test_unknown_agent_type_returns_422( + self, + client: TestClient, + headers: dict[str, str], + anthropic_config: Config, + ) -> None: + resp = self._create( + client, + headers, + anthropic_config, + f"contact-{random_lower_string()}", + agent_type="faq", + ) + + assert resp.status_code == 422 + + def test_config_without_knowledge_bases_returns_422( + self, + client: TestClient, + db: Session, + headers: dict[str, str], + anthropic_config: Config, + ) -> None: + contact_ref = f"contact-{random_lower_string()}" + + resp = self._create( + client, headers, anthropic_config, contact_ref, agent_type="rag" + ) + + assert resp.status_code == 422 + assert "non-empty completion.params.knowledge_base_ids" in resp.json()["error"] + assert _sessions_for_contact(db, contact_ref) == [] + + @pytest.mark.parametrize( + "seed_knowledge_base_id", + [ + lambda _db, _auth: "vs_does_not_exist", + _deleted_collection_id, + _non_openai_collection_id, + ], + ids=["unknown", "deleted", "non-openai"], + ) + def test_unusable_knowledge_base_returns_422_naming_it( + self, + client: TestClient, + db: Session, + headers: dict[str, str], + user_api_key: TestAuthContext, + seed_knowledge_base_id: Any, + ) -> None: + valid_id = _project_collection_id(db, user_api_key) + bad_id = seed_knowledge_base_id(db, user_api_key) + config = _rag_config(db, user_api_key.project_id, [valid_id, bad_id]) + contact_ref = f"contact-{random_lower_string()}" + + resp = self._create(client, headers, config, contact_ref, agent_type="rag") + + assert resp.status_code == 422 + error = resp.json()["error"] + assert f"collections: {bad_id};" in error + assert valid_id not in error + assert _sessions_for_contact(db, contact_ref) == [] + + def test_project_without_openai_credentials_returns_422( + self, + client: TestClient, + db: Session, + headers: dict[str, str], + user_api_key: TestAuthContext, + ) -> None: + config = _rag_config( + db, user_api_key.project_id, [_project_collection_id(db, user_api_key)] + ) + # Rolled back with the test transaction, so the seeded credential returns. + remove_provider_credential( + db, + org_id=user_api_key.organization_id, + project_id=user_api_key.project_id, + provider=Provider.OPENAI.value, + ) + contact_ref = f"contact-{random_lower_string()}" + + resp = self._create(client, headers, config, contact_ref, agent_type="rag") + + assert resp.status_code == 422 + assert "OpenAI credentials" in resp.json()["error"] + assert _sessions_for_contact(db, contact_ref) == [] + + class TestCreateChatbotSessionChatbotRef: def _create( self, @@ -444,6 +649,7 @@ def test_returns_contacts_newest_session_in_callers_project( assert body["contact_ref"] == contact_ref assert body["chatbot_ref"] == chatbot_ref assert body["system_prompt"] == "Collect the student's name and grade." + assert body["agent_type"] == "onboarding" assert body["status"] == "active" assert body["completion_status"] == "in_progress" assert body["collected_fields"] == {} diff --git a/backend/app/tests/services/chatbot/agents/onboarding/test_graph.py b/backend/app/tests/services/chatbot/agents/onboarding/test_graph.py index 01f8a9990..ffda540b5 100644 --- a/backend/app/tests/services/chatbot/agents/onboarding/test_graph.py +++ b/backend/app/tests/services/chatbot/agents/onboarding/test_graph.py @@ -53,10 +53,8 @@ async def test_first_call_returns_greeting_without_kickoff( model_settings=_model_settings(), ) - assert result.transcript == [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "assistant", "content": "Hi! What's your name?"}, - ] + assert result.reply == "Hi! What's your name?" + assert result.sources == [] # Anthropic rejects a conversation that doesn't open with a user turn. sent = chat_model.calls[0] assert isinstance(sent[0], SystemMessage) @@ -81,12 +79,7 @@ async def test_follow_up_turn_appends_user_message_and_reply( model_settings=_model_settings(), ) - assert result.transcript == [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "assistant", "content": "Hi! What's your name?"}, - {"role": "user", "content": "Asha"}, - {"role": "assistant", "content": "Nice to meet you, Asha."}, - ] + assert result.reply == "Nice to meet you, Asha." sent_contents = [m.content for m in chat_model.calls[0]] assert "stale prompt from the previous call" not in sent_contents @@ -215,13 +208,8 @@ async def test_long_history_over_budget_is_summarized( "bot 3", "user 4", ] - # Summarization only trims the model's view; the caller keeps the full transcript. - assert result.transcript == [ - {"role": "system", "content": SYSTEM_PROMPT}, - *history, - {"role": "user", "content": "user 4"}, - {"role": "assistant", "content": "Got it."}, - ] + # The reply is the turn's answer, never the summary text. + assert result.reply == "Got it." @pytest.mark.asyncio async def test_long_history_within_budget_is_sent_whole( diff --git a/backend/app/tests/services/chatbot/agents/rag/__init__.py b/backend/app/tests/services/chatbot/agents/rag/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/backend/app/tests/services/chatbot/agents/rag/test_graph.py b/backend/app/tests/services/chatbot/agents/rag/test_graph.py new file mode 100644 index 000000000..3b2891608 --- /dev/null +++ b/backend/app/tests/services/chatbot/agents/rag/test_graph.py @@ -0,0 +1,319 @@ +import anthropic +import httpx +import openai +import pytest +from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage + +from app.services.chatbot.agents.rag.graph import run_rag_turn +from app.services.chatbot.agents.rag.prompts import RAG_INSTRUCTION +from app.services.chatbot.exceptions import ChatbotInputError, ChatbotProviderError +from app.services.chatbot.utils import ChatbotModelSettings +from app.tests.utils.chatbot import ( + FakeOpenAIClient, + FakeVectorStores, + ScriptedChatModel, + patch_chat_model, + search_hit, + tool_call_reply, +) + +GOAL = "Answer questions about school fees only from the documents." +SUMMARY_TEXT = "User asked about admissions." + +_ANTHROPIC_REQUEST = httpx.Request("POST", "https://api.anthropic.com/v1/messages") +_OPENAI_REQUEST = httpx.Request("POST", "https://api.openai.com/v1/vector_stores") + + +def _model_settings( + *, instructions: str | None = None, history_token_limit: int = 4000 +) -> ChatbotModelSettings: + return ChatbotModelSettings( + model="claude-sonnet-4-6", + instructions=instructions, + knowledge_base_ids=("vs_fees",), + history_token_limit=history_token_limit, + ) + + +def _fees_client() -> FakeOpenAIClient: + return FakeOpenAIClient( + vector_stores=FakeVectorStores( + hits={ + "vs_fees": [ + search_hit( + file_id="file-fees", + filename="fees.pdf", + score=0.9, + texts=("Tuition is 100 per term.",), + ), + search_hit(file_id="file-faq", filename="faq.pdf", score=0.4), + ] + } + ) + ) + + +async def _run( + chat_model: ScriptedChatModel, + monkeypatch: pytest.MonkeyPatch, + *, + client: FakeOpenAIClient | None = None, + user_message: str | None = "What is the tuition?", + history: list[dict[str, str]] | None = None, + model_settings: ChatbotModelSettings | None = None, + collected_fields: dict[str, str] | None = None, +): + patch_chat_model(monkeypatch, chat_model) + return await run_rag_turn( + system_prompt=GOAL, + history=history or [], + user_message=user_message, + model_settings=model_settings or _model_settings(), + openai_client=client or _fees_client(), + collected_fields=collected_fields, + ) + + +class TestRunRagTurn: + @pytest.mark.asyncio + async def test_answers_without_searching( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + client = _fees_client() + chat_model = ScriptedChatModel(replies=["Hello! Ask me about fees."]) + + result = await _run( + chat_model, + monkeypatch, + client=client, + user_message="hi", + collected_fields={"name": "Asha"}, + ) + + assert result.reply == "Hello! Ask me about fees." + assert result.sources == [] + assert result.is_complete is False + assert result.collected_fields == {"name": "Asha"} + assert client.vector_stores.calls == [] + assert chat_model.bound_tool_names[0] == ["search_knowledge_base"] + + @pytest.mark.asyncio + async def test_tool_round_feeds_passages_back_and_returns_deduped_sources( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + client = _fees_client() + chat_model = ScriptedChatModel( + replies=[ + tool_call_reply("tuition fee", "fee payment deadline"), + "Tuition is 100 per term.", + ] + ) + + result = await _run(chat_model, monkeypatch, client=client) + + assert result.reply == "Tuition is 100 per term." + # Both searches hit the same two files; each file is reported once. + assert result.sources == [ + {"file_id": "file-fees", "filename": "fees.pdf"}, + {"file_id": "file-faq", "filename": "faq.pdf"}, + ] + assert [call["query"] for call in client.vector_stores.calls] == [ + "tuition fee", + "fee payment deadline", + ] + tool_results = [m for m in chat_model.calls[1] if isinstance(m, ToolMessage)] + assert len(tool_results) == 2 + assert ( + '\nTuition is 100 per term.\n' + in tool_results[0].content + ) + + @pytest.mark.asyncio + async def test_fourth_agent_call_is_forced_to_answer_after_three_tool_rounds( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + client = _fees_client() + chat_model = ScriptedChatModel( + replies=[ + tool_call_reply("tuition"), + tool_call_reply("tuition 2024"), + tool_call_reply("tuition per term"), + "Tuition is 100 per term.", + ] + ) + + result = await _run(chat_model, monkeypatch, client=client) + + assert result.reply == "Tuition is 100 per term." + assert chat_model.call_tool_choices == [None, None, None, {"type": "none"}] + assert len(client.vector_stores.calls) == 3 + + @pytest.mark.asyncio + async def test_invalid_tool_args_are_handed_back_to_the_model( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + client = _fees_client() + chat_model = ScriptedChatModel( + replies=[tool_call_reply(""), "Could you rephrase that?"] + ) + + result = await _run(chat_model, monkeypatch, client=client) + + assert result.reply == "Could you rephrase that?" + assert result.sources == [] + assert client.vector_stores.calls == [] + [tool_result] = [m for m in chat_model.calls[1] if isinstance(m, ToolMessage)] + assert tool_result.status == "error" + + @pytest.mark.asyncio + async def test_search_failure_fails_the_turn_instead_of_reaching_the_model( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + client = FakeOpenAIClient( + vector_stores=FakeVectorStores( + error=openai.APITimeoutError(request=_OPENAI_REQUEST) + ) + ) + chat_model = ScriptedChatModel( + replies=[tool_call_reply("tuition"), "should never be used"] + ) + + with pytest.raises(ChatbotProviderError, match="search timed out"): + await _run(chat_model, monkeypatch, client=client) + + assert len(chat_model.calls) == 1 + + @pytest.mark.asyncio + async def test_empty_final_reply_is_a_provider_error( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + with pytest.raises(ChatbotProviderError, match="EmptyReply"): + await _run(ScriptedChatModel(replies=[" "]), monkeypatch) + + @pytest.mark.asyncio + async def test_missing_message_with_history_rejected( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + chat_model = ScriptedChatModel() + + with pytest.raises(ChatbotInputError, match="`message` is required"): + await _run( + chat_model, + monkeypatch, + user_message=None, + history=[{"role": "assistant", "content": "Hi!"}], + ) + assert chat_model.calls == [] + + +class TestRagSystemPrompt: + @pytest.mark.asyncio + async def test_config_instructions_and_goal_follow_the_rag_instruction( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + chat_model = ScriptedChatModel(replies=["Namaste!"]) + + await _run( + chat_model, + monkeypatch, + model_settings=_model_settings(instructions="Answer in Hindi."), + ) + + system_message = chat_model.calls[0][0] + assert isinstance(system_message, SystemMessage) + assert system_message.content == ( + f"{RAG_INSTRUCTION}" + "\n\n\nAnswer in Hindi.\n" + f"\n\n\n{GOAL}\n" + ) + + @pytest.mark.asyncio + async def test_without_config_instructions_only_goal_is_added( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + chat_model = ScriptedChatModel(replies=["Hello!"]) + + await _run(chat_model, monkeypatch) + + system_message = chat_model.calls[0][0] + assert system_message.content == f"{RAG_INSTRUCTION}\n\n\n{GOAL}\n" + + +class TestRagHistorySummarization: + @pytest.mark.asyncio + async def test_history_over_budget_is_summarized_before_the_tool_loop( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + history: list[dict[str, str]] = [] + for i in range(4): + history.append({"role": "user", "content": f"user {i}"}) + history.append({"role": "assistant", "content": f"bot {i}"}) + chat_model = ScriptedChatModel( + replies=[SUMMARY_TEXT, tool_call_reply("tuition"), "Tuition is 100."], + token_count=5000, + ) + + result = await _run( + chat_model, + monkeypatch, + history=history, + model_settings=_model_settings(history_token_limit=4000), + ) + + assert result.reply == "Tuition is 100." + # One summary call, then two agent calls; the tool round is never summarized. + _summary_call, first_agent_call, second_agent_call = chat_model.calls + assert first_agent_call[1].content == ( + f"[Earlier conversation summary] {SUMMARY_TEXT}" + ) + assert isinstance(first_agent_call[-1], HumanMessage) + assert first_agent_call[-1].content == "What is the tuition?" + assert isinstance(second_agent_call[-1], ToolMessage) + assert "user 0" not in [m.content for m in second_agent_call] + + +class TestRagProviderErrorMapping: + @pytest.mark.parametrize( + ("error", "expected_message"), + [ + ( + anthropic.APITimeoutError(request=_ANTHROPIC_REQUEST), + "(code: APITimeoutError)", + ), + ( + anthropic.APIConnectionError(request=_ANTHROPIC_REQUEST), + "(code: APIConnectionError)", + ), + ( + ValueError("thinking.type=disabled is not supported"), + "[KAAPI] Unsupported model configuration for 'claude-sonnet-4-6': " + "thinking.type=disabled is not supported", + ), + ( + anthropic.APIStatusError( + "Overloaded", + response=httpx.Response( + 529, + request=_ANTHROPIC_REQUEST, + headers={"request-id": "req_123"}, + ), + body=None, + ), + "[ANTHROPIC] Chat model call failed: Overloaded " + "(code: 529, request_id: req_123)", + ), + ], + ids=["timeout", "connection", "value_error", "status"], + ) + @pytest.mark.asyncio + async def test_anthropic_errors_become_provider_errors( + self, + monkeypatch: pytest.MonkeyPatch, + error: Exception, + expected_message: str, + ) -> None: + with pytest.raises(ChatbotProviderError) as exc_info: + await _run(ScriptedChatModel(error=error), monkeypatch) + + assert expected_message in str(exc_info.value) + assert exc_info.value.__cause__ is error diff --git a/backend/app/tests/services/chatbot/agents/rag/test_retriever.py b/backend/app/tests/services/chatbot/agents/rag/test_retriever.py new file mode 100644 index 000000000..2a1b6a05b --- /dev/null +++ b/backend/app/tests/services/chatbot/agents/rag/test_retriever.py @@ -0,0 +1,243 @@ +import httpx +import openai +import pytest +from langchain_core.messages import ToolMessage +from sqlmodel import Session + +from app.services.chatbot.agents.rag.prompts import NO_PASSAGES_FOUND +from app.services.chatbot.agents.rag.retriever import ( + build_search_tool, + format_passages, + get_knowledge_base_client, + search_knowledge_bases, +) +from app.services.chatbot.agents.rag.types import KnowledgeBasePassage +from app.services.chatbot.exceptions import ChatbotInputError, ChatbotProviderError +from app.tests.utils.chatbot import FakeOpenAIClient, FakeVectorStores, search_hit +from app.tests.utils.test_data import create_test_credential, create_test_project + +_OPENAI_REQUEST = httpx.Request("POST", "https://api.openai.com/v1/vector_stores") + + +def _client(**hits: list) -> FakeOpenAIClient: + return FakeOpenAIClient(vector_stores=FakeVectorStores(hits=hits)) + + +def _hits(prefix: str, count: int) -> list: + return [ + search_hit(file_id=f"{prefix}-{i}", filename=f"{prefix}-{i}.pdf", score=i / 100) + for i in range(count) + ] + + +class TestSearchKnowledgeBases: + def test_merges_knowledge_bases_and_keeps_overall_top_by_score(self) -> None: + client = _client( + vs_a=[ + search_hit(file_id="a-low", filename="a.pdf", score=0.2), + search_hit(file_id="a-high", filename="a.pdf", score=0.9), + ], + vs_b=[ + search_hit(file_id="b-mid", filename="b.pdf", score=0.5), + search_hit(file_id="b-high", filename="b.pdf", score=0.7), + ], + ) + + passages = search_knowledge_bases( + client=client, + knowledge_base_ids=("vs_a", "vs_b"), + query="fees", + max_num_results=3, + ) + + assert [p.file_id for p in passages] == ["a-high", "b-high", "b-mid"] + assert client.vector_stores.calls == [ + {"vector_store_id": "vs_a", "query": "fees", "max_num_results": 3}, + {"vector_store_id": "vs_b", "query": "fees", "max_num_results": 3}, + ] + + def test_defaults_to_five_results(self) -> None: + client = _client(vs_a=_hits("a", 7)) + + passages = search_knowledge_bases( + client=client, knowledge_base_ids=("vs_a",), query="q", max_num_results=None + ) + + assert len(passages) == 5 + assert client.vector_stores.calls[0]["max_num_results"] == 5 + + def test_caps_requested_results_at_fifty(self) -> None: + client = _client(vs_a=_hits("a", 60)) + + passages = search_knowledge_bases( + client=client, knowledge_base_ids=("vs_a",), query="q", max_num_results=80 + ) + + assert len(passages) == 50 + assert client.vector_stores.calls[0]["max_num_results"] == 50 + + def test_chunks_of_a_hit_are_joined_into_one_passage(self) -> None: + client = _client( + vs_a=[ + search_hit( + file_id="f1", + filename="guide.pdf", + score=0.8, + texts=("first chunk", "second chunk"), + ) + ] + ) + + passages = search_knowledge_bases( + client=client, knowledge_base_ids=("vs_a",), query="q", max_num_results=None + ) + + assert passages == [ + KnowledgeBasePassage( + file_id="f1", + filename="guide.pdf", + score=0.8, + text="first chunk\nsecond chunk", + ) + ] + + @pytest.mark.parametrize( + ("error", "expected_message"), + [ + ( + openai.APITimeoutError(request=_OPENAI_REQUEST), + "[KAAPI] Knowledge base search timed out — retry the turn. If " + "persistent, contact Kaapi. (code: APITimeoutError)", + ), + ( + openai.APIConnectionError(request=_OPENAI_REQUEST), + "[KAAPI] Network/DNS issue reaching OpenAI for knowledge base search — " + "check connectivity. If persistent, contact Kaapi. " + "(code: APIConnectionError)", + ), + ( + openai.NotFoundError( + "No vector store found", + response=httpx.Response( + 404, + request=_OPENAI_REQUEST, + headers={"x-request-id": "req_123"}, + ), + body=None, + ), + "[OPENAI] Knowledge base search failed for 'vs_a': No vector store " + "found (code: 404, request_id: req_123)", + ), + ( + openai.OpenAIError("unexpected payload"), + "[OPENAI] Knowledge base search failed for 'vs_a': unexpected payload " + "— retry the turn. If persistent, contact Kaapi. (code: OpenAIError)", + ), + ], + ids=["timeout", "connection", "status", "other"], + ) + def test_openai_errors_become_provider_errors( + self, error: openai.OpenAIError, expected_message: str + ) -> None: + client = FakeOpenAIClient(vector_stores=FakeVectorStores(error=error)) + + with pytest.raises(ChatbotProviderError) as exc_info: + search_knowledge_bases( + client=client, + knowledge_base_ids=("vs_a",), + query="q", + max_num_results=None, + ) + + assert str(exc_info.value) == expected_message + assert exc_info.value.__cause__ is error + + +class TestFormatPassages: + def test_no_passages_tells_the_model_nothing_was_found(self) -> None: + assert format_passages([]) == NO_PASSAGES_FOUND + + def test_passages_are_wrapped_with_escaped_filename(self) -> None: + passages = [ + KnowledgeBasePassage( + file_id="f1", filename='fees "2024".pdf', score=0.9, text="Fee is 100." + ), + KnowledgeBasePassage( + file_id="f2", filename="a&b.pdf", score=0.5, text="Due in May." + ), + ] + + assert format_passages(passages) == ( + '\nFee is 100.\n' + "\n\n" + '\nDue in May.\n' + ) + + +class TestBuildSearchTool: + def test_tool_returns_passages_as_content_and_hits_as_artifact(self) -> None: + client = _client( + vs_a=[search_hit(file_id="f1", filename="fees.pdf", score=0.9)] + ) + tool = build_search_tool( + client=client, knowledge_base_ids=("vs_a",), max_num_results=None + ) + + message = tool.invoke( + { + "type": "tool_call", + "name": "search_knowledge_base", + "args": {"query": "fees"}, + "id": "toolu_1", + } + ) + + assert isinstance(message, ToolMessage) + assert message.content == '\ntext\n' + assert message.artifact == [ + KnowledgeBasePassage( + file_id="f1", filename="fees.pdf", score=0.9, text="text" + ) + ] + assert client.vector_stores.calls[0]["query"] == "fees" + + def test_no_hits_returns_not_found_text_and_empty_artifact(self) -> None: + tool = build_search_tool( + client=_client(), knowledge_base_ids=("vs_a",), max_num_results=None + ) + + message = tool.invoke( + { + "type": "tool_call", + "name": "search_knowledge_base", + "args": {"query": "fees"}, + "id": "toolu_1", + } + ) + + assert message.content == NO_PASSAGES_FOUND + assert message.artifact == [] + + +class TestGetKnowledgeBaseClient: + def test_project_without_openai_credentials_rejected(self, db: Session) -> None: + project = create_test_project(db) + + with pytest.raises(ChatbotInputError, match="OpenAI credentials"): + get_knowledge_base_client( + session=db, + organization_id=project.organization_id, + project_id=project.id, + ) + + def test_client_uses_project_api_key(self, db: Session) -> None: + _credentials, project = create_test_credential(db) + + client = get_knowledge_base_client( + session=db, + organization_id=project.organization_id, + project_id=project.id, + ) + + assert client.api_key.startswith("sk-") + assert client.max_retries == 1 diff --git a/backend/app/tests/services/chatbot/agents/test_registry.py b/backend/app/tests/services/chatbot/agents/test_registry.py new file mode 100644 index 000000000..0c92a5f04 --- /dev/null +++ b/backend/app/tests/services/chatbot/agents/test_registry.py @@ -0,0 +1,92 @@ +import pytest + +from app.models import ChatbotAgentTypeEnum +from app.services.chatbot.agents.registry import run_agent_turn +from app.services.chatbot.exceptions import ChatbotInputError +from app.services.chatbot.utils import ChatbotModelSettings +from app.tests.utils.chatbot import ( + FakeOpenAIClient, + FakeVectorStores, + ScriptedChatModel, + patch_chat_model, + search_hit, + tool_call_reply, +) + +SYSTEM_PROMPT = "Collect the user's name and city." +_MODEL_SETTINGS = ChatbotModelSettings( + model="claude-sonnet-4-6", knowledge_base_ids=("vs_fees",) +) + + +class TestRunAgentTurn: + @pytest.mark.asyncio + async def test_rag_without_openai_client_rejected_before_calling_model( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + chat_model = ScriptedChatModel(replies=["should never be used"]) + patch_chat_model(monkeypatch, chat_model) + + with pytest.raises(ChatbotInputError, match="needs an OpenAI client"): + await run_agent_turn( + agent_type=ChatbotAgentTypeEnum.RAG, + system_prompt=SYSTEM_PROMPT, + history=[], + user_message="hello", + model_settings=_MODEL_SETTINGS, + ) + assert chat_model.calls == [] + + @pytest.mark.asyncio + async def test_onboarding_runs_field_extraction_and_ignores_client( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + chat_model = ScriptedChatModel(replies=["Which city?"]) + patch_chat_model(monkeypatch, chat_model) + client = FakeOpenAIClient() + + result = await run_agent_turn( + agent_type=ChatbotAgentTypeEnum.ONBOARDING, + system_prompt=SYSTEM_PROMPT, + history=[], + user_message="I'm Asha", + model_settings=_MODEL_SETTINGS, + openai_client=client, + ) + + assert result.reply == "Which city?" + # Only the onboarding agent runs the structured-output extraction step. + assert len(chat_model.extraction_calls) == 1 + assert chat_model.bound_tool_names == [] + assert client.vector_stores.calls == [] + + @pytest.mark.asyncio + async def test_rag_searches_with_the_given_client( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + chat_model = ScriptedChatModel( + replies=[tool_call_reply("tuition"), "Tuition is 100."] + ) + patch_chat_model(monkeypatch, chat_model) + client = FakeOpenAIClient( + vector_stores=FakeVectorStores( + hits={ + "vs_fees": [ + search_hit(file_id="file-fees", filename="fees.pdf", score=0.9) + ] + } + ) + ) + + result = await run_agent_turn( + agent_type=ChatbotAgentTypeEnum.RAG, + system_prompt=SYSTEM_PROMPT, + history=[], + user_message="What is the tuition?", + model_settings=_MODEL_SETTINGS, + openai_client=client, + ) + + assert result.reply == "Tuition is 100." + assert result.sources == [{"file_id": "file-fees", "filename": "fees.pdf"}] + assert chat_model.extraction_calls == [] diff --git a/backend/app/tests/services/chatbot/test_config.py b/backend/app/tests/services/chatbot/test_config.py index bf29597fa..6e8e20d5c 100644 --- a/backend/app/tests/services/chatbot/test_config.py +++ b/backend/app/tests/services/chatbot/test_config.py @@ -5,14 +5,20 @@ from fastapi import HTTPException from sqlmodel import Session +from app.crud import CollectionCrud from app.crud.config import ConfigVersionCrud -from app.models import Config -from app.services.chatbot.config import resolve_chatbot_model_settings +from app.models import Config, Project, ProviderType +from app.services.chatbot.config import ( + prepare_rag_search_client, + resolve_chatbot_model_settings, +) from app.services.chatbot.exceptions import ChatbotInputError from app.services.chatbot.utils import ChatbotModelSettings +from app.tests.utils.collection import get_vector_store_collection from app.tests.utils.test_data import ( build_anthropic_text_config_blob, create_test_config, + create_test_credential, create_test_project, create_test_version, ) @@ -73,6 +79,31 @@ def test_latest_version_with_server_defaults(self, db: Session) -> None: history_token_limit=4000, ) + def test_rag_params_are_carried_into_settings(self, db: Session) -> None: + config = create_test_config( + db, + config_blob=build_anthropic_text_config_blob( + instructions="Answer in Hindi.", + knowledge_base_ids=["vs_alpha", "vs_beta"], + max_num_results=8, + ), + ) + + model_settings = _resolve(db, config) + + assert model_settings.instructions == "Answer in Hindi." + assert model_settings.knowledge_base_ids == ("vs_alpha", "vs_beta") + assert model_settings.max_num_results == 8 + + def test_absent_rag_params_leave_empty_defaults(self, db: Session) -> None: + config = create_test_config(db, config_blob=build_anthropic_text_config_blob()) + + model_settings = _resolve(db, config) + + assert model_settings.instructions is None + assert model_settings.knowledge_base_ids == () + assert model_settings.max_num_results is None + def test_enabled_thinking_with_budget_below_max_tokens_is_accepted( self, db: Session ) -> None: @@ -215,3 +246,95 @@ def test_non_anthropic_effort_rejected(self, db: Session, effort: str) -> None: with pytest.raises(ChatbotInputError, match=f"'{effort}' is not an Anthropic"): _resolve(db, config) + + +def _prepare(project: Project, *knowledge_base_ids: str, db: Session): + return prepare_rag_search_client( + session=db, + organization_id=project.organization_id, + project_id=project.id, + model_settings=ChatbotModelSettings( + model="claude-sonnet-4-6", knowledge_base_ids=knowledge_base_ids + ), + ) + + +@pytest.fixture +def project_with_openai_creds(db: Session) -> Project: + _credentials, project = create_test_credential(db) + return project + + +class TestPrepareRagSearchClient: + def test_known_collections_and_credentials_return_project_client( + self, db: Session, project_with_openai_creds: Project + ) -> None: + first = get_vector_store_collection(db, project_with_openai_creds) + second = get_vector_store_collection(db, project_with_openai_creds) + + client = _prepare( + project_with_openai_creds, + first.knowledge_base_id, + second.knowledge_base_id, + db=db, + ) + + assert client.api_key.startswith("sk-") + + def test_empty_knowledge_base_ids_rejected( + self, db: Session, project_with_openai_creds: Project + ) -> None: + with pytest.raises(ChatbotInputError, match="non-empty"): + _prepare(project_with_openai_creds, db=db) + + def test_unknown_id_is_named_and_known_id_is_not( + self, db: Session, project_with_openai_creds: Project + ) -> None: + known = get_vector_store_collection(db, project_with_openai_creds) + + with pytest.raises(ChatbotInputError) as exc_info: + _prepare( + project_with_openai_creds, + known.knowledge_base_id, + "vs_missing", + db=db, + ) + + assert str(exc_info.value).endswith( + "collections: vs_missing; use the knowledge_base_id of an existing collection." + ) + assert known.knowledge_base_id not in str(exc_info.value) + + def test_deleted_collection_rejected( + self, db: Session, project_with_openai_creds: Project + ) -> None: + collection = get_vector_store_collection(db, project_with_openai_creds) + CollectionCrud(db, project_with_openai_creds.id).delete_by_id(collection.id) + + with pytest.raises(ChatbotInputError, match=collection.knowledge_base_id): + _prepare(project_with_openai_creds, collection.knowledge_base_id, db=db) + + def test_non_openai_collection_rejected( + self, db: Session, project_with_openai_creds: Project + ) -> None: + collection = get_vector_store_collection( + db, project_with_openai_creds, provider=ProviderType.google_aistudio + ) + + with pytest.raises(ChatbotInputError, match=collection.knowledge_base_id): + _prepare(project_with_openai_creds, collection.knowledge_base_id, db=db) + + def test_other_projects_collection_rejected( + self, db: Session, project_with_openai_creds: Project + ) -> None: + foreign = get_vector_store_collection(db, create_test_project(db)) + + with pytest.raises(ChatbotInputError, match=foreign.knowledge_base_id): + _prepare(project_with_openai_creds, foreign.knowledge_base_id, db=db) + + def test_missing_openai_credentials_rejected(self, db: Session) -> None: + project = create_test_project(db) + collection = get_vector_store_collection(db, project) + + with pytest.raises(ChatbotInputError, match="OpenAI credentials"): + _prepare(project, collection.knowledge_base_id, db=db) diff --git a/backend/app/tests/services/chatbot/test_session.py b/backend/app/tests/services/chatbot/test_session.py index f238a5e5c..5048b79a4 100644 --- a/backend/app/tests/services/chatbot/test_session.py +++ b/backend/app/tests/services/chatbot/test_session.py @@ -5,35 +5,46 @@ import anthropic import httpx +import openai import pytest from sqlmodel import Session +from app.crud import CollectionCrud from app.crud.chatbot import ( list_chatbot_session_messages, update_chatbot_session_progress, ) from app.crud.config import ConfigVersionCrud from app.models import ( + ChatbotAgentTypeEnum, ChatbotCompletionStatusEnum, ChatbotMessageRoleEnum, ChatbotSession, Config, Project, ) +from app.services.chatbot import config as config_service from app.services.chatbot import session as session_service from app.services.chatbot.agents.onboarding.types import FieldExtraction from app.services.chatbot.session import process_chatbot_session_message from app.tests.utils.chatbot import ( DEFAULT_SYSTEM_PROMPT, + FakeOpenAIClient, + FakeVectorStores, ScriptedChatModel, add_test_chatbot_messages, create_test_chatbot_session, extracted_fields, patch_chat_model, + search_hit, + tool_call_reply, ) +from app.tests.utils.collection import get_vector_store_collection from app.tests.utils.test_data import ( build_anthropic_text_config_blob, create_test_config, + create_test_credential, + create_test_project, ) WEBHOOK_URL = "https://hooks.example.com/chatbot" @@ -130,10 +141,12 @@ def test_first_call_stores_greeting_and_posts_success_webhook( "data": { "session_id": str(chatbot_session.id), "chatbot_ref": chatbot_session.chatbot_ref, + "agent_type": "onboarding", "reply": "Hi! What's your name?", "status": "active", "completion_status": "in_progress", "collected_fields": {}, + "sources": [], }, }, None, @@ -404,10 +417,12 @@ def test_turn_persists_merged_fields_and_completion_and_webhook_reports_them( "data": { "session_id": str(chatbot_session.id), "chatbot_ref": chatbot_session.chatbot_ref, + "agent_type": "onboarding", "reply": "Thanks, all set!", "status": "active", "completion_status": "complete", "collected_fields": {"name": "Asha", "city": "Pune"}, + "sources": [], }, } ] @@ -487,3 +502,158 @@ def test_complete_session_still_answers_new_messages( [payload] = callbacks.payloads assert payload["success"] is True assert payload["data"]["completion_status"] == "complete" + + +RAG_GOAL = "Answer questions about school fees from the documents." +_OPENAI_REQUEST = httpx.Request("POST", "https://api.openai.com/v1/vector_stores") + + +@pytest.fixture +def vector_stores(monkeypatch: pytest.MonkeyPatch) -> FakeVectorStores: + stores = FakeVectorStores() + real_get_client = config_service.get_knowledge_base_client + + def _get_client(**kwargs: Any) -> FakeOpenAIClient: + # The real lookup still enforces the project's OpenAI credentials; only the + # network-facing client is swapped. + real_get_client(**kwargs) + return FakeOpenAIClient(vector_stores=stores) + + monkeypatch.setattr(config_service, "get_knowledge_base_client", _get_client) + return stores + + +def _rag_session(db: Session, project: Project) -> tuple[ChatbotSession, str]: + collection = get_vector_store_collection(db, project) + rag_config = create_test_config( + db, + project_id=project.id, + config_blob=build_anthropic_text_config_blob( + knowledge_base_ids=[collection.knowledge_base_id] + ), + ) + chatbot_session = create_test_chatbot_session( + db, + config=rag_config, + system_prompt=RAG_GOAL, + agent_type=ChatbotAgentTypeEnum.RAG, + ) + add_test_chatbot_messages( + db, chatbot_session, [(ChatbotMessageRoleEnum.USER, "What is the tuition?")] + ) + return chatbot_session, collection.knowledge_base_id + + +class TestRagSessionTurn: + def test_answer_with_sources_is_stored_and_posted( + self, + db: Session, + callbacks: CallbackRecorder, + vector_stores: FakeVectorStores, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _credentials, project = create_test_credential(db) + chatbot_session, knowledge_base_id = _rag_session(db, project) + vector_stores.hits[knowledge_base_id] = [ + search_hit( + file_id="file-fees", + filename="fees.pdf", + score=0.9, + texts=("Tuition is 100 per term.",), + ) + ] + chat_model = ScriptedChatModel( + replies=[tool_call_reply("tuition fee"), "Tuition is 100 per term."] + ) + patch_chat_model(monkeypatch, chat_model) + + _process(chatbot_session, db) + + assert callbacks.payloads == [ + { + "success": True, + "data": { + "session_id": str(chatbot_session.id), + "chatbot_ref": chatbot_session.chatbot_ref, + "agent_type": "rag", + "reply": "Tuition is 100 per term.", + "status": "active", + "completion_status": "in_progress", + "collected_fields": {}, + "sources": [{"file_id": "file-fees", "filename": "fees.pdf"}], + }, + } + ] + assert _transcript(db, chatbot_session) == [ + ("user", "What is the tuition?"), + ("assistant", "Tuition is 100 per term."), + ] + assert vector_stores.calls[0]["vector_store_id"] == knowledge_base_id + + def test_missing_openai_credentials_post_failure_without_calling_model( + self, + db: Session, + callbacks: CallbackRecorder, + vector_stores: FakeVectorStores, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + chatbot_session, _kb_id = _rag_session(db, create_test_project(db)) + chat_model = ScriptedChatModel(replies=["should never be used"]) + patch_chat_model(monkeypatch, chat_model) + + _process(chatbot_session, db) + + assert chat_model.calls == [] + assert vector_stores.calls == [] + [payload] = callbacks.payloads + assert payload["success"] is False + assert "OpenAI credentials" in payload["error"] + assert payload["data"] == {"session_id": str(chatbot_session.id)} + assert _transcript(db, chatbot_session) == [("user", "What is the tuition?")] + + def test_knowledge_base_deleted_after_create_posts_failure( + self, + db: Session, + callbacks: CallbackRecorder, + vector_stores: FakeVectorStores, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _credentials, project = create_test_credential(db) + chatbot_session, knowledge_base_id = _rag_session(db, project) + crud = CollectionCrud(db, project.id) + [collection] = crud.read_all() + crud.delete_by_id(collection.id) + chat_model = ScriptedChatModel(replies=["should never be used"]) + patch_chat_model(monkeypatch, chat_model) + + _process(chatbot_session, db) + + assert chat_model.calls == [] + [payload] = callbacks.payloads + assert payload["success"] is False + assert knowledge_base_id in payload["error"] + + def test_search_failure_posts_failure_and_stores_no_reply( + self, + db: Session, + callbacks: CallbackRecorder, + vector_stores: FakeVectorStores, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _credentials, project = create_test_credential(db) + chatbot_session, _kb_id = _rag_session(db, project) + vector_stores.error = openai.APIConnectionError(request=_OPENAI_REQUEST) + chat_model = ScriptedChatModel( + replies=[tool_call_reply("tuition fee"), "should never be used"] + ) + patch_chat_model(monkeypatch, chat_model) + + _process(chatbot_session, db) + + [payload] = callbacks.payloads + assert payload["success"] is False + assert "(code: APIConnectionError)" in payload["error"] + assert "knowledge base search" in payload["error"] + assert payload["data"] == {"session_id": str(chatbot_session.id)} + assert _transcript(db, chatbot_session) == [("user", "What is the tuition?")] + assert len(chat_model.calls) == 1 diff --git a/backend/app/tests/utils/chatbot.py b/backend/app/tests/utils/chatbot.py index 7137bc766..4f2a77f9f 100644 --- a/backend/app/tests/utils/chatbot.py +++ b/backend/app/tests/utils/chatbot.py @@ -1,4 +1,5 @@ from collections.abc import Sequence +from dataclasses import dataclass, field from typing import Any import pytest @@ -6,12 +7,15 @@ from langchain_core.messages import AIMessage, BaseMessage from langchain_core.outputs import ChatGeneration, ChatResult from langchain_core.runnables import Runnable, RunnableLambda +from openai.types import VectorStoreSearchResponse +from openai.types.vector_store_search_response import Content from pydantic import Field from sqlmodel import Session from app.crud.chatbot import create_chatbot_session, create_chatbot_session_message from app.models import ( DEFAULT_CHATBOT_NAME, + ChatbotAgentTypeEnum, ChatbotMessageRoleEnum, ChatbotSession, ChatbotSessionMessage, @@ -19,6 +23,7 @@ Project, ) from app.services.chatbot.agents.onboarding import graph as onboarding_graph +from app.services.chatbot.agents.rag import graph as rag_graph from app.services.chatbot.agents.onboarding.types import ( ExtractedField, FieldExtraction, @@ -40,7 +45,8 @@ def extracted_fields(fields: dict[str, str]) -> list[ExtractedField]: class ScriptedChatModel(BaseChatModel): """Stands in for ChatAnthropic: replies from a script and records each call.""" - replies: list[str] = Field(default_factory=list) + # A str becomes a plain AIMessage; pass an AIMessage to script tool calls. + replies: list[str | AIMessage] = Field(default_factory=list) token_count: int = 0 error: Exception | None = None calls: list[list[BaseMessage]] = Field(default_factory=list) @@ -49,6 +55,9 @@ class ScriptedChatModel(BaseChatModel): extractions: list[Any] = Field(default_factory=list) extraction_calls: list[list[BaseMessage]] = Field(default_factory=list) structured_output_methods: list[str | None] = Field(default_factory=list) + # tool_choice each call was bound with (None = unbound or default "auto"). + call_tool_choices: list[Any] = Field(default_factory=list) + bound_tool_names: list[list[str]] = Field(default_factory=list) @property def _llm_type(self) -> str: @@ -62,11 +71,19 @@ def _generate( **kwargs: Any, ) -> ChatResult: self.calls.append(list(messages)) + self.call_tool_choices.append(kwargs.get("tool_choice")) if self.error is not None: raise self.error - reply = AIMessage(content=self.replies.pop(0)) + scripted = self.replies.pop(0) + reply = scripted if isinstance(scripted, AIMessage) else AIMessage(scripted) return ChatResult(generations=[ChatGeneration(message=reply)]) + def bind_tools( + self, tools: Sequence[Any], *, tool_choice: Any = None, **kwargs: Any + ) -> Runnable[Any, BaseMessage]: + self.bound_tool_names.append([tool.name for tool in tools]) + return self.bind(tool_choice=tool_choice) + def get_num_tokens_from_messages( self, messages: list[BaseMessage], tools: Sequence[Any] | None = None ) -> int: @@ -91,11 +108,72 @@ def patch_chat_model( monkeypatch: pytest.MonkeyPatch, chat_model: ScriptedChatModel ) -> None: # The Anthropic SDK is the external boundary; everything above it runs for real. - monkeypatch.setattr( - onboarding_graph, "get_chat_model", lambda _settings: chat_model + for graph_module in (onboarding_graph, rag_graph): + monkeypatch.setattr( + graph_module, "get_chat_model", lambda _settings: chat_model + ) + + +def tool_call_reply(*queries: str, text: str = "") -> AIMessage: + """An agent reply that calls the knowledge-base search once per query.""" + return AIMessage( + content=text, + tool_calls=[ + { + "name": "search_knowledge_base", + "args": {"query": query}, + "id": f"toolu_{index}_{query.replace(' ', '_')}", + } + for index, query in enumerate(queries) + ], ) +def search_hit( + *, file_id: str, filename: str, score: float, texts: Sequence[str] = ("text",) +) -> VectorStoreSearchResponse: + return VectorStoreSearchResponse( + attributes=None, + content=[Content(text=text, type="text") for text in texts], + file_id=file_id, + filename=filename, + score=score, + ) + + +@dataclass +class _SearchPage: + data: list[VectorStoreSearchResponse] + + +@dataclass +class FakeVectorStores: + hits: dict[str, list[VectorStoreSearchResponse]] = field(default_factory=dict) + error: Exception | None = None + calls: list[dict[str, Any]] = field(default_factory=list) + + def search( + self, vector_store_id: str, *, query: str, max_num_results: int + ) -> _SearchPage: + self.calls.append( + { + "vector_store_id": vector_store_id, + "query": query, + "max_num_results": max_num_results, + } + ) + if self.error is not None: + raise self.error + return _SearchPage(data=list(self.hits.get(vector_store_id, []))) + + +@dataclass +class FakeOpenAIClient: + """Stands in for openai.OpenAI; only vector-store search is used by the rag agent.""" + + vector_stores: FakeVectorStores = field(default_factory=FakeVectorStores) + + def create_test_chatbot_session( db: Session, *, @@ -105,6 +183,7 @@ def create_test_chatbot_session( chatbot_ref: str | None = None, config_version: int | None = None, system_prompt: str = DEFAULT_SYSTEM_PROMPT, + agent_type: ChatbotAgentTypeEnum = ChatbotAgentTypeEnum.ONBOARDING, ) -> ChatbotSession: project = db.get(Project, config.project_id) assert project is not None @@ -118,6 +197,7 @@ def create_test_chatbot_session( config_id=config.id, config_version=config_version, system_prompt=system_prompt, + agent_type=agent_type, ) diff --git a/backend/app/tests/utils/collection.py b/backend/app/tests/utils/collection.py index d662f339d..212b11ce5 100644 --- a/backend/app/tests/utils/collection.py +++ b/backend/app/tests/utils/collection.py @@ -31,6 +31,7 @@ def get_vector_store_collection( *, vector_store_id: Optional[str] = None, collection_id: Optional[UUID] = None, + provider: ProviderType = ProviderType.openai, ) -> Collection: """ Create a Collection configured for the Vector Store path. @@ -42,9 +43,9 @@ def get_vector_store_collection( collection = Collection( id=collection_id or uuid4(), project_id=project.id, - knowledge_base_provider=get_service_name("openai"), + knowledge_base_provider=get_service_name(provider.value), knowledge_base_id=vector_store_id, - provider=ProviderType.openai, + provider=provider, ) return CollectionCrud(db, project.id).create(collection) diff --git a/docs/wiki/domain-map.md b/docs/wiki/domain-map.md index 412da4a73..6f4277f30 100644 --- a/docs/wiki/domain-map.md +++ b/docs/wiki/domain-map.md @@ -37,7 +37,7 @@ APIKey → Organization, Project, User # programmatic access | TTSResult | tts_evaluation.py | Run, Org, Project | human annotation UI (logical) | | Assessment / AssessmentRun | assessment.py | Config, AssessmentSubmission, BatchJob, Job, Org, Project | console UI (logical); webhook consumers (logical) | | Document | document.py | Project, Document (parent) | DocumentCollection, DocTransformationJob, FineTuning, ModelEvaluation | -| Collection | collection.py | Project | DocumentCollection, CollectionJob; provider vector stores (logical) | +| Collection | collection.py | Project | DocumentCollection, CollectionJob; provider vector stores (logical); rag chatbot turns search OpenAI collections by `knowledge_base_id` from the config (logical) | | DocumentCollection | document_collection.py | Document, Collection | RAG lookups (logical) | | CollectionJob | collection_job.py | Collection, Project | async collection create/delete (logical) | | DocTransformationJob | doc_transformation_job.py | Document | doctransform service (logical) | diff --git a/docs/wiki/modules/chatbot.md b/docs/wiki/modules/chatbot.md index 4b482448a..f37907bbc 100644 --- a/docs/wiki/modules/chatbot.md +++ b/docs/wiki/modules/chatbot.md @@ -1,18 +1,18 @@ # Module: Chatbot -Prompt-driven chatbot ("Beneficiary Info Extractor") over a two-node LangGraph and a LangChain chat model (Anthropic). The admin's free-text `system_prompt` defines the goal, the fields to collect, and the rules; each turn extracts the fields the user has given so far and then replies. The model and reasoning settings come from a saved Kaapi config (`config_id` + optional `config_version`, latest if omitted); only Anthropic text completion configs are supported. Conversations are server-side sessions whose turns run on Celery and reply via webhook. No deep-dive doc yet. +Prompt-driven chatbot over LangGraph and a LangChain chat model (Anthropic), with two agents picked per session by `agent_type`: `onboarding` ("Beneficiary Info Extractor"; the admin's free-text `system_prompt` defines the goal, the fields to collect, and the rules; each turn extracts the fields the user has given so far and then replies) and `rag` (answers questions from the config's knowledge bases, i.e. OpenAI vector stores searched with the project's OpenAI credentials; `system_prompt` is the goal and decides scope). The model and reasoning settings come from a saved Kaapi config (`config_id` + optional `config_version`, latest if omitted); only Anthropic text completion configs are supported. Conversations are server-side sessions whose turns run on Celery and reply via webhook. No deep-dive doc yet. All paths relative to `backend/app/`. ## Routes - `api/routes/chatbot/session.py` - - `POST /chatbot/sessions` (create; validates the config; `chatbot_ref` omitted → new chatbot, ref generated by `services/chatbot/chatbot_ref.py:generate_chatbot_ref`; given → existing chatbot, supersedes the contact's active sessions for it → `{session_id, chatbot_ref}`; 409 on a concurrent create, 503 if no unused ref could be generated) - - `GET /chatbot/sessions?contact_ref=&chatbot_ref=` (latest session for a contact in the caller's project; `chatbot_ref` optional — omitted means newest across all chatbots → `ChatbotSessionSummary`, incl. `chatbot_ref`, `system_prompt`, `status`, `completion_status`, `collected_fields`) + - `POST /chatbot/sessions` (create; validates the config, and for `agent_type=rag` also the knowledge bases + OpenAI credentials via `prepare_rag_search_client` → 422; `chatbot_ref` omitted → new chatbot, ref generated by `services/chatbot/chatbot_ref.py:generate_chatbot_ref`; given → existing chatbot, supersedes the contact's active sessions for it → `{session_id, chatbot_ref}`; 409 on a concurrent create, 503 if no unused ref could be generated) + - `GET /chatbot/sessions?contact_ref=&chatbot_ref=` (latest session for a contact in the caller's project; `chatbot_ref` optional — omitted means newest across all chatbots → `ChatbotSessionSummary`, incl. `chatbot_ref`, `system_prompt`, `agent_type`, `status`, `completion_status`, `collected_fields`) - `POST /chatbot/sessions/{session_id}/message` (stores the user message, enqueues the turn → 202; reply POSTed to `webhook_url`; 409 if the session is inactive) - All require project scope and the `CHATBOT` feature flag (`require_feature` on the `/chatbot` router → 403 when off). Migration `089` only adds the enum value — no rows are seeded, so the flag is off until a superuser creates it per project via `POST /feature-flags` (`enabled: true`). ## Tables (SQLModel, `models/chatbot.py`) -- `chatbot_session` (`ChatbotSession`) — `name` (display only), `contact_ref` (end user), `chatbot_ref` (stable chatbot id, generated as `LLMResponseWithPrompt-<5 hex>` on the first create and unique per project; caller sends it back to start new sessions for the same chatbot), `config_id` (FK `config`, RESTRICT) / `config_version`, `system_prompt`, `status` (`ChatbotSessionStatusEnum` active/inactive = session in use), `completion_status` (`ChatbotCompletionStatusEnum` in_progress/complete/abandoned), `collected_fields` (JSON flat `{field_name: value}`), org/project FKs. Composite index `(project_id, chatbot_ref, contact_ref)` for GET-latest and the generated-ref collision check, plus a partial unique index on `(project_id, contact_ref, chatbot_ref) WHERE status = 'active'`. Migration `088`. +- `chatbot_session` (`ChatbotSession`) — `name` (display only), `contact_ref` (end user), `chatbot_ref` (stable chatbot id, generated as `LLMResponseWithPrompt-<5 hex>` on the first create and unique per project; caller sends it back to start new sessions for the same chatbot), `config_id` (FK `config`, RESTRICT) / `config_version`, `system_prompt`, `agent_type` (`ChatbotAgentTypeEnum` onboarding/rag, PG enum `chatbotagenttypeenum`, server default `onboarding`; migration `090`), `status` (`ChatbotSessionStatusEnum` active/inactive = session in use), `completion_status` (`ChatbotCompletionStatusEnum` in_progress/complete/abandoned), `collected_fields` (JSON flat `{field_name: value}`), org/project FKs. Composite index `(project_id, chatbot_ref, contact_ref)` for GET-latest and the generated-ref collision check, plus a partial unique index on `(project_id, contact_ref, chatbot_ref) WHERE status = 'active'`. Migration `088`. - `chatbot_session_message` (`ChatbotSessionMessage`) — ordered transcript rows (`role` user/assistant; the prompt lives on the session). - Session lifecycle: one active session per contact + chatbot. `active`+`in_progress` → `active`+`complete` as fields are collected. Creating a new session for the same `contact_ref` + `chatbot_ref` (same transaction as the insert) moves the previous active ones to `inactive`+`complete`, or `inactive`+`abandoned` if still in progress. Inactive sessions reject new messages (409); a turn already queued when the session was superseded is skipped by the worker, which posts a failure webhook instead. A partial unique index `uq_chatbot_session_active_contact_chatbot` (`WHERE status = 'active'`) makes a concurrent double-create fail; crud raises `ChatbotSessionConflictError` → 409. - Request/response shapes: `CreateChatbotSessionRequest/Response`, `ChatbotSessionSummary`, `SendChatbotSessionMessageRequest/Response`. @@ -20,16 +20,29 @@ All paths relative to `backend/app/`. ## Turn flow (onboarding agent, `services/chatbot/agents/onboarding/graph.py`) 1. `extract_fields_node` — skipped until there is a real user turn (opening greeting). Otherwise calls the model with structured output (`FieldExtraction{collected_fields, missing_fields, is_complete}`) over the admin prompt, the previously collected map, and the (summarized) conversation. Uses Anthropic native structured output (`with_structured_output(..., method="json_schema")`) with the config's own thinking/effort, because forced tool use is rejected with thinking on and Opus 5.5 can't disable thinking. Merged in code: `{**prior, **extracted}`; `is_complete` only when nothing is missing. 2. `chatbot_node` — replies using the same summarized view plus a system note listing collected/missing fields, so it asks only for what's missing and wraps up when complete. -3. `run_chatbot_turn(..., collected_fields=prior)` → `ChatbotTurnResult{transcript, collected_fields, is_complete}` (`.completion_status` maps to the enum). Anthropic errors from either call, and langchain_anthropic's client-side model-config `ValueError`s ("[KAAPI] Unsupported model configuration"), map to `ChatbotProviderError`, so the session path always sends a failure webhook. +3. `run_chatbot_turn(..., collected_fields=prior)` → `ChatbotTurnResult{reply, collected_fields, is_complete, sources=[]}` (`.completion_status` maps to the enum). Anthropic errors from either call, and langchain_anthropic's client-side model-config `ValueError`s ("[KAAPI] Unsupported model configuration"), map to `ChatbotProviderError` via the shared `provider_errors.to_provider_error`, so the session path always sends a failure webhook. + +## Turn flow (rag agent, `services/chatbot/agents/rag/graph.py`) +1. `prepare_context_node` — runs `summarize_if_needed` once and replaces the state messages (`RemoveMessage(REMOVE_ALL_MESSAGES)`), so the tool loop is never summarized mid tool_use/tool_result pair. +2. `agent_node` — chat model with `search_knowledge_base` bound (`tool_choice` auto; forced tool use is rejected with thinking). System prompt = `RAG_INSTRUCTION` + `` (config `params.instructions`, if set) + `` (session `system_prompt`). After `MAX_TOOL_ROUNDS` (3) tool-calling AIMessages it is re-bound with `tool_choice={"type": "none"}` so it must answer. +3. `tools_condition` → `ToolNode` (`tools`) → back to `agent_node`; END once the model stops calling tools. Only bad tool args (`ToolInvocationError`) are fed back to the model; a search failure (`ChatbotProviderError`) fails the turn. +4. `search_knowledge_base(query)` (`retriever.py`, `content_and_artifact`) — `client.vector_stores.search` on every `knowledge_base_ids` entry, merged, sorted by score, top `max_num_results` (default 5, capped at 50); content is `text` passages, artifact the `KnowledgeBasePassage` list. OpenAI errors → `ChatbotProviderError` (`[OPENAI] Knowledge base search failed ...`; timeouts/connection `[KAAPI]`). +5. `run_rag_turn(..., openai_client)` → `ChatbotTurnResult{reply, collected_fields (unchanged), is_complete=False, sources}`; `sources` are the retrieved files deduped by `file_id`. ## Services / CRUD - `services/chatbot/agents/onboarding/graph.py` — `run_chatbot_turn` (public entry point), `build_onboarding_graph`, the extraction/reply nodes - `services/chatbot/agents/onboarding/prompts.py` — extraction instruction/templates, progress-note templates, `EXTRACTION_METHOD` - `services/chatbot/agents/onboarding/types.py` — `OnboardingState`, `FieldExtraction` -- `services/chatbot/history.py` — agent-agnostic history plumbing: `summarize_if_needed`, `to_langchain_messages` / `to_message_dicts`, kickoff message, `has_user_turn`, `insert_after_leading_system` -- `services/chatbot/types.py` — shared `ChatbotTurnResult`, `ChatbotMessageDict` -- `services/chatbot/session.py` — `process_chatbot_session_message` (Celery worker body: load session + history; an inactive session gets a failure webhook and no turn; otherwise run the turn with the stored `collected_fields`, store the reply and update `collected_fields` / `completion_status`, POST the webhook) -- `services/chatbot/config.py` — `resolve_chatbot_model_settings` (saved config + request overrides → `ChatbotModelSettings`; 404 via `ConfigVersionCrud`, 422 for non-Anthropic/non-text/no-model configs) +- `services/chatbot/agents/rag/graph.py` — `run_rag_turn` (public entry point), `build_rag_graph`, `build_rag_system_prompt`, `MAX_TOOL_ROUNDS` +- `services/chatbot/agents/rag/retriever.py` — `get_knowledge_base_client` (project OpenAI credentials → client; missing → `ChatbotInputError`), `search_knowledge_bases`, `format_passages`, `build_search_tool` +- `services/chatbot/agents/rag/prompts.py` — `RAG_INSTRUCTION`, tool name/description, passage templates +- `services/chatbot/agents/rag/types.py` — `RagState`, `KnowledgeBasePassage`, `SearchKnowledgeBaseArgs` +- `services/chatbot/agents/registry.py` — `run_agent_turn` (dispatches on `ChatbotAgentTypeEnum`; rag needs the OpenAI client) +- `services/chatbot/provider_errors.py` — `to_provider_error`, `CHAT_MODEL_ERRORS` (Anthropic error mapping shared by both agents) +- `services/chatbot/history.py` — agent-agnostic history plumbing: `summarize_if_needed`, `to_langchain_messages`, `extract_final_reply` (empty → `ChatbotProviderError`), `validate_turn_input`, kickoff message, `has_user_turn`, `insert_after_leading_system` +- `services/chatbot/types.py` — shared `ChatbotTurnResult`, `ChatbotSource`, `ChatbotMessageDict` +- `services/chatbot/session.py` — `process_chatbot_session_message` (Celery worker body: load session + history; an inactive session gets a failure webhook and no turn; for rag sessions builds the OpenAI client in the first DB block (missing credentials / bad knowledge bases → failure webhook); runs `run_agent_turn` with the stored `collected_fields`, stores the reply and updates `collected_fields` / `completion_status`, POSTs the webhook) +- `services/chatbot/config.py` — `resolve_chatbot_model_settings` (saved config + request overrides → `ChatbotModelSettings`, incl. `instructions`, `knowledge_base_ids`, `max_num_results`; 404 via `ConfigVersionCrud`, 422 for non-Anthropic/non-text/no-model configs), `prepare_rag_search_client` (non-empty `knowledge_base_ids`, each a non-deleted OpenAI collection's `knowledge_base_id` via `CollectionCrud.read_all`, plus OpenAI credentials → client) - `services/chatbot/utils.py` — `ChatbotModelSettings`, defaults (`DEFAULT_MAX_OUTPUT_TOKENS`, `DEFAULT_HISTORY_TOKEN_LIMIT`, `FLOW_KEEP_RECENT_MESSAGES`), `get_chat_model` - `services/chatbot/chatbot_ref.py` — `generate_chatbot_ref` (random `LLMResponseWithPrompt-<5 hex>`, retried against `chatbot_ref_exists` up to `CHATBOT_REF_MAX_ATTEMPTS`, else `ChatbotRefGenerationError`) - `crud/chatbot.py` — `chatbot_ref_exists`, `create_chatbot_session` (also supersedes prior active sessions via `_deactivate_active_chatbot_sessions`; raises `ChatbotSessionConflictError` on the active-session unique index), `get_chatbot_session`, `get_latest_chatbot_session`, `update_chatbot_session_progress`, `create_chatbot_session_message`, `list_chatbot_session_messages` @@ -37,8 +50,9 @@ All paths relative to `backend/app/`. ## Async - `celery/tasks/job_execution.py` — `run_chatbot_session_message` (enqueued via `celery/utils.py:start_chatbot_session_message`) → `process_chatbot_session_message`. -- Webhook success payload: `{"success": true, "data": {"session_id", "chatbot_ref", "reply", "status", "completion_status", "collected_fields"}}`; failure: `{"success": false, "error", "data": {"session_id"}}`. +- Webhook success payload: `{"success": true, "data": {"session_id", "chatbot_ref", "agent_type", "reply", "status", "completion_status", "collected_fields", "sources": [{"file_id", "filename"}]}}`; failure: `{"success": false, "error", "data": {"session_id"}}`. ## External -- Anthropic via `langchain-anthropic` (`ANTHROPIC_API_KEY` setting) — two calls per turn once the user has spoken (extraction + reply), plus one summary call when history exceeds the token budget. +- Anthropic via `langchain-anthropic` (`ANTHROPIC_API_KEY` setting) — onboarding: two calls per turn once the user has spoken (extraction + reply); rag: one call per tool round (≤ 4); plus one summary call when history exceeds the token budget. +- OpenAI vector store search (`vector_stores.search`, rag only) with the project's OpenAI credentials, one call per knowledge base per tool call; `OPENAI_TIMEOUT_SECONDS`, one SDK retry. - Caller webhook (`send_callback`, HMAC-signed with the project webhook secret).