helo-ayush/Diarization_VoiceFingerprinted

1# ==============================================================================
2# GEMINI LLM PROCESSOR
3# Handles sentiment analysis, technical summary generation, and formatting.
4# ==============================================================================
5import os
6import json
7import time
8import asyncio
9import traceback
10from langchain_google_genai import ChatGoogleGenerativeAI
11from langchain_core.messages import HumanMessage
12
13
14# Initialize LangChain Gemini model using API key from environment
15# Temperature is set to 0.1 to force deterministic, predictable output (less hallucination)
16llm = ChatGoogleGenerativeAI(
17    model="gemini-2.5-flash",
18    google_api_key=os.getenv("GEMINI_API_KEY"),
19    temperature=0.1,
20)
21
22PROMPT_TEMPLATE = """You are a Senior Technical Dialogue Editor and Data Analyst. Your job is to format a highly-accurate diarized transcript and generate metadata for a RAG search system.
23
24═══════════════════════════════════
25TRANSCRIPT (highly accurate):
26{raw_transcript}
27═══════════════════════════════════
28
29IMPORTANT CONTEXT: This transcript is generated by an upgraded, highly accurate ML model. Your job is ONLY to apply very minimal formatting. DO NOT over-change any lines or skip any lines. 
30
31YOUR WORKFLOW & RULES:
32
331. **Role Analysis & Labeling:**
34{pynote_context}
35   - **CRITICAL:** If a speaker is ALREADY named (e.g., "Agent Ayush" or "Ayush" or "Customer"), their role has already been explicitly identified. Keep their explicit name exactly as it appears. 
36   - **CRITICAL**: When returning the `refinedTranscript`, the `role` field MUST be exactly the speaker's name if they have one (e.g. "Ayush"). Otherwise use "Agent" or "Client".
37
382. **Vector-Optimized Summary:**
39   - Write a dense summary mentioning all technical entities (app names, features, error messages).
40   - Include what the issue was and how it was resolved.
41
423. **Smart & Light Transcript Cleanup:**
43   - **DO NOT drop ANY lines.** Every single line of dialogue must be kept to retain the audio context.
44   - **Correct phonetic & contextual mistranscriptions.** You MUST fix words that were transcribed poorly but clearly refer to IT/support terms based on how they sound (e.g., change "index" to "AnyDesk", "PV" to "PC", etc.).
45   - **Avoid heavy rewriting.** Fix the wrong words to make sentences contextually correct for tech support, but do not completely paraphrase the sentence structure.
46   - **Fix speaker assignments** only where logically backward.
47   - **Merge fragments** when a speaker's thought is split across consecutive lines.
48
494. **Formatting & Numbers:**
50   - Convert any phonetically spelled-out numbers (e.g., 'pachchis', 'chhabbees', 'twenty five') back into numeric digits (25, 26, 25).
51   - Keep technical English words exactly as they are.
52
535. **Satisfaction Score:**
54   - Rate the Client's satisfaction from 1-10.
55
56STRICT JSON OUTPUT FORMAT (return ONLY this JSON, nothing else):
57{{
58  "summary": "Dense technical summary for embeddings...",
59  "satisfactionScore": 10,
60  "detectedRoles": {{ "speaker0": "Agent", "speaker1": "Client" }},
61  "tags": ["ERP", "Technical Support"], 
62  "refinedTranscript": [
63      {{ "role": "Ayush", "text": "Ultra-lightly cleaned Hinglish text..." }}
64  ]
65}}
66
67Speaker count is {speaker_count}. Return ONLY the JSON object."""
68
69
70async def refine_transcript(raw_transcript: str, speaker_count: int, pynnote_diarized: bool = False) -> dict:
71    """
72    Refine a raw diarized transcript using LangChain + Gemini.
73    """
74    print("   📝 Building Gemini prompt...")
75    
76    # Modify the prompt rules slightly if Pyannote already identified the agent
77    if pynnote_diarized:
78        pynote_context = (
79            "   - **CRITICAL:** Our Pyannote voice fingerprinting model has ALREADY identified the Agent's name in the transcript. "
80            "You DO NOT need to guess or figure out who the Agent is vs the Customer. "
81            "Assume the named person is the Agent, and any generic 'Speaker X' is the Customer. "
82            "When returning the `refinedTranscript`, the `role` field MUST be exactly the agent's actual name (e.g. 'Ayush'). DO NOT just return 'Agent'."
83        )
84    else:
85        pynote_context = (
86            "   - If the transcript simply says 'Speaker 0' or 'Speaker X', determine who is the 'Agent' (Provider) "
87            "and who is the 'Client' (User) based on the context of the conversation and their dialogue."
88        )
89
90    prompt = PROMPT_TEMPLATE.format(
91        raw_transcript=raw_transcript,
92        speaker_count=speaker_count,
93        pynote_context=pynote_context
94    )
95    print(f"   📝 Prompt length: {len(prompt)} chars")
96
97    try:
98        start_time = time.time()
99
100        # Step 1: Call Gemini via LangChain (use thread + timeout to prevent hanging)
101        print("   🔄 Calling Gemini via LangChain...")
102        message = HumanMessage(content=prompt)
103        response = await asyncio.to_thread(llm.invoke, [message])
104        
105        elapsed = int((time.time() - start_time) * 1000)
106        print(f"   ✅ Gemini responded in {elapsed}ms")
107
108        # Step 2: Extract text
109        text = response.content.strip()
110        print(f"   📄 Response length: {len(text)} chars")
111        print(f"   📄 First 200 chars: {text[:200]}")
112
113        # Step 3: Clean markdown fences
114        if text.startswith("```"):
115            lines = text.split("\n")
116            # Remove first line (```json) and last line (```)
117            text = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])
118            text = text.strip()
119            print("   🧹 Stripped markdown code fences")
120
121        # Step 4: Parse JSON
122        print("   🔍 Parsing JSON...")
123        data = json.loads(text)
124        print(f"   ✅ JSON parsed successfully — keys: {list(data.keys())}")
125
126        # Step 5: Build transcript string
127        transcript_parts = data.get("refinedTranscript", [])
128        transcript_string = "\n".join(
129            f"{s['role']}: {s['text']}" for s in transcript_parts
130        )
131
132        print(f"   🔧 Gemini produced {len(transcript_parts)} refined segments")
133
134        return {
135            "summary": data.get("summary", ""),
136            "satisfaction_score": data.get("satisfactionScore", 0),
137            "tags": data.get("tags", []),
138            "detected_roles": data.get("detectedRoles", {}),
139            "transcript": transcript_string,
140            "processing_time": elapsed,
141        }
142
143    except json.JSONDecodeError as e:
144        print(f"   ❌ JSON parse error: {e}")
145        print(f"   ❌ Raw text was: {text[:500] if 'text' in dir() else 'N/A'}")
146        return _error_result()
147    except Exception as err:
148        print(f"   ❌ Gemini processing failed: {err}")
149        traceback.print_exc()
150        return _error_result()
151
152
153def _error_result():
154    return {
155        "transcript": None,
156        "summary": "Could not generate summary",
157        "satisfaction_score": 0,
158        "tags": [],
159        "detected_roles": {},
160        "processing_time": 0,
161    }
162