Multi-Agent Earnings Call Analysis System (MAECAS)

import os
import json
import hashlib
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Any, Literal, Optional, Union
from abc import ABC, abstractmethod

from pydantic import BaseModel, Field
from mistralai import Mistral
from IPython.display import display, Markdown

os.environ['MISTRAL_API_KEY'] = '<YOUR MISTRALAI API KEY>'  # Get your API key from https://console.mistral.ai/api-keys/
api_key = os.environ.get('MISTRAL_API_KEY')

mistral_client = Mistral(api_key=os.environ.get("MISTRAL_API_KEY"))

DEFAULT_MODEL = "mistral-small-latest"
STRUCTURED_MODEL = "mistral-large-latest"
OCR_MODEL = "mistral-ocr-latest"

class FinancialInsight(BaseModel):
    """Financial insights extracted from transcript"""
    metric_name: str = Field(description="Name of the financial metric")
    value: Optional[str] = Field(description="Numerical or textual value of the metric")
    context: str = Field(description="Surrounding context for the metric")
    quarter: Literal["Q1", "Q2", "Q3", "Q4"] = Field(description="Quarter the insight relates to (e.g., Q1, Q2)")
    confidence: float = Field(description="Confidence score for the insight with limits ge=0.0, le=1.0")

class FinancialInsightsResponse(BaseModel):
    """Wrapper for list of financial insights"""
    insights: List[FinancialInsight] = Field(description="Collection of financial insights")

class StrategicInsight(BaseModel):
    """Strategic insights about business direction"""
    initiative: str = Field(description="Name of the strategic initiative")
    description: str = Field(description="Details about the strategic initiative")
    timeframe: Optional[str] = Field(description="Expected timeline for implementation")
    quarter: Literal["Q1", "Q2", "Q3", "Q4"] = Field(description="Quarter the insight relates to (e.g., Q1, Q2, Q3, Q4)")
    importance: int = Field(description="Importance rating with limits ge=1, le=5")

class StrategicInsightsResponse(BaseModel):
    """Wrapper for list of strategic insights"""
    insights: List[StrategicInsight] = Field(description="Collection of strategic insights")

class SentimentInsight(BaseModel):
    """Insights about management sentiment"""
    topic: str = Field(description="Subject matter being discussed")
    sentiment: Literal["very negative", "negative", "neutral", "positive", "very positive"] = Field(description="Tone expressed by management")
    evidence: str = Field(description="Quote or context supporting the sentiment analysis")
    speaker: str = Field(description="Person who expressed the sentiment")
    quarter: Literal["Q1", "Q2", "Q3", "Q4"] = Field(description="Quarter the insight relates to (e.g., Q1, Q2, Q3, Q4)")

class SentimentInsightsResponse(BaseModel):
    """Wrapper for list of sentiment insights"""
    insights: List[SentimentInsight] = Field(description="Collection of sentiment insights")

class RiskInsight(BaseModel):
    """Identified risks or challenges"""
    risk_factor: str = Field(description="Name or type of risk identified")
    description: str = Field(description="Details about the risk")
    potential_impact: str = Field(description="Possible consequences of the risk")
    mitigation_mentioned: Optional[str] = Field(description="Strategies to address the risk")
    quarter: Literal["Q1", "Q2", "Q3", "Q4"] = Field(description="Quarter the insight relates to (e.g., Q1, Q2, Q3, Q4)")
    severity: int = Field(description="Severity rating with limits ge=1, le=5")

class RiskInsightsResponse(BaseModel):
    """Wrapper for list of risk insights"""
    insights: List[RiskInsight] = Field(description="Collection of risk insights")

class CompetitorInsight(BaseModel):
    """Insights about competitive positioning"""
    competitor: Optional[str] = Field(description="Name of the competitor company")
    market_segment: str = Field(description="Specific market area being discussed")
    positioning: str = Field(description="Competitive stance or market position")
    quarter: Literal["Q1", "Q2", "Q3", "Q4"] = Field(description="Quarter the insight relates to (e.g., Q1, Q2, Q3, Q4)")
    mentioned_by: str = Field(description="Person who mentioned the competitive information")

class CompetitorInsightsResponse(BaseModel):
    """Wrapper for list of competitor insights"""
    insights: List[CompetitorInsight] = Field(description="Collection of competitor insights")

class TemporalInsight(BaseModel):
    """Insights about trends across quarters"""
    trend_type: Literal["growth", "decline", "stable", "volatile", "emerging", "fading"] = Field(description="Direction or pattern of the trend")
    topic: str = Field(description="Subject matter of the trend")
    description: str = Field(description="Explanation of the trend's significance")
    quarters_observed: List[Literal["Q1", "Q2", "Q3", "Q4"]] = Field(description="Quarters where the trend appears")
    supporting_evidence: str = Field(description="Data or quotes supporting the trend identification")

