import logging import os from typing import Dict, Tuple, Optional import psycopg2 import psycopg2.extras import requests from dotenv import load_dotenv # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') # Load environment variables load_dotenv() class DatabaseConnectionError(Exception): """Custom exception for database connection errors.""" pass class AIVAError(Exception): """Custom exception for AIVA related errors.""" pass class Summarizer: """ A class for summarizing documents using AIVA and caching summaries in PostgreSQL. """ def __init__(self, db_host: str, db_port: int, db_name: str, db_user: str, db_password: str, aiva_url: str): """ Initializes the Summarizer with database connection details and AIVA URL. Args: db_host: The hostname of the PostgreSQL server. db_port: The port number of the PostgreSQL server. db_name: The name of the PostgreSQL database. db_user: The username for connecting to the database. db_password: The password for connecting to the database. aiva_url: The URL of the AIVA service. """ self.db_host = db_host self.db_port = db_port self.db_name = db_name self.db_user = db_user self.db_password = db_password self.aiva_url = aiva_url self.conn = None # Initialize connection attribute def connect_to_db(self): """ Establishes a connection to the PostgreSQL database. """ try: self.conn = psycopg2.connect( host=self.db_host, port=self.db_port, dbname=self.db_name, user=self.db_user, password=self.db_password, ) logging.info("Successfully connected to the database.") except psycopg2.Error as e: logging.error(f"Error connecting to the database: {e}") raise DatabaseConnectionError(f"Failed to connect to the database: {e}") def close_db_connection(self): """ Closes the connection to the PostgreSQL database. """ if self.conn: self.conn.close() logging.info("Database connection closed.") def get_summary_from_db(self, document_content: str) -> Optional[Dict[str, str]]: """ Retrieves a summary from the database if it exists. Args: document_content: The content of the document to summarize. Returns: A dictionary containing the summary components (one_line_summary, key_points, technical_terms) if found in the database, otherwise None. """ if not self.conn: logging.error("Database connection not established.") return None try: cur = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) cur.execute( "SELECT one_line_summary, key_points, technical_terms FROM summaries WHERE document_content = %s", (document_content,), ) result = cur.fetchone() cur.close() if result: logging.info("Summary found in the database.") return dict(result) # Convert psycopg2.extras.DictRow to a standard dictionary else: logging.info("Summary not found in the database.") return None except psycopg2.Error as e: logging.error(f"Error retrieving summary from the database: {e}") return None def save_summary_to_db(self, document_content: str, summary: Dict[str, str]): """ Saves the generated summary to the database. Args: document_content: The content of the document. summary: A dictionary containing the summary components (one_line_summary, key_points, technical_terms). """ if not self.conn: logging.error("Database connection not established.") return try: cur = self.conn.cursor() cur.execute( "INSERT INTO summaries (document_content, one_line_summary, key_points, technical_terms) VALUES (%s, %s, %s, %s)", ( document_content, summary['one_line_summary'], summary['key_points'], summary['technical_terms'], ), ) self.conn.commit() cur.close() logging.info("Summary saved to the database.") except psycopg2.Error as e: logging.error(f"Error saving summary to the database: {e}") self.conn.rollback() # Rollback the transaction in case of an error def generate_summary_with_aiva(self, document_content: str) -> Dict[str, str]: """ Generates a summary of the document content using the AIVA service. Args: document_content: The content of the document to summarize. Returns: A dictionary containing the summary components (one_line_summary, key_points, technical_terms). Raises: AIVAError: If there is an error communicating with the AIVA service. """ try: payload = {"document_content": document_content} response = requests.post(f"{self.aiva_url}/summarize", json=payload, timeout=60) # Added timeout response.raise_for_status() # Raise HTTPError for bad responses (4xx or 5xx) summary = response.json() if not isinstance(summary, dict) or not all( key in summary for key in ("one_line_summary", "key_points", "technical_terms") ): raise ValueError(f"Invalid AIVA response format: {summary}") logging.info("Summary generated by AIVA.") return summary except requests.exceptions.RequestException as e: logging.error(f"Error communicating with AIVA: {e}") raise AIVAError(f"Failed to communicate with AIVA: {e}") except ValueError as e: logging.error(f"Error parsing AIVA response: {e}") raise AIVAError(f"Failed to parse AIVA response: {e}") def summarize(self, document_content: str) -> Dict[str, str]: """ Summarizes the given document content, using the database cache or AIVA. Args: document_content: The content of the document to summarize. Returns: A dictionary containing the summary components (one_line_summary, key_points, technical_terms). """ try: self.connect_to_db() # Check if summary exists in the database cached_summary = self.get_summary_from_db(document_content) if cached_summary: return cached_summary # If not in the database, generate summary with AIVA summary = self.generate_summary_with_aiva(document_content) # Save the summary to the database self.save_summary_to_db(document_content, summary) return summary except (DatabaseConnectionError, AIVAError) as e: logging.error(f"Error during summarization: {e}") raise finally: self.close_db_connection() if __name__ == '__main__': # Example Usage (replace with your actual values) DB_HOST = os.getenv("DB_HOST", "postgresql-genesis-u50607.vm.elestio.app") DB_PORT = int(os.getenv("DB_PORT", "25432")) DB_NAME = os.getenv("DB_NAME", "your_db_name") DB_USER = os.getenv("DB_USER", "your_db_user") DB_PASSWORD = os.getenv("DB_PASSWORD", "your_db_password") AIVA_URL = os.getenv("AIVA_URL", "http://localhost:23405") summarizer = Summarizer(DB_HOST, DB_PORT, DB_NAME, DB_USER, DB_PASSWORD, AIVA_URL) document_content = "This is a sample document about the importance of documentation. Good documentation helps developers understand the code and maintain it easily. It also helps new team members get up to speed quickly." try: summary = summarizer.summarize(document_content) print("Summary:") print(f" One-Line Summary: {summary['one_line_summary']}") print(f" Key Points: {summary['key_points']}") print(f" Technical Terms: {summary['technical_terms']}") except Exception as e: print(f"An error occurred: {e}")