296 lines
10 KiB
Python
296 lines
10 KiB
Python
import os
|
|
from typing import cast
|
|
|
|
from dotenv import load_dotenv
|
|
from langchain.agents import create_agent
|
|
from langchain.chat_models import BaseChatModel
|
|
from langchain.tools import tool
|
|
from langchain_openai import ChatOpenAI
|
|
from tavily import AsyncTavilyClient
|
|
|
|
from blueprints.rag.logic import query_vector_store
|
|
from utils.ynab_service import YNABService
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
# Configure LLM with llama-server or OpenAI fallback
|
|
llama_url = os.getenv("LLAMA_SERVER_URL")
|
|
if llama_url:
|
|
llama_chat = ChatOpenAI(
|
|
base_url=llama_url,
|
|
api_key="not-needed",
|
|
model=os.getenv("LLAMA_MODEL_NAME", "llama-3.1-8b-instruct"),
|
|
)
|
|
else:
|
|
llama_chat = None
|
|
|
|
openai_fallback = ChatOpenAI(model="gpt-5-mini")
|
|
model_with_fallback = cast(
|
|
BaseChatModel,
|
|
llama_chat.with_fallbacks([openai_fallback]) if llama_chat else openai_fallback,
|
|
)
|
|
client = AsyncTavilyClient(api_key=os.getenv("TAVILY_API_KEY", ""))
|
|
|
|
# Initialize YNAB service (will only work if YNAB_ACCESS_TOKEN is set)
|
|
try:
|
|
ynab_service = YNABService()
|
|
ynab_enabled = True
|
|
except (ValueError, Exception) as e:
|
|
print(f"YNAB service not initialized: {e}")
|
|
ynab_enabled = False
|
|
|
|
|
|
@tool
|
|
async def web_search(query: str) -> str:
|
|
"""Search the web for current information using Tavily.
|
|
|
|
Use this tool when you need to:
|
|
- Find current information not in the knowledge base
|
|
- Look up recent events, news, or updates
|
|
- Verify facts or get additional context
|
|
- Search for information outside of Simba's documents
|
|
|
|
Args:
|
|
query: The search query to look up on the web
|
|
|
|
Returns:
|
|
Search results from the web with titles, content, and source URLs
|
|
"""
|
|
response = await client.search(query=query, search_depth="basic")
|
|
results = response.get("results", [])
|
|
|
|
if not results:
|
|
return "No results found for the query."
|
|
|
|
formatted = "\n\n".join(
|
|
[
|
|
f"**{result['title']}**\n{result['content']}\nSource: {result['url']}"
|
|
for result in results[:5]
|
|
]
|
|
)
|
|
return formatted
|
|
|
|
|
|
@tool(response_format="content_and_artifact")
|
|
async def simba_search(query: str):
|
|
"""Search through Simba's medical records, veterinary documents, and personal information.
|
|
|
|
Use this tool whenever the user asks questions about:
|
|
- Simba's health history, medical records, or veterinary visits
|
|
- Medications, treatments, or diagnoses
|
|
- Weight, diet, or physical characteristics over time
|
|
- Veterinary recommendations or advice
|
|
- Ryan's (the owner's) information related to Simba
|
|
- Any factual information that would be found in documents
|
|
|
|
Args:
|
|
query: The user's question or information need about Simba
|
|
|
|
Returns:
|
|
Relevant information from Simba's documents
|
|
"""
|
|
print(f"[SIMBA SEARCH] Tool called with query: {query}")
|
|
serialized, docs = await query_vector_store(query=query)
|
|
print(f"[SIMBA SEARCH] Found {len(docs)} documents")
|
|
print(f"[SIMBA SEARCH] Serialized result length: {len(serialized)}")
|
|
print(f"[SIMBA SEARCH] First 200 chars: {serialized[:200]}")
|
|
return serialized, docs
|
|
|
|
|
|
@tool
|
|
def ynab_budget_summary() -> str:
|
|
"""Get overall budget summary and health status from YNAB.
|
|
|
|
Use this tool when the user asks about:
|
|
- Overall budget health or status
|
|
- How much money is to be budgeted
|
|
- Total budget amounts or spending
|
|
- General budget overview questions
|
|
|
|
Returns:
|
|
Summary of budget health, to-be-budgeted amount, total budgeted,
|
|
total activity, and available amounts.
|
|
"""
|
|
if not ynab_enabled:
|
|
return "YNAB integration is not configured. Please set YNAB_ACCESS_TOKEN environment variable."
|
|
|
|
try:
|
|
summary = ynab_service.get_budget_summary()
|
|
return summary["summary"]
|
|
except Exception as e:
|
|
return f"Error fetching budget summary: {str(e)}"
|
|
|
|
|
|
@tool
|
|
def ynab_search_transactions(
|
|
start_date: str = "",
|
|
end_date: str = "",
|
|
category_name: str = "",
|
|
payee_name: str = "",
|
|
) -> str:
|
|
"""Search YNAB transactions by date range, category, or payee.
|
|
|
|
Use this tool when the user asks about:
|
|
- Specific transactions or purchases
|
|
- Spending at a particular store or payee
|
|
- Transactions in a specific category
|
|
- What was spent during a time period
|
|
|
|
Args:
|
|
start_date: Start date in YYYY-MM-DD format (optional, defaults to 30 days ago)
|
|
end_date: End date in YYYY-MM-DD format (optional, defaults to today)
|
|
category_name: Filter by category name (optional, partial match)
|
|
payee_name: Filter by payee/store name (optional, partial match)
|
|
|
|
Returns:
|
|
List of matching transactions with dates, amounts, categories, and payees.
|
|
"""
|
|
if not ynab_enabled:
|
|
return "YNAB integration is not configured. Please set YNAB_ACCESS_TOKEN environment variable."
|
|
|
|
try:
|
|
result = ynab_service.get_transactions(
|
|
start_date=start_date or None,
|
|
end_date=end_date or None,
|
|
category_name=category_name or None,
|
|
payee_name=payee_name or None,
|
|
)
|
|
|
|
if result["count"] == 0:
|
|
return "No transactions found matching the specified criteria."
|
|
|
|
# Format transactions for readability
|
|
txn_list = []
|
|
for txn in result["transactions"][:10]: # Limit to 10 for readability
|
|
txn_list.append(
|
|
f"- {txn['date']}: {txn['payee']} - ${abs(txn['amount']):.2f} ({txn['category'] or 'Uncategorized'})"
|
|
)
|
|
|
|
return (
|
|
f"Found {result['count']} transactions from {result['start_date']} to {result['end_date']}. "
|
|
f"Total: ${abs(result['total_amount']):.2f}\n\n"
|
|
+ "\n".join(txn_list)
|
|
+ (
|
|
f"\n\n(Showing first 10 of {result['count']} transactions)"
|
|
if result["count"] > 10
|
|
else ""
|
|
)
|
|
)
|
|
except Exception as e:
|
|
return f"Error searching transactions: {str(e)}"
|
|
|
|
|
|
@tool
|
|
def ynab_category_spending(month: str = "") -> str:
|
|
"""Get spending breakdown by category for a specific month.
|
|
|
|
Use this tool when the user asks about:
|
|
- Spending by category
|
|
- What categories were overspent
|
|
- Monthly spending breakdown
|
|
- Budget vs actual spending for a month
|
|
|
|
Args:
|
|
month: Month in YYYY-MM format (optional, defaults to current month)
|
|
|
|
Returns:
|
|
Spending breakdown by category with budgeted, spent, and available amounts.
|
|
"""
|
|
if not ynab_enabled:
|
|
return "YNAB integration is not configured. Please set YNAB_ACCESS_TOKEN environment variable."
|
|
|
|
try:
|
|
result = ynab_service.get_category_spending(month=month or None)
|
|
|
|
summary = (
|
|
f"Budget spending for {result['month']}:\n"
|
|
f"Total budgeted: ${result['total_budgeted']:.2f}\n"
|
|
f"Total spent: ${result['total_spent']:.2f}\n"
|
|
f"Total available: ${result['total_available']:.2f}\n"
|
|
)
|
|
|
|
if result["overspent_categories"]:
|
|
summary += (
|
|
f"\nOverspent categories ({len(result['overspent_categories'])}):\n"
|
|
)
|
|
for cat in result["overspent_categories"][:5]:
|
|
summary += f"- {cat['name']}: Budgeted ${cat['budgeted']:.2f}, Spent ${cat['spent']:.2f}, Over by ${cat['overspent_by']:.2f}\n"
|
|
|
|
# Add top spending categories
|
|
summary += "\nTop spending categories:\n"
|
|
for cat in result["categories"][:10]:
|
|
if cat["activity"] < 0: # Only show spending (negative activity)
|
|
summary += f"- {cat['category']}: ${abs(cat['activity']):.2f} (budgeted: ${cat['budgeted']:.2f}, available: ${cat['available']:.2f})\n"
|
|
|
|
return summary
|
|
except Exception as e:
|
|
return f"Error fetching category spending: {str(e)}"
|
|
|
|
|
|
@tool
|
|
def ynab_insights(months_back: int = 3) -> str:
|
|
"""Generate insights about spending patterns and budget health over time.
|
|
|
|
Use this tool when the user asks about:
|
|
- Spending trends or patterns
|
|
- Budget recommendations
|
|
- Which categories are frequently overspent
|
|
- How current spending compares to past months
|
|
- Overall budget health analysis
|
|
|
|
Args:
|
|
months_back: Number of months to analyze (default 3, max 6)
|
|
|
|
Returns:
|
|
Insights about spending trends, frequently overspent categories,
|
|
and personalized recommendations.
|
|
"""
|
|
if not ynab_enabled:
|
|
return "YNAB integration is not configured. Please set YNAB_ACCESS_TOKEN environment variable."
|
|
|
|
try:
|
|
# Limit to reasonable range
|
|
months_back = min(max(1, months_back), 6)
|
|
result = ynab_service.get_spending_insights(months_back=months_back)
|
|
|
|
if "error" in result:
|
|
return result["error"]
|
|
|
|
summary = (
|
|
f"Spending insights for the last {months_back} months:\n\n"
|
|
f"Average monthly spending: ${result['average_monthly_spending']:.2f}\n"
|
|
f"Current month spending: ${result['current_month_spending']:.2f}\n"
|
|
f"Spending trend: {result['spending_trend']}\n"
|
|
)
|
|
|
|
if result["frequently_overspent_categories"]:
|
|
summary += "\nFrequently overspent categories:\n"
|
|
for cat in result["frequently_overspent_categories"][:5]:
|
|
summary += f"- {cat['category']}: overspent in {cat['months_overspent']} of {months_back} months\n"
|
|
|
|
if result["recommendations"]:
|
|
summary += "\nRecommendations:\n"
|
|
for rec in result["recommendations"]:
|
|
summary += f"- {rec}\n"
|
|
|
|
return summary
|
|
except Exception as e:
|
|
return f"Error generating insights: {str(e)}"
|
|
|
|
|
|
# Create tools list based on what's available
|
|
tools = [simba_search, web_search]
|
|
if ynab_enabled:
|
|
tools.extend(
|
|
[
|
|
ynab_budget_summary,
|
|
ynab_search_transactions,
|
|
ynab_category_spending,
|
|
ynab_insights,
|
|
]
|
|
)
|
|
|
|
# Llama 3.1 supports native function calling via OpenAI-compatible API
|
|
main_agent = create_agent(model=model_with_fallback, tools=tools)
|