diff --git a/cmd/wsh/cmd/wshcmd-ai.go b/cmd/wsh/cmd/wshcmd-ai.go index 71f175089..2972e1eaa 100644 --- a/cmd/wsh/cmd/wshcmd-ai.go +++ b/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 +} diff --git a/frontend/app/store/wshclientapi.ts b/frontend/app/store/wshclientapi.ts index fba651911..2a9e5bffa 100644 --- a/frontend/app/store/wshclientapi.ts +++ b/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 { + return client.wshRpcCall("aigetmessages", data, opts); + } + // command "aisendmessage" [call] AiSendMessageCommand(client: WshClient, data: AiMessageData, opts?: RpcOpts): Promise { return client.wshRpcCall("aisendmessage", data, opts); diff --git a/frontend/app/view/waveai/waveai.tsx b/frontend/app/view/waveai/waveai.tsx index 3ed852a78..a632bb7ce 100644 --- a/frontend/app/view/waveai/waveai.tsx +++ b/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 { diff --git a/frontend/types/gotypes.d.ts b/frontend/types/gotypes.d.ts index 2570f4b4d..59ad4e3fb 100644 --- a/frontend/types/gotypes.d.ts +++ b/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; diff --git a/pkg/wshrpc/wshclient/wshclient.go b/pkg/wshrpc/wshclient/wshclient.go index 243065b9c..de9b44855 100644 --- a/pkg/wshrpc/wshclient/wshclient.go +++ b/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) diff --git a/pkg/wshrpc/wshrpctypes.go b/pkg/wshrpc/wshrpctypes.go index 3d629533c..569e3ad58 100644 --- a/pkg/wshrpc/wshrpctypes.go +++ b/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"`