pembot 0.0.3__py2.py3-none-any.whl

pembot/pdf2markdown/requirements.txt

@@ -0,0 +1,8 @@
+PyMuPDF
+pdfplumber
+pytesseract
+opencv-python
+numpy
+transformers
+torch
+Pillow

pembot/pem.py

@@ -0,0 +1,157 @@
+from pathlib import Path
+
+from huggingface_hub import InferenceClient
+import ollama
+from pymongo import MongoClient, ReturnDocument
+from pembot.TextEmbedder.mongodb_embedder import process_document_and_embed
+from pembot.TextEmbedder.mongodb_index_creator import create_vector_index
+import os
+from enum import Enum
+from datetime import datetime
+
+class Category(Enum):
+    WEBDEV     = "webdev"
+    ANDROID    = "android"
+    PHYSICS    = "physics"
+    SPACE      = "space"
+    PHILOSOPHY = "philosophy"
+    CULTURE    = "culture"
+    HISTORY    = "history"
+    MATHS      = "maths"
+    SOCIOLOGY  = "sociology"
+    CODING     = "coding"
+
+def get_title(input_text: str) -> str:
+    #   1) splits `input_text` on "\n" → an array
+    #   2) takes the [0]th element → the first line
+    #   3) trims that element
+
+    s= input_text.split('\n')[0].strip()
+
+    if s.startswith("#"):
+        return s[1:].strip()
+    return s.strip()
+
+
+def embed_all_docs(database, llm_client, inference_client, chunk_size, embedding_model, embeddings_collection_name, index_name, docs_collection):
+    for doc in docs_collection.find({}, {"_id": 1, "content": 1}):
+
+        print("on doc: ", doc)
+
+        docs_embed= process_document_and_embed(db_client= database, llm_client= llm_client,
+            inference_client= inference_client, file_path= Path(), chunk_size= chunk_size,
+            embedding_model= embedding_model, embeddings_collection_name= embeddings_collection_name,
+            use_custom_id= str(doc["_id"]), use_custom_input= doc["content"])
+
+        doc_embedding_length= len(docs_embed[-1]["embedding"])
+
+        # ensuring
+        create_vector_index(database[embeddings_collection_name], index_name, num_dimensions= doc_embedding_length)
+
+def update_from_local_fs_using_title(database, llm_client, inference_client, docs_dir, chunk_size, embedding_model, embeddings_collection_name, index_name, docs_collection):
+
+    for docfile in docs_dir.iterdir():
+        # for each blog in blog dir, update it if its not alr there in blogs coll,
+        # that that created/updated doc's id
+        # and update the embeddings in blog_embeddings collection with docId as that id
+        # the docId of the embedding doc must be id of the corresponding blog
+        if docfile.is_file :
+            # TODO: create/update doc and get doc._id, pass that as the optional string param "use_custom_id"
+
+            input_text= ""
+            with open(str(docfile), "r") as md_file:
+                input_text = md_file.read()
+
+            categories= list([[i + 1, x, x.value] for i, x in enumerate(Category)])
+            category= "culture" # default
+            while True:
+                category_index = input("Enter category no. for {}...\n{}\n: ".format(
+                                    input_text[:300],
+                                    '\n'.join(map(lambda x: str(x[0]) + ". " + str(x[2]), categories))
+                                ))
+                if category_index.isdigit():
+                    category= categories[int(category_index) - 1][2]
+                    break
+                else:
+                    print("pls enter a number.")
+
+            set_on_insert= {
+                "name": os.getenv("PROFILE_NAME"),
+                "email": os.getenv("PROFILE_EMAIL"),
+                "category": category,
+                "dateAdded": datetime.now(),
+                "views": 0,
+                "likes": 0
+            }
+
+            title= get_title(input_text)
+
+            updated= docs_collection.find_one_and_update(
+                {'title': title},
+                {
+                    '$set': {'content': input_text},
+                    '$setOnInsert': set_on_insert
+                },
+                upsert=True,
+                return_document=ReturnDocument.AFTER
+            )
+
+            docs_embed= process_document_and_embed(db_client= database, llm_client= llm_client,
+                inference_client= inference_client, file_path= docfile, chunk_size= chunk_size,
+                embedding_model= embedding_model, embeddings_collection_name= embeddings_collection_name,
+                use_custom_id= str(updated["_id"]))
+            doc_embedding_length= len(docs_embed[-1]["embedding"])
+
+            # ensuring
+            create_vector_index(database[embeddings_collection_name], index_name, num_dimensions= doc_embedding_length)
+            print("in the db now!")
+
+if __name__ == "__main__":
+
+    mongodb_uri= os.environ['MONGODB_PEM']
+    mc = MongoClient(mongodb_uri)
+
+    llm_client= ollama.Client()
+
+    #### FOR USING JINA INSTEAD OF HUGGINGFACE SDK, REPLACE WITH THE InferenceClient TOP IMPORT
+    # from pembot.utils.inference_client import InferenceClient
+    # JINA_API_KEY= os.environ['JINA_API_KEY']
+    # inference_client= InferenceClient(
+    #     provider="Jina AI",
+    #     api_key= JINA_API_KEY,
+    # )
+
+    inference_client= InferenceClient(
+        provider="hf-inference",
+        api_key= os.environ["HF_TOKEN"],
+    )
+
+    mc.admin.command('ping')
+    print("ping test ok")
+    database = mc["pem"]
+    print("dbs and cols loaded")
+
+    embeddings_collection_name: str= "blog_chunks"
+    docs_collection_name: str= "blogs"
+    docs_collection= database[docs_collection_name]
+
+    # nerfed, but provided by hf serverless inference: BAAI/bge-small-en-v1.5
+    # Worth mentioning:
+    # jinaai/jina-embeddings-v3
+    # BAAI/bge-base-en-v1.5
+    # nomic-ai/nomic-embed-text-v1.5
+    # embedding_model: str= 'BAAI/bge-base-en-v1.5'
+    embedding_model: str= 'gemini-embedding-exp-03-07'
+
+    index_name: str=  "gemini_vectors"
+
+    # the number depends on the amount of chunks that will go into LLMs in the end
+    total_input_chars_allowed= 2_50_000 # we got 63k tokens => ~2.5 lac characters
+    no_of_chunks_to_be_sent = 2 #len(required_fields)
+    chunk_size= int( total_input_chars_allowed / no_of_chunks_to_be_sent)
+
+    # curdir <- project dir <- dev -> (~) -> Documents -> notes -> blogs
+    docs_dir = Path.cwd().parent.parent.parent / "Documents" / "notes" / "Obsidian" / "blogs"
+
+    embed_all_docs(database, llm_client, inference_client, chunk_size, embedding_model, embeddings_collection_name, index_name, docs_collection)
+    # update_from_local_fs_using_title(database, llm_client, inference_client, docs_dir, chunk_size, embedding_model, embeddings_collection_name, index_name, docs_collection)

