mb-rag 1.1.58__tar.gz → 1.1.66__tar.gz

{mb_rag-1.1.58 → mb_rag-1.1.66}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mb_rag
-Version: 1.1.58
+Version: 1.1.66
 Summary: RAG function file
 Author: ['Malav Bateriwala']
 Requires-Python: >=3.8

{mb_rag-1.1.58 → mb_rag-1.1.66}/README.md

@@ -73,6 +73,9 @@ ollama_model = ModelFactory(
 )
 response = ollama_model.invoke_query("What is the meaning of life?")
 
+## Running in threads 
+response = model.invoke_query_threads(query_list=['q1','q2'],input_data=[[images_data],[images_data]],n_workers=4)
+
 
 ## check example_conversation.ipynb for more details
 

{mb_rag-1.1.58 → mb_rag-1.1.66}/mb_rag/basic.py

@@ -2,6 +2,7 @@
 
 import os
 from langchain_core.messages import HumanMessage
+import torch
 from mb_rag.utils.extra import check_package
 import base64
 from .utils.extra import check_package
@@ -41,7 +42,7 @@ class ModelFactory:
         
         self.model_type = model_type
         self.model_name = model_name
-        model_data = creators.get(model_type)
+        model_data = creators[model_type] if model_type in creators else None
         if not model_data:
             raise ValueError(f"Unsupported model type: {model_type}")
         
@@ -195,7 +196,7 @@ class ModelFactory:
         from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, AutoModelForImageTextToText,AutoProcessor
         import torch
         
-        device = torch.device(device) if torch.cuda.is_available() else torch.device("cpu")
+        device = torch.device("cuda" if torch.cuda.is_available() and device == "cuda" else "cpu")
         
         temperature = kwargs.pop("temperature", 0.7)
         max_length = kwargs.pop("max_length", 1024)
@@ -301,6 +302,7 @@ class ModelFactory:
 
         df = pd.DataFrame(query_list, columns=["query"])
         df["response"] = None
+        df["input_data"] = None if input_data is None else input_data
 
         structured_model = None
         if pydantic_model is not None:
@@ -355,6 +357,8 @@ class ModelFactory:
         Returns:
             str: Output from the model
         """
+        if not isinstance(images, list):
+            images = [images]
         base64_images = [self._image_to_base64(image) for image in images]
         image_prompt_create = [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_images[i]}"}} for i in range(len(images))]
         prompt_new = [{"type": "text", "text": prompt}, *image_prompt_create]

mb_rag-1.1.66/mb_rag/prompts_bank.py

@@ -0,0 +1,77 @@
+from langchain_core.prompts.chat import ChatPromptTemplate
+
+__all__ = ["PromptManager"]
+
+class PromptManager:
+    """
+    Central class for storing and invoking prompt templates.
+
+    Example:
+        pm = PromptManager()
+        prompt_text = pm.render_prompt("greeting")
+        print(prompt_text)
+
+        pm = PromptManager()
+        prompt_text = pm.render_prompt("todo_task", {"task": "Plan a deep learning project for image recognition"})
+        print(prompt_text)
+    """
+
+    def __init__(self):
+        self.templates = {
+            "coding_python": """You are a Python developer.
+Human: {question}
+Assistant:""",
+
+            "greeting": """You are a friendly assistant.
+Human: Hello!
+Assistant: Hi there! How can I assist you today?""",
+
+            "goodbye": """You are a friendly assistant.
+Human: Goodbye!
+Assistant: Goodbye! Have a great day!""",
+
+            "todo_task": """You are a helpful assistant.
+Human: Please create a to-do list for the following task: {task}
+Assistant:""",
+
+            "map_function": "*map(lambda x: image_url, baseframes_list)",
+            
+            "SQL_AGENT_SYS_PROMPT": """You are an expert SQL agent. Your task is to generate and execute SQL queries based on user requests.
+RULES:
+- THINK step by step before answering.
+- Use the provided database schema to inform your queries.
+- When you need to retrieve data, generate a SQL query and execute it using the provided tools.
+- Read-only mode: Do not attempt to modify the database.
+- NO INSERT/UPDATE/DELETE/ALTER/DROP/CREATE/REPLACE/TRUNCATE statements allowed.
+- LIMIT your results to 10 rows. Unless specified otherwise.
+- If you encounter an error while executing a query, analyze the error message and adjust your query accordingly.
+- Prefer using explicit column names instead of SELECT * for better performance.    
+- Always ensure your SQL syntax is correct. """
+        }
+
+    def get_template(self, name: str) -> str:
+        """
+        Get a prompt template by name.
+        Args:
+            name (str): The key name of the prompt.
+        Returns:
+            str: The prompt template string.
+        """
+        template = self.templates.get(name)
+        if not template:
+            raise ValueError(f"Prompt '{name}' not found. Available prompts: {list(self.templates.keys())}")
+        return template
+
+    def render_prompt(self, name: str, context: dict = None) -> str:
+        """
+        Fill and return a rendered prompt string.
+        Args:
+            name (str): The key name of the prompt.
+            context (dict): Variables to fill into the template.
+        Returns:
+            str: The final rendered prompt text.
+        """
+        template = self.get_template(name)
+        chat_prompt = ChatPromptTemplate.from_template(template)
+        rendered = chat_prompt.invoke(context or {})
+        return rendered.to_string()

mb_rag-1.1.66/mb_rag/utils/llm_wrapper.py

@@ -0,0 +1,31 @@
+## simple llm wrapper to replace invoke with invoke_query/own model query
+
+__all__ = ["LLMWrapper"]
+
+class LLMWrapper:
+    """A simple wrapper for the language model to standardize the invoke method.
+    """
+
+    def __init__(self, llm):
+        self.llm = llm
+
+    def __getattr__(self, name):
+        """Get all attributes from llm module. (invoke_query, invoke_query_threads, etc.)"""
+        return getattr(self.llm, name)
+
+    def invoke(self, use_threads=False,**kwargs) -> str:
+        """
+        Invoke the language model with a list of messages.
+        Using invoke_query method of the underlying model.
+        Check ModelFactory for more details.
+        
+        Args:
+            use_threads (bool): Whether to use threading for invocation. Defaults to False.
+            **kwargs: Keyword arguments for the model invocation.
+
+        Returns:
+            str: The generated response.
+        """
+        if use_threads:
+            return self.llm.invoke_query_threads(**kwargs)
+        return self.llm.invoke_query(**kwargs)

mb_rag-1.1.66/mb_rag/utils/viewer.py

@@ -0,0 +1,8 @@
+from IPython.display import Image, display
+
+__all__ = ["display_graph_png"]
+
+
+def display_graph_png(graph):
+    """Display a graph using IPython's display capabilities."""
+    display(Image(graph.get_graph().draw_mermaid_png()))

