aboutsummaryrefslogtreecommitdiff
path: root/R2R/r2r/examples/scripts/run_web_multi_rag.py
diff options
context:
space:
mode:
Diffstat (limited to 'R2R/r2r/examples/scripts/run_web_multi_rag.py')
-rwxr-xr-xR2R/r2r/examples/scripts/run_web_multi_rag.py54
1 files changed, 54 insertions, 0 deletions
diff --git a/R2R/r2r/examples/scripts/run_web_multi_rag.py b/R2R/r2r/examples/scripts/run_web_multi_rag.py
new file mode 100755
index 00000000..3ba70b74
--- /dev/null
+++ b/R2R/r2r/examples/scripts/run_web_multi_rag.py
@@ -0,0 +1,54 @@
+import fire
+
+from r2r import R2RBuilder, SerperClient, WebSearchPipe
+from r2r.base.abstractions.llm import GenerationConfig
+from r2r.main.assembly.factory_extensions import R2RPipeFactoryWithMultiSearch
+
+
+def run_rag_pipeline(query="Who was Aristotle?"):
+ # Initialize a web search pipe
+ web_search_pipe = WebSearchPipe(serper_client=SerperClient())
+
+ # Define a new synthetic query generation template
+ synthetic_query_generation_template = {
+ "name": "synthetic_query_generation_template",
+ "template": """
+ ### Instruction:
+ Given the following query, write a double newline separated list of up to {num_outputs} advanced queries meant to help answer the original query.
+ DO NOT generate any single query which is likely to require information from multiple distinct documents.
+ EACH single query will be used to carry out a cosine similarity semantic search over distinct indexed documents.
+ FOR EXAMPLE, if asked `how do the key themes of Great Gatsby compare with 1984`, the two queries would be
+ `What are the key themes of Great Gatsby?` and `What are the key themes of 1984?`.
+ Here is the original user query to be transformed into answers:
+
+ ### Query:
+ {message}
+
+ ### Response:
+ """,
+ "input_types": {"num_outputs": "int", "message": "str"},
+ }
+
+ # Build the R2R application with the custom pipeline
+ app = (
+ R2RBuilder()
+ .with_pipe_factory(R2RPipeFactoryWithMultiSearch)
+ .build(
+ # override inputs consumed in building the MultiSearchPipe
+ multi_inner_search_pipe_override=web_search_pipe,
+ query_generation_template_override=synthetic_query_generation_template,
+ )
+ )
+
+ # Run the RAG pipeline through the R2R application
+ result = app.rag(
+ query,
+ rag_generation_config=GenerationConfig(model="gpt-4o"),
+ )
+
+ print(f"Search Results:\n\n{result.search_results}")
+ print(f"RAG Results:\n\n{result.completion}")
+
+
+if __name__ == "__main__":
+ fire.Fire(run_rag_pipeline)