update git pull and provide retry logic

src/ai_generators/ollama_md_generator.py

@@ -1,11 +1,17 @@
-import os, re, json, random, time, string
+import json
-from ollama import Client
+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)
+        self.chroma = chromadb.HttpClient(
-        ollama_url = f"{os.environ["OLLAMA_PROTOCOL"]}://{os.environ["OLLAMA_HOST"]}:{os.environ["OLLAMA_PORT"]}"
+            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'''
+        """Split text into chunks of size chunk_size"""
-        words = re.findall(r'\S+', text)
+        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'''
+        """Generate a draft blog post using the specified model"""
-        try:
+
+        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'''
+        """Load the generated blog drafts into a vector database"""
-        collection_name = f"blog_{self.title.lower().replace(" ", "_")}_{self.id_generator()}"
+        collection_name = (
-        collection = self.chroma.get_or_create_collection(name=collection_name)#, metadata={"hnsw:space": "cosine"})
+            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')
+            print(f"loading into collection")
-            collection.add(documents=draft_chunks, embeddings=embeds, ids=ids, metadatas=metadata)
+            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),]
+            messages = [
-            self.response = self.llm.invoke(messages).text()
+                ("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}"
+                    )