agents.py

#!/usr/bin/env python3
"""
Agents Orchestrator
Monitors `assemblyai.log` and runs three agents:
 - SpeakerIdentificationAgent → writes labeled transcript lines to `transcripts.log`
 - CustomerMatchAgent → writes best-match CRM data to `customer.log` (idempotent, only when changes)
 - RecommendationAgent → continuously rewrites `recommendations.log` with immediate guidance for STAFF

Lightweight LangGraph-style pipeline using Python classes and a coordinator loop.
"""

import os
import sys
import time
from datetime import datetime
from typing import List, Optional, Tuple

# Cerebras Cloud SDK (LLM for all agent reasoning)
try:
    from cerebras.cloud.sdk import Cerebras  # QuickStart import path
except Exception:
    Cerebras = None

ASSEMBLYAI_LOG = "./logs/assemblyai.log"
TRANSCRIPTS_LOG = "./logs/transcripts.log"
CUSTOMER_LOG = "./logs/customer.log"
RECOMMENDATIONS_LOG = "./logs/recommendations.log"
CRM_DIR = "crm"
KB_DIR = "knowledge_base"


def read_new_lines(path: str, last_size: int) -> Tuple[List[str], int]:
    try:
        size = os.path.getsize(path)
        if size < last_size:
            # Rotated or truncated
            last_size = 0
        if size == last_size:
            return [], last_size
        with open(path, "r", encoding="utf-8") as f:
            f.seek(last_size)
            data = f.read()
        lines = [ln.strip() for ln in data.splitlines() if ln.strip()]
        return lines, size
    except FileNotFoundError:
        return [], last_size


def reset_output_logs():
    now = datetime.now().isoformat()
    try:
        with open(TRANSCRIPTS_LOG, "w", encoding="utf-8") as f:
            f.write(f"# Transcripts (speaker-labeled)\n# Reset: {now}\n\n")
        with open(CUSTOMER_LOG, "w", encoding="utf-8") as f:
            f.write(f"# Customer match log\n# Reset: {now}\n\n")
        with open(RECOMMENDATIONS_LOG, "w", encoding="utf-8") as f:
            f.write(f"# Recommendations\n# Reset: {now}\n\n")
        print("🧹 Reset transcripts.log, customer.log, recommendations.log")
    except Exception as e:
        print(f"❌ Failed to reset logs: {e}")


def parse_assemblyai_line(line: str) -> Optional[Tuple[str, str, str]]:
    # Expect: ISO8601 | Speaker N | transcript
    parts = [p.strip() for p in line.split("|")]
    if len(parts) < 3:
        return None
    ts, speaker, transcript = parts[0], parts[1], "|".join(parts[2:]).strip()
    return ts, speaker, transcript


class SpeakerIdentificationAgent:
    def __init__(self, llm):
        self.llm = llm

    def _classify_role_llm(self, speaker_label: str, transcript: str) -> str:
        prompt = (
            "You are a call analysis assistant. Based on this transcript line, "
            "classify the current speaker into one of: STAFF, CUSTOMER, or SPEAKER.\n"
            f"Speaker Label: {speaker_label}\n"
            f"Transcript: {transcript}\n\n"
            "Return ONLY one of: STAFF, CUSTOMER, SPEAKER."
        )
        try:
            res = self.llm.chat.completions.create(
                model=os.getenv("CEREBRAS_MODEL", "llama3.1-8b"),
                messages=[{"role": "user", "content": prompt}],
                temperature=0.0,
            )
            out = (res.choices[0].message.content or "").strip().upper()
            if out.startswith("STAFF"):
                return "STAFF"
            if out.startswith("CUSTOMER"):
                return "CUSTOMER"
            return "SPEAKER"
        except Exception:
            return "SPEAKER"

    def process(self, ts: str, speaker_label: str, transcript: str) -> str:
        role = self._classify_role_llm(speaker_label, transcript)
        return f"{ts} | {role} | {transcript}"


