pembot 0.0.3__py2.py3-none-any.whl

pembot/TextEmbedder/gemini_embedder.py

@@ -0,0 +1,27 @@
+from google.genai import types
+import requests
+import json
+import os
+from google import genai
+
+
+def get_embedding(client, text):
+
+    result = client.models.embed_content(
+        model="gemini-embedding-exp-03-07",
+        contents= text,
+        config= types.EmbedContentConfig(task_type= "RETRIEVAL_DOCUMENT"),
+    )
+
+    if result:
+        return result.embeddings[0].values
+    else:
+        print(f"Error: {result}")
+        #raise ValueError("Error communicating with the gemini embedding API")
+
+
+if __name__ == "__main__":
+
+    api_key = os.environ['GEMINI_API_KEY']
+    genai_client= genai.Client(api_key= api_key)
+    print(get_embedding(genai_client, "send bobs"))

pembot/TextEmbedder/mongodb_embedder.py

@@ -0,0 +1,258 @@
+from pathlib import Path
+import uuid
+from google import genai
+from google.genai import types
+from pymongo import MongoClient
+import os
+import re
+import time
+import numpy as np
+
+from pembot.AnyToText.convertor import chunk_text
+from pembot.utils.string_tools import make_it_an_id
+
+
+
+def clean_text(text):
+    # Use regex to find words with at least one alphanumeric character
+    cleaned_words = [word for word in text.split() if re.search(r'\w', word)]
+# Join the cleaned words back into a single string
+    cleaned_text = ' '.join(cleaned_words)
+    return cleaned_text
+
+
+
+def search_within_document(
+    db_client,
+    aggregate_query_embedding,
+    document_name_id: str,
+    limit: int = 5,
+    index_name: str = "test_search",
+    embeddings_collection_name: str= "doc_chunks",
+):
+    """
+    Performs a vector similarity search within the chunks of a specific document
+    in the 'embeddings_collection' MongoDB collection.
+
+    Args:
+        db_client: An initialized PyMongo Database instance.
+        aggregate_query_embedding: The np.mean of queries vectors of your search query.
+        document_name_id: This will be used to filter by the 'docId'.
+        limit: The maximum number of similar chunks to return.
+        index_name: The name of your MongoDB Atlas Vector Search index.
+                    You MUST have a vector search index created on the 'embedding' field
+                    of the 'embeddings_collection' collection for this to work efficiently.
+
+    Returns:
+        A list of dictionaries, where each dictionary represents a matching chunk
+        from the specified document, including its text, docId, and score.
+    """
+    embeddings_collection = db_client[embeddings_collection_name]
+
+    print(f"Searching within document (docId: {document_name_id})...")
+
+    # MongoDB Atlas Vector Search aggregation pipeline
+    # The 'path' should point to the field containing the embeddings.
+    # The 'filter' stage is crucial for searching within a specific document.
+    pipeline = [
+        {
+            '$vectorSearch': {
+                'queryVector': aggregate_query_embedding,
+                'path': 'embedding',
+
+                #number of nearest neighbors to consider
+                'numCandidates': 100,
+                'limit': limit,
+                'index': index_name,
+
+                #filter to search only within the specified document
+                'filter': {
+                    'docId': document_name_id
+                }
+            }
+        },
+
+        # to exclude the MongoDB internal _id
+        {
+            '$project': {
+                '_id': 0,
+                'docId': 1,
+                'chunk_number': 1,
+                'chunk_text': 1,
+                'score': { '$meta': 'vectorSearchScore' } # Get the similarity score
+            }
+        }
+    ]
+
+    results = list(embeddings_collection.aggregate(pipeline))
+
+    if not results:
+        print(f"No relevant chunks found for document '{document_name_id}' with the given query.")
+    else:
+        print(f"Found {len(results)} relevant chunks in document '{document_name_id}':")
+        for i, res in enumerate(results):
+            print(f"  Result {i+1} (Score: {res['score']:.4f}):")
+            print(f"    Chunk Number: {res['chunk_number']}")
+            print(f"    Text: '{res['chunk_text'][:100]}...'") # Print first 100 chars
+            print("-" * 30)
+
+    return results
+
+
+
+def process_document_and_embed(db_client,
+    llm_client,
+    inference_client,
+    file_path: Path,
+    chunk_size: int,
+    embedding_model: str = 'nomic-embed-text:v1.5',
+    embeddings_collection_name= "doc_chunks",
+    use_custom_id: str | None = None,
+    use_custom_input: str | None = None
+) -> list[dict]:
+    """
+    Processes an input document by chunking its text, generating embeddings using
+    Ollama's specified embedding model, and storing these embeddings and chunks
+    in a MongoDB collection.
+
+    Args:
+        db_client: An initialized PyMongo Database instance.
+        file_path: The original path of the document being processed.
+                   This path will be used to create a sanitized ID for the
+                   document.
+        chunk_size: The desired chunk size in words for text segmentation.
+        embedding_model: The name of the Ollama embedding model to use.
+    """
+
+    input_text= None
+    if use_custom_input is not None:
+        input_text= use_custom_input
+    else:
+        # Read the input text from the file
+        with open(str(file_path), "r") as md_file:
+            input_text = md_file.read()
+
+
+    document_name_id= None
+    if use_custom_id is not None:
+        document_name_id= use_custom_id
+    else:
+        # Create a valid ID for the document from the file name (without extension)
+        file_root = os.path.splitext(file_path.name)[0]
+        document_name_id = make_it_an_id(file_root)
+
+    # Reference the MongoDB collection where chunks will be stored
+    # This single collection will serve as the global 'embeddings_collection'
+    # and document-specific data can be queried using 'docId'.
+    embeddings_collection = db_client[embeddings_collection_name]
+
+    # Check if this document's embeddings already exist in MongoDB
+    # We check if any document with this docId exists in the 'embeddings_collection' collection.
+    an_existing_chunk= embeddings_collection.find_one({'docId': document_name_id})
+    if an_existing_chunk:
+        print(f"Document '{file_path.name}' (ID: {document_name_id}) already processed. Skipping.")
+        return [an_existing_chunk]
+
+    print(f"Processing document '{file_path.name}' (ID: {document_name_id})...")
+
+    embed_locally= False
+    try:
+        models = llm_client.list()
+
+        for model in models.models:
+            if model.model == embedding_model:
+                embed_locally= True
+    except Exception as e:
+        print("local model list error: ", e)
+
+
+    # Chunk the input text into smaller segments
+    chunks = chunk_text(input_text, chunk_size)
+    print(f"Text chunked into {len(chunks)} segments.")
+
+    # Process each chunk: generate embedding and add to MongoDB
+    res = []
+    for i, chunk in enumerate(chunks):
+        try:
+            print(f"Processing chunk {i+1}/{len(chunks)} for document '{file_path.name}'...")
+            # Generate embedding using the specified Ollama model
+            print("embedding_model is: ", embedding_model)
+            print("if statement is", 'gemini' in embedding_model)
+
+            if 'gemini' in embedding_model:
+
+                client = genai.Client(api_key= os.environ['GEMINI_API_KEY'])
+                result = client.models.embed_content(
+                        model= embedding_model,
+                        contents= chunk,
+                        config=types.EmbedContentConfig(task_type="RETRIEVAL_DOCUMENT")
+                )
+                if result is not None and result.embeddings:
+                    embedding= result.embeddings[0].values
+                else:
+                    raise ValueError("Gemini not givingz embeddingzzz")
+
+                # API safety
+                time.sleep(1)
+
+            elif embed_locally:
+                response = llm_client.embeddings(model=embedding_model, prompt=chunk)
+                embedding= response['embedding']
+
+            else:
+                embedding = inference_client.feature_extraction(chunk, model=embedding_model)
+
+                # API rate limiting safety
+                time.sleep(1)
+                print("zzzzzzzzz")
+
+            if isinstance(embedding, np.ndarray):
+                embedding = embedding.tolist()
+            elif not isinstance(embedding, list):
+                raise TypeError("Embedding is not a list or numpy array. Cannot store in MongoDB.")
+
+
+            # Generate a random suffix for unique chunk IDs
+            random_suffix = uuid.uuid4().hex[:8]
+
+            # Create the unique ID for the chunk for the global 'embeddings_collection' collection
+            # This ID can be stored in the document if needed for external reference
+            chunk_id_global = f"{document_name_id}_chunk_{i + 1}_{random_suffix}"
+            # Create the unique ID for the chunk (local to the document's chunking)
+            chunk_id_doc_specific = f"chunk_{i}_{random_suffix}"
+
+            # Store the chunk data in MongoDB using update_one with upsert=True
+            # This will insert a new document if 'docId' and 'chunk_number' don't match,
+            # or update an existing one if they do.
+            doc_set = {
+                'chunk_text': chunk,
+                'embedding': embedding,
+                'chunk_id_global': chunk_id_global,
+                'chunk_id_doc_specific': chunk_id_doc_specific
+            }
+            embeddings_collection.update_one(
+                {'docId': document_name_id, 'chunk_number': i + 1},
+                {'$set': doc_set},
+                upsert=True
+            )
+            print(f"Successfully stored chunk {i+1} for '{file_path.name}' in MongoDB.")
+            res.append({**doc_set, "docId": document_name_id, "chunk_number": i + 1})
+
+        except Exception as e:
+            print(f"Error processing chunk {i+1} for '{file_path.name}': {e}")
+            # Continue to the next chunk even if one fails
+            continue
+
+    print(f"Finished processing document '{file_path.name}'. All chunks embedded and stored in MongoDB.")
+    return res
+
+
+
+
+if __name__ == "__main__":
+    client = MongoClient(os.environ['MONGODB_PEM'])
+    db = client["pem"]
+    collection = db["blogs"]
+    embeddings_collection = db["blog_embeddings"]
+
+    api_key = os.environ['GEMINI_API_KEY']

