Skip to content

Commit 0efb175

Browse files
yekkhan-liftoffwyangsun
authored andcommitted
feat: saturn query pipeline for rag optimisation (#23)
* feat: saturn query pipeline for rag optimisation * feat: remove hardcoded limit * feat: remove unused metadata * feat: todo comments * feat: todo comments * feat: decouple query rewriting and rag search * chore: remove unused comments * fix: fix missing s3 config, make embedding model configurable * feat: debug voyage api key * feat: substitute rag embedding provider env var, remove debug log * feat: add IRSA support to service account template * feat: add logs * feat: add observability for query enhancement * fix: fix empty input and tool name for tool-execution span * feat: added embedding span and fixed incorrect token usage * feat: vector search span * feat: make date filter field configurable * feat: let llm handles the date window * feat: inject query enhancement prompt * feat: handle corrupted metadata * fix: fix race condition in S3Provider.Initialize() * perf(rag): optimize result sorting from O(n²) to O(n log n) * fix: sort dates in descending order, better for LLM * fix: fix test * fix: fix golangci lint err * fix: remove redundant metadata filtering * refactor: dates filter are stored as int * refactor: dates filter are stored as int * fix: fix lint
1 parent 8f6115d commit 0efb175

File tree

20 files changed

+1830
-122
lines changed

20 files changed

+1830
-122
lines changed

go.mod

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ module github.com/tuannvm/slack-mcp-client
33
go 1.24.4
44

55
require (
6+
github.com/aws/aws-sdk-go-v2/config v1.29.4
7+
github.com/aws/aws-sdk-go-v2/service/s3vectors v1.4.10
68
github.com/joho/godotenv v1.5.1
79
github.com/mark3labs/mcp-go v0.42.0
810
github.com/openai/openai-go v1.8.2
@@ -24,6 +26,18 @@ require (
2426
github.com/Masterminds/sprig/v3 v3.2.3 // indirect
2527
github.com/PuerkitoBio/goquery v1.8.1 // indirect
2628
github.com/andybalholm/cascadia v1.3.2 // indirect
29+
github.com/aws/aws-sdk-go-v2 v1.39.4 // indirect
30+
github.com/aws/aws-sdk-go-v2/credentials v1.17.57 // indirect
31+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.27 // indirect
32+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.11 // indirect
33+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.11 // indirect
34+
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.2 // indirect
35+
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.3 // indirect
36+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 // indirect
37+
github.com/aws/aws-sdk-go-v2/service/sso v1.24.14 // indirect
38+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.13 // indirect
39+
github.com/aws/aws-sdk-go-v2/service/sts v1.33.12 // indirect
40+
github.com/aws/smithy-go v1.23.1 // indirect
2741
github.com/aymerick/douceur v0.2.0 // indirect
2842
github.com/bahlo/generic-list-go v0.2.0 // indirect
2943
github.com/beorn7/perks v1.0.1 // indirect

go.sum

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,34 @@ github.com/airbrake/gobrake v3.6.1+incompatible/go.mod h1:wM4gu3Cn0W0K7GUuVWnlXZ
3434
github.com/andybalholm/cascadia v1.3.1/go.mod h1:R4bJ1UQfqADjvDa4P6HZHLh/3OxWWEqc0Sk8XGwHqvA=
3535
github.com/andybalholm/cascadia v1.3.2 h1:3Xi6Dw5lHF15JtdcmAHD3i1+T8plmv7BQ/nsViSLyss=
3636
github.com/andybalholm/cascadia v1.3.2/go.mod h1:7gtRlve5FxPPgIgX36uWBX58OdBsSS6lUvCFb+h7KvU=
37+
github.com/aws/aws-sdk-go-v2 v1.39.4 h1:qTsQKcdQPHnfGYBBs+Btl8QwxJeoWcOcPcixK90mRhg=
38+
github.com/aws/aws-sdk-go-v2 v1.39.4/go.mod h1:yWSxrnioGUZ4WVv9TgMrNUeLV3PFESn/v+6T/Su8gnM=
39+
github.com/aws/aws-sdk-go-v2/config v1.29.4 h1:ObNqKsDYFGr2WxnoXKOhCvTlf3HhwtoGgc+KmZ4H5yg=
40+
github.com/aws/aws-sdk-go-v2/config v1.29.4/go.mod h1:j2/AF7j/qxVmsNIChw1tWfsVKOayJoGRDjg1Tgq7NPk=
41+
github.com/aws/aws-sdk-go-v2/credentials v1.17.57 h1:kFQDsbdBAR3GZsB8xA+51ptEnq9TIj3tS4MuP5b+TcQ=
42+
github.com/aws/aws-sdk-go-v2/credentials v1.17.57/go.mod h1:2kerxPUUbTagAr/kkaHiqvj/bcYHzi2qiJS/ZinllU0=
43+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.27 h1:7lOW8NUwE9UZekS1DYoiPdVAqZ6A+LheHWb+mHbNOq8=
44+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.27/go.mod h1:w1BASFIPOPUae7AgaH4SbjNbfdkxuggLyGfNFTn8ITY=
45+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.11 h1:7AANQZkF3ihM8fbdftpjhken0TP9sBzFbV/Ze/Y4HXA=
46+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.11/go.mod h1:NTF4QCGkm6fzVwncpkFQqoquQyOolcyXfbpC98urj+c=
47+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.11 h1:ShdtWUZT37LCAA4Mw2kJAJtzaszfSHFb5n25sdcv4YE=
48+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.11/go.mod h1:7bUb2sSr2MZ3M/N+VyETLTQtInemHXb/Fl3s8CLzm0Y=
49+
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.2 h1:Pg9URiobXy85kgFev3og2CuOZ8JZUBENF+dcgWBaYNk=
50+
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.2/go.mod h1:FbtygfRFze9usAadmnGJNc8KsP346kEe+y2/oyhGAGc=
51+
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.3 h1:eAh2A4b5IzM/lum78bZ590jy36+d/aFLgKF/4Vd1xPE=
52+
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.3/go.mod h1:0yKJC/kb8sAnmlYa6Zs3QVYqaC8ug2AbnNChv5Ox3uA=
53+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 h1:dM9/92u2F1JbDaGooxTq18wmmFzbJRfXfVfy96/1CXM=
54+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15/go.mod h1:SwFBy2vjtA0vZbjjaFtfN045boopadnoVPhu4Fv66vY=
55+
github.com/aws/aws-sdk-go-v2/service/s3vectors v1.4.10 h1:hgJrhznAL6SjFZAqNIexiE9L7Zjc5PMGmwPWNtTE3zc=
56+
github.com/aws/aws-sdk-go-v2/service/s3vectors v1.4.10/go.mod h1:gJNoydxeaa5Av62mqcKTcA/9oFJnnZRseWfDmPKfGv8=
57+
github.com/aws/aws-sdk-go-v2/service/sso v1.24.14 h1:c5WJ3iHz7rLIgArznb3JCSQT3uUMiz9DLZhIX+1G8ok=
58+
github.com/aws/aws-sdk-go-v2/service/sso v1.24.14/go.mod h1:+JJQTxB6N4niArC14YNtxcQtwEqzS3o9Z32n7q33Rfs=
59+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.13 h1:f1L/JtUkVODD+k1+IiSJUUv8A++2qVr+Xvb3xWXETMU=
60+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.13/go.mod h1:tvqlFoja8/s0o+UruA1Nrezo/df0PzdunMDDurUfg6U=
61+
github.com/aws/aws-sdk-go-v2/service/sts v1.33.12 h1:fqg6c1KVrc3SYWma/egWue5rKI4G2+M4wMQN2JosNAA=
62+
github.com/aws/aws-sdk-go-v2/service/sts v1.33.12/go.mod h1:7Yn+p66q/jt38qMoVfNvjbm3D89mGBnkwDcijgtih8w=
63+
github.com/aws/smithy-go v1.23.1 h1:sLvcH6dfAFwGkHLZ7dGiYF7aK6mg4CgKA/iDKjLDt9M=
64+
github.com/aws/smithy-go v1.23.1/go.mod h1:LEj2LM3rBRQJxPZTB4KuzZkaZYnZPnvgIhb4pu07mx0=
3765
github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
3866
github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
3967
github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=

helm-chart/slack-mcp-client/templates/deployment.yaml

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,10 @@ spec:
2424
{{- end }}
2525
securityContext:
2626
{{- toYaml .Values.podSecurityContext | nindent 8 }}
27-
{{- if and .Values.serviceAccount.create .Values.serviceAccount.clusterRoleName }}
28-
serviceAccountName: {{ include "slack-mcp-client.fullname" . }}
27+
{{- if .Values.serviceAccount.create }}
28+
serviceAccountName: {{ .Values.serviceAccount.name | default (include "slack-mcp-client.fullname" .) }}
2929
{{- end }}
30+
3031
{{- if .Values.initContainers }}
3132
initContainers:
3233
{{- range .Values.initContainers }}
Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
1-
{{- if and .Values.serviceAccount.create .Values.serviceAccount.clusterRoleName }}
1+
{{- if .Values.serviceAccount.create }}
22
apiVersion: v1
33
kind: ServiceAccount
44
metadata:
5-
name: {{ include "slack-mcp-client.fullname" . }}
5+
name: {{ .Values.serviceAccount.name | default (include "slack-mcp-client.fullname" .) }}
66
labels:
77
{{- include "slack-mcp-client.labels" . | nindent 4 }}
8+
{{- with .Values.serviceAccount.annotations }}
9+
annotations:
10+
{{- toYaml . | nindent 4 }}
11+
{{- end }}
812
{{- end }}

internal/config/config.go

Lines changed: 31 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -22,17 +22,19 @@ const (
2222

2323
// Config represents the main application configuration
2424
type Config struct {
25-
Version string `json:"version"`
26-
Slack SlackConfig `json:"slack"`
27-
LLM LLMConfig `json:"llm"`
28-
MCPServers map[string]MCPServerConfig `json:"mcpServers"`
29-
RAG RAGConfig `json:"rag,omitempty"`
30-
Monitoring MonitoringConfig `json:"monitoring,omitempty"`
31-
Timeouts TimeoutConfig `json:"timeouts,omitempty"`
32-
Retry RetryConfig `json:"retry,omitempty"`
33-
Reload ReloadConfig `json:"reload,omitempty"`
34-
Observability ObservabilityConfig `json:"observability,omitempty"`
35-
UseStdIOClient bool `json:"useStdIOClient,omitempty"` // Use terminal client instead of a real slack bot, for local development
25+
Version string `json:"version"`
26+
Slack SlackConfig `json:"slack"`
27+
LLM LLMConfig `json:"llm"`
28+
MCPServers map[string]MCPServerConfig `json:"mcpServers"`
29+
QueryEnhancementProvider string `json:"queryEnhancementProvider,omitempty"` // Optional: LLM provider for query enhancement (applies to all queries)
30+
QueryEnhancementPromptFile string `json:"queryEnhancementPromptFile,omitempty"` // Optional: Path to custom query enhancement prompt file
31+
RAG RAGConfig `json:"rag,omitempty"`
32+
Monitoring MonitoringConfig `json:"monitoring,omitempty"`
33+
Timeouts TimeoutConfig `json:"timeouts,omitempty"`
34+
Retry RetryConfig `json:"retry,omitempty"`
35+
Reload ReloadConfig `json:"reload,omitempty"`
36+
Observability ObservabilityConfig `json:"observability,omitempty"`
37+
UseStdIOClient bool `json:"useStdIOClient,omitempty"` // Use terminal client instead of a real slack bot, for local development
3638
}
3739

3840
// SlackConfig contains Slack-specific configuration
@@ -111,26 +113,37 @@ type MCPToolsConfig struct {
111113

112114
// RAGConfig contains RAG system configuration
113115
type RAGConfig struct {
114-
Enabled bool `json:"enabled,omitempty"`
115-
Provider string `json:"provider,omitempty"`
116-
ChunkSize int `json:"chunkSize,omitempty"`
117-
Providers map[string]RAGProviderConfig `json:"providers,omitempty"`
116+
Enabled bool `json:"enabled,omitempty"`
117+
Provider string `json:"provider,omitempty"`
118+
ChunkSize int `json:"chunkSize,omitempty"`
119+
EmbeddingProvider string `json:"embeddingProvider,omitempty"` // Optional: Embedding provider (voyage, openai, cohere, etc.)
120+
Providers map[string]RAGProviderConfig `json:"providers,omitempty"`
121+
EmbeddingProviders map[string]RAGEmbeddingProviderConfig `json:"embeddingProviders,omitempty"` // Embedding provider configs
118122
}
119123

120124
// RAGProviderConfig contains RAG provider-specific settings
121125
// TODO: Refactor this to use a common interface for all RAG providers, can use environment variables to configure the different providers
122126
type RAGProviderConfig struct {
123127
DatabasePath string `json:"databasePath,omitempty"` // Simple provider: path to JSON database
124-
IndexName string `json:"indexName,omitempty"` // OpenAI provider: vector store name
128+
IndexName string `json:"indexName,omitempty"` // OpenAI/S3 provider: vector store/index name
125129
VectorStoreID string `json:"vectorStoreId,omitempty"` // OpenAI provider: existing vector store ID
126130
Dimensions int `json:"dimensions,omitempty"` // OpenAI provider: embedding dimensions
127131
SimilarityMetric string `json:"similarityMetric,omitempty"` // OpenAI provider: similarity metric
128-
MaxResults int `json:"maxResults,omitempty"` // OpenAI provider: maximum search results
129-
ScoreThreshold float64 `json:"scoreThreshold,omitempty"` // OpenAI provider: score threshold
132+
MaxResults int `json:"maxResults,omitempty"` // OpenAI/S3 provider: maximum search results
133+
ScoreThreshold float64 `json:"scoreThreshold,omitempty"` // OpenAI/S3 provider: score threshold
130134
RewriteQuery bool `json:"rewriteQuery,omitempty"` // OpenAI provider: rewrite query
131135
VectorStoreNameRegex string `json:"vectorStoreNameRegex,omitempty"` // OpenAI provider: vector store name regex
132136
VectorStoreMetadataKey string `json:"vectorStoreMetadataKey,omitempty"` // OpenAI provider: vector store metadata key
133137
VectorStoreMetadataValue string `json:"vectorStoreMetadataValue,omitempty"` // OpenAI provider: vector store metadata value
138+
BucketName string `json:"bucketName,omitempty"` // S3 provider: S3 bucket name
139+
Region string `json:"region,omitempty"` // S3 provider: AWS region
140+
DateFilterField string `json:"dateFilterField,omitempty"` // Date filter metadata field name
141+
DateRangeWindowDays int `json:"dateRangeWindowDays,omitempty"` // Days to expand date range backward (default: 7)
142+
}
143+
144+
// RAGEmbeddingProviderConfig contains embedding provider-specific settings
145+
type RAGEmbeddingProviderConfig struct {
146+
APIKey string `json:"apiKey,omitempty"` // API key for the embedding provider
134147
}
135148

136149
// MonitoringConfig contains monitoring and observability settings

internal/config/validation.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,12 @@ func (c *Config) SubstituteEnvironmentVariables() {
194194
c.Observability.ServiceName = substituteEnvVars(c.Observability.ServiceName)
195195
c.Observability.ServiceVersion = substituteEnvVars(c.Observability.ServiceVersion)
196196

197+
// Substitute in RAG Embedding Providers configuration
198+
for name, provider := range c.RAG.EmbeddingProviders {
199+
provider.APIKey = substituteEnvVars(provider.APIKey)
200+
c.RAG.EmbeddingProviders[name] = provider
201+
}
202+
197203
}
198204

199205
// substituteEnvVars replaces ${VAR_NAME} patterns with environment variable values

internal/handlers/llm_mcp_bridge.go

Lines changed: 37 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -224,54 +224,62 @@ func NewLLMMCPBridgeFromClientsWithLogLevel(mcpClients interface{}, stdLogger *l
224224
return NewLLMMCPBridgeWithLogLevel(interfaceClients, stdLogger, discoveredTools, logLevel, llmRegistry, cfg)
225225
}
226226

227-
// ProcessLLMResponse processes an LLM response, expecting a specific JSON tool call format.
228-
// It no longer uses natural language detection.
229-
func (b *LLMMCPBridge) ProcessLLMResponse(ctx context.Context, llmResponse *llms.ContentChoice, _ string, extraArgs map[string]interface{}) (string, error) {
227+
// ExtractToolCall extracts tool call information from LLM response
228+
// Returns nil if no tool call is detected
229+
func (b *LLMMCPBridge) ExtractToolCall(llmResponse *llms.ContentChoice) (*ToolCall, error) {
230230
var toolCall *ToolCall
231231
var err error
232+
233+
// Check for native tool calls first
232234
funcCall := llmResponse.FuncCall
233-
// Check for a tool call in JSON format
234235
if len(llmResponse.ToolCalls) > 0 {
235236
funcCall = llmResponse.ToolCalls[0].FunctionCall
236237
}
237238

238239
if funcCall != nil {
239240
toolCall, err = b.getToolCall(funcCall)
240241
if err != nil {
241-
return "", err
242+
return nil, err
242243
}
243244
} else {
245+
// Fallback: try to detect JSON tool call in Content
244246
toolCall = b.detectSpecificJSONToolCall(llmResponse.Content)
245247
}
246248

247-
if toolCall != nil {
248-
// Execute the tool call
249-
result, err := b.executeToolCall(ctx, toolCall, extraArgs)
250-
if err != nil {
251-
// Check if it's already a domain error
252-
var errorMessage string
253-
if customErrors.IsDomainError(err) {
254-
// Extract structured information from the domain error
255-
code, _ := customErrors.GetErrorCode(err)
256-
b.logger.ErrorKV("Failed to execute tool call",
257-
"error", err.Error(),
258-
"error_code", code,
259-
"tool", toolCall.Tool)
260-
errorMessage = fmt.Sprintf("Error executing tool call: %v (code: %s)", err, code)
261-
} else {
262-
b.logger.ErrorKV("Failed to execute tool call",
263-
"error", err.Error(),
264-
"tool", toolCall.Tool)
265-
errorMessage = fmt.Sprintf("Error executing tool call: %v", err)
266-
}
249+
return toolCall, nil
250+
}
267251

268-
return errorMessage, nil
252+
// ExecuteToolCall executes a tool call and returns the result
253+
func (b *LLMMCPBridge) ExecuteToolCall(ctx context.Context, toolCall *ToolCall, extraArgs map[string]interface{}) (string, error) {
254+
if toolCall == nil {
255+
return "", fmt.Errorf("toolCall cannot be nil")
256+
}
257+
258+
// Execute the tool call
259+
result, err := b.executeToolCall(ctx, toolCall, extraArgs)
260+
if err != nil {
261+
// Check if it's already a domain error
262+
var errorMessage string
263+
if customErrors.IsDomainError(err) {
264+
// Extract structured information from the domain error
265+
code, _ := customErrors.GetErrorCode(err)
266+
b.logger.ErrorKV("Failed to execute tool call",
267+
"error", err.Error(),
268+
"error_code", code,
269+
"tool", toolCall.Tool)
270+
errorMessage = err.Error()
271+
} else {
272+
// Wrap as domain error
273+
domainErr := customErrors.WrapMCPError(err, "tool_execution_failed",
274+
fmt.Sprintf("Failed to execute tool '%s'", toolCall.Tool))
275+
b.logger.ErrorKV("Failed to execute tool call", "error", domainErr.Error(), "tool", toolCall.Tool)
276+
errorMessage = domainErr.Error()
269277
}
270-
return result, nil
278+
return "", fmt.Errorf("%s", errorMessage)
271279
}
272280

273-
// Just return the LLM response as-is if no tool call was detected
274-
return llmResponse.Content, nil
281+
b.logger.DebugKV("Tool call executed successfully", "tool", toolCall.Tool, "result_length", len(result))
282+
return result, nil
275283
}
276284

277285
// ToolCall represents the expected JSON structure for a tool call from the LLM

internal/observability/langfuse.go

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -197,17 +197,23 @@ func (p *LangfuseProvider) SetOutput(span OtelTrace.Span, output string) {
197197
}
198198

199199
func (p *LangfuseProvider) SetTokenUsage(span OtelTrace.Span, promptTokens, completionTokens, reasoningTokens, totalTokens int) {
200-
// Langfuse usage format
200+
// Langfuse usage format - uses "input", "output", "total" field names
201+
// Map our standard token names to Langfuse's expected format
201202
usageDetails := map[string]int{
202-
"prompt_tokens": promptTokens,
203-
"completion_tokens": completionTokens,
204-
"total_tokens": totalTokens,
205-
"reasoning_tokens": reasoningTokens,
203+
"input": promptTokens,
204+
"output": completionTokens,
205+
"total": totalTokens,
206+
}
207+
208+
// Add reasoning tokens as a separate metadata field since Langfuse doesn't have a standard field for it
209+
if reasoningTokens > 0 {
210+
usageDetails["reasoning_tokens"] = reasoningTokens
206211
}
207212

208213
if usageJSON, err := json.Marshal(usageDetails); err == nil {
209214
span.SetAttributes(
210215
attribute.String("langfuse.observation.usage_details", string(usageJSON)),
216+
// Also set OpenTelemetry standard fields for compatibility
211217
attribute.Int("llm.token_count.prompt_tokens", promptTokens),
212218
attribute.Int("llm.token_count.completion_tokens", completionTokens),
213219
attribute.Int("llm.token_count.total_tokens", totalTokens),

0 commit comments

Comments
 (0)