6e68efaba3
- Add gRPC server configuration to the application, allowing for gRPC communication. - Introduce new Makefile targets for gRPC unit testing and proto code generation. - Update CI workflows to include gRPC tests with SQLite as the transport layer. - Refactor the sandbox design to support multi-node capabilities and improve isolation. - Enhance the service layer to facilitate internal request forwarding for gRPC APIs. This commit lays the groundwork for integrating gRPC into the Yao SDK, improving performance and scalability.
93 lines
2.7 KiB
Go
93 lines
2.7 KiB
Go
package agent
|
|
|
|
import (
|
|
"net/http"
|
|
"strings"
|
|
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/status"
|
|
|
|
"github.com/yaoapp/yao/agent"
|
|
"github.com/yaoapp/yao/agent/assistant"
|
|
agentContext "github.com/yaoapp/yao/agent/context"
|
|
"github.com/yaoapp/yao/grpc/auth"
|
|
"github.com/yaoapp/yao/grpc/pb"
|
|
)
|
|
|
|
// Handler implements the AgentStream gRPC method.
|
|
type Handler struct{}
|
|
|
|
// AgentStream resolves an assistant by ID and streams agent output as AgentChunk messages.
|
|
// Mirrors openapi/chat/completions.go GinCreateCompletions flow via context.GetGRPCAgentRequest.
|
|
func (h *Handler) AgentStream(req *pb.AgentRequest, stream grpc.ServerStreamingServer[pb.AgentChunk]) error {
|
|
ctx := stream.Context()
|
|
|
|
if req.AssistantId == "" {
|
|
return status.Error(codes.InvalidArgument, "assistant_id is required")
|
|
}
|
|
|
|
agentDSL := agent.GetAgent()
|
|
if agentDSL == nil {
|
|
return status.Error(codes.Internal, "agent DSL not initialized")
|
|
}
|
|
|
|
cache, err := agentDSL.GetCacheStore()
|
|
if err != nil {
|
|
return status.Errorf(codes.Internal, "failed to get cache store: %v", err)
|
|
}
|
|
|
|
messages, agentCtx, opts, err := agentContext.GetGRPCAgentRequest(ctx, agentContext.GRPCAgentInput{
|
|
AssistantID: req.AssistantId,
|
|
Messages: req.Messages,
|
|
Options: req.Options,
|
|
AuthInfo: auth.GetAuthorizedInfo(ctx),
|
|
Cache: cache,
|
|
Writer: &grpcStreamWriter{stream: stream, header: make(http.Header)},
|
|
})
|
|
if err != nil {
|
|
return toGRPCError(err)
|
|
}
|
|
defer agentCtx.Release()
|
|
|
|
ast, err := assistant.Get(agentCtx.AssistantID)
|
|
if err != nil {
|
|
return status.Errorf(codes.NotFound, "assistant not found: %v", err)
|
|
}
|
|
|
|
_, err = ast.Stream(agentCtx, messages, opts)
|
|
if err != nil {
|
|
return status.Errorf(codes.Internal, "agent stream failed: %v", err)
|
|
}
|
|
|
|
return stream.Send(&pb.AgentChunk{Done: true})
|
|
}
|
|
|
|
func toGRPCError(err error) error {
|
|
msg := err.Error()
|
|
if strings.Contains(msg, "is required") ||
|
|
strings.Contains(msg, "must not be empty") ||
|
|
strings.Contains(msg, "invalid") {
|
|
return status.Error(codes.InvalidArgument, msg)
|
|
}
|
|
return status.Error(codes.Internal, msg)
|
|
}
|
|
|
|
// grpcStreamWriter bridges agent/context.Writer (http.ResponseWriter) to gRPC stream.
|
|
type grpcStreamWriter struct {
|
|
stream grpc.ServerStreamingServer[pb.AgentChunk]
|
|
header http.Header
|
|
code int
|
|
}
|
|
|
|
func (w *grpcStreamWriter) Header() http.Header { return w.header }
|
|
func (w *grpcStreamWriter) WriteHeader(statusCode int) { w.code = statusCode }
|
|
func (w *grpcStreamWriter) Write(data []byte) (int, error) {
|
|
if err := w.stream.Send(&pb.AgentChunk{Data: data}); err != nil {
|
|
return 0, err
|
|
}
|
|
return len(data), nil
|
|
}
|
|
|
|
// Flush implements http.Flusher for streaming compatibility.
|
|
func (w *grpcStreamWriter) Flush() {}
|