CallCenterSummarizationAgent / src /agents /TranscriptionAgent.py
Fade0510's picture
Score improvement changes
0aa8e87
Raw
History Blame Contribute Delete
1.02 kB
import whisper
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from src.agents.CallState import CallState
class TranscriptionAgent:
def __init__(self):
self.llm = ChatOpenAI(model="gpt-4o", temperature=0)
def __call__(self, state: CallState) -> CallState:
"""Converts audio to text using whisper."""
file_path = state["file_path"]
# Load whisper model - using 'base' for faster processing
model = whisper.load_model("base")
result = model.transcribe(file_path)
state["content"] = result["text"]
# Clean the transcript
prompt = PromptTemplate.from_template(
"Clean the following text of any profanity and fix basic grammatical errors. Return only the clean text:\n{text}"
)
chain = prompt | self.llm
clean_text = chain.invoke({"text": state["content"]}).content
state["clean_content"] = clean_text
return state