integrating agentic chroma

src/ai_generators/ollama_md_generator.py

@@ -1,5 +1,6 @@
-import os
+import os, re
 from ollama import Client
+import chromadb, time
 
 
 class OllamaGenerator:
@@ -7,16 +8,16 @@ class OllamaGenerator:
     def __init__(self, title: str, content: str, model: str):
         self.title = title
         self.content = content
+        self.chroma = chromadb.HttpClient(host="172.18.0.2", port=8000)
         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 = model
-
-    def generate_markdown(self) -> str:
-
-        prompt = f"""
-            You are a Software Developer and DevOps expert
+        self.embed_model = "snowflake-arctic-embed2:latest"
+        self.agent_models = ["openthinker:7b", "deepseek-r1:7b", "qwen2.5:7b", "deepseek-coder-v2:16b"]
+        self.prompt_inject = f"""
+            You are a journalist Software Developer and DevOps expert
             who has transistioned in Developer Relations 
-            writing a 1000 word blog for other tech enthusiast.
+            writing a 1000 word draft blog for other tech enthusiasts.
             You like to use almost no code examples and prefer to talk
             in a light comedic tone. You are also Australian 
             As this person write this blog as a markdown document.
@@ -24,14 +25,101 @@ class OllamaGenerator:
             Do not output the title in the markdown.
             The basis for the content of the blog is:
                 {self.content}
-            Only output markdown DO NOT GENERATE AN EXPLANATION
+            """
+
+    def split_into_chunks(self, text, chunk_size=100):
+        '''Split text into chunks of size chunk_size'''
+        words = re.findall(r'\S+', text)
+
+        chunks = []
+        current_chunk = []
+        word_count = 0
+
+        for word in words:
+            current_chunk.append(word)
+            word_count += 1
+
+        if word_count >= chunk_size:
+            chunks.append(' '.join(current_chunk))
+            current_chunk = []
+            word_count = 0
+
+        if 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:
+            self.response = self.ollama_client.chat(model=model,
+                                                    messages=[
+                    {
+                        'role': 'user',
+                        'content': f'{self.prompt_inject}',
+                    },
+                ])
+            return self.response['message']['content']
+
+        except Exception as e:
+            raise Exception(f"Failed to generate blog draft: {e}")
+        
+    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', [])
+
+
+    def load_to_vector_db(self):
+        '''Load the generated blog drafts into a vector database'''
+        collection_name = f"blog_{self.title.lower().replace(" ", "_")}"
+        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")
+            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)
+
+        return collection
+    
+
+    def generate_markdown(self) -> str:
+        
+        prompt = f"""
+            You are an editor taking information from {len(self.agent_models)} Software 
+            Developers and Data experts 
+            who have transistioned into Developer Relations 
+            writing a 3000 word blog for other tech enthusiasts.
+            You like when they use almost no code examples and the 
+            voice is in a light comedic tone. You are also Australian 
+            As this person produce and an amalgamtion of this blog as a markdown document.
+            The title for the blog is {self.title}.
+            Do not output the title in the markdown.
+            The basis for the content of the blog is:
+                {self.content}
             """
         try:
+            query_embed = self.ollama_client.embed(model=self.embed_model, input=prompt)['embeddings']
+            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")
+            for document in collection_query:
+                print (document)
+            pertinent_draft_info = '\n\n'.join(collection.query(query_embeddings=query_embed, n_results=100)['documents'][0])
+            prompt_enhanced = f"{prompt} - Generate the final document using this information from the drafts: {pertinent_draft_info} - ONLY OUTPUT THE MARKDOWN"
+            print("Generating final document")
             self.response = self.ollama_client.chat(model=self.ollama_model,
                                                     messages=[
                     {
                         'role': 'user',
-                        'content': f'{prompt}',
+                        'content': f'{prompt_enhanced}',
                     },
                 ])
             return self.response['message']['content']

src/main.py

@@ -17,6 +17,6 @@ for note in tril_notes:
     print("Generating Document")
     ai_gen = omg.OllamaGenerator(tril_notes[note]['title'],
                                  tril_notes[note]['content'],
-                                 "openthinker:7b")
+                                 "qwen2.5:7b")
     os_friendly_title = convert_to_lowercase_with_underscores(tril_notes[note]['title'])
     ai_gen.save_to_file(f"/blog_creator/generated_files/{os_friendly_title}.md")