class TemporalInsightsResponse(BaseModel):
    """Wrapper for list of temporal insights"""
    insights: List[TemporalInsight] = Field(description="Collection of temporal insights")

class QueryAnalysis(BaseModel):
    """Analysis of user query to determine required components"""
    quarters: List[str] = Field(description="List of quarters to analyze")
    agent_types: List[str] = Field(description="List of agent types to use")
    temporal_analysis_required: bool = Field(description="Whether temporal analysis across quarters is needed")
    query_intent: str = Field(description="Brief description of user's intent")

class ReportSection(BaseModel):
    """Section of the final report"""
    title: str = Field(description="Heading for the report section")
    content: str = Field(description="Main text content of the section")
    subsections: Optional[List["ReportSection"]] = Field(description="Nested sections within this section.")

class PDFParser:
    """Parse a transcript PDF file and extract text from all pages using Mistral OCR."""

    CACHE_DIR = Path("transcript_cache")

    @staticmethod
    def _ensure_cache_dir():
        """Make sure cache directory exists"""
        PDFParser.CACHE_DIR.mkdir(exist_ok=True)

    @staticmethod
    def _get_cache_path(file_path: str) -> Path:
        """Get the path for a cached transcript file"""
        # Create a hash of the file path to use as the cache filename
        file_hash = hashlib.md5(file_path.encode()).hexdigest()
        return PDFParser.CACHE_DIR / f"{file_hash}.txt"

    @staticmethod
    def read_transcript(file_path: str, mistral_client: Mistral) -> str:
        """Extract text from PDF transcript using Mistral OCR"""
        print(f"Processing PDF file: {file_path}")

        uploaded_pdf = mistral_client.files.upload(
            file={
                "file_name": file_path,
                "content": open(file_path, "rb"),
            },
            purpose="ocr"
        )

        signed_url = mistral_client.files.get_signed_url(file_id=uploaded_pdf.id)

        ocr_response = mistral_client.ocr.process(
            model=OCR_MODEL,
            document={
                "type": "document_url",
                "document_url": signed_url.url,
            }
        )

        text = "\n".join([x.markdown for x in (ocr_response.pages)])
        return text

    @staticmethod
    def get_transcript_by_quarter(company: str, quarter: str, year: str, mistral_client: Mistral) -> str:
        """Get the transcript for a specific quarter"""
        company_lower = company.lower()
        file_path = f"{company_lower}_earnings_{year}_{quarter}.pdf"

        PDFParser._ensure_cache_dir()
        cache_path = PDFParser._get_cache_path(file_path)

        # Check if transcript is in cache
        if cache_path.exists():
            print(f"Using cached transcript for {company} {year} {quarter}")
            with open(cache_path, "r", encoding="utf-8") as f:
                return f.read()
        else:
            try:
                print(f"Parsing transcript for {company} {year} {quarter}")
                transcript = PDFParser.read_transcript(file_path, mistral_client)

                # Store in cache for future use
                with open(cache_path, "w", encoding="utf-8") as f:
                    f.write(transcript)

                print(f"Cached transcript for {company} {year} {quarter}")
                return transcript
            except Exception as e:
                print(f"Error processing transcript: {str(e)}")
                raise

class InsightsStore:
    """Centralized storage for insights across all quarters and analysis types"""

    def __init__(self, company: str, year: str):
        self.company = company.lower()
        self.year = year
        self.db_path = Path(f"{self.company}_{self.year}_insights.json")
        self.insights = self._load_insights()

    def _load_insights(self) -> Dict:
        """Load insights from database file or initialize if not exists"""
        if self.db_path.exists():
            with open(self.db_path, "r", encoding="utf-8") as f:
                return json.load(f)
        else:
            return {
                "financial": {},
                "strategic": {},
                "sentiment": {},
                "risk": {},
                "competitor": {},
                "temporal": {}
            }

    def save_insights(self):
        """Save insights to database file"""
        with open(self.db_path, "w", encoding="utf-8") as f:
            json.dump(self.insights, f, indent=2)

    def add_insights(self, insight_type: str, quarter: str, insights: List):
        """Add insights for a specific type and quarter"""
        if quarter not in self.insights[insight_type]:
            self.insights[insight_type][quarter] = []

        # Convert insights to dictionaries for storage
        insight_dicts = []
        for insight in insights:
            if hasattr(insight, "dict"):
                insight_dicts.append(insight.dict())
            elif isinstance(insight, dict):
                insight_dicts.append(insight)
            else:
                insight_dicts.append({"content": str(insight)})

        # Append new insights
        self.insights[insight_type][quarter] = insight_dicts
        self.save_insights()

    def get_insights(self, insight_type=None, quarters=None):
        """Retrieve insights, optionally filtered by type and quarters"""
        if insight_type is None:
            return self.insights

        if quarters is None:
            return self.insights[insight_type]

        filtered = {}
        for q in quarters:
            if q in self.insights[insight_type]:
                filtered[q] = self.insights[insight_type][q]

        return filtered

