[Improvement] return all the metadata when citations flag is True (#1059)

Co-authored-by: Deven Patel <deven298@yahoo.com>

embedchain/pipeline.py

@@ -237,13 +237,7 @@ class Pipeline(EmbedChain):
             )
             result = []
             for c in context:
-                result.append(
-                    {
-                        "context": c[0],
-                        "source": c[1],
-                        "document_id": c[2],
-                    }
-                )
+                result.append({"context": c[0], "metadata": c[1]})
             return result
         else:
             # Make API call to the backend to get the results

embedchain/vectordb/chroma.py

@@ -250,9 +250,7 @@ class ChromaDB(BaseVectorDB):
             context = result[0].page_content
             if citations:
                 metadata = result[0].metadata
-                source = metadata["url"]
-                doc_id = metadata["doc_id"]
-                contexts.append((context, source, doc_id))
+                contexts.append((context, metadata))
             else:
                 contexts.append(context)
         return contexts

embedchain/vectordb/elasticsearch.py

@@ -202,7 +202,7 @@ class ElasticsearchDB(BaseVectorDB):
         if "app_id" in where:
             app_id = where["app_id"]
             query["script_score"]["query"] = {"match": {"metadata.app_id": app_id}}
-        _source = ["text", "metadata.url", "metadata.doc_id"]
+        _source = ["text", "metadata"]
         response = self.client.search(index=self._get_index(), query=query, _source=_source, size=n_results)
         docs = response["hits"]["hits"]
         contexts = []
@@ -210,9 +210,7 @@ class ElasticsearchDB(BaseVectorDB):
             context = doc["_source"]["text"]
             if citations:
                 metadata = doc["_source"]["metadata"]
-                source = metadata["url"]
-                doc_id = metadata["doc_id"]
-                contexts.append(tuple((context, source, doc_id)))
+                contexts.append(tuple((context, metadata)))
             else:
                 contexts.append(context)
         return contexts

embedchain/vectordb/opensearch.py

@@ -218,9 +218,7 @@ class OpenSearchDB(BaseVectorDB):
         for doc in docs:
             context = doc.page_content
             if citations:
-                source = doc.metadata["url"]
-                doc_id = doc.metadata["doc_id"]
-                contexts.append(tuple((context, source, doc_id)))
+                contexts.append(tuple((context, doc.metadata)))
             else:
                 contexts.append(context)
         return contexts

embedchain/vectordb/pinecone.py

@@ -154,9 +154,7 @@ class PineconeDB(BaseVectorDB):
             metadata = doc["metadata"]
             context = metadata["text"]
             if citations:
-                source = metadata["url"]
-                doc_id = metadata["doc_id"]
-                contexts.append(tuple((context, source, doc_id)))
+                contexts.append(tuple((context, metadata)))
             else:
                 contexts.append(context)
         return contexts

embedchain/vectordb/qdrant.py

@@ -219,9 +219,7 @@ class QdrantDB(BaseVectorDB):
             context = result.payload["text"]
             if citations:
                 metadata = result.payload["metadata"]
-                source = metadata["url"]
-                doc_id = metadata["doc_id"]
-                contexts.append(tuple((context, source, doc_id)))
+                contexts.append(tuple((context, metadata)))
             else:
                 contexts.append(context)
         return contexts

embedchain/vectordb/weaviate.py

@@ -271,9 +271,7 @@ class WeaviateDB(BaseVectorDB):
             context = doc["text"]
             if citations:
                 metadata = doc["metadata"][0]
-                source = metadata["url"]
-                doc_id = metadata["doc_id"]
-                contexts.append((context, source, doc_id))
+                contexts.append((context, metadata))
             else:
                 contexts.append(context)
         return contexts

embedchain/vectordb/zilliz.py

@@ -187,9 +187,7 @@ class ZillizVectorDB(BaseVectorDB):
             data = query[0]["entity"]
             context = data["text"]
             if citations:
-                source = data["url"]
-                doc_id = data["doc_id"]
-                contexts.append(tuple((context, source, doc_id)))
+                contexts.append(tuple((context, data)))
             else:
                 contexts.append(context)
         return contexts