app.py

"""
LlamaIndex Agent for HuggingFace AI Agents Course - Unit 2
Simple ReAct agent using LlamaIndex FunctionTool.
"""

import os
import re
import gradio as gr
from llama_index.core.tools import FunctionTool

# ============================================
# TOOLS (using LlamaIndex FunctionTool)
# ============================================

def calculate(operation: str) -> str:
    """Performs basic math calculations. Args: operation - A math expression like '2 + 2'"""
    try:
        allowed = set('0123456789+-*/(). ')
        if all(c in allowed for c in operation):
            return f"Result: {eval(operation)}"
        return "Invalid operation"
    except Exception as e:
        return f"Error: {e}"

def get_current_time() -> str:
    """Gets the current date and time in UTC."""
    from datetime import datetime
    return f"Current time (UTC): {datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')}"

def get_weather(city: str) -> str:
    """Gets weather for a city (mock data)."""
    import random
    conditions = ["sunny", "cloudy", "rainy", "windy"]
    return f"Weather in {city}: {random.randint(15,30)}°C, {random.choice(conditions)}"

def count_words(text: str) -> str:
    """Counts words in text."""
    return f"Word count: {len(text.split())}"

# Create LlamaIndex FunctionTools
calc_tool = FunctionTool.from_defaults(fn=calculate, name="calculate")
time_tool = FunctionTool.from_defaults(fn=get_current_time, name="get_time")
weather_tool = FunctionTool.from_defaults(fn=get_weather, name="get_weather")
words_tool = FunctionTool.from_defaults(fn=count_words, name="count_words")

TOOLS = {
    "calculate": calculate,
    "get_time": get_current_time,
    "get_weather": get_weather,
    "count_words": count_words,
}

# ============================================
# AGENT (Simple ReAct with direct API call)
# ============================================

def call_hf_api(prompt: str) -> str:
    """Call HuggingFace Inference API."""
    import requests
    
    token = os.environ.get("HF_TOKEN")
    
    # Use the OpenAI-compatible endpoint
    response = requests.post(
        "https://router.huggingface.co/novita/v3/openai/chat/completions",
        headers={
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        },
        json={
            "model": "meta-llama/llama-3.1-8b-instruct",
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": 500,
            "temperature": 0.7
        },
        timeout=60
    )
    
    if response.status_code == 200:
        return response.json()["choices"][0]["message"]["content"]
    else:
        raise Exception(f"API Error {response.status_code}: {response.text}")


def run_agent(query: str) -> str:
    """Run the ReAct agent loop."""
    
    tools_desc = """Available tools:
- calculate: Do math (e.g., "2+2", "10*5")
- get_time: Get current UTC time
- get_weather: Get weather for a city
- count_words: Count words in text"""

    prompt = f"""{tools_desc}

To use a tool: Action: <name> | Input: <value>
For final answer: Answer: <response>

Query: {query}
Response:"""

    for _ in range(3):
        try:
            response = call_hf_api(prompt)
            
            # Check for final answer
            if "Answer:" in response:
                return response.split("Answer:")[-1].strip()
            
            # Check for tool use
            action = re.search(r"Action:\s*(\w+)", response)
            inp = re.search(r"Input:\s*(.+?)(?:\n|$)", response)
            
            if action:
                tool_name = action.group(1).lower()
                tool_input = inp.group(1).strip() if inp else ""
                
                # Find and call tool
                for name, func in TOOLS.items():
                    if tool_name in name.lower():
                        if name == "get_time":
                            result = func()
                        else:
                            result = func(tool_input)
                        
                        prompt += f"\n{response}\nObservation: {result}\nResponse:"
                        break
                else:
                    return response
            else:
                return response
                
        except Exception as e:
            return f"Error: {str(e)}"
    
    return "Could not complete request."


# ============================================
# GRADIO UI
# ============================================

def chat(message, history):
    if not message.strip():
        return "Please ask something!"
    return run_agent(message)

with gr.Blocks(title="LlamaIndex Agent") as demo:
    gr.Markdown("""
# 🦙 LlamaIndex ReAct Agent - Unit 2

Built with **LlamaIndex FunctionTool** for the HuggingFace AI Agents Course.

**Tools:** Calculator | Time | Weather | Word Count
---
""")
    
    gr.ChatInterface(
        fn=chat,
        examples=[
            "What is 25 * 4 + 10?",
            "What time is it?",
            "Weather in Tokyo?",
            "Count words: hello world test",
        ],
    )
    
    gr.Markdown("---\n*Unit 2: LlamaIndex Framework*")

if __name__ == "__main__":
    demo.launch()