From 3cd0654306cb1b3ac7d2da4f7c57f2cb2fb027e6 Mon Sep 17 00:00:00 2001 From: Andrew Ridgway Date: Tue, 3 Feb 2026 12:21:17 +1000 Subject: [PATCH] syntax fixes for retry logic --- src/ai_generators/ollama_md_generator.py | 141 ++++++++++++++++++----- 1 file changed, 113 insertions(+), 28 deletions(-) diff --git a/src/ai_generators/ollama_md_generator.py b/src/ai_generators/ollama_md_generator.py index 4c93198..bb2b4cc 100644 --- a/src/ai_generators/ollama_md_generator.py +++ b/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: {pertinent_draft_info} - 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