Skip to content

Instantly share code, notes, and snippets.

@zilto
Created October 30, 2024 14:05
Show Gist options
  • Save zilto/af5f99c1832f00ec73f55f2f1a58fb53 to your computer and use it in GitHub Desktop.
Save zilto/af5f99c1832f00ec73f55f2f1a58fb53 to your computer and use it in GitHub Desktop.
import os
os.environ["OPENAI_API_KEY"] = "sk-..."
from haystack.components.embedders import SentenceTransformersTextEmbedder
from haystack.components.builders import PromptBuilder
from haystack.components.generators import OpenAIGenerator
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack import Pipeline
# 1. create components
document_store = InMemoryDocumentStore()
text_embedder = SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
prompt_builder = PromptBuilder(template="Document: {{documents}} Question: {{question}}")
retriever = InMemoryEmbeddingRetriever(document_store)
generator = OpenAIGenerator(model="gpt-4o-mini")
# 2. create pipeline
rag_pipeline = Pipeline()
# 3. add components to the pipeline
rag_pipeline.add_component("text_embedder", text_embedder)
rag_pipeline.add_component("retriever", retriever)
rag_pipeline.add_component("prompt_builder", prompt_builder)
rag_pipeline.add_component("llm", generator)
# 4. connect components
rag_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
rag_pipeline.connect("retriever", "prompt_builder.documents")
rag_pipeline.connect("prompt_builder", "llm")
# 5. visualize the pipeline
rag_pipeline.show()
# 6. run the pipeline
user_query = "What is the capital of France?"
rag_pipeline.run(text=user_query, question=user_query)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment