From 58ffa562c23877d514fa0c3e870e417039c5039a Mon Sep 17 00:00:00 2001
From: Alistair Rogers <alistairlr112@gmail.com>
Date: Thu, 4 Jan 2024 11:53:58 +0000
Subject: [PATCH] change example to reference the greatest politician of all
 time - Super Mario

---
 integrations/ollama/example/example.py | 28 ++++++++++++++++++++------
 integrations/ollama/pyproject.toml     |  2 ++
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/integrations/ollama/example/example.py b/integrations/ollama/example/example.py
index 73206fc9a..6e310a3f0 100644
--- a/integrations/ollama/example/example.py
+++ b/integrations/ollama/example/example.py
@@ -1,22 +1,36 @@
-from haystack import Pipeline
+from haystack import Document, Pipeline
 from haystack.components.builders.prompt_builder import PromptBuilder
 from haystack.components.retrievers import InMemoryBM25Retriever
-from haystack.document_stores import InMemoryDocumentStore
+from haystack.document_stores.in_memory import InMemoryDocumentStore
 
 from ollama_haystack import OllamaGenerator
 
+document_store = InMemoryDocumentStore()
+document_store.write_documents(
+    [
+        Document(content="Super Mario was an important politician"),
+        Document(content="Mario owns several castles and uses them to conduct important political business"),
+        Document(
+            content="Super Mario was a successful military leader who fought off several invasion attempts by "
+            "his arch rival - Bowser"
+        ),
+    ]
+)
+
+query = "Who is Super Mario?"
+
 template = """
-Given the following information, answer the question.
+Given only the following information, answer the question.
+Ignore your own knowledge.
 
 Context:
 {% for document in documents %}
     {{ document.content }}
 {% endfor %}
 
-Question: What's the official language of {{ country }}?
+Question: {{ query }}?
 """
 pipe = Pipeline()
-document_store = InMemoryDocumentStore()
 
 pipe.add_component("retriever", InMemoryBM25Retriever(document_store=document_store))
 pipe.add_component("prompt_builder", PromptBuilder(template=template))
@@ -24,4 +38,6 @@
 pipe.connect("retriever", "prompt_builder.documents")
 pipe.connect("prompt_builder", "llm")
 
-pipe.run({"prompt_builder": {"country": "Ghana"}})
+response = pipe.run({"prompt_builder": {"query": query}, "retriever": {"query": query}})
+
+print(response["llm"]["replies"])
diff --git a/integrations/ollama/pyproject.toml b/integrations/ollama/pyproject.toml
index 4505138d9..551ee299b 100644
--- a/integrations/ollama/pyproject.toml
+++ b/integrations/ollama/pyproject.toml
@@ -145,6 +145,8 @@ ban-relative-imports = "all"
 [tool.ruff.per-file-ignores]
 # Tests can use magic values, assertions, and relative imports
 "tests/**/*" = ["PLR2004", "S101", "TID252"]
+# Examples can print their output
+"example/**" = ["T201"]
 
 
 [tool.coverage.run]