tools: add perplexity AI

tools/perplexity/doc.go

@@ -0,0 +1,40 @@
+// Package perplexity provides integration with Perplexity AI's API for AI agents.
+//
+// Perplexity AI functions as an AI-powered search engine that indexes, analyzes,
+// and summarizes content from across the internet. This package allows you to
+// integrate Perplexity's capabilities into your AI agents to enrich them with
+// up-to-date web data.
+//
+// Example usage:
+//
+//	llm, err := openai.New(
+//		openai.WithModel("gpt-4-mini"),
+//		openai.WithCallback(callbacks.LogHandler{}),
+//	)
+//	if err != nil {
+//		return err
+//	}
+//
+//	// Create a new Perplexity instance
+//	perpl, err := perplexity.New(
+//		perplexity.WithModel(perplexity.ModelLlamaSonarSmall),
+//		perplexity.WithAPIKey("your-api-key"), // Optional: defaults to PERPLEXITY_API_KEY env var
+//	)
+//	if err != nil {
+//		return err
+//	}
+//
+//	// Add Perplexity as a tool for your agent
+//	agentTools := []tools.Tool{
+//		perpl,
+//	}
+//
+//	// Create and use the agent
+//	toolAgent := agents.NewOneShotAgent(llm,
+//		agentTools,
+//		agents.WithMaxIterations(2),
+//	)
+//	executor := agents.NewExecutor(toolAgent)
+//
+//	answer, err := chains.Run(context.Background(), executor, "your question here")
+package perplexity

tools/perplexity/perplexity.go

@@ -0,0 +1,124 @@
+package perplexity
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"github.com/tmc/langchaingo/callbacks"
+	"github.com/tmc/langchaingo/llms"
+	"github.com/tmc/langchaingo/llms/openai"
+	"github.com/tmc/langchaingo/tools"
+)
+
+// Model represents a Perplexity AI model type.
+type Model string
+
+// Model pricing overview: https://docs.perplexity.ai/guides/pricing
+const (
+	// ModelLlamaSonarSmall is the small version of the Llama Sonar model.
+	ModelLlamaSonarSmall Model = "llama-3.1-sonar-small-128k-online"
+	// ModelLlamaSonarLarge is the large version of the Llama Sonar model.
+	ModelLlamaSonarLarge Model = "llama-3.1-sonar-large-128k-online"
+	// ModelLlamaSonarHuge is the huge version of the Llama Sonar model.
+	ModelLlamaSonarHuge Model = "llama-3.1-sonar-huge-128k-online"
+)
+
+// Option is a function that modifies the options for the Perplexity AI tool.
+type Option func(*options)
+
+type options struct {
+	apiKey string
+	model  Model
+}
+
+// WithAPIKey sets the API key for Perplexity AI.
+func WithAPIKey(apiKey string) Option {
+	return func(o *options) {
+		o.apiKey = apiKey
+	}
+}
+
+// WithModel sets the model to be used by Perplexity AI.
+func WithModel(model Model) Option {
+	return func(o *options) {
+		o.model = model
+	}
+}
+
+// Tool implements the Perplexity AI integration.
+type Tool struct {
+	llm              *openai.LLM
+	CallbacksHandler callbacks.Handler
+}
+
+var _ tools.Tool = (*Tool)(nil)
+
+// New creates a new instance of the Perplexity AI tool with the given options.
+func New(opts ...Option) (*Tool, error) {
+	options := &options{
+		apiKey: os.Getenv("PERPLEXITY_API_KEY"),
+		model:  ModelLlamaSonarSmall, // Default model
+	}
+
+	for _, opt := range opts {
+		opt(options)
+	}
+
+	if options.apiKey == "" {
+		return nil, fmt.Errorf("PERPLEXITY_API_KEY key not set")
+	}
+
+	llm, err := openai.New(
+		openai.WithModel(string(options.model)),
+		openai.WithBaseURL("https://api.perplexity.ai"),
+		openai.WithToken(options.apiKey),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Tool{
+		llm: llm,
+	}, nil
+}
+
+// Name returns the name of the tool.
+func (t *Tool) Name() string {
+	return "PerplexityAI"
+}
+
+// Description returns a description of the Perplexity AI tool's capabilities.
+func (t *Tool) Description() string {
+	return "Perplexity AI has access to a wide range of information, as it functions as an AI-powered search engine that indexes, analyzes, and summarizes content from across the internet."
+}
+
+// Call executes a query against the Perplexity AI model and returns the response.
+func (t *Tool) Call(ctx context.Context, input string) (string, error) {
+	if t.CallbacksHandler != nil {
+		t.CallbacksHandler.HandleToolStart(ctx, input)
+	}
+
+	content := []llms.MessageContent{
+		llms.TextParts(llms.ChatMessageTypeHuman, input),
+	}
+
+	var generatedText string
+	_, err := t.llm.GenerateContent(ctx, content,
+		llms.WithStreamingFunc(func(_ context.Context, chunk []byte) error {
+			generatedText += string(chunk)
+			return nil
+		}))
+	if err != nil {
+		if t.CallbacksHandler != nil {
+			t.CallbacksHandler.HandleToolError(ctx, err)
+		}
+		return "", err
+	}
+
+	if t.CallbacksHandler != nil {
+		t.CallbacksHandler.HandleToolEnd(ctx, generatedText)
+	}
+
+	return generatedText, nil
+}

tools/perplexity/perplexity_test.go

@@ -0,0 +1,32 @@
+package perplexity
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTool_Integration(t *testing.T) {
+	t.Parallel()
+
+	apiKey := os.Getenv("PERPLEXITY_API_KEY")
+	if apiKey == "" {
+		t.Skip("PERPLEXITY_API_KEY not set")
+	}
+
+	tool, err := New()
+	require.NoError(t, err)
+	require.NotNil(t, tool)
+
+	assert.Equal(t, "PerplexityAI", tool.Name())
+	assert.NotEmpty(t, tool.Description())
+
+	// Test Call functionality
+	ctx := context.Background()
+	response, err := tool.Call(ctx, "what is the largest country in the world by total area?")
+	require.NoError(t, err)
+	assert.Contains(t, response, "Russia")
+}