update git pull and provide retry logic
This commit is contained in:
parent
7b160be3b7
commit
a4fb413151
GPG Key ID:
C8D9EAC514B47EF1
@ -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,38 +80,67 @@ 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']
|
||||
|
||||
# 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:
|
||||
raise Exception(f"Failed to generate blog draft: {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"})
|
||||
"""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:
|
||||
@ -111,14 +151,14 @@ class OllamaGenerator:
|
||||
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,23 +173,70 @@ class OllamaGenerator:
|
||||
The basis for the content of the blog is:
|
||||
<blog>{self.content}</blog>
|
||||
"""
|
||||
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])
|
||||
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: <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
|
||||
"""
|
||||
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}',
|
||||
# },
|
||||
# ])
|
||||
@ -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),]
|
||||
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}"
|
||||
)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user