Simplify the BaseComponent inteface (#64)
This change remove `BaseComponent`'s: - run_raw - run_batch_raw - run_document - run_batch_document - is_document - is_batch Each component is expected to support multiple types of inputs and a single type of output. Since we want the component to work out-of-the-box with both standardized and customized use cases, supporting multiple types of inputs are expected. At the same time, to reduce the complexity of understanding how to use a component, we restrict a component to only have a single output type. To accommodate these changes, we also refactor some components to remove their run_raw, run_batch_raw... methods, and to decide the common output interface for those components. Tests are updated accordingly. Commit changes: * Add kwargs to vector store's query * Simplify the BaseComponent * Update tests * Remove support for Python 3.8 and 3.9 * Bump version 0.3.0 * Fix github PR caching still use old environment after bumping version --------- Co-authored-by: ian <ian@cinnamon.is>
This commit is contained in:
Nguyen Trung Duc (john)
committed by
GitHub
GitHub
parent
6095526dc7
commit
d79b3744cb
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import List, Union
|
||||
|
||||
from theflow import Node, Param
|
||||
|
||||
@@ -26,44 +27,34 @@ class IndexVectorStoreFromDocumentPipeline(BaseComponent):
|
||||
vector_store: Param[BaseVectorStore] = Param()
|
||||
doc_store: Param[BaseDocumentStore] = Param()
|
||||
embedding: Node[BaseEmbeddings] = Node()
|
||||
|
||||
# TODO: refer to llama_index's storage as well
|
||||
|
||||
def run_raw(self, text: str) -> None:
|
||||
document = Document(text=text, id_=str(uuid.uuid4()))
|
||||
self.run_batch_document([document])
|
||||
def run(self, text: str | list[str] | Document | list[Document]) -> None:
|
||||
input_: list[Document] = []
|
||||
if not isinstance(text, list):
|
||||
text = [text]
|
||||
|
||||
def run_batch_raw(self, text: List[str]) -> None:
|
||||
documents = [Document(text=t, id_=str(uuid.uuid4())) for t in text]
|
||||
self.run_batch_document(documents)
|
||||
for item in text:
|
||||
if isinstance(item, str):
|
||||
input_.append(Document(text=item, id_=str(uuid.uuid4())))
|
||||
elif isinstance(item, Document):
|
||||
input_.append(item)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid input type {type(item)}, should be str or Document"
|
||||
)
|
||||
|
||||
def run_document(self, text: Document) -> None:
|
||||
self.run_batch_document([text])
|
||||
|
||||
def run_batch_document(self, text: List[Document]) -> None:
|
||||
embeddings = self.embedding(text)
|
||||
embeddings = self.embedding(input_)
|
||||
self.vector_store.add(
|
||||
embeddings=embeddings,
|
||||
ids=[t.id_ for t in text],
|
||||
ids=[t.id_ for t in input_],
|
||||
)
|
||||
if self.doc_store:
|
||||
self.doc_store.add(text)
|
||||
|
||||
def is_document(self, text) -> bool:
|
||||
if isinstance(text, Document):
|
||||
return True
|
||||
elif isinstance(text, List) and isinstance(text[0], Document):
|
||||
return True
|
||||
return False
|
||||
|
||||
def is_batch(self, text) -> bool:
|
||||
if isinstance(text, list):
|
||||
return True
|
||||
return False
|
||||
self.doc_store.add(input_)
|
||||
|
||||
def save(
|
||||
self,
|
||||
path: Union[str, Path],
|
||||
path: str | Path,
|
||||
vectorstore_fname: str = VECTOR_STORE_FNAME,
|
||||
docstore_fname: str = DOC_STORE_FNAME,
|
||||
):
|
||||
@@ -80,7 +71,7 @@ class IndexVectorStoreFromDocumentPipeline(BaseComponent):
|
||||
|
||||
def load(
|
||||
self,
|
||||
path: Union[str, Path],
|
||||
path: str | Path,
|
||||
vectorstore_fname: str = VECTOR_STORE_FNAME,
|
||||
docstore_fname: str = DOC_STORE_FNAME,
|
||||
):
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from abc import abstractmethod
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import List, Union
|
||||
|
||||
from theflow import Node, Param
|
||||
|
||||
@@ -14,31 +14,7 @@ VECTOR_STORE_FNAME = "vectorstore"
|
||||
DOC_STORE_FNAME = "docstore"
|
||||
|
||||
|
||||
class BaseRetrieval(BaseComponent):
|
||||
"""Define the base interface of a retrieval pipeline"""
|
||||
|
||||
@abstractmethod
|
||||
def run_raw(self, text: str, top_k: int = 1) -> List[RetrievedDocument]:
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def run_batch_raw(
|
||||
self, text: List[str], top_k: int = 1
|
||||
) -> List[List[RetrievedDocument]]:
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def run_document(self, text: Document, top_k: int = 1) -> List[RetrievedDocument]:
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def run_batch_document(
|
||||
self, text: List[Document], top_k: int = 1
|
||||
) -> List[List[RetrievedDocument]]:
|
||||
...
|
||||
|
||||
|
||||
class RetrieveDocumentFromVectorStorePipeline(BaseRetrieval):
|
||||
class RetrieveDocumentFromVectorStorePipeline(BaseComponent):
|
||||
"""Retrieve list of documents from vector store"""
|
||||
|
||||
vector_store: Param[BaseVectorStore] = Param()
|
||||
@@ -46,53 +22,33 @@ class RetrieveDocumentFromVectorStorePipeline(BaseRetrieval):
|
||||
embedding: Node[BaseEmbeddings] = Node()
|
||||
# TODO: refer to llama_index's storage as well
|
||||
|
||||
def run_raw(self, text: str, top_k: int = 1) -> List[RetrievedDocument]:
|
||||
return self.run_batch_raw([text], top_k=top_k)[0]
|
||||
def run(self, text: str | Document, top_k: int = 1) -> list[RetrievedDocument]:
|
||||
"""Retrieve a list of documents from vector store
|
||||
|
||||
def run_batch_raw(
|
||||
self, text: List[str], top_k: int = 1
|
||||
) -> List[List[RetrievedDocument]]:
|
||||
Args:
|
||||
text: the text to retrieve similar documents
|
||||
|
||||
Returns:
|
||||
list[RetrievedDocument]: list of retrieved documents
|
||||
"""
|
||||
if self.doc_store is None:
|
||||
raise ValueError(
|
||||
"doc_store is not provided. Please provide a doc_store to "
|
||||
"retrieve the documents"
|
||||
)
|
||||
|
||||
result = []
|
||||
for each_text in text:
|
||||
emb = self.embedding(each_text)
|
||||
_, scores, ids = self.vector_store.query(embedding=emb, top_k=top_k)
|
||||
docs = self.doc_store.get(ids)
|
||||
each_result = [
|
||||
RetrievedDocument(**doc.to_dict(), score=score)
|
||||
for doc, score in zip(docs, scores)
|
||||
]
|
||||
result.append(each_result)
|
||||
emb: list[float] = self.embedding(text)[0]
|
||||
_, scores, ids = self.vector_store.query(embedding=emb, top_k=top_k)
|
||||
docs = self.doc_store.get(ids)
|
||||
result = [
|
||||
RetrievedDocument(**doc.to_dict(), score=score)
|
||||
for doc, score in zip(docs, scores)
|
||||
]
|
||||
return result
|
||||
|
||||
def run_document(self, text: Document, top_k: int = 1) -> List[RetrievedDocument]:
|
||||
return self.run_raw(text.text, top_k)
|
||||
|
||||
def run_batch_document(
|
||||
self, text: List[Document], top_k: int = 1
|
||||
) -> List[List[RetrievedDocument]]:
|
||||
return self.run_batch_raw(text=[t.text for t in text], top_k=top_k)
|
||||
|
||||
def is_document(self, text, *args, **kwargs) -> bool:
|
||||
if isinstance(text, Document):
|
||||
return True
|
||||
elif isinstance(text, List) and isinstance(text[0], Document):
|
||||
return True
|
||||
return False
|
||||
|
||||
def is_batch(self, text, *args, **kwargs) -> bool:
|
||||
if isinstance(text, list):
|
||||
return True
|
||||
return False
|
||||
|
||||
def save(
|
||||
self,
|
||||
path: Union[str, Path],
|
||||
path: str | Path,
|
||||
vectorstore_fname: str = VECTOR_STORE_FNAME,
|
||||
docstore_fname: str = DOC_STORE_FNAME,
|
||||
):
|
||||
@@ -109,7 +65,7 @@ class RetrieveDocumentFromVectorStorePipeline(BaseRetrieval):
|
||||
|
||||
def load(
|
||||
self,
|
||||
path: Union[str, Path],
|
||||
path: str | Path,
|
||||
vectorstore_fname: str = VECTOR_STORE_FNAME,
|
||||
docstore_fname: str = DOC_STORE_FNAME,
|
||||
):
|
||||
|
||||
@@ -92,7 +92,7 @@ class BaseTool(BaseComponent):
|
||||
"""Convert this tool to Langchain format to use with its agent"""
|
||||
return LCTool(name=self.name, description=self.description, func=self.run)
|
||||
|
||||
def run_raw(
|
||||
def run(
|
||||
self,
|
||||
tool_input: Union[str, Dict],
|
||||
verbose: Optional[bool] = None,
|
||||
@@ -110,23 +110,6 @@ class BaseTool(BaseComponent):
|
||||
else:
|
||||
return observation
|
||||
|
||||
def run_document(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def run_batch_raw(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def run_batch_document(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def is_document(self, *args, **kwargs) -> bool:
|
||||
"""Tool does not support processing document"""
|
||||
return False
|
||||
|
||||
def is_batch(self, *args, **kwargs) -> bool:
|
||||
"""Tool does not support processing batch"""
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def from_langchain_format(cls, langchain_tool: LCTool) -> "BaseTool":
|
||||
"""Wrapper for Langchain Tool"""
|
||||
|
||||
Reference in New Issue
Block a user