surfacing footnotes from ragclient subtask

[jm-nab]

With the following setup, how can I surface subtask references? I can see the footnotes in the subtask, but not in the last toolkit task.

I've include the output logs, the first text value which is from the subtask has the footnote, however the toolkit task which returns, no longer has the footnote.

rag_engine = engine = RagEngine(
        retrieval_stage=RetrievalRagStage(
            retrieval_modules=[
                VectorStoreRetrievalRagModule(
                    vector_store_driver=vector_store,
                    query_params={
                        "namespace": collection_name,
                        "top_n": 50,
                        "count": 50,
                    },
                    process_query_output_fn=get_entries,
                )
            ],
            rerank_module=VertexAIRerankModule(
                rerank_driver=VertexAIRerankDriver(
                        project_id="my-project-1234",
                        location="global",
                        ranking_config="default_ranking_config",
                        model="semantic-ranker-512@latest",
                        top_n=10
                )
            )
        ),
        response_stage=ResponseRagStage(
            response_module=FootnotePromptResponseRagModule(
                prompt_driver=GooglePromptDriver(model=model_name, temperature=0.3)
            )
        )
    )
agent = Agent(
        tools = [
            RagClient(
                description="VertexAI rag engine use it to ask questions about documentation",
                rag_engine=rag_engine,
            ),
        ],
        config=StructureConfig(
            prompt_driver=GooglePromptDriver(model=model_name, temperature=0.3),
        ),
        conversation_memory=ConversationMemory(
            driver=RedisConversationMemoryDriver(
                index=conversation_id,
                conversation_id=conversation_id,
                host=host,
                port=port,
            )
        ),
        input=prompt_template,
        rules=[
            Rule(
                value="Generate english output."
            )
        ],
    )

With the example output from the logs:

07/24/24 19:33:27] INFO     Subtask 2c62a2ea570641d190b01842eb7cf675
                             Response: The xyz in abc was created in 1927.
                             Footnotes:
                             [1] https://my-source-document.xyz/item.pdf

[07/24/24 19:33:29] INFO     ToolkitTask 0e5c720c0c0148c5be41e2bdfa324d03
                             Output: The xyz in abc was created in 1927.

And the source code that I threw together for google related stuff:

from typing import Optional
from attrs import define, field
from google.api_core.exceptions import ResourceExhausted
from griptape.drivers import BaseEmbeddingDriver
import logging
import time
from griptape.artifacts import TextArtifact

from vertexai.language_models import TextEmbeddingModel
from google.cloud import discoveryengine_v1alpha as discoveryengine
from attrs import define, field
from typing import Optional
from griptape.artifacts import TextArtifact
from griptape.drivers import BaseRerankDriver
from collections.abc import Sequence
from attrs import define
from griptape.artifacts import BaseArtifact
from griptape.engines.rag import RagContext
from griptape.engines.rag.modules import BaseRerankRagModule
from griptape.artifacts import TextArtifact

log = logging.getLogger(__name__)

@define
class VertexEmbeddingDriver(BaseEmbeddingDriver):
    """
    Attributes:
        api_key: Google API key.
        model: Google model name.
        task_type: Embedding model task type (https://ai.google.dev/tutorials/python_quickstart#use_embeddings). Defaults to `retrieval_document`.
        title: Optional title for the content. Only works with `retrieval_document` task type.
    """

    DEFAULT_MODEL = "text-embedding-004"

    model: str = field(default=DEFAULT_MODEL, kw_only=True, metadata={"serializable": True})
    api_key: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": False})
    task_type: str = field(default="retrieval_document", kw_only=True, metadata={"serializable": True})
    title: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})

    def try_embed_chunk(self, chunk: str) -> list[float]:
        return self.generate_embeddings(chunk)

    def _params(self, chunk: str) -> dict:
        return {"input": chunk, "model": self.model}

    def generate_embeddings(self, chunk: str):
        init_gcloud_credentials()
        model = TextEmbeddingModel.from_pretrained(self.model)
        retry_count = 0
        while retry_count < 10:
            try:
                embedding = model.get_embeddings([chunk])[0].values
                break
            except ResourceExhausted as e:
                log.info(e)
                log.info(f"Waiting for embeding api, retry: {retry_count} because of {str(e)}")
                time.sleep(20)
                retry_count += 1
        return embedding


@define(kw_only=True)
class VertexAIRerankDriver(BaseRerankDriver):
    project_id: str = field(metadata={"serializable": True})
    location: str = field(default="global", metadata={"serializable": True})
    ranking_config: str = field(default="default_ranking_config", metadata={"serializable": True})
    model: str = field(default="semantic-ranker-512@latest", metadata={"serializable": True})
    top_n: Optional[int] = field(default=None)

    client: discoveryengine.RankServiceClient = field(init=False)
    _ranking_config: str = field(init=False)

    def __attrs_post_init__(self):
        self.client = discoveryengine.RankServiceClient()
        self._ranking_config = self.client.ranking_config_path(
            project=self.project_id,
            location=self.location,
            ranking_config=self.ranking_config,
        )

    def run(self, query: str, artifacts: list[TextArtifact]) -> list[TextArtifact]:
        records = [
            discoveryengine.RankingRecord(
                id=str(i),
                title=f"Document {i}",
                content=artifact.value
            )
            for i, artifact in enumerate(artifacts)
        ]

        request = discoveryengine.RankRequest(
            ranking_config=self._ranking_config,
            model=self.model,
            top_n=self.top_n or len(artifacts),
            query=query,
            records=records,
        )

        response = self.client.rank(request=request)

        # Create a mapping of content to original artifacts
        content_to_artifact = {artifact.value: artifact for artifact in artifacts}

        # Sort the artifacts based on the ranking results
        ranked_artifacts = []
        for record in response.records:
            original_artifact = content_to_artifact.get(record.content)
            if original_artifact:
                ranked_artifacts.append(original_artifact)

        return ranked_artifacts

@define(kw_only=True)
class VertexAIRerankModule(BaseRerankRagModule):
    def run(self, context: RagContext) -> Sequence[BaseArtifact]:
        # Ensure that all artifacts in text_chunks are TextArtifacts
        text_artifacts = [chunk for chunk in context.text_chunks if isinstance(chunk, TextArtifact)]
        
        if len(text_artifacts) != len(context.text_chunks):
            print("Warning: Some chunks were not TextArtifacts and were filtered out.")
        
        # Use the rerank_driver to rerank the text artifacts
        reranked_artifacts = self.rerank_driver.run(context.query, text_artifacts)
        return reranked_artifacts

[vasinov]

Looks like the LLM chooses not to include footnotes in the final response. A quick way to fix it would be to add a rule to the agent with something like this:

Always preserve footnotes in the final answer.

You might need to experiment with the wording of the rule, since I'm not sure how it would respond when there are no footnotes returned by subtasks. Let me know if that helps!

[jm-nab]

Thank you, I'll take a look at doing that.

[jm-nab]

Yes, that appears to work! Thank you!