74 lines
2.4 KiB
Go
74 lines
2.4 KiB
Go
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
|
|
}
|