Repository: truera/trulens
Issue: Custom Retriever Class for AWS Bedrock + Retriever Class 
Root cause: Validation Gap

I created a retriever implementing the langchain_community.retrievers.bedrock AmazonKnowledgeBasesRetriever. This is then ran through the RetrievalQA chain from LangChain, but when being wrapped by the Tru app, it gives me back a "no BaseRetriever found" 

```# select context to be used in feedback. the location of context is app specific.

from trulens_eval.app import App

import numpy as np

tru = Tru()

context = App.select_context(qa_chain)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[40], line 5
      3 import numpy as np
      4 tru = Tru()
----> 5 context = App.select_context(qa_chain)
      6 # OpenAI as feedback provider
      7 # hugs = feedback.Huggingface()
      9 f_lang_match = Feedback(hugs.language_match).on_input_output()

File /opt/conda/lib/python3.10/site-packages/trulens_eval/app.py:474, in App.select_context(cls, app)
    472 if type(app).__module__.startswith("langchain"):
    473     from trulens_eval.tru_chain import TruChain
--> 474     return TruChain.select_context(app)
    475 elif type(app).__module__.startswith("llama_index"):
    476     from trulens_eval.tru_llama import TruLlama

File /opt/conda/lib/python3.10/site-packages/trulens_eval/tru_chain.py:246, in TruChain.select_context(cls, app)
    243         pass
    245 if len(retrievers) == 0:
--> 246     raise ValueError("Cannot find any `BaseRetriever` in app.")
    248 if len(retrievers) > 1:
    249     raise ValueError(
    250         "Found more than one `BaseRetriever` in app:\n\t" + \
    251         ("\n\t".join(map(
    252             lambda lr: f"{type(lr[1])} at {lr[0]}", 
    253             retrievers)))
    254     )

ValueError: Cannot find any `BaseRetriever` in app.
```

For reference, my RetrievalQA chain is setup as such: 

```
# Configure the AWS credentials and region
# config = Config(region_name='us-east-1')
    
retriever = MyKnowledgeBaseRetriever(
    knowledge_base_id="MY KNOWLEDGE BASE ID",
    region_name="us-east-1",
    retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 6}}
)

# Create an instance of the Bedrock LLM
llm = Bedrock(model_id='meta.llama2-70b-chat-v1', model_kwargs={"temperature": 0.1, "top_p": 0.9, "max_gen_len": 1200})

# Create an instance of the RetrievalQA chain
qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever, return_source_documents=True)

qa_chain.invoke("What steps should be taken during the demo ride in the buying cycle to improve the likelihood of closing the deal?")
```

Chatting with the LangchainChat bot, it's telling me to wrap it around a class that inherits from trulens_eval.retrievers, BaseRetriever. But I don't believe that's a part of the package and it's possibly a hallucination. 

Anything helps! Thank you. 

Hi, Luis; can you share your installed versions of langchain and langchain_core ? I see that `langchain.retrievers.bedrock.AmazonKnowledgeBasesRetriever` is a subclass of `langchain.schema.BaseRetriever` which we use for finding retrievers on my installation but perhaps that is not the case in your combination of versions?

Also, does `MyKnowledgeBaseRetriever` extend `AmazonKnowledgeBasesRetriever` ?

@piotrm0 surely! I'm on `langchain-0.0.354` and ` langchain-core==0.1.6`

And yes, MyKnowledgeBaseRetriever extends AmazonKnowledgeBasesRetriever. Here is my class. 

```
import boto3
import json
from langchain_community.retrievers.bedrock import (
    AmazonKnowledgeBasesRetriever,
    RetrievalConfig,
)
from langchain.docstore.document import Document


class MyKnowledgeBaseRetriever(AmazonKnowledgeBasesRetriever): 
    def __init__(self, knowledge_base_id, region_name, retrieval_config):
        super().__init__(
            knowledge_base_id=knowledge_base_id,
            region_name=region_name,
            retrieval_config=retrieval_config
        )
        self.knowledge_base_id = knowledge_base_id
        self.region_name = region_name
        self.retrieval_config = retrieval_config

    def get_relevant_documents(self, query):
        # Create a session using your AWS credentials
        session = boto3.Session(region_name=self.region_name)

        # Create a client for the Bedrock Agent Runtime service
        bedrock_client = session.client('bedrock-agent-runtime')

        # Retrieve relevant documents from the Knowledge Base
        response = bedrock_client.retrieve(
            knowledgeBaseId=self.knowledge_base_id,
            retrievalQuery={
                'text': query
            },
            retrievalConfiguration=self.retrieval_config
        )
        
        relevant_documents = response['retrievalResults']
        docs = []

        for result in relevant_documents:
            content = result['content']['text']
            location = result['location']['s3Location']['uri']
            score = result['score']

            doc = Document(page_content=content, metadata={'location': location, 'score': score})
            docs.append(doc)

        # Return the relevant documents
        return docs
```

I think I found the issue, the `retriever` attribute of qa_chain is marked as "excluded" so we don't dump it. I made a branch/pr to ignore this exclusion marking. Please try it out: https://github.com/truera/trulens/pull/768 .

Note I had to adjust the PR a bit as including "excluded" sometimes cause big issues so it needs to be a bit more selective.

Thanks @piotrm0 I'm very privy to anything that might help, even including using a Custom TruApp or something around that. Let me know if there is anything more information wise I can provide to help diagnose the issue. 

I have a similiar issues where the retrieve is pinecone index : 

```from langchain.vectorstores import Pinecone
vectorstore = Pinecone(index, embed_model.embed_query, 'text')
rag_pipeline = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=vectorstore.as_retriever(search_kwargs={"k":5}),
    return_source_documents=True)
```


> Note I had to adjust the PR a bit as including "excluded" sometimes cause big issues so it needs to be a bit more selective.

I believe the change was pulled in about an hour ago. How would I got about pulling that into my installation? 

Spoke with @luisramirez7 separately on installing. For @luisramirez7 and others, this change is also now available in trulens_eval==0.20.2