wehub-resource-sync
314 lines
14 KiB
Python
314 lines
14 KiB
Python
from pydantic import BaseModel
|
|
from typing import List, Tuple, Literal
|
|
from pydantic_ai import Agent
|
|
from pydantic_ai.models.anthropic import AnthropicModel
|
|
from pydantic_ai.models.groq import GroqModel
|
|
from pydantic_ai.models.openai import OpenAIModel
|
|
from pydantic_ai.providers.openai import OpenAIProvider
|
|
from pydantic_ai.providers.groq import GroqProvider
|
|
from pydantic_ai.providers.anthropic import AnthropicProvider
|
|
|
|
import logging
|
|
import os
|
|
from dotenv import load_dotenv
|
|
from db import DatabaseManager
|
|
from ollama import chat
|
|
import asyncio
|
|
from ollama import AsyncClient
|
|
|
|
|
|
|
|
|
|
|
|
# Set up logging
|
|
logging.basicConfig(
|
|
level=logging.DEBUG,
|
|
format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
load_dotenv() # Load environment variables from .env file
|
|
|
|
db = DatabaseManager()
|
|
|
|
class Block(BaseModel):
|
|
"""Represents a block of content in a section.
|
|
|
|
Block types must align with frontend rendering capabilities:
|
|
- 'text': Plain text content
|
|
- 'bullet': Bulleted list item
|
|
- 'heading1': Large section heading
|
|
- 'heading2': Medium section heading
|
|
|
|
Colors currently supported:
|
|
- 'gray': Gray text color
|
|
- '' or any other value: Default text color
|
|
"""
|
|
id: str
|
|
type: Literal['bullet', 'heading1', 'heading2', 'text']
|
|
content: str
|
|
color: str # Frontend currently only uses 'gray' or default
|
|
|
|
class Section(BaseModel):
|
|
"""Represents a section in the meeting summary"""
|
|
title: str
|
|
blocks: List[Block]
|
|
|
|
class MeetingNotes(BaseModel):
|
|
"""Represents the meeting notes"""
|
|
meeting_name: str
|
|
sections: List[Section]
|
|
|
|
class People(BaseModel):
|
|
"""Represents the people in the meeting. Always have this part in the output. Title - Person Name (Role, Details)"""
|
|
title: str
|
|
blocks: List[Block]
|
|
|
|
class SummaryResponse(BaseModel):
|
|
"""Represents the meeting summary response based on a section of the transcript"""
|
|
MeetingName : str
|
|
People : People
|
|
SessionSummary : Section
|
|
CriticalDeadlines: Section
|
|
KeyItemsDecisions: Section
|
|
ImmediateActionItems: Section
|
|
NextSteps: Section
|
|
MeetingNotes: MeetingNotes
|
|
|
|
# --- Main Class Used by main.py ---
|
|
|
|
class TranscriptProcessor:
|
|
"""Handles the processing of meeting transcripts using AI models."""
|
|
def __init__(self):
|
|
"""Initialize the transcript processor."""
|
|
logger.info("TranscriptProcessor initialized.")
|
|
self.db = DatabaseManager()
|
|
self.active_clients = [] # Track active Ollama client sessions
|
|
async def process_transcript(self, text: str, model: str, model_name: str, chunk_size: int = 5000, overlap: int = 1000, custom_prompt: str = "") -> Tuple[int, List[str]]:
|
|
"""
|
|
Process transcript text into chunks and generate structured summaries for each chunk using an AI model.
|
|
|
|
Args:
|
|
text: The transcript text.
|
|
model: The AI model provider ('claude', 'ollama', 'groq', 'openai').
|
|
model_name: The specific model name.
|
|
chunk_size: The size of each text chunk.
|
|
overlap: The overlap between consecutive chunks.
|
|
custom_prompt: A custom prompt to use for the AI model.
|
|
|
|
Returns:
|
|
A tuple containing:
|
|
- The number of chunks processed.
|
|
- A list of JSON strings, where each string is the summary of a chunk.
|
|
"""
|
|
|
|
logger.info(f"Processing transcript (length {len(text)}) with model provider={model}, model_name={model_name}, chunk_size={chunk_size}, overlap={overlap}")
|
|
|
|
all_json_data = []
|
|
agent = None # Define agent variable
|
|
llm = None # Define llm variable
|
|
|
|
try:
|
|
# Select and initialize the AI model and agent
|
|
if model == "claude":
|
|
api_key = await db.get_api_key("claude")
|
|
if not api_key: raise ValueError("ANTHROPIC_API_KEY environment variable not set")
|
|
llm = AnthropicModel(model_name, provider=AnthropicProvider(api_key=api_key))
|
|
logger.info(f"Using Claude model: {model_name}")
|
|
elif model == "ollama":
|
|
# Use environment variable for Ollama host configuration
|
|
ollama_host = os.getenv('OLLAMA_HOST', 'http://localhost:11434')
|
|
ollama_base_url = f"{ollama_host}/v1"
|
|
ollama_model = OpenAIModel(
|
|
model_name=model_name, provider=OpenAIProvider(base_url=ollama_base_url)
|
|
)
|
|
llm = ollama_model
|
|
if model_name.lower().startswith("phi4") or model_name.lower().startswith("llama"):
|
|
chunk_size = 10000
|
|
overlap = 1000
|
|
else:
|
|
chunk_size = 30000
|
|
overlap = 1000
|
|
logger.info(f"Using Ollama model: {model_name}")
|
|
elif model == "groq":
|
|
api_key = await db.get_api_key("groq")
|
|
if not api_key: raise ValueError("GROQ_API_KEY environment variable not set")
|
|
llm = GroqModel(model_name, provider=GroqProvider(api_key=api_key))
|
|
logger.info(f"Using Groq model: {model_name}")
|
|
# --- ADD OPENAI SUPPORT HERE ---
|
|
elif model == "openai":
|
|
api_key = await db.get_api_key("openai")
|
|
if not api_key: raise ValueError("OPENAI_API_KEY environment variable not set")
|
|
llm = OpenAIModel(model_name, provider=OpenAIProvider(api_key=api_key))
|
|
logger.info(f"Using OpenAI model: {model_name}")
|
|
# --- END OPENAI SUPPORT ---
|
|
else:
|
|
logger.error(f"Unsupported model provider requested: {model}")
|
|
raise ValueError(f"Unsupported model provider: {model}")
|
|
|
|
# Initialize the agent with the selected LLM
|
|
agent = Agent(
|
|
llm,
|
|
result_type=SummaryResponse,
|
|
result_retries=2,
|
|
)
|
|
logger.info("Pydantic-AI Agent initialized.")
|
|
|
|
# Split transcript into chunks
|
|
step = chunk_size - overlap
|
|
if step <= 0:
|
|
logger.warning(f"Overlap ({overlap}) >= chunk_size ({chunk_size}). Adjusting overlap.")
|
|
overlap = max(0, chunk_size - 100)
|
|
step = chunk_size - overlap
|
|
|
|
chunks = [text[i:i+chunk_size] for i in range(0, len(text), step)]
|
|
num_chunks = len(chunks)
|
|
logger.info(f"Split transcript into {num_chunks} chunks.")
|
|
|
|
for i, chunk in enumerate(chunks):
|
|
logger.info(f"Processing chunk {i+1}/{num_chunks}...")
|
|
try:
|
|
# Run the agent to get the structured summary for the chunk
|
|
if model != "ollama":
|
|
summary_result = await agent.run(
|
|
f"""Given the following meeting transcript chunk, extract the relevant information according to the required JSON structure. If a specific section (like Critical Deadlines) has no relevant information in this chunk, return an empty list for its 'blocks'. Ensure the output is only the JSON data.
|
|
|
|
IMPORTANT: Block types must be one of: 'text', 'bullet', 'heading1', 'heading2'
|
|
- Use 'text' for regular paragraphs
|
|
- Use 'bullet' for list items
|
|
- Use 'heading1' for major headings
|
|
- Use 'heading2' for subheadings
|
|
|
|
For the color field, use 'gray' for less important content or '' (empty string) for default.
|
|
|
|
Transcript Chunk:
|
|
---
|
|
{chunk}
|
|
---
|
|
|
|
Please capture all relevant action items. Transcription can have spelling mistakes. correct it if required. context is important.
|
|
|
|
While generating the summary, please add the following context:
|
|
---
|
|
{custom_prompt}
|
|
---
|
|
Make sure the output is only the JSON data.
|
|
""",
|
|
)
|
|
else:
|
|
logger.info(f"Using Ollama model: {model_name} and chunk size: {chunk_size} with overlap: {overlap}")
|
|
response = await self.chat_ollama_model(model_name, chunk, custom_prompt)
|
|
|
|
# Check if response is already a SummaryResponse object or a string that needs validation
|
|
if isinstance(response, SummaryResponse):
|
|
summary_result = response
|
|
else:
|
|
# If it's a string (JSON), validate it
|
|
summary_result = SummaryResponse.model_validate_json(response)
|
|
|
|
logger.info(f"Summary result for chunk {i+1}: {summary_result}")
|
|
logger.info(f"Summary result type for chunk {i+1}: {type(summary_result)}")
|
|
|
|
if hasattr(summary_result, 'data') and isinstance(summary_result.data, SummaryResponse):
|
|
final_summary_pydantic = summary_result.data
|
|
elif isinstance(summary_result, SummaryResponse):
|
|
final_summary_pydantic = summary_result
|
|
else:
|
|
logger.error(f"Unexpected result type from agent for chunk {i+1}: {type(summary_result)}")
|
|
continue # Skip this chunk
|
|
|
|
# Convert the Pydantic model to a JSON string
|
|
chunk_summary_json = final_summary_pydantic.model_dump_json()
|
|
all_json_data.append(chunk_summary_json)
|
|
logger.info(f"Successfully generated summary for chunk {i+1}.")
|
|
|
|
except Exception as chunk_error:
|
|
logger.error(f"Error processing chunk {i+1}: {chunk_error}", exc_info=True)
|
|
|
|
logger.info(f"Finished processing all {num_chunks} chunks.")
|
|
return num_chunks, all_json_data
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error during transcript processing: {str(e)}", exc_info=True)
|
|
raise
|
|
|
|
async def chat_ollama_model(self, model_name: str, transcript: str, custom_prompt: str):
|
|
message = {
|
|
'role': 'system',
|
|
'content': f'''
|
|
Given the following meeting transcript chunk, extract the relevant information according to the required JSON structure. If a specific section (like Critical Deadlines) has no relevant information in this chunk, return an empty list for its 'blocks'. Ensure the output is only the JSON data.
|
|
|
|
Transcript Chunk:
|
|
---
|
|
{transcript}
|
|
---
|
|
Please capture all relevant action items. Transcription can have spelling mistakes. correct it if required. context is important.
|
|
|
|
While generating the summary, please add the following context:
|
|
---
|
|
{custom_prompt}
|
|
---
|
|
|
|
Make sure the output is only the JSON data.
|
|
|
|
''',
|
|
}
|
|
|
|
# Create a client and track it for cleanup
|
|
ollama_host = os.getenv('OLLAMA_HOST', 'http://127.0.0.1:11434')
|
|
client = AsyncClient(host=ollama_host)
|
|
self.active_clients.append(client)
|
|
|
|
try:
|
|
response = await client.chat(model=model_name, messages=[message], stream=True, format=SummaryResponse.model_json_schema())
|
|
|
|
full_response = ""
|
|
async for part in response:
|
|
content = part['message']['content']
|
|
print(content, end='', flush=True)
|
|
full_response += content
|
|
|
|
try:
|
|
summary = SummaryResponse.model_validate_json(full_response)
|
|
print("\n", summary.model_dump_json(indent=2), type(summary))
|
|
return summary
|
|
except Exception as e:
|
|
print(f"\nError parsing response: {e}")
|
|
return full_response
|
|
except asyncio.CancelledError:
|
|
logger.info("Ollama request was cancelled during shutdown")
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Error in Ollama chat: {e}")
|
|
raise
|
|
finally:
|
|
# Remove the client from active clients list
|
|
if client in self.active_clients:
|
|
self.active_clients.remove(client)
|
|
|
|
def cleanup(self):
|
|
"""Clean up resources used by the TranscriptProcessor."""
|
|
logger.info("Cleaning up TranscriptProcessor resources")
|
|
try:
|
|
# Close database connections if any
|
|
if hasattr(self, 'db') and self.db is not None:
|
|
# self.db.close()
|
|
logger.info("Database connection cleanup (using context managers)")
|
|
|
|
# Cancel any active Ollama client sessions
|
|
if hasattr(self, 'active_clients') and self.active_clients:
|
|
logger.info(f"Terminating {len(self.active_clients)} active Ollama client sessions")
|
|
for client in self.active_clients:
|
|
try:
|
|
# Close the client's underlying connection
|
|
if hasattr(client, '_client') and hasattr(client._client, 'close'):
|
|
asyncio.create_task(client._client.aclose())
|
|
except Exception as client_error:
|
|
logger.error(f"Error closing Ollama client: {client_error}", exc_info=True)
|
|
# Clear the list
|
|
self.active_clients.clear()
|
|
logger.info("All Ollama client sessions terminated")
|
|
except Exception as e:
|
|
logger.error(f"Error during TranscriptProcessor cleanup: {str(e)}", exc_info=True)
|
|
|
|
|