Fix ChatPromptTemplate variable parsing issue and update tests
- Replace ChatPromptTemplate with direct HumanMessage/SystemMessage to avoid template variable parsing - Fix f-string formatting issues in prompt strings - Update test_general_summary_fallback_on_error to properly mock fallback chain - Add tests directory with comprehensive test coverage
This commit is contained in:
+4
-1
@@ -25,4 +25,7 @@ moviepy
|
||||
yt-dlp
|
||||
ffmpeg-python
|
||||
reportlab
|
||||
anthropic
|
||||
anthropic
|
||||
pytest
|
||||
pytest-mock
|
||||
langchain-anthropic
|
||||
+222
-31
@@ -1,61 +1,252 @@
|
||||
import anthropic
|
||||
import os
|
||||
import logging
|
||||
from typing import Dict, List, Any, Optional
|
||||
from dotenv import load_dotenv
|
||||
import json
|
||||
from pydantic import BaseModel
|
||||
from langchain_anthropic import ChatAnthropic
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from src.prompt import advanced_summary_prompt, basic_summary_prompt, custom_template_prompt
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Setup logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ============================================================================
|
||||
# Pydantic Models for Structured Outputs
|
||||
# ============================================================================
|
||||
|
||||
# Basic Summary Models (Freemium Plan)
|
||||
class KeyPoint(BaseModel):
|
||||
"""A key point from the meeting"""
|
||||
text: str
|
||||
timestamp: float
|
||||
|
||||
class Summary(BaseModel):
|
||||
"""Overall summary of the meeting"""
|
||||
text: str
|
||||
duration_minutes: float
|
||||
|
||||
class BasicSummary(BaseModel):
|
||||
"""Basic summary structure for freemium plan"""
|
||||
Key_Points: List[KeyPoint]
|
||||
Summary: Summary
|
||||
|
||||
|
||||
# Advanced Summary Models (Pro Plan)
|
||||
class Purpose(BaseModel):
|
||||
"""Purpose of the meeting"""
|
||||
text: str
|
||||
|
||||
class ChapterContent(BaseModel):
|
||||
"""Content item within a chapter"""
|
||||
text: str
|
||||
original_transcript_start: float
|
||||
original_transcript_end: float
|
||||
|
||||
class WordTimestamp(BaseModel):
|
||||
"""Word-level timestamp"""
|
||||
word: str
|
||||
timestamp: float
|
||||
|
||||
class TimeStamp(BaseModel):
|
||||
"""Time range"""
|
||||
start: float
|
||||
end: float
|
||||
|
||||
class Chapter(BaseModel):
|
||||
"""A chapter in the meeting"""
|
||||
chapter: str
|
||||
time_stamp: TimeStamp
|
||||
content: List[ChapterContent]
|
||||
words_time_stamp: List[WordTimestamp]
|
||||
|
||||
class Chapters(BaseModel):
|
||||
"""Chapters section"""
|
||||
minutes_total: float
|
||||
content: List[Chapter]
|
||||
|
||||
class OutcomeContent(BaseModel):
|
||||
"""Content item in outcomes"""
|
||||
text: str
|
||||
time_stamp: TimeStamp
|
||||
words_time_stamp: List[WordTimestamp]
|
||||
|
||||
class Outcomes(BaseModel):
|
||||
"""Outcomes section"""
|
||||
minutes_total: float
|
||||
content: List[OutcomeContent]
|
||||
|
||||
class ActionItem(BaseModel):
|
||||
"""An action item"""
|
||||
text: str
|
||||
time_stamp: TimeStamp
|
||||
words_time_stamp: List[WordTimestamp]
|
||||
|
||||
class ActionItemsPerUser(BaseModel):
|
||||
"""Action items for a specific user"""
|
||||
speaker: str
|
||||
minutes_total: float
|
||||
action_items: List[ActionItem]
|
||||
|
||||
class AdvancedSummary(BaseModel):
|
||||
"""Advanced summary structure for pro plan"""
|
||||
Purpose: Purpose
|
||||
Chapters: Chapters
|
||||
Outcomes: Outcomes
|
||||
Action_Items_Per_User: List[ActionItemsPerUser]
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Summary Generation Functions
|
||||
# ============================================================================
|
||||
|
||||
def general_summary(transcription, plan_tier="pro"):
|
||||
"""
|
||||
Generate a summary of the transcription based on the user's plan tier.
|
||||
Uses LangChain Anthropic with structured outputs.
|
||||
|
||||
Args:
|
||||
transcription: The transcription to summarize
|
||||
transcription: The transcription to summarize (dict or JSON string)
|
||||
plan_tier: The user's plan tier ("freemium" or "pro")
|
||||
|
||||
Returns:
|
||||
A JSON object containing the summary
|
||||
A dict containing the summary (structured output)
|
||||
"""
|
||||
client = anthropic.Anthropic(
|
||||
api_key=os.getenv("ANTHTROPIC_API_KEY"),
|
||||
)
|
||||
# Get API key (note: original code had typo ANTHTROPIC_API_KEY)
|
||||
api_key = os.getenv("ANTHROPIC_API_KEY") or os.getenv("ANTHTROPIC_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("ANTHROPIC_API_KEY environment variable is required")
|
||||
|
||||
logger.info(f"Generating {plan_tier} summary with structured output")
|
||||
|
||||
# Select the appropriate prompt based on the user's plan tier
|
||||
# Convert transcription to string if it's a dict
|
||||
if isinstance(transcription, dict):
|
||||
transcription_str = json.dumps(transcription)
|
||||
else:
|
||||
transcription_str = str(transcription)
|
||||
|
||||
# Select the appropriate prompt and schema based on the user's plan tier
|
||||
if plan_tier.lower() == "freemium":
|
||||
prompt = basic_summary_prompt
|
||||
max_tokens = 2000 # Reduced token count for basic summaries
|
||||
max_tokens = 2000
|
||||
output_schema = BasicSummary
|
||||
else: # Default to pro
|
||||
prompt = advanced_summary_prompt
|
||||
max_tokens = 4000
|
||||
output_schema = AdvancedSummary
|
||||
|
||||
message = client.messages.create(
|
||||
model="claude-3-5-sonnet-20241022",
|
||||
max_tokens=max_tokens,
|
||||
messages=[
|
||||
{"role": "user", "content": f"{prompt}"},
|
||||
{"role": "user", "content": f"Transcription: {transcription}"}
|
||||
]
|
||||
# Initialize LangChain Anthropic model
|
||||
model = ChatAnthropic(
|
||||
model="claude-sonnet-4-5-20250929",
|
||||
api_key=api_key,
|
||||
temperature=0.2,
|
||||
max_tokens=max_tokens
|
||||
)
|
||||
|
||||
text = message.content[0].text
|
||||
return json.loads(text)
|
||||
|
||||
# Create messages directly to avoid template variable parsing issues
|
||||
messages = [
|
||||
SystemMessage(content="You are an AI meeting transcript summary formatter. Follow the instructions carefully and return structured output."),
|
||||
HumanMessage(content=prompt + "\n\nTranscription: " + transcription_str)
|
||||
]
|
||||
|
||||
# Use structured output
|
||||
structured_model = model.with_structured_output(output_schema)
|
||||
|
||||
try:
|
||||
# Invoke the structured model with messages
|
||||
result = structured_model.invoke(messages)
|
||||
|
||||
# Convert Pydantic model to dict
|
||||
if isinstance(result, BaseModel):
|
||||
logger.info(f"Successfully generated {plan_tier} summary with structured output")
|
||||
return result.model_dump()
|
||||
else:
|
||||
logger.info(f"Successfully generated {plan_tier} summary")
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
# Log error and return fallback
|
||||
logger.error(f"Error generating summary with structured output: {e}")
|
||||
# Fallback: try without structured output
|
||||
try:
|
||||
logger.warning("Falling back to non-structured output")
|
||||
response = model.invoke(messages)
|
||||
text = response.content if hasattr(response, 'content') else str(response)
|
||||
return json.loads(text)
|
||||
except Exception as fallback_error:
|
||||
logger.error(f"Fallback also failed: {fallback_error}")
|
||||
raise
|
||||
|
||||
|
||||
def custom_summary(template, transcription):
|
||||
client = anthropic.Anthropic(
|
||||
api_key=os.getenv("ANTHTROPIC_API_KEY"),
|
||||
)
|
||||
message = client.messages.create(
|
||||
model="claude-3-5-sonnet-20241022",
|
||||
max_tokens=8000,
|
||||
messages=[
|
||||
{"role": "user", "content": f"{custom_template_prompt}"},
|
||||
{"role": "user", "content": f"TEMPLATE : {template}"},
|
||||
{"role": "user", "content": f"Transcription: {transcription}"}
|
||||
]
|
||||
"""
|
||||
Generate a custom summary based on a user-defined template.
|
||||
Uses LangChain Anthropic.
|
||||
|
||||
Args:
|
||||
template: The custom template (dict or JSON string)
|
||||
transcription: The transcription to summarize (dict or JSON string)
|
||||
|
||||
Returns:
|
||||
A dict containing the custom summary
|
||||
"""
|
||||
# Get API key (note: original code had typo ANTHTROPIC_API_KEY)
|
||||
api_key = os.getenv("ANTHROPIC_API_KEY") or os.getenv("ANTHTROPIC_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("ANTHROPIC_API_KEY environment variable is required")
|
||||
|
||||
logger.info("Generating custom summary")
|
||||
|
||||
# Convert to strings if needed
|
||||
if isinstance(template, dict):
|
||||
template_str = json.dumps(template)
|
||||
else:
|
||||
template_str = str(template)
|
||||
|
||||
if isinstance(transcription, dict):
|
||||
transcription_str = json.dumps(transcription)
|
||||
else:
|
||||
transcription_str = str(transcription)
|
||||
|
||||
# Initialize LangChain Anthropic model
|
||||
model = ChatAnthropic(
|
||||
model="claude-sonnet-4-5-20250929", # Using the same model as general_summary
|
||||
api_key=api_key,
|
||||
temperature=0.2,
|
||||
max_tokens=8000
|
||||
)
|
||||
|
||||
text = message.content[0].text
|
||||
return json.loads(text)
|
||||
# Create messages directly to avoid template variable parsing issues
|
||||
messages = [
|
||||
SystemMessage(content="You are an AI meeting transcript summary formatter. Follow the user-defined template structure exactly."),
|
||||
HumanMessage(content=custom_template_prompt + "\n\nTEMPLATE: " + template_str + "\n\nTranscription: " + transcription_str)
|
||||
]
|
||||
|
||||
try:
|
||||
response = model.invoke(messages)
|
||||
text = response.content if hasattr(response, 'content') else str(response)
|
||||
|
||||
# Try to parse as JSON
|
||||
try:
|
||||
return json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
# If it's wrapped in markdown code blocks, try to extract JSON
|
||||
if "```json" in text:
|
||||
json_start = text.find("```json") + 7
|
||||
json_end = text.find("```", json_start)
|
||||
text = text[json_start:json_end].strip()
|
||||
return json.loads(text)
|
||||
elif "```" in text:
|
||||
json_start = text.find("```") + 3
|
||||
json_end = text.find("```", json_start)
|
||||
text = text[json_start:json_end].strip()
|
||||
return json.loads(text)
|
||||
else:
|
||||
raise ValueError(f"Could not parse response as JSON: {text[:200]}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating custom summary: {e}")
|
||||
raise
|
||||
|
||||
+9
-121
@@ -28,111 +28,13 @@ At the end of each section, include a field named "minutes_total" which represen
|
||||
- Ensure that for each sentence you generate, every word in that sentence is assigned the same timestamp—the start timestamp of that sentence.
|
||||
**Example Output JSON:**
|
||||
|
||||
{
|
||||
"Purpose": {
|
||||
"text": "Discuss project progress and define upcoming milestones."
|
||||
},
|
||||
"Chapters": {
|
||||
"minutes_total": 3,
|
||||
"content": [
|
||||
{
|
||||
"chapter": "Project Overview",
|
||||
"time_stamp": {"start": 5.12, "end": 5.68},
|
||||
"content": [
|
||||
{"text":"- overview of the project's objectives.","original_transcript_start":3.4,"original_transcript_end":5.7},
|
||||
{"text":"- It outlines the key milestones achieved so far.", "original_transcript_start":6.7, "original_transcript_end":10.5},
|
||||
{"text":"- main challenges faced during the project.", "original_transcript_start":10.8, "original_transcript_end":11.2}
|
||||
],
|
||||
"words_time_stamp": [
|
||||
{"word": "Project", "timestamp": 5.12},
|
||||
{"word": "Overview", "timestamp": 5.12}
|
||||
]
|
||||
},
|
||||
{
|
||||
"chapter": "Budget Review",
|
||||
"time_stamp": {"start": 10.50, "end": 11.20},
|
||||
"content": [
|
||||
{"text":"- review of the current budget allocations.","original_transcript_start":10.5,"original_transcript_end":11.0},
|
||||
{"text":"- discussion on potential cost-saving measures.", "original_transcript_start":11.1, "original_transcript_end":12.0},
|
||||
{"text":"- approval of the budget for the next quarter.", "original_transcript_start":12.1, "original_transcript_end":13.0}
|
||||
],
|
||||
"words_time_stamp": [
|
||||
{"word": "Budget", "timestamp": 10.50},
|
||||
{"word": "Review", "timestamp": 10.50}
|
||||
]
|
||||
}
|
||||
]
|
||||
"Outcomes": {
|
||||
"minutes_total": 3,
|
||||
"content": [
|
||||
{
|
||||
"text": "Key performance metrics were defined and improvement areas identified.",
|
||||
"time_stamp": {"start": 15.30, "end": 16.00},
|
||||
"words_time_stamp": [
|
||||
{"word": "Key", "timestamp": 15.30},
|
||||
{"word": "performance", "timestamp": 15.30},
|
||||
{"word": "metrics", "timestamp": 15.30},
|
||||
{"word": "were", "timestamp": 15.30},
|
||||
{"word": "defined", "timestamp": 15.30},
|
||||
{"word": "and", "timestamp": 15.30},
|
||||
{"word": "improvement", "timestamp": 15.30},
|
||||
{"word": "areas", "timestamp": 15.30},
|
||||
{"word": "identified", "timestamp": 15.30}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"Action_Items_Per_User": [
|
||||
{
|
||||
"speaker": "Speaker_A"
|
||||
"minutes_total": 3,
|
||||
"action_items": [
|
||||
{
|
||||
"text": "Prepare a detailed budget report for the next meeting.",
|
||||
"time_stamp": {"start": 30.45, "end": 30.45},
|
||||
"words_time_stamp": [
|
||||
{"word": "Prepare", "timestamp": 30.45},
|
||||
{"word": "a", "timestamp": 30.45},
|
||||
{"word": "detailed", "timestamp": 30.45},
|
||||
{"word": "budget", "timestamp": 30.45},
|
||||
{"word": "report", "timestamp": 30.45},
|
||||
{"word": "for", "timestamp": 30.45},
|
||||
{"word": "the", "timestamp": 30.45},
|
||||
{"word": "next", "timestamp": 30.45},
|
||||
{"word": "meeting", "timestamp": 30.45}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"speaker": "unassigned",
|
||||
"minutes_total": 2,
|
||||
"action_items": [
|
||||
{
|
||||
"text": "Follow up with the marketing team for the latest campaign updates.",
|
||||
"time_stamp": {"start": 45.67, "end": 45.67},
|
||||
"words_time_stamp": [
|
||||
{"word": "Follow", "timestamp": 45.67},
|
||||
{"word": "up", "timestamp": 45.67},
|
||||
{"word": "with", "timestamp": 45.67},
|
||||
{"word": "the", "timestamp": 45.67},
|
||||
{"word": "marketing", "timestamp": 45.67},
|
||||
{"word": "team", "timestamp": 45.67},
|
||||
{"word": "for", "timestamp": 45.67},
|
||||
{"word": "the", "timestamp": 45.67},
|
||||
{"word": "latest", "timestamp": 45.67},
|
||||
{"word": "campaign", "timestamp": 45.67},
|
||||
{"word": "updates", "timestamp": 45.67}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
NOTE: Action points to the person who is to take the action and if not specified use "unassigned"
|
||||
NOTE: The content under each chapter provides a detailed bulleted explanation of the chapter. It includes "original_transcript_start" and "original_transcript_end," which indicate the timestamps for each bulleted point, referencing where to find it in the original transcript.
|
||||
Remember, every word in each sentence must have a single timestamp equal to the start timestamp of that sentence. Your output must strictly adhere to the provided structure, and the "minutes_total" for each section must be correctly calculated based on the start time of the first sentence and the end time of the last sentence, expressed as a decimal if necessary.
|
||||
NOTE : start and end time are in seconds , so take that into considerations when calculating the total time in mins
|
||||
NOTE: When creating action items per user, if the assigned user is among the speakers, use their associated speaker key that was presented in the sentence (do not infer names from context). If you can't determine the action item is for one of the speakers, make it "unassigned."
|
||||
NOTE: OUTPUT ONLY JSON, NO OTHER TEXT OR COMMENTS DO NOT ADD ```json before or after the text , just return the json output format please
|
||||
"""
|
||||
|
||||
# Basic (Freemium Plan) summary prompt
|
||||
@@ -158,28 +60,11 @@ Create a simple JSON response with just two sections:
|
||||
|
||||
**Example Output JSON:**
|
||||
|
||||
{
|
||||
"Key_Points": [
|
||||
{
|
||||
"text": "Team discussed Q3 marketing strategy.",
|
||||
"timestamp": 120.5
|
||||
},
|
||||
{
|
||||
"text": "Budget approval needed by Friday.",
|
||||
"timestamp": 360.2
|
||||
},
|
||||
{
|
||||
"text": "New product launch delayed until September.",
|
||||
"timestamp": 480.7
|
||||
}
|
||||
],
|
||||
"Summary": {
|
||||
"text": "Marketing team meeting to review Q3 plans and budget requirements. Team agreed on strategy but product launch delayed.",
|
||||
"duration_minutes": 15.5
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Remember to keep your output extremely simple and concise, focusing only on the most important information from the meeting.
|
||||
|
||||
NOTE: OUTPUT ONLY JSON, NO OTHER TEXT OR COMMENTS DO NOT ADD ```json before or after the text , just return the json output format please
|
||||
"""
|
||||
|
||||
# Keeping the original as general_summary_prompt for backward compatibility
|
||||
@@ -203,4 +88,7 @@ Example Output JSON:
|
||||
{ "Key_Points": { "minutes_total": 3.5, "content": [ { "text": "Introductions between Diane Taylor and Cody Smith.", "time_stamp": {"start": 5.12, "end": 5.68}, "words_time_stamp": [ {"word": "Introductions", "timestamp": 5.12}, {"word": "between", "timestamp": 5.12}, {"word": "Diane", "timestamp": 5.12}, {"word": "Taylor", "timestamp": 5.12}, {"word": "and", "timestamp": 5.12}, {"word": "Cody", "timestamp": 5.12}, {"word": "Smith.", "timestamp": 5.12} ] } ] }, "Summary": { "minutes_total": 3.5, "content": [ { "text": "The meeting started with introductions, followed by a discussion of key topics.", "time_stamp": {"start": 5.12, "end": 10.12}, "words_time_stamp": [ {"word": "The", "timestamp": 5.12}, {"word": "meeting", "timestamp": 5.12}, {"word": "started", "timestamp": 5.12}, {"word": "with", "timestamp": 5.12}, {"word": "introductions,", "timestamp": 5.12}, {"word": "followed", "timestamp": 5.12}, {"word": "by", "timestamp": 5.12}, {"word": "a", "timestamp": 5.12}, {"word": "discussion", "timestamp": 5.12}, {"word": "of", "timestamp": 5.12}, {"word": "key", "timestamp": 5.12}, {"word": "topics.", "timestamp": 5.12} ] } ] }, "Next_Steps": { "minutes_total": 2.0, "content": [ { "text": "Diane will follow up with Cody regarding office management tasks.", "time_stamp": {"start": 30.45, "end": 30.45}, "words_time_stamp": [ {"word": "Diane", "timestamp": 30.45}, {"word": "will", "timestamp": 30.45}, {"word": "follow", "timestamp": 30.45}, {"word": "up", "timestamp": 30.45}, {"word": "with", "timestamp": 30.45}, {"word": "Cody", "timestamp": 30.45}, {"word": "regarding", "timestamp": 30.45}, {"word": "office", "timestamp": 30.45}, {"word": "management", "timestamp": 30.45}, {"word": "tasks.", "timestamp": 30.45} ] } ] } }
|
||||
|
||||
Remember, for every sentence generated in any section, every word must be assigned the sentence’s start timestamp as its "timestamp" value. Additionally, calculate the "minutes_total" for each section by using the start time of the first sentence and the end time of the last sentence; if the result is not a whole number, express it as a decimal (e.g., 0.5 mins). Your output must strictly adhere to the provided structure.
|
||||
NOTE : start and end time are in seconds , so take that into considerations when calculating the total time in mins"""
|
||||
NOTE : start and end time are in seconds , so take that into considerations when calculating the total time in mins
|
||||
|
||||
NOTE: OUTPUT ONLY JSON, NO OTHER TEXT OR COMMENTS DO NOT ADD ```json before or after the text , just return the json output format please
|
||||
"""
|
||||
@@ -0,0 +1,2 @@
|
||||
# Tests package
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
"""
|
||||
Shared fixtures and configuration for tests
|
||||
"""
|
||||
import pytest
|
||||
import os
|
||||
from unittest.mock import Mock, MagicMock
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_anthropic_api_key():
|
||||
"""Mock Anthropic API key for testing"""
|
||||
return "test-api-key-12345"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_transcription_dict():
|
||||
"""Sample transcription dictionary for testing"""
|
||||
return {
|
||||
"sentences": [
|
||||
{
|
||||
"text": "Hello, welcome to the meeting.",
|
||||
"start": 0.5,
|
||||
"end": 3.2,
|
||||
"speaker": "Speaker_A",
|
||||
"words": [
|
||||
{"word": "Hello", "start": 0.5, "end": 0.8},
|
||||
{"word": "welcome", "start": 0.9, "end": 1.3},
|
||||
{"word": "to", "start": 1.4, "end": 1.5},
|
||||
{"word": "the", "start": 1.6, "end": 1.7},
|
||||
{"word": "meeting", "start": 1.8, "end": 2.3}
|
||||
]
|
||||
},
|
||||
{
|
||||
"text": "Let's discuss the project timeline.",
|
||||
"start": 4.0,
|
||||
"end": 7.5,
|
||||
"speaker": "Speaker_B",
|
||||
"words": [
|
||||
{"word": "Let's", "start": 4.0, "end": 4.3},
|
||||
{"word": "discuss", "start": 4.4, "end": 5.0},
|
||||
{"word": "the", "start": 5.1, "end": 5.2},
|
||||
{"word": "project", "start": 5.3, "end": 5.8},
|
||||
{"word": "timeline", "start": 5.9, "end": 6.5}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_basic_summary_response():
|
||||
"""Sample basic summary response (freemium plan)"""
|
||||
return {
|
||||
"Key_Points": [
|
||||
{"text": "Team discussed project timeline.", "timestamp": 4.0},
|
||||
{"text": "Meeting started with introductions.", "timestamp": 0.5}
|
||||
],
|
||||
"Summary": {
|
||||
"text": "Brief meeting to discuss project timeline and introductions.",
|
||||
"duration_minutes": 7.5
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_advanced_summary_response():
|
||||
"""Sample advanced summary response (pro plan)"""
|
||||
return {
|
||||
"Purpose": {
|
||||
"text": "Discuss project timeline and team introductions."
|
||||
},
|
||||
"Chapters": {
|
||||
"minutes_total": 0.125,
|
||||
"content": [
|
||||
{
|
||||
"chapter": "Introduction",
|
||||
"time_stamp": {"start": 0.5, "end": 3.2},
|
||||
"content": [
|
||||
{
|
||||
"text": "- Welcome to the meeting.",
|
||||
"original_transcript_start": 0.5,
|
||||
"original_transcript_end": 3.2
|
||||
}
|
||||
],
|
||||
"words_time_stamp": [
|
||||
{"word": "Introduction", "timestamp": 0.5}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"Outcomes": {
|
||||
"minutes_total": 0.125,
|
||||
"content": [
|
||||
{
|
||||
"text": "Project timeline discussed.",
|
||||
"time_stamp": {"start": 4.0, "end": 7.5},
|
||||
"words_time_stamp": [
|
||||
{"word": "Project", "timestamp": 4.0},
|
||||
{"word": "timeline", "timestamp": 4.0}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"Action_Items_Per_User": [
|
||||
{
|
||||
"speaker": "Speaker_B",
|
||||
"minutes_total": 0.125,
|
||||
"action_items": [
|
||||
{
|
||||
"text": "Review project timeline.",
|
||||
"time_stamp": {"start": 4.0, "end": 7.5},
|
||||
"words_time_stamp": [
|
||||
{"word": "Review", "timestamp": 4.0}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_template():
|
||||
"""Sample custom template"""
|
||||
return {
|
||||
"Key_Points": "Summarize the most critical discussion points from the meeting.",
|
||||
"Summary": "Provide a brief overall summary of what was discussed.",
|
||||
"Next_Steps": "List the next steps decided during the meeting."
|
||||
}
|
||||
|
||||
@@ -0,0 +1,417 @@
|
||||
"""
|
||||
Unit tests for generate_summary.py module
|
||||
"""
|
||||
import pytest
|
||||
import os
|
||||
import json
|
||||
from unittest.mock import Mock, MagicMock, patch, call
|
||||
from pydantic import ValidationError
|
||||
|
||||
# Import the functions and models to test
|
||||
import sys
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
|
||||
|
||||
from scripts.generate_summary import (
|
||||
general_summary,
|
||||
custom_summary,
|
||||
BasicSummary,
|
||||
AdvancedSummary,
|
||||
KeyPoint,
|
||||
Summary,
|
||||
Purpose,
|
||||
Chapters,
|
||||
Outcomes,
|
||||
ActionItemsPerUser
|
||||
)
|
||||
|
||||
|
||||
class TestPydanticModels:
|
||||
"""Test Pydantic model validation"""
|
||||
|
||||
def test_key_point_model(self):
|
||||
"""Test KeyPoint model validation"""
|
||||
key_point = KeyPoint(text="Test point", timestamp=10.5)
|
||||
assert key_point.text == "Test point"
|
||||
assert key_point.timestamp == 10.5
|
||||
|
||||
def test_summary_model(self):
|
||||
"""Test Summary model validation"""
|
||||
summary = Summary(text="Test summary", duration_minutes=15.5)
|
||||
assert summary.text == "Test summary"
|
||||
assert summary.duration_minutes == 15.5
|
||||
|
||||
def test_basic_summary_model(self, sample_basic_summary_response):
|
||||
"""Test BasicSummary model validation"""
|
||||
basic_summary = BasicSummary(**sample_basic_summary_response)
|
||||
assert len(basic_summary.Key_Points) == 2
|
||||
assert basic_summary.Summary.text is not None
|
||||
assert basic_summary.Summary.duration_minutes > 0
|
||||
|
||||
def test_advanced_summary_model(self, sample_advanced_summary_response):
|
||||
"""Test AdvancedSummary model validation"""
|
||||
advanced_summary = AdvancedSummary(**sample_advanced_summary_response)
|
||||
assert advanced_summary.Purpose.text is not None
|
||||
assert len(advanced_summary.Chapters.content) > 0
|
||||
assert len(advanced_summary.Outcomes.content) > 0
|
||||
assert len(advanced_summary.Action_Items_Per_User) > 0
|
||||
|
||||
|
||||
class TestGeneralSummary:
|
||||
"""Test general_summary function"""
|
||||
|
||||
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
||||
@patch('scripts.generate_summary.ChatAnthropic')
|
||||
@patch('scripts.generate_summary.ChatPromptTemplate')
|
||||
def test_general_summary_freemium_plan(
|
||||
self,
|
||||
mock_prompt_template,
|
||||
mock_chat_anthropic,
|
||||
sample_transcription_dict,
|
||||
sample_basic_summary_response
|
||||
):
|
||||
"""Test general_summary with freemium plan"""
|
||||
# Setup mocks
|
||||
mock_model = MagicMock()
|
||||
mock_chain = MagicMock()
|
||||
mock_structured_model = MagicMock()
|
||||
|
||||
# Create a mock BasicSummary instance
|
||||
from scripts.generate_summary import BasicSummary
|
||||
mock_result = BasicSummary(**sample_basic_summary_response)
|
||||
|
||||
mock_structured_model.invoke.return_value = mock_result
|
||||
mock_model.with_structured_output.return_value = mock_structured_model
|
||||
mock_prompt_template.from_messages.return_value = MagicMock()
|
||||
mock_chat_anthropic.return_value = mock_model
|
||||
|
||||
# Mock the pipe operator
|
||||
mock_prompt_instance = MagicMock()
|
||||
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
||||
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
||||
mock_chain.invoke.return_value = mock_result
|
||||
|
||||
# Call the function
|
||||
result = general_summary(sample_transcription_dict, plan_tier="freemium")
|
||||
|
||||
# Verify results
|
||||
assert result is not None
|
||||
assert "Key_Points" in result
|
||||
assert "Summary" in result
|
||||
assert len(result["Key_Points"]) > 0
|
||||
|
||||
# Verify model was initialized with correct parameters
|
||||
mock_chat_anthropic.assert_called_once()
|
||||
call_args = mock_chat_anthropic.call_args
|
||||
assert call_args.kwargs["model"] == "claude-sonnet-4-5-20250929"
|
||||
assert call_args.kwargs["max_tokens"] == 2000
|
||||
assert call_args.kwargs["temperature"] == 0.2
|
||||
|
||||
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
||||
@patch('scripts.generate_summary.ChatAnthropic')
|
||||
@patch('scripts.generate_summary.ChatPromptTemplate')
|
||||
def test_general_summary_pro_plan(
|
||||
self,
|
||||
mock_prompt_template,
|
||||
mock_chat_anthropic,
|
||||
sample_transcription_dict,
|
||||
sample_advanced_summary_response
|
||||
):
|
||||
"""Test general_summary with pro plan"""
|
||||
# Setup mocks
|
||||
mock_model = MagicMock()
|
||||
mock_chain = MagicMock()
|
||||
|
||||
# Create a mock AdvancedSummary instance
|
||||
from scripts.generate_summary import AdvancedSummary
|
||||
mock_result = AdvancedSummary(**sample_advanced_summary_response)
|
||||
|
||||
mock_structured_model = MagicMock()
|
||||
mock_structured_model.invoke.return_value = mock_result
|
||||
mock_model.with_structured_output.return_value = mock_structured_model
|
||||
mock_prompt_template.from_messages.return_value = MagicMock()
|
||||
mock_chat_anthropic.return_value = mock_model
|
||||
|
||||
# Mock the pipe operator
|
||||
mock_prompt_instance = MagicMock()
|
||||
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
||||
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
||||
mock_chain.invoke.return_value = mock_result
|
||||
|
||||
# Call the function
|
||||
result = general_summary(sample_transcription_dict, plan_tier="pro")
|
||||
|
||||
# Verify results
|
||||
assert result is not None
|
||||
assert "Purpose" in result
|
||||
assert "Chapters" in result
|
||||
assert "Outcomes" in result
|
||||
assert "Action_Items_Per_User" in result
|
||||
|
||||
# Verify model was initialized with correct parameters
|
||||
mock_chat_anthropic.assert_called_once()
|
||||
call_args = mock_chat_anthropic.call_args
|
||||
assert call_args.kwargs["model"] == "claude-sonnet-4-5-20250929"
|
||||
assert call_args.kwargs["max_tokens"] == 4000
|
||||
|
||||
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
||||
@patch('scripts.generate_summary.ChatAnthropic')
|
||||
@patch('scripts.generate_summary.ChatPromptTemplate')
|
||||
def test_general_summary_with_string_transcription(
|
||||
self,
|
||||
mock_prompt_template,
|
||||
mock_chat_anthropic,
|
||||
sample_basic_summary_response
|
||||
):
|
||||
"""Test general_summary with string transcription"""
|
||||
transcription_str = json.dumps({"sentences": []})
|
||||
|
||||
mock_model = MagicMock()
|
||||
mock_chain = MagicMock()
|
||||
from scripts.generate_summary import BasicSummary
|
||||
mock_result = BasicSummary(**sample_basic_summary_response)
|
||||
|
||||
mock_structured_model = MagicMock()
|
||||
mock_structured_model.invoke.return_value = mock_result
|
||||
mock_model.with_structured_output.return_value = mock_structured_model
|
||||
mock_prompt_template.from_messages.return_value = MagicMock()
|
||||
mock_chat_anthropic.return_value = mock_model
|
||||
|
||||
mock_prompt_instance = MagicMock()
|
||||
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
||||
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
||||
mock_chain.invoke.return_value = mock_result
|
||||
|
||||
result = general_summary(transcription_str, plan_tier="freemium")
|
||||
|
||||
assert result is not None
|
||||
assert "Key_Points" in result
|
||||
|
||||
def test_general_summary_missing_api_key(self, sample_transcription_dict):
|
||||
"""Test general_summary raises error when API key is missing"""
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
with pytest.raises(ValueError, match="ANTHROPIC_API_KEY"):
|
||||
general_summary(sample_transcription_dict, plan_tier="freemium")
|
||||
|
||||
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
||||
@patch('scripts.generate_summary.ChatAnthropic')
|
||||
@patch('scripts.generate_summary.ChatPromptTemplate')
|
||||
def test_general_summary_fallback_on_error(
|
||||
self,
|
||||
mock_prompt_template,
|
||||
mock_chat_anthropic,
|
||||
sample_transcription_dict,
|
||||
sample_basic_summary_response
|
||||
):
|
||||
"""Test general_summary falls back to non-structured output on error"""
|
||||
mock_model = MagicMock()
|
||||
mock_chain = MagicMock()
|
||||
mock_fallback_chain = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.content = json.dumps(sample_basic_summary_response)
|
||||
|
||||
# First call (structured) raises error
|
||||
mock_structured_model = MagicMock()
|
||||
mock_structured_model.invoke.side_effect = Exception("Structured output failed")
|
||||
mock_model.with_structured_output.return_value = mock_structured_model
|
||||
|
||||
# Set up the chain so that structured chain raises exception
|
||||
# and fallback chain returns the mock_response
|
||||
mock_prompt_instance = MagicMock()
|
||||
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
||||
|
||||
# When __or__ is called with structured_model, return chain that raises exception
|
||||
# When __or__ is called with model (fallback), return fallback_chain that succeeds
|
||||
def or_handler(self, other):
|
||||
if other == mock_structured_model:
|
||||
# Structured chain - should raise exception
|
||||
mock_chain.invoke.side_effect = Exception("Structured output failed")
|
||||
return mock_chain
|
||||
elif other == mock_model:
|
||||
# Fallback chain - should succeed
|
||||
mock_fallback_chain.invoke.return_value = mock_response
|
||||
return mock_fallback_chain
|
||||
return mock_chain
|
||||
|
||||
mock_prompt_instance.__or__ = or_handler
|
||||
mock_chat_anthropic.return_value = mock_model
|
||||
|
||||
result = general_summary(sample_transcription_dict, plan_tier="freemium")
|
||||
|
||||
assert result is not None
|
||||
assert "Key_Points" in result
|
||||
|
||||
|
||||
class TestCustomSummary:
|
||||
"""Test custom_summary function"""
|
||||
|
||||
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
||||
@patch('scripts.generate_summary.ChatAnthropic')
|
||||
@patch('scripts.generate_summary.ChatPromptTemplate')
|
||||
def test_custom_summary_success(
|
||||
self,
|
||||
mock_prompt_template,
|
||||
mock_chat_anthropic,
|
||||
sample_transcription_dict,
|
||||
sample_template
|
||||
):
|
||||
"""Test custom_summary with successful response"""
|
||||
mock_model = MagicMock()
|
||||
mock_chain = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
|
||||
expected_result = {
|
||||
"Key_Points": {"content": []},
|
||||
"Summary": {"content": []},
|
||||
"Next_Steps": {"content": []}
|
||||
}
|
||||
mock_response.content = json.dumps(expected_result)
|
||||
|
||||
mock_prompt_instance = MagicMock()
|
||||
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
||||
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
||||
mock_chain.invoke.return_value = mock_response
|
||||
mock_chat_anthropic.return_value = mock_model
|
||||
|
||||
result = custom_summary(sample_template, sample_transcription_dict)
|
||||
|
||||
assert result is not None
|
||||
assert "Key_Points" in result or "Summary" in result
|
||||
|
||||
# Verify model was initialized correctly
|
||||
mock_chat_anthropic.assert_called_once()
|
||||
call_args = mock_chat_anthropic.call_args
|
||||
assert call_args.kwargs["model"] == "claude-sonnet-4-5-20250929"
|
||||
assert call_args.kwargs["max_tokens"] == 8000
|
||||
|
||||
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
||||
@patch('scripts.generate_summary.ChatAnthropic')
|
||||
@patch('scripts.generate_summary.ChatPromptTemplate')
|
||||
def test_custom_summary_with_markdown_wrapper(
|
||||
self,
|
||||
mock_prompt_template,
|
||||
mock_chat_anthropic,
|
||||
sample_transcription_dict,
|
||||
sample_template
|
||||
):
|
||||
"""Test custom_summary handles markdown-wrapped JSON"""
|
||||
mock_model = MagicMock()
|
||||
mock_chain = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
|
||||
expected_result = {"result": "test"}
|
||||
wrapped_json = f"```json\n{json.dumps(expected_result)}\n```"
|
||||
mock_response.content = wrapped_json
|
||||
|
||||
mock_prompt_instance = MagicMock()
|
||||
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
||||
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
||||
mock_chain.invoke.return_value = mock_response
|
||||
mock_chat_anthropic.return_value = mock_model
|
||||
|
||||
result = custom_summary(sample_template, sample_transcription_dict)
|
||||
|
||||
assert result == expected_result
|
||||
|
||||
def test_custom_summary_missing_api_key(self, sample_transcription_dict, sample_template):
|
||||
"""Test custom_summary raises error when API key is missing"""
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
with pytest.raises(ValueError, match="ANTHROPIC_API_KEY"):
|
||||
custom_summary(sample_template, sample_transcription_dict)
|
||||
|
||||
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
||||
@patch('scripts.generate_summary.ChatAnthropic')
|
||||
@patch('scripts.generate_summary.ChatPromptTemplate')
|
||||
def test_custom_summary_invalid_json(
|
||||
self,
|
||||
mock_prompt_template,
|
||||
mock_chat_anthropic,
|
||||
sample_transcription_dict,
|
||||
sample_template
|
||||
):
|
||||
"""Test custom_summary handles invalid JSON gracefully"""
|
||||
mock_model = MagicMock()
|
||||
mock_chain = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.content = "This is not valid JSON"
|
||||
|
||||
mock_prompt_instance = MagicMock()
|
||||
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
||||
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
||||
mock_chain.invoke.return_value = mock_response
|
||||
mock_chat_anthropic.return_value = mock_model
|
||||
|
||||
with pytest.raises(ValueError, match="Could not parse response as JSON"):
|
||||
custom_summary(sample_template, sample_transcription_dict)
|
||||
|
||||
|
||||
class TestSchemaSwitching:
|
||||
"""Test schema switching based on plan tier"""
|
||||
|
||||
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
||||
@patch('scripts.generate_summary.ChatAnthropic')
|
||||
@patch('scripts.generate_summary.ChatPromptTemplate')
|
||||
def test_schema_switching_freemium(
|
||||
self,
|
||||
mock_prompt_template,
|
||||
mock_chat_anthropic,
|
||||
sample_transcription_dict,
|
||||
sample_basic_summary_response
|
||||
):
|
||||
"""Test that freemium plan uses BasicSummary schema"""
|
||||
mock_model = MagicMock()
|
||||
mock_chain = MagicMock()
|
||||
from scripts.generate_summary import BasicSummary
|
||||
mock_result = BasicSummary(**sample_basic_summary_response)
|
||||
|
||||
mock_structured_model = MagicMock()
|
||||
mock_structured_model.invoke.return_value = mock_result
|
||||
mock_model.with_structured_output = MagicMock(return_value=mock_structured_model)
|
||||
|
||||
mock_prompt_instance = MagicMock()
|
||||
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
||||
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
||||
mock_chain.invoke.return_value = mock_result
|
||||
mock_chat_anthropic.return_value = mock_model
|
||||
|
||||
general_summary(sample_transcription_dict, plan_tier="freemium")
|
||||
|
||||
# Verify BasicSummary schema was used
|
||||
mock_model.with_structured_output.assert_called_once()
|
||||
call_args = mock_model.with_structured_output.call_args
|
||||
from scripts.generate_summary import BasicSummary
|
||||
assert call_args[0][0] == BasicSummary
|
||||
|
||||
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
||||
@patch('scripts.generate_summary.ChatAnthropic')
|
||||
@patch('scripts.generate_summary.ChatPromptTemplate')
|
||||
def test_schema_switching_pro(
|
||||
self,
|
||||
mock_prompt_template,
|
||||
mock_chat_anthropic,
|
||||
sample_transcription_dict,
|
||||
sample_advanced_summary_response
|
||||
):
|
||||
"""Test that pro plan uses AdvancedSummary schema"""
|
||||
mock_model = MagicMock()
|
||||
mock_chain = MagicMock()
|
||||
from scripts.generate_summary import AdvancedSummary
|
||||
mock_result = AdvancedSummary(**sample_advanced_summary_response)
|
||||
|
||||
mock_structured_model = MagicMock()
|
||||
mock_structured_model.invoke.return_value = mock_result
|
||||
mock_model.with_structured_output = MagicMock(return_value=mock_structured_model)
|
||||
|
||||
mock_prompt_instance = MagicMock()
|
||||
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
||||
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
||||
mock_chain.invoke.return_value = mock_result
|
||||
mock_chat_anthropic.return_value = mock_model
|
||||
|
||||
general_summary(sample_transcription_dict, plan_tier="pro")
|
||||
|
||||
# Verify AdvancedSummary schema was used
|
||||
mock_model.with_structured_output.assert_called_once()
|
||||
call_args = mock_model.with_structured_output.call_args
|
||||
from scripts.generate_summary import AdvancedSummary
|
||||
assert call_args[0][0] == AdvancedSummary
|
||||
|
||||
Reference in New Issue
Block a user