Retrieve messages from AI block via wsh

cmd/wsh/cmd/wshcmd-ai.go

@@ -14,6 +14,7 @@ import (
 	"github.com/wavetermdev/waveterm/pkg/wshrpc"
 	"github.com/wavetermdev/waveterm/pkg/wshrpc/wshclient"
 	"github.com/wavetermdev/waveterm/pkg/wshutil"
+	"encoding/json"
 )
 
 var aiCmd = &cobra.Command{
@@ -159,3 +160,59 @@ func aiRun(cmd *cobra.Command, args []string) (rtnErr error) {
 
 	return nil
 }
+
+// AI Get Command implementation section
+var aigetCmd = &cobra.Command{
+    Use:                   "aiget [blockid]",
+    Short:                 "Get messages from an AI block",
+    Long:                  "Get messages from an AI block. Use --limit or -n to specify maximum number of messages to retrieve (default: 10)",
+    RunE:                  aigetRun,
+    PreRunE:               preRunSetupRpcClient,
+    DisableFlagsInUseLine: true,
+}
+
+var aigetLimit int
+
+func init() {
+    rootCmd.AddCommand(aigetCmd)
+    aigetCmd.Flags().IntVarP(&aigetLimit, "limit", "n", 10, "maximum number of messages to retrieve")
+}
+
+func aigetRun(cmd *cobra.Command, args []string) (rtnErr error) {
+    defer func() {
+        sendActivity("aiget", rtnErr == nil)
+    }()
+
+    // Default to "waveai" block
+    isDefaultBlock := blockArg == ""
+    if isDefaultBlock {
+        blockArg = "view@waveai"
+    }
+
+    fullORef, err := resolveSimpleId(blockArg)
+    if err != nil {
+        return fmt.Errorf("resolving block: %w", err)
+    }
+
+    // Create the route for this block
+    route := wshutil.MakeFeBlockRouteId(fullORef.OID)
+
+    messageData := wshrpc.AiGetMessagesData{
+        Limit: aigetLimit,
+    }
+
+    response, err := wshclient.AiGetMessagesCommand(RpcClient, messageData, &wshrpc.RpcOpts{
+        Route:   route,
+        Timeout: 2000,
+    })
+    if err != nil {
+        return fmt.Errorf("getting messages: %w", err)
+    }
+    jsonBytes, err := json.Marshal(response)
+    if err != nil {
+        return fmt.Errorf("marshalling response: %w", err)
+    }
+    fmt.Print(string(jsonBytes))
+
+    return nil
+}

frontend/app/store/wshclientapi.ts

@@ -12,6 +12,11 @@ class RpcApiType {
         return client.wshRpcCall("activity", data, opts);
     }
 
+    // command "aigetmessages" [call]
+    AiGetMessagesCommand(client: WshClient, data: AiGetMessagesData, opts?: RpcOpts): Promise<AiMessage[]> {
+        return client.wshRpcCall("aigetmessages", data, opts);
+    }
+
     // command "aisendmessage" [call]
     AiSendMessageCommand(client: WshClient, data: AiMessageData, opts?: RpcOpts): Promise<void> {
         return client.wshRpcCall("aisendmessage", data, opts);

frontend/app/view/waveai/waveai.tsx

@@ -59,6 +59,16 @@ class AiWshClient extends WshClient {
         }
         this.model.sendMessage(data.message);
     }
+
+    handle_aigetmessages(rh: RpcResponseHelper, data: AiGetMessagesData) {
+        const messages = globalStore.get(this.model.messagesAtom);
+        const limit = data.limit || 10;
+        const limitedMessages = messages.slice(-limit).map((msg) => ({
+            role: msg.user,
+            content: msg.text,
+        }));
+        rh.sendResponse({ data: limitedMessages });
+    }
 }
 
 export class WaveAiModel implements ViewModel {

frontend/types/gotypes.d.ts

@@ -40,6 +40,17 @@ declare global {
         conn?: {[key: string]: number};
     };
 
+    // wshrpc.AiGetMessagesData
+    type AiGetMessagesData = {
+        limit?: number;
+    };
+
+    // wshrpc.AiMessage
+    type AiMessage = {
+        role: string;
+        content: string;
+    };
+
     // wshrpc.AiMessageData
     type AiMessageData = {
         message?: string;

pkg/wshrpc/wshclient/wshclient.go

@@ -22,6 +22,12 @@ func ActivityCommand(w *wshutil.WshRpc, data wshrpc.ActivityUpdate, opts *wshrpc
 	return err
 }
 
+// command "aigetmessages", wshserver.AiGetMessagesCommand
+func AiGetMessagesCommand(w *wshutil.WshRpc, data wshrpc.AiGetMessagesData, opts *wshrpc.RpcOpts) ([]wshrpc.AiMessage, error) {
+	resp, err := sendRpcRequestCallHelper[[]wshrpc.AiMessage](w, "aigetmessages", data, opts)
+	return resp, err
+}
+
 // command "aisendmessage", wshserver.AiSendMessageCommand
 func AiSendMessageCommand(w *wshutil.WshRpc, data wshrpc.AiMessageData, opts *wshrpc.RpcOpts) error {
 	_, err := sendRpcRequestCallHelper[any](w, "aisendmessage", data, opts)

pkg/wshrpc/wshrpctypes.go

@@ -137,6 +137,7 @@ const (
 	Command_VDomUrlRequest      = "vdomurlrequest"
 
 	Command_AiSendMessage = "aisendmessage"
+	Command_AiGetMessages = "aigetmessages"
 )
 
 type RespOrErrorUnion[T any] struct {
@@ -255,6 +256,7 @@ type WshRpcInterface interface {
 
 	// ai
 	AiSendMessageCommand(ctx context.Context, data AiMessageData) error
+	AiGetMessagesCommand(ctx context.Context, data AiGetMessagesData) ([]AiMessage, error)
 
 	// proc
 	VDomRenderCommand(ctx context.Context, data vdom.VDomFrontendUpdate) chan RespOrErrorUnion[*vdom.VDomBackendUpdate]
@@ -681,6 +683,15 @@ type AiMessageData struct {
 	Message string `json:"message,omitempty"`
 }
 
+type AiMessage struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+type AiGetMessagesData struct {
+	Limit int `json:"limit,omitempty"` // If 0, defaults to 10
+}
+
 type CommandVarData struct {
 	Key      string `json:"key"`
 	Val      string `json:"val,omitempty"`