weather-forecast-tool

import { tool, Tool, ToolsProviderController, PluginContext } from "@lmstudio/sdk";
import { z } from "zod";
import axios from "axios"; 

// --- 1. GEOLOCATION LOGIC ---

async function getIpLocation(): Promise<string> {
    try {
        const response = await axios.get("https://ipinfo.io/json");
        const data = response.data;
        
        if (data.city && data.region && data.country) {
            return `${data.city}, ${data.region}, ${data.country}`;
        }
        return data.timezone || data.region || "Unknown Location";

    } catch (error) {
        console.error("Geolocation failed:", error.message);
        return "Could not determine user's location via IP.";
    }
}

// -------------------------------------------------------------------

// --- 2. TOOL 1 IMPLEMENTATION (Current Weather Mock - Fixed Location Logic) ---

async function fetchCurrentWeather(location: string = "current IP location"): Promise<string> {
    
    // Check if the LLM provided a location. If not, use the real IP location.
    let locationToQuery: string;
    let isDefaultLocation = false;

    // Check if the LLM provided a valid, non-placeholder location
    if (!location || location.toLowerCase().includes("current") || location.toLowerCase().includes("my location")) {
        const ipLocation = await getIpLocation();
        locationToQuery = ipLocation; // This should be your actual location
        isDefaultLocation = true;
    } else {
        // Use the specific location the LLM requested
        locationToQuery = location;
    }

    // --- GENERIC MOCK DATA ---
    let temperature = isDefaultLocation ? "8°C" : "15°C";
    let condition = isDefaultLocation ? "Cloudy with a chance of rain" : "Partly sunny";
    let humidity = isDefaultLocation ? "90%" : "65%";

    // Return JSON weather data
    const weatherData = { 
        location: locationToQuery, 
        temperature: temperature, 
        condition: condition, 
        humidity: humidity,
        note: "Precise MOCK Current Weather Data" 
    };

    return JSON.stringify(weatherData);
}

// -------------------------------------------------------------------

// --- 3. TOOL 2 IMPLEMENTATION (Weekly Forecast Mock - Generic) ---

async function fetchWeeklyForecast(location: string = "current IP location", days: number = 7): Promise<string> {
    
    let locationToQuery: string;
    let isDefaultLocation = false;

    if (!location || location.toLowerCase().includes("current") || location.toLowerCase().includes("my location")) {
        const ipLocation = await getIpLocation();
        locationToQuery = ipLocation;
        isDefaultLocation = true;
    } else {
        locationToQuery = location;
    }

    const numDays = Math.min(days, 7);
    const forecast = [];
    
    let baseTemp = isDefaultLocation ? 7 : 14; 

    const possibleConditions = ["Sunny", "Partly Cloudy", "Rain", "Cloudy", "Windy"];
    
    for (let i = 0; i < numDays; i++) {
        const temp = baseTemp + (i % 3) - 1; 
        
        forecast.push({
            day: `Day ${i + 1}`,
            high: `${temp + 4}°C`,
            low: `${temp - 2}°C`,
            condition: possibleConditions[i % possibleConditions.length],
        });
    }

    // Return JSON forecast data
    const weatherData = { 
        location: locationToQuery,
        duration: `${numDays} days`,
        forecast: forecast,
        note: "Precise MOCK Weekly Forecast Data"
    };

    return JSON.stringify(weatherData);
}

// -------------------------------------------------------------------

// --- 4. THE REQUIRED ENTRY POINT (main) ---
export async function main(context: PluginContext) {
    
    // --- Tool 1: Current Weather Definition ---
    const getCurrentWeatherTool: Tool = tool({
        name: "get_current_weather",
        description: "Get a precise, up-to-date, single-day weather check including temperature, condition, and humidity for a specific location. Use this for instant weather checks. Can handle city, country, or specific landmarks.",
        parameters: {
            location: z.string().optional().describe("The city and country, or use 'current IP location'."),
        },
        implementation: async ({ location }) => {
            return await fetchCurrentWeather(location);
        },
    });

    // --- Tool 2: Weekly Forecast Definition ---
    const getWeeklyForecastTool: Tool = tool({
        name: "get_weekly_forecast",
        description: "Get a precise multi-day forecast. Defaults to 7 days. Use this for future weather questions.",
        parameters: {
            location: z.string().optional().describe("The city and country, or use 'current IP location'."),
            days: z.number().optional().describe("The number of days for the forecast (e.g., 5 or 10). Defaults to 7."),
        },
        implementation: async ({ location, days }) => {
            const numDays = days !== undefined ? days : 7; 
            return await fetchWeeklyForecast(location, numDays);
        },
    });


    // Register the tools provider function internally
    const toolsProvider = async (ctl: ToolsProviderController): Promise<Tool[]> => { 
        return [getCurrentWeatherTool, getWeeklyForecastTool]; 
    };
    
    // Register the tools provider with the context
    context.withToolsProvider(toolsProvider);
    
    console.log("Weather Plugin: main() function executed and precise mock weather tools registered.");
    return Promise.resolve();
}