haystack rag

haystack/h1-rag.py

@@ -0,0 +1,51 @@
+from dotenv import load_dotenv
+load_dotenv()
+
+import os
+from haystack import Pipeline, Document
+from haystack.document_stores.in_memory import InMemoryDocumentStore
+from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
+from haystack.components.generators import OpenAIGenerator
+from haystack.components.builders.answer_builder import AnswerBuilder
+from haystack.components.builders.prompt_builder import PromptBuilder
+
+# Write documents to InMemoryDocumentStore
+document_store = InMemoryDocumentStore()
+document_store.write_documents([
+    Document(content="My name is Jean and I live in Paris."), 
+    Document(content="My name is Mark and I live in Berlin."), 
+    Document(content="My name is Giorgio and I live in Rome.")
+])
+
+# Build a RAG pipeline
+prompt_template = """
+Given these documents, answer the question.
+Documents:
+{% for doc in documents %}
+    {{ doc.content }}
+{% endfor %}
+Question: {{question}}
+Answer:
+"""
+
+retriever = InMemoryBM25Retriever(document_store=document_store)
+prompt_builder = PromptBuilder(template=prompt_template)
+llm = OpenAIGenerator()
+
+rag_pipeline = Pipeline()
+rag_pipeline.add_component("retriever", retriever)
+rag_pipeline.add_component("prompt_builder", prompt_builder)
+rag_pipeline.add_component("llm", llm)
+rag_pipeline.connect("retriever", "prompt_builder.documents")
+rag_pipeline.connect("prompt_builder", "llm")
+
+# Ask a question
+question = "Who lives in Paris?"
+results = rag_pipeline.run(
+    {
+        "retriever": {"query": question},
+        "prompt_builder": {"question": question},
+    }
+)
+
+print(results["llm"]["replies"])