Repository: truera/trulens
Issue: [BUG]  Selector __record__.app.retrieve.args.query does not exist in source data.
Root cause: Validation Gap

**Bug Description**
When trying to create a custom RAG retriever instrumentation for AzureAI search, a query argument of the method doesn't seem to get recognized

**To Reproduce**
Create a rag class, declare the three instrument methods like in the examples

**Expected behavior**
It Should find the query parameter of the instrument method, as it exists

**Relevant Logs/Tracebacks**
Run of run_and_call_callback in <Thread(TP.submit with debug timeout_0, started 24764)> failed with: Selector __record__.app.retrieve.args.query does not exist in source data.
Run of run_and_call_callback in <Thread(TP.submit with debug timeout_1, started 6000)> failed with: Selector __record__.app.retrieve.args.query does not exist in source data.      
Run of run_and_call_callback in <Thread(TP.submit with debug timeout_0, started 24764)> failed with: Selector __record__.app.retrieve.rets.collect() does not exist in source data. 
Error calling wrapped function query.


**Environment:**
 - OS: Windows
 - Python 3.9
 - Trulens-eval==0.30.1

**Additional context**
Instrument class: 
class RAG_cluster:
    @instrument
    def retrieve(self, query: str) -> list:
        completion = client.chat.completions.create(
            model=OPENAI_CHAT_DEPLOYMENT,
            messages=[
                {
                    "role": "user",
                    "content": query
                }
            ],
            extra_body={
                "data_sources": [
                    {
                        "type": "azure_search",
                        "parameters": {
                            "endpoint": search_endpoint,
                            "index_name": search_index,
                            "authentication": {
                                "type": "api_key",
                                "key": AZURE_SEARCH_SERVICE_KEY}
                        }
                    }
                ]
            }
        )
        content_blobs = []
        for content in completion.choices[0].message.context["citations"]:
            content_blobs.append(content["content"])
        return content_blobs
(Other instrument methods here, the method itself tested without instrument works fine)


rag = RAG_cluster()
df = pd.read_csv("questions_ai_act.csv")
df_set = [{"query": item["question"], "response": item["answer"]} for index, item in df.iterrows()]

openai_provider = fAzureOpenai(deployment_name="chat_test")
f_qa_relevance = (
    Feedback(
        openai_provider.relevance_with_cot_reasons,
        name = "Answer Relevance",
    )
    .on(Select.RecordCalls.retrieve.args.query)
    .on_output()
)
(Other result functions like groundedness here)


tru_rag = TruCustomApp(
    rag,
    app_id="custom-rag-1",
    feedbacks=[
        f_qa_relevance,
        f_context_relevance,
        f_groundedness,
    ],
)

with tru_rag as recording:
    for i in df_set:
        rag.query(i['query'])

<!-- Answer -->
To resolve the issue where the selector `__record__.app.retrieve.args.query` does not exist in the source data, ensure that the `query` argument is being correctly captured and logged. Here are the steps to verify and fix this:

1. **Verify Instrumentation**: Ensure that the `retrieve` method is properly instrumented and that the `query` argument is being passed and recorded. The `@instrument` decorator should handle this, but double-check the method definition and its usage.

2. **Check Method Definition**: Ensure the `retrieve` method is defined correctly and the `query` argument is being passed properly.

Here's the relevant part of the code that sets up the feedback functions:

```python
from trulens_eval import Feedback
from trulens_eval import Select
from trulens_eval.feedback.provider.openai import OpenAI

provider = OpenAI()

# Define a groundedness feedback function
f_groundedness = (
    Feedback(
        provider.groundedness_measure_with_cot_reasons, name="Groundedness"
    ).on(Select.RecordCalls.retrieve.rets.collect()).on_output()
)

# Question/answer relevance between overall question and answer.
f_answer_relevance = (
    Feedback(provider.relevance_with_cot_reasons, name="Answer Relevance").on(
        Select.RecordCalls.retrieve.args.query
    ).on_output()
)

# Context relevance between question and each context chunk.
f_context_relevance = (
    Feedback(
        provider.context_relevance_with_cot_reasons, name="Context Relevance"
    ).on(Select.RecordCalls.retrieve.args.query).on(
        Select.RecordCalls.retrieve.rets
    ).aggregate(np.mean)  # choose a different aggregation method if you wish
)
```

