local running, more fixes

src/generator/generator.py

@@ -7,9 +7,10 @@ from langchain.chains import ConversationalRetrievalChain
 from langchain.prompts import PromptTemplate
 from huggingface_hub import snapshot_download, hf_hub_download, HfApi
 from config.settings import Settings
-from tenacity import retry, stop_after_attempt, wait_exponential
+from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
 from datasets import load_dataset
 import sqlite3
+from openai import APIConnectionError, RateLimitError
 
 
 class LyricGenerator:
@@ -30,21 +31,36 @@ class LyricGenerator:
         print(f"Embeddings directory: {self.embeddings_dir}")
         print(f"Chroma directory: {self.chroma_dir}")
 
-        # Initialize OpenAI embeddings
+        # Initialize OpenAI embeddings with retry
         print("Setting up OpenAI embeddings...")
         if not Settings.OPENAI_API_KEY:
             raise RuntimeError(
                 "OpenAI API key is not set. Please configure it in your environment variables or HuggingFace Secrets.")
-        self.embeddings = OpenAIEmbeddings(
-            openai_api_key=Settings.OPENAI_API_KEY
-        )
-
+        
+        self.embeddings = self._create_embeddings_with_retry()
+        
         self.vector_store = None
         self.qa_chain = None
 
         # Load embeddings
         self._load_embeddings()
 
+    @retry(
+        retry=retry_if_exception_type((APIConnectionError, RateLimitError)),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        stop=stop_after_attempt(5)
+    )
+    def _create_embeddings_with_retry(self):
+        """Create OpenAI embeddings with retry logic"""
+        try:
+            return OpenAIEmbeddings(
+                openai_api_key=Settings.OPENAI_API_KEY,
+                timeout=60  # Increase timeout
+            )
+        except Exception as e:
+            print(f"Error creating embeddings: {type(e).__name__}: {str(e)}")
+            raise
+
     def _setup_embeddings_from_hf(self) -> None:
         """Download and setup embeddings from HuggingFace dataset"""
         print("\n=== Setting up embeddings from HuggingFace dataset ===")
@@ -293,19 +309,19 @@ class LyricGenerator:
         )
 
     @retry(
-        stop=stop_after_attempt(3),
-        wait=wait_exponential(multiplier=1, min=4, max=10)
+        retry=retry_if_exception_type((APIConnectionError, RateLimitError)),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        stop=stop_after_attempt(5)
     )
-    def _embed_query(self, text: str):
-        """Embed query with retry logic"""
+    def _similarity_search_with_retry(self, query: str, k: int = 5):
+        """Perform similarity search with retry logic"""
         try:
-            return self.embeddings.embed_query(text)
+            return self.vector_store.similarity_search_with_score(
+                query,
+                k=k
+            )
         except Exception as e:
-            print(f"Embedding error: {str(e)}")
-            if "api_key" in str(e).lower():
-                raise RuntimeError(
-                    "OpenAI API key validation failed. Please check your API key."
-                )
+            print(f"Similarity search error: {type(e).__name__}: {str(e)}")
             raise
 
     def generate_lyrics(
@@ -330,51 +346,23 @@ class LyricGenerator:
             print("Starting lyrics generation process...")
             print(f"Using OpenAI model: {Settings.LLM_MODEL}")
 
-            # Get source documents with scores first
+            # Get source documents
             print("Searching for similar documents...")
             try:
-                # Test embeddings function with better error handling
-                print("Testing embeddings function...")
-                try:
-                    test_embedding = self._embed_query("test")
-                    print(
-                        f"Embeddings function working (vector size: {len(test_embedding)})")
-                except Exception as e:
-                    print(f"OpenAI API Connection Error: {str(e)}")
-                    print("Checking API key...")
-                    if not Settings.OPENAI_API_KEY:
-                        raise RuntimeError(
-                            "OpenAI API key not found in environment")
-                    raise RuntimeError(
-                        f"OpenAI API connection failed: {str(e)}")
-
-                # Now try similarity search
-                docs_and_scores = self.vector_store.similarity_search_with_score(
-                    prompt,
-                    k=20
-                )
-                print(f"Found {len(docs_and_scores)} similar documents")
-
-                if not docs_and_scores:
-                    print(
-                        "Warning: No similar documents found. This may affect generation quality.")
-
+                results = self._similarity_search_with_retry(prompt)
+                # results is a list of (Document, score) tuples
+                docs_and_scores = [(doc[0], doc[1]) for doc in results]  # Unpack tuples correctly
             except Exception as e:
                 print(f"Error during similarity search: {str(e)}")
                 raise RuntimeError(f"Failed to search vector store: {str(e)}")
 
-            # Sort by similarity (convert distance to similarity)
-            docs_and_scores.sort(key=lambda x: x[1], reverse=False)
-
             # Create detailed context log
             context_details = []
             for doc, score in docs_and_scores[:5]:  # Log top 5 for brevity
-                similarity = round((1 - score) * 100, 2)
                 context_details.append({
                     'artist': doc.metadata['artist'],
                     'song': doc.metadata['song_title'],
-                    'similarity': similarity,
-                    # First 200 chars
+                    'similarity': f"{score:.2f}",
                     'content': doc.page_content[:200] + "..."
                 })