Skip to content

Commit 4a8c531

Browse files
author
Marek Safarik
committed
Agentic loop + refactored code
Signed-off-by: Marek Safarik <[email protected]>
1 parent f4a6836 commit 4a8c531

File tree

3 files changed

+211
-91
lines changed

3 files changed

+211
-91
lines changed

‎backend/helpers.go‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,10 @@ func fetchAllAgentUUIDs() ([]string, error) {
2222
return nil, err
2323
}
2424

25+
if agents.Results.UUIDs == nil {
26+
return []string{}, nil
27+
}
28+
2529
return agents.Results.UUIDs, nil
2630
}
2731

@@ -66,16 +70,30 @@ func mapAgentToOutput(agentUUID string, agentStatus keylimeAgentStatusResponse)
6670
}
6771

6872
func mapAgentToPolicies(agentUUID string, agentStatus keylimeAgentStatusResponse) getAgentPoliciesOutput {
73+
// Ensure slices are never nil
74+
hashAlgs := agentStatus.Results.AcceptTPMHashAlgs
75+
if hashAlgs == nil {
76+
hashAlgs = []string{}
77+
}
78+
encAlgs := agentStatus.Results.AcceptTPMEncryptionAlgs
79+
if encAlgs == nil {
80+
encAlgs = []string{}
81+
}
82+
signAlgs := agentStatus.Results.AcceptTPMSigningAlgs
83+
if signAlgs == nil {
84+
signAlgs = []string{}
85+
}
86+
6987
return getAgentPoliciesOutput{
7088
AgentUUID: agentUUID,
7189
TPMPolicy: parseJSONString(agentStatus.Results.TPMPolicy),
7290
VTPMPolicy: parseJSONString(agentStatus.Results.VTPMPolicy),
7391
MetaData: parseJSONString(agentStatus.Results.MetaData),
7492
HasMeasuredBootPolicy: agentStatus.Results.HasMbRefstate != 0,
7593
HasRuntimePolicy: agentStatus.Results.HasRuntimePolicy != 0,
76-
AcceptedTPMHashAlgs: agentStatus.Results.AcceptTPMHashAlgs,
77-
AcceptedTPMEncryptionAlgs: agentStatus.Results.AcceptTPMEncryptionAlgs,
78-
AcceptedTPMSigningAlgs: agentStatus.Results.AcceptTPMSigningAlgs,
94+
AcceptedTPMHashAlgs: hashAlgs,
95+
AcceptedTPMEncryptionAlgs: encAlgs,
96+
AcceptedTPMSigningAlgs: signAlgs,
7997
}
8098
}
8199

‎backend/tools.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,9 @@ func getFailedAgents(ctx context.Context, req *mcp.CallToolRequest, input getFai
4545
return nil, getFailedAgentsOutput{}, err
4646
}
4747

