mirror of https://github.com/0xplaygrounds/rig
76 lines
3.0 KiB
Rust
76 lines
3.0 KiB
Rust
use std::env;
|
|
|
|
use rig::{
|
|
embeddings::EmbeddingsBuilder,
|
|
parallel,
|
|
pipeline::{self, agent_ops::lookup, passthrough, Op},
|
|
providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
|
|
vector_store::in_memory_store::InMemoryVectorStore,
|
|
};
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<(), anyhow::Error> {
|
|
// Create OpenAI client
|
|
let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
|
|
let openai_client = Client::new(&openai_api_key);
|
|
|
|
let embedding_model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);
|
|
|
|
// Create embeddings for our documents
|
|
let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
|
|
.document("Definition of a *flurbo*: A flurbo is a green alien that lives on cold planets")?
|
|
.document("Definition of a *glarb-glarb*: A glarb-glarb is a ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.")?
|
|
.document("Definition of a *linglingdong*: A term used by inhabitants of the far side of the moon to describe humans.")?
|
|
.build()
|
|
.await?;
|
|
|
|
// Create vector store with the embeddings
|
|
let vector_store = InMemoryVectorStore::from_documents(embeddings);
|
|
|
|
// Create vector store index
|
|
let index = vector_store.index(embedding_model);
|
|
|
|
let agent = openai_client.agent("gpt-4")
|
|
.preamble("
|
|
You are a dictionary assistant here to assist the user in understanding the meaning of words.
|
|
")
|
|
.build();
|
|
|
|
let chain = pipeline::new()
|
|
// Chain a parallel operation to the current chain. The parallel operation will
|
|
// perform a lookup operation to retrieve additional context from the user prompt
|
|
// while simultaneously applying a passthrough operation. The latter will allow
|
|
// us to forward the initial prompt to the next operation in the chain.
|
|
.chain(parallel!(
|
|
passthrough(),
|
|
lookup::<_, _, String>(index, 1), // Required to specify document type
|
|
))
|
|
// Chain a "map" operation to the current chain, which will combine the user
|
|
// prompt with the retrieved context documents to create the final prompt.
|
|
// If an error occurs during the lookup operation, we will log the error and
|
|
// simply return the initial prompt.
|
|
.map(|(prompt, maybe_docs)| match maybe_docs {
|
|
Ok(docs) => format!(
|
|
"Non standard word definitions:\n{}\n\n{}",
|
|
docs.into_iter()
|
|
.map(|(_, _, doc)| doc)
|
|
.collect::<Vec<_>>()
|
|
.join("\n"),
|
|
prompt,
|
|
),
|
|
Err(err) => {
|
|
println!("Error: {}! Prompting without additional context", err);
|
|
format!("{prompt}")
|
|
}
|
|
})
|
|
// Chain a "prompt" operation which will prompt out agent with the final prompt
|
|
.prompt(agent);
|
|
|
|
// Prompt the agent and print the response
|
|
let response = chain.call("What does \"glarb-glarb\" mean?").await?;
|
|
|
|
println!("{:?}", response);
|
|
|
|
Ok(())
|
|
}
|