-
Notifications
You must be signed in to change notification settings - Fork 98
/
Copy pathllm_chain.rs
89 lines (78 loc) · 3 KB
/
llm_chain.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
use langchain_rust::{
chain::{Chain, LLMChainBuilder},
fmt_message, fmt_placeholder, fmt_template,
language_models::llm::LLM,
llm::openai::OpenAI,
message_formatter,
prompt::HumanMessagePromptTemplate,
prompt_args,
schemas::messages::Message,
template_fstring,
};
#[tokio::main]
async fn main() {
//We can then initialize the model:
// If you'd prefer not to set an environment variable you can pass the key in directly via the `openai_api_key` named parameter when initiating the OpenAI LLM class:
//let open_ai = OpenAI::default().with_api_key("...");
let open_ai = OpenAI::default();
//Once you've installed and initialized the LLM of your choice, we can try using it! Let's ask it what LangSmith is - this is something that wasn't present in the training data so it shouldn't have a very good response.
let resp = open_ai.invoke("What is rust").await.unwrap();
println!("{}", resp);
// We can also guide it's response with a prompt template. Prompt templates are used to convert raw user input to a better input to the LLM.
let prompt = message_formatter![
fmt_message!(Message::new_system_message(
"You are world class technical documentation writer."
)),
fmt_template!(HumanMessagePromptTemplate::new(template_fstring!(
"{input}", "input"
)))
];
//We can now combine these into a simple LLM chain:
let chain = LLMChainBuilder::new()
.prompt(prompt)
.llm(open_ai.clone())
.build()
.unwrap();
//We can now invoke it and ask the same question. It still won't know the answer, but it should respond in a more proper tone for a technical writer!
match chain
.invoke(prompt_args! {
"input" => "Quien es el escritor de 20000 millas de viaje submarino",
})
.await
{
Ok(result) => {
println!("Result: {:?}", result);
}
Err(e) => panic!("Error invoking LLMChain: {:?}", e),
}
//If you want to prompt to have a list of messages you could use the `fmt_placeholder` macro
let prompt = message_formatter![
fmt_message!(Message::new_system_message(
"You are world class technical documentation writer."
)),
fmt_placeholder!("history"),
fmt_template!(HumanMessagePromptTemplate::new(template_fstring!(
"{input}", "input"
))),
];
let chain = LLMChainBuilder::new()
.prompt(prompt)
.llm(open_ai)
.build()
.unwrap();
match chain
.invoke(prompt_args! {
"input" => "Who is the writer of 20,000 Leagues Under the Sea, and what is my name?",
"history" => vec![
Message::new_human_message("My name is: luis"),
Message::new_ai_message("Hi luis"),
],
})
.await
{
Ok(result) => {
println!("Result: {:?}", result);
}
Err(e) => panic!("Error invoking LLMChain: {:?}", e),
}
}