File Agent

package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
)

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type ToolDef struct {
	Type     string `json:"type"`
	Function struct {
		Name        string          `json:"name"`
		Description string          `json:"description"`
		Parameters  json.RawMessage `json:"parameters"`
	} `json:"function"`
}

type ToolCall struct {
	Function struct {
		Name      string          `json:"name"`
		Arguments json.RawMessage `json:"arguments"`
	} `json:"function"`
}

type ChatResponse struct {
	Message struct {
		Role      string     `json:"role"`
		Content   string     `json:"content"`
		ToolCalls []ToolCall `json:"tool_calls,omitempty"`
	} `json:"message"`
}

func callOllama(messages []Message, tools []ToolDef) ChatResponse {
	body, _ := json.Marshal(map[string]any{
		"model":    "llama3.2",
		"messages": messages,
		"stream":   false,
		"tools":    tools,
	})
	resp, err := http.Post("http://localhost:11434/api/chat",
		"application/json", bytes.NewReader(body))
	if err != nil {
		fmt.Println("Error:", err)
		return ChatResponse{}
	}
	data, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	var result ChatResponse
	json.Unmarshal(data, &result)
	return result
}

func executeToolCall(tc ToolCall) string {
	var args map[string]string
	json.Unmarshal(tc.Function.Arguments, &args)

	switch tc.Function.Name {
	case "read_file":
		content, err := os.ReadFile(args["path"])
		if err != nil {
			return fmt.Sprintf("error: %s", err)
		}
		// Limit output to avoid flooding the context
		s := string(content)
		if len(s) > 3000 {
			s = s[:3000] + "\n... (truncated)"
		}
		return s

	case "list_directory":
		path := args["path"]
		if path == "" {
			path = "."
		}
		entries, err := os.ReadDir(path)
		if err != nil {
			return fmt.Sprintf("error: %s", err)
		}
		var names []string
		for _, e := range entries {
			name := e.Name()
			if e.IsDir() {
				name += "/"
			}
			names = append(names, name)
		}
		return strings.Join(names, "\n")
	}
	return fmt.Sprintf("unknown tool: %s", tc.Function.Name)
}

func makeToolDef(name, desc, params string) ToolDef {
	td := ToolDef{Type: "function"}
	td.Function.Name = name
	td.Function.Description = desc
	td.Function.Parameters = json.RawMessage(params)
	return td
}

func agentLoop(messages []Message, tools []ToolDef) string {
	for i := 0; i < 10; i++ {
		result := callOllama(messages, tools)

		if len(result.Message.ToolCalls) == 0 {
			return result.Message.Content
		}

		messages = append(messages, Message{
			Role:    "assistant",
			Content: result.Message.Content,
		})

		for _, tc := range result.Message.ToolCalls {
			fmt.Printf("  → %s(%s)\n", tc.Function.Name, summarizeArgs(tc))
			toolResult := executeToolCall(tc)
			messages = append(messages, Message{Role: "tool", Content: toolResult})
		}
	}
	return "Reached max iterations."
}

func summarizeArgs(tc ToolCall) string {
	var args map[string]string
	json.Unmarshal(tc.Function.Arguments, &args)
	if p, ok := args["path"]; ok {
		return p
	}
	return string(tc.Function.Arguments)
}

func main() {
	tools := []ToolDef{
		makeToolDef("read_file", "Read the contents of a file", `{
			"type": "object",
			"properties": {"path": {"type": "string", "description": "File path"}},
			"required": ["path"]
		}`),
		makeToolDef("list_directory", "List files and folders in a directory", `{
			"type": "object",
			"properties": {"path": {"type": "string", "description": "Directory path"}},
			"required": ["path"]
		}`),
	}

	system := Message{
		Role:    "system",
		Content: "You are a code assistant. Use your tools to answer questions about code. Always start by calling list_directory with path \".\" to see what files exist. Use real file paths like \".\" or \"main.go\". Never guess or make up paths.",
	}

	fmt.Println("File Agent (type 'exit' to quit)")
	fmt.Println("Ask questions about the code in this directory.")

	scanner := bufio.NewScanner(os.Stdin)
	messages := []Message{system}

	for {
		fmt.Print("\n> ")
		if !scanner.Scan() {
			break
		}
		input := scanner.Text()
		if input == "exit" || input == "" {
			break
		}

		messages = append(messages, Message{Role: "user", Content: input})
		answer := agentLoop(messages, tools)
		messages = append(messages, Message{Role: "assistant", Content: answer})
		fmt.Println("\n" + answer)
	}
}

// File Agent (type 'exit' to quit)
// Ask questions about the code in this directory.
//
// > What does main.go do?
//   → list_directory(.)
//   → read_file(main.go)
//
// The main function initializes a config and starts an HTTP server on port 8080...