Building a CLI RAG Tool in Golang Using Qdrant and LangChain
This is a simple and fast implementation of a CLI-based RAG (Retrieval-Augmented Generation) tool in Golang.
This tool leverages:
- LangChain for embeddings and text processing
- Qdrant as the vector database for storing and retrieving documents
- OpenAI for generating responses
- go-fitz for extracting text from PDFs which is a wrapper for a MuPDF library
- urfave/cli for creating a CLI interface
Before starting this tool, we need to install the Qdrant vector database. The easiest way to do this is by using the Docker image.
docker pull qdrant/qdrant
docker run -p 6333:6333 -p 6334:6334 \
-v "$(pwd)/qdrant_storage:/qdrant/storage:z" \
qdrant/qdrant
Development Process
1. Setting Up Dependencies
The necessary packages are imported at the beginning, including OpenAI for embeddings, Qdrant for vector storage, and go-fitz for handling PDFs.
import (
"context"
"fmt"
"log"
"net/http"
"net/url"
"os"
"strings"
"github.com/gen2brain/go-fitz"
"github.com/google/uuid"
"github.com/tmc/langchaingo/embeddings"
"github.com/tmc/langchaingo/llms"
"github.com/tmc/langchaingo/llms/openai"
"github.com/tmc/langchaingo/schema"
"github.com/tmc/langchaingo/textsplitter"
"github.com/tmc/langchaingo/vectorstores"
"github.com/tmc/langchaingo/vectorstores/qdrant"
"github.com/urfave/cli/v3"
)
2. Initializing OpenAI and Qdrant
The OpenAI model is initialized with specific settings, including an embedding model. The Qdrant API URL is also set up.
opts := []openai.Option{
openai.WithToken("YOUR-OPENAI-TOKEN"),
openai.WithModel("gpt-4o-mini"),
openai.WithEmbeddingModel("text-embedding-3-small"),
}
llm, err := openai.New(opts...)
if err != nil {
log.Fatal(err)
}
// Initialize embedding model
e, err := embeddings.NewEmbedder(llm)
if err != nil {
log.Fatal(err)
}
// Define Qdrant API endpoint
urlAPI, err := url.Parse("http://localhost:6333")
if err != nil {
log.Fatal(err)
}
3. CLI Command Structure
Using urfave/cli, the application supports two commands: index for document storage and query for information retrieval.
cmd := &cli.Command{
Name: "Simple CLI RAG Tool",
Usage: "A simple CLI RAG Tool with index and query commands",
4. Indexing a PDF Document
The index command extracts text from a PDF, splits it into chunks, and stores it in a Qdrant collection.
a. Extracting Text from the PDF
doc, err := fitz.New(fileName)
if err != nil {
panic(err)
}
defer doc.Close()
b. Splitting the Text into Chunks
Text is chunked into smaller parts to optimize retrieval performance.
reqCharacterSplitter := textsplitter.NewRecursiveCharacter() reqCharacterSplitter.ChunkSize = 1000 reqCharacterSplitter.ChunkOverlap = 200
c. Storing in Qdrant
A new Qdrant collection is created and documents are added.
collectionName := uuid.NewString()
urlCollection := urlAPI.JoinPath("collections", collectionName)
_, _, err = qdrant.DoRequest(ctx, *urlCollection, "", http.MethodPut, collectionConfig)
5. Querying the Indexed Documents
The query command searches for relevant document chunks and generates responses using OpenAI.
a. Retrieving Similar Documents
docs, err := store.SimilaritySearch(ctx,
question, 2,
vectorstores.WithScoreThreshold(0))
b. Generating a Response
A response is generated based on the retrieved context.
output, err := llm.GenerateContent(ctx, content,
llms.WithMaxTokens(1024),
llms.WithTemperature(0),
)
fmt.Println(output.Choices[0].Content)
Conclusion
This CLI tool demonstrates how to use Golang to implement a RAG-based retrieval system. By combining LangChain, Qdrant, and OpenAI, it enables efficient document indexing and querying. Future improvements could include enhanced error handling and a web interface for broader usability.
Complete source code:
main.go
package main
import (
"context"
"fmt"
"log"
"net/http"
"net/url"
"os"
"strings"
"github.com/gen2brain/go-fitz"
"github.com/google/uuid"
"github.com/tmc/langchaingo/embeddings"
"github.com/tmc/langchaingo/llms"
"github.com/tmc/langchaingo/llms/openai"
"github.com/tmc/langchaingo/schema"
"github.com/tmc/langchaingo/textsplitter"
"github.com/tmc/langchaingo/vectorstores"
"github.com/tmc/langchaingo/vectorstores/qdrant"
"github.com/urfave/cli/v3"
)
func main() {
opts := []openai.Option{
openai.WithToken("YOUR-OPENAI-TOKEN"),
openai.WithModel("gpt-4o-mini"),
openai.WithEmbeddingModel("text-embedding-3-small"),
}
llm, err := openai.New(opts...)
if err != nil {
log.Fatal(err)
}
e, err := embeddings.NewEmbedder(llm)
if err != nil {
log.Fatal(err)
}
urlAPI, err := url.Parse("http://localhost:6333")
if err != nil {
log.Fatal(err)
}
cmd := &cli.Command{
Name: "Simple CLI RAG Tool",
Usage: "A simple CLI RAG Tool with index and query commands",
Commands: []*cli.Command{
{
Name: "index",
Usage: "Prints the file name",
Flags: []cli.Flag{
&cli.StringFlag{
Name: "file",
Aliases: []string{"f"},
Usage: "File to be indexed",
Required: true,
},
},
Action: func(ctx context.Context, cmd *cli.Command) error {
fileName := os.Args[3]
fmt.Println("Indexing file:", fileName)
doc, err := fitz.New(fileName)
if err != nil {
panic(err)
}
defer doc.Close()
reqCharacterSplitter := textsplitter.NewRecursiveCharacter()
reqCharacterSplitter.ChunkSize = 1000
reqCharacterSplitter.ChunkOverlap = 200
reqCharacterSplitter.LenFunc = func(s string) int { return len(s) }
pagesList := make([]schema.Document, 0)
for idx := range doc.NumPage() {
text, _ := doc.Text(idx)
text = strings.ReplaceAll(text, "\n", " ")
text = strings.ToLower(text)
newDoc := schema.Document{PageContent: text}
pagesList = append(pagesList, newDoc)
}
chunksDocList, _ := textsplitter.SplitDocuments(reqCharacterSplitter, pagesList)
collectionName := uuid.NewString()
fmt.Println("Collection name:", collectionName)
collectionConfig := map[string]interface{}{
"vectors": map[string]interface{}{
"size": 1536,
"distance": "Cosine",
},
}
urlCollection := urlAPI.JoinPath("collections", collectionName)
_, _, err = qdrant.DoRequest(ctx, *urlCollection, "", http.MethodPut, collectionConfig)
if err != nil {
log.Fatal(err)
}
store, err := qdrant.New(
qdrant.WithURL(*urlAPI),
qdrant.WithCollectionName(collectionName),
qdrant.WithEmbedder(e),
)
if err != nil {
log.Fatal(err)
}
_, err = store.AddDocuments(ctx, chunksDocList)
if err != nil {
log.Fatal(err)
}
return nil
},
},
{
Name: "query",
Usage: "Prints the query string",
Flags: []cli.Flag{
&cli.StringFlag{
Name: "collection",
Aliases: []string{"c"},
Usage: "Collection name",
Required: true,
},
&cli.StringFlag{
Name: "string",
Aliases: []string{"s"},
Usage: "Query string",
Required: true,
},
},
Action: func(ctx context.Context, cmd *cli.Command) error {
collectionName := os.Args[3]
queryString := os.Args[5]
fmt.Println("Querying:", queryString)
question := strings.ReplaceAll(queryString, "\"", "")
question = strings.ToLower(question)
collectionConfig := map[string]interface{}{
"vectors": map[string]interface{}{
"size": 1536,
"distance": "Cosine",
},
}
urlCollection := urlAPI.JoinPath("collections", collectionName)
_, _, err = qdrant.DoRequest(ctx, *urlCollection, "", http.MethodPut, collectionConfig)
if err != nil {
log.Fatal(err)
}
store, err := qdrant.New(
qdrant.WithURL(*urlAPI),
qdrant.WithCollectionName(collectionName),
qdrant.WithEmbedder(e),
)
if err != nil {
log.Fatal(err)
}
docs, err := store.SimilaritySearch(ctx,
question, 2,
vectorstores.WithScoreThreshold(0))
if err != nil {
log.Fatal(err)
}
stringContext := ""
for i := range len(docs) {
stringContext += docs[i].PageContent
}
content := []llms.MessageContent{
llms.TextParts(llms.ChatMessageTypeSystem, "You are a helpful assistant."),
llms.TextParts(llms.ChatMessageTypeHuman, `Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
`+stringContext+`
Question:`+question+`
Helpful Answer:",
"context","question")))
`),
}
output, err := llm.GenerateContent(ctx, content,
llms.WithMaxTokens(1024),
llms.WithTemperature(0),
)
if err != nil {
log.Fatal(err)
}
fmt.Println(output.Choices[0].Content)
return nil
},
},
},
}
if err := cmd.Run(context.Background(), os.Args); err != nil {
log.Fatal(err)
}
}