48-
var failedAgents getFailedAgentsOutput
48+
failedAgents := getFailedAgentsOutput{
49+
FailedAgents: []getAgentStatusOutput{},
50+
}
4951
for _, agentUUID := range uuids {
5052
agentStatus, err := fetchAgentDetails(agentUUID)
5153
if err != nil {

‎mcp-client/main.go‎

Lines changed: 187 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package main
22

33
import (
44
"context"
5+
"fmt"
56
"log"
67
"os"
78
"os/exec"
@@ -13,133 +14,232 @@ import (
1314
"github.com/modelcontextprotocol/go-sdk/mcp"
1415
)
1516

16-
const serverPath = "../backend/server"
17+
const (
18+
serverPath = "../backend/server"
19+
mcpClientName = "mcp-client"
20+
mcpClientVersion = "v1.0.0"
21+
22+
claudeModel = anthropic.ModelClaude_3_Haiku_20240307
23+
maxTokens = 1024
24+
maxAgentTurns = 5
25+
26+
systemPrompt = `You are an autonomous agent with access to Keylime system management tools. Your goal is to help users manage and monitor their Keylime infrastructure.
27+
28+
When given a task:
29+
1. Break it down into steps if needed
30+
2. Use available tools to gather information and take actions
31+
3. Chain multiple tool calls together to accomplish complex tasks
32+
4. Provide clear explanations of what you're doing and what you found
33+
5. If you encounter failures, investigate and suggest solutions`
34+
)
1735

1836
func main() {
1937
ctx := context.Background()
20-
godotenv.Load("../.env")
38+
39+
if err := godotenv.Load("../.env"); err != nil {
40+
log.Printf("Warning: .env file not loaded: %v", err)
41+
}
2142

2243
apiKey := os.Getenv("ANTHROPIC_API_KEY")
2344
if apiKey == "" {
2445
log.Fatal("ANTHROPIC_API_KEY environment variable not set")
2546
}
2647

48+
if len(os.Args) <= 1 {
49+
log.Fatal("Usage: go run main.go <content>")
50+
}
51+
userQuery := strings.Join(os.Args[1:], " ")
52+
53+
session, err := connectToMCPServer(ctx)
54+
if err != nil {
55+
log.Fatalf("Failed to connect to MCP server: %v", err)
56+
}
57+
defer session.Close()
58+
59+
claudeTools, err := getMCPTools(ctx, session)
60+
if err != nil {
61+
log.Fatalf("Failed to get MCP tools: %v", err)
62+
}
63+
64+
anthropicClient := anthropic.NewClient(option.WithAPIKey(apiKey))
65+
66+
if err := runAgentLoop(ctx, anthropicClient, session, claudeTools, userQuery); err != nil {
67+
log.Fatalf("Agent loop failed: %v", err)
68+
}
69+
}
70+
71+
// connectToMCPServer establishes connection to the MCP server
72+
func connectToMCPServer(ctx context.Context) (*mcp.ClientSession, error) {
2773
if _, err := os.Stat(serverPath); os.IsNotExist(err) {
28-
log.Fatalf("Server not found: %s", serverPath)
74+
return nil, fmt.Errorf("server not found: %s", serverPath)
2975
}
3076

31-
// Connect to MCP server
32-
client := mcp.NewClient(&mcp.Implementation{Name: "mcp-client", Version: "v1.0.0"}, nil)
77+
client := mcp.NewClient(&mcp.Implementation{
78+
Name: mcpClientName,
79+
Version: mcpClientVersion,
80+
}, nil)
81+
3382
transport := &mcp.CommandTransport{Command: exec.Command(serverPath)}
3483
session, err := client.Connect(ctx, transport, nil)
3584
if err != nil {
36-
log.Fatal(err)
85+
return nil, fmt.Errorf("failed to connect: %w", err)
3786
}
38-
defer session.Close()
3987

40-
tools, _ := session.ListTools(ctx, &mcp.ListToolsParams{})
88+
log.Printf("Connected to MCP server: %s", serverPath)
89+
return session, nil
90+
}
4191

42-
claudeTools := make([]anthropic.ToolUnionParam, 0)
92+
// getMCPTools retrieves and converts MCP tools to Claude format
93+
func getMCPTools(ctx context.Context, session *mcp.ClientSession) ([]anthropic.ToolUnionParam, error) {
94+
tools, err := session.ListTools(ctx, &mcp.ListToolsParams{})
95+
if err != nil {
96+
return nil, fmt.Errorf("failed to list tools: %w", err)
97+
}
98+
99+
var claudeTools []anthropic.ToolUnionParam
43100
for _, tool := range tools.Tools {
44-
inputSchemaMap, ok := tool.InputSchema.(map[string]interface{})
45-
if !ok || inputSchemaMap == nil {
46-
inputSchemaMap = map[string]interface{}{}
47-
}
101+
claudeTool := convertMCPToolToClaudeTool(tool)
102+
claudeTools = append(claudeTools, claudeTool)
103+
}
48104

49-
properties := inputSchemaMap["properties"]
50-
required, _ := inputSchemaMap["required"].([]string)
105+
return claudeTools, nil
106+
}
51107

52-
toolParam := anthropic.ToolUnionParamOfTool(
53-
anthropic.ToolInputSchemaParam{
54-
Type: "object",
55-
Properties: properties,
56-
Required: required,
57-
},
58-
tool.Name,
59-
)
108+
// convertMCPToolToClaudeTool converts a single MCP tool to Claude format
109+
func convertMCPToolToClaudeTool(tool *mcp.Tool) anthropic.ToolUnionParam {
110+
inputSchemaMap, ok := tool.InputSchema.(map[string]any)
111+
if !ok || inputSchemaMap == nil {
112+
inputSchemaMap = map[string]any{}
113+
}
60114

61-
toolParam.OfTool.Description = anthropic.String(tool.Description)
115+
properties := inputSchemaMap["properties"]
116+
required, _ := inputSchemaMap["required"].([]string)
62117

63-
claudeTools = append(claudeTools, toolParam)
64-
}
65-
anthropicClient := anthropic.NewClient(
66-
option.WithAPIKey(apiKey),
118+
toolParam := anthropic.ToolUnionParamOfTool(
119+
anthropic.ToolInputSchemaParam{
120+
Type: "object",
121+
Properties: properties,
122+
Required: required,
123+
},
124+
tool.Name,
67125
)
68-
messages := []anthropic.MessageParam{}
69-
if len(os.Args) <= 1 {
70-
log.Fatal("Usage: go run main.go <content>")
71-
return
126+
127+
toolParam.OfTool.Description = anthropic.String(tool.Description)
128+
return toolParam
129+
}
130+
131+
// runAgentLoop executes the main agent conversation loop
132+
func runAgentLoop(ctx context.Context, client anthropic.Client, session *mcp.ClientSession, tools []anthropic.ToolUnionParam, userQuery string) error {
133+
messages := []anthropic.MessageParam{
134+
anthropic.NewUserMessage(anthropic.NewTextBlock(userQuery)),
72135
}
73-
content := strings.Join(os.Args[1:], " ")
74136

75-
messages = append(messages, anthropic.NewUserMessage(anthropic.NewTextBlock(content)))
76-
message, err := anthropicClient.Messages.New(ctx, anthropic.MessageNewParams{
77-
Model: anthropic.ModelClaude_3_Haiku_20240307,
78-
MaxTokens: 256,
79-
Messages: messages,
80-
Tools: claudeTools,
81-
})
82-
if err != nil {
83-
log.Fatal(err)
137+
for _ = range maxAgentTurns {
138+
139+
message, err := client.Messages.New(ctx, anthropic.MessageNewParams{
140+
Model: claudeModel,
141+
MaxTokens: maxTokens,
142+
System: []anthropic.TextBlockParam{{Type: "text", Text: systemPrompt}},
143+
Messages: messages,
144+
Tools: tools,
145+
})
146+
if err != nil {
147+
return fmt.Errorf("claude API error: %w", err)
148+
}
149+
150+
assistantContent, toolResults, hasToolUse := processClaudeResponse(ctx, session, message)
151+
152+
if !hasToolUse {
153+
break
154+
}
155+
156+
messages = append(messages, anthropic.NewAssistantMessage(assistantContent...))
157+
messages = append(messages, anthropic.NewUserMessage(toolResults...))
84158
}
85159

86-
assistantMessageContent := []anthropic.ContentBlockParamUnion{}
87-
toolResults := []anthropic.ContentBlockParamUnion{}
160+
return nil
161+
}
162+
163+
// processClaudeResponse handles Claude's response and executes tool calls
164+
func processClaudeResponse(
165+
ctx context.Context,
166+
session *mcp.ClientSession,
167+
message *anthropic.Message,
168+
) (
169+
assistantContent []anthropic.ContentBlockParamUnion,
170+
toolResults []anthropic.ContentBlockParamUnion,
171+
hasToolUse bool,
172+
) {
173+
assistantContent = []anthropic.ContentBlockParamUnion{}
174+
toolResults = []anthropic.ContentBlockParamUnion{}
88175

89176
for _, block := range message.Content {
90177
switch block := block.AsAny().(type) {
91178
case anthropic.TextBlock:
92-
println(block.Text)
93-
assistantMessageContent = append(assistantMessageContent, anthropic.NewTextBlock(block.Text))
179+
fmt.Println(block.Text)
180+
assistantContent = append(assistantContent, anthropic.NewTextBlock(block.Text))
181+
94182
case anthropic.ToolUseBlock:
95-
assistantMessageContent = append(assistantMessageContent, anthropic.NewToolUseBlock(block.ID, block.Input, block.Name))
96-
97-
params := &mcp.CallToolParams{
98-
Name: block.Name,
99-
Arguments: block.Input,
100-
}
101-
res, err := session.CallTool(ctx, params)
102-
if err != nil {
103-
log.Fatalf("CallTool failed: %v", err)
104-
}
105-
if res.IsError {
106-
log.Fatal("tool failed")
107-
}
108-
109-
for _, c := range res.Content {
110-
log.Print(c.(*mcp.TextContent).Text)
111-
println()
112-
}
113-
114-
toolResults = append(toolResults, anthropic.NewToolResultBlock(
115-
block.ID,
116-
res.Content[0].(*mcp.TextContent).Text,
117-
false,
118-
))
183+
hasToolUse = true
184+
log.Printf("\n[Tool Use] %s", block.Name)
185+
186+
assistantContent = append(assistantContent,
187+
anthropic.NewToolUseBlock(block.ID, block.Input, block.Name))
188+
189+
toolResult := executeToolCall(ctx, session, block)
190+
toolResults = append(toolResults, toolResult)
119191
}
120192
}
121193

122-
if len(toolResults) > 0 {
194+
return assistantContent, toolResults, hasToolUse
195+
}
123196

124-
messages = append(messages, anthropic.NewAssistantMessage(assistantMessageContent...))
197+
// executeToolCall calls a tool via MCP and returns the result
198+
func executeToolCall(
199+
ctx context.Context,
200+
session *mcp.ClientSession,
201+
toolUse anthropic.ToolUseBlock,
202+
) anthropic.ContentBlockParamUnion {
125203

126-
messages = append(messages, anthropic.NewUserMessage(toolResults...))
204+
result, err := session.CallTool(ctx, &mcp.CallToolParams{
205+
Name: toolUse.Name,
206+
Arguments: toolUse.Input,
207+
})
127208

128-
message, err = anthropicClient.Messages.New(ctx, anthropic.MessageNewParams{
129-
Model: anthropic.ModelClaude_3_Haiku_20240307,
130-
MaxTokens: 256,
131-
Messages: messages,
132-
Tools: claudeTools,
133-
})
134-
if err != nil {
135-
log.Fatal(err)
136-
}
209+
if err != nil {
210+
log.Printf("[Error] CallTool failed: %v", err)
211+
return anthropic.NewToolResultBlock(
212+
toolUse.ID,
213+
fmt.Sprintf("Error: %v", err),
214+
true,
215+
)
216+
}
217+
218+
if result.IsError {
219+
log.Printf("[Error] Tool execution failed")
220+
return anthropic.NewToolResultBlock(
221+
toolUse.ID,
222+
"Tool execution failed",
223+
true,
224+
)
225+
}
226+
227+
resultText := extractTextContent(result.Content)
228+
log.Printf("================================================")
229+
log.Printf("[Tool Result]\n%s", resultText)
230+
log.Printf("================================================")
137231

138-
for _, block := range message.Content {
139-
if textBlock, ok := block.AsAny().(anthropic.TextBlock); ok {
140-
println("\nFinal response:")
141-
println(textBlock.Text)
142-
}
232+
return anthropic.NewToolResultBlock(toolUse.ID, resultText, false)
233+
}
234+
235+
func extractTextContent(content []mcp.Content) string {
236+
var resultText strings.Builder
237+
238+
for _, c := range content {
239+
if textContent, ok := c.(*mcp.TextContent); ok {
240+
resultText.WriteString(textContent.Text)
143241
}
144242
}
243+
244+
return resultText.String()
145245
}

0 commit comments

Comments
 (0)