agent-starter-pack 0.3.3__py3-none-any.whl → 0.21.0__py3-none-any.whl

agent_starter_pack/deployment_targets/cloud_run/tests/load_test/.results/.placeholder

@@ -0,0 +1,321 @@
+import datetime
+import logging
+import os
+import re
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass
+from typing import Literal
+
+import google.auth
+from google.adk.agents import LlmAgent, LoopAgent, SequentialAgent
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.planners import BuiltInPlanner
+from google.adk.tools import google_search
+from google.genai import types as genai_types
+from pydantic import BaseModel, Field
+
+from google.adk.tools.agent_tool import AgentTool
+from google.adk.agents import BaseAgent
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.events import Event, EventActions
+
+
+_, project_id = google.auth.default()
+os.environ.setdefault("GOOGLE_CLOUD_PROJECT", project_id)
+os.environ.setdefault("GOOGLE_CLOUD_LOCATION", "global")
+os.environ.setdefault("GOOGLE_GENAI_USE_VERTEXAI", "True")
+
+
+# --- Centralized Configuration ---
+@dataclass
+class ResearchConfiguration:
+    critic_model: str = "gemini-2.5-pro"
+    worker_model: str = "gemini-2.5-flash"
+    max_search_iterations: int = 5
+
+
+config = ResearchConfiguration()
+
+
+# --- Structured Output Models ---
+class SearchQuery(BaseModel):
+    search_query: str = Field(
+        description="A highly specific and targeted query for web search."
+    )
+
+
+class Feedback(BaseModel):
+    grade: Literal["pass", "fail"] = Field(
+        description="Evaluation result. 'pass' if the research is sufficient, 'fail' if it needs revision."
+    )
+    comment: str = Field(
+        description="Detailed explanation of the evaluation, highlighting strengths and/or weaknesses of the research."
+    )
+    follow_up_queries: list[SearchQuery] | None = Field(
+        default=None,
+        description="A list of specific, targeted follow-up search queries needed to fix research gaps. This should be null or empty if the grade is 'pass'.",
+    )
+
+# --- Callbacks ---
+def collect_research_sources_callback(callback_context: CallbackContext) -> None:
+    session = callback_context._invocation_context.session
+    url_to_short_id = callback_context.state.get("url_to_short_id", {})
+    sources = callback_context.state.get("sources", {})
+    id_counter = len(url_to_short_id) + 1
+    for event in session.events:
+        if not (event.grounding_metadata and event.grounding_metadata.grounding_chunks):
+            continue
+        chunks_info = {}
+        for idx, chunk in enumerate(event.grounding_metadata.grounding_chunks):
+            if not chunk.web:
+                continue
+            url = chunk.web.uri
+            title = (
+                chunk.web.title
+                if chunk.web.title != chunk.web.domain
+                else chunk.web.domain
+            )
+            if url not in url_to_short_id:
+                short_id = f"src-{id_counter}"
+                url_to_short_id[url] = short_id
+                sources[short_id] = {
+                    "short_id": short_id,
+                    "title": title,
+                    "url": url,
+                    "domain": chunk.web.domain,
+                    "supported_claims": [],
+                }
+                id_counter += 1
+            chunks_info[idx] = url_to_short_id[url]
+        if event.grounding_metadata.grounding_supports:
+            for support in event.grounding_metadata.grounding_supports:
+                confidence_scores = support.confidence_scores or []
+                chunk_indices = support.grounding_chunk_indices or []
+                for i, chunk_idx in enumerate(chunk_indices):
+                    if chunk_idx in chunks_info:
+                        short_id = chunks_info[chunk_idx]
+                        confidence = (
+                            confidence_scores[i] if i < len(confidence_scores) else 0.5
+                        )
+                        text_segment = support.segment.text if support.segment else ""
+                        sources[short_id]["supported_claims"].append(
+                            {
+                                "text_segment": text_segment,
+                                "confidence": confidence,
+                            }
+                        )
+    callback_context.state["url_to_short_id"] = url_to_short_id
+    callback_context.state["sources"] = sources
+
+
+def citation_replacement_callback(callback_context: CallbackContext) -> genai_types.Content:
+    final_report = callback_context.state.get("final_cited_report", "")
+    sources = callback_context.state.get("sources", {})
+    def tag_replacer(match: re.Match) -> str:
+        short_id = match.group(1)
+        if not (source_info := sources.get(short_id)):
+            logging.warning(f"Invalid citation tag found and removed: {match.group(0)}")
+            return ""
+        display_text = source_info.get("title", source_info.get("domain", short_id))
+        return f" [{display_text}]({source_info['url']})"
+    processed_report = re.sub(
+        r'<cite\s+source\s*=\s*["\']?\s*(src-\d+)\s*["\']?\s*/>', tag_replacer, final_report
+    )
+    processed_report = re.sub(r'\s+([.,;:])', r'\1', processed_report)
+    callback_context.state["final_report_with_citations"] = processed_report
+    return genai_types.Content(parts=[genai_types.Part(text=processed_report)])
+
+
+# --- Custom Agent for Loop Control ---
+class EscalationChecker(BaseAgent):
+    """Checks research evaluation and escalates to stop the loop if grade is 'pass'."""
+    def __init__(self, name: str):
+        super().__init__(name=name)
+    async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event, None]:
+        evaluation_result = ctx.session.state.get("research_evaluation")
+        if evaluation_result and evaluation_result.get("grade") == "pass":
+            logging.info(f"[{self.name}] Research evaluation passed. Escalating to stop loop.")
+            yield Event(author=self.name, actions=EventActions(escalate=True))
+        else:
+            logging.info(f"[{self.name}] Research evaluation failed or not found. Loop will continue.")
+            # Yielding an event without content or actions just lets the flow continue.
+            yield Event(author=self.name)
+
+
+# --- AGENT DEFINITIONS ---
+
+plan_generator = LlmAgent(
+    model=config.worker_model,
+    name="plan_generator",
+    description="Generates a 4-5 line action-oriented research plan, using minimal search only for topic clarification.",
+    instruction=f"""
+    You are a research strategist. Your job is to create a high-level RESEARCH PLAN, not a summary.
+    **RULE: Your output MUST be a bulleted list of 4-5 action-oriented research goals or key questions.**
+    - A good goal starts with a verb like "Analyze," "Identify," "Investigate."
+    - A bad output is a statement of fact like "The event was in April 2024."
+    **TOOL USE IS STRICTLY LIMITED:**
+    Your goal is to create a generic, high-quality plan *without searching*.
+    Only use `google_search` if a topic is ambiguous or time-sensitive and you absolutely cannot create a plan without a key piece of identifying information.
+    You are explicitly forbidden from researching the *content* or *themes* of the topic. That is the next agent's job. Your search is only to identify the subject, not to investigate it.
+    Current date: {datetime.datetime.now().strftime("%Y-%m-%d")}
+    """,
+    tools=[google_search],
+)
+
+
+section_planner = LlmAgent(
+    model=config.worker_model,
+    name="section_planner",
+    description="Breaks down the research plan into a structured markdown outline of report sections.",
+    instruction="""
+    You are an expert report architect. Using the research topic and the plan from the 'research_plan' state key, design a logical structure for the final report.
+    Your task is to create a markdown outline with 4-6 distinct sections that cover the topic comprehensively without overlap.
+    You can use any markdown format you prefer, but here's a suggested structure:
+    # Section Name
+    A brief overview of what this section covers
+    Feel free to add subsections or bullet points if needed to better organize the content.
+    Make sure your outline is clear and easy to follow.
+    Do not include a "References" or "Sources" section in your outline. Citations will be handled in-line.
+    """,
+    output_key="report_sections",
+)
+
+
+section_researcher = LlmAgent(
+    model=config.worker_model,
+    name="section_researcher",
+    description="Performs the crucial first pass of web research.",
+    planner=BuiltInPlanner(
+        thinking_config=genai_types.ThinkingConfig(include_thoughts=True)
+    ),
+    instruction="""
+    You are a diligent and exhaustive researcher. Your task is to perform the initial, broad information gathering for a report.
+    You will be provided with a list of sections in the 'report_sections' state key.
+    For each section where 'research' is marked as 'true', generate a comprehensive list of 4-5 targeted search queries to cover the topic from multiple angles.
+    Execute all of these queries using the 'google_search' tool and synthesize the results into a detailed summary for that section.
+    """,
+    tools=[google_search],
+    output_key="section_research_findings",
+    after_agent_callback=collect_research_sources_callback,
+)
+
+research_evaluator = LlmAgent(
+    model=config.critic_model,
+    name="research_evaluator", 
+    description="Critically evaluates research and generates follow-up queries.",
+    instruction=f"""
+    You are a meticulous quality assurance analyst evaluating the research findings in 'section_research_findings'.
+    
+    **CRITICAL RULES:**
+    1. Assume the given research topic is correct. Do not question or try to verify the subject itself.
+    2. Your ONLY job is to assess the quality, depth, and completeness of the research provided *for that topic*.
+    3. Focus on evaluating: Comprehensiveness of coverage, logical flow and organization, use of credible sources, depth of analysis, and clarity of explanations.
+    4. Do NOT fact-check or question the fundamental premise or timeline of the topic.
+    5. If suggesting follow-up queries, they should dive deeper into the existing topic, not question its validity.
+
+    Be very critical about the QUALITY of research. If you find significant gaps in depth or coverage, assign a grade of "fail", 
+    write a detailed comment about what's missing, and generate 5-7 specific follow-up queries to fill those gaps. 
+    If the research thoroughly covers the topic, grade "pass".
+
+    Current date: {datetime.datetime.now().strftime("%Y-%m-%d")}
+    Your response must be a single, raw JSON object validating against the 'Feedback' schema.
+    """,
+    output_schema=Feedback,
+    output_key="research_evaluation",
+)
+
+enhanced_search_executor = LlmAgent(
+    model=config.worker_model,
+    name="enhanced_search_executor",
+    description="Executes follow-up searches and integrates new findings.",
+    planner=BuiltInPlanner(
+        thinking_config=genai_types.ThinkingConfig(include_thoughts=True)
+    ),
+    instruction="""
+    You are a specialist researcher executing a refinement pass.
+    You have been activated because the previous research was graded as 'fail'.
+
+    1.  Review the 'research_evaluation' state key to understand the feedback and required fixes.
+    2.  Execute EVERY query listed in 'follow_up_queries' using the 'google_search' tool.
+    3.  Synthesize the new findings and COMBINE them with the existing information in 'section_research_findings'.
+    4.  Your output MUST be the new, complete, and improved set of research findings.
+    """,
+    tools=[google_search],
+    output_key="section_research_findings",
+    after_agent_callback=collect_research_sources_callback,
+)
+
+report_composer = LlmAgent(
+    model=config.critic_model,
+    name="report_composer_with_citations",
+    include_contents="none",
+    description="Transforms research data and a markdown outline into a final, cited report.",
+    instruction="""
+    Transform the provided data into a polished, professional, and meticulously cited research report.
+
+    ---
+    ### INPUT DATA
+    *   Research Plan: `{research_plan}`
+    *   Research Findings: `{section_research_findings}`
+    *   Citation Sources: `{sources}`
+    *   Report Structure: `{report_sections}`
+
+    ---
+    ### CRITICAL: Citation System
+    To cite a source, you MUST insert a special citation tag directly after the claim it supports.
+
+    **The only correct format is:** `<cite source="src-ID_NUMBER" />`
+
+    ---
+    ### Final Instructions
+    Generate a comprehensive report using ONLY the `<cite source="src-ID_NUMBER" />` tag system for all citations.
+    The final report must strictly follow the structure provided in the **Report Structure** markdown outline.
+    Do not include a "References" or "Sources" section; all citations must be in-line.
+    """,
+    output_key="final_cited_report",
+    after_agent_callback=citation_replacement_callback,
+)
+
+research_pipeline = SequentialAgent(
+    name="research_pipeline",
+    description="Executes a pre-approved research plan. It performs iterative research, evaluation, and composes a final, cited report.",
+    sub_agents=[
+        section_planner,
+        section_researcher,
+        LoopAgent(
+            name="iterative_refinement_loop",
+            max_iterations=config.max_search_iterations,
+            sub_agents=[
+                research_evaluator,
+                EscalationChecker(name="escalation_checker"),
+                enhanced_search_executor,
+            ],
+        ),
+        report_composer,
+    ],
+)
+
+interactive_planner_agent = LlmAgent(
+    name="interactive_planner_agent",
+    model=config.worker_model,
+    description="The primary research assistant. It collaborates with the user to create a research plan, and then executes it upon approval.",
+    instruction=f"""
+    You are a research planning assistant. Your primary function is to convert ANY user request into a research plan.
+
+    **CRITICAL RULE: Never answer a question directly or refuse a request.** Your one and only first step is to use the `plan_generator` tool to propose a research plan for the user's topic.
+    If the user asks a question, you MUST immediately call `plan_generator` to create a plan to answer the question.
+
+    Your workflow is:
+    1.  **Plan:** Use `plan_generator` to create a draft plan and present it to the user.
+    2.  **Refine:** Incorporate user feedback until the plan is approved.
+    3.  **Execute:** Once the user gives EXPLICIT approval (e.g., "looks good, run it"), you MUST delegate the task to the `research_pipeline` agent, passing the approved plan.
+
+    Current date: {datetime.datetime.now().strftime("%Y-%m-%d")}
+    Do not perform any research yourself. Your job is to Plan, Refine, and Delegate.
+    """,
+    sub_agents=[research_pipeline],
+    tools=[AgentTool(plan_generator)],
+    output_key="research_plan",
+)
+
+root_agent = interactive_planner_agent

