tools/tools_on_modal_labs.py

"""
This module provides tools for searching and retrieving context from a knowledge base,
and for conducting a research workflow that includes searching, writing, and reviewing reports.
The tools are designed to be used with Modal Labs for scalable and efficient processing.
The technology stack includes FastAPI for the API interface, GroundX for knowledge base search,
LlamaIndex for LLM workflows, Nebius for LLM, and Modal Labs for tool execution.
"""

import os
import asyncio

import modal
from pydantic import BaseModel

image = modal.Image.debian_slim().pip_install(
    "fastapi[standard]",
    "groundx",
    "llama-index",
    "llama-index-llms-nebius",
    "tavily-python")

app = modal.App(name="hackathon-mcp-tools", image=image)

class QueryInput(BaseModel):
    query: str

@app.function(secrets=[
        modal.Secret.from_name("hackathon-secret", required_keys=["GROUNDX_API_KEY"])
    ])
@modal.fastapi_endpoint(docs=True, method="POST")
def search_rag_context(queryInput: QueryInput) -> str:
    """
    Searches and retrieves relevant context from a knowledge base,
    based on the user's query.
    Args:
        query: The search query supplied by the user.
    Returns:
        str: Relevant text content that can be used by the LLM to answer the query.
    """

    result = search_groundx_for_rag_context(queryInput.query)

    print("\n\n=============================")
    print(f"RAG Search Result: {result}")
    print("=============================\n")

    return result

def search_groundx_for_rag_context(query) -> str:
    from groundx import GroundX

    client = GroundX(api_key=os.getenv("GROUNDX_API_KEY") or '')
    response = client.search.content(
        id=os.getenv("GROUNDX_BUCKET_ID"),
        query=query,
        n=10,
    )

    return response.search.text or "No relevant context found"

from llama_index.llms.nebius import NebiusLLM

# llama-index workflow classes
from llama_index.core.workflow import Context
from llama_index.core.agent.workflow import (
    FunctionAgent,
    AgentWorkflow,
    AgentOutput,
    ToolCall,
    ToolCallResult,
)

from tavily import AsyncTavilyClient

@app.function(secrets=[
        modal.Secret.from_name("hackathon-secret", required_keys=["NEBIUS_API_KEY", "AGENT_MODEL"])
    ])
@modal.fastapi_endpoint(docs=True, method="POST")
def run_research_workflow(queryInput: QueryInput) -> str:
    handler = asyncio.run(execute_research_workflow(queryInput.query))
    result = asyncio.run(final_report(handler))
    return result

NEBIUS_API_KEY = os.getenv("NEBIUS_API_KEY")
AGENT_MODEL = os.getenv("AGENT_MODEL", "Qwen/Qwen2.5-Coder-32B-Instruct")
print(f"Using LLM model: {AGENT_MODEL}")

# Load an LLM
llm = NebiusLLM(
    api_key=NEBIUS_API_KEY,
    model=AGENT_MODEL,
    is_function_calling_model=True
)

TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")

# Search tools using Tavily
async def search_web(query: str) -> str:
    """Useful for using the web to answer questions."""
    client = AsyncTavilyClient(api_key=TAVILY_API_KEY)
    return str(await client.search(query))

# Research tools
async def save_research(ctx: Context, notes: str, notes_title: str) -> str:
    """
    Store research notes under a given title in the shared context.
    """

    current_state = await ctx.get("state")
    if "research_notes" not in current_state:
        current_state["research_notes"] = {}
    current_state["research_notes"][notes_title] = notes
    await ctx.set("state", current_state)
    return "Notes saved."

# Report tools
async def write_report(ctx: Context, report_content: str) -> str:
    """
    Write a report in markdown, storing it in the shared context.
    """

    current_state = await ctx.get("state")
    current_state["report_content"] = report_content
    await ctx.set("state", current_state)
    return "Report written."

# Review tools
async def review_report(ctx: Context, review: str) -> str:
    """
    Review the report and store feedback in the shared context.
    """

    current_state = await ctx.get("state")
    current_state["review"] = review
    await ctx.set("state", current_state)
    return "Report reviewed."


# We have three agents with distinct responsibilities:
# - The ResearchAgent is responsible for gathering information from the web.
# - The WriteAgent is responsible for writing the report.
# - The ReviewAgent is responsible for reviewing the report.

