Получение метаданных об использовании токенов из вызовов Gemini LLM в LangChain RAG RunnableSequence ⇐ Python

[Anonymous]

Я хотел бы, чтобы токен использовался моей цепочкой RAG каждый раз, когда она вызывается.
Что бы я ни делал, я не могу найти правильный способ вывода общее количество токенов модели Gemini, которую я использую.

Код: Выделить всё

import vertexai
from langchain_google_vertexai import VertexAI
from vertexai.generative_models import GenerativeModel

vertexai.init(
project='MY_PROJECT',
location="MY_LOCATION",
)

question = "What is the meaning of life"

llm = VertexAI(model_name="gemini-1.5-pro-001",)
response1 = llm.invoke(question)

llm2 = GenerativeModel("gemini-1.5-pro-001",)
response2 = llm2.generate_content(question)

Код: Выделить всё

response1

выше — это просто строка.

Код: Выделить всё

response2

— это то, что мне нужно, т. е. словарь, содержащий use_metadata, Safety_rating, Finish_reason и т. д. Но мне не удалось запустить мою цепочку RAG, используя этот подход.
Моя цепочка RAG это RunnableSequence (из langchain_core.runnables), а также я пробовал использовать обратные вызовы, поскольку цепочка не поддерживает класс vertexai.generative_models.GenerativeModel

Код: Выделить всё

from langchain_google_vertexai import VertexAI
from langchain.callbacks.base import BaseCallbackHandler
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.outputs import LLMResult
from langchain_core.messages import BaseMessage

class LoggingHandler(BaseCallbackHandler):

def on_llm_start(self, serialized, prompts, **kwargs) -> None:

print('On LLM Start: {}'.format(prompts))

def on_llm_end(self, response: LLMResult, **kwargs) -> None:

print('On LLM End: {}'.format(response))

callbacks = [LoggingHandler()]
llm = VertexAI(model_name="gemini-1.5-pro-001",)
prompt = ChatPromptTemplate.from_template("What is 1 + {number}?")

chain = prompt | llm

chain_with_callbacks = chain.with_config(callbacks=callbacks)
response = chain_with_callbacks.invoke({"number": "2"})

Это содержимое показано ниже

Код: Выделить всё

On LLM Start: ['Human: What is 1 + 2?']
On LLM End: generations=[[GenerationChunk(text='Human: What is 1 + 2?\nAssistant: 3 \n', generation_info={'is_blocked': False, 'safety_ratings': [{'category': 'HARM_CATEGORY_HATE_SPEECH', 'probability_label': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', 'probability_label': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_HARASSMENT', 'probability_label': 'NEGLIGIBLE', 'blocked': False}, {'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 'probability_label': 'NEGLIGIBLE', 'blocked': False}], 'citation_metadata': })]] llm_output=None run=None

Т.е. нет метаданных об использовании.
Есть идеи, как получить метаданные об использовании для каждого вызова цепочки RAG?

Подробнее здесь: https://stackoverflow.com/questions/790 ... unnableseq