knowledge-graph-agent/backend/internal/llm/agent.go

package llm

import (
	"context"
	"fmt"
	"log"

	"github.com/tmc/langchaingo/llms"
)

type Agent struct {
	client   *Client
	executor *ToolExecutor
}

func NewAgent(client *Client, executor *ToolExecutor) *Agent {
	return &Agent{
		client:   client,
		executor: executor,
	}
}

func (a *Agent) ChatWithTools(ctx context.Context, systemPrompt string, messages []llms.MessageContent) (string, error) {
	tools := GetKnowledgeGraphTools()

	builder := NewMessageBuilder().
		WithSystemPrompt(systemPrompt).
		WithMessages(messages)

	maxIterations := 10

	for i := 0; i < maxIterations; i++ {
		resp, err := a.client.Generate(ctx, builder.Build(), llms.WithTools(tools))
		if err != nil {
			log.Printf("[ERROR] Agent GenerateContent failed: %v", err)
			return "", fmt.Errorf("agent generate content error: %w", err)
		}

		if resp == nil || len(resp.Choices) == 0 {
			log.Printf("[WARN] Agent response is empty or has no choices")
			return "", fmt.Errorf("no content choices returned")
		}

		choice := resp.Choices[0]

		if len(choice.ToolCalls) == 0 {
			log.Printf("[INFO] Agent chat completed successfully, choice length: %d", len(choice.Content))
			return choice.Content, nil
		}

		builder.AppendAIMessage(choice.Content, choice.ToolCalls)

		for _, tc := range choice.ToolCalls {
			log.Printf("[INFO] Executing tool call: %s, Args: %s", tc.FunctionCall.Name, tc.FunctionCall.Arguments)

			toolResult, err := a.executor.Execute(tc)
			if err != nil {
				log.Printf("[WARN] Tool execution failed: %v", err)
				toolResult = fmt.Sprintf(`{"error": "%s"}`, err.Error())
			}

			log.Printf("[INFO] Tool call %s completed, result length: %d", tc.FunctionCall.Name, len(toolResult))

			builder.AppendToolResult(tc.ID, tc.FunctionCall.Name, toolResult)
		}
	}

	return "", fmt.Errorf("agent reached maximum tool call iterations (%d)", maxIterations)
}