# The ResearchAgent uses the Tavily Client to search the web.

research_agent = FunctionAgent(
    name="ResearchAgent",
    description=(
        "Useful for searching the web for information on a given topic and recording notes on the topic."
        "It must not exceed 3 searches total and must avoid repeating the same query. "
        "Once sufficient information is collected, you should hand off to the WriteAgent."
    ),
    system_prompt=(
        "You are the ResearchAgent. Your goal is to search the web for information on a given topic and record notes on the topic."
        "Only perform at most 3 distinct searches. If you have enough information or have reached 3 searches, "
        "handoff to the WriteAgent. Avoid infinite loops! If the search throws an error, stop further work and skip WriteAgent and ReviewAgent and return."
        "You should have at least some notes on a topic before handing off control to the WriteAgent."
        "Respect invocation limits and cooldown periods."
    ),
    llm=llm,
    tools=[
        search_web,
        save_research
    ],
    max_iterations=3,  # Limit to 3 iterations to prevent infinite loops
    cooldown=5,  # Cooldown to prevent rapid re-querying
    can_handoff_to=["WriteAgent"]
)

write_agent = FunctionAgent(
    name="WriteAgent",
    description=(
        "Writes a markdown report based on the research notes. "
        "Then hands off to the ReviewAgent for feedback."
    ),
    system_prompt=(
        "You are the WriteAgent. Draft a structured markdown report based on the notes. "
        "If there is no report content or research notes, stop further work and skip ReviewAgent."
        "Do not attempt more than one write attempt. "
        "After writing, hand off to the ReviewAgent."
        "Respect invocation limits and cooldown periods."
    ),
    llm=llm,
    tools=[write_report],
    max_iterations=2,  # Limit to 2 iterations to prevent infinite loops
    cooldown=5,  # Cooldown to prevent rapid re-querying
    can_handoff_to=["ReviewAgent", "ResearchAgent"]
)

review_agent = FunctionAgent(
    name="ReviewAgent",
    description=(
        "Reviews the final report for correctness. Approves or requests changes."
    ),
    system_prompt=(
        "You are the ReviewAgent. If there is no research notes or report content, skip this step and return."
        "Do not attempt more than one review attempt. "
        "Read the report, provide feedback, and either approve "
        "or request revisions. If revisions are needed, handoff to WriteAgent."
        "Respect invocation limits and cooldown periods."
    ),
    llm=llm,
    tools=[review_report],
    max_iterations=2,  # Limit to 2 iterations to prevent infinite loops
    cooldown=5,  # Cooldown to prevent rapid re-querying
    can_handoff_to=["WriteAgent"]
)

agent_workflow = AgentWorkflow(
    agents=[research_agent, write_agent, review_agent],
    root_agent=research_agent.name,  # Start with the ResearchAgent
    initial_state={
        "research_notes": {},
        "report_content": "No report has been generated after the search.",
        "review": "Review required.",
    },
)

async def execute_research_workflow(query: str):
    handler = agent_workflow.run(
        user_msg=(
            query
        )
    )

    current_agent = None

    async for event in handler.stream_events():
        if hasattr(event, "current_agent_name") and event.current_agent_name != current_agent:
            current_agent = event.current_agent_name
            print(f"\n{'='*50}")
            print(f"🤖 Agent: {current_agent}")
            print(f"{'='*50}\n")

        # Print outputs or tool calls
        if isinstance(event, AgentOutput):
            if event.response.content:
                print("📤 Output:", event.response.content)
            if event.tool_calls:
                print("🛠️  Planning to use tools:", [call.tool_name for call in event.tool_calls])

        elif isinstance(event, ToolCall):
            print(f"🔨 Calling Tool: {event.tool_name}")
            print(f"  With arguments: {event.tool_kwargs}")

        elif isinstance(event, ToolCallResult):
            print(f"🔧 Tool Result ({event.tool_name}):")
            print(f"  Arguments: {event.tool_kwargs}")
            print(f"  Output: {event.tool_output}")

    return handler

async def final_report(handler) -> str:
    """Retrieve the final report from the context."""
    final_state = await handler.ctx.get("state")
    print("\n\n=============================")
    print("FINAL REPORT:\n")
    print(final_state["report_content"])
    print("=============================\n")

    return final_state["report_content"]