src/ai_generators/ollama_md_generator.py

@@ -91,7 +91,9 @@ class OllamaGenerator:
                 ("human", self.prompt_inject),
             ]
             response = agent_llm.invoke(messages)
-            return response.text()  # ['message']['content']
+            return (
+                response.text if hasattr(response, "text") else str(response)
+            )  # ['message']['content']
 
         # Retry mechanism with 30-minute timeout
         timeout_seconds = 30 * 60  # 30 minutes
@@ -127,8 +129,26 @@ class OllamaGenerator:
 
     def get_draft_embeddings(self, draft_chunks):
         """Get embeddings for the draft chunks"""
-        embeds = self.ollama_client.embed(model=self.embed_model, input=draft_chunks)
-        return embeds.get("embeddings", [])
+        try:
+            # Handle empty draft chunks
+            if not draft_chunks:
+                print("Warning: No draft chunks to embed")
+                return []
+
+            embeds = self.ollama_client.embed(
+                model=self.embed_model, input=draft_chunks
+            )
+            embeddings = embeds.get("embeddings", [])
+
+            # Check if embeddings were generated successfully
+            if not embeddings:
+                print("Warning: No embeddings generated")
+                return []
+
+            return embeddings
+        except Exception as e:
+            print(f"Error generating embeddings: {e}")
+            return []
 
     def id_generator(self, size=6, chars=string.ascii_uppercase + string.digits):
         return "".join(random.choice(chars) for _ in range(size))
@@ -145,16 +165,52 @@ class OllamaGenerator:
         #    self.chroma.delete_collection("blog_creator")
         for model in self.agent_models:
             print(f"Generating draft from {model} for load into vector database")
-            draft_chunks = self.split_into_chunks(self.generate_draft(model))
-            print(f"generating embeds")
-            embeds = self.get_draft_embeddings(draft_chunks)
-            ids = [model + str(i) for i in range(len(draft_chunks))]
-            chunknumber = list(range(len(draft_chunks)))
-            metadata = [{"model_agent": model} for index in chunknumber]
-            print(f"loading into collection")
-            collection.add(
-                documents=draft_chunks, embeddings=embeds, ids=ids, metadatas=metadata
-            )
+            try:
+                draft_content = self.generate_draft(model)
+                draft_chunks = self.split_into_chunks(draft_content)
+
+                # Skip if no content was generated
+                if not draft_chunks or all(
+                    chunk.strip() == "" for chunk in draft_chunks
+                ):
+                    print(f"Skipping {model} - no content generated")
+                    continue
+
+                print(f"generating embeds for {model}")
+                embeds = self.get_draft_embeddings(draft_chunks)
+
+                # Skip if no embeddings were generated
+                if not embeds:
+                    print(f"Skipping {model} - no embeddings generated")
+                    continue
+
+                # Ensure we have the same number of embeddings as chunks
+                if len(embeds) != len(draft_chunks):
+                    print(
+                        f"Warning: Mismatch between chunks ({len(draft_chunks)}) and embeddings ({len(embeds)}) for {model}"
+                    )
+                    # Truncate or pad to match
+                    min_length = min(len(embeds), len(draft_chunks))
+                    draft_chunks = draft_chunks[:min_length]
+                    embeds = embeds[:min_length]
+                    if min_length == 0:
+                        print(f"Skipping {model} - no valid content/embeddings pairs")
+                        continue
+
+                ids = [model + str(i) for i in range(len(draft_chunks))]
+                chunknumber = list(range(len(draft_chunks)))
+                metadata = [{"model_agent": model} for index in chunknumber]
+                print(f"loading into collection for {model}")
+                collection.add(
+                    documents=draft_chunks,
+                    embeddings=embeds,
+                    ids=ids,
+                    metadatas=metadata,
+                )
+            except Exception as e:
+                print(f"Error processing model {model}: {e}")
+                # Continue with other models rather than failing completely
+                continue
 
         return collection
 
@@ -175,19 +231,47 @@ class OllamaGenerator:
             """
 
         def _generate_final_document():
-            query_embed = self.ollama_client.embed(
-                model=self.embed_model, input=prompt_human
-            )["embeddings"]
+            try:
+                embed_result = self.ollama_client.embed(
+                    model=self.embed_model, input=prompt_human
+                )
+                query_embed = embed_result.get("embeddings", [])
+                if not query_embed:
+                    print(
+                        "Warning: Failed to generate query embeddings, using empty list"
+                    )
+                    query_embed = [[]]  # Use a single empty embedding as fallback
+            except Exception as e:
+                print(f"Error generating query embeddings: {e}")
+                # Generate empty embeddings as fallback
+                query_embed = [[]]  # Use a single empty embedding as fallback
+
             collection = self.load_to_vector_db()
-            collection_query = collection.query(
-                query_embeddings=query_embed, n_results=100
-            )
-            print("Showing pertinent info from drafts used in final edited edition")
-            pertinent_draft_info = "\n\n".join(
-                collection.query(query_embeddings=query_embed, n_results=100)[
-                    "documents"
-                ][0]
-            )
+
+            # Try to query the collection, with fallback for empty collections
+            try:
+                collection_query = collection.query(
+                    query_embeddings=query_embed, n_results=100
+                )
+                print("Showing pertinent info from drafts used in final edited edition")
+
+                # Get documents with error handling
+                query_result = collection.query(
+                    query_embeddings=query_embed, n_results=100
+                )
+                documents = query_result.get("documents", [])
+
+                if documents and len(documents) > 0 and len(documents[0]) > 0:
+                    pertinent_draft_info = "\n\n".join(documents[0])
+                else:
+                    print("Warning: No relevant documents found in collection")
+                    pertinent_draft_info = "No relevant information found in drafts."
+
+            except Exception as query_error:
+                print(f"Error querying collection: {query_error}")
+                pertinent_draft_info = (
+                    "No relevant information found in drafts due to query error."
+                )
             # print(pertinent_draft_info)
             prompt_system = f"""Generate the final, 5000 word, draft of the blog using this information from the drafts: <context>{pertinent_draft_info}</context>
                             - Only output in markdown, do not wrap in markdown tags, Only provide the draft not a commentary on the drafts in the context
@@ -197,8 +281,8 @@ class OllamaGenerator:
                 ("system", prompt_system),
                 ("human", prompt_human),
             ]
-            response = self.llm.invoke(messages).text()
-            return response
+            response = self.llm.invoke(messages)
+            return response.text if hasattr(response, "text") else str(response)
 
         try:
             # Retry mechanism with 30-minute timeout
@@ -257,7 +341,8 @@ class OllamaGenerator:
                 ("system", prompt_system),
                 ("human", prompt_human),
             ]
-            ai_message = self.llm.invoke(messages).text()
+            response = self.llm.invoke(messages)
+            ai_message = response.text if hasattr(response, "text") else str(response)
             return ai_message
 
         # Retry mechanism with 30-minute timeout