Packages

fnord

0.1.0

AI code archaeology

Current section

Files

Jump to
fnord lib search.ex
Raw

lib/search.ex

defmodule Search do
defstruct [
:project,
:query,
:limit,
:detail,
:store
]
def run(opts) do
%{
project: project,
query: query,
limit: limit,
detail: detail
} = opts
search = %Search{
project: project,
query: query,
limit: limit,
detail: detail,
store: Store.new(project)
}
needle = get_query_embeddings(query)
search
|> list_files()
|> Enum.map(fn file ->
with {:ok, data} <- get_file_data(search, file) do
score = get_score(needle, data)
{file, score, data}
else
_ -> nil
end
end)
|> Enum.reject(&is_nil/1)
|> Enum.sort(fn {_, score1, _}, {_, score2, _} -> score1 >= score2 end)
|> Enum.take(limit)
|> Enum.each(fn {file, score, data} -> output_file(search, file, score, data) end)
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 -> Store.cosine_similarity(needle, emb) end)
|> Enum.max()
end
defp get_query_embeddings(query) do
{:ok, [needle]} = AI.get_embeddings(AI.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
end