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}"
+ )