Add an RPC to generate completions - Support FrontEnd using the conne…

…ct protocol (#198)

* We want the frontend to use the connect protocol and connect SDKs to
send RPCs to the backend.
* Currently its using gRPC and the service is defined in the vscode
protos

https://github.com/stateful/runme/blob/main/pkg/api/proto/runme/ai/v1alpha1/ai.proto

* This PR adds a new RPC defined in the Foyle protos that the front end
can use to send completion RPCs.
* The handler implements the connect protocol so the client will be able
to use the connect protocol to connect to the server.
* Related to #173

app/pkg/agent/agent.go

@@ -6,6 +6,8 @@ import (
 	"strings"
 	"sync"
 
+	"google.golang.org/protobuf/encoding/protojson"
+
 	"go.opentelemetry.io/otel/attribute"
 
 	"connectrpc.com/connect"
@@ -414,6 +416,49 @@ func (a *Agent) StreamGenerate(ctx context.Context, stream *connect.BidiStream[v
 	}
 }
 
+func (a *Agent) GenerateCells(ctx context.Context, req *connect.Request[v1alpha1.GenerateCellsRequest]) (*connect.Response[v1alpha1.GenerateCellsResponse], error) {
+	span := trace.SpanFromContext(ctx)
+	log := logs.FromContext(ctx)
+	// We don't update the logger in the context because that will happen in the agent.Generate method and we
+	// would end up duplicating the traceId key
+	log = log.WithValues("traceId", span.SpanContext().TraceID())
+
+	log.Info("Runme.Generate")
+
+	// Convert the request to the agent format
+	doc, err := converters.NotebookToDoc(req.Msg.Notebook)
+	if err != nil {
+		reqJson, jsonErr := protojson.Marshal(req.Msg)
+		if err != nil {
+			log.Error(jsonErr, "Failed to marshal request")
+		}
+		log.Error(err, "Failed to convert runme notebook to doc", "request", reqJson)
+		return nil, err
+	}
+	agentReq := &v1alpha1.GenerateRequest{
+		Doc: doc,
+	}
+
+	// Call the agent
+	agentResp, err := a.Generate(ctx, agentReq)
+	if err != nil {
+		log.Error(err, "Agent.Generate failed")
+		return nil, err
+	}
+
+	// Convert the agent response to the runme format
+	cells, err := converters.BlocksToCells(agentResp.GetBlocks())
+	if err != nil {
+		log.Error(err, "Failed to convert agent blocks to cells")
+		return nil, err
+	}
+	resp := &v1alpha1.GenerateCellsResponse{
+		Cells: cells,
+	}
+
+	return connect.NewResponse[v1alpha1.GenerateCellsResponse](resp), nil
+}
+
 // createCompletion is a helper function to create a single completion as part of a stream.
 func (a *Agent) createCompletion(ctx context.Context, generateRequest *v1alpha1.GenerateRequest, notebookUri string, selectedCell int32, contextID string) (*v1alpha1.StreamGenerateResponse, error) {
 	span := trace.SpanFromContext(ctx)

app/pkg/runme/proxy.go

@@ -14,6 +14,8 @@ import (
 )
 
 // Proxy is a proxy to the agent that converts the agent responses into the runme protocol.
+// TODO(https://github.com/jlewi/foyle/issues/173): We can get rid of this once runme has a release which
+// migrates to using the connect protocol.
 type Proxy struct {
 	agent *agent.Agent
 	aiv1alpha1.UnimplementedAIServiceServer

protos/README.md

@@ -112,7 +112,9 @@ buf push
 There are two ways I think we can iterate on the vscode extension when making changes to the proto
 
 1. Publish to the [BSR using labels](https://buf.build/docs/bsr/module/publish#pushing-from-a-local-workspace)
+
    * I think we can use labels to create the equivalent of a development branch of the SDK
 
 2. We can update `buf.gen.yaml` to output to a local path inside the vscode extension directory
+
    * See [buf.gen.yaml](https://github.com/jlewi/foyle/blob/9663fb81a36ab63876c33873cf4726dc8ef80092/protos/buf.gen.yaml#L28)

protos/foyle/v1alpha1/agent.proto

@@ -61,6 +61,9 @@ service AIService {
   // StreamGenerate is a bidirectional streaming RPC for generating completions
   rpc StreamGenerate (stream StreamGenerateRequest) returns (stream StreamGenerateResponse) {}
 
+  // GenerateCells uses the AI to generate cells to insert into the notebook.
+  rpc GenerateCells(GenerateCellsRequest) returns (GenerateCellsResponse) {}
+
   // N.B. This is for testing only. Wanted to add a non streaming response which we can use to verify things are working.
   rpc Status(StatusRequest) returns (StatusResponse) {}
 }
@@ -116,6 +119,14 @@ message StreamGenerateResponse {
   string context_id = 5;
 }
 
+message GenerateCellsRequest {
+  runme.parser.v1.Notebook notebook = 1;
+}
+
+message GenerateCellsResponse {
+  repeated runme.parser.v1.Cell cells = 1;
+}
+
 enum AIServiceStatus {
   UNKNOWN = 0;
   OK = 1;