pembot/TextEmbedder/mongodb_index_creator.py

@@ -0,0 +1,133 @@
+from pymongo import MongoClient
+from pymongo.collection import Collection
+from pymongo.operations import SearchIndexModel
+import time
+import os
+
+def create_vector_index(collection: Collection, index_name: str, num_dimensions: int = 768):
+    """
+    Creates a MongoDB Atlas Vector Search index if it does not already exist.
+
+    Args:
+        collection: The PyMongo Collection object on which to create the index.
+        index_name: The desired name for the vector search index.
+        num_dimensions: The number of dimensions for the embedding vectors.
+    """
+    
+    # 1. Check if the index already exists
+    existing_indexes = list(collection.list_search_indexes())
+    
+    for index in existing_indexes:
+        if index.get('name') == index_name:
+            print(f"Search index '{index_name}' already exists. Skipping creation.")
+            
+            # Optional: You can also check if the existing index is "READY"
+            if index.get('status') == 'READY':
+                print(f"Index '{index_name}' is already ready for querying.")
+            else:
+                print(f"Index '{index_name}' is currently {index.get('status')}. Polling for readiness...")
+                # If it exists but is not ready, proceed to the polling loop
+                _wait_for_index_ready(collection, index_name)
+            return # Exit the function if the index already exists
+
+    # 2. If the index does not exist, proceed to create it
+    print(f"Search index '{index_name}' does not exist. Creating it now...")
+
+    search_index_model = SearchIndexModel(definition={
+            "fields": [
+                {
+                    "type": "vector",
+                    "path": "embedding",
+                    "similarity": "dotProduct", # Or "cosine", "euclidean"
+                    "numDimensions": num_dimensions,
+                    "quantization": "scalar" # Or "none"
+                },
+                {
+                    "type": "filter",
+                    "path": "docId"
+                }
+            ]
+        },
+        name=index_name,
+        type="vectorSearch"
+    )
+
+    try:
+        # Create the search index
+        result = collection.create_search_index(model=search_index_model)
+        print("New search index named " + result + " is building.")
+
+        # Wait for initial sync to complete
+        _wait_for_index_ready(collection, index_name)
+        print(result + " is ready for querying.")
+
+    except Exception as e:
+        print(f"Error creating or waiting for search index '{index_name}': {e}")
+        # Depending on the error, you might want to re-raise or handle it differently.
+
+def _wait_for_index_ready(collection: Collection, index_name: str):
+    """
+    Helper function to poll the index status until it's ready.
+    """
+    print("Polling to check if the index is ready. This may take some time (up to a few minutes for large indexes).")
+    
+    start_time = time.time()
+    timeout = 300 # 5 minutes timeout, adjust as needed
+
+    while True:
+        indices= None
+        try:
+            indices = list(collection.list_search_indexes(name=index_name))
+            if len(indices) == 1 and indices[0].get("status") == "READY":
+                print(f"Index '{index_name}' is READY for querying.")
+                break
+            elif len(indices) == 0:
+                print(f"Warning: Index '{index_name}' not found during polling, might have failed creation or name mismatch.")
+                break # Exit if index disappears
+            else:
+                current_status = indices[0].get("status", "UNKNOWN")
+                print(f"Index '{index_name}' status: {current_status}. Waiting...")
+        except Exception as e:
+            print(f"Error while polling index status: {e}. Retrying...")
+        
+        if time.time() - start_time > timeout:
+            status= indices[0].get('status') if indices else 'N/A'
+            print(f"Timeout: Index '{index_name}' did not become ready within {timeout} seconds. Current status: {status}")
+            break # Exit on timeout
+
+        time.sleep(10) # Poll less frequently, every 10 seconds
+
+# --- Example Usage ---
+if __name__ == "__main__":
+    
+    # Replace with your database and collection names
+    DATABASE_NAME = "pembot" 
+    COLLECTION_NAME = "doc_chunks"
+    VECTOR_INDEX_NAME = "test_search"
+
+    # Connect to MongoDB
+    mongo_client= None
+
+    try:
+        mongo_client = MongoClient(os.environ["MONGODB_PEM"])
+        db = mongo_client[DATABASE_NAME]
+        collection = db[COLLECTION_NAME]
+        print("Connected to MongoDB successfully.")
+
+        # Define dimensions (e.g., for nomic-embed-text:v1.5)
+        EMBEDDING_DIMENSIONS = 768 # Check your model's output dimension
+
+        # Call the function to create the index, with existence check
+        create_vector_index(collection, VECTOR_INDEX_NAME, num_dimensions=EMBEDDING_DIMENSIONS)
+        
+        # Test calling it again to see the "already exists" message
+        create_vector_index(collection, VECTOR_INDEX_NAME, num_dimensions=EMBEDDING_DIMENSIONS)
+
+    except Exception as e:
+        print(f"Failed to connect to MongoDB or process: {e}")
+    finally:
+        if 'mongo_client' in locals() and mongo_client:
+            mongo_client.close()
+            print("MongoDB connection closed.")
+
+

