From 91112b75aadfadc45fc46b39508a9a656bdc031a Mon Sep 17 00:00:00 2001
From: Arjun Attam <arjunattam@gmail.com>
Date: Fri, 12 Apr 2024 19:00:59 +0530
Subject: [PATCH] feat: update rag example to use parameters (#122)

---
 examples/rag/empiricalrc.json | 17 ++++++++++++++++-
 examples/rag/rag.py           |  7 +++----
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/examples/rag/empiricalrc.json b/examples/rag/empiricalrc.json
index a732dfd1..42903766 100644
--- a/examples/rag/empiricalrc.json
+++ b/examples/rag/empiricalrc.json
@@ -2,9 +2,24 @@
   "$schema": "https://assets.empirical.run/config/schema/v0.4.0.json",
   "runs": [
     {
-      "name": "rag script run",
       "type": "py-script",
       "path": "rag.py",
+      "parameters": {
+        "model": "gpt-3.5-turbo"
+      },
+      "scorers": [
+        {
+          "type": "py-script",
+          "path": "score.py"
+        }
+      ]
+    },
+    {
+      "type": "py-script",
+      "path": "rag.py",
+      "parameters": {
+        "model": "gpt-4-turbo-preview"
+      },
       "scorers": [
         {
           "type": "py-script",
diff --git a/examples/rag/rag.py b/examples/rag/rag.py
index 84a73a14..c6b40a4c 100644
--- a/examples/rag/rag.py
+++ b/examples/rag/rag.py
@@ -1,10 +1,7 @@
-import nest_asyncio
 from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
 from llama_index.llms.openai import OpenAI as LlamaOpenAI
 from llama_index.core.node_parser import SentenceSplitter
 
-nest_asyncio.apply()
-
 
 def execute(inputs, parameters):
     # load documents
@@ -21,7 +18,9 @@ def build_query_engine(llm):
         query_engine = vector_index.as_query_engine(similarity_top_k=2, llm=llm)
         return query_engine
 
-    query_engine = build_query_engine(llm=LlamaOpenAI(model="gpt-3.5-turbo"))
+    query_engine = build_query_engine(
+        llm=LlamaOpenAI(model=parameters.get("model", "gpt-3.5-turbo"))
+    )
     response = query_engine.query(question)
     output = response.response
     contexts = [c.node.get_content() for c in response.source_nodes]