import json
import random
import threading
import time
import datetime
import os
from typing import Dict, List, Any, Optional
from sentence_transformers import SentenceTransformer
import numpy as np
from transformers import pipeline
class SocialGraphManager:
"""Manages the social graph and provides context for the AAC system."""
def __init__(self, graph_path: str = "social_graph.json"):
"""Initialize the social graph manager.
Args:
graph_path: Path to the social graph JSON file
"""
self.graph_path = graph_path
self.graph = self._load_graph()
# Initialize sentence transformer for semantic matching
try:
self.sentence_model = SentenceTransformer(
"sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)
self.embeddings_cache = {}
self._initialize_embeddings()
except Exception as e:
self.sentence_model = None
def _load_graph(self) -> Dict[str, Any]:
"""Load the social graph from the JSON file."""
try:
with open(self.graph_path, "r") as f:
return json.load(f)
except Exception:
return {"people": {}, "places": [], "topics": []}
def _initialize_embeddings(self):
"""Initialize embeddings for topics and phrases in the social graph."""
if not self.sentence_model:
return
# Create embeddings for topics
topics = self.graph.get("topics", [])
for topic in topics:
if topic not in self.embeddings_cache:
self.embeddings_cache[topic] = self.sentence_model.encode(topic)
# Create embeddings for common phrases
for person_id, person_data in self.graph.get("people", {}).items():
for phrase in person_data.get("common_phrases", []):
if phrase not in self.embeddings_cache:
self.embeddings_cache[phrase] = self.sentence_model.encode(phrase)
# Create embeddings for common utterances
for category, utterances in self.graph.get("common_utterances", {}).items():
for utterance in utterances:
if utterance not in self.embeddings_cache:
self.embeddings_cache[utterance] = self.sentence_model.encode(
utterance
)
def get_people_list(self) -> List[Dict[str, str]]:
"""Get a list of people from the social graph with their names and roles."""
people = []
for person_id, person_data in self.graph.get("people", {}).items():
people.append(
{
"id": person_id,
"name": person_data.get("name", person_id),
"role": person_data.get("role", ""),
}
)
return people
def get_person_context(self, person_id: str) -> Dict[str, Any]:
"""Get context information for a specific person."""
# Check if the person_id contains a display name (e.g., "Emma (wife)")
# and try to extract the actual ID
if person_id not in self.graph.get("people", {}):
# Try to find the person by name
for pid, pdata in self.graph.get("people", {}).items():
name = pdata.get("name", "")
role = pdata.get("role", "")
if f"{name} ({role})" == person_id:
person_id = pid
break
# If still not found, return empty dict
if person_id not in self.graph.get("people", {}):
return {}
person_data = self.graph["people"][person_id]
return person_data
def get_relevant_phrases(
self, person_id: str, user_input: Optional[str] = None
) -> List[str]:
"""Get relevant phrases for a specific person based on user input."""
if person_id not in self.graph.get("people", {}):
return []
person_data = self.graph["people"][person_id]
phrases = person_data.get("common_phrases", [])
# If no user input, return random phrases
if not user_input or not self.sentence_model:
return random.sample(phrases, min(3, len(phrases)))
# Use semantic search to find relevant phrases
user_embedding = self.sentence_model.encode(user_input)
phrase_scores = []
for phrase in phrases:
if phrase in self.embeddings_cache:
phrase_embedding = self.embeddings_cache[phrase]
else:
phrase_embedding = self.sentence_model.encode(phrase)
self.embeddings_cache[phrase] = phrase_embedding
similarity = np.dot(user_embedding, phrase_embedding) / (
np.linalg.norm(user_embedding) * np.linalg.norm(phrase_embedding)
)
phrase_scores.append((phrase, similarity))
# Sort by similarity score and return top phrases
phrase_scores.sort(key=lambda x: x[1], reverse=True)
return [phrase for phrase, _ in phrase_scores[:3]]
def get_common_utterances(self, category: Optional[str] = None) -> List[str]:
"""Get common utterances from the social graph, optionally filtered by category."""
utterances = []
if "common_utterances" not in self.graph:
return utterances
if category and category in self.graph["common_utterances"]:
return self.graph["common_utterances"][category]
# If no category specified, return a sample from each category
for category_utterances in self.graph["common_utterances"].values():
utterances.extend(
random.sample(category_utterances, min(2, len(category_utterances)))
)
return utterances
def get_conversation_history(
self, person_id: str, max_conversations: int = 2
) -> List[Dict[str, Any]]:
"""Get recent conversation history for a specific person.
Args:
person_id: ID of the person to get conversation history for
max_conversations: Maximum number of recent conversations to return
Returns:
List of conversation history entries, most recent first
"""
if person_id not in self.graph.get("people", {}):
return []
person_data = self.graph["people"][person_id]
conversation_history = person_data.get("conversation_history", [])
# Sort by timestamp (most recent first)
sorted_history = sorted(
conversation_history, key=lambda x: x.get("timestamp", ""), reverse=True
)
# Return the most recent conversations
return sorted_history[:max_conversations]
def add_conversation(self, person_id: str, messages: List[Dict[str, str]]) -> bool:
"""Add a new conversation to a person's history.
Args:
person_id: ID of the person to add conversation for
messages: List of message objects with "speaker" and "text" fields
Returns:
True if successful, False otherwise
"""
if person_id not in self.graph.get("people", {}):
return False
# Create a new conversation entry
import datetime
new_conversation = {
"timestamp": datetime.datetime.now().isoformat(),
"messages": messages,
}
# Add to the person's conversation history
if "conversation_history" not in self.graph["people"][person_id]:
self.graph["people"][person_id]["conversation_history"] = []
self.graph["people"][person_id]["conversation_history"].append(new_conversation)
# Save the updated graph
return self._save_graph()
def _save_graph(self) -> bool:
"""Save the social graph to the JSON file.
Returns:
True if successful, False otherwise
"""
try:
print(f"Saving social graph to {self.graph_path}")
# Check if the file is writable
if os.path.exists(self.graph_path):
if not os.access(self.graph_path, os.W_OK):
print(f"Error: No write permission for {self.graph_path}")
return False
# Save the graph
with open(self.graph_path, "w") as f:
json.dump(self.graph, f, indent=2)
print("Social graph saved successfully")
return True
except Exception as e:
print(f"Error saving social graph: {e}")
import traceback
traceback.print_exc()
return False
def summarize_conversation(self, conversation: Dict[str, Any]) -> str:
"""Generate a summary of a conversation.
Args:
conversation: Conversation entry with timestamp and messages
Returns:
A summary string of the conversation
"""
if not conversation or "messages" not in conversation:
return "No conversation data available"
messages = conversation.get("messages", [])
if not messages:
return "No messages in conversation"
# Extract timestamp and format it
timestamp = conversation.get("timestamp", "")
try:
dt = datetime.datetime.fromisoformat(timestamp)
formatted_date = dt.strftime("%B %d, %Y at %I:%M %p")
except (ValueError, TypeError):
formatted_date = timestamp
# Create a brief summary
topic_keywords = set()
for message in messages:
# Extract potential keywords from messages
text = message.get("text", "").lower()
# Simple keyword extraction - could be improved with NLP
words = [
w
for w in text.split()
if len(w) > 4
and w
not in [
"about",
"would",
"could",
"should",
"their",
"there",
"these",
"those",
"where",
"which",
"today",
"tomorrow",
]
]
topic_keywords.update(words[:3]) # Add up to 3 keywords per message
# Limit to 5 most representative keywords
topic_keywords = list(topic_keywords)[:5]
# Create summary
first_speaker = messages[0].get("speaker", "Unknown") if messages else "Unknown"
message_count = len(messages)
summary = f"Conversation on {formatted_date}: {first_speaker} initiated a {message_count}-message conversation"
if topic_keywords:
summary += f" about {', '.join(topic_keywords)}"
return summary
class SuggestionGenerator:
"""Generates contextual suggestions for the AAC system."""
def __init__(self, model_name: str = "distilgpt2"):
"""Initialize the suggestion generator.
Args:
model_name: Name of the HuggingFace model to use
"""
self.model_name = model_name
self.model_loaded = False
self.generator = None
self.aac_user_info = None
self.loaded_models = {} # Cache for loaded models
# Load AAC user information from social graph
try:
with open("social_graph.json", "r") as f:
social_graph = json.load(f)
self.aac_user_info = social_graph.get("aac_user", {})
except Exception as e:
print(f"Error loading AAC user info from social graph: {e}")
self.aac_user_info = {}
# Try to load the model
self.load_model(model_name)
# Fallback responses if model fails to load or generate
self.fallback_responses = [
"I'm not sure how to respond to that.",
"That's interesting. Tell me more.",
"I'd like to talk about that further.",
"I appreciate you sharing that with me.",
"Could we talk about something else?",
"I need some time to think about that.",
]
def load_model(self, model_name: str) -> bool:
"""Load a model (either Hugging Face model or API-based model).
Args:
model_name: Name of the model to use (HuggingFace model name or API identifier)
Returns:
bool: True if model loaded successfully, False otherwise
"""
self.model_name = model_name
self.model_loaded = False
# Check if model is already loaded in cache
if model_name in self.loaded_models:
print(f"Using cached model: {model_name}")
self.generator = self.loaded_models[model_name]
self.model_loaded = True
return True
# Check if this is a Gemini API model
if model_name.startswith("gemini-api:"):
try:
import os
import google.generativeai as genai
# Get API key from environment
api_key = os.environ.get("GEMINI_API_KEY")
if not api_key:
print("No GEMINI_API_KEY found in environment variables.")
print("Please set the GEMINI_API_KEY environment variable.")
return False
# Configure the Gemini API
genai.configure(api_key=api_key)
# Extract the specific model name after the prefix
gemini_model = model_name.split(":", 1)[1]
print(f"Using Gemini API with model: {gemini_model}")
# Store the model name and API client in the generator
self.generator = {
"type": "gemini-api",
"model": gemini_model,
"client": genai,
}
# Cache the API client
self.loaded_models[model_name] = self.generator
self.model_loaded = True
print(f"Gemini API configured successfully for model: {gemini_model}")
return True
except Exception as e:
print(f"Error configuring Gemini API: {e}")
self.model_loaded = False
return False
# Otherwise, try to load a Hugging Face model
try:
print(f"Loading Hugging Face model: {model_name}")
# Check if this is a gated model that requires authentication
is_gated_model = any(
name in model_name.lower()
for name in ["gemma", "llama", "mistral", "qwen", "phi"]
)
if is_gated_model:
# Try to get token from environment
import os
import torch
import time
from transformers import BitsAndBytesConfig
from requests.exceptions import ConnectionError, Timeout, HTTPError
token = os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get(
"HF_TOKEN"
)
if token:
print(f"Using token for gated model: {model_name}")
from huggingface_hub import login
login(token=token, add_to_git_credential=False)
# Explicitly pass token to pipeline
from transformers import AutoTokenizer, AutoModelForCausalLM
# Implement retry mechanism for network issues
max_retries = 3
retry_delay = 2 # seconds
for attempt in range(max_retries):
try:
print(
f"Attempt {attempt+1}/{max_retries} to load model: {model_name}"
)
# First try to load just the tokenizer to check connectivity
print(f"Loading tokenizer for {model_name}...")
tokenizer = AutoTokenizer.from_pretrained(
model_name,
token=token,
use_fast=True,
local_files_only=False,
)
print(f"Tokenizer loaded successfully for {model_name}")
# Configure 4-bit quantization to save memory
print("Configuring quantization settings...")
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
)
# Load model with quantization
print(f"Loading model {model_name} with quantization...")
model = AutoModelForCausalLM.from_pretrained(
model_name,
token=token,
quantization_config=quantization_config,
device_map="auto",
low_cpu_mem_usage=True,
)
print(
f"Model {model_name} loaded successfully with quantization"
)
# Create pipeline
print("Creating text generation pipeline...")
self.generator = {
"type": "huggingface",
"pipeline": pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
torch_dtype=torch.float16,
),
}
print("Pipeline created successfully")
# If we got here, loading succeeded
break
except (ConnectionError, Timeout, HTTPError) as network_error:
# Handle network-related errors with retries
print(
f"Network error loading model (attempt {attempt+1}/{max_retries}): {network_error}"
)
if attempt < max_retries - 1:
print(f"Retrying in {retry_delay} seconds...")
time.sleep(retry_delay)
retry_delay *= 2 # Exponential backoff
else:
print(
"Maximum retries reached, falling back to alternative loading method"
)
raise network_error
except (RuntimeError, ValueError, OSError) as e:
# Handle memory errors or other issues
print(
f"Error loading gated model with token (attempt {attempt+1}/{max_retries}): {e}"
)
print(
"This may be due to memory limitations, network issues, or insufficient permissions."
)
if "CUDA out of memory" in str(
e
) or "DefaultCPUAllocator" in str(e):
print(
"Memory error detected. Trying with more aggressive memory optimization..."
)
break # Skip to non-quantized version with CPU offloading
if attempt < max_retries - 1:
print(f"Retrying in {retry_delay} seconds...")
time.sleep(retry_delay)
retry_delay *= 2 # Exponential backoff
else:
print(
"Maximum retries reached, falling back to alternative loading method"
)
# If the loop completed without success, try alternative loading methods
if not hasattr(self, "generator") or self.generator is None:
# Try loading without quantization as fallback
try:
print(
"Trying to load model without quantization (CPU only)..."
)
tokenizer = AutoTokenizer.from_pretrained(
model_name, token=token, use_fast=True
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
token=token,
device_map="cpu",
low_cpu_mem_usage=True,
)
self.generator = {
"type": "huggingface",
"pipeline": pipeline(
"text-generation", model=model, tokenizer=tokenizer
),
}
print(
"Successfully loaded model on CPU without quantization"
)
except Exception as e2:
print(f"Fallback loading also failed: {e2}")
print(
"All loading attempts failed. Please try a different model or check your connection."
)
raise RuntimeError(
f"Failed to load model after multiple attempts: {str(e2)}"
)
else:
print("No Hugging Face token found in environment variables.")
print(
"To use gated models like Gemma, you need to set up a token with the right permissions."
)
print("1. Create a token at https://huggingface.co/settings/tokens")
print(
"2. Make sure to enable 'Access to public gated repositories'"
)
print(
"3. Set it as an environment variable: export HUGGING_FACE_HUB_TOKEN=your_token_here"
)
raise ValueError("Authentication token required for gated model")
else:
# For non-gated models, use the standard pipeline
from transformers import pipeline
self.generator = {
"type": "huggingface",
"pipeline": pipeline("text-generation", model=model_name),
}
# Cache the loaded model
self.loaded_models[model_name] = self.generator
self.model_loaded = True
print(f"Model loaded successfully: {model_name}")
return True
except Exception as e:
print(f"Error loading model: {e}")
self.model_loaded = False
return False
def _clean_small_model_response(self, response: str) -> str:
"""Clean up responses from small models that often repeat instructions or generate nonsense.
Args:
response: The raw response from the model
Returns:
A cleaned response
"""
# If response is too short, return as is
if len(response) < 5:
return response
# Remove common instruction repetitions
patterns_to_remove = [
"I want to respond to what",
"I'll use language appropriate for our relationship",
"I should speak in first person",
"I should use language appropriate",
"I want to respond directly",
"I'll speak as myself",
"I want to initiate a conversation",
"My response should be natural",
"My response to",
"Will's response to",
"Will says to",
]
# Check for and remove these patterns
cleaned_response = response
for pattern in patterns_to_remove:
if pattern in cleaned_response:
# Find the first occurrence and remove everything from there
index = cleaned_response.find(pattern)
if index > 10: # Keep some beginning text if available
cleaned_response = cleaned_response[:index].strip()
else:
# If pattern is at the beginning, remove just that pattern
parts = cleaned_response.split(pattern, 1)
if len(parts) > 1:
cleaned_response = parts[1].strip()
# Remove any lines that are just the name repeated
lines = cleaned_response.split("\n")
cleaned_lines = []
for line in lines:
# Skip lines that are just a name repeated
if line.strip() and not all(
word == line.split()[0] for word in line.split()
):
cleaned_lines.append(line)
cleaned_response = "\n".join(cleaned_lines).strip()
# If we've removed too much, use a fallback
if len(cleaned_response) < 5:
return "I'm not sure what to say about that."
# Limit to first 2 sentences to avoid rambling
sentences = cleaned_response.split(".")
if len(sentences) > 2:
cleaned_response = ".".join(sentences[:2]) + "."
return cleaned_response
def _get_mood_description(self, mood_value: int) -> str:
"""Convert mood value (1-5) to a descriptive string.
Args:
mood_value: Integer from 1-5 representing mood (1=sad, 5=happy)
Returns:
String description of the mood
"""
mood_descriptions = {
1: "I'm feeling quite down and sad today. My responses might be more subdued.",
2: "I'm feeling a bit low today. I might be less enthusiastic than usual.",
3: "I'm feeling okay today - neither particularly happy nor sad.",
4: "I'm feeling pretty good today. I'm in a positive mood.",
5: "I'm feeling really happy and upbeat today! I'm in a great mood.",
}
# Default to neutral if value is out of range
return mood_descriptions.get(mood_value, mood_descriptions[3])
def test_model(self) -> str:
"""Test if the model is working correctly."""
if not self.model_loaded:
return "Model not loaded"
try:
# Create a more explicit test prompt that clearly establishes Will's identity and role
test_prompt = """I am Will, a 38-year-old with MND (Motor Neuron Disease).
I am talking to my 7-year-old son Billy.
Billy just asked me about football.
I want to respond to Billy in a natural, brief way.
My response to Billy:"""
print(f"Testing model with prompt: {test_prompt}")
# Check if we're using the Gemini API or a Hugging Face model
if (
isinstance(self.generator, dict)
and self.generator.get("type") == "gemini-api"
):
try:
# Use Gemini API
genai = self.generator["client"]
model_name = self.generator["model"]
# Create a generative model
model = genai.GenerativeModel(model_name)
# Generate content with timeout
print("Sending test request to Gemini API...")
# Set a timeout for the test
import threading
import time
result = ["No response received yet"]
generation_complete = [False]
def generate_with_timeout():
try:
print("Starting Gemini API test request...")
response = model.generate_content(test_prompt)
print(f"Received response from Gemini API: {response}")
if response and hasattr(response, "text"):
result[0] = response.text
print(f"Extracted text from response: {result[0]}")
else:
result[0] = "No text in Gemini API response"
print("Response object has no text attribute")
generation_complete[0] = True
except Exception as e:
print(f"Error in Gemini test generation: {e}")
result[0] = f"Error: {str(e)}"
generation_complete[0] = True
# Start generation in a separate thread
generation_thread = threading.Thread(target=generate_with_timeout)
generation_thread.daemon = True
generation_thread.start()
# Wait for up to 10 seconds
timeout = 10
start_time = time.time()
while (
not generation_complete[0]
and time.time() - start_time < timeout
):
print(
f"Waiting for Gemini API response... ({int(time.time() - start_time)}s)"
)
time.sleep(1)
if not generation_complete[0]:
print("Gemini API test request timed out")
return "Gemini API test timed out after 10 seconds"
print(f"Test response from Gemini API: {result[0]}")
return f"Gemini API test successful: {result[0]}"
except Exception as e:
print(f"Error testing Gemini API: {e}")
return f"Gemini API test failed: {str(e)}"
elif (
isinstance(self.generator, dict)
and self.generator.get("type") == "huggingface"
):
# Use Hugging Face pipeline
pipeline = self.generator["pipeline"]
response = pipeline(test_prompt, max_new_tokens=30, do_sample=True)
full_text = response[0]["generated_text"]
if len(test_prompt) < len(full_text):
result = full_text[len(test_prompt) :].strip()
# Check if this is a small model that needs cleaning
is_small_model = any(
name in self.model_name.lower()
for name in ["distilgpt2", "gpt2-small", "tiny"]
)
if is_small_model:
result = self._clean_small_model_response(result)
else:
result = "No additional text generated"
print(f"Test response from Hugging Face: {result}")
return f"Hugging Face model test successful: {result}"
else:
# Legacy format (for backward compatibility)
response = self.generator(
test_prompt, max_new_tokens=30, do_sample=True
)
full_text = response[0]["generated_text"]
if len(test_prompt) < len(full_text):
result = full_text[len(test_prompt) :].strip()
# Check if this is a small model that needs cleaning
is_small_model = any(
name in self.model_name.lower()
for name in ["distilgpt2", "gpt2-small", "tiny"]
)
if is_small_model:
result = self._clean_small_model_response(result)
else:
result = "No additional text generated"
print(f"Test response: {result}")
return f"Model test successful: {result}"
except Exception as e:
print(f"Error testing model: {e}")
return f"Model test failed: {str(e)}"
def generate_suggestion(
self,
person_context: Dict[str, Any],
user_input: Optional[str] = None,
max_length: int = 50,
temperature: float = 0.7,
) -> str:
"""Generate a contextually appropriate suggestion.
Args:
person_context: Context information about the person
user_input: Optional user input to consider
max_length: Maximum length of the generated suggestion
temperature: Controls randomness in generation (higher = more random)
Returns:
A generated suggestion string
"""
if not self.model_loaded:
# Use fallback responses if model isn't loaded
import random
print("Model not loaded, using fallback responses")
return random.choice(self.fallback_responses)
# Extract context information
name = person_context.get("name", "")
role = person_context.get("role", "")
topics = person_context.get("topics", [])
context = person_context.get("context", "")
selected_topic = person_context.get("selected_topic", "")
common_phrases = person_context.get("common_phrases", [])
frequency = person_context.get("frequency", "")
mood = person_context.get("mood", 3) # Default to neutral mood (3)
# Get AAC user information
aac_user = self.aac_user_info
# Build enhanced prompt
prompt = f"""I am {aac_user.get('name', 'Will')}, a {aac_user.get('age', 38)}-year-old with MND (Motor Neuron Disease) from {aac_user.get('location', 'Manchester')}.
{aac_user.get('background', '')}
My communication needs: {aac_user.get('communication_needs', '')}
I am talking to {name}, who is my {role}.
About {name}: {context}
We typically talk about: {', '.join(topics)}
We communicate {frequency}.
My current mood: {self._get_mood_description(mood)}
"""
# Add communication style based on relationship
if role in ["wife", "son", "daughter", "mother", "father"]:
prompt += "I communicate with my family in a warm, loving way, sometimes using inside jokes.\n"
elif role in ["doctor", "therapist", "nurse"]:
prompt += "I communicate with healthcare providers in a direct, informative way.\n"
elif role in ["best mate", "friend"]:
prompt += "I communicate with friends casually, often with humor and sometimes swearing.\n"
elif role in ["work colleague", "boss"]:
prompt += (
"I communicate with colleagues professionally but still friendly.\n"
)
# Add topic information if provided
if selected_topic:
prompt += f"\nWe are currently discussing {selected_topic}.\n"
# Add specific context about this topic with this person
if selected_topic == "football" and "Manchester United" in context:
prompt += "We both support Manchester United and often discuss recent matches.\n"
elif selected_topic == "programming" and "software developer" in context:
prompt += "We both work in software development and share technical interests.\n"
elif selected_topic == "family plans" and role in ["wife", "husband"]:
prompt += (
"We make family decisions together, considering my condition.\n"
)
elif selected_topic == "old scout adventures" and role == "best mate":
prompt += "We often reminisce about our Scout camping trips in South East London.\n"
elif selected_topic == "cycling" and "cycling" in context:
prompt += "I miss being able to cycle but enjoy talking about past cycling adventures.\n"
# Add the user's message if provided, or set up for conversation initiation
if user_input:
# If user input is provided, we're responding to something
prompt += f'\n{name} just said to me: "{user_input}"\n'
prompt += f"I want to respond directly to what {name} just said.\n"
else:
# No user input means we're initiating a conversation
if selected_topic:
# If a topic is selected, initiate conversation about that topic
prompt += f"\nI'm about to start a conversation with {name} about {selected_topic}.\n"
# Add specific context about initiating this topic with this person
if selected_topic == "football" and "Manchester United" in context:
prompt += (
"We both support Manchester United and often discuss matches.\n"
)
elif selected_topic == "family" and role in [
"wife",
"husband",
"son",
"daughter",
]:
prompt += (
"I want to check in about our family plans or activities.\n"
)
elif selected_topic == "health" and role in [
"doctor",
"nurse",
"therapist",
]:
prompt += "I want to discuss my health condition or symptoms.\n"
elif selected_topic == "work" and role in ["work colleague", "boss"]:
prompt += "I want to discuss a work-related matter.\n"
prompt += f"I want to initiate a conversation about {selected_topic} in a natural way.\n"
elif common_phrases:
# Use context about our typical conversations if no specific topic
prompt += f"\nI'm about to start a conversation with {name}.\n"
default_message = common_phrases[0]
prompt += f'{name} typically says things like: "{default_message}"\n'
prompt += f"We typically talk about: {', '.join(topics)}\n"
prompt += "I want to initiate a conversation in a natural way based on our relationship.\n"
else:
# Generic conversation starter
prompt += f"\nI'm about to start a conversation with {name}.\n"
prompt += "I want to initiate a conversation in a natural way based on our relationship.\n"
# Add the response prompt with specific guidance
# Check if this is an instruction-tuned model
is_instruction_model = any(
marker in self.model_name.lower()
for marker in ["-it", "instruct", "chat", "phi-3", "phi-2"]
)
# Check if this is a very small model that needs simpler prompts
is_small_model = any(
name in self.model_name.lower()
for name in ["distilgpt2", "gpt2-small", "tiny"]
)
if is_small_model:
# Use a much simpler format for very small models
if user_input:
# Responding to something
prompt += f"""
{name} said: "{user_input}"
Will's response:"""
else:
# Initiating a conversation
if selected_topic:
prompt += f"""
Will starts a conversation with {name} about {selected_topic}.
Will says:"""
else:
prompt += f"""
Will starts a conversation with {name}.
Will says:"""
elif is_instruction_model:
# Use instruction format for instruction-tuned models
if user_input:
# Responding to something
prompt += f"""
I am Will, the person with MND. I need to respond to {name}'s message: "{user_input}"
My response should be natural, brief (1-2 sentences), and directly relevant to what {name} just said.
I should use language appropriate for our relationship.
I should speak in first person as myself (Will).
My response to {name}:"""
else:
# Initiating a conversation
prompt += f"""
I am Will, the person with MND. I need to start a conversation with {name}.
My conversation starter should be natural, brief (1-2 sentences), and appropriate for our relationship.
If a topic was selected, I should focus on that topic.
I should speak in first person as myself (Will).
My conversation starter to {name}:"""
else:
# Use standard format for other models
if user_input:
# Responding to something
prompt += f"""
I am Will, the person with MND. I want to respond to {name}'s message: "{user_input}"
My response should be natural, brief (1-2 sentences), and directly relevant to what {name} just said.
I'll use language appropriate for our relationship and speak as myself (Will).
My response to {name}:"""
else:
# Initiating a conversation
prompt += f"""
I am Will, the person with MND. I want to start a conversation with {name}.
My conversation starter should be natural, brief (1-2 sentences), and appropriate for our relationship.
I'll speak in first person as myself (Will).
My conversation starter to {name}:"""
# Generate suggestion
try:
print(f"Generating suggestion with prompt: {prompt}")
# Check if we're using the Gemini API or a Hugging Face model
if (
isinstance(self.generator, dict)
and self.generator.get("type") == "gemini-api"
):
try:
# Use Gemini API
try:
genai = self.generator["client"]
model_name = self.generator["model"]
# Create a generative model
model = genai.GenerativeModel(model_name)
# Set generation config
generation_config = {
"temperature": temperature,
"top_p": 0.92,
"top_k": 50,
"max_output_tokens": 100,
}
# Generate content with timeout
result = [
"I'm thinking about what to say..."
] # Default response
generation_complete = [False]
def generate_with_gemini():
try:
response = model.generate_content(
prompt, generation_config=generation_config
)
if response and hasattr(response, "text"):
result[0] = response.text.strip()
print(f"Gemini API response: {result[0]}")
else:
print("No response from Gemini API")
generation_complete[0] = True
except Exception as e:
print(f"Error in Gemini generation thread: {e}")
generation_complete[0] = True
# Start generation in a separate thread
generation_thread = threading.Thread(
target=generate_with_gemini
)
generation_thread.daemon = True
generation_thread.start()
# Wait for up to 10 seconds
timeout = 10
start_time = time.time()
while (
not generation_complete[0]
and time.time() - start_time < timeout
):
time.sleep(0.1)
if not generation_complete[0]:
print("Gemini API request timed out")
return "I'm thinking about what to say... (API timeout)"
return result[0]
except Exception as e:
print(f"Error setting up Gemini API: {e}")
return (
"I'm having trouble connecting to the Gemini API right now."
)
except Exception as e:
print(f"Error generating with Gemini API: {e}")
return "Could not generate a suggestion with Gemini API. Please try again."
elif (
isinstance(self.generator, dict)
and self.generator.get("type") == "huggingface"
):
# Use Hugging Face pipeline
pipeline = self.generator["pipeline"]
# Generate with Hugging Face
response = pipeline(
prompt,
max_new_tokens=100, # Generate more tokens to ensure we get a response
temperature=temperature,
do_sample=True,
top_p=0.92,
top_k=50,
truncation=False,
)
# Extract only the generated part, not the prompt
full_text = response[0]["generated_text"]
print(f"Full generated text length: {len(full_text)}")
print(f"Prompt length: {len(prompt)}")
# Make sure we're not trying to slice beyond the text length
if len(prompt) < len(full_text):
result = full_text[len(prompt) :].strip()
# Post-process the result for small models
if is_small_model:
result = self._clean_small_model_response(result)
print(f"Generated response: {result}")
return result
else:
# If the model didn't generate anything beyond the prompt
print("Model didn't generate text beyond prompt")
return "I'm thinking about what to say..."
else:
# Legacy format (for backward compatibility)
response = self.generator(
prompt,
max_new_tokens=100,
temperature=temperature,
do_sample=True,
top_p=0.92,
top_k=50,
truncation=False,
)
# Extract only the generated part, not the prompt
full_text = response[0]["generated_text"]
print(f"Full generated text length: {len(full_text)}")
print(f"Prompt length: {len(prompt)}")
# Make sure we're not trying to slice beyond the text length
if len(prompt) < len(full_text):
result = full_text[len(prompt) :].strip()
# Post-process the result for small models
if is_small_model:
result = self._clean_small_model_response(result)
print(f"Generated response: {result}")
return result
else:
# If the model didn't generate anything beyond the prompt
print("Model didn't generate text beyond prompt")
return "I'm thinking about what to say..."
except Exception as e:
print(f"Error generating suggestion: {e}")
return "Could not generate a suggestion. Please try again."