agent-starter-pack 0.0.1b0__py3-none-any.whl

agents/agentic_rag_vertexai_search/app/agent.py

@@ -0,0 +1,145 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# mypy: disable-error-code="arg-type"
+import os
+
+import google
+import vertexai
+from langchain_core.documents import Document
+from langchain_core.messages import BaseMessage
+from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import tool
+from langchain_google_vertexai import ChatVertexAI, VertexAIEmbeddings
+from langgraph.graph import END, MessagesState, StateGraph
+from langgraph.prebuilt import ToolNode
+
+from app.retrievers import get_compressor, get_retriever
+from app.templates import format_docs, inspect_conversation_template, rag_template
+
+EMBEDDING_MODEL = "text-embedding-005"
+LOCATION = "us-central1"
+LLM = "gemini-2.0-flash-001"
+EMBEDDING_COLUMN = "embedding"
+TOP_K = 5
+
+data_store_region = os.getenv("DATA_STORE_REGION", "us")
+data_store_id = os.getenv("DATA_STORE_ID", "sample-datastore")
+
+# Initialize Google Cloud and Vertex AI
+credentials, project_id = google.auth.default()
+vertexai.init(project=project_id, location=LOCATION)
+
+embedding = VertexAIEmbeddings(
+    project=project_id, location=LOCATION, model_name=EMBEDDING_MODEL
+)
+
+retriever = get_retriever(
+    project_id=project_id,
+    data_store_id=data_store_id,
+    data_store_region=data_store_region,
+    embedding=embedding,
+    embedding_column=EMBEDDING_COLUMN,
+    max_documents=10,
+)
+compressor = get_compressor(
+    project_id=project_id,
+)
+
+
+@tool(response_format="content_and_artifact")
+def retrieve_docs(query: str) -> tuple[str, list[Document]]:
+    """
+    Useful for retrieving relevant documents based on a query.
+    Use this when you need additional information to answer a question.
+
+    Args:
+        query (str): The user's question or search query.
+
+    Returns:
+        List[Document]: A list of the top-ranked Document objects, limited to TOP_K (5) results.
+    """
+    # Use the retriever to fetch relevant documents based on the query
+    retrieved_docs = retriever.invoke(query)
+    # Re-rank docs with Vertex AI Rank for better relevance
+    ranked_docs = compressor.compress_documents(documents=retrieved_docs, query=query)
+    # Format ranked documents into a consistent structure for LLM consumption
+    formatted_docs = format_docs.format(docs=ranked_docs)
+    return (formatted_docs, ranked_docs)
+
+
+@tool
+def should_continue() -> None:
+    """
+    Use this tool if you determine that you have enough context to respond to the questions of the user.
+    """
+    return None
+
+
+tools = [retrieve_docs, should_continue]
+
+llm = ChatVertexAI(model=LLM, temperature=0, max_tokens=1024, streaming=True)
+
+# Set up conversation inspector
+inspect_conversation = inspect_conversation_template | llm.bind_tools(
+    tools, tool_choice="any"
+)
+
+# Set up response chain
+response_chain = rag_template | llm
+
+
+def inspect_conversation_node(
+    state: MessagesState, config: RunnableConfig
+) -> dict[str, BaseMessage]:
+    """Inspects the conversation state and returns the next message using the conversation inspector."""
+    response = inspect_conversation.invoke(state, config)
+    return {"messages": response}
+
+
+def generate_node(
+    state: MessagesState, config: RunnableConfig
+) -> dict[str, BaseMessage]:
+    """Generates a response using the RAG template and returns it as a message."""
+    response = response_chain.invoke(state, config)
+    return {"messages": response}
+
+
+# Flow:
+# 1. Start with agent node that inspects conversation using inspect_conversation_node
+# 2. Agent node connects to tools node which can either:
+#    - Retrieve relevant docs using retrieve_docs tool
+#    - End tool usage with should_continue tool
+# 3. Tools node connects to generate node which produces final response
+# 4. Generate node connects to END to complete the workflow
+
+workflow = StateGraph(MessagesState)
+workflow.add_node("agent", inspect_conversation_node)
+workflow.add_node("generate", generate_node)
+workflow.set_entry_point("agent")
+
+workflow.add_node(
+    "tools",
+    ToolNode(
+        tools=tools,
+        # With False, tool errors won't be caught by LangGraph
+        handle_tool_errors=False,
+    ),
+)
+workflow.add_edge("agent", "tools")
+workflow.add_edge("tools", "generate")
+
+workflow.add_edge("generate", END)
+
+agent = workflow.compile()