class CustomerMatchAgent:
    def __init__(self, crm_dir: str, llm):
        self.crm_dir = crm_dir
        self.llm = llm
        self._current_customer_id: Optional[str] = None

    def _read_crm_files(self) -> List[Tuple[str, str]]:
        results = []
        if not os.path.isdir(self.crm_dir):
            return results
        for name in os.listdir(self.crm_dir):
            if not name.lower().endswith(".txt"):
                continue
            path = os.path.join(self.crm_dir, name)
            try:
                with open(path, "r", encoding="utf-8") as f:
                    results.append((name, f.read()))
            except Exception:
                continue
        return results

    def _best_match(self, transcript_window: List[str]) -> Optional[Tuple[str, str]]:
        """LLM-only customer matching (no heuristics).
        Returns (filename, content) or None.
        """
        candidates = self._read_crm_files()
        if not candidates:
            # No candidates
            return None
        convo = "\n".join(transcript_window[-60:])
        file_list = ", ".join([name for name, _ in candidates])
        # Include content excerpts so the LLM can reason over actual CRM data
        catalog_lines = []
        for name, content in candidates:
            excerpt = content[:1200].replace("\n", " ")
            catalog_lines.append(f"- {name}: {excerpt}")
        catalog_str = "\n".join(catalog_lines)

        # Ask the LLM to return EXACT filename present in crm/ or NONE.
        system_msg = (
            "You are a CRM matching assistant. You MUST select the single best matching customer file "
            "from the provided list, based ONLY on the conversation. If there is no clear match, respond 'NONE'."
        )
        user_msg = (
            f"Customer files available (filenames): {file_list}\n\n"
            f"Customer file excerpts (filename: excerpt):\n{catalog_str}\n\n"
            f"Recent conversation (most recent first):\n{convo}\n\n"
            "Return EXACTLY one of the filenames above or 'NONE'. No explanations."
        )
        try:
            res = self.llm.chat.completions.create(
                model=os.getenv("CEREBRAS_MODEL", "llama3.1-8b"),
                messages=[
                    {"role": "system", "content": system_msg},
                    {"role": "user", "content": user_msg},
                ],
                temperature=0.0,
            )
            choice = (res.choices[0].message.content or "").strip()
            if not choice or choice.upper().startswith("NONE"):
                return None
            # Match exactly against available filenames
            for n, content in candidates:
                if n.strip() == choice.strip():
                    return (n, content)
            # If LLM deviates (e.g., extra text), do not fallback heuristically; declare no match
            return None
        except Exception as e:
            return None

    def maybe_update_customer_log(self, transcript_window: List[str]) -> Optional[str]:
        match = self._best_match(transcript_window)
        if not match:
            return None
        fname, content = match
        customer_id = os.path.splitext(fname)[0]
        if customer_id == self._current_customer_id:
            return None
        # Update log only on change
        with open(CUSTOMER_LOG, "w", encoding="utf-8") as f:
            f.write(f"# Customer match selected at {datetime.now().isoformat()}\n")
            f.write(f"# File: {fname}\n\n")
            f.write(content.strip() + "\n")
        self._current_customer_id = customer_id
        return customer_id


