diff --git a/docs/administrator/upgrade_and_migration/database_change_log.md b/docs/administrator/upgrade_and_migration/database_change_log.md index 9547b21d51..63b0bbbe74 100644 --- a/docs/administrator/upgrade_and_migration/database_change_log.md +++ b/docs/administrator/upgrade_and_migration/database_change_log.md @@ -16,7 +16,7 @@ A record of the schema and data changes the Go backend applies to a RAGFlow data - The marker is a row in `system_settings` with the key `mysql_migration.database.version`. - For version-gated steps, the marker prevents a migration step from running again after that step has completed successfully. -- Go defines the versions as constants: `modelMigrationBaseVersion` and `modelMigrationTargetVersion` in [migration_version.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/migration_version.go), `conversationHistoryTargetVersion` in [conversation_history_migration.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/conversation_history_migration.go), `knowledgebaseParserConfigTargetVersion` in [migration.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/migration.go), and `ingestionParserConfigTargetVersion` in [ingestion_parser_config_migration.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/ingestion_parser_config_migration.go). The two parser-config constants share `v1.0.0-rc2`. The DSL multi-version backfill is gated by `dslVersionTargetVersion` (`v1.0.0-rc4`) in [dsl_version_migration.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/dsl_version_migration.go). +- Go defines the versions as constants: `modelMigrationBaseVersion` and `modelMigrationTargetVersion` in [migration_version.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/migration_version.go), `conversationHistoryTargetVersion` in [conversation_history_migration.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/conversation_history_migration.go), `knowledgebaseParserConfigTargetVersion` in [migration.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/migration.go), and `ingestionParserConfigTargetVersion` in [ingestion_parser_config_migration.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/ingestion_parser_config_migration.go). The two parser-config constants share `v1.0.0-rc2`. The DSL multi-version backfill is gated by `dslVersionTargetVersion` (`v1.0.0-rc3.dev6`) in [dsl_version_migration.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/dsl_version_migration.go). - The current permission migration targets `permissionMigrationTargetVersion` (`v1.0.0-rc3.dev4`) in [permission_migration.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/permission_migration.go). Tenant default-model migration targets `tenantDefaultModelsTargetVersion` (`v1.0.0-rc3.dev5`) in [tenant_default_models_migration.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/tenant_default_models_migration.go). - The compilation queue split records `knowledgeCompileQueueTargetVersion` (`v1.0.0-rc3.dev5`), defined in [knowledge_compile_queue_migration.go](https://github.com/infiniflow/ragflow/blob/main/internal/dao/knowledge_compile_queue_migration.go), after both queues finish. Source-column presence controls the queue copy even when the marker already contains this or a newer version. - A version-gated step runs when the stored version is lower than the step's version. Versions are compared with Python PEP 440 semantics, where a `dev` prerelease sorts *before* the release it precedes (`v1.0.0-rc1.dev1` < `v1.0.0-rc1`). A missing or unparsable marker never skips a step. @@ -35,6 +35,8 @@ Startup order in `InitDB` for `ragflow_server --migrate` ([database.go](https:// 9. `migrateTenantDefaultModels` — moves tenant-wide model selections to owner memberships, drops the tenant columns, and records `v1.0.0-rc3.dev5`. 10. `migrateTenantModelProviderUserUniqueIndex` — reconciles provider ownership indexes. 11. `migrateKnowledgeCompileQueues` — copies the existing compilation queues and drops their source columns, then records `v1.0.0-rc3.dev5` if the marker is older. It runs after tenant default-model migration so the queue split cannot make that step skip its data changes. Queue copying remains gated on source-column presence even when the preceding step has already recorded `dev5`. +12. `MigrateDSLVersionColumns` and `BackfillDSLVersionBindings` — add pipeline version columns and bind existing dataflow canvases, then record `v1.0.0-rc3.dev6`. +13. `migrateConversationReferences` — move answer references into message rows and record `v1.0.0-rc3.dev7`. ## Version overview @@ -50,7 +52,8 @@ Startup order in `InitDB` for `ragflow_server --migrate` ([database.go](https:// | `v1.0.0-rc3.dev4` | Normalize legacy tenant invitation roles and statuses | `permissionMigrationTargetVersion` | | `v1.0.0-rc3.dev5` | Move tenant default-model selections to owner memberships and drop the tenant columns | `tenantDefaultModelsTargetVersion` | | `v1.0.0-rc3.dev5` | Split compilation queues into `knowledge_compile_docs_backlog` and `knowledge_compile_docs_inflight`, then drop their source columns | `knowledgeCompileQueueTargetVersion`; source-column presence controls copying | -| `v1.0.0-rc4` | Lock existing dataflow-canvas dataset/document bindings to the canvas latest version (the `pipeline_version_id` columns are added by the version-less step below) | `dslVersionTargetVersion` | +| `v1.0.0-rc3.dev6` | Lock existing dataflow-canvas dataset/document bindings to the canvas latest version (the `pipeline_version_id` columns are added by the version-less step below) | `dslVersionTargetVersion` | +| `v1.0.0-rc3.dev7` | Store conversation references in each assistant message | `conversationReferenceTargetVersion` | The two tenant model steps are cumulative: a database at `v0.26.0` still runs the `v1.0.0-rc1` step, and a database at or above `v1.0.0-rc1` runs neither. @@ -257,7 +260,7 @@ This migration requires a coordinated service stop: Starting a server while either source column remains fails with `knowledge compilation queues need migration`. A database version marker at or above the current release does not bypass this requirement. If migration fails, keep services stopped, resolve the reported error, and rerun `ragflow_server --migrate`. -## v1.0.0-rc4 +## v1.0.0-rc3.dev6 Locks every pre-existing dataflow-canvas binding to the latest version of that canvas, so already-configured datasets and documents keep parsing against the version that was current when the upgrade ran (design section 5/B4 sticky snapshot). The `pipeline_version_id` columns it writes are added by the version-less step below; this step only backfills their values and is gated by the `dslVersionTargetVersion` marker. @@ -270,9 +273,50 @@ Locks every pre-existing dataflow-canvas binding to the latest version of that c Execution details: - Rows are read in id-ordered batches of 256, each batch in its own transaction, so an interrupted run resumes cleanly. -- The marker advances to `v1.0.0-rc4` only after both tables are processed. Databases already marked `v1.0.0-rc4` or later skip the backfill. +- The marker advances to `v1.0.0-rc3.dev6` only after both tables are processed. Databases already marked `v1.0.0-rc3.dev6` or later skip the backfill. -On success the step writes `v1.0.0-rc4` into the version marker. +On success the step writes `v1.0.0-rc3.dev6` into the version marker. + +--- + +## v1.0.0-rc3.dev7 + +Adds a nullable `reference` LONGTEXT column to `conversation_message` and +`api_4_conversation_message`. New chat, agent, and channel answers store their +reference object directly on their message row. + +Run `ragflow_server --migrate` with writers stopped before starting the new +services. The migration copies `conversation_reference` and +`api_4_conversation_reference` into the corresponding assistant messages. It +uses `message_position` when present; older array entries follow assistant +order after the first user message, excluding the prologue. Rows are read in +batches of 100 and written in transactions. Reruns preserve reference objects +already stored on messages. References duplicated in message metadata are moved +into the column and removed from metadata, so generation reads avoid that data. + +After both message tables are copied successfully, the migration drops +`conversation_reference` and `api_4_conversation_reference`. Automatic schema +creation no longer recreates these tables. Source-table presence triggers copying +and cleanup even if the marker already contains `v1.0.0-rc3.dev7` or a newer version; +a newer marker is preserved. + +References without a surviving assistant message are skipped and counted in a +warning. These include references left behind by deleted messages, empty histories, +and unsuccessful turns. They are removed with the source tables and are never +attached to another answer. Invalid JSON or a copy error stops the migration before +either source table is dropped. Resolve the reported issue and rerun it. Completed +batches are safe to retry. The marker advances to `v1.0.0-rc3.dev7` only after both +copies and table drops finish. A failed table drop can be retried even if the +other source table has already been dropped. + +Session responses return `reference` within each answer, with no session-level +reference array when `include_reference=true`. Clients must read +`messages[i].reference` (or `message[i].reference` for agent sessions). +`include_reference` and `include_prompt` both default to `false`; omitted fields +are excluded from SQL as well as the response. The UI requests the fields it +displays. Generation history SQL omits +`reference`, `prompt`, `usage`, and `elapsed_ms`; session detail reads request +the fields they display. --- diff --git a/docs/references/http_api_reference.md b/docs/references/http_api_reference.md index dc496e46cb..accff398b9 100644 --- a/docs/references/http_api_reference.md +++ b/docs/references/http_api_reference.md @@ -3203,6 +3203,11 @@ curl --request GET \ When `id` or `name` is provided, exact filtering takes precedence over `keywords`. +- `include_reference`: (*Query parameter*) `boolean`, defaults to `false`. + Include reference objects in answer messages when history is included. +- `include_prompt`: (*Query parameter*) `boolean`, defaults to `false`. + Include recorded prompt strings in answer messages when history is included. + #### Response Success: @@ -3490,7 +3495,6 @@ Success: } ], "name": "new session", - "reference": [], "update_date": "Fri, 11 Oct 2024 08:46:43 GMT", "update_time": 1728636403974, "user_id": "" @@ -3516,13 +3520,26 @@ Failure: Gets a specific session of a specified chat assistant, including its messages, references, and avatar. -Completed assistant entries in `messages` include the recorded `prompt` (string), -`usage` (object with `prompt_tokens`, `completion_tokens`, and `total_tokens`), and -`elapsed_ms` (number, in milliseconds). These fields are also returned by the -session list when history is included. Older messages and turns without recorded -usage or timing omit the unavailable fields. They are used for displaying answer -details and are excluded from history passed to the LLM for subsequent turns. -Agentic answers store and return `usage` and `elapsed_ms` without a `prompt`. +Completed assistant entries in `messages` include recorded `usage` (an object +with `prompt_tokens`, `completion_tokens`, and `total_tokens`) and `elapsed_ms` +(a number in milliseconds). Large answer fields are opt-in: + +- `include_reference`: boolean, defaults to `false`. Set to `true` to return a + `reference` object containing `chunks` and `doc_aggs` inside each answer. +- `include_prompt`: boolean, defaults to `false`. Set to `true` to return the + recorded prompt string inside each answer. + +These query parameters apply only to chat and agent session list and detail +endpoints. Fields omitted by these flags are excluded from SQL as well as the +response. Invalid boolean values return an argument error on these read endpoints. +Create, update, message deletion, and feedback endpoints ignore these query +parameters and omit `reference` and `prompt` from their responses. +There is no session-level reference array. + +The UI requests references for citations and source previews, and prompts for +answer details. Older messages omit unavailable usage, timing, or prompt fields. +Generation history SQL omits `reference`, `prompt`, `usage`, and `elapsed_ms`, and +these fields are excluded from LLM input. Agentic answers have no recorded prompt. #### Request @@ -3535,7 +3552,7 @@ Agentic answers store and return `usage` and `elapsed_ms` without a `prompt`. ```bash curl --request GET \ - --url http://{address}/api/v1/chats/{chat_id}/sessions/{session_id} \ + --url 'http://{address}/api/v1/chats/{chat_id}/sessions/{session_id}?include_reference=true&include_prompt=true' \ --header 'Authorization: Bearer ' ``` @@ -3562,9 +3579,19 @@ Success: { "content": "Hi! I am your assistant, can I help you?", "role": "assistant" + }, + { + "id": "turn-1", + "role": "user", + "content": "What does the document say?" + }, + { + "id": "turn-1", + "role": "assistant", + "content": "Here is the answer.", + "reference": {"chunks": [], "doc_aggs": []} } - ], - "reference": [] + ] } } ``` @@ -3620,8 +3647,7 @@ Success: returns the updated session object. "data": { "chat_id": "2ca4b22e878011ef88fe0242ac120005", "id": "4606b4ec87ad11efbc4f0242ac120006", - "messages": [], - "reference": [] + "messages": [] } } ``` diff --git a/internal/dao/api_for_conversation.go b/internal/dao/api_for_conversation.go index 989eae57ba..0ba9fb8cab 100644 --- a/internal/dao/api_for_conversation.go +++ b/internal/dao/api_for_conversation.go @@ -56,10 +56,7 @@ func (dao *API4ConversationDAO) Create(ctx context.Context, db *gorm.DB, conv *e if err := tx.Create(conv).Error; err != nil { return err } - if err := createHistory(ctx, tx, apiConversationMessageTable, "message", conv.ID, conv.Message); err != nil { - return err - } - return createHistory(ctx, tx, apiConversationReferenceTable, "reference", conv.ID, conv.Reference) + return createHistory(ctx, tx, apiConversationMessageTable, "message", conv.ID, conv.Message) }) } @@ -85,7 +82,7 @@ func (dao *API4ConversationDAO) UpdateHistory(ctx context.Context, db *gorm.DB, if err := query.Updates(updates).Error; err != nil { return err } - return updateConversationHistory(ctx, tx, apiConversationMessageTable, apiConversationReferenceTable, sessionID, history) + return updateConversationHistory(ctx, tx, apiConversationMessageTable, sessionID, history) }) } @@ -118,12 +115,16 @@ func (dao *API4ConversationDAO) Stats(ctx context.Context, db *gorm.DB, tenantID return rows, err } -func (dao *API4ConversationDAO) GetBySessionID(ctx context.Context, db *gorm.DB, sessionID, agentID string) (*entity.API4Conversation, error) { +func (dao *API4ConversationDAO) GetBySessionID(ctx context.Context, db *gorm.DB, sessionID, agentID string, options ...ConversationHistoryOptions) (*entity.API4Conversation, error) { + historyOptions := ConversationHistoryOptions{} + if len(options) > 0 { + historyOptions = options[0] + } result, err := dao.GetMetadataBySessionID(ctx, db, sessionID, agentID) if err != nil || result == nil { return result, err } - if err := hydrateAPIConversations(ctx, db, []*entity.API4Conversation{result}); err != nil { + if err := hydrateAPIConversations(ctx, db, []*entity.API4Conversation{result}, historyOptions); err != nil { return nil, err } return result, nil @@ -144,7 +145,11 @@ func (dao *API4ConversationDAO) GetMetadataBySessionID(ctx context.Context, db * // GetByID returns a conversation without requiring the caller to know its // agent. It is used when the session itself is the authorization resource. -func (dao *API4ConversationDAO) GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.API4Conversation, error) { +func (dao *API4ConversationDAO) GetByID(ctx context.Context, db *gorm.DB, id string, options ...ConversationHistoryOptions) (*entity.API4Conversation, error) { + historyOptions := ConversationHistoryOptions{} + if len(options) > 0 { + historyOptions = options[0] + } var result entity.API4Conversation tx := db.WithContext(ctx).Session(&gorm.Session{QueryFields: true}).Take(&result, "id = ?", id) if tx.Error != nil { @@ -153,7 +158,7 @@ func (dao *API4ConversationDAO) GetByID(ctx context.Context, db *gorm.DB, id str } return nil, tx.Error } - if err := hydrateAPIConversations(ctx, db, []*entity.API4Conversation{&result}); err != nil { + if err := hydrateAPIConversations(ctx, db, []*entity.API4Conversation{&result}, historyOptions); err != nil { return nil, err } return &result, nil @@ -174,7 +179,7 @@ func (dao *API4ConversationDAO) DeleteBySessionIDAndAgentID(ctx context.Context, if err := tx.Model(&entity.API4Conversation{}).Clauses(clause.Locking{Strength: "UPDATE"}).Where("id = ? AND dialog_id = ?", sessionID, agentID).Order("id").Pluck("id", &ids).Error; err != nil { return err } - if err := deleteHistory(ctx, tx, []string{apiConversationMessageTable, apiConversationReferenceTable}, ids); err != nil { + if err := deleteHistory(ctx, tx, []string{apiConversationMessageTable}, ids); err != nil { return err } result := tx.Where("id = ? AND dialog_id = ?", sessionID, agentID).Delete(&entity.API4Conversation{}) @@ -195,7 +200,7 @@ func (dao *API4ConversationDAO) DeleteByDialogIDs(ctx context.Context, db *gorm. if err := tx.Model(&entity.API4Conversation{}).Clauses(clause.Locking{Strength: "UPDATE"}).Where("dialog_id IN ?", dialogIDs).Order("id").Pluck("id", &ids).Error; err != nil { return err } - if err := deleteHistory(ctx, tx, []string{apiConversationMessageTable, apiConversationReferenceTable}, ids); err != nil { + if err := deleteHistory(ctx, tx, []string{apiConversationMessageTable}, ids); err != nil { return err } result := tx.Where("dialog_id IN ?", dialogIDs).Delete(&entity.API4Conversation{}) diff --git a/internal/dao/api_token_test.go b/internal/dao/api_token_test.go index 79ab9c5738..ea2a1a074d 100644 --- a/internal/dao/api_token_test.go +++ b/internal/dao/api_token_test.go @@ -48,11 +48,10 @@ func createAPI4ConversationForDAOTest(t *testing.T, id, agentID string) { t.Helper() ctx := t.Context() if err := NewAPI4ConversationDAO().Create(ctx, DB, &entity.API4Conversation{ - ID: id, - DialogID: agentID, - UserID: "user-1", - Message: json.RawMessage(`[]`), - Reference: json.RawMessage(`[]`), + ID: id, + DialogID: agentID, + UserID: "user-1", + Message: json.RawMessage(`[]`), }); err != nil { t.Fatalf("failed to create api conversation %s: %v", id, err) } diff --git a/internal/dao/chat_session.go b/internal/dao/chat_session.go index db5d5c4f44..210b299771 100644 --- a/internal/dao/chat_session.go +++ b/internal/dao/chat_session.go @@ -33,20 +33,21 @@ import ( type ChatSessionDAO struct{} type ListAgentSessionsParams struct { - AgentID string - Page int - PageSize int - OrderBy string - Desc bool - SessionID string - UserID string - TenantID string - IncludeDSL bool - NoHistory bool - Keywords string - FromDate *time.Time - ToDate *time.Time - ExpUserID string + AgentID string + Page int + PageSize int + OrderBy string + Desc bool + SessionID string + UserID string + TenantID string + IncludeDSL bool + NoHistory bool + HistoryOptions ConversationHistoryOptions + Keywords string + FromDate *time.Time + ToDate *time.Time + ExpUserID string } // NewChatSessionDAO create chat session DAO @@ -55,26 +56,34 @@ func NewChatSessionDAO() *ChatSessionDAO { } // GetByID gets chat session by ID -func (dao *ChatSessionDAO) GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.ChatSession, error) { +func (dao *ChatSessionDAO) GetByID(ctx context.Context, db *gorm.DB, id string, options ...ConversationHistoryOptions) (*entity.ChatSession, error) { + historyOptions := ConversationHistoryOptions{} + if len(options) > 0 { + historyOptions = options[0] + } var conv entity.ChatSession err := db.WithContext(ctx).Session(&gorm.Session{QueryFields: true}).Take(&conv, "id = ?", id).Error if err != nil { return nil, err } - if err = hydrateChatSessions(ctx, db, []*entity.ChatSession{&conv}); err != nil { + if err = hydrateChatSessions(ctx, db, []*entity.ChatSession{&conv}, historyOptions); err != nil { return nil, err } return &conv, nil } // GetBySessionIDAndChatID gets a chat session by session ID and chat ID. -func (dao *ChatSessionDAO) GetBySessionIDAndChatID(ctx context.Context, db *gorm.DB, sessionID, chatID string) (*entity.ChatSession, error) { +func (dao *ChatSessionDAO) GetBySessionIDAndChatID(ctx context.Context, db *gorm.DB, sessionID, chatID string, options ...ConversationHistoryOptions) (*entity.ChatSession, error) { + historyOptions := ConversationHistoryOptions{} + if len(options) > 0 { + historyOptions = options[0] + } var conv entity.ChatSession err := db.WithContext(ctx).Session(&gorm.Session{QueryFields: true}).Where("id = ? AND dialog_id = ?", sessionID, chatID).First(&conv).Error if err != nil { return nil, err } - if err = hydrateChatSessions(ctx, db, []*entity.ChatSession{&conv}); err != nil { + if err = hydrateChatSessions(ctx, db, []*entity.ChatSession{&conv}, historyOptions); err != nil { return nil, err } return &conv, nil @@ -86,10 +95,7 @@ func (dao *ChatSessionDAO) Create(ctx context.Context, db *gorm.DB, conv *entity if err := tx.Create(conv).Error; err != nil { return err } - if err := createHistory(ctx, tx, conversationMessageTable, "message", conv.ID, conv.Message); err != nil { - return err - } - return createHistory(ctx, tx, conversationReferenceTable, "reference", conv.ID, conv.Reference) + return createHistory(ctx, tx, conversationMessageTable, "message", conv.ID, conv.Message) }) } @@ -127,7 +133,7 @@ func (dao *ChatSessionDAO) UpdateByID(ctx context.Context, db *gorm.DB, id strin } } if targeted { - return updateConversationHistory(ctx, tx, conversationMessageTable, conversationReferenceTable, id, historyUpdate) + return updateConversationHistory(ctx, tx, conversationMessageTable, id, historyUpdate) } return nil }) @@ -140,7 +146,7 @@ func (dao *ChatSessionDAO) DeleteByID(ctx context.Context, db *gorm.DB, id strin if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id").Where("id = ?", id).Take(&parent).Error; err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { return err } - if err := deleteHistory(ctx, tx, []string{conversationMessageTable, conversationReferenceTable}, []string{id}); err != nil { + if err := deleteHistory(ctx, tx, []string{conversationMessageTable}, []string{id}); err != nil { return err } return tx.Where("id = ?", id).Delete(&entity.ChatSession{}).Error @@ -148,7 +154,7 @@ func (dao *ChatSessionDAO) DeleteByID(ctx context.Context, db *gorm.DB, id strin } // ListByResourceIDs lists sessions for a chat, optionally restricted to sessionIDs. -func (dao *ChatSessionDAO) ListByResourceIDs(ctx context.Context, db *gorm.DB, chatID string, sessionIDs []string, sessionID, name string, terms []OrderTerm, page, pageSize int, includeHistory ...bool) ([]*entity.ChatSession, int64, error) { +func (dao *ChatSessionDAO) ListByResourceIDs(ctx context.Context, db *gorm.DB, chatID string, sessionIDs []string, sessionID, name string, terms []OrderTerm, page, pageSize int, options ...ConversationHistoryOptions) ([]*entity.ChatSession, int64, error) { if sessionIDs != nil && len(sessionIDs) == 0 { return []*entity.ChatSession{}, 0, nil } @@ -181,8 +187,12 @@ func (dao *ChatSessionDAO) ListByResourceIDs(ctx context.Context, db *gorm.DB, c if err != nil { return nil, 0, err } - if len(includeHistory) == 0 || includeHistory[0] { - if err = hydrateChatSessions(ctx, db, chatSessions); err != nil { + historyOptions := ConversationHistoryOptions{} + if len(options) > 0 { + historyOptions = options[0] + } + if !historyOptions.NoHistory { + if err = hydrateChatSessions(ctx, db, chatSessions, historyOptions); err != nil { return nil, 0, err } } @@ -210,7 +220,7 @@ func (dao *ChatSessionDAO) DeleteByDialogIDs(ctx context.Context, db *gorm.DB, d if err := tx.Model(&entity.ChatSession{}).Clauses(clause.Locking{Strength: "UPDATE"}).Where("dialog_id IN ?", dialogIDs).Order("id").Pluck("id", &ids).Error; err != nil { return err } - if err := deleteHistory(ctx, tx, []string{conversationMessageTable, conversationReferenceTable}, ids); err != nil { + if err := deleteHistory(ctx, tx, []string{conversationMessageTable}, ids); err != nil { return err } result := tx.Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.ChatSession{}) @@ -341,7 +351,7 @@ func (dao *ChatSessionDAO) ListAgentSessions(ctx context.Context, db *gorm.DB, p return 0, nil, err } if !params.NoHistory { - if err = hydrateAPIConversations(ctx, db, sessions); err != nil { + if err = hydrateAPIConversations(ctx, db, sessions, params.HistoryOptions); err != nil { return 0, nil, err } } diff --git a/internal/dao/chat_session_test.go b/internal/dao/chat_session_test.go index e3ea6c56b7..d9a6c9c89a 100644 --- a/internal/dao/chat_session_test.go +++ b/internal/dao/chat_session_test.go @@ -50,11 +50,10 @@ func createAgentSessionForDAOTest(t *testing.T, db *gorm.DB, id, agentID, userID updateDate := time.UnixMilli(updateTime).Local() session := &entity.API4Conversation{ - ID: id, - DialogID: agentID, - UserID: userID, - Message: json.RawMessage(`[{"role":"assistant","content":"hello"}]`), - Reference: json.RawMessage(`[]`), + ID: id, + DialogID: agentID, + UserID: userID, + Message: json.RawMessage(`[{"role":"assistant","content":"hello"}]`), BaseModel: entity.BaseModel{ CreateTime: &updateTime, CreateDate: &updateDate, @@ -72,12 +71,11 @@ func createNamedAgentSessionForDAOTest(t *testing.T, db *gorm.DB, id, agentID, u updateDate := time.UnixMilli(updateTime).Local() session := &entity.API4Conversation{ - ID: id, - Name: &name, - DialogID: agentID, - UserID: userID, - Message: message, - Reference: json.RawMessage(`[]`), + ID: id, + Name: &name, + DialogID: agentID, + UserID: userID, + Message: message, BaseModel: entity.BaseModel{ CreateTime: &updateTime, CreateDate: &updateDate, @@ -95,11 +93,10 @@ func createChatSessionForDAOTest(t *testing.T, db *gorm.DB, id, chatID, name str updateDate := time.UnixMilli(updateTime).Local() session := &entity.ChatSession{ - ID: id, - DialogID: chatID, - Name: &name, - Message: json.RawMessage(`[{"role":"assistant","content":"hello"}]`), - Reference: json.RawMessage(`[]`), + ID: id, + DialogID: chatID, + Name: &name, + Message: json.RawMessage(`[{"role":"assistant","content":"hello"}]`), BaseModel: entity.BaseModel{ CreateTime: &updateTime, CreateDate: &updateDate, diff --git a/internal/dao/conversation_history.go b/internal/dao/conversation_history.go index b27a673f79..d300d8a73c 100644 --- a/internal/dao/conversation_history.go +++ b/internal/dao/conversation_history.go @@ -31,19 +31,21 @@ import ( ) const ( - conversationMessageTable = "conversation_message" - conversationReferenceTable = "conversation_reference" - apiConversationMessageTable = "api_4_conversation_message" - apiConversationReferenceTable = "api_4_conversation_reference" - conversationHistoryIDColumn = "conversation_id" - conversationHistoryOrderColumn = "position" + conversationMessageTable = "conversation_message" + conversationReferenceTable = "conversation_reference" + apiConversationMessageTable = "api_4_conversation_message" + apiConversationReferenceTable = "api_4_conversation_reference" + conversationHistoryIDColumn = "conversation_id" ) -type conversationHistoryRow struct { - ConversationID string `gorm:"column:conversation_id"` - Position int `gorm:"column:position"` - Payload string `gorm:"column:payload"` - MessagePos *int `gorm:"column:message_position"` +// ConversationHistoryOptions selects message payloads needed by the caller. +// The zero value loads history without answer details. +type ConversationHistoryOptions struct { + NoHistory bool + IncludeReference bool + IncludePrompt bool + IncludeUsage bool + IncludeElapsedMs bool } // ConversationHistoryUpdate changes only the addressed turn, never synchronizing @@ -51,39 +53,12 @@ type conversationHistoryRow struct { type ConversationHistoryUpdate struct { Message map[string]interface{} QuestionID string - Reference map[string]interface{} - AppendReference bool DeleteMessageID string FeedbackMessageID string Feedback map[string]interface{} } -func historyReferenceQuery(ctx context.Context, db *gorm.DB, messageTable, referenceTable, conversationID string, messagePosition int) (*gorm.DB, error) { - query := db.WithContext(ctx).Table(referenceTable).Where("conversation_id = ?", conversationID) - var row conversationHistoryRow - err := query.Session(&gorm.Session{}).Select("position").Where("position = ? AND message_position = ?", messagePosition, messagePosition).Take(&row).Error - if err == nil { - return query.Where("position = ?", row.Position), nil - } - if !errors.Is(err, gorm.ErrRecordNotFound) { - return nil, err - } - // Histories explicitly supplied as arrays have no position association. - // Resolve their remaining references by assistant order, not array pairs. - var ordinal int64 - err = db.WithContext(ctx).Table(messageTable). - Where("conversation_id = ? AND role = ? AND position < ? AND position >= (SELECT MIN(position) FROM "+messageTable+" WHERE conversation_id = ? AND role = ?)", conversationID, "assistant", messagePosition, conversationID, "user").Count(&ordinal).Error - if err != nil { - return nil, err - } - err = query.Session(&gorm.Session{}).Select("position").Where("message_position IS NULL").Order("position").Offset(int(ordinal)).Take(&row).Error - if err != nil { - return nil, err - } - return query.Where("position = ?", row.Position), nil -} - -func updateConversationHistory(ctx context.Context, db *gorm.DB, messageTable, referenceTable, conversationID string, update ConversationHistoryUpdate) error { +func updateConversationHistory(ctx context.Context, db *gorm.DB, messageTable, conversationID string, update ConversationHistoryUpdate) error { if update.DeleteMessageID != "" { var question entity.ConversationMessage query := db.WithContext(ctx).Table(messageTable).Where("conversation_id = ? AND message_id = ?", conversationID, update.DeleteMessageID) @@ -98,15 +73,6 @@ func updateConversationHistory(ctx context.Context, db *gorm.DB, messageTable, r } if err == nil { positions = append(positions, answer.Position) - refQuery, err := historyReferenceQuery(ctx, db, messageTable, referenceTable, conversationID, answer.Position) - if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { - return err - } - if err == nil { - if err := refQuery.Delete(map[string]interface{}{}).Error; err != nil { - return err - } - } } return query.Where("position IN ?", positions).Delete(map[string]interface{}{}).Error } @@ -147,7 +113,7 @@ func updateConversationHistory(ctx context.Context, db *gorm.DB, messageTable, r update.Feedback["metadata"] = string(raw) return query.Where("position = ?", message.Position).Updates(update.Feedback).Error } - if update.Message == nil && !update.AppendReference { + if update.Message == nil { return nil } var fields entity.ConversationMessageFields @@ -169,14 +135,11 @@ func updateConversationHistory(ctx context.Context, db *gorm.DB, messageTable, r } position = question.Position + 1 } else { - var lastMessage, lastReference int + var lastMessage int if err := db.WithContext(ctx).Table(messageTable).Select("COALESCE(MAX(position), -2)").Where("conversation_id = ?", conversationID).Scan(&lastMessage).Error; err != nil { return err } - if err := db.WithContext(ctx).Table(referenceTable).Select("COALESCE(MAX(position), -2)").Where("conversation_id = ?", conversationID).Scan(&lastReference).Error; err != nil { - return err - } - position = max(lastMessage, lastReference) + 2 + position = lastMessage + 2 } if update.Message != nil { values := messageValues(fields) @@ -186,17 +149,6 @@ func updateConversationHistory(ctx context.Context, db *gorm.DB, messageTable, r return err } } - if update.AppendReference { - raw, err := json.Marshal(update.Reference) - if err != nil { - return err - } - row := map[string]interface{}{"conversation_id": conversationID, "position": position, "reference": string(raw)} - if update.Message != nil { - row["message_position"] = position - } - return db.WithContext(ctx).Table(referenceTable).Create(row).Error - } return nil } @@ -266,7 +218,7 @@ func flattenMessage(item json.RawMessage) (entity.ConversationMessageFields, err for key, target := range map[string]interface{}{ "id": &fields.MessageID, "role": &fields.Role, "status": &fields.Status, "thumbup": &fields.ThumbUp, "feedback": &fields.Feedback, "created_at": &fields.CreatedAt, - "prompt": &fields.Prompt, "usage": &fields.Usage, "elapsed_ms": &fields.ElapsedMs, + "prompt": &fields.Prompt, "usage": &fields.Usage, "elapsed_ms": &fields.ElapsedMs, "reference": &fields.Reference, } { if value, ok := metadata[key]; ok && !bytes.Equal(bytes.TrimSpace(value), []byte("null")) { if err := json.Unmarshal(value, target); err == nil { @@ -294,7 +246,7 @@ func flattenMessage(item json.RawMessage) (entity.ConversationMessageFields, err return fields, err } -func marshalMessage(fields entity.ConversationMessageFields) (json.RawMessage, error) { +func marshalMessage(fields entity.ConversationMessageFields, options ConversationHistoryOptions) (json.RawMessage, error) { metadata := make(map[string]json.RawMessage) if len(fields.Metadata) > 0 { if err := json.Unmarshal([]byte(fields.Metadata), &metadata); err != nil { @@ -318,14 +270,26 @@ func marshalMessage(fields entity.ConversationMessageFields) (json.RawMessage, e if fields.Content != nil && fields.ContentType == "json" { metadata["content"] = json.RawMessage(*fields.Content) } + for key, include := range map[string]bool{ + "reference": options.IncludeReference, "prompt": options.IncludePrompt, + "usage": options.IncludeUsage, "elapsed_ms": options.IncludeElapsedMs, + } { + if !include { + delete(metadata, key) + } + } return json.Marshal(metadata) } func messageValues(fields entity.ConversationMessageFields) map[string]interface{} { + var reference interface{} + if len(fields.Reference) > 0 { + reference = string(fields.Reference) + } return map[string]interface{}{ "message_id": fields.MessageID, "role": fields.Role, "content": fields.Content, "content_type": fields.ContentType, "status": fields.Status, "thumb_up": fields.ThumbUp, "feedback": fields.Feedback, "created_at": fields.CreatedAt, "metadata": fields.Metadata, - "prompt": fields.Prompt, "usage": fields.Usage, "elapsed_ms": fields.ElapsedMs, + "prompt": fields.Prompt, "usage": fields.Usage, "elapsed_ms": fields.ElapsedMs, "reference": reference, } } @@ -363,84 +327,78 @@ func historyRows(kind, conversationID string, raw json.RawMessage) ([]map[string return rows, nil } -func loadHistory(ctx context.Context, db *gorm.DB, table, payloadColumn string, conversationIDs []string) (map[string]json.RawMessage, error) { +// loadHistory excludes answer details unless the caller needs them for display. +func loadHistory(ctx context.Context, db *gorm.DB, table string, conversationIDs []string, options ConversationHistoryOptions) (map[string]json.RawMessage, error) { result := make(map[string]json.RawMessage, len(conversationIDs)) if len(conversationIDs) == 0 { return result, nil } - itemsByConversation := make(map[string][]json.RawMessage, len(conversationIDs)) - query := db.WithContext(ctx).Table(table).Where(conversationHistoryIDColumn+" IN ?", conversationIDs).Order(conversationHistoryIDColumn + ", " + conversationHistoryOrderColumn) - if payloadColumn == "message" { - var rows []entity.ConversationMessage - if err := query.Find(&rows).Error; err != nil { - return nil, err - } - for _, row := range rows { - raw, err := marshalMessage(row.ConversationMessageFields) - if err != nil { - return nil, fmt.Errorf("marshal message for conversation %s: %w", row.ConversationID, err) - } - itemsByConversation[row.ConversationID] = append(itemsByConversation[row.ConversationID], raw) - } - } else { - var rows []conversationHistoryRow - if err := query.Select(conversationHistoryIDColumn + ", " + conversationHistoryOrderColumn + ", " + payloadColumn + " AS payload").Find(&rows).Error; err != nil { - return nil, err - } - for _, row := range rows { - itemsByConversation[row.ConversationID] = append(itemsByConversation[row.ConversationID], json.RawMessage(row.Payload)) + query := db.WithContext(ctx).Session(&gorm.Session{QueryFields: true}).Table(table).Where("conversation_id IN ?", conversationIDs).Order("conversation_id, position") + var omitted []string + for column, include := range map[string]bool{ + "reference": options.IncludeReference, "prompt": options.IncludePrompt, + "usage": options.IncludeUsage, "elapsed_ms": options.IncludeElapsedMs, + } { + if !include { + omitted = append(omitted, column) } } - for _, conversationID := range conversationIDs { - items := itemsByConversation[conversationID] + if len(omitted) > 0 { + query = query.Omit(omitted...) + } + var rows []entity.ConversationMessage + if err := query.Find(&rows).Error; err != nil { + return nil, err + } + itemsByConversation := make(map[string][]json.RawMessage, len(conversationIDs)) + for _, row := range rows { + fields := row.ConversationMessageFields + raw, err := marshalMessage(fields, options) + if err != nil { + return nil, fmt.Errorf("marshal message for conversation %s: %w", row.ConversationID, err) + } + itemsByConversation[row.ConversationID] = append(itemsByConversation[row.ConversationID], raw) + } + for _, id := range conversationIDs { + items := itemsByConversation[id] if items == nil { items = []json.RawMessage{} } raw, err := json.Marshal(items) if err != nil { - return nil, fmt.Errorf("marshal history for conversation %s: %w", conversationID, err) + return nil, fmt.Errorf("marshal history for conversation %s: %w", id, err) } - result[conversationID] = raw + result[id] = raw } return result, nil } -func hydrateChatSessions(ctx context.Context, db *gorm.DB, sessions []*entity.ChatSession) error { +func hydrateChatSessions(ctx context.Context, db *gorm.DB, sessions []*entity.ChatSession, options ConversationHistoryOptions) error { ids := make([]string, 0, len(sessions)) for _, session := range sessions { ids = append(ids, session.ID) } - messages, err := loadHistory(ctx, db, conversationMessageTable, "message", ids) - if err != nil { - return err - } - references, err := loadHistory(ctx, db, conversationReferenceTable, "reference", ids) + messages, err := loadHistory(ctx, db, conversationMessageTable, ids, options) if err != nil { return err } for _, session := range sessions { session.Message = messages[session.ID] - session.Reference = references[session.ID] } return nil } -func hydrateAPIConversations(ctx context.Context, db *gorm.DB, sessions []*entity.API4Conversation) error { +func hydrateAPIConversations(ctx context.Context, db *gorm.DB, sessions []*entity.API4Conversation, options ConversationHistoryOptions) error { ids := make([]string, 0, len(sessions)) for _, session := range sessions { ids = append(ids, session.ID) } - messages, err := loadHistory(ctx, db, apiConversationMessageTable, "message", ids) - if err != nil { - return err - } - references, err := loadHistory(ctx, db, apiConversationReferenceTable, "reference", ids) + messages, err := loadHistory(ctx, db, apiConversationMessageTable, ids, options) if err != nil { return err } for _, session := range sessions { session.Message = messages[session.ID] - session.Reference = references[session.ID] } return nil } diff --git a/internal/dao/conversation_history_benchmark_test.go b/internal/dao/conversation_history_benchmark_test.go index 5aedf12cd0..9733a41469 100644 --- a/internal/dao/conversation_history_benchmark_test.go +++ b/internal/dao/conversation_history_benchmark_test.go @@ -68,21 +68,15 @@ func BenchmarkConversationHistoryOperations(b *testing.B) { {ConversationID: conversation.ID, Position: 0, ConversationMessageFields: entity.ConversationMessageFields{Role: &assistantRole, Content: &prologue, ContentType: "text", Metadata: "{}"}}, {ConversationID: conversation.ID, Position: 2, ConversationMessageFields: entity.ConversationMessageFields{Role: &userRole, MessageID: &pendingID, Content: &content, ContentType: "text", Metadata: "{}"}}, } - var references []entity.API4ConversationReference for i := 0; i < (size-2)/2; i++ { id, position := fmt.Sprintf("turn-%d", i), 4+i*4 messages = append(messages, entity.API4ConversationMessage{ConversationID: conversation.ID, Position: position, ConversationMessageFields: entity.ConversationMessageFields{Role: &userRole, MessageID: &id, Content: &content, ContentType: "text", Metadata: "{}"}}, - entity.API4ConversationMessage{ConversationID: conversation.ID, Position: position + 1, ConversationMessageFields: entity.ConversationMessageFields{Role: &assistantRole, MessageID: &id, Content: &content, ContentType: "text", Metadata: "{}"}}) - answerPosition := position + 1 - references = append(references, entity.API4ConversationReference{ConversationID: conversation.ID, Position: answerPosition, MessagePos: &answerPosition, Reference: []byte(`{"chunks":[]}`)}) + entity.API4ConversationMessage{ConversationID: conversation.ID, Position: position + 1, ConversationMessageFields: entity.ConversationMessageFields{Role: &assistantRole, MessageID: &id, Content: &content, ContentType: "text", Metadata: "{}", Reference: []byte(`{"chunks":[]}`)}}) } if err := db.CreateInBatches(messages, 100).Error; err != nil { b.Fatal(err) } - if err := db.CreateInBatches(references, 100).Error; err != nil { - b.Fatal(err) - } if operation == "search_miss" { for i := 1; i < 20; i++ { id := fmt.Sprintf("history-benchmark-%d", i) @@ -100,16 +94,13 @@ func BenchmarkConversationHistoryOperations(b *testing.B) { store := NewAPI4ConversationDAO() update := ConversationHistoryUpdate{Message: map[string]interface{}{"role": "user", "id": "new-question", "content": content, "created_at": 1.0}} if answer { - update = ConversationHistoryUpdate{Message: map[string]interface{}{"role": "assistant", "id": pendingID, "content": content, "created_at": 2.0}, QuestionID: pendingID, Reference: map[string]interface{}{"chunks": []interface{}{}}, AppendReference: true} + update = ConversationHistoryUpdate{Message: map[string]interface{}{"role": "assistant", "id": pendingID, "content": content, "created_at": 2.0, "reference": map[string]interface{}{"chunks": []interface{}{}}}, QuestionID: pendingID} } b.ResetTimer() for i := 0; i < b.N; i++ { if answer && i > 0 { b.StopTimer() stats.enabled = false - if err := db.Where("conversation_id = ? AND position = ?", conversation.ID, 3).Delete(&entity.API4ConversationReference{}).Error; err != nil { - b.Fatal(err) - } if err := db.Where("conversation_id = ? AND position = ?", conversation.ID, 3).Delete(&entity.API4ConversationMessage{}).Error; err != nil { b.Fatal(err) } @@ -146,7 +137,7 @@ func BenchmarkConversationHistoryOperations(b *testing.B) { if err := db.Model(&entity.API4ConversationMessage{}).Where("conversation_id = ?", conversation.ID).Count(&messageCount).Error; err != nil { b.Fatal(err) } - if err := db.Model(&entity.API4ConversationReference{}).Where("conversation_id = ?", conversation.ID).Count(&referenceCount).Error; err != nil { + if err := db.Model(&entity.API4ConversationMessage{}).Where("conversation_id = ? AND reference IS NOT NULL", conversation.ID).Count(&referenceCount).Error; err != nil { b.Fatal(err) } if messageCount != int64(remaining) || referenceCount != int64((size-2)/2) { diff --git a/internal/dao/conversation_history_migration.go b/internal/dao/conversation_history_migration.go index 966bb57162..417d040196 100644 --- a/internal/dao/conversation_history_migration.go +++ b/internal/dao/conversation_history_migration.go @@ -25,6 +25,7 @@ import ( "gorm.io/gorm" "ragflow/internal/common" + "ragflow/internal/entity" ) // conversationHistoryTargetVersion is the database version a completed split @@ -102,6 +103,9 @@ func migrateConversationHistory(ctx context.Context, db *gorm.DB) error { zap.String("target_version", conversationHistoryTargetVersion)) return nil } + if err := db.WithContext(ctx).AutoMigrate(&entity.ConversationReference{}, &entity.API4ConversationReference{}); err != nil { + return fmt.Errorf("prepare reference tables for history split: %w", err) + } splits := []conversationHistorySplit{ { parentTable: "conversation", diff --git a/internal/dao/conversation_history_migration_test.go b/internal/dao/conversation_history_migration_test.go index a055e84e4a..c237e1d4a7 100644 --- a/internal/dao/conversation_history_migration_test.go +++ b/internal/dao/conversation_history_migration_test.go @@ -116,7 +116,7 @@ func TestConversationMessageDetailsSchemaMigration(t *testing.T) { if _, ok := model.(*entity.API4ConversationMessage); ok { table, parent = apiConversationMessageTable, "api_4_conversation" } - for _, column := range []string{"prompt", "usage", "elapsed_ms"} { + for _, column := range []string{"prompt", "usage", "elapsed_ms", "reference"} { if err := db.Migrator().DropColumn(model, column); err != nil { t.Fatal(err) } @@ -132,7 +132,7 @@ func TestConversationMessageDetailsSchemaMigration(t *testing.T) { t.Fatal(err) } } - for _, column := range []string{"prompt", "usage", "elapsed_ms"} { + for _, column := range []string{"prompt", "usage", "elapsed_ms", "reference"} { if !db.Migrator().HasColumn(model, column) { t.Fatalf("missing column %s", column) } @@ -141,7 +141,7 @@ func TestConversationMessageDetailsSchemaMigration(t *testing.T) { if err := db.Table(table).Where("conversation_id = ?", "session-1").Take(&row).Error; err != nil { t.Fatal(err) } - if row.Content == nil || *row.Content != "old answer" || row.Prompt != nil || len(row.Usage) != 0 || row.ElapsedMs != nil { + if row.Content == nil || *row.Content != "old answer" || row.Prompt != nil || len(row.Usage) != 0 || row.ElapsedMs != nil || len(row.Reference) != 0 { t.Fatalf("migration changed old answer: %+v", row) } }) diff --git a/internal/dao/conversation_history_test.go b/internal/dao/conversation_history_test.go index f87363be37..1dfb135a36 100644 --- a/internal/dao/conversation_history_test.go +++ b/internal/dao/conversation_history_test.go @@ -13,6 +13,8 @@ import ( func TestFlattenMessageRoundTrip(t *testing.T) { for _, item := range []string{ + `{"role":"assistant","content":"answer","reference":{"chunks":[{"id":"chunk-1"}],"doc_aggs":[]}}`, + `{"role":"assistant","reference":{}}`, `{"role":"assistant","content":"agentic answer","usage":{"prompt_tokens":896,"completion_tokens":71,"total_tokens":967},"elapsed_ms":4168.774506}`, `{"role":"assistant","content":"answer","prompt":"system prompt","usage":{"prompt_tokens":896,"completion_tokens":71,"total_tokens":967},"elapsed_ms":4168.774506}`, `{"role":"assistant","prompt":"","usage":null,"elapsed_ms":0}`, @@ -27,7 +29,7 @@ func TestFlattenMessageRoundTrip(t *testing.T) { if err != nil { t.Fatal(err) } - raw, err := marshalMessage(fields) + raw, err := marshalMessage(fields, ConversationHistoryOptions{IncludeReference: true, IncludePrompt: true, IncludeUsage: true, IncludeElapsedMs: true}) if err != nil { t.Fatal(err) } @@ -92,7 +94,7 @@ func TestFlattenedMessagePersistence(t *testing.T) { if agenticRow.Prompt != nil || !reflect.DeepEqual(agenticRow.Usage, row.Usage) || agenticRow.ElapsedMs == nil || *agenticRow.ElapsedMs != *row.ElapsedMs || agenticRow.Metadata != "{}" { t.Fatalf("unexpected agentic answer details: %+v", agenticRow.ConversationMessageFields) } - histories, err := loadHistory(t.Context(), db, table, "message", []string{"session-1"}) + histories, err := loadHistory(t.Context(), db, table, []string{"session-1"}, ConversationHistoryOptions{IncludeReference: true, IncludePrompt: true, IncludeUsage: true, IncludeElapsedMs: true}) if err != nil { t.Fatal(err) } diff --git a/internal/dao/conversation_reference_migration.go b/internal/dao/conversation_reference_migration.go new file mode 100644 index 0000000000..cb502a19ba --- /dev/null +++ b/internal/dao/conversation_reference_migration.go @@ -0,0 +1,190 @@ +package dao + +import ( + "context" + "encoding/json" + "errors" + "fmt" + + "go.uber.org/zap" + "gorm.io/gorm" + + "ragflow/internal/common" +) + +const conversationReferenceTargetVersion = "v1.0.0-rc3.dev7" + +// migrateConversationReferences runs after the history split and all earlier +// versioned migrations. Source tables are dropped only after both copies succeed. +func migrateConversationReferences(ctx context.Context, db *gorm.DB) error { + version, err := GetDatabaseMigrationVersion(ctx, db) + if err != nil { + return err + } + if shouldSkipMigration(version, conversationReferenceTargetVersion) && + !db.WithContext(ctx).Migrator().HasTable(conversationReferenceTable) && + !db.WithContext(ctx).Migrator().HasTable(apiConversationReferenceTable) { + return nil + } + for _, tables := range [][2]string{ + {conversationMessageTable, conversationReferenceTable}, + {apiConversationMessageTable, apiConversationReferenceTable}, + } { + if err := migrateMessageReferences(ctx, db, tables[0], tables[1]); err != nil { + return fmt.Errorf("migrate %s references: %w", tables[0], err) + } + if err := migrateMessageReferenceMetadata(ctx, db, tables[0]); err != nil { + return fmt.Errorf("migrate %s reference metadata: %w", tables[0], err) + } + } + for _, table := range []string{conversationReferenceTable, apiConversationReferenceTable} { + if db.WithContext(ctx).Migrator().HasTable(table) { + if err := db.WithContext(ctx).Migrator().DropTable(table); err != nil { + return fmt.Errorf("drop migrated table %s: %w", table, err) + } + } + } + if shouldSkipMigration(version, conversationReferenceTargetVersion) { + return nil + } + return setDatabaseMigrationVersion(ctx, db, conversationReferenceTargetVersion) +} + +func migrateMessageReferences(ctx context.Context, db *gorm.DB, messageTable, referenceTable string) error { + db = db.WithContext(ctx) + if !db.Migrator().HasTable(referenceTable) { + return nil + } + type referenceRow struct { + ConversationID string + Position int + MessagePos *int `gorm:"column:message_position"` + Reference string + } + lastConversation, activeConversation := "", "" + lastPosition, lastAssistant := -1, -1 + unmatched := 0 + for { + var rows []referenceRow + if err := db.Table(referenceTable). + Select("conversation_id", "position", "message_position", "reference"). + Where("conversation_id > ? OR (conversation_id = ? AND position > ?)", lastConversation, lastConversation, lastPosition). + Order("conversation_id, position").Limit(conversationHistoryBatchSize).Find(&rows).Error; err != nil { + return err + } + if len(rows) == 0 { + break + } + if err := db.Transaction(func(tx *gorm.DB) error { + for _, row := range rows { + if row.ConversationID != activeConversation { + activeConversation, lastAssistant = row.ConversationID, -1 + } + if !json.Valid([]byte(row.Reference)) { + return fmt.Errorf("invalid reference JSON for conversation %s position %d", row.ConversationID, row.Position) + } + position := row.MessagePos + if position == nil { + // Array references follow assistant order after the first user. + // Explicit associations belong to later targeted writes and must + // not consume an entry from the old reference array. + var answer struct{ Position int } + err := tx.Table(messageTable).Select("position"). + Where("conversation_id = ? AND role = ? AND position > ?", row.ConversationID, "assistant", lastAssistant). + Where("position >= (SELECT MIN(position) FROM "+messageTable+" WHERE conversation_id = ? AND role = ?)", row.ConversationID, "user"). + Where("NOT EXISTS (SELECT 1 FROM "+referenceTable+" r WHERE r.conversation_id = ? AND r.message_position = "+messageTable+".position)", row.ConversationID). + Order("position").Take(&answer).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + unmatched++ + continue + } + if err != nil { + return err + } + lastAssistant = answer.Position + position = &answer.Position + } + query := tx.Table(messageTable).Where("conversation_id = ? AND position = ? AND role = ?", row.ConversationID, *position, "assistant") + result := query.Session(&gorm.Session{}).Where("reference IS NULL OR reference = ? OR reference = ? OR TRIM(reference) LIKE ?", "", "null", "[%").Update("reference", row.Reference) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + var count int64 + if err := query.Count(&count).Error; err != nil { + return err + } + if count == 0 { + unmatched++ + } + } + } + return nil + }); err != nil { + return err + } + last := rows[len(rows)-1] + lastConversation, lastPosition = last.ConversationID, last.Position + } + if unmatched > 0 { + // References can outlive deleted messages or unsuccessful turns. They + // cannot be attached to another answer without changing its citations. + common.Warn("Skipped references without assistant messages during migration", zap.String("table", referenceTable), zap.Int("references", unmatched)) + } + return nil +} + +// Older message rows can carry a second copy in metadata. Move that value when +// there is no reference column value, then remove it so lean SQL reads stay lean. +func migrateMessageReferenceMetadata(ctx context.Context, db *gorm.DB, table string) error { + db = db.WithContext(ctx) + if !db.Migrator().HasTable(table) { + return nil + } + lastConversation, lastPosition := "", -1 + for { + var rows []struct { + ConversationID string + Position int + Metadata string + } + if err := db.Table(table).Select("conversation_id", "position", "metadata"). + Where("metadata LIKE ?", `%"reference"%`). + Where("conversation_id > ? OR (conversation_id = ? AND position > ?)", lastConversation, lastConversation, lastPosition). + Order("conversation_id, position").Limit(conversationHistoryBatchSize).Find(&rows).Error; err != nil { + return err + } + if len(rows) == 0 { + return nil + } + if err := db.Transaction(func(tx *gorm.DB) error { + for _, row := range rows { + var metadata map[string]json.RawMessage + if err := json.Unmarshal([]byte(row.Metadata), &metadata); err != nil { + return fmt.Errorf("decode metadata for conversation %s position %d: %w", row.ConversationID, row.Position, err) + } + reference, ok := metadata["reference"] + if !ok { + continue + } + delete(metadata, "reference") + raw, err := json.Marshal(metadata) + if err != nil { + return err + } + query := tx.Table(table).Where("conversation_id = ? AND position = ?", row.ConversationID, row.Position) + if err := query.Session(&gorm.Session{}).Where("reference IS NULL OR reference = ? OR reference = ?", "", "null").Update("reference", string(reference)).Error; err != nil { + return err + } + if err := query.Update("metadata", string(raw)).Error; err != nil { + return err + } + } + return nil + }); err != nil { + return err + } + last := rows[len(rows)-1] + lastConversation, lastPosition = last.ConversationID, last.Position + } +} diff --git a/internal/dao/conversation_reference_migration_test.go b/internal/dao/conversation_reference_migration_test.go new file mode 100644 index 0000000000..bfc2ceb8bd --- /dev/null +++ b/internal/dao/conversation_reference_migration_test.go @@ -0,0 +1,315 @@ +package dao + +import ( + "encoding/json" + "fmt" + "strings" + "testing" + + "gorm.io/gorm" + + "ragflow/internal/entity" +) + +func TestMigrateConversationReferences(t *testing.T) { + for _, tables := range [][3]string{ + {"conversation", conversationMessageTable, conversationReferenceTable}, + {"api_4_conversation", apiConversationMessageTable, apiConversationReferenceTable}, + } { + t.Run(tables[0], func(t *testing.T) { + db := setupConversationHistoryMigrationDB(t, false) + insertConversation(t, db, tables[0], "session") + messages := []map[string]interface{}{ + {"role": "system", "content": "system"}, + {"role": "assistant", "content": "prologue"}, + {"role": "user", "id": "unanswered"}, + } + // Cross a reference batch boundary within the same conversation. + for i := 0; i < conversationHistoryBatchSize+3; i++ { + messages = append(messages, map[string]interface{}{"role": "user"}, map[string]interface{}{"role": "assistant"}) + } + messages = append(messages, + map[string]interface{}{"role": "user", "id": "delayed"}, + map[string]interface{}{"role": "assistant", "id": "delayed"}, + map[string]interface{}{"role": "user", "id": "later"}, + map[string]interface{}{"role": "assistant", "id": "later", "reference": map[string]interface{}{"chunks": []interface{}{}, "saved": true}}, + ) + raw, err := json.Marshal(messages) + if err != nil { + t.Fatal(err) + } + if err := createHistory(t.Context(), db, tables[1], "message", "session", raw); err != nil { + t.Fatal(err) + } + for i := 0; i < conversationHistoryBatchSize+3; i++ { + if err := db.Table(tables[2]).Create(map[string]interface{}{ + "conversation_id": "session", "position": i, "reference": fmt.Sprintf(`{"source":%d}`, i), + }).Error; err != nil { + t.Fatal(err) + } + } + // Completion order differs from message order; message_position wins. + for i, position := range []int{len(messages) - 1, len(messages) - 3} { + if err := db.Table(tables[2]).Create(map[string]interface{}{ + "conversation_id": "session", "position": 500 + i, "message_position": position, "reference": `{"explicit":true}`, + }).Error; err != nil { + t.Fatal(err) + } + } + if err := migrateMessageReferences(t.Context(), db, tables[1], tables[2]); err != nil { + t.Fatal(err) + } + for range 2 { + // Retry copied rows, then retry after table deletion before the marker. + writeDatabaseVersion(t, db, dslVersionTargetVersion) + if err := migrateConversationReferences(t.Context(), db); err != nil { + t.Fatal(err) + } + if got := databaseVersion(t, db); got != conversationReferenceTargetVersion { + t.Fatalf("version = %s", got) + } + for _, table := range []string{conversationReferenceTable, apiConversationReferenceTable} { + if db.Migrator().HasTable(table) { + t.Fatalf("source table still exists: %s", table) + } + } + var rows []entity.ConversationMessage + if err := db.Table(tables[1]).Order("position").Find(&rows).Error; err != nil { + t.Fatal(err) + } + for i, row := range rows { + want := "" + if i >= 4 && i < len(messages)-4 && i%2 == 0 { + want = fmt.Sprintf(`{"source":%d}`, (i-4)/2) + } else if i == len(messages)-3 { + want = `{"explicit":true}` + } else if i == len(messages)-1 { + want = `{"chunks":[],"saved":true}` + } + if string(row.Reference) != want { + t.Fatalf("message %d reference = %s, want %s", i, row.Reference, want) + } + } + } + writeDatabaseVersion(t, db, "v1.0.0-rc3.dev8") + if err := migrateConversationReferences(t.Context(), db); err != nil { + t.Fatal(err) + } + if got := databaseVersion(t, db); got != "v1.0.0-rc3.dev8" { + t.Fatalf("downgraded marker: %s", got) + } + }) + } +} + +func TestConversationReferenceCleanupForMigratedDatabase(t *testing.T) { + for _, version := range []string{conversationReferenceTargetVersion, "v1.0.0-rc3.dev8"} { + t.Run(version, func(t *testing.T) { + db := setupConversationHistoryMigrationDB(t, false) + insertConversation(t, db, "conversation", "session") + if err := createHistory(t.Context(), db, conversationMessageTable, "message", "session", json.RawMessage(`[{"role":"user"},{"role":"assistant"}]`)); err != nil { + t.Fatal(err) + } + if err := db.Table(conversationReferenceTable).Create(map[string]interface{}{ + "conversation_id": "session", "position": 0, "reference": `{"copied":true}`, + }).Error; err != nil { + t.Fatal(err) + } + writeDatabaseVersion(t, db, version) + if err := migrateConversationReferences(t.Context(), db); err != nil { + t.Fatal(err) + } + if got := databaseVersion(t, db); got != version { + t.Fatalf("changed migration version: %s", got) + } + var answer entity.ConversationMessage + if err := db.Where("conversation_id = ? AND position = ?", "session", 1).Take(&answer).Error; err != nil { + t.Fatal(err) + } + if string(answer.Reference) != `{"copied":true}` { + t.Fatalf("version marker prevented copying: %s", answer.Reference) + } + if err := autoMigrateRuntimeModels(t.Context(), db); err != nil { + t.Fatal(err) + } + for _, table := range []string{conversationReferenceTable, apiConversationReferenceTable} { + if db.Migrator().HasTable(table) { + t.Fatalf("server recreated source table: %s", table) + } + } + }) + } +} + +func TestConversationReferenceCleanupSkipsUnmatchedReferences(t *testing.T) { + db := setupConversationHistoryMigrationDB(t, false) + insertConversation(t, db, "api_4_conversation", "unmatched") + insertConversation(t, db, "conversation", "surviving") + if err := createHistory(t.Context(), db, conversationMessageTable, "message", "surviving", json.RawMessage(`[{"role":"user"},{"role":"assistant"}]`)); err != nil { + t.Fatal(err) + } + missingPosition := 9 + for _, row := range []struct { + table, conversationID string + position int + messagePosition *int + reference string + }{ + {apiConversationReferenceTable, "unmatched", 0, nil, `{"chunks":[{"id":"empty-history"}]}`}, + {conversationReferenceTable, "surviving", 0, nil, `{"chunks":[{"id":"surviving"}]}`}, + {conversationReferenceTable, "surviving", 1, nil, `{"chunks":[{"id":"surplus"}]}`}, + {conversationReferenceTable, "surviving", 9, &missingPosition, `{"chunks":[{"id":"missing-position"}]}`}, + } { + if err := db.Table(row.table).Create(map[string]interface{}{ + "conversation_id": row.conversationID, "position": row.position, "message_position": row.messagePosition, "reference": row.reference, + }).Error; err != nil { + t.Fatal(err) + } + } + writeDatabaseVersion(t, db, dslVersionTargetVersion) + if err := migrateConversationReferences(t.Context(), db); err != nil { + t.Fatal(err) + } + for _, table := range []string{conversationReferenceTable, apiConversationReferenceTable} { + if db.Migrator().HasTable(table) { + t.Fatalf("stale references prevented source cleanup: %s", table) + } + } + if got := databaseVersion(t, db); got != conversationReferenceTargetVersion { + t.Fatalf("migration version = %s", got) + } + var messages []entity.ConversationMessage + if err := db.Order("position").Find(&messages).Error; err != nil { + t.Fatal(err) + } + if len(messages) != 2 || string(messages[1].Reference) != `{"chunks":[{"id":"surviving"}]}` { + t.Fatalf("stale references changed the surviving history: %+v", messages) + } + if err := migrateConversationReferences(t.Context(), db); err != nil { + t.Fatal(err) + } +} + +func TestMigrateConversationReferencesFailurePreservesVersion(t *testing.T) { + db := setupConversationHistoryMigrationDB(t, false) + insertConversation(t, db, "conversation", "session") + if err := db.Table(conversationReferenceTable).Create(map[string]interface{}{ + "conversation_id": "session", "position": 0, "reference": "invalid JSON", + }).Error; err != nil { + t.Fatal(err) + } + writeDatabaseVersion(t, db, dslVersionTargetVersion) + if err := migrateConversationReferences(t.Context(), db); err == nil { + t.Fatal("invalid reference accepted") + } + if got := databaseVersion(t, db); got != dslVersionTargetVersion { + t.Fatalf("advanced version after failure: %s", got) + } +} + +func TestMigrateConversationReferenceMetadata(t *testing.T) { + for _, tables := range [][3]string{ + {"conversation", conversationMessageTable, conversationReferenceTable}, + {"api_4_conversation", apiConversationMessageTable, apiConversationReferenceTable}, + } { + t.Run(tables[0], func(t *testing.T) { + db := setupConversationHistoryMigrationDB(t, false) + insertConversation(t, db, tables[0], "session") + if err := createHistory(t.Context(), db, tables[1], "message", "session", json.RawMessage(`[{"role":"user"},{"role":"assistant","reference":[{"old":"inline chunks"}]},{"role":"assistant"}]`)); err != nil { + t.Fatal(err) + } + // Existing flattened rows stored references in their metadata column. + if err := db.Table(tables[1]).Where("conversation_id = ?", "session").Update("metadata", `{"reference":{"chunks":[]},"attachment":"file"}`).Error; err != nil { + t.Fatal(err) + } + if err := db.Table(tables[2]).Create(map[string]interface{}{ + "conversation_id": "session", "position": 0, "reference": `{"chunks":[{"id":"source"}]}`, + }).Error; err != nil { + t.Fatal(err) + } + if err := migrateConversationReferences(t.Context(), db); err != nil { + t.Fatal(err) + } + var messages []entity.ConversationMessage + if err := db.Table(tables[1]).Order("position").Find(&messages).Error; err != nil { + t.Fatal(err) + } + for i, message := range messages { + if message.Metadata != `{"attachment":"file"}` { + t.Fatalf("metadata still contains reference: %s", message.Metadata) + } + want := `{"chunks":[]}` + if i == 1 { + want = `{"chunks":[{"id":"source"}]}` + } + if string(message.Reference) != want { + t.Fatalf("message %d reference = %s, want %s", i, message.Reference, want) + } + } + }) + } +} + +func TestMessageReferencePersistenceAndProjection(t *testing.T) { + for _, table := range []string{conversationMessageTable, apiConversationMessageTable} { + t.Run(table, func(t *testing.T) { + db := setupConversationHistoryMigrationDB(t, false) + parent := "conversation" + if table == apiConversationMessageTable { + parent = "api_4_conversation" + } + insertConversation(t, db, parent, "session") + for _, update := range []ConversationHistoryUpdate{ + {Message: map[string]interface{}{"id": "turn", "role": "user", "content": "question"}}, + {QuestionID: "turn", Message: map[string]interface{}{"id": "turn", "role": "assistant", "content": "answer", "prompt": "system prompt", "usage": map[string]interface{}{"total_tokens": 10}, "elapsed_ms": 50, "reference": map[string]interface{}{"chunks": []interface{}{map[string]interface{}{"id": "chunk"}}}}}, + } { + if err := updateConversationHistory(t.Context(), db, table, "session", update); err != nil { + t.Fatal(err) + } + } + for _, options := range []ConversationHistoryOptions{ + {}, + {IncludeUsage: true, IncludeElapsedMs: true}, + {IncludeReference: true}, + {IncludePrompt: true}, + {IncludeReference: true, IncludePrompt: true, IncludeUsage: true, IncludeElapsedMs: true}, + } { + var sql string + capture := db.Session(&gorm.Session{DryRun: true}) + if err := capture.Callback().Query().After("gorm:query").Register("capture_history_sql", func(tx *gorm.DB) { + sql = tx.Statement.SQL.String() + }); err != nil { + t.Fatal(err) + } + if _, err := loadHistory(t.Context(), capture, table, []string{"session"}, options); err != nil { + t.Fatal(err) + } + if err := capture.Callback().Query().Remove("capture_history_sql"); err != nil { + t.Fatal(err) + } + for column, include := range map[string]bool{"reference": options.IncludeReference, "prompt": options.IncludePrompt, "usage": options.IncludeUsage, "elapsed_ms": options.IncludeElapsedMs} { + if strings.Contains(sql, column) != include { + t.Fatalf("options=%+v projection %s: %s", options, column, sql) + } + } + if strings.Contains(sql, "SELECT *") || !strings.Contains(sql, "content") { + t.Fatalf("invalid projection: %s", sql) + } + + history, err := loadHistory(t.Context(), db, table, []string{"session"}, options) + if err != nil { + t.Fatal(err) + } + var messages []map[string]interface{} + if err := json.Unmarshal(history["session"], &messages); err != nil { + t.Fatal(err) + } + for key, include := range map[string]bool{"reference": options.IncludeReference, "prompt": options.IncludePrompt, "usage": options.IncludeUsage, "elapsed_ms": options.IncludeElapsedMs} { + if _, ok := messages[1][key]; ok != include { + t.Fatalf("options=%+v key %s in %s", options, key, history["session"]) + } + } + } + }) + } +} diff --git a/internal/dao/database.go b/internal/dao/database.go index dfd7daa310..2a07533a44 100644 --- a/internal/dao/database.go +++ b/internal/dao/database.go @@ -167,12 +167,10 @@ func InitDB(ctx context.Context, migrateDB bool) error { &entity.ChatChannel{}, &entity.ChatSession{}, &entity.ConversationMessage{}, - &entity.ConversationReference{}, &entity.Task{}, &entity.APIToken{}, &entity.API4Conversation{}, &entity.API4ConversationMessage{}, - &entity.API4ConversationReference{}, &entity.Knowledgebase{}, &entity.InvitationCode{}, &entity.Document{}, @@ -275,7 +273,7 @@ func InitDB(ctx context.Context, migrateDB bool) error { } // Add pipeline_version_id columns and lock existing dataflow bindings to a - // version. Runs LAST so the shared migration marker only advances to + // version. Runs after the lower-target migrations so the marker advances to // v1.0.0-rc3.dev6 after every lower-target migration (notably the // v1.0.0-rc3.dev5 tenant default-model column drop) has already run. if err = MigrateDSLVersionColumns(ctx, DB); err != nil { @@ -284,6 +282,9 @@ func InitDB(ctx context.Context, migrateDB bool) error { if err = BackfillDSLVersionBindings(ctx, DB); err != nil { return fmt.Errorf("failed to backfill DSL version bindings: %w", err) } + if err = migrateConversationReferences(ctx, DB); err != nil { + return fmt.Errorf("failed to move references into messages: %w", err) + } common.Info("Database schema migrated successfully") } else { for _, column := range []string{"backlog_doc_ids", "inflight_doc_ids"} { @@ -394,9 +395,7 @@ func autoMigrateRuntimeModels(ctx context.Context, db *gorm.DB) error { &entity.KnowledgeCompileBacklog{}, &entity.KnowledgeCompileInflight{}, &entity.ConversationMessage{}, - &entity.ConversationReference{}, &entity.API4ConversationMessage{}, - &entity.API4ConversationReference{}, } for _, m := range goRuntimeModels { if err := autoMigrateSafely(ctx, db, m); err != nil { diff --git a/internal/entity/api_for_conversation.go b/internal/entity/api_for_conversation.go index f79d660482..6a7d0a05d9 100644 --- a/internal/entity/api_for_conversation.go +++ b/internal/entity/api_for_conversation.go @@ -18,8 +18,8 @@ package entity import "encoding/json" -// API4Conversation API for conversation model. Message and Reference are -// hydrated from their ordered child tables. +// API4Conversation API for conversation model. Messages are +// hydrated from the ordered message table. type API4Conversation struct { ID string `gorm:"column:id;primaryKey;size:32" json:"id"` Name *string `gorm:"column:name;size:255" json:"name,omitempty"` @@ -27,7 +27,6 @@ type API4Conversation struct { UserID string `gorm:"column:user_id;size:255;not null;index" json:"user_id"` ExpUserID *string `gorm:"column:exp_user_id;size:255;index" json:"exp_user_id,omitempty"` Message json.RawMessage `gorm:"-" json:"message,omitempty"` - Reference json.RawMessage `gorm:"-" json:"reference,omitempty"` Tokens int `gorm:"column:tokens" json:"tokens"` Source *string `gorm:"column:source;size:16" json:"source,omitempty"` DSL JSONMap `gorm:"column:dsl;type:longtext" json:"dsl,omitempty"` diff --git a/internal/entity/chat.go b/internal/entity/chat.go index ae4ea2905b..d82400147b 100644 --- a/internal/entity/chat.go +++ b/internal/entity/chat.go @@ -72,15 +72,14 @@ type ChatListItem struct { TenantAvatar *string `gorm:"column:tenant_avatar" json:"tenant_avatar,omitempty"` } -// ChatSession chat session model. Message and Reference are hydrated from -// their ordered child tables. +// ChatSession chat session model. Messages are hydrated from +// the ordered message table. type ChatSession struct { - ID string `gorm:"column:id;primaryKey;size:32" json:"id"` - DialogID string `gorm:"column:dialog_id;size:32;not null;index" json:"dialog_id"` - Name *string `gorm:"column:name;size:255;index" json:"name,omitempty"` - Message json.RawMessage `gorm:"-" json:"message,omitempty"` - Reference json.RawMessage `gorm:"-" json:"reference"` - UserID *string `gorm:"column:user_id;size:255;index" json:"user_id,omitempty"` + ID string `gorm:"column:id;primaryKey;size:32" json:"id"` + DialogID string `gorm:"column:dialog_id;size:32;not null;index" json:"dialog_id"` + Name *string `gorm:"column:name;size:255;index" json:"name,omitempty"` + Message json.RawMessage `gorm:"-" json:"message,omitempty"` + UserID *string `gorm:"column:user_id;size:255;index" json:"user_id,omitempty"` BaseModel } diff --git a/internal/entity/conversation_history.go b/internal/entity/conversation_history.go index dde3226f8b..a4b9388eb8 100644 --- a/internal/entity/conversation_history.go +++ b/internal/entity/conversation_history.go @@ -20,18 +20,19 @@ import "encoding/json" // ConversationMessageFields contains searchable fields and variable message metadata. type ConversationMessageFields struct { - MessageID *string `gorm:"column:message_id;size:255;index:,composite:message_lookup,priority:2" json:"id,omitempty"` - Role *string `gorm:"column:role;size:32;index:,composite:message_lookup,priority:3" json:"role,omitempty"` - Content *string `gorm:"column:content;type:longtext" json:"content,omitempty"` - ContentType string `gorm:"column:content_type;size:8" json:"-"` - Status *string `gorm:"column:status;size:32" json:"status,omitempty"` - ThumbUp *bool `gorm:"column:thumb_up" json:"thumbup,omitempty"` - Feedback *string `gorm:"column:feedback;type:longtext" json:"feedback,omitempty"` - CreatedAt *float64 `gorm:"column:created_at" json:"created_at,omitempty"` - Prompt *string `gorm:"column:prompt;type:longtext" json:"prompt,omitempty"` - Usage JSONMap `gorm:"column:usage;type:longtext" json:"usage,omitempty"` - ElapsedMs *float64 `gorm:"column:elapsed_ms" json:"elapsed_ms,omitempty"` - Metadata string `gorm:"column:metadata;type:longtext" json:"-"` + MessageID *string `gorm:"column:message_id;size:255;index:,composite:message_lookup,priority:2" json:"id,omitempty"` + Role *string `gorm:"column:role;size:32;index:,composite:message_lookup,priority:3" json:"role,omitempty"` + Content *string `gorm:"column:content;type:longtext" json:"content,omitempty"` + ContentType string `gorm:"column:content_type;size:8" json:"-"` + Status *string `gorm:"column:status;size:32" json:"status,omitempty"` + ThumbUp *bool `gorm:"column:thumb_up" json:"thumbup,omitempty"` + Feedback *string `gorm:"column:feedback;type:longtext" json:"feedback,omitempty"` + CreatedAt *float64 `gorm:"column:created_at" json:"created_at,omitempty"` + Prompt *string `gorm:"column:prompt;type:longtext" json:"prompt,omitempty"` + Usage JSONMap `gorm:"column:usage;type:longtext" json:"usage,omitempty"` + ElapsedMs *float64 `gorm:"column:elapsed_ms" json:"elapsed_ms,omitempty"` + Reference json.RawMessage `gorm:"column:reference;type:longtext;serializer:json" json:"reference,omitempty"` + Metadata string `gorm:"column:metadata;type:longtext" json:"-"` } // ConversationMessage stores one ordered message for a conversation. diff --git a/internal/handler/agent.go b/internal/handler/agent.go index 6aef2957c2..68b578a8b2 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -989,19 +989,26 @@ func (h *AgentHandler) ListAgentSessions(c *gin.Context) { dslParam := c.Query("dsl") includeDSL := dslParam != "false" && dslParam != "False" ctx := c.Request.Context() + options, err := sessionHistoryQuery(c) + if err != nil { + common.ErrorWithCode(c, common.CodeArgumentError, err.Error()) + return + } resp, code, err := h.agentService.ListAgentSessions(ctx, user.ID, user.ID, canvasID, service.ListAgentSessionsRequest{ - Page: page, - PageSize: pageSize, - Keywords: keywords, - FromDate: fromDate, - ToDate: toDate, - OrderBy: orderby, - Desc: desc, - SessionID: sessionID, - UserID: queryUserID, - ExpUserID: expUserID, - IncludeDSL: includeDSL, - NoHistory: c.Query("include_history") == "false" || c.Query("include_history") == "False", + Page: page, + PageSize: pageSize, + Keywords: keywords, + FromDate: fromDate, + ToDate: toDate, + OrderBy: orderby, + Desc: desc, + SessionID: sessionID, + UserID: queryUserID, + ExpUserID: expUserID, + IncludeDSL: includeDSL, + IncludeReference: options.IncludeReference, + IncludePrompt: options.IncludePrompt, + NoHistory: c.Query("include_history") == "false" || c.Query("include_history") == "False", }) if err != nil { common.ErrorWithCode(c, code, err.Error()) @@ -1052,7 +1059,12 @@ func (h *AgentHandler) GetAgentSession(c *gin.Context) { canvasID := c.Param("canvas_id") sessionID := c.Param("session_id") ctx := c.Request.Context() - row, code, err := h.agentService.GetAgentSession(ctx, user.ID, canvasID, sessionID) + options, err := sessionHistoryQuery(c) + if err != nil { + common.ErrorWithCode(c, common.CodeArgumentError, err.Error()) + return + } + row, code, err := h.agentService.GetAgentSession(ctx, user.ID, canvasID, sessionID, options) if err != nil { common.ErrorWithCode(c, code, err.Error()) return diff --git a/internal/handler/chat_session.go b/internal/handler/chat_session.go index 88dd6383cf..7dad6110b2 100644 --- a/internal/handler/chat_session.go +++ b/internal/handler/chat_session.go @@ -28,6 +28,7 @@ import ( "github.com/gin-gonic/gin" + "ragflow/internal/dao" "ragflow/internal/service" ) @@ -45,6 +46,21 @@ func NewChatSessionHandler(chatSessionService *service.ChatSessionService, userS } } +func sessionHistoryQuery(c *gin.Context) (dao.ConversationHistoryOptions, error) { + var options dao.ConversationHistoryOptions + for name, target := range map[string]*bool{ + "include_reference": &options.IncludeReference, + "include_prompt": &options.IncludePrompt, + } { + value, err := strconv.ParseBool(c.DefaultQuery(name, "false")) + if err != nil { + return options, fmt.Errorf("%s must be a boolean", name) + } + *target = value + } + return options, nil +} + // ListChatSessions list chat sessions for a dialog // @Summary List Chat Sessions // @Description Get list of chat sessions for a specific dialog @@ -110,7 +126,13 @@ func (h *ChatSessionHandler) ListChatSessions(c *gin.Context) { // Call service to list chat sessions ctx := c.Request.Context() includeHistory := c.DefaultQuery("include_history", "true") - result, err := h.chatSessionService.ListChatSessions(ctx, userID, chatID, c.Query("id"), c.Query("name"), terms, page, pageSize, includeHistory != "false" && includeHistory != "False") + options, err := sessionHistoryQuery(c) + if err != nil { + common.ErrorWithCode(c, common.CodeArgumentError, err.Error()) + return + } + options.NoHistory = includeHistory == "false" || includeHistory == "False" + result, err := h.chatSessionService.ListChatSessions(ctx, userID, chatID, c.Query("id"), c.Query("name"), terms, page, pageSize, options) if err != nil { if respondPermissionErrorIf(c, err, true) { return @@ -314,7 +336,12 @@ func (h *ChatSessionHandler) GetSession(c *gin.Context) { chatID, sessionID := c.Param("chat_id"), c.Param("session_id") ctx := c.Request.Context() - result, code, err := h.chatSessionService.GetSession(ctx, userID, chatID, sessionID) + options, err := sessionHistoryQuery(c) + if err != nil { + common.ErrorWithCode(c, common.CodeArgumentError, err.Error()) + return + } + result, code, err := h.chatSessionService.GetSession(ctx, userID, chatID, sessionID, options) if err != nil { if respondPermissionErrorIf(c, err, true) { return diff --git a/internal/handler/session_history_test.go b/internal/handler/session_history_test.go new file mode 100644 index 0000000000..061c447eca --- /dev/null +++ b/internal/handler/session_history_test.go @@ -0,0 +1,167 @@ +package handler + +import ( + "encoding/json" + "io" + "net/http" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" + "ragflow/internal/service" +) + +func TestSessionHistoryQueryFlags(t *testing.T) { + for _, endpoint := range []string{"chat", "chat_list", "agent", "agent_list", "create", "update", "delete_message", "feedback"} { + for _, flags := range []struct { + query string + reference bool + prompt bool + invalid bool + }{ + {}, + {query: "include_reference=false&include_prompt=false"}, + {query: "include_reference=true", reference: true}, + {query: "include_prompt=true", prompt: true}, + {query: "include_reference=true&include_prompt=true", reference: true, prompt: true}, + {query: "include_reference=invalid", invalid: true}, + {query: "include_prompt=invalid", invalid: true}, + } { + t.Run(endpoint+"/"+flags.query, func(t *testing.T) { + mutation := endpoint == "create" || endpoint == "update" || endpoint == "delete_message" || endpoint == "feedback" + c, w, db := setupGinContextWithUserAndDB(t, http.MethodGet, "/?"+flags.query) + if mutation { + body := `{"name":"session"}` + if endpoint == "feedback" { + body = `{"thumbup":true}` + } + c.Request.Body = io.NopCloser(strings.NewReader(body)) + } + if err := db.AutoMigrate(&entity.Chat{}, &entity.ChatSession{}, &entity.ConversationMessage{}, &entity.API4ConversationMessage{}); err != nil { + t.Fatal(err) + } + if err := db.Create(&entity.Chat{ID: "chat", TenantID: "user-1", LLMID: "model", LLMSetting: entity.JSONMap{}, PromptConfig: entity.JSONMap{}, KBIDs: entity.JSONSlice{}, DoRefer: "1", Status: sptr("1")}).Error; err != nil { + t.Fatal(err) + } + if err := db.Create(&entity.UserCanvas{ID: "agent", UserID: "user-1"}).Error; err != nil { + t.Fatal(err) + } + message := json.RawMessage(`[{"id":"answer","role":"assistant","content":"answer","reference":{"chunks":[],"doc_aggs":[]},"prompt":"large prompt","usage":{"total_tokens":1},"elapsed_ms":50}]`) + if err := dao.NewChatSessionDAO().Create(c.Request.Context(), db, &entity.ChatSession{ID: "session", DialogID: "chat", UserID: sptr("user-1"), Message: message}); err != nil { + t.Fatal(err) + } + if mutation { + for _, update := range []dao.ConversationHistoryUpdate{ + {Message: map[string]interface{}{"id": "turn", "role": "user", "content": "question"}}, + {QuestionID: "turn", Message: map[string]interface{}{"id": "turn", "role": "assistant", "content": "answer", "prompt": "large prompt", "reference": map[string]interface{}{"chunks": []interface{}{}}}}, + } { + if err := dao.NewChatSessionDAO().UpdateByID(c.Request.Context(), db, "session", map[string]interface{}{"history_update": update}); err != nil { + t.Fatal(err) + } + } + } + if err := dao.NewAPI4ConversationDAO().Create(c.Request.Context(), db, &entity.API4Conversation{ID: "session", DialogID: "agent", UserID: "user-1", Message: message}); err != nil { + t.Fatal(err) + } + var historySQL string + if err := db.Callback().Query().After("gorm:query").Register("capture_session_history", func(tx *gorm.DB) { + sql := tx.Statement.SQL.String() + if strings.Contains(sql, "`content`") && (strings.Contains(sql, "FROM `conversation_message`") || strings.Contains(sql, "FROM `api_4_conversation_message`")) { + historySQL = sql + } + }); err != nil { + t.Fatal(err) + } + chat := NewChatSessionHandler(service.NewChatSessionService(), nil) + agent := NewAgentHandler(c.Request.Context(), service.NewAgentService(), nil) + c.Params = gin.Params{{Key: "chat_id", Value: "chat"}, {Key: "canvas_id", Value: "agent"}, {Key: "session_id", Value: "session"}, {Key: "msg_id", Value: "turn"}} + switch endpoint { + case "chat": + chat.GetSession(c) + case "chat_list": + chat.ListChatSessions(c) + case "agent": + agent.GetAgentSession(c) + case "agent_list": + agent.ListAgentSessions(c) + case "create": + chat.CreateSession(c) + case "update": + chat.UpdateSession(c) + case "delete_message": + chat.DeleteSessionMessage(c) + case "feedback": + chat.UpdateMessageFeedback(c) + } + var response struct { + Code common.ErrorCode `json:"code"` + Data json.RawMessage `json:"data"` + } + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatal(err) + } + if flags.invalid && !mutation { + if response.Code != common.CodeArgumentError || historySQL != "" { + t.Fatalf("invalid flags should fail before history SQL: %s", w.Body.String()) + } + return + } + if response.Code != common.CodeSuccess { + t.Fatalf("session request failed: %s", w.Body.String()) + } + if mutation { + var session struct { + Messages []map[string]interface{} `json:"messages"` + } + if err := json.Unmarshal(response.Data, &session); err != nil || len(session.Messages) == 0 { + t.Fatalf("invalid mutation response: %s, %v", response.Data, err) + } + for _, message := range session.Messages { + for _, field := range []string{"reference", "prompt"} { + if _, exists := message[field]; exists { + t.Fatalf("mutation returned %s: %s", field, response.Data) + } + } + } + if strings.Contains(historySQL, "`prompt`") || strings.Contains(historySQL, "`reference`") != (endpoint == "feedback") { + t.Fatalf("mutation loaded unneeded large fields: %s", historySQL) + } + return + } + data := response.Data + if strings.HasSuffix(endpoint, "_list") { + var rows []json.RawMessage + if err := json.Unmarshal(data, &rows); err != nil || len(rows) != 1 { + t.Fatalf("invalid session list: %s, %v", data, err) + } + data = rows[0] + } + var session map[string]json.RawMessage + if err := json.Unmarshal(data, &session); err != nil { + t.Fatal(err) + } + key := "messages" + if endpoint != "chat" { + key = "message" + } + var messages []map[string]interface{} + if err := json.Unmarshal(session[key], &messages); err != nil || len(messages) != 1 { + t.Fatalf("invalid messages: %s, %v", session[key], err) + } + for field, include := range map[string]bool{"reference": flags.reference, "prompt": flags.prompt, "usage": true, "elapsed_ms": true} { + if _, exists := messages[0][field]; exists != include { + t.Fatalf("unexpected %s in response: %s", field, w.Body.String()) + } + if strings.Contains(historySQL, "`"+field+"`") != include { + t.Fatalf("unexpected %s in SQL: %s", field, historySQL) + } + } + }) + } + } +} diff --git a/internal/service/agent.go b/internal/service/agent.go index 11b7c5af5e..03c61c393f 100644 --- a/internal/service/agent.go +++ b/internal/service/agent.go @@ -2574,7 +2574,6 @@ func (s *AgentService) createAgentRunSession( UserID: userID, ExpUserID: &userID, Message: json.RawMessage(`[]`), - Reference: json.RawMessage(`[]`), Source: &source, DSL: entity.JSONMap(runDSL), } @@ -2702,12 +2701,12 @@ func (s *AgentService) persistAgentRunSession( if sessionID == "" || s == nil || s.api4ConversationDAO == nil || dao.DB == nil { return nil } - history := dao.ConversationHistoryUpdate{Reference: normalizeAgentReferenceEntry(reference), AppendReference: true} + history := dao.ConversationHistoryUpdate{} if stringifyAgentUserInput(agentRunQuery(userInput)) != "" { history.QuestionID = messageID } if appendAssistantMessage { - history.Message = map[string]interface{}{"role": "assistant", "content": agentSessionMessageContent(answer, thinking), "id": messageID, "created_at": now} + history.Message = map[string]interface{}{"role": "assistant", "content": agentSessionMessageContent(answer, thinking), "id": messageID, "created_at": now, "reference": normalizeAgentReferenceEntry(reference)} } updates := map[string]interface{}{ "round": gorm.Expr("COALESCE(round, 0) + 1"), diff --git a/internal/service/agent_cancel_test.go b/internal/service/agent_cancel_test.go index 65fee76430..abfc294c5c 100644 --- a/internal/service/agent_cancel_test.go +++ b/internal/service/agent_cancel_test.go @@ -83,11 +83,10 @@ func TestCancelSessionRunFinishedSessionDoesNotCreateCancelMarker(t *testing.T) testDB := setupServiceTestDB(t) pushServiceDB(t, testDB) if err := dao.NewAPI4ConversationDAO().Create(t.Context(), testDB, &entity.API4Conversation{ - ID: "session-persisted", - DialogID: "agent-1", - UserID: "user-a", - Message: json.RawMessage(`[]`), - Reference: json.RawMessage(`[]`), + ID: "session-persisted", + DialogID: "agent-1", + UserID: "user-a", + Message: json.RawMessage(`[]`), }); err != nil { t.Fatalf("create conversation: %v", err) } diff --git a/internal/service/agent_persistence_test.go b/internal/service/agent_persistence_test.go index 428cb347bf..ef67a93dd8 100644 --- a/internal/service/agent_persistence_test.go +++ b/internal/service/agent_persistence_test.go @@ -39,11 +39,10 @@ func TestAgentRunSessionUpdateFailurePreventsSuccessEvents(t *testing.T) { t.Cleanup(func() { dao.DB = originalDB }) if err := dao.NewAPI4ConversationDAO().Create(t.Context(), testDB, &entity.API4Conversation{ - ID: "session-update-failure", - DialogID: "canvas-update-failure", - UserID: "user-1", - Message: json.RawMessage(`[]`), - Reference: json.RawMessage(`[]`), + ID: "session-update-failure", + DialogID: "canvas-update-failure", + UserID: "user-1", + Message: json.RawMessage(`[]`), }); err != nil { t.Fatalf("create session: %v", err) } @@ -128,11 +127,10 @@ func TestPersistAgentRunSessionPreservesThinking(t *testing.T) { t.Cleanup(func() { dao.DB = originalDB }) if err := dao.NewAPI4ConversationDAO().Create(t.Context(), testDB, &entity.API4Conversation{ - ID: "session-think", - DialogID: "canvas-think", - UserID: "user-1", - Message: json.RawMessage(`[]`), - Reference: json.RawMessage(`[]`), + ID: "session-think", + DialogID: "canvas-think", + UserID: "user-1", + Message: json.RawMessage(`[]`), }); err != nil { t.Fatalf("create session: %v", err) } @@ -145,7 +143,7 @@ func TestPersistAgentRunSessionPreservesThinking(t *testing.T) { t.Fatalf("persist: %v", err) } - conv, err := dao.NewAPI4ConversationDAO().GetByID(t.Context(), testDB, "session-think") + conv, err := dao.NewAPI4ConversationDAO().GetByID(t.Context(), testDB, "session-think", dao.ConversationHistoryOptions{IncludeReference: true}) if err != nil { t.Fatalf("reload session: %v", err) } @@ -222,13 +220,13 @@ func TestPersistAgentRunSessionPreservesThinking(t *testing.T) { t.Fatalf("concurrent run: %v", err) } } - conv, err = dao.NewAPI4ConversationDAO().GetByID(ctx, testDB, "session-think") + conv, err = dao.NewAPI4ConversationDAO().GetByID(ctx, testDB, "session-think", dao.ConversationHistoryOptions{IncludeReference: true}) if err != nil { t.Fatal(err) } messages = parseMessages(conv.Message) - if len(messages) != 6 || len(parseReferenceList(conv.Reference)) != 3 { - t.Fatalf("concurrent turns were lost: messages=%s references=%s", conv.Message, conv.Reference) + if len(messages) != 6 || messages[1]["reference"] == nil || messages[3]["reference"] == nil || messages[5]["reference"] == nil { + t.Fatalf("concurrent turns were lost: messages=%s", conv.Message) } for _, index := range []int{2, 4} { if messages[index]["role"] != "user" || messages[index+1]["role"] != "assistant" || messages[index]["id"] != messages[index+1]["id"] { diff --git a/internal/service/agent_run_e2e_test.go b/internal/service/agent_run_e2e_test.go index a108389226..37427b7be6 100644 --- a/internal/service/agent_run_e2e_test.go +++ b/internal/service/agent_run_e2e_test.go @@ -287,11 +287,10 @@ func TestRunAgent_SessionHistoryFeedsSysHistoryAndPersists(t *testing.T) { } makeCanvasWithDSL(t, "canvas-history", "user-1", "tenant-1", "v-history", dsl) if err := dao.NewAPI4ConversationDAO().Create(t.Context(), testDB, &entity.API4Conversation{ - ID: "session-history", - DialogID: "canvas-history", - UserID: "user-1", - Message: json.RawMessage(`[]`), - Reference: json.RawMessage(`[]`), + ID: "session-history", + DialogID: "canvas-history", + UserID: "user-1", + Message: json.RawMessage(`[]`), }); err != nil { t.Fatalf("create session: %v", err) } @@ -472,12 +471,11 @@ func TestRunAgent_RejectsSessionOwnedByAnotherUser(t *testing.T) { makeCanvasWithDSL(t, "canvas-session-owner", "user-1", "tenant-1", "v-session-owner", dsl) foreignMessage := json.RawMessage(`[{"role":"assistant","content":"foreign"}]`) if err := dao.NewAPI4ConversationDAO().Create(t.Context(), testDB, &entity.API4Conversation{ - ID: "session-foreign", - DialogID: "canvas-session-owner", - UserID: "user-2", - Message: foreignMessage, - Reference: json.RawMessage(`[]`), - DSL: entity.JSONMap(dsl), + ID: "session-foreign", + DialogID: "canvas-session-owner", + UserID: "user-2", + Message: foreignMessage, + DSL: entity.JSONMap(dsl), }); err != nil { t.Fatalf("create foreign session: %v", err) } @@ -571,12 +569,11 @@ func TestRunAgent_RealCanvas_WaitForUserResume(t *testing.T) { } makeCanvasWithDSL(t, "canvas-fillup", "user-1", "tenant-1", "v-fillup", dsl) if err := dao.NewAPI4ConversationDAO().Create(t.Context(), testDB, &entity.API4Conversation{ - ID: "session-fillup", - DialogID: "canvas-fillup", - UserID: "user-1", - Message: json.RawMessage(`[]`), - Reference: json.RawMessage(`[]`), - DSL: entity.JSONMap(dsl), + ID: "session-fillup", + DialogID: "canvas-fillup", + UserID: "user-1", + Message: json.RawMessage(`[]`), + DSL: entity.JSONMap(dsl), }); err != nil { t.Fatalf("create session: %v", err) } @@ -708,12 +705,11 @@ func TestRunAgent_InterruptPersistsPartialAssistantHistory(t *testing.T) { } makeCanvasWithDSL(t, "canvas-partial", "user-1", "tenant-1", "v-partial", dsl) if err := dao.NewAPI4ConversationDAO().Create(t.Context(), testDB, &entity.API4Conversation{ - ID: "session-partial", - DialogID: "canvas-partial", - UserID: "user-1", - Message: json.RawMessage(`[]`), - Reference: json.RawMessage(`[]`), - DSL: entity.JSONMap(dsl), + ID: "session-partial", + DialogID: "canvas-partial", + UserID: "user-1", + Message: json.RawMessage(`[]`), + DSL: entity.JSONMap(dsl), }); err != nil { t.Fatalf("create session: %v", err) } diff --git a/internal/service/agent_sessions.go b/internal/service/agent_sessions.go index d2ea1cbc76..733fa6806d 100644 --- a/internal/service/agent_sessions.go +++ b/internal/service/agent_sessions.go @@ -23,7 +23,6 @@ import ( "fmt" "ragflow/internal/utility" "sort" - "strconv" "strings" "time" @@ -41,18 +40,20 @@ const ( // ListAgentSessionsRequest are the parameters for ListAgentSessions. type ListAgentSessionsRequest struct { - SessionID string - UserID string - Page int - PageSize int - Keywords string - FromDate string - ToDate string - OrderBy string - Desc bool - ExpUserID string - IncludeDSL bool - NoHistory bool + SessionID string + UserID string + Page int + PageSize int + Keywords string + FromDate string + ToDate string + OrderBy string + Desc bool + ExpUserID string + IncludeDSL bool + NoHistory bool + IncludeReference bool + IncludePrompt bool } // ListAgentSessionsResponse is the response body for ListAgentSessions. @@ -118,26 +119,6 @@ func parseAgentSessionDate(value string, isEnd bool) (*time.Time, error) { func normalizeAgentSession(session *entity.API4Conversation, includeDSL bool) map[string]interface{} { messages := parseAgentSessionMessages(session.Message) - references := parseAgentSessionReferences(session.Reference) - - for _, message := range messages { - delete(message, "prompt") - } - - if len(references) > 0 { - assistantMessages := make([]map[string]interface{}, 0) - for i, message := range messages { - role, _ := message["role"].(string) - if i != 0 && role != "user" { - assistantMessages = append(assistantMessages, message) - } - } - - for i := 0; i < len(assistantMessages) && i < len(references); i++ { - rawChunks, _ := references[i]["chunks"].([]interface{}) - assistantMessages[i]["reference"] = normalizeAgentReferenceChunks(rawChunks) - } - } result := map[string]interface{}{ "id": session.ID, @@ -164,48 +145,6 @@ func normalizeAgentSession(session *entity.API4Conversation, includeDSL bool) ma return result } -func parseAgentSessionReferences(raw json.RawMessage) []map[string]interface{} { - if len(raw) == 0 { - return []map[string]interface{}{} - } - - var references []map[string]interface{} - if err := json.Unmarshal(raw, &references); err == nil { - for i, reference := range references { - references[i] = normalizeAgentReferenceEntry(reference) - } - return references - } - - var referenceMap map[string]interface{} - if err := json.Unmarshal(raw, &referenceMap); err != nil { - return []map[string]interface{}{} - } - if _, ok := referenceMap["chunks"]; ok { - return []map[string]interface{}{normalizeAgentReferenceEntry(referenceMap)} - } - - keys := make([]string, 0, len(referenceMap)) - for key := range referenceMap { - keys = append(keys, key) - } - sort.Slice(keys, func(i, j int) bool { - left, _ := strconv.Atoi(keys[i]) - right, _ := strconv.Atoi(keys[j]) - return left < right - }) - - result := make([]map[string]interface{}, 0, len(keys)) - for _, key := range keys { - reference, ok := referenceMap[key].(map[string]interface{}) - if !ok { - continue - } - result = append(result, normalizeAgentReferenceEntry(reference)) - } - return result -} - func parseAgentSessionMessages(raw json.RawMessage) []map[string]interface{} { if len(raw) == 0 { return []map[string]interface{}{} @@ -214,6 +153,12 @@ func parseAgentSessionMessages(raw json.RawMessage) []map[string]interface{} { if err := json.Unmarshal(raw, &messages); err != nil { return []map[string]interface{}{} } + for _, message := range messages { + if reference, ok := message["reference"].(map[string]interface{}); ok { + chunks, _ := reference["chunks"].([]interface{}) + reference["chunks"] = normalizeAgentReferenceChunks(chunks) + } + } return messages } @@ -342,19 +287,20 @@ func (s *AgentService) ListAgentSessions(ctx context.Context, userID, _ string, } total, sessions, err := sessionDAO.ListAgentSessions(ctx, dao.DB, dao.ListAgentSessionsParams{ - AgentID: agentID, - Page: req.Page, - PageSize: req.PageSize, - OrderBy: req.OrderBy, - Desc: req.Desc, - SessionID: req.SessionID, - UserID: req.UserID, - IncludeDSL: req.IncludeDSL, - NoHistory: req.NoHistory, - Keywords: req.Keywords, - FromDate: fromDate, - ToDate: toDate, - ExpUserID: req.ExpUserID, + AgentID: agentID, + Page: req.Page, + PageSize: req.PageSize, + OrderBy: req.OrderBy, + Desc: req.Desc, + SessionID: req.SessionID, + UserID: req.UserID, + IncludeDSL: req.IncludeDSL, + NoHistory: req.NoHistory, + HistoryOptions: sessionHistoryOptions(dao.ConversationHistoryOptions{IncludeReference: req.IncludeReference, IncludePrompt: req.IncludePrompt}), + Keywords: req.Keywords, + FromDate: fromDate, + ToDate: toDate, + ExpUserID: req.ExpUserID, }) if err != nil { return nil, common.CodeServerError, err @@ -371,7 +317,7 @@ func (s *AgentService) ListAgentSessions(ctx context.Context, userID, _ string, } // GetAgentSession fetches a single conversation belonging to agentID. -func (s *AgentService) GetAgentSession(ctx context.Context, userID, agentID, sessionID string) (*entity.API4Conversation, common.ErrorCode, error) { +func (s *AgentService) GetAgentSession(ctx context.Context, userID, agentID, sessionID string, options ...dao.ConversationHistoryOptions) (*entity.API4Conversation, common.ErrorCode, error) { if sessionID == "" { return nil, common.CodeArgumentError, fmt.Errorf("session_id is required") } @@ -386,13 +332,17 @@ func (s *AgentService) GetAgentSession(ctx context.Context, userID, agentID, ses return nil, common.CodeOperatingError, errors.New("agent not found or no permission") } - data, err := s.api4ConversationDAO.GetBySessionID(ctx, dao.DB, sessionID, agentID) + data, err := s.api4ConversationDAO.GetBySessionID(ctx, dao.DB, sessionID, agentID, sessionHistoryOptions(options...)) if err != nil { return nil, common.CodeServerError, fmt.Errorf("failed to fetch session: %w", err) } if data == nil { return nil, common.CodeNotFound, fmt.Errorf("agent session not found") } + data.Message, err = json.Marshal(parseAgentSessionMessages(data.Message)) + if err != nil { + return nil, common.CodeServerError, fmt.Errorf("format session messages: %w", err) + } return data, common.CodeSuccess, nil } @@ -409,7 +359,7 @@ func (s *AgentService) DeleteAgentSessionItem(ctx context.Context, userID, agent return false, common.CodeServerError, fmt.Errorf("failed to check agent access: %w", err) } - conv, cerr := s.api4ConversationDAO.GetBySessionID(ctx, dao.DB, sessionID, agentID) + conv, cerr := s.api4ConversationDAO.GetMetadataBySessionID(ctx, dao.DB, sessionID, agentID) if cerr != nil { return false, common.CodeServerError, cerr } @@ -467,7 +417,7 @@ func (s *AgentService) DeleteAgentSessions(ctx context.Context, userID, agentID continue } - conv, err := s.api4ConversationDAO.GetBySessionID(ctx, dao.DB, sessionID, agentID) + conv, err := s.api4ConversationDAO.GetMetadataBySessionID(ctx, dao.DB, sessionID, agentID) if err != nil { return nil, common.CodeServerError, err } @@ -651,7 +601,6 @@ type CreateAgentSessionRequest struct { // - dialog_id : agent canvas id // - user_id : caller's id // - message : ordered rows in api_4_conversation_message -// - reference : ordered rows in api_4_conversation_reference // - dsl : JSON map; copied from user_canvas.dsl if the // caller did not pass one // - create_time : unix-millis @@ -686,7 +635,6 @@ func (s *AgentService) CreateAgentSession(ctx context.Context, req *CreateAgentS if len(messages) == 0 { messages = json.RawMessage(`[]`) } - reference := json.RawMessage(`[]`) var dsl entity.JSONMap if len(req.DSL) > 0 { @@ -730,7 +678,6 @@ func (s *AgentService) CreateAgentSession(ctx context.Context, req *CreateAgentS UserID: req.UserID, ExpUserID: &req.UserID, Message: messages, - Reference: reference, Source: sourcePtr, DSL: dsl, } diff --git a/internal/service/agent_test.go b/internal/service/agent_test.go index d332c93f78..77fd4e1064 100644 --- a/internal/service/agent_test.go +++ b/internal/service/agent_test.go @@ -948,11 +948,10 @@ func createAgentSessionTestConversation(t *testing.T, id, agentID, userID string t.Helper() updateDate := time.UnixMilli(updateTime) if err := dao.NewAPI4ConversationDAO().Create(t.Context(), dao.DB, &entity.API4Conversation{ - ID: id, - DialogID: agentID, - UserID: userID, - Message: json.RawMessage(`[{"role":"assistant","content":"hello","prompt":"hidden"}]`), - Reference: json.RawMessage(`[]`), + ID: id, + DialogID: agentID, + UserID: userID, + Message: json.RawMessage(`[{"role":"assistant","content":"hello","prompt":"hidden"}]`), BaseModel: entity.BaseModel{ CreateTime: ptr(updateTime), CreateDate: &updateDate, @@ -1061,7 +1060,15 @@ func TestGetAgentSessionServiceSuccess(t *testing.T) { createAgentSessionTestConversation(t, "session-1", "canvas-1", "user-1", 1000) ctx := serviceTestContext(t, "user-1", "user-1") - session, code, err := NewAgentService().GetAgentSession(ctx, "user-1", "canvas-1", "session-1") + if err := dao.NewAPI4ConversationDAO().UpdateHistory(ctx, dao.DB, "session-1", "canvas-1", "user-1", nil, dao.ConversationHistoryUpdate{ + Message: map[string]interface{}{"role": "assistant", "content": "answer", "reference": map[string]interface{}{ + "chunks": []interface{}{map[string]interface{}{"chunk_id": "chunk", "doc_id": "document", "content_with_weight": "evidence"}}, + "doc_aggs": []interface{}{}, + }}, + }); err != nil { + t.Fatal(err) + } + session, code, err := NewAgentService().GetAgentSession(ctx, "user-1", "canvas-1", "session-1", dao.ConversationHistoryOptions{IncludeReference: true}) if err != nil { t.Fatalf("GetAgentSession failed: %v", err) } @@ -1077,6 +1084,27 @@ func TestGetAgentSessionServiceSuccess(t *testing.T) { if session.DialogID != "canvas-1" { t.Fatalf("expected dialog_id canvas-1, got %s", session.DialogID) } + messages := parseMessages(session.Message) + reference, ok := messages[1]["reference"].(map[string]interface{}) + if !ok { + t.Fatalf("answer reference is not an object: %v", messages[1]) + } + chunks := reference["chunks"].([]interface{}) + if len(chunks) != 1 || chunks[0].(map[string]interface{})["document_id"] != "document" || chunks[0].(map[string]interface{})["id"] != "chunk" { + t.Fatalf("unexpected reference chunks: %v", chunks) + } + body, err := json.Marshal(session) + if err != nil { + t.Fatal(err) + } + var payload map[string]interface{} + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatal(err) + } + if _, ok := payload["reference"]; ok { + t.Fatalf("session contains reference array: %s", body) + } + } // TestCreateAgentSessionLeavesVersionTitleNilWithVersions pins the diff --git a/internal/service/bot_completion.go b/internal/service/bot_completion.go index a686111c23..8e013a5085 100644 --- a/internal/service/bot_completion.go +++ b/internal/service/bot_completion.go @@ -604,7 +604,14 @@ func buildChatbotPipelineMessages(raw json.RawMessage, question, messageID strin if role == "assistant" && len(msg) == 0 { continue } - msg = append(msg, m) + message := make(map[string]interface{}, len(m)) + for key, value := range m { + if key == "prompt" || key == "usage" || key == "elapsed_ms" || key == "reference" { + continue + } + message[key] = value + } + msg = append(msg, message) } return msg } @@ -638,7 +645,8 @@ func (s *BotService) persistChatbotTurn( reference = map[string]any{"chunks": []any{}, "doc_aggs": []any{}} } message := map[string]interface{}{"role": "assistant", "content": answer, "id": messageID, "created_at": completedAt} - return s.api4ConversationDAO.UpdateHistory(ctx, dao.DB, session.ID, session.DialogID, session.UserID, nil, dao.ConversationHistoryUpdate{Message: message, QuestionID: messageID, Reference: reference, AppendReference: true}) + message["reference"] = reference + return s.api4ConversationDAO.UpdateHistory(ctx, dao.DB, session.ID, session.DialogID, session.UserID, nil, dao.ConversationHistoryUpdate{Message: message, QuestionID: messageID}) } // normalizeBotBoolFlag coerces the JSON-encoded reasoning / internet diff --git a/internal/service/bot_completion_history_test.go b/internal/service/bot_completion_history_test.go index 032d679dd4..e141b4ade9 100644 --- a/internal/service/bot_completion_history_test.go +++ b/internal/service/bot_completion_history_test.go @@ -66,7 +66,7 @@ func TestBuildChatbotPipelineMessages_DropsLeadingAssistantAndSystem(t *testing. {"role": "assistant", "content": "Hello, how can I help?", "created_at": 1}, {"role": "system", "content": "hidden", "created_at": 2}, {"role": "user", "content": "What is Go?", "created_at": 3}, - {"role": "assistant", "content": "Go is a compiled language.", "created_at": 4}, + {"role": "assistant", "content": "Go is a compiled language.", "created_at": 4, "reference": map[string]any{"chunks": []any{"large chunk"}}, "prompt": "system", "usage": map[string]any{"total_tokens": 10}, "elapsed_ms": 100}, } raw, err := json.Marshal(turns) if err != nil { @@ -82,6 +82,11 @@ func TestBuildChatbotPipelineMessages_DropsLeadingAssistantAndSystem(t *testing. if msgs[1]["role"] != "assistant" || msgs[1]["content"] != "Go is a compiled language." { t.Errorf("turn 1: %+v", msgs[1]) } + for _, key := range []string{"prompt", "usage", "elapsed_ms", "reference"} { + if _, ok := msgs[1][key]; ok { + t.Fatalf("LLM history contains %s: %v", key, msgs[1]) + } + } if msgs[2]["role"] != "user" || msgs[2]["content"] != "and Rust?" || msgs[2]["id"] != "msg-2" { t.Errorf("turn 2: %+v", msgs[2]) } @@ -518,7 +523,7 @@ func TestPersistChatbotTurn_AppendsPairAndReference(t *testing.T) { if err := svc.persistChatbotQuestion(ctx, sess, "What is Go?", "msg-p1", 2); err != nil { t.Fatalf("persist question: %v", err) } - questionRow, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, dao.DB, "sess-p1", "dlg-p1") + questionRow, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, dao.DB, "sess-p1", "dlg-p1", dao.ConversationHistoryOptions{IncludeReference: true}) if err != nil || questionRow == nil { t.Fatalf("read pending question: %v", err) } @@ -529,7 +534,7 @@ func TestPersistChatbotTurn_AppendsPairAndReference(t *testing.T) { t.Fatalf("persistChatbotTurn: %v", err) } - row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, dao.DB, "sess-p1", "dlg-p1") + row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, dao.DB, "sess-p1", "dlg-p1", dao.ConversationHistoryOptions{IncludeReference: true}) if err != nil || row == nil { t.Fatalf("re-read session: row=%v err=%v", row, err) } @@ -550,8 +555,10 @@ func TestPersistChatbotTurn_AppendsPairAndReference(t *testing.T) { } var refs []map[string]any - if err := json.Unmarshal(row.Reference, &refs); err != nil { - t.Fatalf("reference decode: %v", err) + for _, message := range parseChatbotTurns(row.Message) { + if reference, ok := message["reference"].(map[string]any); ok { + refs = append(refs, reference) + } } if len(refs) != 1 { t.Fatalf("want 1 reference entry, got %d: %+v", len(refs), refs) @@ -564,7 +571,7 @@ func TestPersistChatbotTurn_AppendsPairAndReference(t *testing.T) { // TestPersistChatbotTurn_NilReferenceDefaultsToEmpty ensures a turn // without retrieval results still records an empty reference object, -// keeping the reference list index-aligned with the turn pairs. +// attached directly to its assistant message. func TestPersistChatbotTurn_NilReferenceDefaultsToEmpty(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) @@ -583,13 +590,15 @@ func TestPersistChatbotTurn_NilReferenceDefaultsToEmpty(t *testing.T) { t.Fatalf("persistChatbotTurn: %v", err) } - row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, dao.DB, "sess-p2", "dlg-p1") + row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, dao.DB, "sess-p2", "dlg-p1", dao.ConversationHistoryOptions{IncludeReference: true}) if err != nil || row == nil { t.Fatalf("re-read session: row=%v err=%v", row, err) } var refs []map[string]any - if err := json.Unmarshal(row.Reference, &refs); err != nil { - t.Fatalf("reference decode: %v", err) + for _, message := range parseChatbotTurns(row.Message) { + if reference, ok := message["reference"].(map[string]any); ok { + refs = append(refs, reference) + } } if len(refs) != 1 { t.Fatalf("want 1 reference entry, got %+v", refs) @@ -643,7 +652,7 @@ func TestPersistChatbotTurn_ConcurrentSameSession(t *testing.T) { } wg.Wait() - row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, dao.DB, "sess-p3", "dlg-p1") + row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, dao.DB, "sess-p3", "dlg-p1", dao.ConversationHistoryOptions{IncludeReference: true}) if err != nil || row == nil { t.Fatalf("re-read session: row=%v err=%v", row, err) } @@ -651,8 +660,10 @@ func TestPersistChatbotTurn_ConcurrentSameSession(t *testing.T) { t.Fatalf("want both exchanges persisted (4 turns), got %d: %+v", len(turns), turns) } var refs []map[string]any - if err = json.Unmarshal(row.Reference, &refs); err != nil { - t.Fatalf("reference decode: %v", err) + for _, message := range parseChatbotTurns(row.Message) { + if reference, ok := message["reference"].(map[string]any); ok { + refs = append(refs, reference) + } } if len(refs) != 2 { t.Fatalf("want 2 reference entries, got %d: %+v", len(refs), refs) diff --git a/internal/service/chat_session.go b/internal/service/chat_session.go index 453796f6c3..2215f424cd 100644 --- a/internal/service/chat_session.go +++ b/internal/service/chat_session.go @@ -44,12 +44,12 @@ import ( // Interfaces for testability — satisfied by the concrete DAO/pipeline types. type chatSessionStore interface { - GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.ChatSession, error) - GetBySessionIDAndChatID(ctx context.Context, db *gorm.DB, sessionID, chatID string) (*entity.ChatSession, error) + GetByID(ctx context.Context, db *gorm.DB, id string, options ...dao.ConversationHistoryOptions) (*entity.ChatSession, error) + GetBySessionIDAndChatID(ctx context.Context, db *gorm.DB, sessionID, chatID string, options ...dao.ConversationHistoryOptions) (*entity.ChatSession, error) Create(ctx context.Context, db *gorm.DB, conv *entity.ChatSession) error UpdateByID(ctx context.Context, db *gorm.DB, id string, updates map[string]interface{}) error DeleteByID(ctx context.Context, db *gorm.DB, id string) error - ListByResourceIDs(ctx context.Context, db *gorm.DB, chatID string, sessionIDs []string, sessionID, name string, terms []dao.OrderTerm, page, pageSize int, includeHistory ...bool) ([]*entity.ChatSession, int64, error) + ListByResourceIDs(ctx context.Context, db *gorm.DB, chatID string, sessionIDs []string, sessionID, name string, terms []dao.OrderTerm, page, pageSize int, options ...dao.ConversationHistoryOptions) ([]*entity.ChatSession, int64, error) GetDialogByID(ctx context.Context, db *gorm.DB, chatID string) (*entity.Chat, error) } @@ -121,7 +121,6 @@ type ChatSessionPayload struct { ChatID string `json:"chat_id"` Name *string `json:"name,omitempty"` Messages []map[string]interface{} `json:"messages"` - Reference []interface{} `json:"reference"` UserID *string `json:"user_id,omitempty"` Avatar *string `json:"avatar,omitempty"` CreateDate *time.Time `json:"create_date,omitempty"` @@ -131,14 +130,14 @@ type ChatSessionPayload struct { } // ListChatSessions lists chat sessions for a dialog -func (s *ChatSessionService) ListChatSessions(ctx context.Context, userID, chatID, sessionID, name string, terms []dao.OrderTerm, page, pageSize int, includeHistory ...bool) (*ListChatSessionsResponse, error) { +func (s *ChatSessionService) ListChatSessions(ctx context.Context, userID, chatID, sessionID, name string, terms []dao.OrderTerm, page, pageSize int, options ...dao.ConversationHistoryOptions) (*ListChatSessionsResponse, error) { if err := checkChatPermission(ctx, s.permissionChecker, userID, chatID, permission.OperationRead); err != nil { _, normalizedErr := normalizeChatPermissionError(err) return nil, normalizedErr } // List chat sessions - sessions, total, err := s.chatSessionDAO.ListByResourceIDs(ctx, dao.DB, chatID, nil, sessionID, name, terms, page, pageSize, includeHistory...) + sessions, total, err := s.chatSessionDAO.ListByResourceIDs(ctx, dao.DB, chatID, nil, sessionID, name, terms, page, pageSize, sessionHistoryOptions(options...)) if err != nil { return nil, err } @@ -174,13 +173,13 @@ func (s *ChatSessionService) ListChatSessions(ctx context.Context, userID, chatI } // GetSession returns one chat session after ownership validation. -func (s *ChatSessionService) GetSession(ctx context.Context, userID, chatID, sessionID string) (*ChatSessionPayload, common.ErrorCode, error) { +func (s *ChatSessionService) GetSession(ctx context.Context, userID, chatID, sessionID string, options ...dao.ConversationHistoryOptions) (*ChatSessionPayload, common.ErrorCode, error) { if err := checkChatPermission(ctx, s.permissionChecker, userID, chatID, permission.OperationRead); err != nil { code, normalizedErr := normalizeChatPermissionError(err) return nil, code, normalizedErr } - session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID) + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID, sessionHistoryOptions(options...)) if err != nil { if isChatSessionNotFound(err) { return nil, common.CodeDataError, errors.New("session not found") @@ -200,7 +199,7 @@ func (s *ChatSessionService) GetSession(ctx context.Context, userID, chatID, ses return nil, common.CodeServerError, err } - return s.buildSessionPayload(session, dialog, true), common.CodeSuccess, nil + return s.buildSessionPayload(session, dialog, true, options...), common.CodeSuccess, nil } // CreateSession create a session in a dialog @@ -245,22 +244,19 @@ func (s *ChatSessionService) CreateSession(ctx context.Context, userID, chatID s }, }) - referenceJSON, _ := json.Marshal([]interface{}{}) - conv := &entity.ChatSession{ - ID: common.GenerateUUID(), - DialogID: chatID, - Name: &name, - Message: messagesJSON, - UserID: &userID, - Reference: referenceJSON, + ID: common.GenerateUUID(), + DialogID: chatID, + Name: &name, + Message: messagesJSON, + UserID: &userID, } if err = s.chatSessionDAO.Create(ctx, dao.DB, conv); err != nil { return nil, common.CodeDataError, errors.New("fail to create a session") } - session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, conv.ID) + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, conv.ID, sessionHistoryOptions()) if err != nil { return nil, common.CodeDataError, errors.New("fail to create a session") } @@ -291,7 +287,7 @@ func (s *ChatSessionService) DeleteSessions(ctx context.Context, userID, chatID if len(sessionIDs) == 0 { return map[string]interface{}{}, "success", common.CodeSuccess, nil } - sessions, _, err := s.chatSessionDAO.ListByResourceIDs(ctx, dao.DB, chatID, sessionIDs, "", "", []dao.OrderTerm{{Column: "create_time", Desc: true}}, 0, -1, false) + sessions, _, err := s.chatSessionDAO.ListByResourceIDs(ctx, dao.DB, chatID, sessionIDs, "", "", []dao.OrderTerm{{Column: "create_time", Desc: true}}, 0, -1, dao.ConversationHistoryOptions{NoHistory: true}) if err != nil { return nil, "", common.CodeServerError, err } @@ -500,7 +496,7 @@ func (s *ChatSessionService) UpdateSession(ctx context.Context, userID, chatID, } } - session, err = s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID) + session, err = s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID, sessionHistoryOptions()) if err != nil { if isChatSessionNotFound(err) { return nil, common.CodeDataError, errors.New("fail to update a session") @@ -517,7 +513,7 @@ func (s *ChatSessionService) DeleteSessionMessage(ctx context.Context, userID, c return nil, code, normalizedErr } - session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID) + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID, sessionHistoryOptions()) if err != nil || session.DialogID != chatID { if err != nil && !isChatSessionNotFound(err) { return nil, common.CodeServerError, err @@ -530,17 +526,10 @@ func (s *ChatSessionService) DeleteSessionMessage(ctx context.Context, userID, c return nil, code, normalizedErr } - // parseMessages / parseReferenceList return nil for - // malformed input, so the existing `nil` guards below are - // the single source of truth for "this blob is corrupt". messages := parseMessages(session.Message) if len(session.Message) > 0 && messages == nil { return nil, common.CodeDataError, errors.New("invalid session messages") } - references := parseReferenceList(session.Reference) - if len(session.Reference) > 0 && references == nil { - return nil, common.CodeDataError, errors.New("invalid session reference") - } found := false for i, msg := range messages { if msgID != stringValue(msg["id"]) { @@ -551,10 +540,6 @@ func (s *ChatSessionService) DeleteSessionMessage(ctx context.Context, userID, c } end := i + 1 if end < len(messages) && stringValue(messages[end]["role"]) == "assistant" && stringValue(messages[end]["id"]) == msgID { - refIndex := sessionMessageReferenceIndex(messages, end) - if refIndex >= 0 && refIndex < len(references) { - references = append(references[:refIndex], references[refIndex+1:]...) - } end++ } messages = append(messages[:i], messages[end:]...) @@ -569,10 +554,6 @@ func (s *ChatSessionService) DeleteSessionMessage(ctx context.Context, userID, c if err != nil { return nil, common.CodeServerError, err } - referenceRaw, err := json.Marshal(references) - if err != nil { - return nil, common.CodeServerError, err - } if err = s.chatSessionDAO.UpdateByID(ctx, dao.DB, session.ID, map[string]interface{}{ "history_update": dao.ConversationHistoryUpdate{DeleteMessageID: msgID}, }); err != nil { @@ -582,7 +563,6 @@ func (s *ChatSessionService) DeleteSessionMessage(ctx context.Context, userID, c return nil, common.CodeServerError, err } session.Message = messageRaw - session.Reference = referenceRaw return s.buildSessionPayload(session, nil, false), common.CodeSuccess, nil } @@ -593,7 +573,9 @@ func (s *ChatSessionService) UpdateMessageFeedback(ctx context.Context, userID, return nil, code, normalizedErr } - session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID) + historyOptions := sessionHistoryOptions() + historyOptions.IncludeReference = true + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID, historyOptions) if err != nil || session.DialogID != chatID { if err != nil && !isChatSessionNotFound(err) { return nil, common.CodeServerError, err @@ -623,13 +605,6 @@ func (s *ChatSessionService) UpdateMessageFeedback(ctx context.Context, userID, if len(session.Message) > 0 && messages == nil { return nil, common.CodeDataError, errors.New("invalid session messages") } - // References are only used later in this function but a - // malformed blob must surface immediately, not silently - // collapse to an empty slice. - references := parseReferenceList(session.Reference) - if len(session.Reference) > 0 && references == nil { - return nil, common.CodeDataError, errors.New("invalid session reference") - } messageIndex := -1 var priorThumb interface{} applyChunkFeedback := false @@ -657,12 +632,7 @@ func (s *ChatSessionService) UpdateMessageFeedback(ctx context.Context, userID, } if messageIndex != -1 && applyChunkFeedback { - refIndex := sessionMessageReferenceIndex(messages, messageIndex) - if refIndex >= 0 && refIndex < len(references) { - if reference, ok := references[refIndex].(map[string]interface{}); ok && len(reference) > 0 { - feedbackReference = reference - } - } + feedbackReference, _ = messages[messageIndex]["reference"].(map[string]interface{}) } messageRaw, err := json.Marshal(messages) @@ -706,28 +676,6 @@ func (s *ChatSessionService) UpdateMessageFeedback(ctx context.Context, userID, return s.buildSessionPayload(session, nil, false), common.CodeSuccess, nil } -// sessionMessageReferenceIndex counts assistant responses after the first user -// message, excluding the prologue and unanswered questions. -func sessionMessageReferenceIndex(messages []map[string]interface{}, messageIndex int) int { - refIndex := -1 - hasUser := false - for i, msg := range messages { - role := stringValue(msg["role"]) - if role == "user" { - hasUser = true - } else if role == "assistant" && hasUser { - refIndex++ - } - if i == messageIndex { - if role == "assistant" && hasUser { - return refIndex - } - return -1 - } - } - return -1 -} - const ( upvoteWeightIncrement = 1 downvoteWeightDecrement = 1 @@ -1028,7 +976,17 @@ func floatValue(value interface{}) (float64, bool) { } } -func (s *ChatSessionService) buildSessionPayload(session *entity.ChatSession, dialog *entity.Chat, includeAvatar bool) *ChatSessionPayload { +func sessionHistoryOptions(options ...dao.ConversationHistoryOptions) dao.ConversationHistoryOptions { + var result dao.ConversationHistoryOptions + if len(options) > 0 { + result = options[0] + } + result.IncludeUsage = true + result.IncludeElapsedMs = true + return result +} + +func (s *ChatSessionService) buildSessionPayload(session *entity.ChatSession, dialog *entity.Chat, includeAvatar bool, options ...dao.ConversationHistoryOptions) *ChatSessionPayload { var avatar *string if includeAvatar { value := "" @@ -1038,22 +996,24 @@ func (s *ChatSessionService) buildSessionPayload(session *entity.ChatSession, di avatar = &value } - references := parseReferenceList(session.Reference) - for index, ref := range references { - refMap, ok := ref.(map[string]interface{}) - if !ok { - continue + historyOptions := sessionHistoryOptions(options...) + messages := parseMessages(session.Message) + for _, message := range messages { + if !historyOptions.IncludePrompt { + delete(message, "prompt") + } + if !historyOptions.IncludeReference { + delete(message, "reference") + } else if reference, ok := message["reference"].(map[string]interface{}); ok { + reference["chunks"] = formatReferenceChunks(reference) } - refMap["chunks"] = formatReferenceChunks(refMap) - references[index] = refMap } return &ChatSessionPayload{ ID: session.ID, ChatID: session.DialogID, Name: session.Name, - Messages: parseMessages(session.Message), - Reference: references, + Messages: messages, UserID: session.UserID, Avatar: avatar, CreateDate: session.CreateDate, @@ -1118,23 +1078,6 @@ func parseMessages(raw json.RawMessage) []map[string]interface{} { return messages } -// parseReferenceList decodes a session.Reference blob. Same -// nil-on-malformed contract as parseMessages — callers gate on -// `if len(raw) > 0 && references == nil` to reject corruption. -func parseReferenceList(raw json.RawMessage) []interface{} { - references := make([]interface{}, 0) - if len(raw) == 0 { - return references - } - if err := json.Unmarshal(raw, &references); err != nil { - return nil - } - if references == nil { - return make([]interface{}, 0) - } - return references -} - func formatReferenceChunks(reference map[string]interface{}) []FormattedChunk { rawChunks, ok := reference["chunks"].([]interface{}) if !ok { @@ -1262,9 +1205,9 @@ func (s *ChatSessionService) ChatCompletions( // --- 4. Initialize reference --- var reference []interface{} if session != nil { - reference = s.initializeReference(session) + reference = s.initializeReference() } else if !storeHistoryMessages { - reference = s.initializeReference(&entity.ChatSession{}) + reference = s.initializeReference() } // --- 5. LLM override --- @@ -1602,7 +1545,7 @@ func (s *ChatSessionService) normalizeCompletionMessages( for _, msg := range messages { message := make(map[string]interface{}, len(msg)) for key, value := range msg { - if key == "prompt" || key == "usage" || key == "elapsed_ms" { + if key == "prompt" || key == "usage" || key == "elapsed_ms" || key == "reference" { continue } message[key] = value @@ -1649,14 +1592,12 @@ func (s *ChatSessionService) createSessionForCompletion(ctx context.Context, cha {"role": "assistant", "content": prologue, "created_at": float64(time.Now().Unix())}, } msgJSON, _ := json.Marshal(msg) - refJSON, _ := json.Marshal([]interface{}{}) session := &entity.ChatSession{ - DialogID: chatID, - Name: &name, - Message: msgJSON, - UserID: &userID, - Reference: refJSON, + DialogID: chatID, + Name: &name, + Message: msgJSON, + UserID: &userID, } session.ID = common.GenerateUUID() if err := s.chatSessionDAO.Create(ctx, dao.DB, session); err != nil { @@ -1679,17 +1620,15 @@ func (s *ChatSessionService) GetOrCreateForChannel(ctx context.Context, dialogID } messagesJSON, _ := json.Marshal([]map[string]interface{}{}) - referenceJSON, _ := json.Marshal([]interface{}{}) name := channelSessionName(firstQuery, channelName) if name == "" { name = fmt.Sprintf("channel:%s:%s", channelID, chatID) } session = &entity.ChatSession{ - ID: sessionID, - DialogID: dialogID, - Name: &name, - Message: messagesJSON, - Reference: referenceJSON, + ID: sessionID, + DialogID: dialogID, + Name: &name, + Message: messagesJSON, } if err = s.chatSessionDAO.Create(ctx, dao.DB, session); err != nil { session, rereadErr := s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID) @@ -1775,7 +1714,7 @@ func (s *ChatSessionService) filterSystemAndLeadingAssistant(session *entity.Cha } message := make(map[string]interface{}, len(msg)) for key, value := range msg { - if key == "prompt" || key == "usage" || key == "elapsed_ms" { + if key == "prompt" || key == "usage" || key == "elapsed_ms" || key == "reference" { continue } message[key] = value @@ -1838,23 +1777,11 @@ func (s *ChatSessionService) sendSSEError(streamChan chan<- string, errMsg strin // Helper methods -func (s *ChatSessionService) initializeReference(session *entity.ChatSession) []interface{} { - var reference []interface{} - if len(session.Reference) > 0 { - json.Unmarshal(session.Reference, &reference) - } - // Filter out nil entries and append new reference - var filtered []interface{} - for _, r := range reference { - if r != nil { - filtered = append(filtered, r) - } - } - filtered = append(filtered, map[string]interface{}{ +func (s *ChatSessionService) initializeReference() []interface{} { + return []interface{}{map[string]interface{}{ "chunks": []map[string]interface{}{}, "doc_aggs": []interface{}{}, - }) - return filtered + }} } func (s *ChatSessionService) checkTenantLLMAPIKey(ctx context.Context, tenantID, modelName string) (bool, error) { @@ -2092,32 +2019,24 @@ func (s *ChatSessionService) structureAnswer(session *entity.ChatSession, answer } func (s *ChatSessionService) updateSessionMessages(ctx context.Context, session *entity.ChatSession, messages []map[string]interface{}, reference []interface{}) { + if len(messages) == 0 || stringValue(messages[len(messages)-1]["role"]) != "assistant" { + return + } + message := messages[len(messages)-1] + if len(reference) > 0 { + message["reference"] = reference[len(reference)-1] + } messagesJSON, err := json.Marshal(messages) if err != nil { common.Warn("updateSessionMessages: failed to marshal messages", zap.Error(err)) return } - referenceJSON, err := json.Marshal(reference) - if err != nil { - common.Warn("updateSessionMessages: failed to marshal reference", zap.Error(err)) - return - } - - if len(messages) == 0 || stringValue(messages[len(messages)-1]["role"]) != "assistant" { - return - } - message := messages[len(messages)-1] - var latestReference map[string]interface{} - if len(reference) > 0 { - latestReference, _ = reference[len(reference)-1].(map[string]interface{}) - } - updates := map[string]interface{}{"history_update": dao.ConversationHistoryUpdate{Message: message, QuestionID: stringValue(message["id"]), Reference: latestReference, AppendReference: true}} + updates := map[string]interface{}{"history_update": dao.ConversationHistoryUpdate{Message: message, QuestionID: stringValue(message["id"])}} if err := s.chatSessionDAO.UpdateByID(ctx, dao.DB, session.ID, updates); err != nil { common.Warn("updateSessionMessages: DAO update failed", zap.Error(err)) return } session.Message = messagesJSON - session.Reference = referenceJSON } // structureAnswerWithConv structures the answer with conversation update (like Python's structure_answer) diff --git a/internal/service/chat_session_test.go b/internal/service/chat_session_test.go index 95a5a2d915..a863253b22 100644 --- a/internal/service/chat_session_test.go +++ b/internal/service/chat_session_test.go @@ -52,7 +52,7 @@ func newFakeSessionStore() *fakeSessionStore { } } -func (f *fakeSessionStore) GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.ChatSession, error) { +func (f *fakeSessionStore) GetByID(ctx context.Context, db *gorm.DB, id string, options ...dao.ConversationHistoryOptions) (*entity.ChatSession, error) { if f.getByIDErr != nil { return nil, f.getByIDErr } @@ -64,7 +64,7 @@ func (f *fakeSessionStore) GetByID(ctx context.Context, db *gorm.DB, id string) return &clone, nil } -func (f *fakeSessionStore) GetBySessionIDAndChatID(ctx context.Context, db *gorm.DB, sessionID, chatID string) (*entity.ChatSession, error) { +func (f *fakeSessionStore) GetBySessionIDAndChatID(ctx context.Context, db *gorm.DB, sessionID, chatID string, options ...dao.ConversationHistoryOptions) (*entity.ChatSession, error) { s, err := f.GetByID(ctx, db, sessionID) if err != nil { return nil, err @@ -99,7 +99,6 @@ func (f *fakeSessionStore) UpdateByID(ctx context.Context, db *gorm.DB, id strin } if history, ok := updates["history_update"].(dao.ConversationHistoryUpdate); ok { messages := parseMessages(s.Message) - references := parseReferenceList(s.Reference) if history.DeleteMessageID != "" { for i, message := range messages { if stringValue(message["id"]) != history.DeleteMessageID { @@ -107,10 +106,6 @@ func (f *fakeSessionStore) UpdateByID(ctx context.Context, db *gorm.DB, id strin } end := i + 1 if end < len(messages) && stringValue(messages[end]["role"]) == "assistant" && stringValue(messages[end]["id"]) == history.DeleteMessageID { - refIndex := sessionMessageReferenceIndex(messages, end) - if refIndex >= 0 && refIndex < len(references) { - references = append(references[:refIndex], references[refIndex+1:]...) - } end++ } messages = append(messages[:i], messages[end:]...) @@ -131,12 +126,8 @@ func (f *fakeSessionStore) UpdateByID(ctx context.Context, db *gorm.DB, id strin } } else if history.Message != nil { messages = append(messages, history.Message) - if history.AppendReference { - references = append(references, history.Reference) - } } updates["message"], _ = json.Marshal(messages) - updates["reference"], _ = json.Marshal(references) delete(updates, "history_update") } f.updateCalled = append(f.updateCalled, struct { @@ -153,10 +144,6 @@ func (f *fakeSessionStore) UpdateByID(ctx context.Context, db *gorm.DB, id strin if raw, ok := v.([]byte); ok { s.Message = append(json.RawMessage(nil), raw...) } - case "reference": - if raw, ok := v.([]byte); ok { - s.Reference = append(json.RawMessage(nil), raw...) - } } } return nil @@ -167,7 +154,7 @@ func (f *fakeSessionStore) DeleteByID(ctx context.Context, db *gorm.DB, id strin return nil } -func (f *fakeSessionStore) ListByResourceIDs(_ context.Context, _ *gorm.DB, chatID string, sessionIDs []string, sessionID, name string, _ []dao.OrderTerm, _, pageSize int, _ ...bool) ([]*entity.ChatSession, int64, error) { +func (f *fakeSessionStore) ListByResourceIDs(_ context.Context, _ *gorm.DB, chatID string, sessionIDs []string, sessionID, name string, _ []dao.OrderTerm, _, pageSize int, _ ...dao.ConversationHistoryOptions) ([]*entity.ChatSession, int64, error) { allowedIDs := make(map[string]struct{}, len(sessionIDs)) for _, id := range sessionIDs { allowedIDs[id] = struct{}{} @@ -358,11 +345,11 @@ func TestChatCompletionsPersistsAnswerDetailsForHistory(t *testing.T) { store := newFakeSessionStore() store.sessions["session-1"] = &entity.ChatSession{ ID: "session-1", DialogID: "chat-1", UserID: strPtr("user-1"), - Message: json.RawMessage(`[{"role":"assistant","content":"Welcome"},{"role":"user","content":"old question","id":"old"},{"role":"assistant","content":"old answer","id":"old","prompt":"old prompt","usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3},"elapsed_ms":123}]`), + Message: json.RawMessage(`[{"role":"assistant","content":"Welcome"},{"role":"user","content":"old question","id":"old"},{"role":"assistant","content":"old answer","id":"old","prompt":"old prompt","usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3},"elapsed_ms":123,"reference":{"chunks":[{"id":"old-chunk"}]}}]`), } store.dialogs["chat-1"] = &entity.Chat{ID: "chat-1", TenantID: "user-1", LLMID: "chat@factory"} store.dialogExists["user-1|chat-1"] = true - final := AsyncChatResult{Answer: "new answer", Prompt: "system prompt", Final: true} + final := AsyncChatResult{Answer: "new answer", Prompt: "system prompt", Reference: map[string]interface{}{"chunks": []interface{}{map[string]interface{}{"id": "new-chunk"}}}, Final: true} if mode.agentic { final.Answer = "agent reasoningnew answer" final.Prompt = "" @@ -384,13 +371,13 @@ func TestChatCompletionsPersistsAnswerDetailsForHistory(t *testing.T) { t.Fatalf("unexpected LLM history: %v", pipeline.messages) } for _, message := range pipeline.messages { - for _, key := range []string{"prompt", "usage", "elapsed_ms"} { + for _, key := range []string{"prompt", "usage", "elapsed_ms", "reference"} { if _, ok := message[key]; ok { t.Fatalf("LLM history contains %s: %v", key, message) } } } - payload, code, err := svc.GetSession(t.Context(), "user-1", "chat-1", "session-1") + payload, code, err := svc.GetSession(t.Context(), "user-1", "chat-1", "session-1", dao.ConversationHistoryOptions{IncludeReference: true, IncludePrompt: true}) if err != nil || code != common.CodeSuccess { t.Fatalf("GetSession: code=%v err=%v", code, err) } @@ -401,6 +388,21 @@ func TestChatCompletionsPersistsAnswerDetailsForHistory(t *testing.T) { if answer["content"] != final.Answer { t.Fatalf("missing stored answer details: %v", answer) } + ref, ok := answer["reference"].(map[string]interface{}) + if !ok || len(ref["chunks"].([]FormattedChunk)) != 1 || ref["chunks"].([]FormattedChunk)[0].ID != "new-chunk" { + t.Fatalf("missing stored reference: %v", answer) + } + body, err := json.Marshal(payload) + if err != nil { + t.Fatal(err) + } + var data map[string]interface{} + if err := json.Unmarshal(body, &data); err != nil { + t.Fatal(err) + } + if _, ok := data["reference"]; ok { + t.Fatalf("session contains a reference array: %s", body) + } if mode.agentic { if _, ok := answer["prompt"]; ok { t.Fatalf("agentic answer stored a prompt: %v", answer) @@ -427,7 +429,7 @@ func TestChatCompletionsPersistsAnswerDetailsForHistory(t *testing.T) { func TestNormalizeCompletionMessagesExcludesAnswerDetails(t *testing.T) { messages := []map[string]interface{}{ - {"role": "assistant", "content": "answer", "prompt": "system prompt", "usage": map[string]int{"total_tokens": 3}, "elapsed_ms": 123.0}, + {"role": "assistant", "content": "answer", "prompt": "system prompt", "usage": map[string]int{"total_tokens": 3}, "elapsed_ms": 123.0, "reference": map[string]interface{}{"chunks": []interface{}{}}}, {"role": "user", "content": "question", "id": "turn-1"}, } svc := &ChatSessionService{} @@ -583,7 +585,7 @@ func TestListChatSessions_Success(t *testing.T) { ctx := t.Context() terms := []dao.OrderTerm{{Column: "create_time", Desc: true}} - resp, err := svc.ListChatSessions(ctx, "user-1", "chat-1", "", "", terms, 1, 30, false) + resp, err := svc.ListChatSessions(ctx, "user-1", "chat-1", "", "", terms, 1, 30, dao.ConversationHistoryOptions{NoHistory: true}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -596,7 +598,7 @@ func TestListChatSessions_Success(t *testing.T) { } } for page, wantCount := range map[int]int{4: 10, 5: 1, 6: 0} { - resp, err := svc.ListChatSessions(ctx, "user-1", "chat-1", "", "", terms, page, 10, false) + resp, err := svc.ListChatSessions(ctx, "user-1", "chat-1", "", "", terms, page, 10, dao.ConversationHistoryOptions{NoHistory: true}) if err != nil { t.Fatalf("page %d: %v", page, err) } @@ -604,11 +606,11 @@ func TestListChatSessions_Success(t *testing.T) { t.Fatalf("page %d: got total=%d and %d sessions, want total=41 and %d sessions", page, resp.Total, len(resp.Sessions), wantCount) } } - resp, err = svc.ListChatSessions(ctx, "user-1", "chat-1", "", "", terms, 1, 0, false) + resp, err = svc.ListChatSessions(ctx, "user-1", "chat-1", "", "", terms, 1, 0, dao.ConversationHistoryOptions{NoHistory: true}) if err != nil || resp.Total != 41 || len(resp.Sessions) != 0 { t.Fatalf("zero page size: got response=%+v err=%v, want total=41 and no sessions", resp, err) } - resp, err = svc.ListChatSessions(ctx, "user-1", "chat-1", "session-01", "", terms, 1, 30, false) + resp, err = svc.ListChatSessions(ctx, "user-1", "chat-1", "session-01", "", terms, 1, 30, dao.ConversationHistoryOptions{NoHistory: true}) if err != nil || resp.Total != 1 || len(resp.Sessions) != 1 { t.Fatalf("ID filter: got response=%+v err=%v, want total=1 and one session", resp, err) } @@ -640,12 +642,8 @@ func TestGetSession_Success(t *testing.T) { ID: "session-1", DialogID: "chat-1", Name: strPtr("session"), - Message: json.RawMessage(`[{"role":"assistant","content":"hello"}]`), - Reference: json.RawMessage(`[ - {"chunks":[{"chunk_id":"chunk-1","content_with_weight":"hello","doc_id":"doc-1","docnm_kwd":"Doc 1","kb_id":"kb-1"}]}, - [] - ]`), - UserID: strPtr("user-1"), + Message: json.RawMessage(`[{"role":"assistant","content":"hello","reference":{"chunks":[{"chunk_id":"chunk-1","content_with_weight":"hello","doc_id":"doc-1","docnm_kwd":"Doc 1","kb_id":"kb-1"}]}}]`), + UserID: strPtr("user-1"), } icon := "avatar.png" store.dialogs["chat-1"] = &entity.Chat{ID: "chat-1", Icon: &icon} @@ -658,7 +656,7 @@ func TestGetSession_Success(t *testing.T) { } ctx := t.Context() - resp, code, err := svc.GetSession(ctx, "user-1", "chat-1", "session-1") + resp, code, err := svc.GetSession(ctx, "user-1", "chat-1", "session-1", dao.ConversationHistoryOptions{IncludeReference: true, IncludePrompt: true}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -674,12 +672,9 @@ func TestGetSession_Success(t *testing.T) { if len(resp.Messages) != 1 || resp.Messages[0]["content"] != "hello" { t.Fatalf("messages=%#v", resp.Messages) } - if len(resp.Reference) != 2 { - t.Fatalf("reference len=%d", len(resp.Reference)) - } - firstRef, ok := resp.Reference[0].(map[string]interface{}) + firstRef, ok := resp.Messages[0]["reference"].(map[string]interface{}) if !ok { - t.Fatalf("reference[0] type=%T", resp.Reference[0]) + t.Fatalf("answer reference type=%T", resp.Messages[0]["reference"]) } chunks, ok := firstRef["chunks"].([]FormattedChunk) if !ok { @@ -688,9 +683,6 @@ func TestGetSession_Success(t *testing.T) { if len(chunks) != 1 || chunks[0].ID != "chunk-1" { t.Fatalf("chunks=%#v", chunks) } - if _, ok := resp.Reference[1].([]interface{}); !ok { - t.Fatalf("reference[1] changed unexpectedly: %T", resp.Reference[1]) - } } func TestGetSession_NotOwner(t *testing.T) { @@ -842,21 +834,12 @@ func TestUpdateSession_NotFound(t *testing.T) { func TestDeleteSessionMessage_RemovesMessagePairAndReference(t *testing.T) { store := newFakeSessionStore() + store.dialogs["chat-1"] = &entity.Chat{ID: "chat-1"} store.sessions["session-1"] = &entity.ChatSession{ ID: "session-1", DialogID: "chat-1", UserID: strPtr("user-1"), - Message: json.RawMessage(`[ - {"role":"assistant","content":"Welcome!"}, - {"role":"user","content":"first","id":"msg-1"}, - {"role":"assistant","content":"answer 1","id":"msg-1"}, - {"role":"user","content":"second","id":"msg-2"}, - {"role":"assistant","content":"answer 2","id":"msg-2"} - ]`), - Reference: json.RawMessage(`[ - {"chunks":[{"id":"chunk-1","kb_id":"kb-1"}]}, - {"chunks":[{"id":"chunk-2","kb_id":"kb-2"}]} - ]`), + Message: json.RawMessage(`[{"role":"assistant","content":"Welcome!"},{"role":"user","content":"first","id":"msg-1"},{"role":"assistant","content":"answer 1","id":"msg-1","reference":{"chunks":[{"id":"chunk-1","kb_id":"kb-1"}]}},{"role":"user","content":"second","id":"msg-2"},{"role":"assistant","content":"answer 2","id":"msg-2","reference":{"chunks":[{"id":"chunk-2","kb_id":"kb-2"}]}}]`), } store.dialogExists["tenant-1|chat-1"] = true @@ -880,12 +863,16 @@ func TestDeleteSessionMessage_RemovesMessagePairAndReference(t *testing.T) { if resp.Messages[1]["id"] != "msg-2" || resp.Messages[2]["id"] != "msg-2" { t.Fatalf("remaining pair=%#v", resp.Messages) } - if len(resp.Reference) != 1 { - t.Fatalf("reference=%#v", resp.Reference) + if _, ok := resp.Messages[2]["reference"]; ok { + t.Fatalf("delete response contains reference: %#v", resp.Messages[2]) } - ref, ok := resp.Reference[0].(map[string]interface{}) + resp, code, err = svc.GetSession(ctx, "user-1", "chat-1", "session-1", dao.ConversationHistoryOptions{IncludeReference: true}) + if err != nil || code != common.CodeSuccess { + t.Fatalf("GetSession: code=%v err=%v", code, err) + } + ref, ok := resp.Messages[2]["reference"].(map[string]interface{}) if !ok { - t.Fatalf("reference type=%T", resp.Reference[0]) + t.Fatalf("reference type=%T", resp.Messages[2]["reference"]) } chunks, ok := ref["chunks"].([]FormattedChunk) if !ok || len(chunks) != 1 || chunks[0].ID != "chunk-2" { @@ -902,14 +889,7 @@ func TestUpdateMessageFeedback_AppliesChunkFeedbackWithResolvedTenantAndContext( ID: "session-1", DialogID: "chat-1", UserID: strPtr("user-1"), - Message: json.RawMessage(`[ - {"role":"assistant","content":"Welcome!"}, - {"role":"user","content":"question","id":"msg-1"}, - {"role":"assistant","content":"answer","id":"msg-1","feedback":"old"} - ]`), - Reference: json.RawMessage(`[ - {"chunks":[{"id":"chunk-1","kb_id":"kb-1","similarity":0.9}]} - ]`), + Message: json.RawMessage(`[{"role":"assistant","content":"Welcome!"},{"role":"user","content":"question","id":"msg-1"},{"role":"assistant","content":"answer","id":"msg-1","feedback":"old","reference":{"chunks":[{"id":"chunk-1","kb_id":"kb-1","similarity":0.9}]}}]`), } store.dialogExists["tenant-owner|chat-1"] = true store.dialogs["chat-1"] = &entity.Chat{ID: "chat-1", TenantID: "tenant-owner"} @@ -962,14 +942,7 @@ func TestUpdateMessageFeedback_ToggleUsesResolvedTenantForUndoAndApply(t *testin ID: "session-1", DialogID: "chat-1", UserID: strPtr("user-1"), - Message: json.RawMessage(`[ - {"role":"assistant","content":"Welcome!"}, - {"role":"user","content":"question","id":"msg-1"}, - {"role":"assistant","content":"answer","id":"msg-1","thumbup":true} - ]`), - Reference: json.RawMessage(`[ - {"chunks":[{"chunk_id":"chunk-1","dataset_id":"kb-1"}]} - ]`), + Message: json.RawMessage(`[{"role":"assistant","content":"Welcome!"},{"role":"user","content":"question","id":"msg-1"},{"role":"assistant","content":"answer","id":"msg-1","thumbup":true,"reference":{"chunks":[{"chunk_id":"chunk-1","dataset_id":"kb-1"}]}}]`), } store.dialogExists["tenant-owner|chat-1"] = true store.dialogs["chat-1"] = &entity.Chat{ID: "chat-1", TenantID: "tenant-owner"} @@ -1140,8 +1113,7 @@ func TestChatCompletions_AppendOnly(t *testing.T) { store := newFakeSessionStore() session := &entity.ChatSession{ ID: "session-1", DialogID: "dialog-1", - Message: json.RawMessage(`[{"role":"assistant","content":"Welcome!"}]`), - Reference: json.RawMessage(`[]`), + Message: json.RawMessage(`[{"role":"assistant","content":"Welcome!"}]`), } store.dialogExists["user-1|dialog-1"] = true store.sessions["session-1"] = session @@ -1207,7 +1179,7 @@ func TestChatCompletions_AppendOnly(t *testing.T) { t.Run("non-storing model tests", func(t *testing.T) { stored := store.sessions["session-1"] - storedMessages, storedReference := string(stored.Message), string(stored.Reference) + storedMessages := string(stored.Message) payload := []map[string]interface{}{ {"role": "system", "content": "client system prompt"}, {"role": "user", "content": "earlier question"}, @@ -1254,7 +1226,7 @@ func TestChatCompletions_AppendOnly(t *testing.T) { } wg.Wait() stored = store.sessions["session-1"] - if len(store.updateCalled) != 2 || len(store.createCalled) != 0 || string(stored.Message) != storedMessages || string(stored.Reference) != storedReference { + if len(store.updateCalled) != 2 || len(store.createCalled) != 0 || string(stored.Message) != storedMessages { t.Fatal("model tests changed stored history or created a session") } if _, err := svc.ChatCompletions(t.Context(), "user-1", "dialog-1", "", nil, "question without messages", nil, "", nil, map[string]interface{}{"store_history_messages": false}, false, false, nil); err == nil { @@ -1266,11 +1238,10 @@ func TestChatCompletions_AppendOnly(t *testing.T) { func TestChatCompletionsPassesRequestUserIDToPipeline(t *testing.T) { store := newFakeSessionStore() store.sessions["session-1"] = &entity.ChatSession{ - ID: "session-1", - DialogID: "dialog-1", - UserID: strPtr("user-1"), - Message: json.RawMessage(`[{"role":"assistant","content":"Welcome!"}]`), - Reference: json.RawMessage(`[]`), + ID: "session-1", + DialogID: "dialog-1", + UserID: strPtr("user-1"), + Message: json.RawMessage(`[{"role":"assistant","content":"Welcome!"}]`), } store.dialogs["dialog-1"] = &entity.Chat{ ID: "dialog-1", @@ -1324,11 +1295,10 @@ func TestChatCompletionsPassesRequestUserIDToPipeline(t *testing.T) { func TestChatCompletionsStreamFinalCarriesDecoratedReference(t *testing.T) { store := newFakeSessionStore() store.sessions["session-1"] = &entity.ChatSession{ - ID: "session-1", - DialogID: "dialog-1", - UserID: strPtr("user-1"), - Message: json.RawMessage(`[{"role":"assistant","content":"Welcome!"}]`), - Reference: json.RawMessage(`[]`), + ID: "session-1", + DialogID: "dialog-1", + UserID: strPtr("user-1"), + Message: json.RawMessage(`[{"role":"assistant","content":"Welcome!"}]`), } store.dialogs["dialog-1"] = &entity.Chat{ ID: "dialog-1", @@ -1448,11 +1418,10 @@ func TestChatCompletionsStreamFinalCarriesDecoratedReference(t *testing.T) { func TestChatCompletionsModelIDOverrideUsesModelResolver(t *testing.T) { store := newFakeSessionStore() store.sessions["session-1"] = &entity.ChatSession{ - ID: "session-1", - DialogID: "dialog-1", - UserID: strPtr("user-1"), - Message: json.RawMessage(`[]`), - Reference: json.RawMessage(`[]`), + ID: "session-1", + DialogID: "dialog-1", + UserID: strPtr("user-1"), + Message: json.RawMessage(`[]`), } store.dialogs["dialog-1"] = &entity.Chat{ ID: "dialog-1", @@ -1540,18 +1509,14 @@ func TestParseMessages_LegacyWrappedObject(t *testing.T) { func TestBuildSessionPayload_EmptyCollectionsEncodeAsEmptyArrays(t *testing.T) { svc := &ChatSessionService{} payload := svc.buildSessionPayload(&entity.ChatSession{ - ID: "session-1", - DialogID: "chat-1", - Message: nil, - Reference: json.RawMessage(`null`), + ID: "session-1", + DialogID: "chat-1", + Message: nil, }, nil, false) if payload.Messages == nil { t.Fatal("messages is nil") } - if payload.Reference == nil { - t.Fatal("reference is nil") - } body, err := json.Marshal(payload) if err != nil { @@ -1560,8 +1525,8 @@ func TestBuildSessionPayload_EmptyCollectionsEncodeAsEmptyArrays(t *testing.T) { if !strings.Contains(string(body), `"messages":[]`) { t.Fatalf("messages did not encode as empty array: %s", string(body)) } - if !strings.Contains(string(body), `"reference":[]`) { - t.Fatalf("reference did not encode as empty array: %s", string(body)) + if strings.Contains(string(body), `"reference":`) { + t.Fatalf("unexpected session reference: %s", string(body)) } } @@ -1578,16 +1543,6 @@ func TestParseCollections_ReturnEmptySlicesForMissingOrNull(t *testing.T) { } } - referenceInputs := []json.RawMessage{ - nil, - json.RawMessage(`null`), - } - for _, input := range referenceInputs { - got := parseReferenceList(input) - if got == nil || len(got) != 0 { - t.Fatalf("parseReferenceList(%s)=%#v", string(input), got) - } - } } func TestParseCollections_ReturnNilForMalformedData(t *testing.T) { @@ -1601,15 +1556,6 @@ func TestParseCollections_ReturnNilForMalformedData(t *testing.T) { } } - referenceInputs := []json.RawMessage{ - json.RawMessage(`not-json`), - json.RawMessage(`{"unexpected":[]}`), - } - for _, input := range referenceInputs { - if got := parseReferenceList(input); got != nil { - t.Fatalf("parseReferenceList(%s)=%#v, want nil", string(input), got) - } - } } // =================================================================== diff --git a/web/src/hooks/__tests__/use-session-history-request.test.tsx b/web/src/hooks/__tests__/use-session-history-request.test.tsx new file mode 100644 index 0000000000..c2495e1ce2 --- /dev/null +++ b/web/src/hooks/__tests__/use-session-history-request.test.tsx @@ -0,0 +1,60 @@ +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import { act, renderHook } from '@testing-library/react'; +import { useFetchSessionManually as useAgentSession } from '../use-agent-request'; +import { useFetchSessionManually as useChatSession } from '../use-chat-request'; + +jest.mock('react-router', () => ({ + ...jest.requireActual('react-router'), + useParams: () => ({ id: 'resource' }), +})); +jest.mock('@/services/next-chat-service', () => ({ + __esModule: true, + default: { getSession: jest.fn() }, +})); +jest.mock('@/services/agent-service', () => ({ + __esModule: true, + default: {}, + fetchAgentLogsById: jest.fn(), +})); + +const chatService = require('@/services/next-chat-service').default; +const agentService = require('@/services/agent-service'); + +const wrapper = ({ children }: { children: React.ReactNode }) => ( + + {children} + +); + +it('requests references and prompts for chat history displayed in the UI', async () => { + chatService.getSession.mockResolvedValue({ + data: { data: { messages: [] } }, + }); + const { result } = renderHook(() => useChatSession(), { wrapper }); + + await act(async () => { + await result.current.fetchSessionManually('session'); + }); + + expect(chatService.getSession).toHaveBeenCalledWith( + expect.objectContaining({ + params: { include_reference: true, include_prompt: true }, + }), + true, + ); +}); + +it('requests references and prompts for agent history displayed in the UI', async () => { + agentService.fetchAgentLogsById.mockResolvedValue({ data: { data: {} } }); + const { result } = renderHook(() => useAgentSession(), { wrapper }); + + await act(async () => { + await result.current.fetchSessionManually('session'); + }); + + expect(agentService.fetchAgentLogsById).toHaveBeenCalledWith( + 'resource', + 'session', + { include_reference: true, include_prompt: true }, + ); +}); diff --git a/web/src/hooks/use-agent-request.ts b/web/src/hooks/use-agent-request.ts index 6363ba3c9e..b1d112d35f 100644 --- a/web/src/hooks/use-agent-request.ts +++ b/web/src/hooks/use-agent-request.ts @@ -1002,6 +1002,7 @@ export const useFetchAgentLog = (searchParams: IAgentLogsRequest) => { queryFn: async () => { const { data } = await fetchAgentLogsByCanvasId(id as string, { ...searchParams, + include_reference: true, }); return { total: data?.total ?? 0, sessions: data?.data ?? [] }; @@ -1027,6 +1028,7 @@ export const useFetchSessionsByCanvasId = () => { fetchPage: async (params) => { const { data } = await fetchAgentLogsByCanvasId(canvasId!, { ...params, + include_history: false, exp_user_id: tenantInfo.tenant_id, }); @@ -1439,7 +1441,10 @@ export function useFetchSessionManually() { return null; } - const { data } = await fetchAgentLogsById(canvasId, sessionId); + const { data } = await fetchAgentLogsById(canvasId, sessionId, { + include_reference: true, + include_prompt: true, + }); return data?.data; }, diff --git a/web/src/hooks/use-chat-request.ts b/web/src/hooks/use-chat-request.ts index a9f5cc92ab..effa1f1f7b 100644 --- a/web/src/hooks/use-chat-request.ts +++ b/web/src/hooks/use-chat-request.ts @@ -368,7 +368,10 @@ export function useFetchSessionManually() { mutationKey: [ChatApiAction.FetchSessionManually], mutationFn: async (sessionId) => { const { data } = await chatService.getSession( - { url: api.getSession(chatId!, sessionId) }, + { + url: api.getSession(chatId!, sessionId), + params: { include_reference: true, include_prompt: true }, + }, true, ); diff --git a/web/src/interfaces/database/agent.ts b/web/src/interfaces/database/agent.ts index d47d8a190a..fd74882a74 100644 --- a/web/src/interfaces/database/agent.ts +++ b/web/src/interfaces/database/agent.ts @@ -307,7 +307,6 @@ export interface IAgentLogResponse { source: string; user_id: string; dsl: string; - reference: IReference[]; name: string; version_title: string; is_new?: boolean; @@ -317,6 +316,9 @@ export interface IAgentLogsResponse { sessions: IAgentLogResponse[]; } export interface IAgentLogsRequest { + include_history?: boolean; + include_reference?: boolean; + include_prompt?: boolean; keywords?: string; to_date?: string | Date; from_date?: string | Date; @@ -330,6 +332,7 @@ export interface IAgentLogsRequest { export interface IAgentLogMessage { content: string; role: 'user' | 'assistant'; + reference?: IReference; id: string; } diff --git a/web/src/interfaces/database/chat.ts b/web/src/interfaces/database/chat.ts index 00d4906896..8b1a5d9a17 100644 --- a/web/src/interfaces/database/chat.ts +++ b/web/src/interfaces/database/chat.ts @@ -131,7 +131,6 @@ export interface IConversation { id: string; avatar: string; messages: Message[]; - reference: IReference[]; name: string; update_date: string; update_time: number; @@ -143,6 +142,7 @@ export interface Message { role: MessageType; doc_ids?: string[]; prompt?: string; + reference?: IReference; id?: string; audio_binary?: string; data?: any; @@ -155,6 +155,7 @@ export interface Message { } export interface IReferenceChunk { + citation_id?: string; id: string; content: null; document_id: string; @@ -254,7 +255,6 @@ export interface IExternalChatInfo { export interface IMessage extends Message { id: string; awaitingServerId?: boolean; - reference?: IReference; // the latest news has reference conversationId?: string; // To distinguish which conversation the message belongs to } diff --git a/web/src/pages/agent/explore/components/session-chat.tsx b/web/src/pages/agent/explore/components/session-chat.tsx index 250699075d..973d445754 100644 --- a/web/src/pages/agent/explore/components/session-chat.tsx +++ b/web/src/pages/agent/explore/components/session-chat.tsx @@ -8,12 +8,12 @@ import { MessageType } from '@/constants/chat'; import { useUploadAgentFileWithProgress } from '@/hooks/use-agent-request'; import { useFetchUserInfo } from '@/hooks/use-user-setting-request'; import { IAgentLogResponse } from '@/interfaces/database/agent'; -import { IMessage, IReferenceObject } from '@/interfaces/database/chat'; +import { IMessage } from '@/interfaces/database/chat'; import DebugContent from '@/pages/agent/debug-content'; import { useAwaitComponentData } from '@/pages/agent/hooks/use-chat-logic'; import { BeginQuery } from '@/pages/agent/interface'; import { ParameterDialog } from '@/pages/agent/share/parameter-dialog'; -import { buildMessageUuidWithRole } from '@/utils/chat'; +import { buildMessageUuidWithRole, buildReferenceObject } from '@/utils/chat'; import { useCallback, useEffect, useRef } from 'react'; import { useTranslation } from 'react-i18next'; import { useExploreUrlParams } from '../hooks/use-explore-url-params'; @@ -181,10 +181,7 @@ export function SessionChat({ session }: SessionChatProps) { avatarDialog={canvasInfo?.avatar || ''} reference={ findReferenceByMessageId(message.id) || - (session?.reference?.[ - Math.floor((i - 1) / 2) - ] as unknown as IReferenceObject) || - {} + buildReferenceObject(message.reference) } clickDocumentButton={clickDocumentButton} index={i} diff --git a/web/src/pages/agent/explore/hooks/use-select-derived-session-list.ts b/web/src/pages/agent/explore/hooks/use-select-derived-session-list.ts index 333553a274..0dc70770a4 100644 --- a/web/src/pages/agent/explore/hooks/use-select-derived-session-list.ts +++ b/web/src/pages/agent/explore/hooks/use-select-derived-session-list.ts @@ -34,7 +34,6 @@ export const useSelectDerivedSessionList = () => { source: '', user_id: '', dsl: '', - reference: [], is_new: true, name: '', version_title: '', diff --git a/web/src/pages/agent/utils/chat.ts b/web/src/pages/agent/utils/chat.ts index 53d8712b33..6aa8da8c7a 100644 --- a/web/src/pages/agent/utils/chat.ts +++ b/web/src/pages/agent/utils/chat.ts @@ -1,24 +1,4 @@ -import { MessageType } from '@/constants/chat'; import { IEventList, MessageEventType } from '@/hooks/use-send-message'; -import { IMessage, IReference } from '@/interfaces/database/chat'; -import { isEmpty } from 'lodash'; - -export const buildAgentMessageItemReference = ( - conversation: { messages: IMessage[]; reference: IReference[] }, - message: IMessage, -) => { - const assistantMessages = conversation.messages?.filter( - (x) => x.role === MessageType.Assistant, - ); - const referenceIndex = assistantMessages.findIndex( - (x) => x.id === message.id, - ); - const reference = !isEmpty(message?.reference) - ? message?.reference - : (conversation?.reference ?? [])[referenceIndex]; - - return reference ?? { doc_aggs: [], chunks: [], total: 0 }; -}; /** * Determines whether the message should be split into two separate entries: diff --git a/web/src/pages/agents/agent-log-detail-modal.tsx b/web/src/pages/agents/agent-log-detail-modal.tsx index bb1ab8a60c..cd2fa5dc20 100644 --- a/web/src/pages/agents/agent-log-detail-modal.tsx +++ b/web/src/pages/agents/agent-log-detail-modal.tsx @@ -3,86 +3,16 @@ import { Modal } from '@/components/ui/modal/modal'; import { useFetchAgent } from '@/hooks/use-agent-request'; import { useFetchUserInfo } from '@/hooks/use-user-setting-request'; import { IAgentLogMessage } from '@/interfaces/database/agent'; -import { - Docagg, - IMessage, - IReferenceChunk, - IReferenceObject, - Message, -} from '@/interfaces/database/chat'; -import { buildMessageUuidWithRole } from '@/utils/chat'; +import { IMessage, Message } from '@/interfaces/database/chat'; +import { buildMessageUuidWithRole, buildReferenceObject } from '@/utils/chat'; import React, { useMemo } from 'react'; interface CustomModalProps { isOpen: boolean; onClose: () => void; message: IAgentLogMessage[]; - reference: IReferenceObject; } -interface IAgentLogReferenceChunk { - citation_id?: string; - content?: string | null; - document_id: string; - document_name: string; - dataset_id: string; - id: string; - image_id: string; - positions?: number[] | number[][]; - similarity?: number; - vector_similarity?: number; - term_similarity?: number; - doc_type?: string; - url?: string; -} - -// Runtime agent-log messages carry the reference chunks array directly, -// even though IAgentLogMessage does not declare it. -type AgentLogMessageWithReference = IAgentLogMessage & { - reference?: IAgentLogReferenceChunk[]; -}; - -const buildReferenceObject = ( - chunks: IAgentLogReferenceChunk[] | undefined | null, -): IReferenceObject => { - if (!chunks || !Array.isArray(chunks) || chunks.length === 0) { - return { chunks: {}, doc_aggs: {} }; - } - - const chunkMap: Record = {}; - const docAggMap: Record = {}; - - chunks.forEach((chunk, index) => { - const key = chunk.citation_id ?? String(index); - - chunkMap[key] = { - id: chunk.id, - content: (chunk.content ?? null) as IReferenceChunk['content'], - document_id: chunk.document_id, - document_name: chunk.document_name, - dataset_id: chunk.dataset_id, - image_id: chunk.image_id, - similarity: chunk.similarity ?? 0, - vector_similarity: chunk.vector_similarity ?? 0, - term_similarity: chunk.term_similarity ?? 0, - positions: (chunk.positions ?? []) as IReferenceChunk['positions'], - doc_type: chunk.doc_type, - }; - - if (!docAggMap[chunk.document_id]) { - docAggMap[chunk.document_id] = { - count: 0, - doc_id: chunk.document_id, - doc_name: chunk.document_name, - url: chunk.url, - }; - } - docAggMap[chunk.document_id].count++; - }); - - return { chunks: chunkMap, doc_aggs: docAggMap }; -}; - export const AgentLogDetailModal: React.FC = ({ isOpen, onClose, @@ -125,7 +55,6 @@ export const AgentLogDetailModal: React.FC = ({
{derivedMessages?.map((message, i) => { - const msg = message as AgentLogMessageWithReference; return ( = ({ avatar={userInfo.avatar} avatarDialog={canvasInfo.avatar} item={message as IMessage} - reference={buildReferenceObject(msg.reference)} + reference={buildReferenceObject(message.reference)} index={i} showLikeButton={false} showLog={false} diff --git a/web/src/pages/agents/agent-log-page.tsx b/web/src/pages/agents/agent-log-page.tsx index a575a13abf..ca122b638d 100644 --- a/web/src/pages/agents/agent-log-page.tsx +++ b/web/src/pages/agents/agent-log-page.tsx @@ -18,7 +18,6 @@ import { IAgentLogMessage, IAgentLogResponse, } from '@/interfaces/database/agent'; -import { IReferenceObject } from '@/interfaces/database/chat'; import { formatDate } from '@/utils/date'; import React, { useCallback, useEffect, useState } from 'react'; import { useTranslation } from 'react-i18next'; @@ -415,7 +414,6 @@ const AgentLogPage: React.FC = () => { setOpenModal(false)} />
diff --git a/web/src/pages/next-chats/chat/chat-box/next-multiple-chat-box.tsx b/web/src/pages/next-chats/chat/chat-box/next-multiple-chat-box.tsx index 4338dbae32..3f53bca58f 100644 --- a/web/src/pages/next-chats/chat/chat-box/next-multiple-chat-box.tsx +++ b/web/src/pages/next-chats/chat/chat-box/next-multiple-chat-box.tsx @@ -94,7 +94,6 @@ const ChatCard = forwardRef(function ChatCard( addChatBox, chatBoxIds, clickDocumentButton, - conversation, value, setValue, files, @@ -120,10 +119,7 @@ const ChatCard = forwardRef(function ChatCard( const { scrollRef } = useScrollToBottom(derivedMessages, messageContainerRef); - const messageReferences = useMessageReferences( - derivedMessages, - conversation.reference, - ); + const messageReferences = useMessageReferences(derivedMessages); const FormSchema = z.object(LlmSettingSchema); diff --git a/web/src/pages/next-chats/chat/chat-box/single-chat-box.test.tsx b/web/src/pages/next-chats/chat/chat-box/single-chat-box.test.tsx index ef1de85bcb..0775a6f38d 100644 --- a/web/src/pages/next-chats/chat/chat-box/single-chat-box.test.tsx +++ b/web/src/pages/next-chats/chat/chat-box/single-chat-box.test.tsx @@ -75,7 +75,6 @@ function buildConversation(id: string) { update_date: '', update_time: 0, is_new: true as const, - reference: [], messages: [ { id: 'old-message', diff --git a/web/src/pages/next-chats/chat/chat-box/single-chat-box.tsx b/web/src/pages/next-chats/chat/chat-box/single-chat-box.tsx index c0ab072d78..3da381f8e6 100644 --- a/web/src/pages/next-chats/chat/chat-box/single-chat-box.tsx +++ b/web/src/pages/next-chats/chat/chat-box/single-chat-box.tsx @@ -51,10 +51,7 @@ export function SingleChatBox({ conversation }: IProps) { const showInternet = useShowInternet(); - const messageReferences = useMessageReferences( - messages, - conversation.reference, - ); + const messageReferences = useMessageReferences(messages); useEffect(() => { // Skip when the conversation prop is stale — its id doesn't match the diff --git a/web/src/pages/next-chats/hooks/use-message-references.test.tsx b/web/src/pages/next-chats/hooks/use-message-references.test.tsx new file mode 100644 index 0000000000..3dffd90f09 --- /dev/null +++ b/web/src/pages/next-chats/hooks/use-message-references.test.tsx @@ -0,0 +1,27 @@ +import { MessageType } from '@/constants/chat'; +import type { IMessage, IReference } from '@/interfaces/database/chat'; +import { renderHook } from '@testing-library/react'; +import { EmptyReference } from '../utils'; +import { useMessageReferences } from './use-message-references'; + +it('resolves each answer reference independently of shared ids or answer order', () => { + const reference: IReference = { chunks: [], doc_aggs: [], total: 1 }; + const messages: IMessage[] = [ + { id: 'turn', role: MessageType.User, content: 'question' }, + { id: 'turn', role: MessageType.Assistant, content: 'answer', reference }, + { id: 'next', role: MessageType.Assistant, content: '' }, + ]; + const { result, rerender } = renderHook(() => useMessageReferences(messages)); + + expect(result.current.get(messages[0])).toBe(EmptyReference); + expect(result.current.get(messages[1])).toBe(reference); + expect(result.current.get(messages[2])).toBe(EmptyReference); + const resolved = result.current; + rerender(); + expect(result.current).toBe(resolved); +}); + +it('handles missing history', () => { + const { result } = renderHook(() => useMessageReferences(undefined)); + expect(result.current.size).toBe(0); +}); diff --git a/web/src/pages/next-chats/hooks/use-message-references.ts b/web/src/pages/next-chats/hooks/use-message-references.ts index 0d04d41e48..f5b06daa90 100644 --- a/web/src/pages/next-chats/hooks/use-message-references.ts +++ b/web/src/pages/next-chats/hooks/use-message-references.ts @@ -1,71 +1,14 @@ -import { MessageType } from '@/constants/chat'; import { IMessage, IReference } from '@/interfaces/database/chat'; -import { isEmpty } from 'lodash'; import { useMemo } from 'react'; import { EmptyReference } from '../utils'; -/** - * Resolves the reference of every message in one pass and caches the result - * until the message list or the conversation's reference list changes. - * - * Doing this per message inside the render loop (the previous - * `buildMessageItemReference` call site) was quadratic, and — worse — handed - * `MessageItem` a freshly allocated object on every streaming flush, which - * defeated its `memo` and re-parsed the markdown of the whole transcript - * ~16 times a second. - * - * Keyed by message object rather than by id: a question and its answer share an - * id, so an id-keyed map could not tell them apart. - */ -export function useMessageReferences( - messages: IMessage[] | undefined, - reference: IReference[] | undefined, -) { +// Key by object because a question and its answer share the same message id. +export function useMessageReferences(messages: IMessage[] | undefined) { return useMemo(() => { - const list = messages ?? []; - const references = reference ?? []; - - // Index of each assistant message within the answer sequence, skipping the - // prologue (which has no reference) and error answers. First id wins, which - // is what the previous `findIndex` lookup did. - const answerIndexById = new Map(); - let answerCount = 0; - list.forEach((message) => { - if ( - message.role !== MessageType.Assistant || - message.content?.startsWith('**ERROR**:') - ) { - return; - } - if (answerCount > 0 && !answerIndexById.has(message.id)) { - answerIndexById.set(message.id, answerCount - 1); - } - answerCount += 1; - }); - const resolved = new Map(); - list.forEach((message) => { - if (!isEmpty(message.reference)) { - resolved.set(message, message.reference as IReference); - return; - } - // An assistant message that has not received any content yet is still - // being generated. Never resolve its reference from the conversation - // reference list — the indices only align with completed answers, so the - // lookup would surface a stale reference from a previous turn (or from a - // previously opened conversation) while the answer is streaming. - if (message.role === MessageType.Assistant && !message.content) { - resolved.set(message, EmptyReference); - return; - } - const answerIndex = answerIndexById.get(message.id); - resolved.set( - message, - (answerIndex === undefined ? undefined : references[answerIndex]) ?? - EmptyReference, - ); + (messages ?? []).forEach((message) => { + resolved.set(message, message.reference ?? EmptyReference); }); - return resolved; - }, [messages, reference]); + }, [messages]); } diff --git a/web/src/pages/next-chats/share/index.tsx b/web/src/pages/next-chats/share/index.tsx index 75d50d307b..7ffd16c9e2 100644 --- a/web/src/pages/next-chats/share/index.tsx +++ b/web/src/pages/next-chats/share/index.tsx @@ -43,7 +43,7 @@ const ChatContainer = () => { const sendDisabled = useSendButtonDisabled(value); const { data: chatInfo } = useFetchExternalChatInfo(); - const messageReferences = useMessageReferences(derivedMessages, undefined); + const messageReferences = useMessageReferences(derivedMessages); React.useEffect(() => { if (locale && i18n.language !== locale) { diff --git a/web/src/pages/next-chats/utils.ts b/web/src/pages/next-chats/utils.ts index 83e064e7ea..65086c1360 100644 --- a/web/src/pages/next-chats/utils.ts +++ b/web/src/pages/next-chats/utils.ts @@ -26,19 +26,16 @@ export const isConversationIdExist = (conversationId: string) => { }; export const getDocumentIdsFromConversionReference = (data: IConversation) => { - const documentIds = data.reference.reduce( - (pre: Array, cur: IReference) => { - cur.doc_aggs - ?.map((x) => x.doc_id) - .forEach((x) => { - if (pre.every((y) => y !== x)) { - pre.push(x); - } - }); - return pre; - }, - [], - ); + const documentIds = data.messages.reduce((pre: Array, cur) => { + cur.reference?.doc_aggs + ?.map((x) => x.doc_id) + .forEach((x) => { + if (pre.every((y) => y !== x)) { + pre.push(x); + } + }); + return pre; + }, []); return documentIds.join(','); }; diff --git a/web/src/services/agent-service.ts b/web/src/services/agent-service.ts index 22bfad1f37..83896114ef 100644 --- a/web/src/services/agent-service.ts +++ b/web/src/services/agent-service.ts @@ -251,8 +251,14 @@ export const fetchAgentLogsByCanvasId = ( }); }; -export const fetchAgentLogsById = (canvasId: string, sessionId: string) => { - return request.get(api.fetchAgentSessionById(canvasId, sessionId)); +export const fetchAgentLogsById = ( + canvasId: string, + sessionId: string, + params?: Pick, +) => { + return request.get(api.fetchAgentSessionById(canvasId, sessionId), { + params, + }); }; export const fetchPipeLineList = (params: IPipeLineListRequest) => { diff --git a/web/src/utils/__tests__/chat.test.ts b/web/src/utils/__tests__/chat.test.ts index 928a4e566a..da3565cb2d 100644 --- a/web/src/utils/__tests__/chat.test.ts +++ b/web/src/utils/__tests__/chat.test.ts @@ -1,4 +1,8 @@ +import { MessageType } from '@/constants/chat'; +import type { IReference } from '@/interfaces/database/chat'; import { + buildMessageListWithUuid, + buildReferenceObject, decodeAngleBracketEntities, escapeUnmatchedAngleBrackets, mergeAnswerChunk, @@ -168,3 +172,39 @@ describe('replaceThinkToSection', () => { expect(replaceThinkToSection('plain answer')).toBe('plain answer'); }); }); + +describe('message references', () => { + it('preserves persisted references when assigning message ids', () => { + const reference: IReference = { chunks: [], doc_aggs: [], total: 0 }; + const messages = buildMessageListWithUuid([ + { role: MessageType.Assistant, content: 'answer', reference }, + ]); + expect(messages[0].id).toBeTruthy(); + expect(messages[0].reference).toBe(reference); + }); + + it('maps persisted citations and documents for the agent renderer', () => { + const chunk = { + id: 'chunk', + content: null, + citation_id: 'citation', + document_id: 'document', + document_name: 'Document', + dataset_id: 'dataset', + image_id: '', + similarity: 1, + vector_similarity: 1, + term_similarity: 1, + positions: [], + }; + const document = { doc_id: 'document', doc_name: 'Document', count: 1 }; + const reference = buildReferenceObject({ + chunks: [chunk], + doc_aggs: [document], + total: 1, + }); + expect(reference.chunks.citation).toBe(chunk); + expect(reference.doc_aggs.document).toBe(document); + expect(buildReferenceObject()).toEqual({ chunks: {}, doc_aggs: {} }); + }); +}); diff --git a/web/src/utils/chat.ts b/web/src/utils/chat.ts index 095cd8a8f5..17705921b2 100644 --- a/web/src/utils/chat.ts +++ b/web/src/utils/chat.ts @@ -18,9 +18,13 @@ import { ChatVariableEnabledField, EmptyConversationId, } from '@/constants/chat'; -import { IMessage, Message } from '@/interfaces/database/chat'; +import { + IMessage, + IReference, + IReferenceObject, + Message, +} from '@/interfaces/database/chat'; import type { CompletionChunk } from '@/services/chat-completion-stream'; -import { omit } from 'lodash'; import { v4 as uuid } from 'uuid'; import { buildCitationDisplayMap, @@ -44,12 +48,26 @@ export const buildMessageUuid = (message: Partial) => { export const buildMessageListWithUuid = (messages?: Message[]) => { return ( messages?.map((x: Message | IMessage) => ({ - ...omit(x, 'reference'), + ...x, id: buildMessageUuid(x), })) ?? [] ); }; +export const buildReferenceObject = ( + reference?: IReference, +): IReferenceObject => ({ + chunks: Object.fromEntries( + (reference?.chunks ?? []).map((chunk, index) => [ + chunk.citation_id ?? String(index), + chunk, + ]), + ), + doc_aggs: Object.fromEntries( + (reference?.doc_aggs ?? []).map((document) => [document.doc_id, document]), + ), +}); + export const generateConversationId = () => { return uuid().replace(/-/g, ''); };