{mb_rag-1.1.58 → mb_rag-1.1.66}/mb_rag/version.py

@@ -1,5 +1,5 @@
 MAJOR_VERSION = 1
 MINOR_VERSION = 1
-PATCH_VERSION = 58
+PATCH_VERSION = 66
 version = '{}.{}.{}'.format(MAJOR_VERSION, MINOR_VERSION, PATCH_VERSION)
 __all__  = ['MAJOR_VERSION', 'MINOR_VERSION', 'PATCH_VERSION', 'version']

{mb_rag-1.1.58 → mb_rag-1.1.66}/mb_rag.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mb_rag
-Version: 1.1.58
+Version: 1.1.66
 Summary: RAG function file
 Author: ['Malav Bateriwala']
 Requires-Python: >=3.8

{mb_rag-1.1.58 → mb_rag-1.1.66}/mb_rag.egg-info/SOURCES.txt

@@ -3,6 +3,7 @@ pyproject.toml
 setup.py
 mb_rag/__init__.py
 mb_rag/basic.py
+mb_rag/prompts_bank.py
 mb_rag/version.py
 mb_rag.egg-info/PKG-INFO
 mb_rag.egg-info/SOURCES.txt
@@ -12,7 +13,6 @@ mb_rag.egg-info/top_level.txt
 mb_rag/chatbot/__init__.py
 mb_rag/chatbot/chains.py
 mb_rag/chatbot/conversation.py
-mb_rag/chatbot/prompts.py
 mb_rag/rag/__init__.py
 mb_rag/rag/embeddings.py
 mb_rag/utils/__init__.py
@@ -20,4 +20,6 @@ mb_rag/utils/all_data_extract.py
 mb_rag/utils/bounding_box.py
 mb_rag/utils/document_extract.py
 mb_rag/utils/extra.py
-mb_rag/utils/pdf_extract.py
+mb_rag/utils/llm_wrapper.py
+mb_rag/utils/pdf_extract.py
+mb_rag/utils/viewer.py

mb_rag-1.1.58/mb_rag/chatbot/prompts.py

@@ -1,59 +0,0 @@
-## file for storing basic prompts template
-from langchain.prompts import ChatPromptTemplate
-
-__all__ = ["prompts", "invoke_prompt"]
-
-class prompts:
-    """
-    Class to get different prompts example for chatbot and templates
-    """
-
-    def get_code_prompts(self):
-        """
-        Get code prompts
-        Returns:
-            str: Code prompt
-        """
-        list_code_prompts = {'coding_python ': """You are a Python developer.
-                                                Human: {}"""}
-
-    def get_text_prompts(self):
-        """
-        Get text prompts
-        Returns:
-            str: Text prompt
-        """
-        list_text_prompts = {
-            'multiple_placeholders': """You are a helpful assistant.
-                                        Human: Tell me a more about {adjective1} and its relation to {adjective2}.
-                                        Assistant:"""
-        }
-
-    def get_image_prompts(self):
-        """
-        Get image prompts
-        Returns:
-            str: Image prompt
-        """
-        list_image_prompts = {'map_function': "*map(lambda x: image_url, baseframes_list)"} # for passing multiple images from a video or a list of images
-
-    def get_assistant_prompts(self):
-        """
-        Get assistant prompts
-        Returns:
-            str: Assistant prompt
-        """
-        list_assistant_prompts = {}
-
-def invoke_prompt(template: str, input_dict : dict = None):
-    """
-    Invoke a prompt
-    Args:
-        template (str): Template for the prompt
-        input_dict (dict): Input dictionary for the prompt
-    Returns:
-        str: Prompt
-    """
-    prompt_multiple = ChatPromptTemplate.from_template(template)
-    prompt = prompt_multiple.invoke(input_dict)
-    return prompt