class Agent(ABC):
    """Base class for all specialized agents"""

    def __init__(self):
        self.client = mistral_client

    @abstractmethod
    def analyze(self, transcript: str, quarter: str) -> Any:
        """Analyze the transcript and return insights"""
        pass

class FinancialAgent(Agent):
    """Agent for financial analysis of earnings call transcripts"""

    def analyze(self, transcript: str, quarter: str) -> List[FinancialInsight]:
        """Extract financial insights from transcript"""
        system_prompt = """
        You are a financial analyst focused on extracting key financial metrics and performance
        indicators from the earnings call transcripts.

        Focus on:
        - Revenue figures (overall and by segment)
        - Profit margins
        - Growth rates
        - Forward guidance
        - Capital expenditures
        - Cash flow metrics
        - Any financial KPIs mentioned

        Extract only facts that are explicitly stated in the transcript, with their proper context.
        Keep the insights as short as possible.
        """

        print(f"Extracting financial insights for quarter {quarter}...")

        response = self.client.chat.parse(
            model=STRUCTURED_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Extract financial insights from this earnings call transcript for {quarter}:\n\n{transcript}"}
            ],
            response_format=FinancialInsightsResponse,
            temperature=0.1
        )

        print(f"Financial agent completed for {quarter}")
        parsed_response = json.loads(response.choices[0].message.content)
        return parsed_response['insights']

class StrategicAgent(Agent):
    """Agent for strategic analysis of earnings call transcripts"""

    def analyze(self, transcript: str, quarter: str) -> List[StrategicInsight]:
        """Extract strategic insights from transcript"""
        system_prompt = """
        You are a business strategy analyst focused on the company's strategic direction.

        Extract insights about:
        - Product roadmaps
        - Market expansions
        - Strategic partnerships
        - R&D investments
        - Long-term vision
        - Business model changes
        - Market segments of focus

        Focus on extracting concrete strategic initiatives and plans, not general statements.
        Assign an importance score (1-5) based on how central it appears to the company's strategy.
        """

        print(f"Extracting strategic insights for quarter {quarter}...")

        response = self.client.chat.parse(
            model=STRUCTURED_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Extract strategic insights from this earnings call transcript for {quarter}:\n\n{transcript}"}
            ],
            response_format=StrategicInsightsResponse,
            temperature=0.1
        )

        print(f"Strategic agent completed for {quarter}")
        parsed_response = json.loads(response.choices[0].message.content)
        return parsed_response['insights']

class SentimentAgent(Agent):
    """Agent for sentiment analysis of earnings call transcripts"""

    def analyze(self, transcript: str, quarter: str) -> List[SentimentInsight]:
        """Extract sentiment insights from transcript"""
        system_prompt = """
        You are an expert in analyzing sentiment and tone in corporate communications.

        Focus on:
        - Management's confidence level
        - Tone when discussing different business segments
        - Enthusiasm for future prospects
        - Concerns or hesitations
        - Changes in sentiment when answering analyst questions

        Extract specific topics and the sentiment expressed about them by specific speakers.
        Use the transcript to find evidence of the sentiment you identify.
        """

        print(f"Extracting sentiment insights for quarter {quarter}...")

        response = self.client.chat.parse(
            model=STRUCTURED_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Extract sentiment insights from this earnings call transcript for {quarter}:\n\n{transcript}"}
            ],
            response_format=SentimentInsightsResponse,
            temperature=0.1
        )

        print(f"Sentiment agent completed for {quarter}")
        parsed_response = json.loads(response.choices[0].message.content)
        return parsed_response['insights']

class RiskAgent(Agent):
    """Agent for risk analysis of earnings call transcripts"""

    def analyze(self, transcript: str, quarter: str) -> List[RiskInsight]:
        """Extract risk insights from transcript"""
        system_prompt = """
        You are a risk analyst specialized in identifying challenges, uncertainties, and risk factors
        mentioned in earnings calls.

        Focus on:
        - Supply chain challenges
        - Market uncertainties
        - Competitive pressures
        - Regulatory concerns
        - Technical challenges
        - Execution risks
        - Macroeconomic factors

        For each risk, identify its potential impact and any mentioned mitigation strategies.
        Assign a severity score (1-5) based on how serious the risk appears from the transcript.
        """

        print(f"Extracting risk insights for quarter {quarter}...")

        response = self.client.chat.parse(
            model=STRUCTURED_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Extract risk insights from this earnings call transcript for {quarter}:\n\n{transcript}"}
            ],
            response_format=RiskInsightsResponse,
            temperature=0.1
        )

        print(f"Risk agent completed for {quarter}")
        parsed_response = json.loads(response.choices[0].message.content)
        return parsed_response['insights']

