Packages

fnord

0.4.21

AI code archaeology

Current section

Files

Jump to
fnord lib search.ex
Raw

lib/search.ex

defmodule Search do
@moduledoc """
This module provides the functionality for the `search` sub-command.
"""
defstruct [
:project,
:query,
:limit,
:detail,
:store,
:concurrency,
:ai_module
]
@doc """
Creates a new `Search` struct with the given options.
"""
def new(opts, ai_module \\ AI) do
%Search{
project: opts[:project],
query: opts[:query],
limit: opts[:limit],
detail: opts[:detail],
store: Store.new(opts[:project]),
concurrency: opts[:concurrency],
ai_module: ai_module
}
end
@doc """
Searches the given project for previously indexed files (see `Indexing`) that
match the given query. The search results are printed to the console.
Note that the query input is first sent to OpenAI's API to generate an
embedding to match against the vector store.
"""
def run(opts, ai_module \\ AI) do
search = new(opts, ai_module)
search
|> get_results()
|> Enum.each(fn {file, score, data} ->
output_file(search, file, score, data)
end)
end
def get_results(search) do
needle = get_query_embeddings(search.query, search.ai_module)
{:ok, queue} =
Queue.start_link(search.concurrency, fn file ->
with {:ok, data} <- get_file_data(search, file) do
get_score(needle, data)
|> case do
{:ok, score} -> {file, score, data}
{:error, :no_embeddings} -> nil
end
else
_ -> nil
end
end)
results =
search
|> list_files()
|> Queue.map(queue)
|> Enum.reject(&is_nil/1)
|> Enum.sort(fn {_, score1, _}, {_, score2, _} -> score1 >= score2 end)
|> Enum.take(search.limit)
Queue.shutdown(queue)
Queue.join(queue)
results
end
defp output_file(search, file, score, data) do
if search.detail do
summary = Map.get(data, "summary")
IO.puts("""
-----
# File: #{file} | Score: #{score}
#{summary}
""")
else
IO.puts(file)
end
end
defp get_score(needle, data) do
data
|> Map.get("embeddings", [])
|> Enum.map(fn emb -> cosine_similarity(needle, emb) end)
|> case do
[] -> {:error, :no_embeddings}
scores -> {:ok, Enum.max(scores)}
end
end
defp get_query_embeddings(query, ai_module) do
{:ok, [needle]} = ai_module.get_embeddings(ai_module.new(), query)
needle
end
defp list_files(search) do
Store.list_files(search.store)
end
defp get_file_data(search, file) do
if search.detail do
Store.get(search.store, file)
else
with {:ok, embeddings} <- Store.get_embeddings(search.store, file) do
{:ok, %{"embeddings" => embeddings}}
else
{:error, _} -> {:error, :file}
end
end
end
# Computes the cosine similarity between two vectors
def cosine_similarity(vec1, vec2) do
dot_product = Enum.zip(vec1, vec2) |> Enum.reduce(0.0, fn {a, b}, acc -> acc + a * b end)
magnitude1 = :math.sqrt(Enum.reduce(vec1, 0.0, fn x, acc -> acc + x * x end))
magnitude2 = :math.sqrt(Enum.reduce(vec2, 0.0, fn x, acc -> acc + x * x end))
if magnitude1 == 0.0 or magnitude2 == 0.0 do
0.0
else
dot_product / (magnitude1 * magnitude2)
end
end
end