Skip to content

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

kc

parent 2337a5b9
Hide whitespace changes
Inline Side-by-side
Showing with 261 additions and 363 deletions
flow.go
View file @ 8c52cd01
...@@ -32,125 +32,125 @@ type ProviderInfo struct { ...@@ -32,125 +32,125 @@ type ProviderInfo struct {
// DefineModelsFlow 定义 models 流 // DefineModelsFlow 定义 models 流
func DefineModelsFlow(g *genkit.Genkit) { func DefineModelsFlow(g *genkit.Genkit) {
genkit.DefineFlow(g, "models", func(ctx context.Context, provider string) ([]ProviderInfo, error) { genkit.DefineFlow(g, "models", func(ctx context.Context, provider string) ([]ProviderInfo, error) {
log.Info(). log.Info().
Str("method", "DefineModelsFlow"). Str("method", "DefineModelsFlow").
Str("provider", provider). Str("provider", provider).
Msg("Listing models and embedders for provider") Msg("Listing models and embedders for provider")
infos := make([]ProviderInfo, 0, 3) infos := make([]ProviderInfo, 0, 3)
switch provider { switch provider {
case googlegenai.GoogleAIProvider: case googlegenai.GoogleAIProvider:
googleModels, err := googlegenai.ListModels(googlegenai.GoogleAIProvider) googleModels, err := googlegenai.ListModels(googlegenai.GoogleAIProvider)
if err != nil { if err != nil {
log.Error(). log.Error().
Str("method", "DefineModelsFlow"). Str("method", "DefineModelsFlow").
Str("provider", provider). Str("provider", provider).
Err(err). Err(err).
Msg("Failed to list Google AI provider info") Msg("Failed to list Google AI provider info")
return nil, err return nil, err
} }
googleEmbedders, err := googlegenai.ListEmbedders(genai.BackendGeminiAPI) googleEmbedders, err := googlegenai.ListEmbedders(genai.BackendGeminiAPI)
if err != nil { if err != nil {
log.Error(). log.Error().
Str("method", "DefineModelsFlow"). Str("method", "DefineModelsFlow").
Str("provider", provider). Str("provider", provider).
Err(err). Err(err).
Msg("Failed to list Google AI embedders") Msg("Failed to list Google AI embedders")
return nil, err return nil, err
} }
googleInfo := ProviderInfo{ googleInfo := ProviderInfo{
Provider: googlegenai.GoogleAIProvider, Provider: googlegenai.GoogleAIProvider,
Models: googleModels, Models: googleModels,
Embedders: googleEmbedders, Embedders: googleEmbedders,
} }
infos = append(infos, googleInfo) infos = append(infos, googleInfo)
return infos, nil return infos, nil
case deepseek.Provider: case deepseek.Provider:
dsModels := deepseek.ListModels() dsModels := deepseek.ListModels()
infos = append(infos, ProviderInfo{ infos = append(infos, ProviderInfo{
Provider: deepseek.Provider, Provider: deepseek.Provider,
Models: dsModels, Models: dsModels,
Embedders: []string{}, // DeepSeek 不支持嵌入器 Embedders: []string{}, // DeepSeek 不支持嵌入器
}) })
return infos, nil return infos, nil
case ollama.Provider:
ollamaModels, err := ollama.ListModels()
if err != nil {
log.Error().
Str("method", "DefineModelsFlow").
Str("provider", provider).
Err(err).
Msg("Failed to list Ollama provider info")
return nil, err
}
infos = append(infos, ProviderInfo{
Provider: ollama.Provider,
Models: ollamaModels,
})
return infos, nil
default:
// Google AI
googleModels, err := googlegenai.ListModels(googlegenai.GoogleAIProvider)
if err != nil {
log.Error().
Str("method", "DefineModelsFlow").
Str("provider", googlegenai.GoogleAIProvider).
Err(err).
Msg("Failed to list Google AI provider info")
return nil, err
}
googleEmbedders, err := googlegenai.ListEmbedders(genai.BackendGeminiAPI)
if err != nil {
log.Error().
Str("method", "DefineModelsFlow").
Str("provider", googlegenai.GoogleAIProvider).
Err(err).
Msg("Failed to list Google AI embedders")
return nil, err
}
infos = append(infos, ProviderInfo{
Provider: googlegenai.GoogleAIProvider,
Models: googleModels,
Embedders: googleEmbedders,
})
// DeepSeek
dsModels := deepseek.ListModels()
infos = append(infos, ProviderInfo{
Provider: deepseek.Provider,
Models: dsModels,
Embedders: []string{}, // DeepSeek 不支持嵌入器
})
case ollama.Provider:
ollamaModels, err := ollama.ListModels() ollamaModels, err := ollama.ListModels()
if err != nil { if err != nil {
log.Error(). log.Error().
Str("method", "DefineModelsFlow"). Str("method", "DefineModelsFlow").
Str("provider", provider). Str("provider", provider).
Err(err). Err(err).
Msg("Failed to list Ollama provider info") Msg("Failed to list Ollama provider info")
return nil, err return nil, err
} }
infos = append(infos, ProviderInfo{ infos = append(infos, ProviderInfo{
Provider: ollama.Provider, Provider: ollama.Provider,
Models: ollamaModels, Models: ollamaModels,
}) })
return infos, nil
log.Info().
Str("method", "DefineModelsFlow"). default:
Str("provider", provider). // Google AI
Msg("------default---") googleModels, err := googlegenai.ListModels(googlegenai.GoogleAIProvider)
if err != nil {
return infos, nil log.Error().
} Str("method", "DefineModelsFlow").
}) Str("provider", googlegenai.GoogleAIProvider).
Err(err).
Msg("Failed to list Google AI provider info")
return nil, err
}
googleEmbedders, err := googlegenai.ListEmbedders(genai.BackendGeminiAPI)
if err != nil {
log.Error().
Str("method", "DefineModelsFlow").
Str("provider", googlegenai.GoogleAIProvider).
Err(err).
Msg("Failed to list Google AI embedders")
return nil, err
}
infos = append(infos, ProviderInfo{
Provider: googlegenai.GoogleAIProvider,
Models: googleModels,
Embedders: googleEmbedders,
})
// DeepSeek
dsModels := deepseek.ListModels()
infos = append(infos, ProviderInfo{
Provider: deepseek.Provider,
Models: dsModels,
Embedders: []string{}, // DeepSeek 不支持嵌入器
})
ollamaModels, err := ollama.ListModels()
if err != nil {
log.Error().
Str("method", "DefineModelsFlow").
Str("provider", provider).
Err(err).
Msg("Failed to list Ollama provider info")
return nil, err
}
infos = append(infos, ProviderInfo{
Provider: ollama.Provider,
Models: ollamaModels,
})
log.Info().
Str("method", "DefineModelsFlow").
Str("provider", provider).
Msg("------default---")
return infos, nil
}
})
} }
func DefineGraphFlow(g *genkit.Genkit, graphIndexer ai.Indexer) { func DefineGraphFlow(g *genkit.Genkit, graphIndexer ai.Indexer) {
...@@ -229,7 +229,7 @@ func DefineDocumentFlow(g *genkit.Genkit, indexer ai.Indexer) { ...@@ -229,7 +229,7 @@ func DefineDocumentFlow(g *genkit.Genkit, indexer ai.Indexer) {
}) })
} }
func DefineChatFlow(g *genkit.Genkit, retriever ai.Retriever, graphRetriever ai.Retriever, pgConnString string,kc *knowledge.KnowledgeClient ) { func DefineChatFlow(g *genkit.Genkit, retriever ai.Retriever, graphRetriever ai.Retriever, pgConnString string, kc *knowledge.KnowledgeClient) {
qa, err := question.InitQAStore(pgConnString) qa, err := question.InitQAStore(pgConnString)
if err != nil { if err != nil {
...@@ -240,9 +240,7 @@ func DefineChatFlow(g *genkit.Genkit, retriever ai.Retriever, graphRetriever ai. ...@@ -240,9 +240,7 @@ func DefineChatFlow(g *genkit.Genkit, retriever ai.Retriever, graphRetriever ai.
genkit.DefineFlow(g, "chat", func(ctx context.Context, input *ChatInput) (Response, error) { genkit.DefineFlow(g, "chat", func(ctx context.Context, input *ChatInput) (Response, error) {
ctxAsJson, _ := json.Marshal(ctx) ctxAsJson, _ := json.Marshal(ctx)
log.Info().Msgf("input----ctxAsJson----%s", string(ctxAsJson)) log.Info().Msgf("input----ctxAsJson----%s", string(ctxAsJson))
inputAsJson, err := json.Marshal(input) inputAsJson, err := json.Marshal(input)
if err != nil { if err != nil {
...@@ -435,41 +433,52 @@ func defineSimpleQaPrompt(g *genkit.Genkit, promptName string) (*ai.Prompt, erro ...@@ -435,41 +433,52 @@ func defineSimpleQaPrompt(g *genkit.Genkit, promptName string) (*ai.Prompt, erro
return simpleQaPrompt, nil return simpleQaPrompt, nil
} }
func DefineSplitDocFlow(g *genkit.Genkit, kc *knowledge.KnowledgeClient) {
func DefineSplitDocFlow(g *genkit.Genkit,kc *knowledge.KnowledgeClient ) {
// Define a simple flow that generates jokes about a given topic // Define a simple flow that generates jokes about a given topic
genkit.DefineFlow(g, "split/document", func(ctx context.Context, input string) (SplitDocResponse, error) { genkit.DefineFlow(g, "split/document", func(ctx context.Context, input string) (SplitDocResponse, error) {
res,err := kc.ReconstructDocumentSSE(ctx,input) res, err := kc.ReconstructDocumentSSE(ctx, input)
if err != nil{ if err != nil {
return SplitDocResponse{ return SplitDocResponse{
Code: 500, Code: 500,
Msg: err.Error(), Msg: err.Error(),
},nil }, nil
} }
for k,v := range res{ for k, v := range res {
fmt.Println("process",k,*v.Progress) fmt.Println("process", k, *v.Progress)
} }
return SplitDocResponse{ return SplitDocResponse{
Code: 200, Code: 200,
Data: res, Data: res,
},nil }, nil
}) })
} }
type Response struct { type Response struct {
Data string `json:"data"` Data string `json:"data"`
Code int `json:"code"` Code int `json:"code"`
Msg string `json:"msg"` Msg string `json:"msg"`
} }
type SplitDocResponse struct { type SplitDocResponse struct {
Data []knowledge.ReconstructDocumentSSEResponse `json:"data"` Data []knowledge.ReconstructDocumentSSEResponse `json:"data"`
Code int `json:"code"` Code int `json:"code"`
Msg string `json:"msg"` Msg string `json:"msg"`
} }
const simpleQaPromptTemplate = `
You're a helpful agent that answers the user's questions based on the provided context.
Here is the user's query: {{query}}
Here is the context you should use: {{context}} from Milvus
Graph context: {{graph}}
Previous conversation summary: {{summary}}
Instructions:
- If the query is related to a character's personality, adopt the tone and style specified in the Personality context, and generate a response using the Milvus and Graph contexts to inform the personality-driven content.
- For all other queries, provide a clear and accurate response using the Milvus and Graph contexts, without emphasizing the Personality context.
- Ensure responses leverage the Previous conversation summary when relevant.
`
main.go
View file @ 8c52cd01
...@@ -59,39 +59,6 @@ type simpleQaPromptInput struct { ...@@ -59,39 +59,6 @@ type simpleQaPromptInput struct {
Summary string `json:"summary"` Summary string `json:"summary"`
} }
// const simpleQaPromptTemplate = `
// You're a helpful agent that answers the user's questions with a tone and style shaped by the specified personality.
// Here is the user's query: {{query}}
// Here is the context you should use: {{context}} from Milvus
// Graph context: {{graph}}
// Previous conversation summary: {{summary}}
// Personality to adopt: {{personality}}
// Please provide a response that aligns with the given personality while leveraging the provided context, graph, and conversation summary.
// `
const simpleQaPromptTemplate = `
You're a helpful agent that answers the user's questions based on the provided context.
Here is the user's query: {{query}}
Here is the context you should use: {{context}} from Milvus
Graph context: {{graph}}
Previous conversation summary: {{summary}}
Instructions:
- If the query is related to a character's personality, adopt the tone and style specified in the Personality context, and generate a response using the Milvus and Graph contexts to inform the personality-driven content.
- For all other queries, provide a clear and accurate response using the Milvus and Graph contexts, without emphasizing the Personality context.
- Ensure responses leverage the Previous conversation summary when relevant.
`
func main() { func main() {
// Define command-line flags with hardcoded values as defaults // Define command-line flags with hardcoded values as defaults
...@@ -163,8 +130,6 @@ func main() { ...@@ -163,8 +130,6 @@ func main() {
log.Fatal().Msgf("graphrag.DefineIndexerAndRetriever failed: %v", err) log.Fatal().Msgf("graphrag.DefineIndexerAndRetriever failed: %v", err)
} }
DefineDocumentFlow(g, indexer) DefineDocumentFlow(g, indexer)
DefineGraphFlow(g, graphIndexer) DefineGraphFlow(g, graphIndexer)
DefineModelsFlow(g) DefineModelsFlow(g)
...@@ -179,8 +144,8 @@ func main() { ...@@ -179,8 +144,8 @@ func main() {
} }
log.Info().Msg("KnowledgeClient initialized successfully") log.Info().Msg("KnowledgeClient initialized successfully")
DefineChatFlow(g, retriever, graphRetriever, *pgConnString,kc) DefineChatFlow(g, retriever, graphRetriever, *pgConnString, kc)
DefineSplitDocFlow(g,kc) DefineSplitDocFlow(g, kc)
// 配置限速器：每秒 10 次请求，突发容量 20，最大并发 5 // 配置限速器：每秒 10 次请求，突发容量 20，最大并发 5
rl := NewRateLimiter(10, 20, 5) rl := NewRateLimiter(10, 20, 5)
...@@ -207,4 +172,3 @@ func main() { ...@@ -207,4 +172,3 @@ func main() {
log.Fatal().Msgf("Server failed: %v", err) log.Fatal().Msgf("Server failed: %v", err)
} }
} }
plugins/knowledge/example/main.go deleted 100644 → 0
View file @ 2337a5b9
package main
import (
"fmt"
"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"
)
func main() {
// 实例化一个认证对象，入参需要传入腾讯云账户 SecretId 和 SecretKey，此处还需注意密钥对的保密
// 代码泄露可能会导致 SecretId 和 SecretKey 泄露，并威胁账号下所有资源的安全性
// 以下代码示例仅供参考，建议采用更安全的方式来使用密钥
// 请参见：https://cloud.tencent.com/document/product/1278/85305
// 密钥可前往官网控制台 https://console.cloud.tencent.com/cam/capi 进行获取
credential := common.NewCredential(
"AKID64oLfmfLtESUJ6i8LPSM4gCVbiniQuBF",
"rX2JMBnBMJ2YqulOo37xa5OUMSN4Xnpd",
)
// 使用临时密钥示例
// credential := common.NewTokenCredential("SecretId", "SecretKey", "Token")
// 实例化一个client选项，可选的，没有特殊需求可以跳过
cpf := profile.NewClientProfile()
cpf.HttpProfile.Endpoint = "lkeap.tencentcloudapi.com"
// 实例化要请求产品的client对象,clientProfile是可选的
client, _ := lkeap.NewClient(credential, "ap-guangzhou", cpf)
// 实例化一个请求对象,每个接口都会对应一个request对象
request := lkeap.NewReconstructDocumentSSERequest()
request.FileType = common.StringPtr("TXT")
request.FileBase64 = common.StringPtr("data:application/octet-stream;base64,4oCc6L+Z5bCx6LGh6Iqx5LiA5qC344CC5aaC5p6c5L2g54ix5LiK5LqG5LiA5py155Sf6ZW/IArlnKjkuIDpopfmmJ/mmJ/kuIrnmoToirHvvIzpgqPkuYjlpJzpl7TvvIwgCuS9oOeci+edgOWkqeepuuWwseaEn+WIsOeUnOicnOaEiSAK5b+r44CC5omA5pyJ55qE5pif5pif5LiK6YO9CuWlveixoeW8gOedgOiKseOAguKAnQ==")
// 返回的resp是一个ReconstructDocumentSSEResponse的实例，与请求对象对应
response, err := client.ReconstructDocumentSSE(request)
if _, ok := err.(*errors.TencentCloudSDKError); ok {
fmt.Printf("An API error has returned: %s", err)
return
}
if err != nil {
panic(err)
}
// 输出json格式的字符串回包
if response.Response != nil {
// 非流式响应
fmt.Println(response.ToJsonString())
} else {
// 流式响应
for event := range response.Events {
fmt.Println(string(event.Data))
}
}
}
\ No newline at end of file
plugins/knowledge/knowledge.go
View file @ 8c52cd01
...@@ -291,157 +291,138 @@ func (kc *KnowledgeClient) QueryRewriteWithSummary(ctx context.Context, userQues ...@@ -291,157 +291,138 @@ func (kc *KnowledgeClient) QueryRewriteWithSummary(ctx context.Context, userQues
return kc.QueryRewrite(ctx, req) return kc.QueryRewrite(ctx, req)
} }
// ReconstructDocumentSSE performs document reconstruction using the Tencent Cloud LKEAP API.
// It takes a text string, encodes it as Base64, and returns the reconstructed content as a slice of ReconstructDocumentSSEResponse structs for events with Progress "100".
func (kc *KnowledgeClient) ReconstructDocumentSSE(ctx context.Context, text string) ([]ReconstructDocumentSSEResponse, error) {
log.Info().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Str("text_prefix", text[:min(32, len(text))]).
Msg("Starting document reconstruction operation")
if !kc.initted {
log.Error().Str("method", "KnowledgeClient.ReconstructDocumentSSE").Msg("Client not initialized")
return nil, fmt.Errorf("knowledge client not initialized; call Init first")
}
// Validate input
if text == "" {
log.Error().Str("method", "KnowledgeClient.ReconstructDocumentSSE").Msg("Text is required")
return nil, fmt.Errorf("text is required")
}
// Encode text to Base64
base64Data := base64.StdEncoding.EncodeToString([]byte(text))
// Check Base64 size (8MB limit)
if len(base64Data)/4*3 > 8*1024*1024 {
log.Error().Str("method", "KnowledgeClient.ReconstructDocumentSSE").Msg("Encoded text exceeds 8MB limit")
return nil, fmt.Errorf("encoded text exceeds 8MB limit")
}
// Remove Base64 prefix if present
if strings.HasPrefix(base64Data, "data:application/octet-stream;base64,") {
base64Data = strings.TrimPrefix(base64Data, "data:application/octet-stream;base64,")
log.Debug().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Msg("Removed Base64 data URI prefix")
}
// Create SDK request
request := lkeap.NewReconstructDocumentSSERequest()
request.FileType = common.StringPtr("TXT")
request.FileBase64 = common.StringPtr(base64Data)
// Debug request
requestAsJson, _ := json.Marshal(request)
log.Debug().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Str("request_json", string(requestAsJson)).
Msg("Prepared Tencent Cloud request")
// Perform request
response, err := kc.client.ReconstructDocumentSSEWithContext(ctx, request)
if _, ok := err.(*errors.TencentCloudSDKError); ok {
log.Error().
Err(err).
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Msg("Tencent Cloud API error")
return nil, fmt.Errorf("tencent cloud api error: %s", err)
}
if err != nil {
log.Error().
Err(err).
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Msg("Failed to perform document reconstruction")
return nil, fmt.Errorf("document reconstruction failed: %w", err)
}
// Handle streaming response
var results []ReconstructDocumentSSEResponse
for event := range response.Events {
var recon ReconstructDocumentSSEResponse
eventData := event.Data
if len(eventData) == 0 {
log.Warn().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Msg("Empty event data in streaming response")
continue
}
if err := json.Unmarshal(eventData, &recon); err != nil {
log.Error().
Err(err).
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Str("event_data", string(eventData)).
Msg("Failed to parse streaming event")
continue
}
// Only collect events with Progress == "100"
if recon.Progress != nil && *recon.Progress == "100" {
results = append(results, recon)
log.Debug().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Str("event_data", string(eventData)).
Msg("Collected streaming event with Progress 100")
} else {
log.Debug().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Str("event_data", string(eventData)).
Msg("Skipped streaming event with Progress != 100")
}
}
// Check if any valid results were collected
if len(results) == 0 {
log.Error().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Msg("No events with Progress 100 returned")
return nil, fmt.Errorf("no events with progress 100 returned")
}
log.Info().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Int("results_count", len(results)).
Msg("Document reconstruction operation completed successfully")
return results, nil
// ReconstructDocumentSSE performs document reconstruction using the Tencent Cloud LKEAP API.
// It takes a text string, encodes it as Base64, and returns the reconstructed content as a slice of ReconstructDocumentSSEResponse structs for events with Progress "100".
func (kc *KnowledgeClient) ReconstructDocumentSSE(ctx context.Context, text string) ([]ReconstructDocumentSSEResponse, error) {
log.Info().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Str("text_prefix", text[:min(32, len(text))]).
Msg("Starting document reconstruction operation")
if !kc.initted {
log.Error().Str("method", "KnowledgeClient.ReconstructDocumentSSE").Msg("Client not initialized")
return nil, fmt.Errorf("knowledge client not initialized; call Init first")
}
// Validate input
if text == "" {
log.Error().Str("method", "KnowledgeClient.ReconstructDocumentSSE").Msg("Text is required")
return nil, fmt.Errorf("text is required")
}
// Encode text to Base64
base64Data := base64.StdEncoding.EncodeToString([]byte(text))
// Check Base64 size (8MB limit)
if len(base64Data)/4*3 > 8*1024*1024 {
log.Error().Str("method", "KnowledgeClient.ReconstructDocumentSSE").Msg("Encoded text exceeds 8MB limit")
return nil, fmt.Errorf("encoded text exceeds 8MB limit")
}
// Remove Base64 prefix if present
if strings.HasPrefix(base64Data, "data:application/octet-stream;base64,") {
base64Data = strings.TrimPrefix(base64Data, "data:application/octet-stream;base64,")
log.Debug().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Msg("Removed Base64 data URI prefix")
}
// Create SDK request
request := lkeap.NewReconstructDocumentSSERequest()
request.FileType = common.StringPtr("TXT")
request.FileBase64 = common.StringPtr(base64Data)
// Debug request
requestAsJson, _ := json.Marshal(request)
log.Debug().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Str("request_json", string(requestAsJson)).
Msg("Prepared Tencent Cloud request")
// Perform request
response, err := kc.client.ReconstructDocumentSSEWithContext(ctx, request)
if _, ok := err.(*errors.TencentCloudSDKError); ok {
log.Error().
Err(err).
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Msg("Tencent Cloud API error")
return nil, fmt.Errorf("tencent cloud api error: %s", err)
}
if err != nil {
log.Error().
Err(err).
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Msg("Failed to perform document reconstruction")
return nil, fmt.Errorf("document reconstruction failed: %w", err)
}
// Handle streaming response
var results []ReconstructDocumentSSEResponse
for event := range response.Events {
var recon ReconstructDocumentSSEResponse
eventData := event.Data
if len(eventData) == 0 {
log.Warn().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Msg("Empty event data in streaming response")
continue
}
if err := json.Unmarshal(eventData, &recon); err != nil {
log.Error().
Err(err).
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Str("event_data", string(eventData)).
Msg("Failed to parse streaming event")
continue
}
// Only collect events with Progress == "100"
if recon.Progress != nil && *recon.Progress == "100" {
results = append(results, recon)
log.Debug().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Str("event_data", string(eventData)).
Msg("Collected streaming event with Progress 100")
} else {
log.Debug().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Str("event_data", string(eventData)).
Msg("Skipped streaming event with Progress != 100")
}
}
// Check if any valid results were collected
if len(results) == 0 {
log.Error().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Msg("No events with Progress 100 returned")
return nil, fmt.Errorf("no events with progress 100 returned")
}
log.Info().
Str("method", "KnowledgeClient.ReconstructDocumentSSE").
Int("results_count", len(results)).
Msg("Document reconstruction operation completed successfully")
return results, nil
} }
// min returns the minimum of two integers // min returns the minimum of two integers
func min(a, b int) int { func min(a, b int) int {
if a < b { if a < b {
return a return a
} }
return b return b
} }
// ReconstructDocumentSSEResponse defines the response structure for the ReconstructDocumentSSE API. // ReconstructDocumentSSEResponse defines the response structure for the ReconstructDocumentSSE API.
type ReconstructDocumentSSEResponse struct { type ReconstructDocumentSSEResponse struct {
RequestId *string `json:"RequestId,omitempty"` RequestId *string `json:"RequestId,omitempty"`
TaskId *string `json:"TaskId,omitempty"` TaskId *string `json:"TaskId,omitempty"`
ResponseType *string `json:"ResponseType,omitempty"` ResponseType *string `json:"ResponseType,omitempty"`
Progress *string `json:"Progress,omitempty"` Progress *string `json:"Progress,omitempty"`
ProgressMessage *string `json:"ProgressMessage,omitempty"` ProgressMessage *string `json:"ProgressMessage,omitempty"`
DocumentRecognizeResultUrl *string `json:"DocumentRecognizeResultUrl,omitempty"` DocumentRecognizeResultUrl *string `json:"DocumentRecognizeResultUrl,omitempty"`
FailedPages []*int64 `json:"FailedPages,omitempty"` FailedPages []*int64 `json:"FailedPages,omitempty"`
StatusCode *string `json:"StatusCode,omitempty"` StatusCode *string `json:"StatusCode,omitempty"`
} }
plugins/knowledge/knowledge_test.go
View file @ 8c52cd01
...@@ -113,8 +113,6 @@ func TestKnowledgeClient_QueryRewrite(t *testing.T) { ...@@ -113,8 +113,6 @@ func TestKnowledgeClient_QueryRewrite(t *testing.T) {
} }
} }
// func TestKnowledgeClient_ReconstructDocumentSSE(t *testing.T) { // func TestKnowledgeClient_ReconstructDocumentSSE(t *testing.T) {
// // Warning: Do not hardcode credentials in production code. Use environment variables or a secure vault. // // 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_ID", "AKID64oLfmfLtESUJ6i8LPSM4gCVbiniQuBF")
...@@ -208,4 +206,4 @@ func TestKnowledgeClient_QueryRewrite(t *testing.T) { ...@@ -208,4 +206,4 @@ func TestKnowledgeClient_QueryRewrite(t *testing.T) {
// Msg("Document reconstruction successful") // Msg("Document reconstruction 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