Agent-Native/cmd/agent-native/main.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/rpc"
	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/runtime"
	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/runtime/llm"
	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/runtime/session"
)

func main() {
	mode := flag.String("mode", "print", "print, json, or rpc")
	model := flag.String("model", "", "provider/model[:thinking]")
	prompt := flag.String("prompt", "", "user message for print and json modes")
	sessionPath := flag.String("session", "", "transcript path")
	noSession := flag.Bool("no-session", false, "do not read or write a transcript")
	export := flag.String("export", "", "write the active session as HTML and exit")
	flag.Parse()

	if *sessionPath == "" && !*noSession {
		home, _ := os.UserHomeDir()
		*sessionPath = filepath.Join(home, ".agent-native", "session.jsonl")
	}
	var store *session.Store
	if !*noSession {
		st, err := session.Open(*sessionPath)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		store = st
	}
	if *export != "" {
		if store == nil {
			fmt.Fprintln(os.Stderr, "export requires a session")
			os.Exit(1)
		}
		if err := store.ExportHTML(*export); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		return
	}

	ref, err := llm.ParseModelRef(*model)
	if *model != "" && err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	agent := runtime.NewAgentSession(runtime.Config{
		Model:         ref.Model,
		ThinkingLevel: ref.Thinking,
		Session:       store,
		Stream:        liveStream(ref),
		OnEvent: func(ev runtime.Event) {
			if *mode != "json" {
				return
			}
			raw, _ := json.Marshal(map[string]any{"type": ev.Type, "text": ev.Text, "tool": ev.Tool})
			fmt.Println(string(raw))
		},
	})
	srv := &rpc.Server{Agent: agent, Store: store, NoStore: *noSession}
	ctx := context.Background()
	switch *mode {
	case "rpc":
		in := bufio.NewScanner(os.Stdin)
		in.Buffer(make([]byte, 0, 64*1024), 8*1024*1024)
		for in.Scan() {
			resp, err := srv.Handle(ctx, in.Bytes())
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}
			os.Stdout.Write(resp)
		}
	case "print", "json":
		if *prompt == "" {
			*prompt = readAll(os.Stdin)
		}
		res, err := agent.Prompt(ctx, llm.Message{Role: llm.RoleUser, Content: *prompt})
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		if res.Err != nil {
			fmt.Fprintln(os.Stderr, res.Err)
			os.Exit(1)
		}
		if *mode == "print" {
			fmt.Println(res.Text)
		}
	default:
		fmt.Fprintln(os.Stderr, "unknown mode")
		os.Exit(1)
	}
}

func liveStream(ref llm.ModelRef) llm.StreamFunc {
	provider, ok := llm.Lookup(ref.Provider, nil)
	if !ok {
		provider = llm.Provider{Protocol: llm.ProtocolOpenAIChat, BaseURL: os.Getenv("OPENAI_BASE_URL"), APIKey: os.Getenv("OPENAI_API_KEY")}
	}
	client, err := llm.ClientFor(provider)
	if err != nil {
		return func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
			ch := make(chan llm.Event, 1)
			ch <- llm.Event{Kind: llm.KindError, Err: err, Message: llm.Message{StopReason: llm.StopError}}
			close(ch)
			return ch
		}
	}
	return llm.WithRetry(client.Stream, 2, nil)
}

func readAll(r io.Reader) string {
	raw, _ := io.ReadAll(r)
	return string(raw)
}