pembot/query.py

@@ -0,0 +1,204 @@
+from os import environ
+from huggingface_hub.inference._generated.types.chat_completion import ChatCompletionOutputMessage
+from huggingface_hub.inference._providers import PROVIDER_T
+import ollama
+import re
+
+from pydantic_core.core_schema import ErrorType
+from pembot.TextEmbedder.mongodb_embedder import search_within_document
+import numpy as np
+from huggingface_hub import InferenceClient
+from google import genai
+from google.genai import types
+import time
+
+from pembot.TextEmbedder.mongodb_index_creator import create_vector_index
+
+def external_llm(rag_prompt, model_name, llm_provider_name: PROVIDER_T= "novita", inference_client = None) -> str:
+
+    # Here, one can change the provider of the inference LLM if
+    # for embedding we are using one which doesnt have our LLM available
+    # or, is costly, so we choose different, just here in the function header, or from the main()
+
+    if not inference_client:
+        inference_client= InferenceClient(
+                # "nebius" "novita" "hyperbolic"
+                provider= llm_provider_name,
+                api_key= environ["HF_TOKEN"]
+            )
+
+    completion= inference_client.chat.completions.create(
+            model= model_name,
+            messages= [
+                {"role": "user", "content": rag_prompt}
+            ]
+        )
+    response_message: ChatCompletionOutputMessage= completion.choices[0].message
+
+    if response_message.content:
+        return response_message.content
+    else:
+        return '{}'
+
+
+def multi_embedding_average(llm_client, inference_client, descriptions, model= "BAAI/bge-en-icl", embed_locally= False):
+
+    description_embeddings = []
+    for desc in descriptions:
+        try:
+            if 'gemini' in model:
+                client = genai.Client(api_key= environ['GEMINI_API_KEY'])
+                result = client.models.embed_content(
+                        model= model,
+                        contents= desc,
+                        config=types.EmbedContentConfig(task_type="RETRIEVAL_DOCUMENT")
+                )
+                if result is not None and result.embeddings:
+                    description_embeddings.append(result.embeddings[0].values)
+                else:
+                    raise ValueError("Gemini not givingz embeddingzzz")
+            elif embed_locally:
+                response = llm_client.embeddings(model=model, prompt=desc)
+                description_embeddings.append(response['embedding'])
+
+            else:
+                response = inference_client.feature_extraction(desc, model=model)
+                description_embeddings.append(response)
+
+        except Exception as e:
+            print(f"Error generating embedding for description '{desc}': {e}")
+            # Decide how to handle errors: skip, raise, or use a placeholder
+            continue
+        time.sleep(1)
+
+    if not description_embeddings:
+        print("No embeddings could be generated for the descriptions. Aborting search.")
+        return []
+
+    # Aggregate embeddings: A simple approach is to average them.
+    # This creates a single query vector that represents the combined meaning.
+    return np.mean(description_embeddings, axis=0).tolist()
+
+
+
+def rag_query_llm(db_client, llm_client, inference_client, user_query: str, document_id: str, required_fields_descriptions: list[str], model_name: str = "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B", ollama_base_url: str = "http://localhost:11434", no_of_fields= 4, embedding_model= "BAAI/bge-en-icl", llm_provider_name: PROVIDER_T= "novita", index_name: str= "test_search", embeddings_collection= "doc_chunks"):
+    """
+    Performs a RAG (Retrieval Augmented Generation) query using a Hugging Face
+    embedding model, ChromaDB for retrieval, and a local Ollama model for generation.
+
+    Args:
+        db_client: The vector DB client
+        user_query (str): The user's query.
+        required_fields_descriptions: The required fields which are to be queried from context
+        model_name (str): The name of the Ollama model to use (e.g., "llama2", "mistral").
+        no_of_fields (str): number of vectors which are to be retrieved from DB
+
+    Returns:
+        str: The generated response from the Ollama model.
+    """
+
+    models = llm_client.list()
+    found= False
+    embed_locally= False
+
+    for model in models.models:
+        # print(model.model)
+        if model.model == model_name:
+            found= True
+        if model.model == embedding_model:
+            embed_locally= True
+
+
+
+    aggregate_query_embedding= multi_embedding_average(llm_client, inference_client, required_fields_descriptions, model= embedding_model, embed_locally= embed_locally)
+    print("Aggregate query embedding generated. length: ", len(aggregate_query_embedding))
+
+    create_vector_index(db_client[embeddings_collection], index_name, num_dimensions= len(aggregate_query_embedding))
+
+    # check the order of args
+    relevant_chunks= search_within_document(db_client, aggregate_query_embedding, document_id, limit= no_of_fields, index_name= index_name, embeddings_collection_name= embeddings_collection)
+    relevant_chunks= list(map(lambda x: x['chunk_text'], relevant_chunks))
+
+    if not relevant_chunks:
+        context = "No relevant context available."
+    else:
+        # print(f"Found {len(relevant_chunks)} relevant chunks.")
+        # Concatenate relevant chunks into a single context string
+        context = "\n\n".join(relevant_chunks)
+
+    # Construct the RAG prompt
+    rag_prompt = f"""
+    You are a helpful assistant. Use the following context to answer the question.
+    If you don't know the answer, just say that you don't know, don't try to make up an answer.
+
+    Context:
+    {context}
+
+    Question: {user_query}
+    """
+
+    # print("Step 3: Calling Ollama model with RAG prompt...")
+    # print("final prompt: ")
+    # print(rag_prompt)
+    if 'gemini' in model_name:
+
+        client = genai.Client(api_key= environ['GEMINI_API_KEY'])
+        response = client.models.generate_content(
+            model= model_name,
+            contents= rag_prompt,
+        )
+        return response.text
+
+    elif found:
+        try:
+            # You can use ollama.chat or ollama.generate depending on your model and preference
+            # ollama.chat is generally preferred for conversational models.
+            response = llm_client.chat(
+                model=model_name,
+                messages=[{'role': 'user', 'content': rag_prompt}],
+                options={"base_url": ollama_base_url} # Ensure the base URL is set
+            )
+            return response['message']['content']
+        except ollama.ResponseError as e:
+            print(f"Error calling Ollama API: {e}")
+            return f"Error: Could not get a response from Ollama. Please check if Ollama is running and the model '{model_name}' is pulled."
+        except Exception as e:
+            print(f"An unexpected error occurred while calling Ollama: {e}")
+            return "An unexpected error occurred."
+    elif 'qwen' in model_name or 'gemma' in model_name or 'Qwen' in model_name or 'deepseek' in model_name:
+        return external_llm(rag_prompt, model_name= model_name, llm_provider_name= llm_provider_name)
+    else:
+        return '{}'
+
+
+
+def remove_bs(text):
+    """
+    Removes everything between <think></think> tags and any text outside of JSON curly brackets.
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The string with text between <think></think> tags removed and only the
+             content within the outermost JSON curly brackets.
+             Returns an empty string if no valid JSON is found.
+    """
+    # 1. Remove <think></think> tags
+    think_pattern = r'<think>.*?</think>'
+    text_without_think = re.sub(think_pattern, '', text, flags=re.DOTALL)
+
+    # 2. Extract JSON content
+    # This regex looks for the first opening curly brace and the last closing curly brace.
+    # It assumes the JSON structure is well-formed within the string.
+    json_match = re.search(r'\{(.*)\}', text_without_think, re.DOTALL)
+
+    if json_match:
+        json_content_str = "{" + json_match.group(1) + "}"
+        return json_content_str
+    else:
+        return ""
+
+
+if __name__ == "__main__":
+    print("hemlo worls")

