From ccc02d944c84a96c6d2edfbe951af5a1b10d0340 Mon Sep 17 00:00:00 2001
From: exowanderer <jdfraine@gmail.com>
Date: Sat, 9 Nov 2024 00:14:04 +0100
Subject: [PATCH] updated pipeline to avoid corner case with empty item

---
 wikidatachat/api.py                    | 14 ++++++++++----
 wikidatachat/vector_store_interface.py |  2 ++
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/wikidatachat/api.py b/wikidatachat/api.py
index 1689de3..40cc753 100644
--- a/wikidatachat/api.py
+++ b/wikidatachat/api.py
@@ -145,7 +145,13 @@ async def api(
     logger.debug(f'{answer=}')
 
     # Return the processed answer and sources.
-    return {
-        "answer": answer.data.content,
-        "sources": sources
-    }
+    if hasattr(answer.data, 'content'):
+        return {
+            "answer": answer.data.content,
+            "sources": sources
+        }
+    else:
+        return {
+            "answer": '',
+            "sources": sources
+        }
diff --git a/wikidatachat/vector_store_interface.py b/wikidatachat/vector_store_interface.py
index 2f50fab..84709a9 100644
--- a/wikidatachat/vector_store_interface.py
+++ b/wikidatachat/vector_store_interface.py
@@ -4,6 +4,7 @@
 from tqdm import tqdm
 
 from haystack import Document  # , Pipeline
+from haystack.utils import ComponentDevice
 from haystack.components.embedders import SentenceTransformersDocumentEmbedder
 from haystack.document_stores.in_memory import InMemoryDocumentStore
 from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
@@ -189,6 +190,7 @@ def make_embedder(embedding_model: str = EMBEDDING_MODEL, device: str = 'cpu'):
     """
     # Use GPU if available, otherwise fallback to CPU.
     device = "cuda" if torch.cuda.is_available() else device
+    device = ComponentDevice.from_str(device)
 
     logger.info(  # Log the name of the embedding model being used.
         'GPU is available. Using GPU.'