Support images in example app

deepsense-ai · Oct 21, 2024 · d52de07 · d52de07
1 parent 651a567
commit d52de07
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/examples/apps/documents_chat.py b/examples/apps/documents_chat.py
@@ -19,7 +19,6 @@
 from ragbits.core.vector_store.chromadb_store import ChromaDBStore
 from ragbits.document_search import DocumentSearch
 from ragbits.document_search.documents.document import DocumentMeta
-from ragbits.document_search.documents.element import TextElement
 
 
 class QueryWithContext(BaseModel):
@@ -124,9 +123,7 @@ async def _handle_message(
         if not self._documents_ingested:
             yield self.NO_DOCUMENTS_INGESTED_MESSAGE
         results = await self.document_search.search(message[-1])
-        prompt = RAGPrompt(
-            QueryWithContext(query=message, context=[i.content for i in results if isinstance(i, TextElement)])
-        )
+        prompt = RAGPrompt(QueryWithContext(query=message, context=[i.get_key() for i in results]))
         response = await self._llm.generate(prompt)
         yield response.answer