Integrate feedback; change structs and general workflow

Author: iuliadmtru

Project.toml

@@ -7,7 +7,6 @@ version = "0.45.0"
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
-DotEnv = "4dc1fcf4-5e3b-5448-94ab-0c38ec0385c1"
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
@@ -18,7 +17,6 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [weakdeps]
@@ -34,7 +32,7 @@ Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 FlashRankPromptingToolsExt = ["FlashRank"]
 GoogleGenAIPromptingToolsExt = ["GoogleGenAI"]
 MarkdownPromptingToolsExt = ["Markdown"]
-RAGToolsExperimentalExt = ["SparseArrays", "LinearAlgebra", "Unicode"]
+RAGToolsExperimentalExt = ["SparseArrays", "LinearAlgebra", "Unicode", "Pinecone"]
 SnowballPromptingToolsExt = ["Snowball"]
 
 [compat]

src/Experimental/RAGTools/RAGTools.jl

@@ -32,12 +32,12 @@ include("api_services.jl")
 
 include("rag_interface.jl")
 
-export ChunkIndex, ChunkKeywordsIndex, ChunkEmbeddingsIndex, PTPineconeIndex, CandidateChunks, RAGResult
+export ChunkIndex, ChunkKeywordsIndex, ChunkEmbeddingsIndex, PineconeIndex, CandidateChunks, CandidateWithChunks, RAGResult
 export MultiIndex, SubChunkIndex, MultiCandidateChunks
 include("types.jl")
 
 export build_index, get_chunks, get_embeddings, get_keywords, get_tags, SimpleIndexer,
-       KeywordsIndexer, PTPineconeIndexer
+       KeywordsIndexer, PineconeIndexer
 include("preparation.jl")
 
 include("rank_gpt.jl")

src/Experimental/RAGTools/generation.jl