pembot/TextEmbedder/vector_query.py

@@ -0,0 +1,64 @@
+import os
+from google import genai
+from pymongo import MongoClient
+from TextEmbedder import gemini_embedder, mongodb_embedder
+from os import environ
+import re
+
+
+def vectory_query_run(genai_client, docs_collection, embeddings_collection, search_string, database, index_name, chunksize= 170):
+
+    query_embedding = gemini_embedder.get_embedding(genai_client, search_string)
+
+    # Define the aggregation pipeline
+    pipeline = [
+        {
+            "$vectorSearch": {
+                "queryVector": query_embedding,
+                "index": f"{index_name}",
+                "path": "embedding",
+                "numCandidates": 3,
+                "limit": 3,
+                "k": 3  # Number of results to return
+            }
+        }
+    ]
+
+    # Execute the aggregation pipeline
+    results = list(embeddings_collection.aggregate(pipeline))
+    for i, res in enumerate(results):
+        doc = docs_collection.find_one({"_id": res['docId']})
+
+        wordslist = mongodb_embedder.clean_text(doc['content']).split()
+
+        cn = int(res['chunk_number'])
+        search_hit = " ".join(wordslist[cn * chunksize: (cn+1) * chunksize])
+
+        print(f"RESULT #{i + 1}")
+        print(search_hit)
+        print("\n\n")
+        
+
+
+if __name__ == "__main__":
+    mc = MongoClient(environ['MONGODB_PEM'])
+
+    api_key = os.environ['GEMINI_API_KEY']
+    genai_client= genai.Client(api_key= api_key)
+
+    try:
+        mc.admin.command('ping')  # This is a simple way to check the connection
+        database = mc["pem"]
+
+        while True:
+            query = input("enter search query [Enter 'exit' to exit]: ")
+            if query.lower() == 'exit':
+                print('bye!')
+                break
+            docs_collection= database["blogs"]
+            embeddings_collection= database["blog_embeddings"]
+            vectory_query_run(genai_client, docs_collection, embeddings_collection, query, database, index_name= 'vector_index_blog_chunks')
+
+    finally:
+        mc.close()
+

