mirror of
https://github.com/tiennm99/test-haystack.git
synced 2026-06-05 22:13:16 +00:00
47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
from haystack_integrations.components.generators.ollama import OllamaGenerator
|
|
|
|
from haystack import Pipeline, Document
|
|
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
|
|
from haystack.components.builders.prompt_builder import PromptBuilder
|
|
from haystack.document_stores.in_memory import InMemoryDocumentStore
|
|
|
|
# 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 = OllamaGenerator(url = "http://localhost:11434", model="llama3.2:1b")
|
|
|
|
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"]) |