communication_agent.py

"""
Communication Agent - Gemini API Version (Stable)
Backup GPT-2 version available in communication_agent_gpt2.py
"""
import google.generativeai as genai
from typing import List, Dict
import os
import config

class CommunicationAgent:
    def __init__(self, api_key: str = None):
        """Initialize Communication Agent with Gemini API"""
        self.api_key = api_key or os.getenv("GEMINI_API_KEY", "")
        
        if self.api_key:
            genai.configure(api_key=self.api_key)
            self.model = genai.GenerativeModel(config.GEMINI_MODEL)
            self.enabled = True
            print("✅ Communication Agent (Gemini) ready!")
        else:
            self.model = None
            self.enabled = False
            print("⚠️ Warning: GEMINI_API_KEY not set. Please add it to .env file")
        
        # System context - Clear instructions for better responses
        self.system_context = (
            "You are a helpful, friendly AI assistant. "
            "Respond naturally and conversationally. "
            "Keep responses concise (2-3 sentences). "
            "Be warm and engaging. "
            "If you don't understand, ask for clarification politely."
        )
    
    def generate_response(
        self, 
        user_message: str, 
        conversation_history: List[Dict] = None,
        temperature: float = 0.7
    ) -> str:
        """
        Generate a conversational response using Gemini
        
        Args:
            user_message: User's input message
            conversation_history: Previous conversation context
            temperature: Response creativity (0.0-1.0)
            
        Returns:
            Generated response text
        """
        if not self.enabled:
            return "⚠️ Gemini API not configured. Please add GEMINI_API_KEY to .env file."
        
        try:
            # Build context from history
            context = self._build_context(conversation_history)
            
            # Create prompt with system context
            prompt = f"""{self.system_context}

{context}

User: {user_message}
Assistant:"""
            
            # Generate response with Gemini
            generation_config = {
                'temperature': temperature,
                'top_p': 0.95,
                'top_k': 40,
                'max_output_tokens': 200,
            }
            
            response = self.model.generate_content(
                prompt,
                generation_config=generation_config
            )
            
            return response.text.strip()
            
        except Exception as e:
            return f"Sorry, an error occurred: {str(e)}"
    
    def _build_context(self, conversation_history: List[Dict] = None) -> str:
        """Build conversation context from history"""
        if not conversation_history:
            return ""
        
        context = "Previous conversation:\n"
        # Use last 5 messages for context
        for msg in conversation_history[-5:]:
            role = "User" if msg.get('role') == 'user' else "Assistant"
            content = msg.get('content', '')
            context += f"{role}: {content}\n"
        
        return context