src/agents/TaxAgent.py

import logging

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI

from src.core.AgentCommand import AgentCommand
from src.core.FinanceState import FinanceState
from src.core.errors import add_error
from src.core.settings import get_settings
from src.rag.KnowledgeBase import KnowledgeBase
from src.rag.TavilySearchRag import TavilySearchRag
from src.data.SemanticCache import SemanticCache

logger = logging.getLogger(__name__)


class TaxAgent(AgentCommand):
    def __init__(self, state: FinanceState):
        load_dotenv()
        self.state = state
        self.tavily = TavilySearchRag()
        self.kb = KnowledgeBase()
        self.semantic_cache = SemanticCache()
        settings = get_settings()
        self.client = ChatOpenAI(model=settings.models.agent_model)

    def process(self):
        query = self.state.get("user_query", "") or ""
        trace_id = str(self.state.get("trace_id") or "")
        logger.info("[trace=%s] TaxAgent.start query_len=%d", trace_id, len(query))

        profile = self.state.get("user_profile") or {}
        risk = str(profile.get("risk") or "").strip().lower()
        experience = str(profile.get("experience") or "").strip().lower()
        cache_query = (
            "agent=tax\n"
            f"risk={risk}; experience={experience}\n"
            f"query={query.strip().lower()}"
        ).strip()

        # Agent-level semantic cache (prevents cross-intent collisions).
        try:
            cached = self.semantic_cache.check_cache(cache_query, threshold=0.91)
        except Exception as e:
            cached = None
            add_error(
                self.state,
                code="semantic_cache_error",
                message=str(e),
                agent="tax_agent",
            )
        if isinstance(cached, str) and cached.strip():
            logger.info("[trace=%s] TaxAgent.semantic_cache_hit", trace_id)
            self.state["response"] = cached.strip()
            return self.state
        logger.info("[trace=%s] TaxAgent.semantic_cache_miss", trace_id)

        # KB-first for tax topics.
        try:
            sources = self.kb.retrieve(query, k=4, categories=["tax"])
        except Exception as e:
            logger.info("[trace=%s] KB retrieval unavailable: %s", trace_id, e)
            sources = []

        if sources:
            logger.info("[trace=%s] TaxAgent.source=kb hits=%d", trace_id, len(sources))
            snippets = "\n\n".join(
                [
                    f"TITLE: {s.title}\nCATEGORY: {s.category}\nEXCERPT:\n{s.excerpt}"
                    for s in sources
                ]
            )
            citations = "\n".join([f"- {s.title} ({s.source_path})" for s in sources])
            system = (
                "You are a financial education assistant focusing on taxes. "
                "You must avoid personalized tax advice; provide general education and encourage consulting a professional. "
                "Use only the provided KB excerpts as truth. Always include: 'Educational only, not tax or financial advice.'"
            )
            prompt = f"""
User question:
{query}

KB excerpts:
{snippets}

Write a clear, general educational answer. End with a short 'Sources' list.
"""
            try:
                msg = self.client.invoke(
                    [
                        {"role": "system", "content": system},
                        {"role": "user", "content": prompt},
                    ],
                    temperature=0.2,
                )
                text = msg.content.strip()
                text += "\n\nSources:\n" + citations
                self.state["response"] = text
                self.state["retrieved_sources"] = [
                    {
                        "title": s.title,
                        "category": s.category,
                        "source_path": s.source_path,
                    }
                    for s in sources
                ]
                try:
                    self.semantic_cache.save_to_cache(cache_query, self.state["response"])
                except Exception:
                    pass
                return self.state
            except Exception as e:
                logger.info("[trace=%s] LLM synthesis failed, returning excerpts: %s", trace_id, e)
                self.state["response"] = (
                    "Educational only, not tax or financial advice.\n\n"
                    "I found these relevant KB excerpts:\n\n"
                    + snippets
                    + "\n\nSources:\n"
                    + citations
                )
                self.state["retrieved_sources"] = [
                    {
                        "title": s.title,
                        "category": s.category,
                        "source_path": s.source_path,
                    }
                    for s in sources
                ]
                try:
                    self.semantic_cache.save_to_cache(cache_query, self.state["response"])
                except Exception:
                    pass
                return self.state

        # Fallback: web search
        tavily_results = self.tavily.search(query)
        logger.info(
            "[trace=%s] TaxAgent.source=tavily results=%d",
            trace_id,
            len(tavily_results or []),
        )
        if (
            isinstance(tavily_results, list)
            and tavily_results
            and tavily_results[0].get("error")
        ):
            add_error(
                self.state,
                code="tavily_error",
                message=str(tavily_results[0].get("error")),
                agent="tax_agent",
            )
            self.state["response"] = (
                "Educational only, not tax or financial advice.\n\n"
                f"Error retrieving tax information: {tavily_results[0].get('error')}"
            )
            return self.state

        response_text = "Tax information (web search):\n"
        for result in tavily_results or []:
            if result.get("title"):
                response_text += f"Title: {result['title']}\n"
            if result.get("url"):
                response_text += f"URL: {result['url']}\n"
            if result.get("content"):
                response_text += f"Content: {result['content']}...\n"
            response_text += "\n"
        self.state["response"] = response_text.strip()
        try:
            self.semantic_cache.save_to_cache(cache_query, self.state["response"])
        except Exception:
            pass
        return self.state