Spaces:
Sleeping
Sleeping
| """ | |
| 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)}") |