From a4fb413151574b659c68bf6b2a88b4de43a6e003 Mon Sep 17 00:00:00 2001 From: Andrew Ridgway Date: Tue, 3 Feb 2026 11:15:23 +1000 Subject: [PATCH] update git pull and provide retry logic --- src/ai_generators/ollama_md_generator.py | 233 +++++++++++++++++------ 1 file changed, 178 insertions(+), 55 deletions(-) diff --git a/src/ai_generators/ollama_md_generator.py b/src/ai_generators/ollama_md_generator.py index 5c5dd0b..4c93198 100644 --- a/src/ai_generators/ollama_md_generator.py +++ b/src/ai_generators/ollama_md_generator.py @@ -1,11 +1,17 @@ -import os, re, json, random, time, string -from ollama import Client +import json +import os +import random +import re +import string +import time +from concurrent.futures import ThreadPoolExecutor, TimeoutError + import chromadb from langchain_ollama import ChatOllama +from ollama import Client class OllamaGenerator: - def __init__(self, title: str, content: str, inner_title: str): self.title = title self.inner_title = inner_title @@ -14,16 +20,20 @@ class OllamaGenerator: print("In Class") print(os.environ["CONTENT_CREATOR_MODELS"]) try: - chroma_port = int(os.environ['CHROMA_PORT']) + chroma_port = int(os.environ["CHROMA_PORT"]) except ValueError as e: raise Exception(f"CHROMA_PORT is not an integer: {e}") - self.chroma = chromadb.HttpClient(host=os.environ['CHROMA_HOST'], port=chroma_port) - ollama_url = f"{os.environ["OLLAMA_PROTOCOL"]}://{os.environ["OLLAMA_HOST"]}:{os.environ["OLLAMA_PORT"]}" + self.chroma = chromadb.HttpClient( + host=os.environ["CHROMA_HOST"], port=chroma_port + ) + ollama_url = f"{os.environ['OLLAMA_PROTOCOL']}://{os.environ['OLLAMA_HOST']}:{os.environ['OLLAMA_PORT']}" self.ollama_client = Client(host=ollama_url) self.ollama_model = os.environ["EDITOR_MODEL"] self.embed_model = os.environ["EMBEDDING_MODEL"] self.agent_models = json.loads(os.environ["CONTENT_CREATOR_MODELS"]) - self.llm = ChatOllama(model=self.ollama_model, temperature=0.6, top_p=0.5) #This is the level head in the room + self.llm = ChatOllama( + model=self.ollama_model, temperature=0.6, top_p=0.5 + ) # This is the level head in the room self.prompt_inject = f""" You are a journalist, Software Developer and DevOps expert writing a 5000 word draft blog article for other tech enthusiasts. @@ -37,8 +47,8 @@ class OllamaGenerator: """ def split_into_chunks(self, text, chunk_size=100): - '''Split text into chunks of size chunk_size''' - words = re.findall(r'\S+', text) + """Split text into chunks of size chunk_size""" + words = re.findall(r"\S+", text) chunks = [] current_chunk = [] @@ -49,18 +59,19 @@ class OllamaGenerator: word_count += 1 if word_count >= chunk_size: - chunks.append(' '.join(current_chunk)) + chunks.append(" ".join(current_chunk)) current_chunk = [] word_count = 0 if current_chunk: - chunks.append(' '.join(current_chunk)) + chunks.append(" ".join(current_chunk)) return chunks def generate_draft(self, model) -> str: - '''Generate a draft blog post using the specified model''' - try: + """Generate a draft blog post using the specified model""" + + def _generate(): # the idea behind this is to make the "creativity" random amongst the content creators # contorlling temperature will allow cause the output to allow more "random" connections in sentences # Controlling top_p will tighten or loosen the embedding connections made @@ -69,56 +80,85 @@ class OllamaGenerator: temp = random.uniform(0.5, 1.0) top_p = random.uniform(0.4, 0.8) top_k = int(random.uniform(30, 80)) - agent_llm = ChatOllama(model=model, temperature=temp, top_p=top_p, top_k=top_k) + agent_llm = ChatOllama( + model=model, temperature=temp, top_p=top_p, top_k=top_k + ) messages = [ - ("system", "You are a creative writer specialising in writing about technology"), - ("human", self.prompt_inject ) + ( + "system", + "You are a creative writer specialising in writing about technology", + ), + ("human", self.prompt_inject), ] response = agent_llm.invoke(messages) - # self.response = self.ollama_client.chat(model=model, - # messages=[ - # { - # 'role': 'user', - # 'content': f'{self.prompt_inject}', - # }, - # ]) - #print ("draft") - #print (response) - return response.text()#['message']['content'] + return response.text() # ['message']['content'] - except Exception as e: - raise Exception(f"Failed to generate blog draft: {e}") + # Retry mechanism with 30-minute timeout + timeout_seconds = 30 * 60 # 30 minutes + max_retries = 3 + + for attempt in range(max_retries): + try: + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(_generate) + result = future.result(timeout=timeout_seconds) + return result + except TimeoutError: + print( + f"AI call timed out after {timeout_seconds} seconds on attempt {attempt + 1}" + ) + if attempt < max_retries - 1: + print("Retrying...") + time.sleep(5) # Wait 5 seconds before retrying + continue + else: + raise Exception( + f"AI call failed to complete after {max_retries} attempts with {timeout_seconds} second timeouts" + ) + except Exception as e: + if attempt < max_retries - 1: + print(f"Attempt {attempt + 1} failed with error: {e}. Retrying...") + time.sleep(5) # Wait 5 seconds before retrying + continue + else: + raise Exception( + f"Failed to generate blog draft after {max_retries} attempts: {e}" + ) def get_draft_embeddings(self, draft_chunks): - '''Get embeddings for the draft chunks''' + """Get embeddings for the draft chunks""" embeds = self.ollama_client.embed(model=self.embed_model, input=draft_chunks) - return embeds.get('embeddings', []) + return embeds.get("embeddings", []) def id_generator(self, size=6, chars=string.ascii_uppercase + string.digits): - return ''.join(random.choice(chars) for _ in range(size)) + return "".join(random.choice(chars) for _ in range(size)) def load_to_vector_db(self): - '''Load the generated blog drafts into a vector database''' - collection_name = f"blog_{self.title.lower().replace(" ", "_")}_{self.id_generator()}" - collection = self.chroma.get_or_create_collection(name=collection_name)#, metadata={"hnsw:space": "cosine"}) - #if any(collection.name == collectionname for collectionname in self.chroma.list_collections()): + """Load the generated blog drafts into a vector database""" + collection_name = ( + f"blog_{self.title.lower().replace(' ', '_')}_{self.id_generator()}" + ) + collection = self.chroma.get_or_create_collection( + name=collection_name + ) # , metadata={"hnsw:space": "cosine"}) + # if any(collection.name == collectionname for collectionname in self.chroma.list_collections()): # self.chroma.delete_collection("blog_creator") for model in self.agent_models: - print (f"Generating draft from {model} for load into vector database") + 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) + print(f"loading into collection") + collection.add( + documents=draft_chunks, embeddings=embeds, ids=ids, metadatas=metadata + ) return collection - def generate_markdown(self) -> str: - prompt_human = f""" You are an editor taking information from {len(self.agent_models)} Software Developers and Data experts @@ -133,29 +173,76 @@ class OllamaGenerator: The basis for the content of the blog is: {self.content} """ - try: - query_embed = self.ollama_client.embed(model=self.embed_model, input=prompt_human)['embeddings'] + + def _generate_final_document(): + query_embed = self.ollama_client.embed( + model=self.embed_model, input=prompt_human + )["embeddings"] collection = self.load_to_vector_db() - collection_query = collection.query(query_embeddings=query_embed, n_results=100) + 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]) - #print(pertinent_draft_info) + pertinent_draft_info = "\n\n".join( + collection.query(query_embeddings=query_embed, n_results=100)[ + "documents" + ][0] + ) + # 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 """ print("Generating final document") - messages = [("system", prompt_system), ("human", prompt_human),] - self.response = self.llm.invoke(messages).text() + messages = [ + ("system", prompt_system), + ("human", prompt_human), + ] + response = self.llm.invoke(messages).text() + return response + + try: + # Retry mechanism with 30-minute timeout + timeout_seconds = 30 * 60 # 30 minutes + max_retries = 3 + + for attempt in range(max_retries): + try: + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(_generate_final_document) + self.response = future.result(timeout=timeout_seconds) + break # Success, exit the retry loop + except TimeoutError: + print( + f"AI call timed out after {timeout_seconds} seconds on attempt {attempt + 1}" + ) + if attempt < max_retries - 1: + print("Retrying...") + time.sleep(5) # Wait 5 seconds before retrying + continue + else: + raise Exception( + f"AI call failed to complete after {max_retries} attempts with {timeout_seconds} second timeouts" + ) + except Exception as e: + if attempt < max_retries - 1: + print( + f"Attempt {attempt + 1} failed with error: {e}. Retrying..." + ) + time.sleep(5) # Wait 5 seconds before retrying + continue + else: + raise Exception( + f"Failed to generate markdown after {max_retries} attempts: {e}" + ) + # self.response = self.ollama_client.chat(model=self.ollama_model, # messages=[ - # { - # 'role': 'user', # 'content': f'{prompt_enhanced}', # }, # ]) - #print ("Markdown Generated") - #print (self.response) - return self.response#['message']['content'] + # print ("Markdown Generated") + # print (self.response) + return self.response # ['message']['content'] except Exception as e: raise Exception(f"Failed to generate markdown: {e}") @@ -165,6 +252,42 @@ class OllamaGenerator: f.write(self.generate_markdown()) def generate_system_message(self, prompt_system, prompt_human): - messages = [("system", prompt_system), ("human", prompt_human),] - ai_message = self.llm.invoke(messages).text() - return ai_message + def _generate(): + messages = [ + ("system", prompt_system), + ("human", prompt_human), + ] + ai_message = self.llm.invoke(messages).text() + return ai_message + + # Retry mechanism with 30-minute timeout + timeout_seconds = 30 * 60 # 30 minutes + max_retries = 3 + + for attempt in range(max_retries): + try: + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(_generate) + result = future.result(timeout=timeout_seconds) + return result + except TimeoutError: + print( + f"AI call timed out after {timeout_seconds} seconds on attempt {attempt + 1}" + ) + if attempt < max_retries - 1: + print("Retrying...") + time.sleep(5) # Wait 5 seconds before retrying + continue + else: + raise Exception( + f"AI call failed to complete after {max_retries} attempts with {timeout_seconds} second timeouts" + ) + except Exception as e: + if attempt < max_retries - 1: + print(f"Attempt {attempt + 1} failed with error: {e}. Retrying...") + time.sleep(5) # Wait 5 seconds before retrying + continue + else: + raise Exception( + f"Failed to generate system message after {max_retries} attempts: {e}" + )