@@ -37,7 +37,8 @@ context = build_context(ContextEnumerator(), index, candidates; chunks_window_ma
 ```
 """
 function build_context(contexter::ContextEnumerator,
-        index::AbstractDocumentIndex, candidates::AbstractCandidateChunks;
+        index::AbstractManagedIndex,
+        candidates::AbstractCandidateWithChunks;
         verbose::Bool = true,
         chunks_window_margin::Tuple{Int, Int} = (1, 1), kwargs...)
     ## Checks
@@ -63,27 +64,31 @@ function build_context(contexter::ContextEnumerator,
     return context
 end
 
-using Pinecone: Pinecone, query
-using JSON3: JSON3, read
-"""
-    build_context(contexter::ContextEnumerator,
-        index::AbstractPTPineconeIndex;
-        verbose::Bool = true,
-        top_k::Int = 10,
-        kwargs...)
-
-Build context strings by querying Pinecone.
-```
-"""
 function build_context(contexter::ContextEnumerator,
-        index::AbstractPTPineconeIndex;
+        index::AbstractManagedIndex,
+        candidates::AbstractCandidateWithChunks;
         verbose::Bool = true,
-        top_k::Int = 10,
-        kwargs...)
-    pinecone_results = Pinecone.query(index.pinecone_context, index.pinecone_index, index.embedding, top_k, index.namespace, false, true)
-    results_json = JSON3.read(pinecone_results)
-    context = results_json.matches[1].metadata.content
+        chunks_window_margin::Tuple{Int, Int} = (1, 1), kwargs...)
+    ## Checks
+    @assert chunks_window_margin[1] >= 0&&chunks_window_margin[2] >= 0 "Both `chunks_window_margin` values must be non-negative"
+
+    context = String[]
+    for (i, position) in enumerate(positions(candidates))
+        ## select the right index
+        id = candidates isa MultiCandidateChunks ? candidates.index_ids[i] :
+             candidates.index_id
+        index_ = index isa AbstractChunkIndex ? index : index[id]
+        isnothing(index_) && continue
 
+        chunks_ = chunks(candidates)[
+            max(1, position - chunks_window_margin[1]):min(end,
+            position + chunks_window_margin[2])]
+        ## Check if surrounding chunks are from the same source
+        is_same_source = sources(candidates)[
+            max(1, position - chunks_window_margin[1]):min(end,
+            position + chunks_window_margin[2])] .== sources(candidates)[position]
+        push!(context, "$(i). $(join(chunks_[is_same_source], "\n"))")
+    end
     return context
 end
 
@@ -94,7 +99,7 @@ end
 
 # Mutating version that dispatches on the result to the underlying implementation
 function build_context!(contexter::ContextEnumerator,
-        index::AbstractDocumentIndex, result::AbstractRAGResult; kwargs...)
+        index::Union{AbstractDocumentIndex, AbstractManagedIndex}, result::AbstractRAGResult; kwargs...)
     result.context = build_context(contexter, index, result.reranked_candidates; kwargs...)
     return result
 end
@@ -114,6 +119,7 @@ function answer!(
     throw(ArgumentError("Answerer $(typeof(answerer)) not implemented"))
 end
 
+# TODO: update docs signature
 """
     answer!(
         answerer::SimpleAnswerer, index::AbstractDocumentIndex, result::AbstractRAGResult;
@@ -138,7 +144,7 @@ Generates an answer using the `aigenerate` function with the provided `result.co
 
 """
 function answer!(
-        answerer::SimpleAnswerer, index::AbstractDocumentIndex, result::AbstractRAGResult;
+        answerer::SimpleAnswerer, index::Union{AbstractDocumentIndex, AbstractManagedIndex}, result::AbstractRAGResult;
         model::AbstractString = PT.MODEL_CHAT, verbose::Bool = true,
         template::Symbol = :RAGAnswerFromContext,
         cost_tracker = Threads.Atomic{Float64}(0.0),
@@ -186,11 +192,13 @@ Refines the answer by executing a web search using the Tavily API. This method a
 struct TavilySearchRefiner <: AbstractRefiner end
 
 function refine!(
-        refiner::AbstractRefiner, index::AbstractDocumentIndex, result::AbstractRAGResult;
+        refiner::AbstractRefiner, index::Union{AbstractDocumentIndex, AbstractManagedIndex}, result::AbstractRAGResult;
         kwargs...)
     throw(ArgumentError("Refiner $(typeof(refiner)) not implemented"))
 end
 
+
+# TODO: update docs signature
 """
     refine!(
         refiner::NoRefiner, index::AbstractChunkIndex, result::AbstractRAGResult;
@@ -199,7 +207,7 @@ end
 Simple no-op function for `refine!`. It simply copies the `result.answer` and `result.conversations[:answer]` without any changes.
 """
 function refine!(
-        refiner::NoRefiner, index::AbstractDocumentIndex, result::AbstractRAGResult;
+        refiner::NoRefiner, index::Union{AbstractDocumentIndex, AbstractManagedIndex}, result::AbstractRAGResult;
         kwargs...)
     result.final_answer = result.answer
     if haskey(result.conversations, :answer)
@@ -208,6 +216,8 @@ function refine!(
     return result
 end
 
+
+# TODO: update docs signature
 """
     refine!(
         refiner::SimpleRefiner, index::AbstractDocumentIndex, result::AbstractRAGResult;
@@ -234,7 +244,7 @@ This method uses the same context as the original answer, however, it can be mod
 - `cost_tracker`: An atomic counter to track the cost of the operation.
 """
 function refine!(
-        refiner::SimpleRefiner, index::AbstractDocumentIndex, result::AbstractRAGResult;
+        refiner::SimpleRefiner, index::Union{AbstractDocumentIndex, AbstractManagedIndex}, result::AbstractRAGResult;
         verbose::Bool = true,
         model::AbstractString = PT.MODEL_CHAT,
         template::Symbol = :RAGAnswerRefiner,
@@ -262,6 +272,8 @@ function refine!(
     return result
 end
 
+
+# TODO: update docs signature
 """
     refine!(
         refiner::TavilySearchRefiner, index::AbstractDocumentIndex, result::AbstractRAGResult;
@@ -312,7 +324,7 @@ pprint(result)
 ```
 """
 function refine!(
-        refiner::TavilySearchRefiner, index::AbstractDocumentIndex, result::AbstractRAGResult;
+        refiner::TavilySearchRefiner, index::Union{AbstractDocumentIndex, AbstractManagedIndex}, result::AbstractRAGResult;
         verbose::Bool = true,
         model::AbstractString = PT.MODEL_CHAT,
         include_answer::Bool = true,
@@ -377,13 +389,13 @@ Overload this method to add custom postprocessing steps, eg, logging, saving con
 """
 struct NoPostprocessor <: AbstractPostprocessor end
 
-function postprocess!(postprocessor::AbstractPostprocessor, index::AbstractDocumentIndex,
+function postprocess!(postprocessor::AbstractPostprocessor, index::Union{AbstractDocumentIndex, AbstractManagedIndex},
         result::AbstractRAGResult; kwargs...)
     throw(ArgumentError("Postprocessor $(typeof(postprocessor)) not implemented"))
 end
 
 function postprocess!(
-        ::NoPostprocessor, index::AbstractDocumentIndex, result::AbstractRAGResult; kwargs...)
+        ::NoPostprocessor, index::Union{AbstractDocumentIndex, AbstractManagedIndex}, result::AbstractRAGResult; kwargs...)
     return result
 end
 
@@ -416,6 +428,7 @@ It uses `ContextEnumerator`, `SimpleAnswerer`, `SimpleRefiner`, and `NoPostproce
     postprocessor::AbstractPostprocessor = NoPostprocessor()
 end
 
+# TODO: update docs signature
 """
     generate!(
         generator::AbstractGenerator, index::AbstractDocumentIndex, result::AbstractRAGResult;
@@ -483,7 +496,7 @@ result = generate!(index, result)
 ```
 """
 function generate!(
-        generator::AbstractGenerator, index::AbstractDocumentIndex, result::AbstractRAGResult;
+        generator::AbstractGenerator, index::Union{AbstractDocumentIndex, AbstractManagedIndex}, result::AbstractRAGResult;
         verbose::Integer = 1,
         api_kwargs::NamedTuple = NamedTuple(),
         contexter::AbstractContextBuilder = generator.contexter,
@@ -672,7 +685,7 @@ PT.pprint(result)
 
 For easier manipulation of nested kwargs, see utilities `getpropertynested`, `setpropertynested`, `merge_kwargs_nested`.
 """
-function airag(cfg::AbstractRAGConfig, index::AbstractDocumentIndex;
+function airag(cfg::AbstractRAGConfig, index::Union{AbstractDocumentIndex, AbstractManagedIndex};
         question::AbstractString,
         verbose::Integer = 1, return_all::Bool = false,
         api_kwargs::NamedTuple = NamedTuple(),

src/Experimental/RAGTools/preparation.jl

@@ -20,13 +20,6 @@ Chunker when you provide text to `get_chunks` functions. Inputs are directly chu
 """
 struct TextChunker <: AbstractChunker end
 
-"""
-    NoChunker <: AbstractChunker
-
-
-"""
-struct NoChunker <: AbstractChunker end
-
 ### Embedding Types
 """
     NoEmbedder <: AbstractEmbedder
@@ -35,6 +28,13 @@ No-op embedder for `get_embeddings` functions. It returns `nothing`.
 """
 struct NoEmbedder <: AbstractEmbedder end
 
+"""
+    SimpleEmbedder <: AbstractEmbedder
+
+Simply passes the input to `aiembed`.
+"""
+struct SimpleEmbedder <: AbstractEmbedder end
+
 """
     BatchEmbedder <: AbstractEmbedder
 
@@ -142,15 +142,13 @@ It uses `TextChunker`, `KeywordsProcessor`, and `NoTagger` as default chunker, p
 end
 
 """
-    PTPineconeIndexer <: AbstractIndexBuilder
+    PineconeIndexer <: AbstractIndexBuilder
 
 Pinecone index to be returned by `build_index`.
-
-It uses `NoChunker`, `NoEmbedder`, and `NoTagger` as default chunker, embedder, and tagger.
 """
-@kwdef mutable struct PTPineconeIndexer <: AbstractIndexBuilder
-    chunker::AbstractChunker = NoChunker()
-    embedder::AbstractEmbedder = NoEmbedder()
+@kwdef mutable struct PineconeIndexer <: AbstractIndexBuilder
+    chunker::AbstractChunker = TextChunker()
+    embedder::AbstractEmbedder = SimpleEmbedder()
     tagger::AbstractTagger = NoTagger()
 end
 
@@ -186,10 +184,6 @@ function load_text(chunker::TextChunker, input::AbstractString;
     @assert length(source)<=512 "Each `source` should be less than 512 characters long. Detected: $(length(source)) characters. You must provide sources for each text when using `TextChunker`"
     return input, source
 end
-function load_text(chunker::NoChunker, input::AbstractString = "";
-    source::AbstractString = input, kwargs...)
-return input, source
-end
 
 """
     get_chunks(chunker::AbstractChunker,
@@ -251,6 +245,13 @@ function get_embeddings(
     return nothing
 end
 
+function get_embeddings(
+        embedder::SimpleEmbedder, docs::AbstractVector{<:AbstractString};
+        model::AbstractString = PT.MODEL_EMBEDDING,
+        kwargs...)
+    return hcat([Vector{Float32}(aiembed(doc; model).content) for doc in docs]...)
+end
+
 """
     get_embeddings(embedder::BatchEmbedder, docs::AbstractVector{<:AbstractString};
         verbose::Bool = true,
@@ -719,31 +720,31 @@ function build_index(
     return index
 end
 
-using Pinecone: Pinecone, init_v3, Index
+using Pinecone: Pinecone, PineconeContextv3, PineconeIndexv3, init_v3, Index
+# TODO: change docs
 """
     build_index(
-        indexer::PTPineconeIndexer;
+        indexer::PineconeIndexer;
         namespace::AbstractString,
-        schema::AbstractPromptSchema = OpenAISchema();
         verbose::Integer = 1,
         index_id = gensym("PTPineconeIndex"),
         cost_tracker = Threads.Atomic{Float64}(0.0))
 
-Builds a `PTPineconeIndex` containing a Pinecone context (API key, index and namespace).
+Builds a `PineconeIndex` containing a Pinecone context (API key, index and namespace).
 """
 function build_index(
-        indexer::PTPineconeIndexer,
+        indexer::PineconeIndexer,
+        context::Pinecone.PineconeContextv3 = Pinecone.init_v3(""),
+        index::Pinecone.PineconeIndexv3 = "",
         namespace::AbstractString,
-        schema::PromptingTools.AbstractPromptSchema = PromptingTools.OpenAISchema();
         verbose::Integer = 1,
-        index_id = gensym("PTPineconeIndex"),
+        index_id = gensym(namespace),
         cost_tracker = Threads.Atomic{Float64}(0.0))
+    @assert !isempty(context.api_key) && !isempty(index) "Pinecone context and index not set"
 
-    pinecone_context = Pinecone.init_v3(ENV["PINECONE_API_KEY"])
-    pindex = ENV["PINECONE_INDEX"]
-    pinecone_index = pinecone_index = !isempty(pindex) ? Pinecone.Index(pinecone_context, pindex) : nothing
+    # TODO: add chunking, embedding, tags?
 
-    index = PTPineconeIndex(; id = index_id, pinecone_context, pinecone_index, namespace, schema)
+    index = PineconeIndex(; id = index_id, context, index, namespace)
 
     (verbose > 0) && @info "Index built! (cost: \$$(round(cost_tracker[], digits=3)))"
 

src/Experimental/RAGTools/rag_interface.jl

@@ -161,13 +161,6 @@ Main abstract type for storing document chunks and their embeddings. It also sto
 """
 abstract type AbstractChunkIndex <: AbstractDocumentIndex end
 
-"""
-    AbstractPTPineconeIndex <: AbstractDocumentIndex
-
-Abstract type for working with Pinecone. For now, just an empty index.
-"""
-abstract type AbstractPTPineconeIndex <: AbstractDocumentIndex end
-
 # ## Retrieval stage
 
 """
@@ -184,6 +177,8 @@ Return type from `find_closest` and `find_tags` functions.
 """
 abstract type AbstractCandidateChunks end
 
+abstract type AbstractCandidateWithChunks end
+
 # Main supertype for retrieval customizations
 abstract type AbstractRetrievalMethod end