wip improved object retrieval

llama_index/core/base_retriever.py

@@ -1,4 +1,5 @@
 """Base retriever."""
+
 from abc import abstractmethod
 from typing import Any, Dict, List, Optional
 
@@ -72,7 +73,29 @@ def _retrieve_from_object(
                 f"Retrieving from object {obj.__class__.__name__} with query {query_bundle.query_str}\n",
                 color="llama_pink",
             )
-        if isinstance(obj, NodeWithScore):
+
+        if isinstance(obj, str):
+            return [
+                NodeWithScore(
+                    node=TextNode(text=obj),
+                    score=score,
+                )
+            ]
+        elif isinstance(obj, dict):
+            from llama_index.storage.docstore.utils import json_to_doc
+
+            # check if its a node, else assume string
+            try:
+                node = json_to_doc(obj)
+                return [NodeWithScore(node=node, score=score)]
+            except Exception:
+                return [
+                    NodeWithScore(
+                        node=TextNode(text=str(obj)),
+                        score=score,
+                    )
+                ]
+        elif isinstance(obj, NodeWithScore):
             return [obj]
         elif isinstance(obj, BaseNode):
             return [NodeWithScore(node=obj, score=score)]
@@ -144,7 +167,7 @@ def _handle_recursive_retrieval(
             node = n.node
             score = n.score or 1.0
             if isinstance(node, IndexNode):
-                obj = self.object_map.get(node.index_id, None)
+                obj = node.obj or self.object_map.get(node.index_id, None)
                 if obj is not None:
                     if self._verbose:
                         print_text(

llama_index/indices/base.py

@@ -1,4 +1,5 @@
 """Base index classes."""
+
 import logging
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Generic, List, Optional, Sequence, Type, TypeVar, cast
@@ -66,7 +67,11 @@ def __init__(
         self._graph_store = self._storage_context.graph_store
 
         objects = objects or []
-        self._object_map = {obj.index_id: obj.obj for obj in objects}
+        self._object_map = {}
+        for obj in objects:
+            self._object_map[obj.index_id] = obj.obj
+            obj.obj = None  # clear the object avoid serialization issues
+
         with self._service_context.callback_manager.as_trace("index_construction"):
             if index_struct is None:
                 nodes = nodes or []

llama_index/schema.py

@@ -1,4 +1,5 @@
 """Base schema for data structures."""
+
 import json
 import textwrap
 import uuid
@@ -501,7 +502,31 @@ class IndexNode(TextNode):
     """
 
     index_id: str
-    obj: Any = Field(exclude=True)
+    obj: Any = None
+
+    def dict(self, **kwargs: Any) -> Dict[str, Any]:
+        from llama_index.storage.docstore.utils import doc_to_json
+
+        data = super().dict(**kwargs)
+
+        is_obj_serializable = False
+        try:
+            if self.obj is None:
+                data["obj"] = None
+            elif isinstance(self.obj, BaseNode):
+                data["obj"] = doc_to_json(self.obj)
+            elif isinstance(self.obj, BaseModel):
+                data["obj"] = self.obj.dict()
+            else:
+                data["obj"] = json.dumps(self.obj)
+            is_obj_serializable = True
+        except Exception:
+            pass
+
+        if not is_obj_serializable:
+            raise ValueError("IndexNode obj is not serializable: " + str(self.obj))
+
+        return data
 
     @classmethod
     def from_text_node(

llama_index/vector_stores/qdrant.py

@@ -4,6 +4,7 @@
 An index that is built on top of an existing Qdrant collection.
 
 """
+
 import logging
 from typing import Any, List, Optional, Tuple, cast
 
@@ -716,6 +717,9 @@ def parse_to_query_result(self, response: List[Any]) -> VectorStoreQueryResult:
         similarities = []
         ids = []
 
+        import pdb
+
+        pdb.set_trace()
         for point in response:
             payload = cast(Payload, point.payload)
             try: