This commit is contained in:
19
services/knowledge/Dockerfile
Normal file
19
services/knowledge/Dockerfile
Normal file
@@ -0,0 +1,19 @@
|
||||
# See services/iam/Dockerfile for the pattern explanation. This image's
|
||||
# counterpart in docker-compose.yml uses pgvector/pgvector as its
|
||||
# Postgres image (plain postgres lacks the vector extension).
|
||||
FROM golang:1.23-bookworm AS build
|
||||
WORKDIR /src
|
||||
COPY go.mod go.sum* ./
|
||||
COPY gen/ gen/
|
||||
COPY pkg/ pkg/
|
||||
COPY services/knowledge/ services/knowledge/
|
||||
RUN go mod tidy
|
||||
RUN CGO_ENABLED=0 go build -o /out/knowledge-server ./services/knowledge/cmd/server
|
||||
|
||||
FROM gcr.io/distroless/static-debian12:nonroot
|
||||
WORKDIR /app
|
||||
COPY --from=build /out/knowledge-server ./knowledge-server
|
||||
COPY --from=build /src/services/knowledge/migrations ./migrations
|
||||
ENV ENKLY_KNOWLEDGE_GRPC_ADDR=:9005
|
||||
EXPOSE 9005
|
||||
ENTRYPOINT ["./knowledge-server"]
|
||||
109
services/knowledge/cmd/server/main.go
Normal file
109
services/knowledge/cmd/server/main.go
Normal file
@@ -0,0 +1,109 @@
|
||||
// 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()
|
||||
}
|
||||
95
services/knowledge/internal/application/service.go
Normal file
95
services/knowledge/internal/application/service.go
Normal file
@@ -0,0 +1,95 @@
|
||||
// Package application holds knowledge's use cases — the content of
|
||||
// internal/rag/store.go's IndexFile/Search methods, now split from
|
||||
// their SQL (internal/infrastructure/postgres) so this package depends
|
||||
// only on domain.DocumentRepository/domain.Embedder.
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
|
||||
"enkly/services/knowledge/internal/domain"
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
Documents domain.DocumentRepository
|
||||
Embedder domain.Embedder
|
||||
}
|
||||
|
||||
func hashContent(content string) string {
|
||||
sum := sha256.Sum256([]byte(content))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
// upsertChunk replaces rag.Store.UpsertChunk: embeds a single chunk and
|
||||
// stores it, skipping the embedding call entirely if the chunk's
|
||||
// content hasn't changed since it was last indexed.
|
||||
func (s *Service) upsertChunk(ctx context.Context, filePath string, chunkIndex int, content string) error {
|
||||
hash := hashContent(content)
|
||||
|
||||
existing, err := s.Documents.ExistingContentHash(ctx, filePath, chunkIndex)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if existing == hash {
|
||||
return nil // unchanged, skip re-embedding
|
||||
}
|
||||
|
||||
vec, err := s.Embedder.Embed(ctx, content)
|
||||
if err != nil {
|
||||
return fmt.Errorf("knowledge: embed %s chunk %d: %w", filePath, chunkIndex, err)
|
||||
}
|
||||
if len(vec) != domain.EmbeddingDim {
|
||||
return fmt.Errorf("knowledge: embedding model returned %d-dimensional vectors but the schema declares vector(%d) — "+
|
||||
"use a %d-dim model (the default nomic-embed-text is) or change the column and re-index",
|
||||
len(vec), domain.EmbeddingDim, domain.EmbeddingDim)
|
||||
}
|
||||
if err := s.Documents.Upsert(ctx, filePath, chunkIndex, content, hash, vec); err != nil {
|
||||
return fmt.Errorf("knowledge: upsert %s chunk %d: %w", filePath, chunkIndex, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// IndexFile replaces rag.Store.IndexFile.
|
||||
func (s *Service) IndexFile(ctx context.Context, filePath, content string) (indexed, skipped int, err error) {
|
||||
chunks := domain.ChunkText(content)
|
||||
if err := s.Documents.DeleteChunksFrom(ctx, filePath, len(chunks)); err != nil {
|
||||
return 0, 0, fmt.Errorf("knowledge: prune stale chunks for %s: %w", filePath, err)
|
||||
}
|
||||
for i, chunk := range chunks {
|
||||
before, _ := s.Documents.ExistingContentHash(ctx, filePath, i)
|
||||
if err := s.upsertChunk(ctx, filePath, i, chunk); err != nil {
|
||||
return indexed, skipped, err
|
||||
}
|
||||
if before == hashContent(chunk) {
|
||||
skipped++
|
||||
} else {
|
||||
indexed++
|
||||
}
|
||||
}
|
||||
return indexed, skipped, nil
|
||||
}
|
||||
|
||||
// PruneMissing replaces rag.Store.pruneMissing.
|
||||
func (s *Service) PruneMissing(ctx context.Context, keepFilePaths []string) (int, error) {
|
||||
return s.Documents.PruneMissing(ctx, keepFilePaths)
|
||||
}
|
||||
|
||||
// Search replaces rag.Store.Search.
|
||||
func (s *Service) Search(ctx context.Context, query string, topK int) ([]domain.CodeChunk, error) {
|
||||
if topK <= 0 {
|
||||
topK = 6
|
||||
}
|
||||
vec, err := s.Embedder.Embed(ctx, query)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("knowledge: embed query: %w", err)
|
||||
}
|
||||
return s.Documents.Search(ctx, vec, topK)
|
||||
}
|
||||
|
||||
// ListIndexedFiles replaces db.ListIndexedFiles.
|
||||
func (s *Service) ListIndexedFiles(ctx context.Context, limit int) ([]string, error) {
|
||||
return s.Documents.ListIndexedFiles(ctx, limit)
|
||||
}
|
||||
120
services/knowledge/internal/domain/chunk.go
Normal file
120
services/knowledge/internal/domain/chunk.go
Normal file
@@ -0,0 +1,120 @@
|
||||
package domain
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const (
|
||||
// chunkMaxChars is the target upper bound for one chunk. ~1,800 chars
|
||||
// is roughly 450–600 tokens of code: small enough that a similarity
|
||||
// hit points at a specific function rather than a whole file, large
|
||||
// enough that a function usually fits in one chunk. It is well under
|
||||
// nomic-embed-text's 2,048-token context, so nothing is truncated.
|
||||
chunkMaxChars = 1800
|
||||
|
||||
// chunkOverlapChars of trailing context (whole lines) is repeated at
|
||||
// the start of the next chunk, so a match that straddles a boundary
|
||||
// still lands in at least one chunk with its surroundings.
|
||||
chunkOverlapChars = 240
|
||||
)
|
||||
|
||||
// ChunkText splits source text into overlapping, line-aligned chunks of
|
||||
// at most ~chunkMaxChars characters. Whitespace-only input yields nil.
|
||||
//
|
||||
// Chunks break on line boundaries (never mid-line, except for a single
|
||||
// line longer than the limit, e.g. minified JS, which is split at rune
|
||||
// boundaries) so each chunk reads as valid-looking code.
|
||||
func ChunkText(content string) []string {
|
||||
if strings.TrimSpace(content) == "" {
|
||||
return nil
|
||||
}
|
||||
content = strings.ReplaceAll(content, "\r\n", "\n")
|
||||
|
||||
lines := strings.Split(strings.TrimRight(content, "\n"), "\n")
|
||||
|
||||
var (
|
||||
chunks []string
|
||||
cur []string
|
||||
curSize int
|
||||
)
|
||||
|
||||
flush := func() {
|
||||
text := strings.Join(cur, "\n")
|
||||
if strings.TrimSpace(text) != "" {
|
||||
chunks = append(chunks, text)
|
||||
}
|
||||
}
|
||||
|
||||
// startNext seeds the next chunk with up to chunkOverlapChars of the
|
||||
// previous chunk's trailing lines.
|
||||
startNext := func() {
|
||||
var keep []string
|
||||
size := 0
|
||||
for i := len(cur) - 1; i >= 0; i-- {
|
||||
l := len(cur[i]) + 1
|
||||
if size+l > chunkOverlapChars {
|
||||
break
|
||||
}
|
||||
keep = append([]string{cur[i]}, keep...)
|
||||
size += l
|
||||
}
|
||||
cur, curSize = keep, size
|
||||
}
|
||||
|
||||
for _, line := range lines {
|
||||
for _, piece := range splitLongLine(line, chunkMaxChars) {
|
||||
add := len(piece) + 1
|
||||
if curSize+add > chunkMaxChars && len(cur) > 0 {
|
||||
flush()
|
||||
startNext()
|
||||
// If the overlap plus this piece still overflows, drop the overlap:
|
||||
// guarantees forward progress for pathological inputs.
|
||||
if curSize+add > chunkMaxChars {
|
||||
cur, curSize = nil, 0
|
||||
}
|
||||
}
|
||||
cur = append(cur, piece)
|
||||
curSize += add
|
||||
}
|
||||
}
|
||||
if len(cur) > 0 {
|
||||
flush()
|
||||
}
|
||||
return dedupeTail(chunks)
|
||||
}
|
||||
|
||||
// splitLongLine breaks a single line longer than max bytes into pieces of
|
||||
// at most max bytes, cutting only at rune boundaries.
|
||||
func splitLongLine(line string, max int) []string {
|
||||
if len(line) <= max {
|
||||
return []string{line}
|
||||
}
|
||||
var out []string
|
||||
for len(line) > max {
|
||||
cut := max
|
||||
for cut > 0 && !utf8.RuneStart(line[cut]) {
|
||||
cut--
|
||||
}
|
||||
if cut == 0 { // no rune boundary found (invalid UTF-8): hard cut
|
||||
cut = max
|
||||
}
|
||||
out = append(out, line[:cut])
|
||||
line = line[cut:]
|
||||
}
|
||||
if line != "" {
|
||||
out = append(out, line)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// dedupeTail drops a final chunk that is nothing but overlap already
|
||||
// contained in the previous chunk (which happens when the last flush
|
||||
// carried only re-seeded lines).
|
||||
func dedupeTail(chunks []string) []string {
|
||||
n := len(chunks)
|
||||
if n >= 2 && strings.HasSuffix(chunks[n-2], chunks[n-1]) {
|
||||
return chunks[:n-1]
|
||||
}
|
||||
return chunks
|
||||
}
|
||||
45
services/knowledge/internal/domain/document.go
Normal file
45
services/knowledge/internal/domain/document.go
Normal file
@@ -0,0 +1,45 @@
|
||||
package domain
|
||||
|
||||
import "context"
|
||||
|
||||
// EmbeddingDim is the vector width migrations/0001_init.sql declares
|
||||
// for documents.embedding (vector(768), matching nomic-embed-text).
|
||||
const EmbeddingDim = 768
|
||||
|
||||
// CodeChunk is a single embedded chunk of a source file, as returned by
|
||||
// a similarity search — identical in shape to the monolith's
|
||||
// models.CodeChunk.
|
||||
type CodeChunk struct {
|
||||
ID int64
|
||||
FilePath string
|
||||
ChunkIndex int
|
||||
Content string
|
||||
Distance float64 // cosine distance; lower = more relevant
|
||||
}
|
||||
|
||||
// Embedder is the subset of *vllmkit.Client the application layer
|
||||
// needs — defined here (rather than depending on the concrete client
|
||||
// type) purely to keep the dependency explicit and mockable in tests,
|
||||
// exactly like the monolith's rag.Embedder.
|
||||
type Embedder interface {
|
||||
Embed(ctx context.Context, text string) ([]float32, error)
|
||||
}
|
||||
|
||||
// DocumentRepository persists chunked, embedded documents — the
|
||||
// pgvector-backed transplant of internal/rag/store.go +
|
||||
// internal/rag/index.go's pruneMissing.
|
||||
type DocumentRepository interface {
|
||||
// ExistingContentHash returns the content_hash stored for
|
||||
// (filePath, chunkIndex), or "" if no such row exists — used to
|
||||
// skip re-embedding an unchanged chunk (see UpsertChunk's doc
|
||||
// comment in the monolith).
|
||||
ExistingContentHash(ctx context.Context, filePath string, chunkIndex int) (string, error)
|
||||
Upsert(ctx context.Context, filePath string, chunkIndex int, content, contentHash string, embedding []float32) error
|
||||
// DeleteChunksFrom removes any chunk_index >= fromIndex for
|
||||
// filePath — stale trailing chunks left over from a previous,
|
||||
// longer version of the file.
|
||||
DeleteChunksFrom(ctx context.Context, filePath string, fromIndex int) error
|
||||
Search(ctx context.Context, embedding []float32, topK int) ([]CodeChunk, error)
|
||||
PruneMissing(ctx context.Context, keepFilePaths []string) (int, error)
|
||||
ListIndexedFiles(ctx context.Context, limit int) ([]string, error)
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
// Package postgres implements knowledge's DocumentRepository against
|
||||
// the pgvector-backed knowledge database. Transplant of
|
||||
// internal/rag/store.go's SQL (vectorLiteral and all) plus
|
||||
// internal/rag/index.go's pruneMissing, and internal/db.DB.
|
||||
// ListIndexedFiles.
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/lib/pq"
|
||||
|
||||
"enkly/pkg/pgkit"
|
||||
"enkly/services/knowledge/internal/domain"
|
||||
)
|
||||
|
||||
// vectorLiteral formats a float32 embedding as the string pgvector
|
||||
// expects for input, e.g. "[0.011,-0.42,...]" — lib/pq has no native
|
||||
// vector type, so it's sent as text and cast with ::vector in SQL.
|
||||
func vectorLiteral(v []float32) string {
|
||||
parts := make([]string, len(v))
|
||||
for i, f := range v {
|
||||
parts[i] = strconv.FormatFloat(float64(f), 'f', 8, 32)
|
||||
}
|
||||
return "[" + strings.Join(parts, ",") + "]"
|
||||
}
|
||||
|
||||
type DocumentRepo struct {
|
||||
DB *pgkit.DB
|
||||
}
|
||||
|
||||
func NewDocumentRepo(db *pgkit.DB) *DocumentRepo { return &DocumentRepo{DB: db} }
|
||||
|
||||
func (r *DocumentRepo) ExistingContentHash(ctx context.Context, filePath string, chunkIndex int) (string, error) {
|
||||
var hash string
|
||||
err := r.DB.QueryRowContext(ctx,
|
||||
`SELECT content_hash FROM documents WHERE file_path = $1 AND chunk_index = $2`,
|
||||
filePath, chunkIndex).Scan(&hash)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return "", nil
|
||||
}
|
||||
return hash, err
|
||||
}
|
||||
|
||||
func (r *DocumentRepo) Upsert(ctx context.Context, filePath string, chunkIndex int, content, contentHash string, embedding []float32) error {
|
||||
_, err := r.DB.ExecContext(ctx, `
|
||||
INSERT INTO documents (file_path, chunk_index, content, content_hash, embedding, updated_at)
|
||||
VALUES ($1, $2, $3, $4, $5::vector, now())
|
||||
ON CONFLICT (file_path, chunk_index)
|
||||
DO UPDATE SET content = EXCLUDED.content,
|
||||
content_hash = EXCLUDED.content_hash,
|
||||
embedding = EXCLUDED.embedding,
|
||||
updated_at = now()`,
|
||||
filePath, chunkIndex, content, contentHash, vectorLiteral(embedding))
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *DocumentRepo) DeleteChunksFrom(ctx context.Context, filePath string, fromIndex int) error {
|
||||
_, err := r.DB.ExecContext(ctx,
|
||||
`DELETE FROM documents WHERE file_path = $1 AND chunk_index >= $2`, filePath, fromIndex)
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *DocumentRepo) Search(ctx context.Context, embedding []float32, topK int) ([]domain.CodeChunk, error) {
|
||||
lit := vectorLiteral(embedding)
|
||||
rows, err := r.DB.QueryContext(ctx, `
|
||||
SELECT id, file_path, chunk_index, content, embedding <=> $1::vector AS distance
|
||||
FROM documents
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $2`, lit, topK)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var out []domain.CodeChunk
|
||||
for rows.Next() {
|
||||
var c domain.CodeChunk
|
||||
if err := rows.Scan(&c.ID, &c.FilePath, &c.ChunkIndex, &c.Content, &c.Distance); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, c)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (r *DocumentRepo) PruneMissing(ctx context.Context, keep []string) (int, error) {
|
||||
if keep == nil {
|
||||
keep = []string{}
|
||||
}
|
||||
res, err := r.DB.ExecContext(ctx,
|
||||
`DELETE FROM documents WHERE NOT (file_path = ANY($1))`, pq.Array(keep))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
n, err := res.RowsAffected()
|
||||
return int(n), err
|
||||
}
|
||||
|
||||
func (r *DocumentRepo) ListIndexedFiles(ctx context.Context, limit int) ([]string, error) {
|
||||
if limit <= 0 {
|
||||
limit = 200
|
||||
}
|
||||
rows, err := r.DB.QueryContext(ctx,
|
||||
`SELECT DISTINCT file_path FROM documents ORDER BY file_path LIMIT $1`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var out []string
|
||||
for rows.Next() {
|
||||
var p string
|
||||
if err := rows.Scan(&p); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, p)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
73
services/knowledge/internal/interfaces/grpc/server.go
Normal file
73
services/knowledge/internal/interfaces/grpc/server.go
Normal file
@@ -0,0 +1,73 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
knowledgev1 "enkly/gen/knowledge/v1"
|
||||
"enkly/pkg/grpcerr"
|
||||
"enkly/pkg/vllmkit"
|
||||
"enkly/services/knowledge/internal/application"
|
||||
"enkly/services/knowledge/internal/domain"
|
||||
)
|
||||
|
||||
type Server struct {
|
||||
knowledgev1.UnimplementedKnowledgeServiceServer
|
||||
Svc *application.Service
|
||||
}
|
||||
|
||||
func New(svc *application.Service) *Server { return &Server{Svc: svc} }
|
||||
|
||||
// fromErr maps an application-layer error to the right gRPC status: an
|
||||
// unreachable vllm-embed deployment (Pending, cold-starting,
|
||||
// mid-rollout — see vllmkit.ErrUpstreamUnavailable) is
|
||||
// codes.Unavailable, safe for a well-behaved caller to retry, not the
|
||||
// opaque codes.Internal every other failure gets.
|
||||
func fromErr(err error) error {
|
||||
if vllmkit.IsUpstreamUnavailable(err) {
|
||||
return grpcerr.Unavailable(err)
|
||||
}
|
||||
return grpcerr.Internal(err)
|
||||
}
|
||||
|
||||
func toProtoChunk(c domain.CodeChunk) *knowledgev1.CodeChunk {
|
||||
return &knowledgev1.CodeChunk{
|
||||
Id: c.ID, FilePath: c.FilePath, ChunkIndex: int32(c.ChunkIndex),
|
||||
Content: c.Content, Distance: c.Distance,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) IndexFile(ctx context.Context, req *knowledgev1.IndexFileRequest) (*knowledgev1.IndexFileResponse, error) {
|
||||
indexed, skipped, err := s.Svc.IndexFile(ctx, req.GetFilePath(), req.GetContent())
|
||||
if err != nil {
|
||||
return nil, fromErr(err)
|
||||
}
|
||||
return &knowledgev1.IndexFileResponse{ChunksIndexed: int32(indexed), ChunksSkipped: int32(skipped)}, nil
|
||||
}
|
||||
|
||||
func (s *Server) PruneMissing(ctx context.Context, req *knowledgev1.PruneMissingRequest) (*knowledgev1.PruneMissingResponse, error) {
|
||||
n, err := s.Svc.PruneMissing(ctx, req.GetKeepFilePaths())
|
||||
if err != nil {
|
||||
return nil, grpcerr.Internal(err)
|
||||
}
|
||||
return &knowledgev1.PruneMissingResponse{RowsDeleted: int32(n)}, nil
|
||||
}
|
||||
|
||||
func (s *Server) Search(ctx context.Context, req *knowledgev1.SearchRequest) (*knowledgev1.SearchResponse, error) {
|
||||
chunks, err := s.Svc.Search(ctx, req.GetQuery(), int(req.GetTopK()))
|
||||
if err != nil {
|
||||
return nil, fromErr(err)
|
||||
}
|
||||
resp := &knowledgev1.SearchResponse{}
|
||||
for _, c := range chunks {
|
||||
resp.Chunks = append(resp.Chunks, toProtoChunk(c))
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (s *Server) ListIndexedFiles(ctx context.Context, req *knowledgev1.ListIndexedFilesRequest) (*knowledgev1.ListIndexedFilesResponse, error) {
|
||||
files, err := s.Svc.ListIndexedFiles(ctx, int(req.GetLimit()))
|
||||
if err != nil {
|
||||
return nil, grpcerr.Internal(err)
|
||||
}
|
||||
return &knowledgev1.ListIndexedFilesResponse{FilePaths: files}, nil
|
||||
}
|
||||
27
services/knowledge/migrations/0001_init.sql
Normal file
27
services/knowledge/migrations/0001_init.sql
Normal file
@@ -0,0 +1,27 @@
|
||||
-- =====================================================================
|
||||
-- knowledge service — 0001_init.sql
|
||||
-- Database-per-service: owns `documents`, the pgvector-backed RAG
|
||||
-- index. Requires PostgreSQL 15+ with the pgvector extension available
|
||||
-- (see deploy — the knowledge database uses the pgvector/pgvector
|
||||
-- Docker image rather than plain postgres).
|
||||
-- =====================================================================
|
||||
|
||||
CREATE EXTENSION IF NOT EXISTS vector;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
file_path TEXT NOT NULL,
|
||||
chunk_index INT NOT NULL DEFAULT 0,
|
||||
content TEXT NOT NULL,
|
||||
content_hash TEXT NOT NULL, -- sha256 of `content`, used to skip re-embedding unchanged chunks
|
||||
embedding vector(768) NOT NULL, -- 768 = nomic-embed-text's output size; see domain.EmbeddingDim
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
UNIQUE (file_path, chunk_index)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS documents_embedding_hnsw_idx
|
||||
ON documents USING hnsw (embedding vector_cosine_ops);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS documents_file_path_idx
|
||||
ON documents (file_path);
|
||||
Reference in New Issue
Block a user