3. **Ensure Proper Method Call**: Make sure the `query` argument is being passed correctly when calling the `retrieve` method. Here is the `retrieve` method with the `@instrument` decorator:

```python
class RAG_cluster:

    @instrument
    def retrieve(self, query: str) -> list:
        completion = client.chat.completions.create(
            model=OPENAI_CHAT_DEPLOYMENT,
            messages=[{"role": "user", "content": query}],
            extra_body={
                "data_sources": [
                    {
                        "type": "azure_search",
                        "parameters": {
                            "endpoint": search_endpoint,
                            "index_name": search_index,
                            "authentication": {
                                "type": "api_key",
                                "key": AZURE_SEARCH_SERVICE_KEY
                            }
                        }
                    }
                ]
            }
        )
        content_blobs = []
        for content in completion.choices[0].message.context["citations"]:
            content_blobs.append(content["content"])
        return content_blobs
```

4. **Debug Instrumentation**: If the issue persists, add logging or print statements inside the `retrieve` method to verify that the `query` argument is being received and recorded.

5. **Use TruChain Recorder**: Wrap your application in a `TruChain` recorder to provide logging and evaluation upon the application's use:

```python
from trulens_eval import TruChain

# Wrap application
tru_recorder = TruChain(
    rag,
    app_id='custom-rag-1',
    feedbacks=[f_context_relevance, f_answer_relevance, f_groundedness]
)

# Record application runs
with tru_recorder as recording:
    for i in df_set:
        rag.query(i['query'])
```

This setup ensures that the context produced by the application's retriever is correctly utilized and logged, which can help in debugging and resolving the issue with the `query` argument <sup>[[1]](https://github.com/truera/trulens/blob/main/trulens_eval/generated_files/all_tools.py)</sup><sup>[[2]](https://github.com/truera/trulens/blob/main/trulens_eval/trulens_eval/tru_chain.py)</sup>.


