src/analyzer/chat/run_chat_llm.py

#!/usr/bin/env python3
"""
run_chat_llm.py — Natural language chat with LLM-driven function calling

This version uses the LLM's native function calling to handle ALL queries,
making it much better at understanding complex natural language requests.

Usage:
    python -m src.analyzer.chat.run_chat_llm
    python -m src.analyzer.chat.run_chat_llm --verbose
"""
from __future__ import annotations

import argparse
import logging
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

from ..config import load_config
from ..data_loader import load_current_grants, load_past_winners
from ..llm_client import LLMClient
from ..search.hybrid_index import load_index
from ..utils.errors import (
    GrantAnalyzerError,
    ValidationError,
    DataLoadError,
    SearchError,
    LLMError,
    ConfigError
)
from .chat_tools import ChatTools
from .tool_schemas import openai_tools, detect_extended_features


def _startup_diagnostics(cfg: Dict, llm_ok: bool, idx_ok: bool) -> str:
    """Generate startup diagnostics display."""
    provider = cfg.get("llm_provider", "unknown")
    model = cfg.get("llm_model", "unknown")
    api_key = cfg.get("openai_api_key") or cfg.get("anthropic_api_key")

    diag = [
        "=== Grant Analyst Chat — LLM Function Calling Mode ===",
        f"Provider:   {provider}",
        f"Model:      {model}",
        f"LLM Ready:  {llm_ok}",
        f"Index OK:   {idx_ok}",
        f"API Key:    {'✓' if api_key else '✗'}",
        ""
    ]
    return "\n".join(diag)


def _dispatch_tool_call(tools: ChatTools, tool_name: str, tool_args: Dict[str, Any]) -> Any:
    """
    Dispatch a tool call to the appropriate ChatTools method.

    Args:
        tools: ChatTools instance
        tool_name: Name of the tool to call
        tool_args: Arguments for the tool

    Returns:
        Tool result (dict or string)
    """
    try:
        if tool_name == "list_grants":
            return tools.list_grants(
                keyword=tool_args.get("keyword"),
                max_award=tool_args.get("max_award"),
                audience=tool_args.get("audience"),
                status=tool_args.get("status"),  # NEW: Add status filter
                limit=tool_args.get("limit")  # FIXED: Don't default to 5, pass None for all
            )

        elif tool_name == "get_grant":
            return tools.get_grant(tool_args["grant_id"])

        elif tool_name == "summarize_grant":
            return tools.summarize_grant(
                tool_args["grant_id"],
            )

        elif tool_name == "compare_grants":
            return tools.compare_grants(
                tool_args["grant_id_a"],
                tool_args["grant_id_b"]
            )

        elif tool_name == "deadlines_overview":
            return tools.deadlines_overview(tool_args.get("n", 5))

        elif tool_name == "analyze_company_for_grants":
            return tools.analyze_company_for_grants(
                tool_args["company_url"],
                limit=tool_args.get("limit", 3)
            )

        elif tool_name == "search_grants":
            # This would need to be implemented in ChatTools or InsightTools
            return {
                "results": tools.list_grants(
                    keyword=tool_args.get("query"),
                    status=tool_args.get("status"),  # NEW: Add status filter
                    limit=tool_args.get("limit")  # FIXED: Don't default to 10, pass None for all
                ),
                "query": tool_args.get("query")
            }

        elif tool_name == "search_past_winners":
            return {
                "results": tools.search_past_winners(
                    keyword=tool_args.get("keyword"),
                    competition=tool_args.get("competition"),
                    limit=tool_args.get("limit")
                ),
                "query": tool_args.get("keyword") or tool_args.get("competition") or "all"
            }

        else:
            return {"error": f"Unknown tool: {tool_name}"}

    except Exception as e:
        logging.error(f"Tool {tool_name} failed: {e}", exc_info=True)
        return {"error": str(e)}


