langchain-tools.py

"""
LangChain Custom Tools with Frostbyte APIs
============================================
Create LangChain tools backed by real APIs: geolocation, crypto,
DNS, scraping, screenshots, and code execution.

Requirements:
    pip install langchain langchain-openai requests

Usage:
    export FROSTBYTE_API_KEY="your-key"   # Get free: https://ozorown.github.io
    export OPENAI_API_KEY="sk-..."
    python langchain-tools.py
"""

import json
import os
import requests
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate

FROSTBYTE_BASE = "https://api-catalog-three.vercel.app/v1"
API_KEY = os.environ.get("FROSTBYTE_API_KEY", "")
HEADERS = {"Authorization": f"Bearer {API_KEY}"}


# --- LangChain Tool Definitions ---

@tool
def geo_lookup(ip: str) -> str:
    """Look up geolocation data for an IP address. Returns country, city, ISP, coordinates."""
    r = requests.get(f"{FROSTBYTE_BASE}/agent-geo/geo/{ip}", headers=HEADERS)
    return json.dumps(r.json(), indent=2)


@tool
def crypto_price(symbol: str) -> str:
    """Get current cryptocurrency price, 24h change, and market cap. Use coin name like 'bitcoin', 'ethereum', 'solana'."""
    r = requests.get(f"{FROSTBYTE_BASE}/crypto-feeds/prices/{symbol}", headers=HEADERS)
    return json.dumps(r.json(), indent=2)


@tool
def dns_lookup(domain: str) -> str:
    """Resolve DNS records (A, AAAA, MX, NS, TXT) for a domain name."""
    r = requests.get(f"{FROSTBYTE_BASE}/agent-dns/api/resolve/{domain}", headers=HEADERS)
    return json.dumps(r.json(), indent=2)


@tool
def scrape_url(url: str) -> str:
    """Scrape a webpage and extract its text content, title, and links."""
    r = requests.post(f"{FROSTBYTE_BASE}/agent-scraper/api/scrape",
                      headers={**HEADERS, "Content-Type": "application/json"},
                      json={"url": url})
    data = r.json()
    # Return a concise summary to avoid token overflow
    return json.dumps({
        "title": data.get("title", ""),
        "text": data.get("text", "")[:2000],
        "links_count": len(data.get("links", []))
    }, indent=2)


@tool
def run_code(language: str, code: str) -> str:
    """Execute code in a sandboxed environment. Supported languages: python3, javascript, bash, ruby, go, rust."""
    r = requests.post(f"{FROSTBYTE_BASE}/agent-coderunner/api/execute",
                      headers={**HEADERS, "Content-Type": "application/json"},
                      json={"language": language, "code": code})
    return json.dumps(r.json(), indent=2)


@tool
def take_screenshot(url: str) -> str:
    """Take a screenshot of a webpage. Returns the image as base64-encoded PNG."""
    r = requests.post(f"{FROSTBYTE_BASE}/agent-screenshot/api/screenshot",
                      headers={**HEADERS, "Content-Type": "application/json"},
                      json={"url": url})
    if r.headers.get("content-type", "").startswith("image/"):
        return f"Screenshot captured: {len(r.content)} bytes PNG image"
    return r.text[:500]


# --- Agent Setup ---

def create_agent():
    """Create a LangChain agent with Frostbyte tools."""
    tools = [geo_lookup, crypto_price, dns_lookup, scrape_url, run_code, take_screenshot]

    llm = ChatOpenAI(model="gpt-4o", temperature=0)

    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a research assistant with real-time tools for IP geolocation, "
                   "crypto prices, DNS lookups, web scraping, code execution, and screenshots. "
                   "Use tools to gather accurate, current information."),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ])

    agent = create_tool_calling_agent(llm, tools, prompt)
    return AgentExecutor(agent=agent, tools=tools, verbose=True)


if __name__ == "__main__":
    if not API_KEY:
        print("Get a free API key: https://ozorown.github.io")
        print("Then: export FROSTBYTE_API_KEY='your-key'")
        exit(1)

    agent = create_agent()

    # Example: Multi-tool research task
    result = agent.invoke({
        "input": "Research the website github.com: look up its IP geolocation, "
                 "check its DNS records, and tell me what country it's hosted in."
    })
    print(f"\nFinal Answer: {result['output']}")