<!-- Dosu Comment Footer -->
*To continue talking to [Dosu](https://dosu.dev), mention @dosu.*


Is one of my imports maybe wrong? 
from dotenv import load_dotenv
from openai import AzureOpenAI
import numpy as np
import pandas as pd
from trulens_eval import Select, TruCustomApp, TruChain
from trulens_eval import Feedback
from trulens_eval.tru_custom_app import instrument
from trulens_eval.feedback.provider.openai import AzureOpenAI as fAzureOpenai

<!-- Answer -->
Yes, one of your imports is likely wrong. The import statement:

```python
from trulens_eval.tru_custom_app import instrument
```

should be:

```python
from trulens_eval.instruments.tru_chain import LangChainInstrument
```

This is because the `LangChainInstrument` class is located in the `trulens_eval.instruments.tru_chain` module, not directly in `trulens_eval.tru_custom_app` <sup>[[1]](https://github.com/truera/trulens/blob/main/trulens_eval/tests/unit/static/test_static.py)</sup>.


<!-- Dosu Comment Footer -->
*To continue talking to [Dosu](https://dosu.dev), mention @dosu.*


The path you mentioned "instruments.tru_chain" does not exist in the library

<!-- Answer -->
I apologize for the confusion. The correct import path for the `instrument` function is:

```python
from trulens_eval import instrument
```

Ensure that you are using this import statement in your code. Additionally, verify that the `query` argument is correctly passed and recorded in the `retrieve` method. Here is the relevant code snippet for reference:

```python
from trulens_eval import instrument

class RAG_cluster:

    @instrument
    def retrieve(self, query: str) -> list:
        completion = client.chat.completions.create(
            model=OPENAI_CHAT_DEPLOYMENT,
            messages=[{"role": "user", "content": query}],
            extra_body={
                "data_sources": [
                    {
                        "type": "azure_search",
                        "parameters": {
                            "endpoint": search_endpoint,
                            "index_name": search_index,
                            "authentication": {
                                "type": "api_key",
                                "key": AZURE_SEARCH_SERVICE_KEY
                            }
                        }
                    }
                ]
            }
        )
        content_blobs = []
        for content in completion.choices[0].message.context["citations"]:
            content_blobs.append(content["content"])
        return content_blobs
```

This should resolve the issue with the selector `__record__.app.retrieve.args.query` not existing in the source data <sup>[[1]](https://github.com/truera/trulens/blob/main/trulens_eval/generated_files/all_tools.py)</sup>.


<!-- Dosu Comment Footer -->
*To continue talking to [Dosu](https://dosu.dev), mention @dosu.*


That import does also not exist, only from trulens_eval.instruments import instrument, there is not "instrument" in trulens_eval itself

<!-- Answer -->
You are correct. The correct import statement for the `instrument` decorator is:

```python
from trulens_eval.instruments import instrument
```

Ensure you use this import statement in your code. Additionally, verify that the `query` argument is correctly passed and recorded in the `retrieve` method. Here is the relevant code snippet for reference:

```python
from trulens_eval.instruments import instrument

class RAG_cluster:

    @instrument
    def retrieve(self, query: str) -> list:
        completion = client.chat.completions.create(
            model=OPENAI_CHAT_DEPLOYMENT,
            messages=[{"role": "user", "content": query}],
            extra_body={
                "data_sources": [
                    {
                        "type": "azure_search",
                        "parameters": {
                            "endpoint": search_endpoint,
                            "index_name": search_index,
                            "authentication": {
                                "type": "api_key",
                                "key": AZURE_SEARCH_SERVICE_KEY
                            }
                        }
                    }
                ]
            }
        )
        content_blobs = []
        for content in completion.choices[0].message.context["citations"]:
            content_blobs.append(content["content"])
        return content_blobs
```

This should resolve the issue with the selector `__record__.app.retrieve.args.query` not existing in the source data <sup>[[1]](https://github.com/truera/trulens/blob/main/trulens_eval/generated_files/all_tools.py)</sup>.


<!-- Dosu Comment Footer -->
*To continue talking to [Dosu](https://dosu.dev), mention @dosu.*


How would i verify that the query arg is correctly passed?

<!-- Answer -->
To verify that the `query` argument is correctly passed and recorded in the `retrieve` method, you can use the `TruCustomApp` class from the `trulens_eval` library, which provides instrumentation for tracking and logging method calls and their arguments.

Here is an example of how you can use `TruCustomApp` to verify the `query` argument:

```python
from trulens_eval import Tru
from trulens_eval.tru_custom_app import TruCustomApp

class CustomApp:
    def __init__(self, retriever):
        self.retriever = retriever

    def retrieve_chunks(self, data):
        return self.retriever.retrieve_chunks(data)

# Assuming you have a retriever instance
retriever = CustomRetriever()
app = CustomApp(retriever)

# Instrument the retrieve_chunks method
tru_retriever = TruCustomApp(app, app_id="custom_app")

# Synchronous call
response_sync, record_sync = tru_retriever.with_record(app.retrieve_chunks, input="your_query_here", record_metadata="meta1")

# Verify the recorded arguments
print(record_sync)

# Using context manager
with tru_retriever as recording:
    response_wrapped = app.retrieve_chunks("your_query_here")

# Verify the recorded arguments
print(recording.get())
```

In this example, the `TruCustomApp` class is used to instrument the `retrieve_chunks` method of the `CustomApp` class. The `with_record` method is used to call the `retrieve_chunks` method synchronously while recording the arguments passed to it. The recorded arguments can then be printed or inspected to verify that the `query` argument is correctly passed and recorded. Additionally, you can use the context manager to achieve the same result.

The `TruCustomApp` class is specifically designed for custom applications, providing a way to instrument, log, and evaluate their behavior <sup>[[1]](https://github.com/truera/trulens/blob/main/trulens_eval/examples/expositional/end2end_apps/custom_app/custom_app.py)</sup><sup>[[2]](https://github.com/truera/trulens/blob/main/trulens_eval/tests/e2e/test_tru_llama.py)</sup><sup>[[3]](https://github.com/truera/trulens/blob/main/trulens_eval/trulens_eval/tru_llama.py)</sup><sup>[[4]](https://github.com/truera/trulens/blob/main/trulens_eval/tests/unit/test_tru_custom.py)</sup>.


<!-- Dosu Comment Footer -->
*To continue talking to [Dosu](https://dosu.dev), mention @dosu.*


I think i figured it out. The query method needs to employ self.retrieve and self.generate_completion in order for it to function, that was the issue.