Agent + RAG APIs

agent.py

import vertexai
from vertexai.preview import reasoning_engines
from vertexai.preview import rag
from vertexai.preview.generative_models import Tool

from langchain.agents.format_scratchpad.tools import format_to_tool_messages
from langchain.memory import ChatMessageHistory
from langchain_core import prompts
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

PROJECT_ID = "sascha-playground-doit"
LOCATION = "us-central1"
STAGING_BUCKET = "gs://doit-llm"

vertexai.init(project=PROJECT_ID, 
              location=LOCATION,
              staging_bucket=STAGING_BUCKET)

model = "gemini-2.0-flash-exp"
#model = "gemini-1.5-flash-002"


def get_session_history(session_id: str):
    from langchain_google_firestore import FirestoreChatMessageHistory
    from google.cloud import firestore

    client = firestore.Client(project="sascha-playground-doit")
    return FirestoreChatMessageHistory(
        client=client,
        session_id=session_id,
        collection="history",
        encode_message=False,
    )


book_retrieval_tool = Tool.from_retrieval(
    retrieval=rag.Retrieval(
        source=rag.VertexRagStore(
            rag_resources=[rag.RagResource(
                rag_corpus="projects/234439745674/locations/us-central1/ragCorpora/3918694625765752832")],
            similarity_top_k=5,
            vector_distance_threshold=0.5,
        ),
    )
)

prompt = {
    "history": lambda x: x["history"],
    "input": lambda x: x["input"],
    "agent_scratchpad": (lambda x: format_to_tool_messages(x["intermediate_steps"])),
} | prompts.ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """you are a digital librarian if a users says Hi you say: welcome to doits book store I am there to help you answer any questions about our books and help you find the right one for you
            follow instructions precisely:
            - you are always friendly and you answer with engaging converstions
            - use the chat history for the conversation
            
            example
            
            """,
        ),
        prompts.MessagesPlaceholder(variable_name="history"),
        ("user", "{input}"),
        prompts.MessagesPlaceholder(variable_name="agent_scratchpad"),
    ]
)


qa_system_prompt = """you are a digital librarian if a users says Hi you say: welcome to doits book store I am there to help you answer any questions about our books and help you find the right one for you  \
            follow instructions precisely: \
            - you are always friendly and you answer with engaging converstions \
            - use the chat history for the conversation \

{input}"""

qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", qa_system_prompt),
        MessagesPlaceholder("history"),
        ("user", "{input}"),
    ]
)

qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", qa_system_prompt),
        MessagesPlaceholder("history"),
        ("user", "Context: {history}. Query: {input}"),
    ]
)

print(qa_prompt)

agent = reasoning_engines.LangchainAgent(

    prompt=qa_prompt,
    model_kwargs={"temperature": 0},
    model=model,
    chat_history=get_session_history,
    tools=[
        book_retrieval_tool,
    ],
    #enable_tracing=True,
)

remote_agent = reasoning_engines.ReasoningEngine.create(
    agent,
    requirements=[
        "google-cloud-aiplatform[langchain,reasoningengine]",
        "langchain-google-firestore"
   ],
)

response = remote_agent.query(
    input="Hi", 
    config={"configurable": {"session_id": "1"}})

print(response["output"])
#print(response)


response = remote_agent.query(
    input="I am looking for a book but cant remember the name. It was about a woman writing a diary in a room", 
    config={"configurable": {"session_id": "1"}})

print(response["output"])




response = remote_agent.query(
    input="please tell me more about the character in that book", 
    config={"configurable": {"session_id": "1"}})

print(response["output"])
print("full response:")
print(response)

frontend-integration.js

const axios = require('axios');
const { GoogleAuth } = require('google-auth-library');

const PROJECT_ID = "sascha-playground-doit";
const LOCATION = "us-central1";
const REASONING_ENGINE_ID = "6135925793873723392";

const auth = new GoogleAuth({
    scopes: ['https://www.googleapis.com/auth/cloud-platform'],
});

async function queryReasoningEngine(question, sessionId) {
    try {
        const client = await auth.getClient();
        const { token } = await client.getAccessToken(); // Ensure correct token extraction
        console.log('Access Token:', token); // Log to confirm it's valid

        const endpoint = `https://${LOCATION}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}/reasoningEngines/${REASONING_ENGINE_ID}:query`;

        const payload = {
            input: {
                input: question,
                config: {
                    configurable: {
                        session_id: sessionId,
                    },
                },
            },
        };

        const response = await axios.post(endpoint, payload, {
            headers: {
                Authorization: `Bearer ${token}`, // Pass the token correctly
                'Content-Type': 'application/json',
            },
        });

        return response.data; // Assuming output is part of the response
    } catch (error) {
        console.error('Error querying reasoning engine:', error.response?.data || error.message);
        throw new Error('Failed to query reasoning engine');
    }
}

module.exports = { queryReasoningEngine };

// Example Usage
//queryReasoningEngine('can you recommend a gothic book?', 'test-session')
//    .then((output) => console.log('Reasoning Engine Output:', output))
//    .catch((err) => console.error(err));

rag.py

import vertexai
from vertexai.preview import rag

vertexai.init(project="sascha-playground-doit", location="us-central1")

#optional
embedding_model_config = rag.EmbeddingModelConfig(
        publisher_model="publishers/google/models/text-embedding-004")


# create a corpus
corpus_name = "books"
corpus = rag.create_corpus(display_name="books", 
                           embedding_model_config=embedding_model_config)

print(corpus)
corpus_name = corpus.name
print(corpus_name)



import_files = rag.import_files(corpus_name=corpus_name, 
                            paths=["gs://doit-llm/books"], 
                            chunk_size=1500)
print(import_files)


# get the corpus
corpus = rag.get_corpus(name=corpus_name)

# this ID will be used later for the agent!!!!!
print(corpus)


# retrieve documents
response = rag.retrieval_query(
        rag_resources=[
            rag.RagResource(
                rag_corpus=corpus_name,
            )
        ],
        text="I am looking for detective novel book anything you can recommend?",
        similarity_top_k=10,  # Optional
        vector_distance_threshold=0.5,  # lower distance means the document as a high similarity to our query
    )
print(response)