package main

import (
 "context"
 "fmt"
 "log"
 "os"

 "github.com/tmc/langchaingo/chains"
 "github.com/tmc/langchaingo/documentloaders"
 "github.com/tmc/langchaingo/embeddings"
 "github.com/tmc/langchaingo/llms/ollama"
 "github.com/tmc/langchaingo/textsplitter"
 "github.com/tmc/langchaingo/vectorstores"
 "github.com/tmc/langchaingo/vectorstores/redisvector"
)

func main() {

 redisURL := "redis://localhost:10001" //redis-stack
 ollamaURL := "http://localhost:11434"
 index := "test_redis_vectorstore"

 model := "llama3.1:8b"
 embeddingModel := "nomic-embed-text:latest"

 opts_chat := []ollama.Option{
  ollama.WithModel(model),
  ollama.WithServerURL(ollamaURL),
 }

 llm_chat, err := ollama.New(opts_chat...)
 if err != nil {
  fmt.Println(err)
 }

 opts_emb := []ollama.Option{
  ollama.WithModel(embeddingModel),
  ollama.WithServerURL(ollamaURL),
 }

 llm_emb, err := ollama.New(opts_emb...)
 if err != nil {
  fmt.Println(err)
 }

 emb, err := embeddings.NewEmbedder(llm_emb)
 if err != nil {
  log.Fatal(err)
 }

 ctx := context.Background()
 store, err := redisvector.New(ctx,
  redisvector.WithConnectionURL(redisURL),
  redisvector.WithIndexName(index, true),
  redisvector.WithEmbedder(emb),
 )

 f, err := os.Open("./spec/doc.pdf")
 if err != nil {
  fmt.Printf("Failed to open file: %v\n", err)
  return
 }
 defer f.Close()
 finfo, err := f.Stat()
 p := documentloaders.NewPDF(f, finfo.Size())
 split := textsplitter.NewRecursiveCharacter()
 split.ChunkSize = 1000
 split.ChunkOverlap = 30

 //docs, err := p.Load(context.Background())
 data, err := p.LoadAndSplit(context.Background(), split)
 if err != nil {
  fmt.Printf("Failed to load file: %v\n", err)
  return
 }
 _, err = store.AddDocuments(ctx, data)
 if err != nil {
  fmt.Println(err)
  return
 }

 // docs, err := store.SimilaritySearch(ctx, "промпт", 2,
 //  vectorstores.WithScoreThreshold(0.5),
 // )
 // fmt.Println(docs)

 result, err := chains.Run(
  ctx,
  chains.NewRetrievalQAFromLLM(
   llm_chat,
   vectorstores.ToRetriever(store, 5, vectorstores.WithScoreThreshold(0.8)),
  ),
  "промпт",
 )
 fmt.Println(result)

}