agents/agentic_rag_vertexai_search/app/retrievers.py

@@ -0,0 +1,79 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from unittest.mock import MagicMock
+
+from langchain_google_community import VertexAISearchRetriever
+from langchain_google_community.vertex_rank import VertexAIRank
+from langchain_google_vertexai import VertexAIEmbeddings
+
+
+def get_retriever(
+    project_id: str,
+    data_store_id: str,
+    data_store_region: str,
+    embedding: VertexAIEmbeddings,
+    embedding_column: str = "embedding",
+    max_documents: int = 10,
+    custom_embedding_ratio: float = 0.5,
+) -> VertexAISearchRetriever:
+    """
+    Creates and returns an instance of the retriever service.
+
+    Uses mock service if the INTEGRATION_TEST environment variable is set to "TRUE",
+    otherwise initializes real Vertex AI retriever.
+    """
+    if os.getenv("INTEGRATION_TEST") == "TRUE":
+        retriever = MagicMock()
+        retriever.invoke = lambda x: []
+        return retriever
+
+    return VertexAISearchRetriever(
+        project_id=project_id,
+        data_store_id=data_store_id,
+        location_id=data_store_region,
+        engine_data_type=1,
+        # The following parameters are used when you want to search
+        # using custom embeddings in Agent Builder.
+        # The ratio is set to 0.5 by default to use a mix of custom
+        # embeddings but you can adapt the ratio as you need.
+        custom_embedding_ratio=custom_embedding_ratio,
+        custom_embedding=embedding,
+        custom_embedding_field_path=embedding_column,
+        # Extracting 20 documents before re-rank.
+        max_documents=max_documents,
+        beta=True,
+    )
+
+
+def get_compressor(project_id: str, top_n: int = 5) -> VertexAIRank:
+    """
+    Creates and returns an instance of the compressor service.
+
+    Uses mock service if the INTEGRATION_TEST environment variable is set to "TRUE",
+    otherwise initializes real Vertex AI compressor.
+    """
+    if os.getenv("INTEGRATION_TEST") == "TRUE":
+        compressor = MagicMock()
+        compressor.compress_documents = lambda documents, query: []
+        return compressor
+
+    return VertexAIRank(
+        project_id=project_id,
+        location_id="global",
+        ranking_config="default_ranking_config",
+        title_field="id",
+        top_n=top_n,
+    )

agents/agentic_rag_vertexai_search/app/templates.py

@@ -0,0 +1,53 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from langchain_core.prompts import (
+    ChatPromptTemplate,
+    MessagesPlaceholder,
+    PromptTemplate,
+)
+
+format_docs = PromptTemplate.from_template(
+    """## Context provided:
+{% for doc in docs%}
+<Document {{ loop.index0 }}>
+{{ doc.page_content | safe }}
+</Document {{ loop.index0 }}>
+{% endfor %}
+""",
+    template_format="jinja2",
+)
+
+inspect_conversation_template = ChatPromptTemplate.from_messages(
+    [
+        (
+            "system",
+            """You are an AI assistant tasked with analyzing the conversation"""
+            """ and determining the best course of action.""",
+        ),
+        MessagesPlaceholder(variable_name="messages"),
+    ]
+)
+
+rag_template = ChatPromptTemplate.from_messages(
+    [
+        (
+            "system",
+            """You are an AI assistant for question-answering tasks."""
+            """ Answer to the best of your ability using the context provided."""
+            """ Leverage the Tools you are provided to answer questions.""",
+        ),
+        MessagesPlaceholder(variable_name="messages"),
+    ]
+)