first commit
This commit is contained in:
didier committed
commit
2ffae91aa7
10 files changed
+891
No files matched your search
@@ -0,0 +1,239 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
"github.com/mark3labs/mcp-go/server"
|
||||
"github.com/qdrant/go-client/qdrant"
|
||||
)
|
||||
|
||||
var version = "0.1.0"
|
||||
|
||||
func main() {
|
||||
qdrantURL := requireEnv("QDRANT_URL")
|
||||
openAIKey := requireEnv("OPENAI_API_KEY")
|
||||
token := requireEnv("MCP_TOKEN")
|
||||
embeddingModel := os.Getenv("EMBEDDING_MODEL")
|
||||
embeddingBaseURL := os.Getenv("EMBEDDING_BASE_URL")
|
||||
port := envOrDefault("PORT", "8080")
|
||||
|
||||
qdrantPort := 6334
|
||||
if p, err := strconv.Atoi(envOrDefault("QDRANT_PORT", "6334")); err == nil {
|
||||
qdrantPort = p
|
||||
}
|
||||
qdrantAPIKey := os.Getenv("QDRANT_API_KEY")
|
||||
qdrantTLS := os.Getenv("QDRANT_TLS") == "true"
|
||||
|
||||
qdrantClient, err := qdrant.NewClient(&qdrant.Config{
|
||||
Host: qdrantURL,
|
||||
Port: qdrantPort,
|
||||
APIKey: qdrantAPIKey,
|
||||
UseTLS: qdrantTLS,
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatalf("failed to create qdrant client: %v", err)
|
||||
}
|
||||
defer qdrantClient.Close()
|
||||
|
||||
emb := newEmbedder(openAIKey, embeddingModel, embeddingBaseURL)
|
||||
store := newQdrantStore(qdrantClient, emb)
|
||||
|
||||
if err := store.ensureCollection(context.Background()); err != nil {
|
||||
log.Fatalf("failed to ensure collection: %v", err)
|
||||
}
|
||||
|
||||
mcpServer := server.NewMCPServer(
|
||||
"knowledge-mcp",
|
||||
version,
|
||||
server.WithToolCapabilities(true),
|
||||
server.WithLogging(),
|
||||
)
|
||||
|
||||
registerTools(mcpServer, store)
|
||||
|
||||
streamableServer := server.NewStreamableHTTPServer(mcpServer)
|
||||
|
||||
authHandler := authMiddleware(token, streamableServer)
|
||||
|
||||
httpServer := &http.Server{
|
||||
Addr: ":" + port,
|
||||
Handler: corsMiddleware(authHandler),
|
||||
}
|
||||
|
||||
log.Printf("MCP Knowledge Server starting on http://0.0.0.0:%s", port)
|
||||
|
||||
go func() {
|
||||
sigChan := make(chan os.Signal, 1)
|
||||
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
|
||||
sig := <-sigChan
|
||||
log.Printf("received signal %v, shutting down...", sig)
|
||||
shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := streamableServer.Shutdown(shutdownCtx); err != nil {
|
||||
log.Printf("shutdown error: %v", err)
|
||||
}
|
||||
qdrantClient.Close()
|
||||
os.Exit(0)
|
||||
}()
|
||||
|
||||
if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
log.Fatalf("server stopped: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func registerTools(mcpServer *server.MCPServer, store *qdrantStore) {
|
||||
mcpServer.AddTool(
|
||||
mcp.NewTool("store_knowledge",
|
||||
mcp.WithDescription("Store a piece of knowledge for later retrieval. Use this to persist useful information discovered during a session."),
|
||||
mcp.WithString("content",
|
||||
mcp.Required(),
|
||||
mcp.Description("The text content to store"),
|
||||
),
|
||||
mcp.WithObject("metadata",
|
||||
mcp.Description("Optional metadata as key-value pairs (e.g. {\"source\": \"docs\", \"category\": \"golang\"})"),
|
||||
),
|
||||
),
|
||||
handleStoreKnowledge(store),
|
||||
)
|
||||
|
||||
mcpServer.AddTool(
|
||||
mcp.NewTool("search_knowledge",
|
||||
mcp.WithDescription("Search previously stored knowledge using semantic search. Returns the most relevant knowledge entries for the given query."),
|
||||
mcp.WithString("query",
|
||||
mcp.Required(),
|
||||
mcp.Description("The search query"),
|
||||
),
|
||||
mcp.WithNumber("limit",
|
||||
mcp.Description("Maximum number of results to return (default: 5)"),
|
||||
),
|
||||
),
|
||||
handleSearchKnowledge(store),
|
||||
)
|
||||
}
|
||||
|
||||
func handleStoreKnowledge(store *qdrantStore) server.ToolHandlerFunc {
|
||||
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
args, ok := request.Params.Arguments.(map[string]any)
|
||||
if !ok {
|
||||
return mcp.NewToolResultError("invalid arguments"), nil
|
||||
}
|
||||
|
||||
content, ok := args["content"].(string)
|
||||
if !ok || content == "" {
|
||||
return mcp.NewToolResultError("content is required and must be a non-empty string"), nil
|
||||
}
|
||||
|
||||
var metadata map[string]any
|
||||
if rawMeta, ok := args["metadata"]; ok && rawMeta != nil {
|
||||
metaMap, ok := rawMeta.(map[string]any)
|
||||
if !ok {
|
||||
return mcp.NewToolResultError("metadata must be a key-value object"), nil
|
||||
}
|
||||
metadata = metaMap
|
||||
}
|
||||
|
||||
result, err := store.store(ctx, content, metadata)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("failed to store knowledge: %v", err)), nil
|
||||
}
|
||||
|
||||
output, _ := json.MarshalIndent(result, "", " ")
|
||||
return mcp.NewToolResultText(string(output)), nil
|
||||
}
|
||||
}
|
||||
|
||||
func handleSearchKnowledge(store *qdrantStore) server.ToolHandlerFunc {
|
||||
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
args, ok := request.Params.Arguments.(map[string]any)
|
||||
if !ok {
|
||||
return mcp.NewToolResultError("invalid arguments"), nil
|
||||
}
|
||||
|
||||
query, ok := args["query"].(string)
|
||||
if !ok || query == "" {
|
||||
return mcp.NewToolResultError("query is required and must be a non-empty string"), nil
|
||||
}
|
||||
|
||||
limit := uint64(5)
|
||||
if rawLimit := args["limit"]; rawLimit != nil {
|
||||
switch v := rawLimit.(type) {
|
||||
case float64:
|
||||
limit = uint64(v)
|
||||
case int:
|
||||
limit = uint64(v)
|
||||
case int64:
|
||||
limit = uint64(v)
|
||||
}
|
||||
}
|
||||
|
||||
results, err := store.search(ctx, query, limit)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil
|
||||
}
|
||||
|
||||
output, _ := json.MarshalIndent(results, "", " ")
|
||||
return mcp.NewToolResultText(string(output)), nil
|
||||
}
|
||||
}
|
||||
|
||||
func authMiddleware(token string, next http.Handler) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/health" {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("ok"))
|
||||
return
|
||||
}
|
||||
|
||||
authHeader := r.Header.Get("Authorization")
|
||||
if !strings.HasPrefix(authHeader, "Bearer ") || strings.TrimPrefix(authHeader, "Bearer ") != token {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
w.Write([]byte(`{"error":"unauthorized"}`))
|
||||
return
|
||||
}
|
||||
|
||||
next.ServeHTTP(w, r)
|
||||
}
|
||||
}
|
||||
|
||||
func corsMiddleware(next http.Handler) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Access-Control-Allow-Origin", "*")
|
||||
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
|
||||
w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
|
||||
|
||||
if r.Method == http.MethodOptions {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
return
|
||||
}
|
||||
|
||||
next.ServeHTTP(w, r)
|
||||
}
|
||||
}
|
||||
|
||||
func requireEnv(key string) string {
|
||||
val := os.Getenv(key)
|
||||
if val == "" {
|
||||
log.Fatalf("environment variable %s is required", key)
|
||||
}
|
||||
return val
|
||||
}
|
||||
|
||||
func envOrDefault(key, defaultVal string) string {
|
||||
val := os.Getenv(key)
|
||||
if val == "" {
|
||||
return defaultVal
|
||||
}
|
||||
return val
|
||||
}
|
||||
Reference in new issue
Block a user