class CompetitorAgent(Agent):
    """Agent for competitive analysis of earnings call transcripts"""

    def analyze(self, transcript: str, quarter: str) -> List[CompetitorInsight]:
        """Extract competitor insights from transcript"""
        system_prompt = """
        You are a competitive intelligence analyst focused on the company's positioning relative to competitors.

        Focus on:
        - Direct mentions of competitors
        - Market share discussions
        - Competitive advantages or disadvantages
        - Differentiation strategies
        - Responses to competitive threats
        - Emerging competition

        Extract specific insights about the company's competitive positioning in different market segments.
        Note who mentioned the competitive information (CEO, CFO, analyst, etc.)
        """

        print(f"Extracting competitor insights for quarter {quarter}...")

        response = self.client.chat.parse(
            model=STRUCTURED_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Extract competitive insights from this earnings call transcript for {quarter}:\n\n{transcript}"}
            ],
            response_format=CompetitorInsightsResponse,
            temperature=0.1
        )

        print(f"Competitor agent completed for {quarter}")
        parsed_response = json.loads(response.choices[0].message.content)
        return parsed_response['insights']

class TemporalAnalysisAgent(Agent):
    """Agent for analyzing trends across quarters"""

    def analyze(self, all_insights: Dict) -> List[TemporalInsight]:
        """Analyze trends and patterns across quarters"""
        system_prompt = """
        You are a trend analyst specialized in identifying patterns, changes, and developments
        across multiple quarters of earnings calls.

        Focus on:
        - Growing or declining emphasis on specific topics
        - Evolving business priorities
        - Shifts in competitive positioning
        - Changes in risk factors
        - Sentiment trends

        Identify meaningful patterns that show how the business is evolving over time.
        Use specific evidence from multiple quarters to support each trend you identify.
        """

        print("Running temporal analysis across quarters...")

        # Format insights for analysis
        formatted_insights = self._format_insights_for_analysis(all_insights)

        response = self.client.chat.parse(
            model=STRUCTURED_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Analyze these insights across quarters to identify trends and patterns:\n\n{formatted_insights}"}
            ],
            response_format=TemporalInsightsResponse,
            temperature=0.2
        )

        print("Temporal analysis completed")
        parsed_response = json.loads(response.choices[0].message.content)
        return parsed_response['insights']

    def _format_insights_for_analysis(self, all_insights: Dict) -> str:
        """Format all insights for temporal analysis"""
        formatted = ""
        for agent_type, quarters_data in all_insights.items():
            formatted += f"\n## {agent_type.capitalize()} Insights Across Quarters:\n"
            for quarter, insights in quarters_data.items():
                formatted += f"\n### {quarter}:\n"
                if isinstance(insights, list):
                    for insight in insights:
                        # Convert insight object to string representation
                        if isinstance(insight, dict):
                            insight_str = json.dumps(insight)
                        else:
                            insight_str = str(insight)
                        formatted += f"- {insight_str}\n"
                else:
                    formatted += f"{insights}\n"
        return formatted

class QueryProcessor:
    """Processes user queries to determine workflow requirements"""

    def __init__(self):
        self.client = mistral_client

    def analyze_query(self, query: str, company: str) -> QueryAnalysis:
        """Analyze user query to determine required components"""
        system_prompt = f"""
        You are a query analyzer for {company} earnings call transcripts.
        Extract key information about which quarters and agent types are needed.

        For agent_types, select from these options: Financial, Strategic, Sentiment, Risk, Competitor
        For quarters, select from these options: Q1, Q2, Q3, Q4

        Determine if temporal analysis is needed (comparing across quarters).
        Provide a brief description of the user's intent.
        """

        print(f"Analyzing query: {query}")

        response = self.client.chat.parse(
            model=STRUCTURED_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Analyze this query about {company}: {query}"}
            ],
            response_format=QueryAnalysis,
            temperature=0
        )

        print("Query analysis completed")
        parsed_response = json.loads(response.choices[0].message.content)

        # Display query analysis results
        print(f"Quarters needed: {parsed_response['quarters']}")
        print(f"Agent types needed: {parsed_response['agent_types']}")
        print(f"Temporal analysis required: {parsed_response['temporal_analysis_required']}")
        print(f"Query intent: {parsed_response['query_intent']}")

        return parsed_response