Skip to content

Commit

Permalink
fix issue 15
Browse files Browse the repository at this point in the history
  • Loading branch information
@tullytim
tullytim committed Dec 31, 2023
1 parent 0ea2f31 commit 80c954a
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 96 deletions.
17 changes: 7 additions & 10 deletions src/Pinecone.jl
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
__precompile__(true)
__precompile__(false)

module Pinecone

Expand Down Expand Up @@ -215,9 +215,9 @@ julia> Pinecone.query(pinecone_context, pinecone_index, [testvector], 4)
],\"namespace\":\"\"}]}"
```
"""
function query(ctx::PineconeContext, indexobj::PineconeIndex, queries::Vector{PineconeVector}, topk::Int64=10, namespace::String="", includevalues::Bool=false, includemeta::Bool=false, filter::Dict{String, Any}=Dict{String,Any}())
rawvectors = [queries[i].values for i in 1:length(queries)]
query(ctx, indexobj, rawvectors, topk, namespace, includevalues, includemeta, filter)
function query(ctx::PineconeContext, indexobj::PineconeIndex, queries::PineconeVector, topk::Int64=10, namespace::String="", includevalues::Bool=false, includemeta::Bool=false, filter::Dict{String, Any}=Dict{String,Any}())
#rawvectors = [queries[i].values for i in 1:length(queries)]
query(ctx, indexobj, queries.values, topk, namespace, includevalues, includemeta, filter)
end

"""
Expand All @@ -242,7 +242,7 @@ julia> Pinecone.query(pinecone_context, pinecone_index,
```
"""

function query(ctx::PineconeContext, indexobj::PineconeIndex, queries::Vector{Vector{T}}, topk::Int64=10, namespace::String="", includevalues::Bool=false,
function query(ctx::PineconeContext, indexobj::PineconeIndex, vector::Vector{T}, topk::Int64=10, namespace::String="", includevalues::Bool=false,
includemeta::Bool=false, filter::Dict{String, Any}=Dict{String,Any}()) where {T<:AbstractFloat}
if topk > MAX_TOPK
throw(ArgumentError("topk larger than largest topk available of " * string(MAX_TOPK)))
Expand All @@ -254,16 +254,13 @@ function query(ctx::PineconeContext, indexobj::PineconeIndex, queries::Vector{Ve
throw(ArgumentError("topk larger than largest topk available of " * string(MAX_TOPK_WITH_META) * " when including meatadata in results"))
end
url = pineconeMakeURLForIndex(indexobj, ctx, ENDPOINTQUERYINDEX)
body = Dict{String, Any}("topK"=>topk, "vector"=>[], "includeValues"=>includevalues, "includeMetadata"=>includemeta, "namespace"=>namespace)
body = Dict{String, Any}("topK"=>topk, "vector"=>vector, "includeValues"=>includevalues, "includeMetadata"=>includemeta, "namespace"=>namespace)
if(length(filter) > 0)
body["filter"] = filter;
end
for vec in queries
push!(body["vector"], Dict{String, Any}("values"=>vec))
end

postbody = JSON3.write(body)
response = pineconeHTTPPost(url, ctx, postbody)

return String(response.body)
end #query

Expand Down
157 changes: 71 additions & 86 deletions test/testquery.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,127 +21,117 @@ v2 = [0.9, 0.8, 0.7, 0.6, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3]
testvector = Pinecone.PineconeVector("testid", v1, testdict)
testvector2 = Pinecone.PineconeVector("testid2", v2, testdict)

result = Pinecone.upsert(context, index, [testvector, testvector2], "")
result = Pinecone.upsert(context, index, [testvector, testvector2], NAMESPACE)
result = Pinecone.upsert(context, index, [testvector], "")
result = Pinecone.upsert(context, index, [testvector], NAMESPACE)
#basic test for results
result = Pinecone.query(context, index, [testvector], 4)
result = Pinecone.query(context, index, testvector, 4)
@test result !== nothing
@test typeof(result) == String



rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0

#basic test for results with a namespace
result = Pinecone.query(context, index, [testvector], 4, NAMESPACE)
result = Pinecone.query(context, index, testvector, 4, NAMESPACE)
@test result !== nothing
@test typeof(result) == String
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0

#test asking for values
result = Pinecone.query(context, index, [v1, v1], 4, "", true)
result = Pinecone.query(context, index, v1, 4, "", true)
@test result !== nothing
@test typeof(result) == String
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test haskey(rvjson["results"][1]["matches"][1], "values")
@test length(rvjson["results"][1]["matches"][1]["values"]) > 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0


#test asking for values but with namespace
result = Pinecone.query(context, index, [v1, v1], 4, NAMESPACE, true)
result = Pinecone.query(context, index, v1, 4, NAMESPACE, true)
@test result !== nothing
@test typeof(result) == String
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test haskey(rvjson["results"][1]["matches"][1], "values")
@test length(rvjson["results"][1]["matches"][1]["values"]) > 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0


#test no values returned
result = Pinecone.query(context, index, [v1, v1], 4, "", false)
result = Pinecone.query(context, index, v1, 4, "", false)
@test result !== nothing
@test typeof(result) == String
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test haskey(rvjson["results"][1]["matches"][1], "values")
@test length(rvjson["results"][1]["matches"][1]["values"]) == 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0


#test no values returned with namespace
result = Pinecone.query(context, index, [v1, v1], 4, NAMESPACE, false)
result = Pinecone.query(context, index, v1, 4, NAMESPACE, false)
@test result !== nothing
@test typeof(result) == String
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test haskey(rvjson["results"][1]["matches"][1], "values")
@test length(rvjson["results"][1]["matches"][1]["values"]) == 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0

#test metadata returned but values not
result = Pinecone.query(context, index, [v1, v1], 4, "", false, true)
result = Pinecone.query(context, index, v1, 4, "", false, true)
@test result !== nothing
@test typeof(result) == String
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test haskey(rvjson["results"][1]["matches"][1], "metadata")
@test length(rvjson["results"][1]["matches"][1]["metadata"]) > 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0


#test metadata returned but values not with namespace
result = Pinecone.query(context, index, [v1, v1], 4, NAMESPACE, false, true)
result = Pinecone.query(context, index, v1, 4, NAMESPACE, false, true)
@test result !== nothing
@test typeof(result) == String
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test haskey(rvjson["results"][1]["matches"][1], "metadata")
@test length(rvjson["results"][1]["matches"][1]["metadata"]) > 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0


#test metadata returned and values returned
result = Pinecone.query(context, index, [v1, v1], 4, "", true, true)
result = Pinecone.query(context, index, v1, 4, "", true, true)
@test result !== nothing
@test typeof(result) == String
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test haskey(rvjson["results"][1]["matches"][1], "metadata")
@test length(rvjson["results"][1]["matches"][1]["metadata"]) > 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0

#test metadata returned and values returned with namespace
result = Pinecone.query(context, index, [v1, v1], 4, NAMESPACE, true, true)
result = Pinecone.query(context, index, v1, 4, NAMESPACE, true, true)
@test result !== nothing
@test typeof(result) == String
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test haskey(rvjson["results"][1]["matches"][1], "metadata")
@test length(rvjson["results"][1]["matches"][1]["metadata"]) > 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0

#test metadata not returned but values returned
result = Pinecone.query(context, index, [v1, v1], 4, "", true, false)
result = Pinecone.query(context, index, v1, 4, "", true, false)
@test result !== nothing
@test typeof(result) == String
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test !haskey(rvjson["results"][1]["matches"][1], "metadata")
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0

#test metadata not returned but values returned with namespace
result = Pinecone.query(context, index, [v1, v1], 4, NAMESPACE, true, false)
result = Pinecone.query(context, index, v1, 4, NAMESPACE, true, false)
@test result !== nothing
@test typeof(result) == String
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test !haskey(rvjson["results"][1]["matches"][1], "metadata")
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0

@testset "Test Filters" begin
#...means the "genre" takes on both values.
moviemeta = [Dict{String, Any}("genre"=>["comedy","documentary"]), Dict{String, Any}("genre"=>["comedy","documentary"])]
result = Pinecone.upsert(context, index, ["zipA", "zipB"], [[0.1, 0.2, 0.3, 0.4, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3],
[0.9, 0.8, 0.7, 0.6, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3]], moviemeta, "mynamespace")
[0.9, 0.8, 0.7, 0.6, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3]], moviemeta, "mynamespace")
filter = """{
"genre": {
"\$in": [
Expand All @@ -154,73 +144,68 @@ v2 = [0.9, 0.8, 0.7, 0.6, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3]
"\$eq": 2019
}
}"""
result = Pinecone.query(context, index, [v1], 4, "mynamespace", true, true, JSON3.read(filter, Dict{String, Any}))
result = Pinecone.query(context, index, v1, 4, "mynamespace", true, true, JSON3.read(filter, Dict{String, Any}))
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test length(rvjson["results"][1]["matches"]) > 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0

#test other version that takes PineconeVector
result = Pinecone.query(context, index, [testvector], 4, "mynamespace", true, true, JSON3.read(filter, Dict{String, Any}))
result = Pinecone.query(context, index, testvector, 4, "mynamespace", true, true, JSON3.read(filter, Dict{String, Any}))
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test length(rvjson["results"][1]["matches"]) > 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0

#changing to filter to match nothing, now checking that. Should get no matches since no 'action' in the moviemeta
filter = """{
"genre": {
"\$in": [
"action"
"documentary"
]
},
"year": {
"\$eq": 2019
}
}"""
result = Pinecone.query(context, index, [v1], 4, "mynamespace", true, true, JSON3.read(filter, Dict{String, Any}))
result = Pinecone.query(context, index, v1, 4, "mynamespace", true, true, JSON3.read(filter, Dict{String, Any}))
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test length(rvjson["results"][1]["matches"]) == 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0

result = Pinecone.query(context, index, [testvector], 4, "mynamespace", true, true, JSON3.read(filter, Dict{String, Any}))
result = Pinecone.query(context, index, v1, 4, "mynamespace", true, true, JSON3.read(filter, Dict{String, Any}))
println(JSON3.@pretty(result))
rvjson = JSON3.read(result)
@test haskey(rvjson, "results")
@test length(rvjson["results"]) > 0
@test length(rvjson["results"][1]["matches"]) == 0
@test haskey(rvjson, "matches")
@test length(rvjson["matches"]) > 0
end

result = Pinecone.query(context, index, [v1, v1], 4)
result = Pinecone.query(context, index, v1, 4)
@test result !== nothing
@test typeof(result) == String
#should get nothing here w/ this bogus namespace
result = Pinecone.query(context, index, [v1, v1], 4, "bogusempty", true)
result = Pinecone.query(context, index, v1, 4, "bogusempty", true)
#the bogus namespace returns HTTP 400 but with JSON body. query() returns nothing so check
rvjson = JSON3.read(result)
resultsarr = rvjson["results"];
@test length(resultsarr) == 2
#now returns empty results w/ matches that'll have empty array, make sure bogus query is empty
@test length(resultsarr[1]["matches"]) == 0
@test length(resultsarr) == 0

println("***************** Test query() topK *****************")
# test topk exceeded
@test_throws ArgumentError Pinecone.query(context, index, [testvector], 10001)
@test_throws ArgumentError Pinecone.query(context, index, testvector, 10001)

println("***************** Test query() topk exceeded include results*************")
# test topk exceeded when values included in return results
@test_throws ArgumentError Pinecone.query(context, index, [testvector], 10000, "", true)
@test_throws ArgumentError Pinecone.query(context, index, testvector, 30000, "", true)

# test topk exceeded with alternate query form
@test_throws ArgumentError Pinecone.query(context, index, [v1], 10001)
@test_throws ArgumentError Pinecone.query(context, index, v1, 10001)
# test topk exceeded when values included in return results
@test_throws ArgumentError Pinecone.query(context, index, [v1], 10000, "", true)
@test_throws ArgumentError Pinecone.query(context, index, v1, 10000, "", true)

# max top k when including data or metadata is 1000, test for 1001 results
# test topk exceeded when values included in return results with includesvalues=true
@test_throws ArgumentError Pinecone.query(context, index, [v1], 1001, "", true, false)
@test_throws ArgumentError Pinecone.query(context, index, v1, 1001, "", true, false)
# test topk exceeded when values included in return results with includesmeta=true
@test_throws ArgumentError Pinecone.query(context, index, [v1], 1001, "", false, true)
@test_throws ArgumentError Pinecone.query(context, index, v1, 1001, "", false, true)

end
end

0 comments on commit 80c954a

Please sign in to comment.