mirror of
https://github.com/hwchase17/langchain
synced 2024-11-02 09:40:22 +00:00
df234fb171
- **Description:** The current embedchain implementation seems to handle document metadata differently than done in the current implementation of langchain and a KeyError is thrown. I would love for someone else to test this... --------- Co-authored-by: KKUGLER <kai.kugler@mercedes-benz.com> Co-authored-by: Harrison Chase <hw.chase.17@gmail.com> Co-authored-by: Deshraj Yadav <deshraj@gatech.edu>
75 lines
2.0 KiB
Python
75 lines
2.0 KiB
Python
"""Wrapper around Embedchain Retriever."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Iterable, List, Optional
|
|
|
|
from langchain_core.callbacks import CallbackManagerForRetrieverRun
|
|
from langchain_core.documents import Document
|
|
from langchain_core.retrievers import BaseRetriever
|
|
|
|
|
|
class EmbedchainRetriever(BaseRetriever):
|
|
"""`Embedchain` retriever."""
|
|
|
|
client: Any
|
|
"""Embedchain Pipeline."""
|
|
|
|
@classmethod
|
|
def create(cls, yaml_path: Optional[str] = None) -> EmbedchainRetriever:
|
|
"""
|
|
Create a EmbedchainRetriever from a YAML configuration file.
|
|
|
|
Args:
|
|
yaml_path: Path to the YAML configuration file. If not provided,
|
|
a default configuration is used.
|
|
|
|
Returns:
|
|
An instance of EmbedchainRetriever.
|
|
|
|
"""
|
|
from embedchain import Pipeline
|
|
|
|
# Create an Embedchain Pipeline instance
|
|
if yaml_path:
|
|
client = Pipeline.from_config(yaml_path=yaml_path)
|
|
else:
|
|
client = Pipeline()
|
|
return cls(client=client)
|
|
|
|
def add_texts(
|
|
self,
|
|
texts: Iterable[str],
|
|
) -> List[str]:
|
|
"""Run more texts through the embeddings and add to the retriever.
|
|
|
|
Args:
|
|
texts: Iterable of strings/URLs to add to the retriever.
|
|
|
|
Returns:
|
|
List of ids from adding the texts into the retriever.
|
|
"""
|
|
ids = []
|
|
for text in texts:
|
|
_id = self.client.add(text)
|
|
ids.append(_id)
|
|
return ids
|
|
|
|
def _get_relevant_documents(
|
|
self, query: str, *, run_manager: CallbackManagerForRetrieverRun
|
|
) -> List[Document]:
|
|
res = self.client.search(query)
|
|
|
|
docs = []
|
|
for r in res:
|
|
docs.append(
|
|
Document(
|
|
page_content=r["context"],
|
|
metadata={
|
|
"source": r["metadata"]["url"],
|
|
"document_id": r["metadata"]["doc_id"],
|
|
},
|
|
)
|
|
)
|
|
return docs
|