pembot/utils/inference_client.py

@@ -0,0 +1,132 @@
+import os
+import requests
+import json
+from enum import Enum
+
+# Define an Enum for providers if you plan to support more in the future
+class PROVIDER_T(Enum):
+    JINA = "Jina AI"
+    # OLLAMA = "Ollama" # Example for future expansion
+
+class InferenceClient:
+    """
+    A client for making inference calls to various AI model providers.
+    Currently supports Jina AI for feature extraction (embeddings).
+    """
+
+    def __init__(self, provider: PROVIDER_T, api_key: str):
+        """
+        Initializes the InferenceClient.
+
+        Args:
+            provider: The AI model provider (e.g., PROVIDER_T.JINA).
+            api_key: The API key for the specified provider.
+        """
+        self.provider = provider
+        self.api_key = api_key
+        self.base_url = self._get_base_url(provider)
+
+    def _get_base_url(self, provider: PROVIDER_T) -> str:
+        """
+        Returns the base URL for the given provider's API.
+        """
+        if provider == PROVIDER_T.JINA:
+            return "https://api.jina.ai/v1/embeddings"
+        # Add more providers here as needed
+        # elif provider == PROVIDER_T.OLLAMA:
+        #     return "http://localhost:11434/api/embeddings"
+        else:
+            raise ValueError(f"Unsupported provider: {provider}")
+
+    def feature_extraction(self, prompt_texts: list[str], model: str = "jina-embeddings-v3") -> list[list[float]]:
+        """
+        Calls the Jina AI API to get embeddings (feature extraction) for input texts.
+
+        Args:
+            prompt_texts: A list of strings for which to generate embeddings.
+            model: The name of the Jina AI embedding model to use.
+                   Defaults to "jina-embeddings-v3".
+
+        Returns:
+            A list of embedding vectors (list of floats).
+
+        Raises:
+            Exception: If the API call fails or returns an error.
+        """
+        if self.provider != PROVIDER_T.JINA:
+            raise NotImplementedError(f"Feature extraction not implemented for provider: {self.provider}")
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}"
+        }
+        
+        # Jina API specific payload
+        data = {
+            "model": model,
+            "input": prompt_texts,
+            "task": "qa" # Recommended task for general purpose embeddings, can be "text-matching", "retrieval", etc.
+                         # "qa" is often good for general question-answering/retrieval contexts.
+        }
+
+        try:
+            response = requests.post(self.base_url, headers=headers, data=json.dumps(data))
+            response.raise_for_status() # Raise an exception for HTTP errors (4xx or 5xx)
+            
+            response_json = response.json()
+            if "data" in response_json and isinstance(response_json["data"], list):
+                # Extract just the embedding vectors
+                embeddings = [item["embedding"] for item in response_json["data"]]
+                return embeddings
+            else:
+                raise Exception(f"Unexpected API response format: {response_json}")
+
+        except requests.exceptions.HTTPError as http_err:
+            print(f"HTTP error occurred: {http_err} - {response.text}")
+            raise
+        except requests.exceptions.ConnectionError as conn_err:
+            print(f"Connection error occurred: {conn_err}")
+            raise
+        except requests.exceptions.Timeout as timeout_err:
+            print(f"Timeout error occurred: {timeout_err}")
+            raise
+        except requests.exceptions.RequestException as req_err:
+            print(f"An unexpected request error occurred: {req_err}")
+            raise
+        except Exception as e:
+            print(f"An error occurred during Jina AI API call: {e}")
+            raise
+
+# --- Example Usage ---
+if __name__ == "__main__":
+    try:
+        JINA_API_KEY = os.environ['JINA_API_KEY']
+    except KeyError:
+        print("Error: JINA_API_KEY environment variable not set.")
+        print("Please set it (e.g., export JINA_API_KEY='your_api_key') and try again.")
+        exit(1)
+
+    # Initialize the InferenceClient for Jina AI
+    jina_client = InferenceClient(provider=PROVIDER_T.JINA, api_key=JINA_API_KEY)
+
+    # Example input texts
+    input_texts = [
+        "Organic skincare for sensitive skin with aloe vera and chamomile.",
+        "Bio-Hautpflege für empfindliche Haut mit Aloe Vera und Kamille.",
+        "Cuidado de la piel orgánico para piel sensible con aloe vera y manzanilla.",
+        "新しいメイクのトレンドは鮮やかな色と革新的な技術に焦点を当てています。"
+    ]
+
+    print("--- Calling Jina AI for Feature Extraction ---")
+    try:
+        embeddings = jina_client.feature_extraction(prompt_texts=input_texts)
+        print(f"Successfully extracted {len(embeddings)} embeddings.")
+        print(f"First embedding (first 5 values): {embeddings[0][:5]}...")
+        print(f"Embedding dimension: {len(embeddings[0])}")
+
+        # Example with a different model (if available and you want to specify)
+        # embeddings_v2 = jina_client.feature_extraction(prompt_texts=input_texts, model="jina-embeddings-v2")
+        # print(f"\nSuccessfully extracted {len(embeddings_v2)} embeddings with jina-embeddings-v2.")
+
+    except Exception as e:
+        print(f"Failed to extract features: {e}")

