Skip to content
Snippets Groups Projects
Unverified Commit 48a286ab authored by Logan's avatar Logan Committed by GitHub
Browse files

Ensure load and search tool spec loads documents (#11733)

parent 0ae69d46
No related branches found
No related tags found
No related merge requests found
...@@ -4,12 +4,12 @@ Tool that wraps any data loader, and is able to load data on-demand. ...@@ -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 typing import Any, Dict, List, Optional, Type
from llama_index.core.bridge.pydantic import BaseModel from llama_index.core.bridge.pydantic import BaseModel
from llama_index.core.indices.base import BaseIndex from llama_index.core.indices.base import BaseIndex
from llama_index.core.indices.vector_store import VectorStoreIndex 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.function_tool import FunctionTool
from llama_index.core.tools.tool_spec.base import SPEC_FUNCTION_TYPE, BaseToolSpec from llama_index.core.tools.tool_spec.base import SPEC_FUNCTION_TYPE, BaseToolSpec
from llama_index.core.tools.types import ToolMetadata from llama_index.core.tools.types import ToolMetadata
...@@ -123,6 +123,19 @@ class LoadAndSearchToolSpec(BaseToolSpec): ...@@ -123,6 +123,19 @@ class LoadAndSearchToolSpec(BaseToolSpec):
def load(self, *args: Any, **kwargs: Any) -> Any: def load(self, *args: Any, **kwargs: Any) -> Any:
# Call the wrapped tool and save the result in the index # Call the wrapped tool and save the result in the index
docs = self._tool(*args, **kwargs).raw_output 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: if self._index:
for doc in docs: for doc in docs:
self._index.insert(doc, **self._index_kwargs) self._index.insert(doc, **self._index_kwargs)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment