agent.py

"""
CFA AI Agent - LangChain Agent Setup
This module sets up the LangChain agent with Finance-Llama-8B model and financial tools.
"""

import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from langchain_community.llms import HuggingFacePipeline
from langchain.agents import initialize_agent, AgentType
from langchain.memory import ConversationBufferMemory
from langchain.schema import SystemMessage
from langchain.prompts import MessagesPlaceholder
from typing import List, Any, Optional

# Import our custom tools
from tools.finance_tools import (
    calculate_dcf,
    calculate_sharpe_ratio,
    compare_pe_ratios,
    calculate_beta,
    calculate_wacc,
    financial_ratios_analysis
)
from tools.data_fetcher import (
    get_stock_price,
    get_historical_data,
    get_company_info,
    get_financial_statements,
    get_market_indices,
    compare_stocks
)


class CFAAgent:
    """
    CFA AI Agent that combines Finance-Llama-8B model with financial analysis tools.
    """

    def __init__(self, model_name: str = "tarun7r/Finance-Llama-8B"):
        """
        Initialize the CFA Agent with model and tools.

        Args:
            model_name: Hugging Face model name for financial analysis
        """
        self.model_name = model_name
        self.tokenizer = None
        self.model = None
        self.llm = None
        self.agent = None
        self.memory = None
        self._setup_model()
        self._setup_tools()
        self._setup_agent()

    def _setup_model(self):
        """Load and setup the Finance-Llama-8B model."""
        try:
            print(f"Loading model: {self.model_name}")

            # Check if CUDA is available
            device = "cuda" if torch.cuda.is_available() else "cpu"
            print(f"Using device: {device}")

            # Load tokenizer
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model_name,
                trust_remote_code=True
            )

            # Add pad token if not present
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token

            # Load model with appropriate settings and memory optimization
            if device == "cuda":
                self.model = AutoModelForCausalLM.from_pretrained(
                    self.model_name,
                    torch_dtype=torch.float16,
                    device_map="auto",
                    trust_remote_code=True,
                    low_cpu_mem_usage=True,
                    load_in_8bit=True,  # Enable 8-bit quantization for memory efficiency
                    max_memory={0: "6GB"}  # Limit GPU memory usage
                )
            else:
                # For CPU, use aggressive memory optimization
                self.model = AutoModelForCausalLM.from_pretrained(
                    self.model_name,
                    trust_remote_code=True,
                    low_cpu_mem_usage=True,
                    torch_dtype=torch.float32,
                    device_map="cpu",
                    max_memory={"cpu": "8GB"}  # Limit CPU memory usage
                )

            # Create pipeline
            pipe = pipeline(
                "text-generation",
                model=self.model,
                tokenizer=self.tokenizer,
                max_new_tokens=512,
                temperature=0.1,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id,
                repetition_penalty=1.1
            )

            # Wrap in LangChain
            self.llm = HuggingFacePipeline(pipeline=pipe)
            print("✅ Model loaded successfully")

        except Exception as e:
            print(f"❌ Error loading model: {str(e)}")
            # Fallback to a smaller model or OpenAI if Finance-Llama-8B fails
            self._setup_fallback_model()

    def _setup_fallback_model(self):
        """Setup a fallback model if Finance-Llama-8B fails to load."""
        try:
            print("Setting up fallback model...")
            from langchain_community.llms import OpenAI

            # Check for OpenAI API key
            if os.getenv("OPENAI_API_KEY"):
                self.llm = OpenAI(
                    temperature=0.1,
                    model_name="gpt-3.5-turbo-instruct",
                    max_tokens=512
                )
                print("✅ Using OpenAI GPT-3.5 as fallback")
            else:
                raise ValueError("No OpenAI API key found")

        except Exception as e:
            print(f"❌ Fallback model failed: {str(e)}")
            # Last resort: use a very small local model
            try:
                pipe = pipeline(
                    "text-generation",
                    model="distilgpt2",
                    max_new_tokens=256,
                    temperature=0.7
                )
                self.llm = HuggingFacePipeline(pipeline=pipe)
                print("✅ Using DistilGPT2 as emergency fallback")
            except Exception as final_e:
                raise RuntimeError(f"All model loading attempts failed: {final_e}")

    def _setup_tools(self):
        """Setup all available financial analysis tools."""
        self.tools = [
            # Finance calculation tools
            calculate_dcf,
            calculate_sharpe_ratio,
            compare_pe_ratios,
            calculate_beta,
            calculate_wacc,
            financial_ratios_analysis,

            # Data fetching tools
            get_stock_price,
            get_historical_data,
            get_company_info,
            get_financial_statements,
            get_market_indices,
            compare_stocks
        ]
        print(f"✅ Loaded {len(self.tools)} financial analysis tools")

    def _setup_agent(self):
        """Setup the LangChain agent with memory and tools."""
        try:
            # Setup conversation memory
            self.memory = ConversationBufferMemory(
                memory_key="chat_history",
                return_messages=True,
                output_key="output"
            )

            # Initialize agent
            self.agent = initialize_agent(
                tools=self.tools,
                llm=self.llm,
                agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
                memory=self.memory,
                verbose=True,
                handle_parsing_errors=True,
                max_iterations=3,
                early_stopping_method="generate"
            )

            # Add system message for financial context
            system_message = """You are a CFA (Chartered Financial Analyst) AI assistant specialized in financial analysis, investment valuation, and portfolio management.

Your expertise includes:
- Financial statement analysis and ratio calculations
- Valuation models (DCF, comparable company analysis, etc.)
- Risk assessment and portfolio theory
- Market analysis and economic indicators
- Investment recommendations based on fundamental analysis

When answering questions:
1. Use the available financial tools to fetch real data when needed
2. Provide clear, professional explanations suitable for CFA-level analysis
3. Show your calculations and reasoning
4. Consider both quantitative and qualitative factors
5. Acknowledge limitations and assumptions in your analysis

You have access to real-time financial data and calculation tools. Use them to provide accurate, data-driven insights."""

            # Store system message for context
            self.system_message = system_message
            print("✅ Agent initialized successfully")

        except Exception as e:
            print(f"❌ Error setting up agent: {str(e)}")
            raise

    def query(self, question: str) -> str:
        """
        Process a financial query using the CFA agent.

        Args:
            question: User's financial question or request

        Returns:
            Agent's response with analysis and recommendations
        """
        try:
            # Enhance the question with context
            enhanced_question = f"""As a CFA analyst, please help with the following:

{question}

Please provide a thorough analysis using available data and financial tools. Show your work and explain your reasoning."""

            # Get response from agent
            response = self.agent.run(enhanced_question)
            return response

        except Exception as e:
            error_msg = f"Error processing query: {str(e)}"
            print(error_msg)
            return error_msg

    def get_conversation_history(self) -> List[Any]:
        """Get the current conversation history."""
        if self.memory:
            return self.memory.chat_memory.messages
        return []

    def clear_memory(self):
        """Clear the conversation memory."""
        if self.memory:
            self.memory.clear()
            print("✅ Conversation memory cleared")

    def get_available_tools(self) -> List[str]:
        """Get list of available tool names."""
        return [tool.name for tool in self.tools]

    def health_check(self) -> dict:
        """Perform a health check of the agent components."""
        status = {
            "model_loaded": self.model is not None,
            "llm_ready": self.llm is not None,
            "agent_ready": self.agent is not None,
            "memory_ready": self.memory is not None,
            "tools_count": len(self.tools),
            "device": "cuda" if torch.cuda.is_available() else "cpu"
        }
        return status


def create_cfa_agent(model_name: str = "tarun7r/Finance-Llama-8B") -> CFAAgent:
    """
    Factory function to create and return a CFA Agent instance.

    Args:
        model_name: Hugging Face model name for financial analysis

    Returns:
        Initialized CFAAgent instance
    """
    try:
        agent = CFAAgent(model_name=model_name)
        print("🎯 CFA Agent created successfully")
        return agent
    except Exception as e:
        print(f"❌ Failed to create CFA Agent: {str(e)}")
        raise


# Example usage and testing
if __name__ == "__main__":
    print("🚀 Initializing CFA AI Agent...")

    try:
        # Create agent
        cfa_agent = create_cfa_agent()

        # Health check
        health = cfa_agent.health_check()
        print("📊 Health Check Results:")
        for key, value in health.items():
            status = "✅" if value else "❌"
            print(f"  {status} {key}: {value}")

        # Test queries
        test_queries = [
            "What is the current stock price of Apple (AAPL)?",
            "Calculate the PE ratio comparison between Apple and Microsoft",
            "Explain the CAPM model in simple terms"
        ]

        print("\n🧪 Running test queries...")
        for i, query in enumerate(test_queries, 1):
            print(f"\n--- Test Query {i} ---")
            print(f"Q: {query}")
            try:
                response = cfa_agent.query(query)
                print(f"A: {response}")
            except Exception as e:
                print(f"❌ Query failed: {str(e)}")

    except Exception as e:
        print(f"❌ CFA Agent initialization failed: {str(e)}")