discourse-ai/lib/completions/dialects/llama2_classic.rb
Roman Rizzi e0bf6adb5b
DEV: Tool support for the LLM service. (#366)
This PR adds tool support to available LLMs. We'll buffer tool invocations and return them instead of making users of this service parse the response.

It also adds support for conversation context in the generic prompt. It includes bot messages, user messages, and tool invocations, which we'll trim to make sure it doesn't exceed the prompt limit, then translate them to the correct dialect.

Finally, It adds some buffering when reading chunks to handle cases when streaming is extremely slow.:M
2023-12-18 18:06:01 -03:00

77 lines
2.0 KiB
Ruby

# frozen_string_literal: true
module DiscourseAi
module Completions
module Dialects
class Llama2Classic < Dialect
class << self
def can_translate?(model_name)
%w[Llama2-*-chat-hf Llama2-chat-hf].include?(model_name)
end
def tokenizer
DiscourseAi::Tokenizer::Llama2Tokenizer
end
end
def translate
llama2_prompt = +<<~TEXT
[INST]
<<SYS>>
#{prompt[:insts]}
#{build_tools_prompt}#{prompt[:post_insts]}
<</SYS>>
[/INST]
TEXT
if prompt[:examples]
prompt[:examples].each do |example_pair|
llama2_prompt << "[INST]#{example_pair.first}[/INST]\n"
llama2_prompt << "#{example_pair.second}\n"
end
end
llama2_prompt << conversation_context if prompt[:conversation_context].present?
llama2_prompt << "[INST]#{prompt[:input]}[/INST]\n"
end
def conversation_context
return "" if prompt[:conversation_context].blank?
trimmed_context = trim_context(prompt[:conversation_context])
trimmed_context
.reverse
.reduce(+"") do |memo, context|
if context[:type] == "tool"
memo << <<~TEXT
[INST]
<function_results>
<result>
<tool_name>#{context[:name]}</tool_name>
<json>
#{context[:content]}
</json>
</result>
</function_results>
[/INST]
TEXT
elsif context[:type] == "assistant"
memo << "[INST]" << context[:content] << "[/INST]\n"
else
memo << context[:content] << "\n"
end
memo
end
end
def max_prompt_tokens
SiteSetting.ai_hugging_face_token_limit
end
end
end
end
end