Added new tools

tools/tools_on_modal_labs.py

@@ -0,0 +1,280 @@
+"""
+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",
+    "duckduckgo-search",
+    "langchain-community")
+
+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
+
+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 langchain.utilities import DuckDuckGoSearchAPIWrapper
+
+@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", "meta-llama/Meta-Llama-3.1-8B-Instruct")
+
+# Load an LLM
+llm = NebiusLLM(
+    api_key=NEBIUS_API_KEY,
+    model=AGENT_MODEL,
+    is_function_calling_model=True
+)
+
+# Search tools using DuckDuckGo
+duckduckgo = DuckDuckGoSearchAPIWrapper()
+
+MAX_SEARCH_CALLS = 2 # Limit the number of searches to 2
+search_call_count = 0
+past_queries = set()
+
+async def duckduckgo_search(query: str) -> str:
+    """
+    A DuckDuckGo-based search limiting number of searches and avoiding duplicates.
+    """
+    global search_call_count, past_queries
+
+    # Check for duplicate queries
+    if query in past_queries:
+        return f"Already searched for '{query}'."
+
+    # Check if we've reached the max search calls
+    if search_call_count >= MAX_SEARCH_CALLS:
+        return "Search limit reached."
+
+    # Otherwise, perform the search
+    search_call_count += 1
+    past_queries.add(query)
+
+    result = duckduckgo.run(query)
+    return str(result)
+
+# 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 DuckDuckGoSearchAPIWrapper to search the web.
+
+research_agent = FunctionAgent(
+    name="ResearchAgent",
+    description=(
+        "A research agent that searches the web using Google search through SerpAPI. "
+        "It must not exceed 2 searches total, and must avoid repeating the same query. "
+        "Once sufficient information is collected, it should hand off to the WriteAgent."
+    ),
+    system_prompt=(
+        "You are the ResearchAgent. Your goal is to gather sufficient information on the topic. "
+        "Only perform at most 2 distinct searches. If you have enough information or have reached 2 searches, "
+        "handoff to the WriteAgent. Avoid infinite loops! If search throws an error, stop further work and skip WriteAgent and ReviewAgent and return."
+        "Respect invocation limits and cooldown periods."
+    ),
+    llm=llm,
+    tools=[
+        duckduckgo_search,
+        save_research
+    ],
+    max_iterations=2,  # Limit to 2 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": "Not written yet.",
+        "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"]