pembot/__init__.py

@@ -0,0 +1,6 @@
+"""
+A Python Package to convert PEM blog content to usseful information by leveraging LLMs
+"""
+__version__ = '0.0.3'
+from .main import save_to_json_file, make_query
+__all__ = ["save_to_json_file", "make_query"]

pembot/config/config.yaml

@@ -0,0 +1,5 @@
+OUTPUT_DIR: /home/cyto/dev/pembotdir
+PAGE_DELIMITER: ___________________________ NEXT PAGE ___________________________
+app:
+  name: pembot
+  version: 0.0.3

pembot/gartner.py

@@ -0,0 +1,140 @@
+from selenium.webdriver.common.by import By
+from selenium.webdriver.support.ui import WebDriverWait
+from selenium.webdriver.support import expected_conditions as EC
+from bs4 import BeautifulSoup
+import json
+import time
+import random
+import undetected_chromedriver as uc                                                                                                                                                      
+driver = uc.Chrome(headless=False,use_subprocess=False)
+
+# Random delays between actions
+def human_delay():
+    time.sleep(random.uniform(1, 3))
+
+# Use this before each navigation
+driver.get("https://www.gartner.com/en/insights")
+human_delay()
+
+# Wait for the dynamic content to load
+wait = WebDriverWait(driver, 20)
+try:
+    print("waiting")
+    wait.until(EC.presence_of_element_located((By.CLASS_NAME, 'row.dynamic-content')))
+    # wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, 'div.row.dynamic-content')))
+    print("waitdone")
+except Exception as e:
+    print("Timed out waiting for page to load")
+    driver.quit()
+    exit()
+
+# Parse the page with BeautifulSoup
+soup = BeautifulSoup(driver.page_source, 'html.parser')
+main_div = soup.find('div', class_='row dynamic-content')
+category_blocks = main_div.find_all('div', class_='category-block') if main_div else []
+
+print("category blocks: ", category_blocks)
+
+result = {}
+
+for block in category_blocks:
+    # Extract category name
+    a_tag = block.find('a')
+    if not a_tag:
+        continue
+    h4 = a_tag.find('h4', class_='categoryheadline')
+    if not h4:
+        continue
+    category_name = h4.get_text(strip=True)
+    result[category_name] = {}
+    print("result", result)
+    
+    # Extract all article links in the category
+    items_content = block.find('div', class_='categoryItemsContent')
+    if not items_content:
+        continue
+    links = items_content.find_all('a', href=True)
+    
+    for link in links:
+        href = link['href']
+        article_title = link.get_text(strip=True)
+        
+        # Form absolute URL
+        full_url = f'https://www.gartner.com{href}' if href.startswith('/') else href
+        
+        # Open article in a new tab
+        original_window = driver.current_window_handle
+        driver.switch_to.new_window('tab')
+        try:
+            human_delay()
+            driver.get(full_url)
+            human_delay()
+            # Wait for article content to load
+            article_wait = WebDriverWait(driver, 15)
+            article_wait.until(EC.presence_of_element_located((By.TAG_NAME, 'article')))
+            
+            # Parse article content
+            # Modify the article content extraction part of your script:
+
+            # Inside the article processing block after driver.get(full_url):
+            article_soup = BeautifulSoup(driver.page_source, 'html.parser')
+
+            # Extract article title
+            title = article_soup.find('h1').get_text(strip=True) if article_soup.find('h1') else "Untitled"
+
+            # Extract authors and date
+            byline = ""
+            authors = []
+            date = ""
+            byline_element = article_soup.find('article', class_='article-text').find('span', class_='rte')
+            if byline_element:
+                byline_text = byline_element.get_text(strip=True, separator='|').split('|')
+                for part in byline_text:
+                    if 'By ' in part:
+                        authors = [a.strip() for a in part.replace('By ', '').split(' and ')]
+                    if any(word in part.lower() for word in ['jan', 'feb', 'mar', 'apr', 'may', 'jun', 
+                                                            'jul', 'aug', 'sep', 'oct', 'nov', 'dec']):
+                        date = part.strip()
+
+            # Format title with metadata
+            formatted_title = f"{title}"
+            if authors or date:
+                formatted_title += " ("
+                if authors:
+                    formatted_title += f"by:{','.join(authors)}"
+                if date:
+                    if authors: formatted_title += ";"
+                    formatted_title += f"on:{date}"
+                formatted_title += ")"
+
+            # Extract main content
+            content_parts = []
+            main_articles = article_soup.find_all('article', class_='article-text')
+            for article in main_articles:
+                # Skip the byline article
+                if 'By ' in article.get_text(): continue
+                
+                # Get all text elements including headings
+                for element in article.find_all(['p', 'h2', 'h3', 'h4']):
+                    text = element.get_text(strip=True, separator=' ')
+                    if text:
+                        content_parts.append(text)
+
+            full_content = '\n'.join([t for t in content_parts if t.strip()])
+
+            # Store in dictionary structure
+            result[category_name][formatted_title] = full_content
+        except Exception as e:
+            print(f"Error retrieving {full_url}: {str(e)}")
+            result[category_name].append('')
+        finally:
+            driver.close()
+            driver.switch_to.window(original_window)
+            time.sleep(1)  # Polite delay
+
+# Save the results to a JSON file
+with open('gartner_articles.json', 'w', encoding='utf-8') as f:
+    print("dumping...")
+    json.dump(result, f, ensure_ascii=False, indent=4)
+
+driver.quit()