langchaingo_custom_tool.go

// Custom Tool
---
// Prometheus is a tool that can query Prometheus for metrics.
type Prometheus struct {
	CallbacksHandler callbacks.Handler
}

var _ tools.Tool = Prometheus{}

// Description returns a string describing the Prometheus tool.
func (c Prometheus) Description() string {
	return `Useful for getting the result of a PromQL expression. 
	The input to this tool should be a valid PromQL expression that could be executed by Prometheus.
	The input should be a string and might contain chars like '(),=,=~' and many others.`
}

// Name returns the name of the tool.
func (c Prometheus) Name() string {
	return "prometheus"
}

// Call evaluates the input using the Prometheus API and returns the result as a
// string. If the evaluator errors, the error is given in the result to give the
// agent the ability to retry.
func (c Prometheus) Call(ctx context.Context, input string) (string, error) {
	if c.CallbacksHandler != nil {
		c.CallbacksHandler.HandleToolStart(ctx, input)
	}

	// Simulated result; in a real implementation, you would query an actual Prometheus server
	result := "Query result for: " + input

	if c.CallbacksHandler != nil {
		c.CallbacksHandler.HandleToolEnd(ctx, result)
	}

	return result, nil
}

--- 
// main.go
func main() {
llm, err = ollama.New(
		ollama.WithModel("llama3"),
	)

	if err != nil {
		return fmt.Errorf("failed to initialize LLM: %w", err)
	}

	// // Define the tools that the agent will use, including the new Celebrity tool.
	agentTools := []tools.Tool{
		itools.Prometheus{},
	}

	// Create a new one-shot agent with the LLM and the tools.
	agent := agents.NewOneShotAgent(
		llm,
		agentTools,
	)

	executor := agents.NewExecutor(agent, agents.WithMaxIterations(3))
	question := "What is value of 'count(up)'?"
	answer, err := chains.Run(context.Background(), executor, question)
	if err != nil {
		return fmt.Errorf("failed to run the executor: %w", err)
	}
	fmt.Println(answer)
}

main.go

package main

import (
	"context"
	"fmt"
	"github.com/tmc/langchaingo/agents"
	"github.com/tmc/langchaingo/chains"
	"github.com/tmc/langchaingo/llms/ollama"
	"github.com/tmc/langchaingo/tools"
	"log"
)

func main() {
	llm, err := ollama.New(
		ollama.WithModel("llama3"),
	)

	if err != nil {
		log.Fatal(err)
	}

	// // Define the tools that the agent will use, including the new Celebrity tool.
	agentTools := []tools.Tool{
		Prometheus{},
	}

	// Create a new one-shot agent with the LLM and the tools.
	agent := agents.NewOneShotAgent(
		llm,
		agentTools,
	)

	executor := agents.NewExecutor(agent, agents.WithMaxIterations(3))
	question := "What is value of 'count(up)'?"
	answer, err := chains.Run(context.Background(), executor, question)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println(answer)
}

tools.go

package main

import (
	"context"
	"github.com/tmc/langchaingo/callbacks"
	"github.com/tmc/langchaingo/tools"
)

// Prometheus is a tool that can query Prometheus for metrics.
type Prometheus struct {
	CallbacksHandler callbacks.Handler
}

var _ tools.Tool = Prometheus{}

// Description returns a string describing the Prometheus tool.
func (c Prometheus) Description() string {
	return `Useful for getting the result of a PromQL expression. 
	The input to this tool should be a valid PromQL expression that could be executed by Prometheus.
	The input should be a string and might contain chars like '(),=,=~' and many others.`
}

// Name returns the name of the tool.
func (c Prometheus) Name() string {
	return "prometheus"
}

// Call evaluates the input using the Prometheus API and returns the result as a
// string. If the evaluator errors, the error is given in the result to give the
// agent the ability to retry.
func (c Prometheus) Call(ctx context.Context, input string) (string, error) {
	if c.CallbacksHandler != nil {
		c.CallbacksHandler.HandleToolStart(ctx, input)
	}
	// Simulated result; in a real implementation, you would query an actual Prometheus server
	result := "Query result for: " + input

	if c.CallbacksHandler != nil {
		c.CallbacksHandler.HandleToolEnd(ctx, result)
	}

	return result, nil
}

I did some adjustments in main.go