agent_starter_pack/deployment_targets/cloud_run/tests/load_test/README.md

@@ -0,0 +1,165 @@
+# Robust Load Testing for Generative AI Applications
+
+This directory provides a comprehensive load testing framework for your Generative AI application, leveraging the power of [Locust](http://locust.io), a leading open-source load testing tool.
+{%- if cookiecutter.agent_name == "adk_live" %}
+
+## Local Load Testing
+
+Follow these steps to execute load tests on your local machine:
+
+**1. Start the FastAPI Server:**
+
+Launch the FastAPI server in a separate terminal:
+
+```bash
+uv run uvicorn {{cookiecutter.agent_directory}}.fast_api_app:app --host 0.0.0.0 --port 8000 --reload
+```
+
+**2. (In another tab) Create virtual environment with Locust**
+Using another terminal tab, This is suggested to avoid conflicts with the existing application python environment.
+
+```bash
+python3 -m venv .locust_env && source .locust_env/bin/activate && pip install locust==2.31.1 websockets
+```
+
+**3. Execute the Load Test:**
+Trigger the Locust load test with the following command:
+
+```bash
+locust -f tests/load_test/load_test.py \
+-H http://127.0.0.1:8000 \
+--headless \
+-t 30s -u 2 -r 2 \
+--csv=tests/load_test/.results/results \
+--html=tests/load_test/.results/report.html
+```
+
+This command initiates a 30-second load test, simulating 2 users spawning per second, reaching a maximum of 60 concurrent users.
+
+**Results:**
+
+Comprehensive CSV and HTML reports detailing the load test performance will be generated and saved in the `tests/load_test/.results` directory.
+
+## Remote Load Testing (Targeting Cloud Run)
+
+This framework also supports load testing against remote targets, such as a staging Cloud Run instance. This process is seamlessly integrated into the Continuous Delivery (CD) pipeline.
+
+**Prerequisites:**
+
+- **Dependencies:** Ensure your environment has the same dependencies required for local testing.
+- **Cloud Run Invoker Role:** You'll need the `roles/run.invoker` role to invoke the Cloud Run service.
+
+**Steps:**
+
+**1. Start Cloud Run Proxy:**
+
+Start the proxy in a separate terminal to expose your Cloud Run service on localhost. The proxy automatically handles IAM authentication:
+
+```bash
+gcloud run services proxy YOUR_SERVICE_NAME --port=8080 --region us-central1 --quiet
+```
+
+Replace `YOUR_SERVICE_NAME` with your Cloud Run service name. The `--quiet` flag auto-approves component installation prompts. You can optionally specify `--tag` to target a specific traffic tag.
+
+**2. (In another tab) Create virtual environment with Locust:**
+
+Using another terminal tab:
+
+```bash
+python3 -m venv .locust_env && source .locust_env/bin/activate && pip install locust==2.31.1 websockets
+```
+
+**3. Execute the Load Test:**
+
+Execute load tests against the proxied service. The proxy handles authentication automatically:
+
+```bash
+locust -f tests/load_test/load_test.py \
+-H http://127.0.0.1:8080 \
+--headless \
+-t 30s -u 2 -r 2 \
+--csv=tests/load_test/.results/results \
+--html=tests/load_test/.results/report.html
+```
+{%- else %}
+
+## Local Load Testing
+
+Follow these steps to execute load tests on your local machine:
+
+**1. Start the FastAPI Server:**
+
+Launch the FastAPI server in a separate terminal:
+
+```bash
+uv run uvicorn {{cookiecutter.agent_directory}}.fast_api_app:app --host 0.0.0.0 --port 8000 --reload
+```
+
+**2. (In another tab) Create virtual environment with Locust**
+Using another terminal tab, This is suggested to avoid conflicts with the existing application python environment.
+
+```bash
+python3 -m venv .locust_env && source .locust_env/bin/activate && pip install locust==2.31.1{%- if cookiecutter.is_a2a %} a2a-sdk~=0.3.9{%- endif %}
+```
+
+**3. Execute the Load Test:**
+Trigger the Locust load test with the following command:
+
+```bash
+locust -f tests/load_test/load_test.py \
+-H http://127.0.0.1:8000 \
+--headless \
+-t 30s -u 10 -r 2 \
+--csv=tests/load_test/.results/results \
+--html=tests/load_test/.results/report.html
+```
+
+This command initiates a 30-second load test, simulating 2 users spawning per second, reaching a maximum of 60 concurrent users.
+
+**Results:**
+
+Comprehensive CSV and HTML reports detailing the load test performance will be generated and saved in the `tests/load_test/.results` directory.
+
+## Remote Load Testing (Targeting Cloud Run)
+
+This framework also supports load testing against remote targets, such as a staging Cloud Run instance. This process is seamlessly integrated into the Continuous Delivery (CD) pipeline.
+
+**Prerequisites:**
+
+- **Dependencies:** Ensure your environment has the same dependencies required for local testing.
+- **Cloud Run Invoker Role:** You'll need the `roles/run.invoker` role to invoke the Cloud Run service.
+
+**Steps:**
+
+**1. Obtain Cloud Run Service URL:**
+
+Navigate to the Cloud Run console, select your service, and copy the URL displayed at the top. Set this URL as an environment variable:
+
+```bash
+export RUN_SERVICE_URL=https://your-cloud-run-service-url.run.app
+```
+
+**2. Obtain ID Token:**
+
+Retrieve the ID token required for authentication:
+
+```bash
+export _ID_TOKEN=$(gcloud auth print-identity-token -q)
+```
+
+**3. Execute the Load Test:**
+Create virtual environment with Locust:
+```bash
+python3 -m venv .locust_env && source .locust_env/bin/activate && pip install locust==2.31.1{%- if cookiecutter.is_a2a %} a2a-sdk~=0.3.9{%- endif %}
+```
+
+Execute load tests. The following command executes the same load test parameters as the local test but targets your remote Cloud Run instance.
+```bash
+locust -f tests/load_test/load_test.py \
+-H $RUN_SERVICE_URL \
+--headless \
+-t 30s -u 60 -r 2 \
+--csv=tests/load_test/.results/results \
+--html=tests/load_test/.results/report.html
+```
+{%- endif %}