Skip to content

Commit

Permalink
Add tools support
Browse files Browse the repository at this point in the history
  • Loading branch information
riccardobl committed May 12, 2024
1 parent 52abed1 commit 1c685bb
Showing 1 changed file with 38 additions and 6 deletions.
44 changes: 38 additions & 6 deletions src/rag/index.js
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ async function run() {
let cacheDurationHint = undefined;
let noCache = false;
let warmUp = false;
let useTools=false;
let toolsResultTemplate ="{{TOOL_RESULT}}";

const documents = []; // plain text documents
const documentsUrls = []; // urls to documents
Expand All @@ -29,6 +31,12 @@ async function run() {
}
}
}

// Early return if no documents
if (documents.length === 0 && documentsUrls.length === 0) {
Host.outputString("");
return;
}


for(const param of job.param){
Expand All @@ -46,21 +54,40 @@ async function run() {
}else if(param.key=="cache-duration-hint"){
cacheDurationHint=param.value;
}else if(param.key=="no-cache"){
noCache=param.value=="true";
noCache = param.value=="true";
}else if(param.key=="warm-up"){
noCache||=param.value=="true";
warmUp||=param.value=="true";
warmUp = param.value == "true";
}else if(param.key=="use-tools"){
useTools=param.value=="true";
}else if(param.key=="tools-result-template"){
toolsResultTemplate=param.value;
}
}

const cacheParams = [];
if(cacheDurationHint){
cacheParams.push(await Job.newParam("cache-duration-hint", cacheDurationHint));
}
if(noCache){
if(noCache||warmUp){
cacheParams.push(await Job.newParam("no-cache", noCache));
}

// Send tool req
let toolReq;
if (!warmUp && useTools){
Job.log("Send tool request...");
toolReq = Job.subrequest({
runOn: "openagents/tools",
outputFormat: "application/json",
inputs: [
await Job.newInputData(queries, "text", "queries")
],
params: [
...cacheParams
]
});
}

Job.log("Starting rag pipeline with k="+topK+", max-tokens="+maxTokens+", quantize="+quantize+", overlap="+overlap+", cache-duration-hint="+cacheDurationHint+", no-cache="+noCache);
Job.log("Fetch documents...");
const downloadDocumentsReq = Job.subrequest({
Expand Down Expand Up @@ -108,7 +135,7 @@ async function run() {
Job.log("Search...");
const searchReq = Job.subrequest({
runOn: "openagents/search",
outputFormat: "application/hyperdrive+bundle",
outputFormat: "application/json",
inputs: [
await Job.newInputData(documentsEmbeddingBundle, "application/hyperdrive+bundle", "index"),
await Job.newInputData(queriesEmbeddingBundle, quantize?"application/json":"application/hyperdrive+bundle", "query")
Expand All @@ -121,13 +148,18 @@ async function run() {
});

try{

const searchResult = JSON.parse(await Job.waitForContent(searchReq));
Job.log("Merge context... "+searchResult.length+" results found");
let newContext ="";
for(const result of searchResult){
newContext+=result.value+"\n";
}
if(toolReq){
Job.log("Merge tools result...");
let toolResult = await Job.waitForContent(toolReq);
toolResult = toolsResultTemplate.replace("{{TOOL_RESULT}}", toolResult);
newContext+=toolResult+"\n";
}
Host.outputString(newContext);
}catch(e){
await Job.log("Error! "+e.message);
Expand Down

0 comments on commit 1c685bb

Please sign in to comment.