src/tools/dietary_tools.py

import sqlite3
import os
import json
from langchain_core.tools import tool
from src.utils.logger import setup_logger
from supabase import create_client, Client
from sentence_transformers import SentenceTransformer

logger = setup_logger("SupabaseDietaryTools")

DB_PATH = os.path.join(os.path.dirname(__file__), "../../data/dietary_guidelines.db")

# Load model once at module level
logger.info("Loading embedding model for tools...")
model = SentenceTransformer('all-MiniLM-L6-v2')

_supabase_client = None


def _build_supabase_client() -> Client:
    url = os.getenv("SUPABASE_URL")
    key = os.getenv("SUPABASE_KEY")
    if not url or not key:
        raise RuntimeError("SUPABASE_URL and SUPABASE_KEY must be set")
    return create_client(url, key)


def get_supabase_client() -> Client:
    global _supabase_client
    if _supabase_client is None:
        _supabase_client = _build_supabase_client()
    return _supabase_client


@tool
def search_guidelines(query: str):
    """
    Search for relevant medical and dietary guidelines using Supabase pgvector.
    Returns content with source and page information.
    """
    logger.info(f"Searching Supabase guidelines for: {query}")
    client = get_supabase_client()

    # Generate embedding for query
    query_embedding = model.encode(query).tolist()
    
    try:
        # Use RPC to perform similarity search (requires a match_documents function in Postgres)
        # Or use simple rpc if match_documents is defined in Supabase
        # See: https://supabase.com/docs/guides/ai/vector-columns#querying-a-vector-column
        
        rpc_params = {
            "query_embedding": query_embedding,
            "match_threshold": 0.5,
            "match_count": 3,
        }
        
        # If the user hasn't created the RPC yet, we might need to fallback to a basic select 
        # but filtering by embedding in client-side is not possible.
        # I'll assume they added the recommended 'match_documents' function.
        
        response = client.rpc("match_knowledge_base", rpc_params).execute()
        results = response.data
        
        if not results:
            return "No specific guidelines found for this query in the vector store."
        
        output = "Here are some relevant guidelines from Supabase pgvector:\n"
        for doc in results:
            metadata = doc.get("metadata", {})
            source = metadata.get("source", "Unknown")
            page = metadata.get("page_index", metadata.get("page", "Unknown"))
            content = doc.get("content", "")
            output += f"- Source: {source}, Page: {page}\n  Content: {content[:500]}...\n\n"
        return output
    except Exception as e:
        logger.error(f"Error accessing Supabase pgvector: {e}")
        # Fallback to text search if RPC fails
        try:
            response = client.table("knowledge_base").select("*").text_search("content", query).limit(3).execute()
            results = response.data
            if not results: return "No guidelines found."
            output = "Found via text search:\n"
            for doc in results:
                metadata = doc.get("metadata", {})
                output += f"- Source: {metadata.get('source')}, Page: {metadata.get('page_index')}\n  Content: {doc['content'][:500]}...\n"
            return output
        except Exception as e2:
            return f"Error retrieving guidelines: {str(e2)}"

@tool
def get_nutritional_data(food_name: str):
    """Get nutritional information for a specific food item."""
    logger.info(f"Retrieving nutritional data for: {food_name}")
    if not os.path.exists(DB_PATH):
        return "Nutritional database not found."
    
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    query_name = f"%{food_name}%"
    cursor.execute("SELECT food_name, calories, protein, carbs, fat, fiber, vitamins FROM nutritional_data WHERE food_name LIKE ?", (query_name,))
    results = cursor.fetchall()
    conn.close()
    
    if not results:
        return f"No nutritional data found for '{food_name}'."
    
    output = "Nutritional data found:\n"
    for name, cals, protein, carbs, fat, fiber, vitamins in results:
        output += f"- {name}: {cals} kcal, Protein: {protein}g, Carbs: {carbs}g, Fat: {fat}g, Fiber: {fiber}g, Vitamins: {vitamins}\n"
    return output

@tool
def page_indexed_retrieval(query: str):
    """
    Perform a Page Indexing based RAG search using Supabase.
    """
    # For now, we use the same vector search as search_guidelines
    return search_guidelines.invoke(query)