class RecommendationAgent:
    def __init__(self, kb_dir: str, llm):
        self.kb_dir = kb_dir
        self.llm = llm

    def _read_kb(self) -> str:
        buff = []
        if not os.path.isdir(self.kb_dir):
            return ""
        for root, _, files in os.walk(self.kb_dir):
            for name in files:
                if not name.lower().endswith(('.txt', '.md')):
                    continue
                path = os.path.join(root, name)
                try:
                    with open(path, "r", encoding="utf-8") as f:
                        buff.append(f"\n# {name}\n" + f.read())
                except Exception:
                    continue
        return "\n".join(buff)

    def _read_customer_log(self) -> str:
        try:
            with open(CUSTOMER_LOG, "r", encoding="utf-8") as f:
                return f.read()
        except Exception:
            return ""

    def _read_transcripts(self) -> List[str]:
        try:
            with open(TRANSCRIPTS_LOG, "r", encoding="utf-8") as f:
                return [ln.strip() for ln in f.readlines() if ln.strip() and not ln.strip().startswith('#')]
        except Exception:
            return []

    def update_recommendations(self):
        kb = self._read_kb()
        customer = self._read_customer_log()
        transcripts = self._read_transcripts()
        recent = "\n".join(transcripts[-20:])

        prompt = (
            "You are an expert agent coach. Generate a concise, prioritized list of IMMEDIATE, "
            "actionable recommendations for the STAFF member on this call. Use the conversation, "
            "customer data, and knowledge base. Avoid repeating transcript text; be directive.\n\n"
            f"Recent transcripts (labeled):\n{recent}\n\n"
            f"Customer data (if any):\n{customer}\n\n"
            f"Knowledge base excerpts:\n{kb[:8000]}\n\n"
            "Return 3 bullet points of recommendations by top priority first."
            "No preface or introduction."
        )
        items: List[str] = []
        try:
            res = self.llm.chat.completions.create(
                model=os.getenv("CEREBRAS_MODEL", "llama3.1-8b"),
                messages=[{"role": "user", "content": prompt}],
                temperature=0.2,
            )
            content = (res.choices[0].message.content or "").strip()
            for line in content.splitlines():
                line = line.strip()
                if not line:
                    continue
                items.append(line)
        except Exception as e:
            items = [
                "Reassure the customer and summarize next steps clearly.",
                "Reference relevant policy points succinctly.",
            ]

        with open(RECOMMENDATIONS_LOG, "w", encoding="utf-8") as f:
            f.write(f"# Recommendations updated at {datetime.now().isoformat()}\n\n")
            for idx, item in enumerate(items, 1):
                f.write(f"{idx}. {item}\n")


class Orchestrator:
    def __init__(self):
        if Cerebras is None:
            raise RuntimeError("Cerebras Cloud SDK not available. Please install cerebras-cloud-sdk.")
        api_key = os.getenv("CEREBRAS_API_KEY")
        if not api_key:
            print("⚠️ CEREBRAS_API_KEY not set; attempting default client")
            self.llm = Cerebras()
        else:
            self.llm = Cerebras(api_key=api_key)

        self.speaker_agent = SpeakerIdentificationAgent(self.llm)
        self.customer_agent = CustomerMatchAgent(CRM_DIR, self.llm)
        self.reco_agent = RecommendationAgent(KB_DIR, self.llm)
        self._last_size = 0
        self._transcript_buffer: List[str] = []

    def handle_line(self, line: str):
        parsed = parse_assemblyai_line(line)
        if not parsed:
            return
        ts, speaker_label, transcript = parsed

        # 1) Speaker labeling → transcripts.log (append)
        labeled = self.speaker_agent.process(ts, speaker_label, transcript)
        with open(TRANSCRIPTS_LOG, "a", encoding="utf-8") as f:
            f.write(labeled + "\n")

        # Accumulate for CRM matching window
        self._transcript_buffer.append(transcript)
        if len(self._transcript_buffer) > 40:
            self._transcript_buffer = self._transcript_buffer[-40:]

        # 2) CRM matching → customer.log (only on change)
        changed = self.customer_agent.maybe_update_customer_log(self._transcript_buffer)
        if changed:
            print(f"🔎 Customer match updated: {changed}")

        # 3) Recommendations → overwrite recommendations.log
        self.reco_agent.update_recommendations()

    def run(self):
        # Reset outputs each run
        reset_output_logs()
        print("📡 Monitoring assemblyai.log ... (Ctrl+C to stop)")
        try:
            while True:
                lines, self._last_size = read_new_lines(ASSEMBLYAI_LOG, self._last_size)
                for ln in lines:
                    self.handle_line(ln)
                time.sleep(0.25)
        except KeyboardInterrupt:
            print("\n👋 Stopped.")


if __name__ == "__main__":
    Orchestrator().run()