From 66a58f5481ab4d9e0e3a114c22cfb4116ec4fcba Mon Sep 17 00:00:00 2001
From: Simone Vellei <henomis@gmail.com>
Date: Wed, 22 May 2024 15:38:54 +0200
Subject: [PATCH] chore: add tools

---
 tools/duckduckgo/duckduckgo.go   |  4 +-
 tools/llm/llm.go                 | 60 +++++++++++++++++++++++++
 tools/rag/rag.go                 | 54 +++++++++++++++++++++++
 tools/tool_router/tool_router.go | 75 ++++++++++++++++++++++++++++++++
 4 files changed, 191 insertions(+), 2 deletions(-)
 create mode 100644 tools/llm/llm.go
 create mode 100644 tools/rag/rag.go
 create mode 100644 tools/tool_router/tool_router.go

diff --git a/tools/duckduckgo/duckduckgo.go b/tools/duckduckgo/duckduckgo.go
index 3c5f2b7d..80b2da55 100644
--- a/tools/duckduckgo/duckduckgo.go
+++ b/tools/duckduckgo/duckduckgo.go
@@ -15,7 +15,7 @@ type Tool struct {
 }
 
 type Input struct {
-	Query string `json:"python_code" jsonschema:"description=the query to search for"`
+	Query string `json:"query" jsonschema:"description=the query to search for"`
 }
 
 type Output struct {
@@ -57,7 +57,7 @@ func (t *Tool) Name() string {
 }
 
 func (t *Tool) Description() string {
-	return "A tool that searches DuckDuckGo for a query."
+	return "A tool that searches on duckduckgo internet search engine for a query."
 }
 
 func (t *Tool) Fn() any {
diff --git a/tools/llm/llm.go b/tools/llm/llm.go
new file mode 100644
index 00000000..8b05bb55
--- /dev/null
+++ b/tools/llm/llm.go
@@ -0,0 +1,60 @@
+package llm
+
+import (
+	"context"
+
+	"github.com/henomis/lingoose/thread"
+)
+
+type LLM interface {
+	Generate(context.Context, *thread.Thread) error
+}
+
+type Tool struct {
+	llm LLM
+}
+
+func New(llm LLM) *Tool {
+	return &Tool{
+		llm: llm,
+	}
+}
+
+type Input struct {
+	Query string `json:"query" jsonschema:"description=user query"`
+}
+
+type Output struct {
+	Error  string `json:"error,omitempty"`
+	Result string `json:"result,omitempty"`
+}
+
+type FnPrototype func(Input) Output
+
+func (t *Tool) Name() string {
+	return "llm"
+}
+
+func (t *Tool) Description() string {
+	return "A tool that uses a language model to generate a response to a user query."
+}
+
+func (t *Tool) Fn() any {
+	return t.fn
+}
+
+//nolint:gosec
+func (t *Tool) fn(i Input) Output {
+	th := thread.New().AddMessage(
+		thread.NewUserMessage().AddContent(
+			thread.NewTextContent(i.Query),
+		),
+	)
+
+	err := t.llm.Generate(context.Background(), th)
+	if err != nil {
+		return Output{Error: err.Error()}
+	}
+
+	return Output{Result: th.LastMessage().Contents[0].AsString()}
+}
diff --git a/tools/rag/rag.go b/tools/rag/rag.go
new file mode 100644
index 00000000..13f05682
--- /dev/null
+++ b/tools/rag/rag.go
@@ -0,0 +1,54 @@
+package rag
+
+import (
+	"context"
+	"strings"
+
+	"github.com/henomis/lingoose/rag"
+)
+
+type Tool struct {
+	rag   *rag.RAG
+	topic string
+}
+
+func New(rag *rag.RAG, topic string) *Tool {
+	return &Tool{
+		rag:   rag,
+		topic: topic,
+	}
+}
+
+type Input struct {
+	Query string `json:"rag_query" jsonschema:"description=search query"`
+}
+
+type Output struct {
+	Error  string `json:"error,omitempty"`
+	Result string `json:"result,omitempty"`
+}
+
+type FnPrototype func(Input) Output
+
+func (t *Tool) Name() string {
+	return "rag"
+}
+
+func (t *Tool) Description() string {
+	return "A tool that searches information ONLY for this topic: " + t.topic + ". DO NOT use this tool for other topics."
+}
+
+func (t *Tool) Fn() any {
+	return t.fn
+}
+
+//nolint:gosec
+func (t *Tool) fn(i Input) Output {
+	results, err := t.rag.Retrieve(context.Background(), i.Query)
+	if err != nil {
+		return Output{Error: err.Error()}
+	}
+
+	// Return the output as a string.
+	return Output{Result: strings.Join(results, "\n")}
+}
diff --git a/tools/tool_router/tool_router.go b/tools/tool_router/tool_router.go
new file mode 100644
index 00000000..529dfab2
--- /dev/null
+++ b/tools/tool_router/tool_router.go
@@ -0,0 +1,75 @@
+package toolrouter
+
+import (
+	"context"
+
+	"github.com/henomis/lingoose/thread"
+)
+
+type TTool interface {
+	Description() string
+	Name() string
+	Fn() any
+}
+
+type Tool struct {
+	llm   LLM
+	tools []TTool
+}
+
+type LLM interface {
+	Generate(context.Context, *thread.Thread) error
+}
+
+func New(llm LLM, tools ...TTool) *Tool {
+	return &Tool{
+		tools: tools,
+		llm:   llm,
+	}
+}
+
+type Input struct {
+	Query string `json:"query" jsonschema:"description=user query"`
+}
+
+type Output struct {
+	Error  string `json:"error,omitempty"`
+	Result any    `json:"result,omitempty"`
+}
+
+type FnPrototype func(Input) Output
+
+func (t *Tool) Name() string {
+	return "query_router"
+}
+
+func (t *Tool) Description() string {
+	return "A tool that select the right tool to answer to an user query."
+}
+
+func (t *Tool) Fn() any {
+	return t.fn
+}
+
+//nolint:gosec
+func (t *Tool) fn(i Input) Output {
+	query := "Here's a list of available tools:\n\n"
+	for _, tool := range t.tools {
+		query += "Name: " + tool.Name() + "\nDescription: " + tool.Description() + "\n\n"
+	}
+
+	query += "\nPlease select the right tool that can better answer the query '" + i.Query + "'. Give me only the name of the tool, nothing else."
+
+	th := thread.New().AddMessage(
+		thread.NewUserMessage().AddContent(
+			thread.NewTextContent(query),
+		),
+	)
+
+	err := t.llm.Generate(context.Background(), th)
+	if err != nil {
+		return Output{Error: err.Error()}
+	}
+
+	return Output{Result: th.LastMessage().Contents[0].AsString()}
+}