def main(argv: Optional[List[str]] = None):
    """Main chat loop with LLM function calling."""

    # Parse arguments
    ap = argparse.ArgumentParser(description="Grant Analyst Chat with LLM function calling")
    ap.add_argument("--verbose", action="store_true", help="Show startup diagnostics")
    ap.add_argument("--limit", type=int, help="Limit number of grants to load")
    ap.add_argument("--extended-tools", action="store_true", help="Enable extended tools")
    args = ap.parse_args(argv)

    # Setup logging
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        handlers=[
            logging.FileHandler("_out/logs/chat_llm.log"),
            logging.StreamHandler(sys.stderr)
        ]
    )

    # Load configuration
    try:
        cfg = load_config()
    except ConfigError as e:
        print(f"⚠️  Configuration error: {e}")
        sys.exit(1)

    # Load data
    logging.info("Loading data...")
    try:
        current = load_current_grants(Path("data/snapshots"), limit=args.limit)
        past = load_past_winners(
            history_xlsx=Path("data/IUK-141025-InnovateUKFundedProjects-FY2015-16topresent.xlsx")
        )
        logging.info(f"Loaded {len(current)} current grants; {len(past)} past winners")
    except DataLoadError as e:
        print(f"⚠️  Data loading error: {e}")
        sys.exit(1)

    # Initialize LLM client
    try:
        llm_client = LLMClient(cfg)
        llm_ok = llm_client.is_ready()
        if not llm_ok:
            print("⚠️  LLM client not ready. Check your API key configuration.")
            sys.exit(1)
    except Exception as e:
        print(f"⚠️  LLM initialization failed: {e}")
        sys.exit(1)

    # Load search index
    idx_ok = False
    try:
        idx_path = Path("data/index/hybrid_index.pkl")
        if idx_path.exists():
            _ = load_index()
            idx_ok = True
    except Exception as e:
        logging.warning(f"Could not load search index: {e}")

    # Initialize ChatTools
    tools = ChatTools(current, past)

    # Tool registration
    extended_mode = args.extended_tools or detect_extended_features()
    available_tools = openai_tools(extended=extended_mode)
    logging.info(f"Registered {len(available_tools)} tools (extended: {extended_mode})")

    # Startup diagnostics
    if args.verbose:
        print(_startup_diagnostics(cfg, llm_ok, idx_ok))

    print("\n💬 Grant Analyst Chat (LLM Function Calling Mode)")
    print("   Ask me anything in natural language! Type 'exit' to quit.\n")

    # Conversation history
    messages = [
        {
            "role": "system",
            "content": (
                "You are a helpful grant analyst assistant for Innovate UK grants. "
                "Use the available tools to answer user questions about grants, funding, "
                "deadlines, eligibility, and comparisons. Always use tools when you need "
                "to look up grant data - don't make up information. When listing grants, "
                "format them clearly with bullets. When comparing, use tables. Be concise "
                "but informative. Current date: 2025-10-22."
            )
        }
    ]

    # REPL loop
    turn_count = 0
    while True:
        try:
            user_input = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n👋 Goodbye.")
            break

        if not user_input:
            continue

        if user_input.lower() in {"exit", "quit"}:
            print("👋 Goodbye.")
            break

        turn_count += 1
        t0 = time.time()

        # Add user message
        messages.append({"role": "user", "content": user_input})

        try:
            # Call LLM with function calling
            response = llm_client.client.chat.completions.create(
                model=llm_client.model,
                messages=messages,
                tools=available_tools,
                tool_choice="auto",
                temperature=0.1,
            )

            response_message = response.choices[0].message
            tool_calls = response_message.tool_calls

            # If LLM wants to call tools
            if tool_calls:
                # Add assistant's response with tool calls
                messages.append(response_message)

                # Execute each tool call
                for tool_call in tool_calls:
                    function_name = tool_call.function.name
                    function_args = eval(tool_call.function.arguments)

                    logging.info(f"Calling tool: {function_name} with args: {function_args}")

                    # Execute the tool
                    tool_result = _dispatch_tool_call(tools, function_name, function_args)

                    # Add tool result to conversation
                    messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "name": function_name,
                        "content": str(tool_result)
                    })

                # Get final response from LLM
                final_response = llm_client.client.chat.completions.create(
                    model=llm_client.model,
                    messages=messages,
                    temperature=0.1,
                )

                assistant_message = final_response.choices[0].message.content
                messages.append({"role": "assistant", "content": assistant_message})
            else:
                # No tool calls, just use the response
                assistant_message = response_message.content
                messages.append({"role": "assistant", "content": assistant_message})

            # Display response
            print(f"\n{assistant_message}\n")

            latency_ms = int((time.time() - t0) * 1000)
            logging.info(f"Turn {turn_count} completed in {latency_ms}ms")

        except LLMError as e:
            print(f"\n⚠️  LLM error: {e}\n")
            logging.error(f"LLM error: {e}")

        except Exception as e:
            print(f"\n❌ Error: {e}\n")
            logging.error(f"Unexpected error: {e}", exc_info=True)


if __name__ == "__main__":
    main()