Skip to content

A
agentchat
Commit afaa731b authored by Wade's avatar Wade
Browse files
Options

add QueryRewriteWithSummary

parent 27b3a0bb
Hide whitespace changes
Inline Side-by-side
Showing with 763 additions and 718 deletions
plugins/knowledge/knowledge.go
View file @ afaa731b
// // Copyright 2025
// //
// // Licensed under the Apache License, Version 2.0 (the "License");
// // you may not use this file except in compliance with the License.
// // You may obtain a copy of the License at
// //
// // http://www.apache.org/licenses/LICENSE-2.0
// //
// // Unless required by applicable law or agreed to in writing, software
// // distributed under the License is distributed on an "AS IS" BASIS,
// // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// // See the License for the specific language governing permissions and
// // limitations under the License.
// //
// // SPDX-License-Identifier: Apache-2.0
// package knowledge
// import (
// "context"
// "encoding/json"
// "fmt"
// "os"
// "strings"
// "sync"
// "github.com/rs/zerolog/log"
// "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common"
// "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common/errors"
// "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common/profile"
// lkeap "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/lkeap/v20240522"
// )
// // ClientConfig holds configuration options for the Tencent Cloud LKEAP client.
// type ClientConfig struct {
// SecretID string // Tencent Cloud Secret ID
// SecretKey string // Tencent Cloud Secret Key
// Token string // Optional: Temporary token for authentication
// Endpoint string // API endpoint (default: lkeap.tencentcloudapi.com)
// Region string // Tencent Cloud region (optional)
// }
// // KnowledgeClient manages interactions with the Tencent Cloud LKEAP API.
// type KnowledgeClient struct {
// client *lkeap.Client
// config ClientConfig
// mu sync.Mutex
// initted bool
// }
// // NewKnowledgeClient creates a new KnowledgeClient with the given configuration.
// func NewKnowledgeClient(config ClientConfig) *KnowledgeClient {
// log.Info().
// Str("method", "NewKnowledgeClient").
// Str("endpoint", config.Endpoint).
// Str("region", config.Region).
// Str("secret_id", maskCredential(config.SecretID)).
// Str("token", maskCredential(config.Token)).
// Msg("Creating new KnowledgeClient")
// return &KnowledgeClient{
// config: config,
// }
// }
// // Init initializes the KnowledgeClient.
// func (kc *KnowledgeClient) Init(ctx context.Context) error {
// log.Info().Str("method", "KnowledgeClient.Init").Msg("Initializing KnowledgeClient")
// kc.mu.Lock()
// defer kc.mu.Unlock()
// if kc.initted {
// log.Error().Str("method", "KnowledgeClient.Init").Msg("Client already initialized")
// return fmt.Errorf("knowledge client already initialized")
// }
// // Load configuration from environment variables if not set
// if kc.config.SecretID == "" {
// kc.config.SecretID = os.Getenv("TENCENTCLOUD_SECRET_ID")
// }
// if kc.config.SecretKey == "" {
// kc.config.SecretKey = os.Getenv("TENCENTCLOUD_SECRET_KEY")
// }
// if kc.config.Token == "" {
// kc.config.Token = os.Getenv("TENCENTCLOUD_TOKEN")
// }
// if kc.config.Endpoint == "" {
// kc.config.Endpoint = "lkeap.tencentcloudapi.com"
// }
// if kc.config.Region == "" {
// kc.config.Region = "ap-guangzhou"
// }
// // Validate configuration
// if kc.config.SecretID == "" || kc.config.SecretKey == "" {
// log.Error().Str("method", "KnowledgeClient.Init").Msg("SecretID and SecretKey are required")
// return fmt.Errorf("knowledge: SecretID and SecretKey are required")
// }
// // Create credential
// var credential *common.Credential
// if kc.config.Token != "" {
// credential = common.NewTokenCredential(kc.config.SecretID, kc.config.SecretKey, kc.config.Token)
// log.Debug().Str("method", "KnowledgeClient.Init").Msg("Using temporary token credential")
// } else {
// credential = common.NewCredential(kc.config.SecretID, kc.config.SecretKey)
// log.Debug().Str("method", "KnowledgeClient.Init").Msg("Using standard credential")
// }
// // Create client profile
// cpf := profile.NewClientProfile()
// cpf.HttpProfile.Endpoint = kc.config.Endpoint
// // Initialize client
// client, err := lkeap.NewClient(credential, kc.config.Region, cpf)
// if err != nil {
// log.Error().
// Err(err).
// Str("method", "KnowledgeClient.Init").
// Msg("Failed to create LKEAP client")
// return err
// }
// kc.client = client
// kc.initted = true
// log.Info().Str("method", "KnowledgeClient.Init").Msg("Initialization successful")
// return nil
// }
// // QueryRewriteRequest defines the input for a query rewrite operation.
// type QueryRewriteRequest struct {
// Messages []*lkeap.Message // Multi-turn conversation history (up to 4 turns)
// Model string // Model name for query rewriting
// }
// // QueryRewriteResponse defines the output of a query rewrite operation.
// type QueryRewriteResponse struct {
// RewrittenQuery string // The rewritten query
// RawResponse *lkeap.QueryRewriteResponse
// }
// // QueryRewrite performs a query rewrite using the Tencent Cloud LKEAP API.
// func (kc *KnowledgeClient) QueryRewrite(ctx context.Context, req QueryRewriteRequest) (*QueryRewriteResponse, error) {
// log.Info().
// Str("method", "KnowledgeClient.QueryRewrite").
// Int("message_count", len(req.Messages)).
// Str("model", req.Model).
// Msg("Starting query rewrite operation")
// if !kc.initted {
// log.Error().Str("method", "KnowledgeClient.QueryRewrite").Msg("Client not initialized")
// return nil, fmt.Errorf("knowledge client not initialized; call Init first")
// }
// // Validate input
// if len(req.Messages) == 0 {
// log.Error().Str("method", "KnowledgeClient.QueryRewrite").Msg("At least one message is required")
// return nil, fmt.Errorf("at least one message is required")
// }
// if len(req.Messages) > 4 {
// log.Warn().
// Str("method", "KnowledgeClient.QueryRewrite").
// Int("message_count", len(req.Messages)).
// Msg("Message count exceeds 4, truncating to 4")
// req.Messages = req.Messages[:4]
// }
// for i, msg := range req.Messages {
// if msg.Role == nil || *msg.Role == "" {
// log.Error().
// Str("method", "KnowledgeClient.QueryRewrite").
// Int("index", i).
// Msg("Role is required in each message")
// return nil, fmt.Errorf("message at index %d missing role", i)
// }
// if *msg.Role != "user" && *msg.Role != "assistant" {
// log.Error().
// Str("method", "KnowledgeClient.QueryRewrite").
// Int("index", i).
// Str("role", *msg.Role).
// Msg("Invalid role; must be 'user' or 'assistant'")
// return nil, fmt.Errorf("invalid role '%s' at index %d", *msg.Role, i)
// }
// if msg.Content == nil || *msg.Content == "" {
// log.Error().
// Str("method", "KnowledgeClient.QueryRewrite").
// Int("index", i).
// Msg("Content is required in each message")
// return nil, fmt.Errorf("message at index %d missing content", i)
// }
// log.Debug().
// Str("method", "KnowledgeClient.QueryRewrite").
// Int("index", i).
// Str("role", *msg.Role).
// Str("content", *msg.Content).
// Msg("Validated message")
// }
// if req.Model == "" {
// log.Warn().Str("method", "KnowledgeClient.QueryRewrite").Msg("Model not specified, using default")
// req.Model = "lke-query-rewrite-base"
// }
// // Create Tencent Cloud request
// tencentReq := lkeap.NewQueryRewriteRequest()
// tencentReq.Messages = req.Messages
// if req.Model != "" {
// tencentReq.Model = common.StringPtr(req.Model)
// }
// // Debug request
// tencentReqAsJson, _ := json.Marshal(tencentReq)
// log.Debug().
// Str("method", "KnowledgeClient.QueryRewrite").
// Str("request_json", string(tencentReqAsJson)).
// Msg("Prepared Tencent Cloud request")
// // Perform request
// response, err := kc.client.QueryRewriteWithContext(ctx, tencentReq)
// if err != nil {
// if _, ok := err.(*errors.TencentCloudSDKError); ok {
// log.Error().
// Err(err).
// Str("method", "KnowledgeClient.QueryRewrite").
// Msg("Tencent Cloud API error")
// return nil, fmt.Errorf("tencent cloud api error: %w", err)
// }
// log.Error().
// Err(err).
// Str("method", "KnowledgeClient.QueryRewrite").
// Msg("Failed to perform query rewrite")
// return nil, fmt.Errorf("query rewrite failed: %w", err)
// }
// // Extract response fields
// var rewrittenQuery string
// var requestId string
// if response.Response.Content != nil {
// rewrittenQuery = *response.Response.Content
// }
// if response.Response.RequestId != nil {
// requestId = *response.Response.RequestId
// }
// result := &QueryRewriteResponse{
// RewrittenQuery: rewrittenQuery,
// RawResponse: response,
// }
// log.Info().
// Str("method", "KnowledgeClient.QueryRewrite").
// Str("rewritten_query", rewrittenQuery).
// Str("request_id", requestId).
// Interface("usage", response.Response.Usage).
// Str("raw_response", response.ToJsonString()).
// Msg("Query rewrite operation completed successfully")
// return result, nil
// }
// // QueryRewriteWithSummary wraps QueryRewrite to handle a user question, assistant answer, and history summary.
// func (kc *KnowledgeClient) QueryRewriteWithSummary(ctx context.Context, userQuestion, assistantAnswer, historySummary string) (*QueryRewriteResponse, error) {
// log.Info().
// Str("method", "KnowledgeClient.QueryRewriteWithSummary").
// Str("user_question", userQuestion).
// Str("assistant_answer", assistantAnswer).
// Str("history_summary", historySummary).
// Msg("Starting query rewrite with summary operation")
// if userQuestion == "" || assistantAnswer == "" {
// log.Error().Str("method", "KnowledgeClient.QueryRewriteWithSummary").Msg("User question and assistant answer are required")
// return nil, fmt.Errorf("user question and assistant answer are required")
// }
// // Construct messages
// messages := []*lkeap.Message{
// {
// Role: common.StringPtr("user"),
// Content: common.StringPtr(userQuestion),
// },
// {
// Role: common.StringPtr("assistant"),
// Content: common.StringPtr(assistantAnswer),
// },
// }
// // Append history summary as an assistant message if provided
// if historySummary != "" {
// messages = append(messages, &lkeap.Message{
// Role: common.StringPtr("assistant"),
// Content: common.StringPtr(fmt.Sprintf("Conversation summary: %s", historySummary)),
// })
// }
// // Create request
// req := QueryRewriteRequest{
// Messages: messages,
// Model: "lke-query-rewrite-base",
// }
// // Call QueryRewrite
// return kc.QueryRewrite(ctx, req)
// }
// // maskCredential masks sensitive credentials for logging
// func maskCredential(cred string) string {
// if len(cred) <= 8 {
// return strings.Repeat("*", len(cred))
// }
// return cred[:4] + strings.Repeat("*", len(cred)-8) + cred[len(cred)-4:]
// }
// Copyright 2025
//
// Licensed under the Apache License, Version 2.0 (the "License");
......@@ -21,6 +329,7 @@ import (
"encoding/json"
"fmt"
"os"
"strings"
"sync"
"github.com/rs/zerolog/log"
......@@ -53,6 +362,8 @@ func NewKnowledgeClient(config ClientConfig) *KnowledgeClient {
Str("method", "NewKnowledgeClient").
Str("endpoint", config.Endpoint).
Str("region", config.Region).
Str("secret_id", maskCredential(config.SecretID)).
Str("token", maskCredential(config.Token)).
Msg("Creating new KnowledgeClient")
return &KnowledgeClient{
config: config,
......@@ -70,7 +381,7 @@ func (kc *KnowledgeClient) Init(ctx context.Context) error {
return fmt.Errorf("knowledge client already initialized")
}
// Load configuration defaults
// Load configuration from environment variables if not set
if kc.config.SecretID == "" {
kc.config.SecretID = os.Getenv("TENCENTCLOUD_SECRET_ID")
}
......@@ -84,7 +395,7 @@ func (kc *KnowledgeClient) Init(ctx context.Context) error {
kc.config.Endpoint = "lkeap.tencentcloudapi.com"
}
if kc.config.Region == "" {
kc.config.Region = "" // Region can be empty for global APIs
kc.config.Region = "ap-guangzhou" // Default to ap-guangzhou as per curl
}
// Validate configuration
......@@ -192,23 +503,22 @@ func (kc *KnowledgeClient) QueryRewrite(ctx context.Context, req QueryRewriteReq
}
if req.Model == "" {
log.Warn().Str("method", "KnowledgeClient.QueryRewrite").Msg("Model not specified, using default")
req.Model = "lke-query-rewrite-base" // Default as per curl
}
// Create Tencent Cloud request
tencentReq := lkeap.NewQueryRewriteRequest()
tencentReq.Messages = req.Messages // Directly use validated messages
tencentReq.Messages = req.Messages
if req.Model != "" {
tencentReq.Model = common.StringPtr(req.Model)
}
fmt.Println("len(tencentReq.Messages)",len(tencentReq.Messages))
tencentReqAsJson, _ := json.Marshal(tencentReq)
fmt.Println("len(tencentReq.Messages) json marsh",string(tencentReqAsJson))
// Debug request
tencentReqAsJson, _ := json.Marshal(tencentReq)
log.Debug().
Str("method", "KnowledgeClient.QueryRewrite").
Str("request_json", string(tencentReqAsJson)).
Msg("Prepared Tencent Cloud request")
// Perform request
response, err := kc.client.QueryRewriteWithContext(ctx, tencentReq)
......@@ -226,526 +536,83 @@ func (kc *KnowledgeClient) QueryRewrite(ctx context.Context, req QueryRewriteReq
Msg("Failed to perform query rewrite")
return nil, fmt.Errorf("query rewrite failed: %w", err)
}
// Extract response fields
var rewrittenQuery string
var requestId string
if response.Response.Content != nil {
rewrittenQuery = *response.Response.Content
}
if response.Response.RequestId != nil {
requestId = *response.Response.RequestId
}
result := &QueryRewriteResponse{
RewrittenQuery: rewrittenQuery,
// Usage: response.Response.Usage,
// RequestId: requestId,
RawResponse: response,
// Extract response fields
var rewrittenQuery string
var requestId string
if response.Response.Content != nil {
rewrittenQuery = *response.Response.Content
}
if response.Response.RequestId != nil {
requestId = *response.Response.RequestId
}
result := &QueryRewriteResponse{
RewrittenQuery: rewrittenQuery,
RawResponse: response,
}
log.Info().
Str("method", "KnowledgeClient.QueryRewrite").
Str("rewritten_query", rewrittenQuery).
Str("request_id", requestId).
Interface("usage", response.Response.Usage).
Str("raw_response", response.ToJsonString()).
Msg("Query rewrite operation completed successfully")
return result, nil
}
log.Info().
Str("method", "KnowledgeClient.QueryRewrite").
Str("rewritten_query", rewrittenQuery).
Str("request_id", requestId).
Interface("usage", response.Response.Usage).
Str("raw_response", response.ToJsonString()).
Msg("Query rewrite operation completed successfully")
return result, nil
// maskCredential masks sensitive credentials for logging
func maskCredential(cred string) string {
if len(cred) <= 8 {
return strings.Repeat("*", len(cred))
}
return cred[:4] + strings.Repeat("*", len(cred)-8) + cred[len(cred)-4:]
}
// // Extract rewritten query
// var rewrittenQuery string
// if response.Response.RewrittenQuery != nil {
// rewrittenQuery = *response.Response.RewrittenQuery
// }
// result := &QueryRewriteResponse{
// RewrittenQuery: rewrittenQuery,
// RawResponse: response,
// }
// log.Info().
// Str("method", "KnowledgeClient.QueryRewrite").
// Str("rewritten_query", rewrittenQuery).
// Str("raw_response", response.ToJsonString()).
// Msg("Query rewrite operation completed successfully")
// return result, nil
// }
// // Copyright 2025
// //
// // Licensed under the Apache License, Version 2.0 (the "License");
// // you may not use this file except in compliance with the License.
// // You may obtain a copy of the License at
// //
// // http://www.apache.org/licenses/LICENSE-2.0
// //
// // Unless required by applicable law or agreed to in writing, software
// // distributed under the License is distributed on an "AS IS" BASIS,
// // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// // See the License for the specific language governing permissions and
// // limitations under the License.
// //
// // SPDX-License-Identifier: Apache-2.0
// package knowledge
// import (
// "context"
// "fmt"
// "os"
// "sync"
// "github.com/rs/zerolog/log"
// "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common"
// "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common/errors"
// "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common/profile"
// lkeap "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/lkeap/v20240522"
// )
// // ClientConfig holds configuration options for the Tencent Cloud LKEAP client.
// type ClientConfig struct {
// SecretID string // Tencent Cloud Secret ID
// SecretKey string // Tencent Cloud Secret Key
// Token string // Optional: Temporary token for authentication
// Endpoint string // API endpoint (default: lkeap.tencentcloudapi.com)
// Region string // Tencent Cloud region (optional)
// }
// // KnowledgeClient manages interactions with the Tencent Cloud LKEAP API.
// type KnowledgeClient struct {
// client *lkeap.Client
// config ClientConfig
// mu sync.Mutex
// initted bool
// }
// // NewKnowledgeClient creates a new KnowledgeClient with the given configuration.
// func NewKnowledgeClient(config ClientConfig) *KnowledgeClient {
// log.Info().
// Str("method", "NewKnowledgeClient").
// Str("endpoint", config.Endpoint).
// Str("region", config.Region).
// Msg("Creating new KnowledgeClient")
// return &KnowledgeClient{
// config: config,
// }
// }
// // Init initializes the KnowledgeClient.
// func (kc *KnowledgeClient) Init(ctx context.Context) error {
// log.Info().Str("method", "KnowledgeClient.Init").Msg("Initializing KnowledgeClient")
// kc.mu.Lock()
// defer kc.mu.Unlock()
// if kc.initted {
// log.Error().Str("method", "KnowledgeClient.Init").Msg("Client already initialized")
// return errors.New("knowledge client already initialized")
// }
// // Load configuration defaults
// if kc.config.SecretID == "" {
// kc.config.SecretID = os.Getenv("TENCENTCLOUD_SECRET_ID")
// }
// if kc.config.SecretKey == "" {
// kc.config.SecretKey = os.Getenv("TENCENTCLOUD_SECRET_KEY")
// }
// if kc.config.Token == "" {
// kc.config.Token = os.Getenv("TENCENTCLOUD_TOKEN")
// }
// if kc.config.Endpoint == "" {
// kc.config.Endpoint = "lkeap.tencentcloudapi.com"
// }
// if kc.config.Region == "" {
// kc.config.Region = "" // Region can be empty for global APIs
// }
// // Validate configuration
// if kc.config.SecretID == "" || kc.config.SecretKey == "" {
// log.Error().Str("method", "KnowledgeClient.Init").Msg("SecretID and SecretKey are required")
// return fmt.Errorf("knowledge: SecretID and SecretKey are required")
// }
// // Create credential
// var credential *common.Credential
// if kc.config.Token != "" {
// credential = common.NewTokenCredential(kc.config.SecretID, kc.config.SecretKey, kc.config.Token)
// log.Debug().Str("method", "KnowledgeClient.Init").Msg("Using temporary token credential")
// } else {
// credential = common.NewCredential(kc.config.SecretID, kc.config.SecretKey)
// log.Debug().Str("method", "KnowledgeClient.Init").Msg("Using standard credential")
// }
// // Create client profile
// cpf := profile.NewClientProfile()
// cpf.HttpProfile.Endpoint = kc.config.Endpoint
// // Initialize client
// client, err := lkeap.NewClient(credential, kc.config.Region, cpf)
// if err != nil {
// log.Error().
// Err(err).
// Str("method", "KnowledgeClient.Init").
// Msg("Failed to create LKEAP client")
// return err
// }
// kc.client = client
// kc.initted = true
// log.Info().Str("method", "KnowledgeClient.Init").Msg("Initialization successful")
// return nil
// }
// // Message represents a single turn in a conversation.
// type Message struct {
// User string // User's question
// Assistant string // Assistant's response
// }
// // QueryRewriteRequest defines the input for a query rewrite operation.
// type QueryRewriteRequest struct {
// Messages []*Message // Multi-turn conversation history (up to 4 turns)
// Model string // Model name for query rewriting
// }
// // QueryRewriteResponse defines the output of a query rewrite operation.
// type QueryRewriteResponse struct {
// RewrittenQuery string // The rewritten query
// RawResponse *lkeap.QueryRewriteResponse
// }
// // QueryRewrite performs a query rewrite using the Tencent Cloud LKEAP API.
// func (kc *KnowledgeClient) QueryRewrite(ctx context.Context, req QueryRewriteRequest) (*QueryRewriteResponse, error) {
// log.Info().
// Str("method", "KnowledgeClient.QueryRewrite").
// Int("message_count", len(req.Messages)).
// Str("model", req.Model).
// Msg("Starting query rewrite operation")
// if !kc.initted {
// log.Error().Str("method", "KnowledgeClient.QueryRewrite").Msg("Client not initialized")
// return nil, fmt.Errorf("knowledge client not initialized; call Init first")
// }
// // Validate input
// if len(req.Messages) == 0 {
// log.Error().Str("method", "KnowledgeClient.QueryRewrite").Msg("At least one message is required")
// return nil, fmt.Errorf("at least one message is required")
// }
// if len(req.Messages) > 4 {
// log.Warn().
// Str("method", "KnowledgeClient.QueryRewrite").
// Int("message_count", len(req.Messages)).
// Msg("Message count exceeds 4, truncating to 4")
// req.Messages = req.Messages[:4]
// }
// for i, msg := range req.Messages {
// if msg.User == "" || msg.Assistant == "" {
// log.Error().
// Str("method", "KnowledgeClient.QueryRewrite").
// Int("index", i).
// Msg("User and Assistant fields are required in each message")
// return nil, fmt.Errorf("message at index %d missing user or assistant", i)
// }
// }
// if req.Model == "" {
// log.Warn().Str("method", "KnowledgeClient.QueryRewrite").Msg("Model not specified, using default")
// }
// // Create Tencent Cloud request
// tencentReq := lkeap.NewQueryRewriteRequest()
// tencentReq.Messages = make([]*lkeap.Message, len(req.Messages))
// for i, msg := range req.Messages {
// tencentReq.Messages[i] = &lkeap.Message{
// User: common.StringPtr(msg.User),
// Assistant: common.StringPtr(msg.Assistant),
// }
// log.Debug().
// Str("method", "KnowledgeClient.QueryRewrite").
// Int("index", i).
// Str("user", msg.User).
// Str("assistant", msg.Assistant).
// Msg("Added message to request")
// }
// if req.Model != "" {
// tencentReq.Model = common.StringPtr(req.Model)
// }
// // Perform request
// response, err := kc.client.QueryRewriteWithContext(ctx, tencentReq)
// if err != nil {
// if _, ok := err.(*errors.TencentCloudSDKError); ok {
// log.Error().
// Err(err).
// Str("method", "KnowledgeClient.QueryRewrite").
// Msg("Tencent Cloud API error")
// return nil, fmt.Errorf("tencent cloud api error: %w", err)
// }
// log.Error().
// Err(err).
// Str("method", "KnowledgeClient.QueryRewrite").
// Msg("Failed to perform query rewrite")
// return nil, fmt.Errorf("query rewrite failed: %w", err)
// }
// // Extract rewritten query
// var rewrittenQuery string
// if response.Response.RewrittenQuery != nil {
// rewrittenQuery = *response.Response.RewrittenQuery
// }
// result := &QueryRewriteResponse{
// RewrittenQuery: rewrittenQuery,
// RawResponse: response,
// }
// log.Info().
// Str("method", "KnowledgeClient.QueryRewrite").
// Str("rewritten_query", rewrittenQuery).
// Str("raw_response", response.ToJsonString()).
// Msg("Query rewrite operation completed successfully")
// return result, nil
// }
// // Copyright 2025
// //
// // Licensed under the Apache License, Version 2.0 (the "License");
// // you may not use this file except in compliance with the License.
// // You may obtain a copy of the License at
// //
// // http://www.apache.org/licenses/LICENSE-2.0
// //
// // Unless required by applicable law or agreed to in writing, software
// // distributed under the License is distributed on an "AS IS" BASIS,
// // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// // See the License for the specific language governing permissions and
// // limitations under the License.
// //
// // SPDX-License-Identifier: Apache-2.0
// package knowledge
// import (
// "context"
// "fmt"
// "os"
// "sync"
// "github.com/rs/zerolog/log"
// "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common"
// "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common/errors"
// "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common/profile"
// lkeap "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/lkeap/v20240522"
// )
// // ClientConfig holds configuration options for the Tencent Cloud LKEAP client.
// type ClientConfig struct {
// SecretID string // Tencent Cloud Secret ID
// SecretKey string // Tencent Cloud Secret Key
// Token string // Optional: Temporary token for authentication
// Endpoint string // API endpoint (default: lkeap.tencentcloudapi.com)
// Region string // Tencent Cloud region (optional)
// }
// // KnowledgeClient manages interactions with the Tencent Cloud LKEAP API.
// type KnowledgeClient struct {
// client *lkeap.Client
// config ClientConfig
// mu sync.Mutex
// initted bool
// }
// // NewKnowledgeClient creates a new KnowledgeClient with the given configuration.
// func NewKnowledgeClient(config ClientConfig) *KnowledgeClient {
// log.Info().
// Str("method", "NewKnowledgeClient").
// Str("endpoint", config.Endpoint).
// Str("region", config.Region).
// Msg("Creating new KnowledgeClient")
// return &KnowledgeClient{
// config: config,
// }
// }
// // Init initializes the KnowledgeClient.
// func (kc *KnowledgeClient) Init(ctx context.Context) error {
// log.Info().Str("method", "KnowledgeClient.Init").Msg("Initializing KnowledgeClient")
// kc.mu.Lock()
// defer kc.mu.Unlock()
// if kc.initted {
// log.Error().Str("method", "KnowledgeClient.Init").Msg("Client already initialized")
// return fmt.Errorf("knowledge client already initialized")
// }
// // Load configuration defaults
// if kc.config.SecretID == "" {
// kc.config.SecretID = os.Getenv("TENCENTCLOUD_SECRET_ID")
// }
// if kc.config.SecretKey == "" {
// kc.config.SecretKey = os.Getenv("TENCENTCLOUD_SECRET_KEY")
// }
// if kc.config.Token == "" {
// kc.config.Token = os.Getenv("TENCENTCLOUD_TOKEN")
// }
// if kc.config.Endpoint == "" {
// kc.config.Endpoint = "lkeap.tencentcloudapi.com"
// }
// if kc.config.Region == "" {
// kc.config.Region = "" // Region can be empty for global APIs
// }
// // Validate configuration
// if kc.config.SecretID == "" || kc.config.SecretKey == "" {
// log.Error().Str("method", "KnowledgeClient.Init").Msg("SecretID and SecretKey are required")
// //return errors.New("knowledge: SecretID and SecretKey are required")
// return fmt.Errorf("knowledge: SecretID and SecretKey are required")
// }
// // Create credential
// var credential *common.Credential
// if kc.config.Token != "" {
// credential = common.NewTokenCredential(kc.config.SecretID, kc.config.SecretKey, kc.config.Token)
// log.Debug().Str("method", "KnowledgeClient.Init").Msg("Using temporary token credential")
// } else {
// credential = common.NewCredential(kc.config.SecretID, kc.config.SecretKey)
// log.Debug().Str("method", "KnowledgeClient.Init").Msg("Using standard credential")
// }
// // Create client profile
// cpf := profile.NewClientProfile()
// cpf.HttpProfile.Endpoint = kc.config.Endpoint
// // Initialize client
// client, err := lkeap.NewClient(credential, kc.config.Region, cpf)
// if err != nil {
// log.Error().
// Err(err).
// Str("method", "KnowledgeClient.Init").
// Msg("Failed to create LKEAP client")
// return err
// }
// kc.client = client
// kc.initted = true
// log.Info().Str("method", "KnowledgeClient.Init").Msg("Initialization successful")
// return nil
// }
// // QueryRewriteRequest defines the input for a query rewrite operation.
// type QueryRewriteRequest struct {
// Query string // The input query to rewrite
// }
// // QueryRewriteResponse defines the output of a query rewrite operation.
// type QueryRewriteResponse struct {
// RewrittenQuery string // The rewritten query
// RawResponse *lkeap.QueryRewriteResponse
// }
// // QueryRewrite performs a query rewrite using the Tencent Cloud LKEAP API.
// func (kc *KnowledgeClient) QueryRewrite(ctx context.Context, req QueryRewriteRequest) (*QueryRewriteResponse, error) {
// log.Info().
// Str("method", "KnowledgeClient.QueryRewrite").
// Str("query", req.Query).
// Msg("Starting query rewrite operation")
// if !kc.initted {
// log.Error().Str("method", "KnowledgeClient.QueryRewrite").Msg("Client not initialized")
// //return nil, errors.New("knowledge client not initialized; call Init first")
// return nil,fmt.Errorf("knowledge client not initialized; call Init first")
// }
// if req.Query == "" {
// log.Error().Str("method", "KnowledgeClient.QueryRewrite").Msg("Query is required")
// //return nil, errors.New("query is required")
// return nil, fmt.Errorf("query is required")
// }
// // Create request
// tencentReq := lkeap.NewQueryRewriteRequest()
// tencentReq.Query = common.StringPtr(req.Query)
// // Perform request
// response, err := kc.client.QueryRewriteWithContext(ctx, tencentReq)
// if err != nil {
// if _, ok := err.(*errors.TencentCloudSDKError); ok {
// log.Error().
// Err(err).
// Str("method", "KnowledgeClient.QueryRewrite").
// Msg("Tencent Cloud API error")
// return nil, fmt.Errorf("tencent cloud api error: %w", err)
// }
// log.Error().
// Err(err).
// Str("method", "KnowledgeClient.QueryRewrite").
// Msg("Failed to perform query rewrite")
// return nil, fmt.Errorf("query rewrite failed: %w", err)
// }
// // Extract rewritten query
// var rewrittenQuery string
// if response.Response.RewrittenQuery != nil {
// rewrittenQuery = *response.Response.RewrittenQuery
// }
// result := &QueryRewriteResponse{
// RewrittenQuery: rewrittenQuery,
// RawResponse: response,
// }
// log.Info().
// Str("method", "KnowledgeClient.QueryRewrite").
// Str("rewritten_query", rewrittenQuery).
// Str("raw_response", response.ToJsonString()).
// Msg("Query rewrite operation completed successfully")
// return result, nil
// }
// QueryRewriteWithSummary wraps QueryRewrite to handle a user question, assistant answer, and history summary.
func (kc *KnowledgeClient) QueryRewriteWithSummary(ctx context.Context, userQuestion, assistantAnswer, historySummary string) (*QueryRewriteResponse, error) {
log.Info().
Str("method", "KnowledgeClient.QueryRewriteWithSummary").
Str("user_question", userQuestion).
Str("assistant_answer", assistantAnswer).
Str("history_summary", historySummary).
Msg("Starting query rewrite with summary operation")
if userQuestion == "" || assistantAnswer == "" {
log.Error().Str("method", "KnowledgeClient.QueryRewriteWithSummary").Msg("User question and assistant answer are required")
return nil, fmt.Errorf("user question and assistant answer are required")
}
// Construct messages
messages := []*lkeap.Message{
{
Role: common.StringPtr("user"),
Content: common.StringPtr(userQuestion),
},
{
Role: common.StringPtr("assistant"),
Content: common.StringPtr(assistantAnswer),
},
}
// Append history summary as an assistant message if provided
if historySummary != "" {
messages = append(messages, &lkeap.Message{
Role: common.StringPtr("user"),
Content: common.StringPtr(fmt.Sprintf("Conversation summary: %s", historySummary)),
})
}
// Create request
req := QueryRewriteRequest{
Messages: messages,
Model: "lke-query-rewrite-base",
}
// Call QueryRewrite
return kc.QueryRewrite(ctx, req)
}
plugins/knowledge/knowledge_test.go
View file @ afaa731b
// Copyright 2025
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0
package knowledge
import (
"context"
"os"
"testing"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common"
lkeap "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/lkeap/v20240522"
)
func TestMain(m *testing.M) {
// Configure zerolog for human-readable console output during tests
zerolog.TimeFieldFormat = zerolog.TimeFormatUnix
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})
// Run tests
os.Exit(m.Run())
}
func TestKnowledgeClient_QueryRewrite(t *testing.T) {
// Warning: Do not hardcode credentials in production code. Use environment variables or a secure vault.
// The credentials below are assumed to be placeholders for testing purposes.
os.Setenv("TENCENTCLOUD_SECRET_ID", "AKID64oLfmfLtESUJ6i8LPSM4gCVbiniQuBF")
os.Setenv("TENCENTCLOUD_SECRET_KEY", "rX2JMBnBMJ2YqulOo37xa5OUMSN4Xnpd")
defer func() {
os.Unsetenv("TENCENTCLOUD_SECRET_ID")
os.Unsetenv("TENCENTCLOUD_SECRET_KEY")
}()
// Create client configuration
config := ClientConfig{
Endpoint: "lkeap.tencentcloudapi.com",
Region: "ap-guangzhou",
}
// Initialize client
client := NewKnowledgeClient(config)
ctx := context.Background()
// Test cases
tests := []struct {
name string
messages []*lkeap.Message
model string
expectError bool
}{
{
name: "ValidMultiTurnConversation",
messages: []*lkeap.Message{
{
Role: common.StringPtr("user"),
Content: common.StringPtr("What is the capital of France?"),
},
{
Role: common.StringPtr("assistant"),
Content: common.StringPtr("The capital of France is Paris."),
},
{
Role: common.StringPtr("user"),
Content: common.StringPtr("Tell me more about Paris."),
},
{
Role: common.StringPtr("assistant"),
Content: common.StringPtr("Paris is known for its art, culture, and landmarks like the Eiffel Tower."),
},
},
model: "",
expectError: true, // Expect error unless valid credentials are used
},
{
name: "EmptyMessages",
messages: []*lkeap.Message{},
model: "",
expectError: true,
},
{
name: "InvalidRole",
messages: []*lkeap.Message{
{
Role: common.StringPtr("invalid-role"),
Content: common.StringPtr("Test query"),
},
},
model: "",
expectError: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Initialize client for each test
if err := client.Init(ctx); err != nil {
t.Fatalf("Failed to initialize KnowledgeClient: %v", err)
}
// Perform query rewrite
req := QueryRewriteRequest{
Messages: tt.messages,
Model: tt.model,
}
resp, err := client.QueryRewrite(ctx, req)
// Check error expectation
if tt.expectError {
if err == nil {
t.Error("Expected error, got none")
} else {
log.Debug().
Str("method", "TestKnowledgeClient_QueryRewrite").
Str("test_name", tt.name).
Err(err).
Msg("Received expected error")
}
return
}
// Check response (only for non-error cases)
if err != nil {
t.Fatalf("QueryRewrite failed: %v", err)
}
if resp.RewrittenQuery == "" {
t.Error("Expected non-empty rewritten query")
}
// if resp.RequestId == "" {
// t.Error("Expected non-empty request ID")
// }
// if resp.Usage == nil {
// t.Error("Expected non-nil usage")
// }
log.Info().
Str("method", "TestKnowledgeClient_QueryRewrite").
Str("test_name", tt.name).
Str("rewritten_query", resp.RewrittenQuery).
// Str("request_id", resp.RequestId).
// Interface("usage", resp.Usage).
Msg("Query rewrite successful")
})
}
}
// // Copyright 2025
// //
// // Licensed under the Apache License, Version 2.0 (the "License");
......@@ -217,15 +36,8 @@ func TestKnowledgeClient_QueryRewrite(t *testing.T) {
// os.Exit(m.Run())
// }
// /*
// SecretId:AKID64oLfmfLtESUJ6i8LPSM4gCVbiniQuBF
// SecretKey:rX2JMBnBMJ2YqulOo37xa5OUMSN4Xnpd
// */
// func TestKnowledgeClient_QueryRewrite(t *testing.T) {
// // Set up environment variables for testing (mock credentials)
// // Warning: Do not hardcode credentials in production code. Use environment variables or a secure vault.
// os.Setenv("TENCENTCLOUD_SECRET_ID", "AKID64oLfmfLtESUJ6i8LPSM4gCVbiniQuBF")
// os.Setenv("TENCENTCLOUD_SECRET_KEY", "rX2JMBnBMJ2YqulOo37xa5OUMSN4Xnpd")
// defer func() {
......@@ -251,6 +63,25 @@ func TestKnowledgeClient_QueryRewrite(t *testing.T) {
// expectError bool
// }{
// {
// name: "CurlPayload",
// messages: []*lkeap.Message{
// {
// Role: common.StringPtr("user"),
// Content: common.StringPtr("你的家在哪里"),
// },
// {
// Role: common.StringPtr("assistant"),
// Content: common.StringPtr("国内"),
// },
// {
// Role: common.StringPtr("user"),
// Content: common.StringPtr("国内哪里"),
// },
// },
// model: "lke-query-rewrite-base",
// expectError: true,
// },
// {
// name: "ValidMultiTurnConversation",
// messages: []*lkeap.Message{
// {
......@@ -270,13 +101,13 @@ func TestKnowledgeClient_QueryRewrite(t *testing.T) {
// Content: common.StringPtr("Paris is known for its art, culture, and landmarks like the Eiffel Tower."),
// },
// },
// model: "default-model",
// expectError: true, // Expect error due to mock credentials
// model: "",
// expectError: true,
// },
// {
// name: "EmptyMessages",
// messages: []*lkeap.Message{},
// model: "default-model",
// model: "",
// expectError: true,
// },
// {
......@@ -287,7 +118,7 @@ func TestKnowledgeClient_QueryRewrite(t *testing.T) {
// Content: common.StringPtr("Test query"),
// },
// },
// model: "default-model",
// model: "",
// expectError: true,
// },
// }
......@@ -320,9 +151,9 @@ func TestKnowledgeClient_QueryRewrite(t *testing.T) {
// return
// }
// // Check response (only for non-error cases)
// // Check response
// if err != nil {
// t.Fatalf("QueryRewrite failed: %v", err)
// t.Errorf("QueryRewrite failed: %v", err)
// }
// if resp.RewrittenQuery == "" {
// t.Error("Expected non-empty rewritten query")
......@@ -335,4 +166,351 @@ func TestKnowledgeClient_QueryRewrite(t *testing.T) {
// Msg("Query rewrite successful")
// })
// }
// }
\ No newline at end of file
// }
// func TestKnowledgeClient_QueryRewriteWithSummary(t *testing.T) {
// // Warning: Do not hardcode credentials in production code. Use environment variables or a secure vault.
// os.Setenv("TENCENTCLOUD_SECRET_ID", "AKID64oLfmfLtESUJ6i8LPSM4gCVbiniQuBF")
// os.Setenv("TENCENTCLOUD_SECRET_KEY", "rX2JMBnBMJ2YqulOo37xa5OUMSN4Xnpd")
// defer func() {
// os.Unsetenv("TENCENTCLOUD_SECRET_ID")
// os.Unsetenv("TENCENTCLOUD_SECRET_KEY")
// }()
// // Create client configuration
// config := ClientConfig{
// Endpoint: "lkeap.tencentcloudapi.com",
// Region: "ap-guangzhou",
// }
// // Initialize client
// client := NewKnowledgeClient(config)
// ctx := context.Background()
// // Test cases
// tests := []struct {
// name string
// userQuestion string
// assistantAnswer string
// historySummary string
// expectError bool
// }{
// {
// name: "ValidWithSummary",
// userQuestion: "你的家在哪里",
// assistantAnswer: "国内",
// historySummary: "User asked about location preferences earlier.",
// expectError: true, // Expect error due to potentially invalid credentials
// },
// {
// name: "ValidWithoutSummary",
// userQuestion: "你的家在哪里",
// assistantAnswer: "国内",
// historySummary: "",
// expectError: true,
// },
// {
// name: "EmptyQuestion",
// userQuestion: "",
// assistantAnswer: "国内",
// historySummary: "Summary",
// expectError: true,
// },
// {
// name: "EmptyAnswer",
// userQuestion: "你的家在哪里",
// assistantAnswer: "",
// historySummary: "Summary",
// expectError: true,
// },
// }
// for _, tt := range tests {
// t.Run(tt.name, func(t *testing.T) {
// // Initialize client for each test
// if err := client.Init(ctx); err != nil {
// t.Fatalf("Failed to initialize KnowledgeClient: %v", err)
// }
// // Perform query rewrite with summary
// resp, err := client.QueryRewriteWithSummary(ctx, tt.userQuestion, tt.assistantAnswer, tt.historySummary)
// // Check error expectation
// if tt.expectError {
// if err == nil {
// t.Error("Expected error, got none")
// } else {
// log.Debug().
// Str("method", "TestKnowledgeClient_QueryWithSummary").
// Str("test_name", tt.name).
// Str("error", err.Error()).
// Msg("Received expected error")
// }
// return
// }
// // Check response
// if err != nil {
// t.Errorf("QueryRewriteWithSummary failed: %v", err)
// }
// if resp.RewrittenQuery == "" {
// t.Error("Expected non-empty rewritten query")
// }
// log.Info().
// Str("method", "TestKnowledgeClient_QueryWithSummary").
// Str("test_name", tt.name).
// Str("rewritten_query", resp.RewrittenQuery).
// Msg("Query rewrite with summary successful")
// })
// }
// }
// Copyright 2025
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0
package knowledge
import (
"context"
"os"
"testing"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common"
lkeap "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/lkeap/v20240522"
)
func TestMain(m *testing.M) {
// Configure zerolog for human-readable console output during tests
zerolog.TimeFieldFormat = zerolog.TimeFormatUnix
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})
// Run tests
os.Exit(m.Run())
}
func TestKnowledgeClient_QueryRewrite(t *testing.T) {
// Warning: Do not hardcode credentials in production code. Use environment variables or a secure vault.
// The credentials below are placeholders for testing purposes.
os.Setenv("TENCENTCLOUD_SECRET_ID", "AKID64oLfmfLtESUJ6i8LPSM4gCVbiniQuBF")
os.Setenv("TENCENTCLOUD_SECRET_KEY", "rX2JMBnBMJ2YqulOo37xa5OUMSN4Xnpd")
defer func() {
os.Unsetenv("TENCENTCLOUD_SECRET_ID")
os.Unsetenv("TENCENTCLOUD_SECRET_KEY")
}()
// Create client configuration
config := ClientConfig{
Endpoint: "lkeap.tencentcloudapi.com",
Region: "ap-guangzhou",
}
// Initialize client
client := NewKnowledgeClient(config)
ctx := context.Background()
// Test cases
tests := []struct {
name string
messages []*lkeap.Message
model string
expectError bool
}{
{
name: "CurlPayload",
messages: []*lkeap.Message{
{
Role: common.StringPtr("user"),
Content: common.StringPtr("你的家在哪里"),
},
{
Role: common.StringPtr("assistant"),
Content: common.StringPtr("国内"),
},
{
Role: common.StringPtr("user"),
Content: common.StringPtr("国内哪里"),
},
},
model: "lke-query-rewrite-base",
expectError: true, // Expect error due to potentially invalid credentials
},
// {
// name: "ValidMultiTurnConversation",
// messages: []*lkeap.Message{
// {
// Role: common.StringPtr("user"),
// Content: common.StringPtr("What is the capital of France?"),
// },
// {
// Role: common.StringPtr("assistant"),
// Content: common.StringPtr("The capital of France is Paris."),
// },
// {
// Role: common.StringPtr("user"),
// Content: common.StringPtr("Tell me more about Paris."),
// },
// {
// Role: common.StringPtr("assistant"),
// Content: common.StringPtr("Paris is known for its art, culture, and landmarks like the Eiffel Tower."),
// },
// },
// model: "",
// expectError: true,
// },
// {
// name: "EmptyMessages",
// messages: []*lkeap.Message{},
// model: "",
// expectError: true,
// },
// {
// name: "InvalidRole",
// messages: []*lkeap.Message{
// {
// Role: common.StringPtr("invalid-role"),
// Content: common.StringPtr("Test query"),
// },
// },
// model: "",
// expectError: true,
// },
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Initialize client for each test
if err := client.Init(ctx); err != nil {
t.Fatalf("Failed to initialize KnowledgeClient: %v", err)
}
// Perform query rewrite
req := QueryRewriteRequest{
Messages: tt.messages,
Model: tt.model,
}
resp, err := client.QueryRewrite(ctx, req)
// Check error expectation
if tt.expectError {
if err == nil {
t.Error("Expected error, got none")
} else {
log.Debug().
Str("method", "TestKnowledgeClient_QueryRewrite").
Str("test_name", tt.name).
Err(err).
Msg("Received expected error")
}
return
}
// Check response (only for non-error cases)
if err != nil {
t.Fatalf("QueryRewrite failed: %v", err)
}
if resp.RewrittenQuery == "" {
t.Error("Expected non-empty rewritten query")
}
log.Info().
Str("method", "TestKnowledgeClient_QueryRewrite").
Str("test_name", tt.name).
Str("rewritten_query", resp.RewrittenQuery).
Msg("Query rewrite successful")
})
}
}
func TestKnowledgeClient_QueryRewriteWithSummary(t *testing.T) {
// Warning: Do not hardcode credentials in production code. Use environment variables or a secure vault.
os.Setenv("TENCENTCLOUD_SECRET_ID", "AKID64oLfmfLtESUJ6i8LPSM4gCVbiniQuBF")
os.Setenv("TENCENTCLOUD_SECRET_KEY", "rX2JMBnBMJ2YqulOo37xa5OUMSN4Xnpd")
defer func() {
os.Unsetenv("TENCENTCLOUD_SECRET_ID")
os.Unsetenv("TENCENTCLOUD_SECRET_KEY")
}()
// Create client configuration
config := ClientConfig{
Endpoint: "lkeap.tencentcloudapi.com",
Region: "ap-guangzhou",
}
// Initialize client
client := NewKnowledgeClient(config)
ctx := context.Background()
// Test cases
tests := []struct {
name string
userQuestion string
assistantAnswer string
historySummary string
expectError bool
}{
{
name: "ValidWithSummary",
userQuestion: "你的家在哪里",
assistantAnswer: "国内",
historySummary: "User asked about location preferences earlier.",
expectError: true, // Expect error due to potentially invalid credentials
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Initialize client for each test
if err := client.Init(ctx); err != nil {
t.Fatalf("Failed to initialize KnowledgeClient: %v", err)
}
// Perform query rewrite with summary
resp, err := client.QueryRewriteWithSummary(ctx, tt.userQuestion, tt.assistantAnswer, tt.historySummary)
// Check error expectation
if tt.expectError {
if err == nil {
t.Error("Expected error, got none")
} else {
log.Debug().
Str("method", "TestKnowledgeClient_QueryWithSummary").
Str("test_name", tt.name).
Str("error", err.Error()).
Msg("Received expected error")
}
return
}
// Check response
if err != nil {
t.Errorf("QueryRewriteWithSummary failed: %v", err)
}
if resp.RewrittenQuery == "" {
t.Error("Expected non-empty rewritten query")
}
log.Info().
Str("method", "TestKnowledgeClient_QueryWithSummary").
Str("test_name", tt.name).
Str("rewritten_query", resp.RewrittenQuery).
Msg("Query rewrite with summary successful")
})
}
}
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment