1package tools
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "os"
8 "runtime"
9 "time"
10
11 "github.com/cloudwego/eino/components/tool"
12 "github.com/cloudwego/eino/schema"
13
14 "github.com/cloudwego/eino/compose"
15 "github.com/cloudwego/eino/flow/agent/react"
16 "github.com/kujtimiihoxha/termai/internal/llm/models"
17 "github.com/spf13/viper"
18)
19
20type agentTool struct {
21 workingDir string
22}
23
24const (
25 AgentToolName = "agent"
26)
27
28type AgentParams struct {
29 Prompt string `json:"prompt"`
30}
31
32func taskAgentTools() []tool.BaseTool {
33 wd := viper.GetString("wd")
34 return []tool.BaseTool{
35 NewBashTool(wd),
36 NewLsTool(wd),
37 NewGlobTool(wd),
38 NewViewTool(wd),
39 NewWriteTool(wd),
40 NewEditTool(wd),
41 }
42}
43
44func NewTaskAgent(ctx context.Context) (*react.Agent, error) {
45 model, err := models.GetModel(ctx, models.ModelID(viper.GetString("models.big")))
46 if err != nil {
47 return nil, err
48 }
49 reactAgent, err := react.NewAgent(ctx, &react.AgentConfig{
50 Model: model,
51 ToolsConfig: compose.ToolsNodeConfig{
52 Tools: taskAgentTools(),
53 },
54 MaxStep: 1000,
55 })
56 if err != nil {
57 return nil, err
58 }
59
60 return reactAgent, nil
61}
62
63func TaskAgentSystemPrompt() string {
64 agentPrompt := `You are an agent for Orbitowl. Given the user's prompt, you should use the tools available to you to answer the user's question.
65
66Notes:
671. IMPORTANT: You should be concise, direct, and to the point, since your responses will be displayed on a command line interface. Answer the user's question directly, without elaboration, explanation, or details. One word answers are best. Avoid introductions, conclusions, and explanations. You MUST avoid text before/after your response, such as "The answer is <answer>.", "Here is the content of the file..." or "Based on the information provided, the answer is..." or "Here is what I will do next...".
682. When relevant, share file names and code snippets relevant to the query
693. Any file paths you return in your final response MUST be absolute. DO NOT use relative paths.
70
71Here is useful information about the environment you are running in:
72<env>
73Working directory: %s
74Platform: %s
75Today's date: %s
76</env>`
77
78 cwd, err := os.Getwd()
79 if err != nil {
80 cwd = "unknown"
81 }
82
83 platform := runtime.GOOS
84
85 switch platform {
86 case "darwin":
87 platform = "macos"
88 case "windows":
89 platform = "windows"
90 case "linux":
91 platform = "linux"
92 }
93 return fmt.Sprintf(agentPrompt, cwd, platform, time.Now().Format("1/2/2006"))
94}
95
96func (b *agentTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
97 return &schema.ToolInfo{
98 Name: AgentToolName,
99 Desc: "Launch a new agent that has access to the following tools: GlobTool, GrepTool, LS, View, ReadNotebook. When you are searching for a keyword or file and are not confident that you will find the right match on the first try, use the Agent tool to perform the search for you. For example:\n\n- If you are searching for a keyword like \"config\" or \"logger\", or for questions like \"which file does X?\", the Agent tool is strongly recommended\n- If you want to read a specific file path, use the View or GlobTool tool instead of the Agent tool, to find the match more quickly\n- If you are searching for a specific class definition like \"class Foo\", use the GlobTool tool instead, to find the match more quickly\n\nUsage notes:\n1. Launch multiple agents concurrently whenever possible, to maximize performance; to do that, use a single message with multiple tool uses\n2. When the agent is done, it will return a single message back to you. The result returned by the agent is not visible to the user. To show the user the result, you should send a text message back to the user with a concise summary of the result.\n3. Each agent invocation is stateless. You will not be able to send additional messages to the agent, nor will the agent be able to communicate with you outside of its final report. Therefore, your prompt should contain a highly detailed task description for the agent to perform autonomously and you should specify exactly what information the agent should return back to you in its final and only message to you.\n4. The agent's outputs should generally be trusted\n5. IMPORTANT: The agent can not use Bash, Replace, Edit, NotebookEditCell, so can not modify files. If you want to use these tools, use them directly instead of going through the agent.",
100 ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
101 "prompt": {
102 Type: "string",
103 Desc: "The task for the agent to perform",
104 Required: true,
105 },
106 }),
107 }, nil
108}
109
110func (b *agentTool) InvokableRun(ctx context.Context, args string, opts ...tool.Option) (string, error) {
111 var params AgentParams
112 if err := json.Unmarshal([]byte(args), ¶ms); err != nil {
113 return "", err
114 }
115 if params.Prompt == "" {
116 return "prompt is required", nil
117 }
118
119 a, err := NewTaskAgent(ctx)
120 if err != nil {
121 return "", err
122 }
123 out, err := a.Generate(
124 ctx,
125 []*schema.Message{
126 schema.SystemMessage(TaskAgentSystemPrompt()),
127 schema.UserMessage(params.Prompt),
128 },
129 )
130 if err != nil {
131 return "", err
132 }
133
134 return out.Content, nil
135}
136
137func NewAgentTool(wd string) tool.InvokableTool {
138 return &agentTool{
139 workingDir: wd,
140 }
141}