62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
|
import anthropic
|
||
|
|
import os
|
||
|
|
from dotenv import load_dotenv
|
||
|
|
import json
|
||
|
|
from src.prompt import advanced_summary_prompt, basic_summary_prompt, custom_template_prompt
|
||
|
|
load_dotenv()
|
||
|
|
|
||
|
|
def general_summary(transcription, plan_tier="pro"):
|
||
|
|
"""
|
||
|
|
Generate a summary of the transcription based on the user's plan tier.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
transcription: The transcription to summarize
|
||
|
|
plan_tier: The user's plan tier ("freemium" or "pro")
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A JSON object containing the summary
|
||
|
|
"""
|
||
|
|
client = anthropic.Anthropic(
|
||
|
|
api_key=os.getenv("ANTHTROPIC_API_KEY"),
|
||
|
|
)
|
||
|
|
|
||
|
|
# Select the appropriate prompt 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
|
||
|
|
else: # Default to pro
|
||
|
|
prompt = advanced_summary_prompt
|
||
|
|
max_tokens = 4000
|
||
|
|
|
||
|
|
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}"}
|
||
|
|
]
|
||
|
|
)
|
||
|
|
|
||
|
|
text = message.content[0].text
|
||
|
|
return json.loads(text)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
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}"}
|
||
|
|
]
|
||
|
|
)
|
||
|
|
|
||
|
|
text = message.content[0].text
|
||
|
|
return json.loads(text)
|