From 48a286ab95d128c034cad3122fec0c27ea09883f Mon Sep 17 00:00:00 2001
From: Logan <logan.markewich@live.com>
Date: Thu, 7 Mar 2024 11:23:57 -0600
Subject: [PATCH] Ensure load and search tool spec loads documents (#11733)
---
.../core/tools/tool_spec/load_and_search/base.py | 15 ++++++++++++++-
1 file changed, 14 insertions(+), 1 deletion(-)
diff --git a/llama-index-core/llama_index/core/tools/tool_spec/load_and_search/base.py b/llama-index-core/llama_index/core/tools/tool_spec/load_and_search/base.py
index 935cd56b67..6626f9e8b3 100644
--- a/llama-index-core/llama_index/core/tools/tool_spec/load_and_search/base.py
+++ b/llama-index-core/llama_index/core/tools/tool_spec/load_and_search/base.py
@@ -4,12 +4,12 @@ Tool that wraps any data loader, and is able to load data on-demand.
"""
-
from typing import Any, Dict, List, Optional, Type
from llama_index.core.bridge.pydantic import BaseModel
from llama_index.core.indices.base import BaseIndex
from llama_index.core.indices.vector_store import VectorStoreIndex
+from llama_index.core.schema import Document
from llama_index.core.tools.function_tool import FunctionTool
from llama_index.core.tools.tool_spec.base import SPEC_FUNCTION_TYPE, BaseToolSpec
from llama_index.core.tools.types import ToolMetadata
@@ -123,6 +123,19 @@ class LoadAndSearchToolSpec(BaseToolSpec):
def load(self, *args: Any, **kwargs: Any) -> Any:
# Call the wrapped tool and save the result in the index
docs = self._tool(*args, **kwargs).raw_output
+
+ # convert to Document if necessary
+ if isinstance(docs, list):
+ for i, doc in enumerate(docs):
+ if not isinstance(doc, Document):
+ docs[i] = Document(text=str(doc))
+ elif isinstance(docs, str):
+ docs = [Document(text=docs)]
+ elif isinstance(docs, Document):
+ docs = [docs]
+ else:
+ doc = [Document(text=str(docs))]
+
if self._index:
for doc in docs:
self._index.insert(doc, **self._index_kwargs)
--
GitLab