pembot/utils/string_tools.py

@@ -0,0 +1,45 @@
+import re
+import random
+
+def make_it_an_id(file_name):
+    """
+    Input: A file name mixed with spaces, periods, etc.
+    Output: '_' separated smallcase alphabetic id, with random filler if less than 5 chars
+    """
+    # 1. Convert to lowercase
+    file_name = file_name.lower()
+
+    # 2. Replace non-alphanumeric characters (except periods and spaces) with underscores
+    # Keep periods for now to split by later, and spaces for initial conversion
+    cleaned_name = re.sub(r'[^a-z0-9\s\.]', '_', file_name)
+    
+    # 3. Replace spaces and periods with underscores
+    cleaned_name = re.sub(r'[\s\.]+', '_', cleaned_name)
+
+    # 4. Remove leading/trailing underscores and multiple consecutive underscores
+    cleaned_name = re.sub(r'_{2,}', '_', cleaned_name).strip('_')
+
+    # Ensure it only contains alphabetic characters (after previous cleaning)
+    # If the file_name was something like "123.pdf", this step ensures we only keep alphabetic parts.
+    # We will filter out non-alphabetic parts after initial cleaning to retain some structure.
+    # Let's refine this to ensure we only keep alphabetic parts before padding.
+    alphabetic_parts = re.findall(r'[a-z]+', cleaned_name)
+    
+    # Join alphabetic parts with underscores
+    result_id = '_'.join(alphabetic_parts)
+
+    # 5. Add random filler if less than 5 chars
+    if len(result_id) < 5:
+        # Generate random lowercase alphabetic characters
+        filler_length = 5 - len(result_id)
+        random_filler = ''.join(random.choice('abcdefghijklmnopqrstuvwxyz') for _ in range(filler_length))
+        
+        # Append filler with an underscore if result_id is not empty
+        if result_id:
+            result_id += '_' + random_filler
+        else: # If the result_id is empty (e.g., from "123.txt" or "$.*"), just use the filler
+            result_id = random_filler
+            
+    return result_id
+
+

pembot-0.0.3.dist-info/METADATA

@@ -0,0 +1,8 @@
+Metadata-Version: 2.4
+Name: pembot
+Version: 0.0.3
+Summary: A Python Package to convert PEM blog content to usseful information by leveraging LLMs
+Author-email: cyto <aryan_sidhwani@protonmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Project-URL: Home, https://github.com/silverstone-git/pem-rag-chatbot