110 lines
3.6 KiB
Go
110 lines
3.6 KiB
Go
// Command server runs the knowledge microservice: a gRPC server
|
|
// exposing api/proto/knowledge/v1, backed by its own pgvector-enabled
|
|
// Postgres database and a vLLM-served embedding model. Replaces
|
|
// internal/rag/*.go.
|
|
package main
|
|
|
|
import (
|
|
"context"
|
|
"log/slog"
|
|
"net"
|
|
"os"
|
|
"os/signal"
|
|
"syscall"
|
|
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/health"
|
|
healthpb "google.golang.org/grpc/health/grpc_health_v1"
|
|
"google.golang.org/grpc/reflection"
|
|
|
|
knowledgev1 "enkly/gen/knowledge/v1"
|
|
"enkly/pkg/envkit"
|
|
"enkly/pkg/pgkit"
|
|
"enkly/pkg/vllmkit"
|
|
"enkly/services/knowledge/internal/application"
|
|
"enkly/services/knowledge/internal/infrastructure/postgres"
|
|
grpcsrv "enkly/services/knowledge/internal/interfaces/grpc"
|
|
)
|
|
|
|
func main() {
|
|
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
|
|
|
if err := envkit.LoadDotEnv(".env"); err != nil {
|
|
logger.Error("load .env", "err", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
var (
|
|
grpcAddr = envkit.String("ENKLY_KNOWLEDGE_GRPC_ADDR", ":9005")
|
|
postgresDSN = envkit.String("ENKLY_KNOWLEDGE_POSTGRES_DSN", "postgres://enkly:enkly@localhost:5432/enkly_knowledge?sslmode=disable")
|
|
migrationsPath = envkit.String("ENKLY_KNOWLEDGE_MIGRATIONS", "migrations/0001_init.sql")
|
|
|
|
// vllmEmbedURL points at the vLLM OpenAI-compatible deployment
|
|
// serving the embedding model — its own vLLM engine process,
|
|
// independent of (and typically scaled very differently from)
|
|
// the chat deployment the agent service talks to, since vLLM
|
|
// serves exactly one model per server. See k8s/11-vllm.yaml.
|
|
// This service never calls the chat model, so it only ever
|
|
// constructs a vllmkit.Client with chatBaseURL left empty.
|
|
vllmEmbedURL = envkit.String("ENKLY_VLLM_EMBED_URL", "http://localhost:8001/v1")
|
|
embedModel = envkit.String("ENKLY_EMBED_MODEL", vllmkit.DefaultEmbedModel)
|
|
// vllmAPIKey is sent as a bearer token on every embed request
|
|
// (vLLM's --api-key flag) — see k8s/02-secrets.yaml.
|
|
vllmAPIKey = envkit.String("ENKLY_VLLM_API_KEY", "")
|
|
)
|
|
|
|
db, err := pgkit.Open(postgresDSN)
|
|
if err != nil {
|
|
logger.Error("open postgres", "err", err)
|
|
os.Exit(1)
|
|
}
|
|
defer db.Close()
|
|
migrationSQL, err := os.ReadFile(migrationsPath)
|
|
if err != nil {
|
|
logger.Error("read migrations", "path", migrationsPath, "err", err)
|
|
os.Exit(1)
|
|
}
|
|
if err := db.Migrate(context.Background(), string(migrationSQL)); err != nil {
|
|
logger.Error("migrate", "err", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
// chatBaseURL is deliberately left "" (defaults to vllmkit's own
|
|
// unused-but-harmless placeholder) — this service never calls
|
|
// ChatStream, only Embed, so only embedBaseURL matters here.
|
|
llmClient := vllmkit.New("", vllmEmbedURL, "", embedModel)
|
|
llmClient.APIKey = vllmAPIKey
|
|
|
|
svc := &application.Service{
|
|
Documents: postgres.NewDocumentRepo(db),
|
|
Embedder: llmClient,
|
|
}
|
|
|
|
lis, err := net.Listen("tcp", grpcAddr)
|
|
if err != nil {
|
|
logger.Error("listen", "addr", grpcAddr, "err", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
grpcServer := grpc.NewServer()
|
|
knowledgev1.RegisterKnowledgeServiceServer(grpcServer, grpcsrv.New(svc))
|
|
healthSrv := health.NewServer()
|
|
healthpb.RegisterHealthServer(grpcServer, healthSrv)
|
|
healthSrv.SetServingStatus("", healthpb.HealthCheckResponse_SERVING)
|
|
reflection.Register(grpcServer)
|
|
|
|
go func() {
|
|
logger.Info("knowledge service listening", "addr", grpcAddr, "embed_model", embedModel)
|
|
if err := grpcServer.Serve(lis); err != nil {
|
|
logger.Error("serve", "err", err)
|
|
os.Exit(1)
|
|
}
|
|
}()
|
|
|
|
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
|
defer stop()
|
|
<-ctx.Done()
|
|
logger.Info("shutting down")
|
|
grpcServer.GracefulStop()
|
|
}
|