Skip to content

Commit

Permalink
updated data extraction for for changes
Browse files Browse the repository at this point in the history
  • Loading branch information
brainlid committed Apr 27, 2024
1 parent e8640a1 commit 6f3eef2
Showing 1 changed file with 14 additions and 3 deletions.
17 changes: 14 additions & 3 deletions lib/chains/data_extraction_chain.ex
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ defmodule LangChain.Chains.DataExtractionChain do
require Logger
alias LangChain.PromptTemplate
alias LangChain.Message
alias LangChain.Message.ToolCall
alias LangChain.Chains.LLMChain

@function_name "information_extraction"
Expand Down Expand Up @@ -101,15 +102,19 @@ Passage:
{:ok, chain} = LLMChain.new(%{llm: llm, verbose: verbose})

chain
|> LLMChain.add_functions(build_extract_function(json_schema))
|> LLMChain.add_tools(build_extract_function(json_schema))
|> LLMChain.add_messages(messages)
|> LLMChain.run()
|> case do
{:ok, _updated_chain,
%Message{
role: :assistant,
function_name: @function_name,
arguments: %{"info" => info}
tool_calls: [
%ToolCall{
name: @function_name,
arguments: %{"info" => info}
}
]
}}
when is_list(info) ->
{:ok, info}
Expand All @@ -136,6 +141,12 @@ Passage:
LangChain.Function.new!(%{
name: @function_name,
description: "Extracts the relevant information from the passage.",
function: fn args, _context ->
# NOTE: The function is not executed here because we won't be returning
# this to the LLM. The LLMChain does not run the function, but stops at
# the request for it.
{:ok, args}
end,
parameters_schema: %{
type: "object",
properties: %{
Expand Down

0 comments on commit 6f3eef2

Please sign in to comment.