Brave Search Rag
Mix.install(
[
{:req, "> 0.0.0"},
{:instructor, github: "agoodway/instructor_ex", branch: "fixes"},
{:kino, "~> 0.12.0"},
{:memoize, "~> 1.4"}
],
config: [
instructor: [
adapter: Instructor.Adapters.OpenAI,
openai: [api_key: System.fetch_env!("LB_OPENAI_API_KEY")]
]
]
)
Requirments
This LiveBook Notebook requires 2 API keys:
-
Brave Search API: https://brave.com/search/api/
- The Brave API Key must have access to the Data for AI Pro plan.
-
PRO_BRAVE_API_KEY
secret needs to be set
-
OpenAI API: https://openai.com/index/openai-api/
-
OPENAI_API_KEY
secret needs to be set
-
Define Instructor Schema
defmodule Researcher do
use Ecto.Schema
use Instructor.Validator
@doc """
## Field Descriptions:
- response: answer to the prompt
- related_questions: a list of questions related to the prompt, limit to a max of 3
## Additional Instructions
- set all id fields to nil
"""
@primary_key false
embedded_schema do
field(:response, :string)
embeds_many :related_questions, Question do
field(:content, :string)
end
end
@impl true
def validate_changeset(changeset) do
changeset
end
def prompt(text, context) do
Instructor.chat_completion(
model: "gpt-3.5-turbo",
response_model: Researcher,
max_retries: 3,
messages: [
%{
role: "system",
content: "You are an expert researcher."
},
%{
role: "user",
content: """
Use the following context to enrich your answer to the prompt:
```
#{context}
```
"""
},
%{
role: "user",
content: """
Prompt:
```
#{text}
```
"""
}
]
)
end
end
defmodule ReqFileCache do
@moduledoc """
A simple HTTP cache for `req` that do not use headers.
If the file is not found on disk, the download will occur,
otherwise response will be read from disk.
"""
require Logger
def attach(%Req.Request{} = request, options \\ []) do
request
|> Req.Request.register_options([:cache_dir])
|> Req.Request.merge_options(options)
|> Req.Request.append_request_steps(custom_cache: &request_local_cache_step/1)
|> Req.Request.prepend_response_steps(custom_cache: &response_local_cache_step/1)
end
def request_local_cache_step(request) do
# TODO: handle a form of expiration - for now it is
# acceptable to wipe out the whole folder manually for me
# NOTE: race condition here, for parallel queries
if File.exists?(path = cache_path(request)) do
Logger.info("File found in cache (#{path})")
# a request step can return a {req,resp} tuple,
# and this will bypass the remaining request steps
{request, load_cache(path)}
else
request
end
end
def response_local_cache_step({request, response}) do
unless File.exists?(path = cache_path(request)) do
if response.status == 200 do
Logger.info("Saving file to cache (#{path})")
write_cache(path, response)
else
Logger.info("Status is #{response.status}, not saving file to disk")
end
end
{request, response}
end
# https://github.com/wojtekmach/req/blob/102b9aa6c6ff66f00403054a0093c4f06f6abc2f/lib/req/steps.ex#L1268
def cache_path(cache_dir, request) do
cache_key =
Enum.join(
[
request.url.host,
Atom.to_string(request.method),
:crypto.hash(:sha256, :erlang.term_to_binary(request.url))
|> Base.encode16(case: :lower)
],
"-"
)
Path.join(cache_dir, cache_key)
end
def cache_path(request) do
cache_path(request.options[:cache_dir], request)
end
# https://github.com/wojtekmach/req/blob/102b9aa6c6ff66f00403054a0093c4f06f6abc2f/lib/req/steps.ex#L1288-L1290
def load_cache(path) do
path |> File.read!() |> :erlang.binary_to_term()
end
# https://github.com/wojtekmach/req/blob/102b9aa6c6ff66f00403054a0093c4f06f6abc2f/lib/req/steps.ex#L1283-L1286
def write_cache(path, response) do
File.mkdir_p!(Path.dirname(path))
File.write!(path, :erlang.term_to_binary(response))
end
end
Let’s Go!
defmodule BraveSearch do
@cache_dir Path.join(__ENV__.file, "../cache") |> Path.expand()
def web_search(text) do
params = [summary: 1, q: text]
response =
Req.get!(new(), url: "/res/v1/web/search", params: params, cache_dir: @cache_dir)
response.body
end
def summarizer_search(summarizer_key) do
params = [entity_info: 1, key: summarizer_key]
response =
Req.get!(new(),
url: "/res/v1/summarizer/search",
params: params,
cache_dir: @cache_dir
)
response.body
end
defp new do
Req.new(
base_url: "https://api.search.brave.com/",
headers: [
{"X-Subscription-Token", System.fetch_env!("LB_PRO_BRAVE_API_KEY")},
{"Accept", "application/json"}
]
)
|> ReqFileCache.attach()
end
end
text_input = Kino.Input.textarea("Text", default: "What is the best way to cook chicken?")
form = Kino.Control.form([text: text_input], submit: "Run")
frame = Kino.Frame.new()
Kino.listen(form, fn %{data: %{text: text}} ->
Kino.Frame.render(frame, Kino.Text.new("Running..."))
summarizer_key = BraveSearch.web_search(text)["summarizer"]["key"]
search_context =
BraveSearch.summarizer_search(summarizer_key)["summary"] |> List.first() |> Map.get("data")
{:ok, research} = Researcher.prompt(text, search_context)
html =
"""
Answer
<%= @research.response %>
Related questions:
<%= for question <- @research.related_questions do %>
- <%= question.content %>
<% end %>
"""
|> EEx.eval_string(assigns: [research: research])
Kino.Frame.render(frame, Kino.HTML.new(html))
end)
Kino.Layout.grid([form, frame], boxed: true, gap: 16)