mb-rag 1.0.117__tar.gz → 1.0.124__tar.gz

{mb_rag-1.0.117 → mb_rag-1.0.124}/PKG-INFO

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

mb_rag-1.0.124/mb_rag/chatbot/basic.py

@@ -0,0 +1,395 @@
+import os
+from dotenv import load_dotenv
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from IPython.display import display, HTML
+from typing import Optional, List, Dict, Any, Union
+from mb_rag.utils.extra import check_package
+import base64
+
+__all__ = [
+    'ChatbotBase',
+    'ModelFactory',
+    'ConversationModel',
+    'IPythonStreamHandler'
+]
+
+class ChatbotBase:
+    """Base class for chatbot functionality"""
+    
+    @staticmethod
+    def load_env(file_path: str) -> None:
+        """
+        Load environment variables from a file
+        Args:
+            file_path (str): Path to the environment file
+        """
+        load_dotenv(file_path)
+    
+    @staticmethod
+    def add_os_key(name: str, key: str) -> None:
+        """
+        Add an API key to the environment
+        Args:
+            name (str): Name of the API key
+            key (str): API key
+        """
+        os.environ[name] = key
+
+    @staticmethod
+    def get_client():
+        """
+        Returns a boto3 client for S3
+        """
+        if not check_package("boto3"):
+            raise ImportError("Boto3 package not found. Please install it using: pip install boto3")
+        
+        import boto3
+        return boto3.client('s3')
+
+class ModelFactory:
+    """Factory class for creating different types of chatbot models"""
+    
+    def __init__(self, model_type: str = 'openai', model_name: str = "gpt-4o", **kwargs) -> Any:
+        """
+        Factory method to create any type of model
+        Args:
+            model_type (str): Type of model to create
+            model_name (str): Name of the model
+            **kwargs: Additional arguments
+        Returns:
+            Any: Chatbot model
+        """
+        creators = {
+            'openai': self.create_openai,
+            'anthropic': self.create_anthropic,
+            'google': self.create_google,
+            'ollama': self.create_ollama
+        }
+        
+        model_data = creators.get(model_type)
+        if not model_data:
+            raise ValueError(f"Unsupported model type: {model_type}")
+        
+        try:
+            self.model = model_data(model_name, **kwargs)
+        except Exception as e:
+            raise ValueError(f"Error creating {model_type} model: {str(e)}")
+        
+    @classmethod
+    def create_openai(cls, model_name: str = "gpt-4o", **kwargs) -> Any:
+        """
+        Create OpenAI chatbot model
+        Args:
+            model_name (str): Name of the model
+            **kwargs: Additional arguments
+        Returns:
+            ChatOpenAI: Chatbot model
+        """
+        if not check_package("openai"):
+            raise ImportError("OpenAI package not found. Please install it using: pip install openai langchain-openai")
+        
+        from langchain_openai import ChatOpenAI
+        kwargs["model_name"] = model_name
+        return ChatOpenAI(**kwargs)
+
+    @classmethod
+    def create_anthropic(cls, model_name: str = "claude-3-opus-20240229", **kwargs) -> Any:
+        """
+        Create Anthropic chatbot model
+        Args:
+            model_name (str): Name of the model
+            **kwargs: Additional arguments
+        Returns:
+            ChatAnthropic: Chatbot model
+        """
+        if not check_package("anthropic"):
+            raise ImportError("Anthropic package not found. Please install it using: pip install anthropic langchain-anthropic")
+        
+        from langchain_anthropic import ChatAnthropic
+        kwargs["model_name"] = model_name
+        return ChatAnthropic(**kwargs)
+
+    @classmethod
+    def create_google(cls, model_name: str = "gemini-1.5-flash", **kwargs) -> Any:
+        """
+        Create Google chatbot model
+        Args:
+            model_name (str): Name of the model
+            **kwargs: Additional arguments
+        Returns:
+            ChatGoogleGenerativeAI: Chatbot model
+        """
+        if not check_package("google.generativeai"):
+            raise ImportError("Google Generative AI package not found. Please install it using: pip install google-generativeai langchain-google-genai")
+        
+        from langchain_google_genai import ChatGoogleGenerativeAI
+        kwargs["model"] = model_name
+        return ChatGoogleGenerativeAI(**kwargs)
+
+    @classmethod
+    def create_ollama(cls, model_name: str = "llama3", **kwargs) -> Any:
+        """
+        Create Ollama chatbot model
+        Args:
+            model_name (str): Name of the model
+            **kwargs: Additional arguments
+        Returns:
+            Ollama: Chatbot model
+        """
+        if not check_package("langchain_community"):
+            raise ImportError("Langchain Community package not found. Please install it using: pip install langchain-community")
+        
+        from langchain_community.llms import Ollama
+        kwargs["model"] = model_name
+        return Ollama(**kwargs)
+
+    def invoke_query(self,query: str,get_content_only: bool = True,images: list = None,pydantic_model = None) -> str:
+        """
+        Invoke the model
+        Args:
+            query (str): Query to send to the model
+            get_content_only (bool): Whether to return only content
+            images (list): List of images to send to the model
+            pydantic_model: Pydantic model for structured output
+        Returns:
+            str: Response from the model
+        """
+        
+        if pydantic_model is not None:
+            try:
+                self.model = self.model.with_structured_output(pydantic_model)
+            except Exception as e:
+                raise ValueError(f"Error with pydantic_model: {e}")
+        if images:
+            res = self._model_invoke_images(images=images,prompt=query,pydantic_model=pydantic_model)
+        else:
+            res = self.model.invoke(query)
+            if get_content_only:
+                try:
+                    return res.content
+                except Exception:
+                    return res
+        return res
+
+    def _image_to_base64(self,image):
+        with open(image, "rb") as f:
+            return base64.b64encode(f.read()).decode('utf-8')
+
+    def _model_invoke_images(self,images: list, prompt: str,pydantic_model = None):
+        """
+        Function to invoke the model with images
+        Args:
+            model (ChatOpenAI): Chatbot model
+            images (list): List of images
+            prompt (str): Prompt
+            pydantic_model (PydanticModel): Pydantic model
+    Returns:
+        str: Output from the model
+    """
+        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,]
+        if pydantic_model is not None:
+            try:
+                self.model = self.model.with_structured_output(pydantic_model)
+            except Exception as e:
+                print(f"Error with pydantic_model: {e}")
+                print("Continuing without structured output")
+        message= HumanMessage(content=prompt_new,)
+        response = self.model.invoke([message])
+        return response.content
+
+
+class ConversationModel:
+    """
+    A class to handle conversation with AI models
+    
+    Attributes:
+        chatbot: The AI model for conversation
+        message_list (List): List of conversation messages
+        file_path (str): Path to save/load conversations. Can be local or S3
+    """
+    
+    def __init__(self, 
+                 model_name: str = "gpt-4o",
+                 model_type: str = 'openai',
+                 **kwargs) -> None:
+        """Initialize conversation model"""
+        self.chatbot = ModelFactory(model_type, model_name, **kwargs)
+
+    def initialize_conversation(self,
+                                question: Optional[str],
+                                context: Optional[str] = None,
+                                file_path: Optional[str]=None) -> None:
+        """Initialize conversation state"""
+        if file_path:
+            self.file_path = file_path  
+            self.load_conversation(file_path)
+                
+        else:
+            if not question:
+                raise ValueError("Question is required.")
+            
+            if context:
+                self.context = context
+            else:
+                self.context = "Answer question to the point and don't hallucinate."
+            self.message_list = [
+                SystemMessage(content=context),
+                HumanMessage(content=question)
+            ]
+            
+            res = self._ask_question(self.message_list)
+            print(res)
+            self.message_list.append(AIMessage(content=res))
+
+    def _ask_question(self,messages: List[Union[SystemMessage, HumanMessage, AIMessage]], 
+                     get_content_only: bool = True) -> str:
+        """
+        Ask a question and get response
+        Args:
+            messages: List of messages
+            get_content_only: Whether to return only content
+        Returns:
+            str: Response from the model
+        """
+        res = self.chatbot.invoke_query(messages)
+        if get_content_only:
+            try:
+                return res.content
+            except Exception:
+                return res
+        return res
+
+    def add_message(self, message: str) -> str:
+        """
+        Add a message to the conversation
+        Args:
+            message (str): Message to add
+        Returns:
+            str: Response from the chatbot
+        """
+        self.message_list.append(HumanMessage(content=message))
+        res = self._ask_question(self.message_list)
+        self.message_list.append(AIMessage(content=res))
+        return res
+
+    @property
+    def all_messages(self) -> List[Union[SystemMessage, HumanMessage, AIMessage]]:
+        """Get all messages"""
+        return self.message_list
+
+    @property
+    def last_message(self) -> str:
+        """Get the last message"""
+        return self.message_list[-1].content
+
+    @property
+    def all_messages_content(self) -> List[str]:
+        """Get content of all messages"""
+        return [message.content for message in self.message_list]
+
+    def _is_s3_path(self, path: str) -> bool:
+        """
+        Check if path is an S3 path
+        Args:
+            path (str): Path to check
+        Returns:
+            bool: True if S3 path
+        """
+        return path.startswith("s3://")
+
+    def save_conversation(self, file_path: Optional[str] = None, **kwargs) -> bool:
+        """
+        Save the conversation
+        Args:
+            file_path: Path to save the conversation
+            **kwargs: Additional arguments for S3
+        Returns:
+            bool: Success status
+        """
+        if self._is_s3_path(file_path or self.file_path):
+            print("Saving conversation to S3.")
+            self.save_file_path = file_path
+            return self._save_to_s3(self.file_path,**kwargs)
+        return self._save_to_file(file_path or self.file_path)
+
+    def _save_to_s3(self,**kwargs) -> bool:
+        """Save conversation to S3"""
+        try:
+            client = kwargs.get('client', self.client)
+            bucket = kwargs.get('bucket', self.bucket)
+            client.put_object(
+                Body=str(self.message_list),
+                Bucket=bucket,
+                Key=self.save_file_path
+            )
+            print(f"Conversation saved to s3_path: {self.s3_path}")
+            return True
+        except Exception as e:
+            raise ValueError(f"Error saving conversation to s3: {e}")
+
+    def _save_to_file(self, file_path: str) -> bool:
+        """Save conversation to file"""
+        try:
+            with open(file_path, 'w') as f:
+                for message in self.message_list:
+                    f.write(f"{message.content}\n")
+            print(f"Conversation saved to file: {file_path}")
+            return True
+        except Exception as e:
+            raise ValueError(f"Error saving conversation to file: {e}")
+
+    def load_conversation(self, file_path: Optional[str] = None, **kwargs) -> List[Any]:
+        """
+        Load a conversation
+        Args:
+            file_path: Path to load from
+            **kwargs: Additional arguments for S3
+        Returns:
+            List: Loaded messages
+        """
+        self.message_list = []
+        if self._is_s3_path(file_path or self.file_path):
+            print("Loading conversation from S3.")
+            self.file_path = file_path
+            return self._load_from_s3(**kwargs)
+        return self._load_from_file(file_path or self.file_path)
+
+    def _load_from_s3(self, **kwargs) -> List[Any]:
+        """Load conversation from S3"""
+        try:
+            client = kwargs.get('client', self.client)
+            bucket = kwargs.get('bucket', self.bucket)
+            res = client.get_response(client, bucket, self.s3_path)
+            res_str = eval(res['Body'].read().decode('utf-8'))
+            self.message_list = [SystemMessage(content=res_str)]
+            print(f"Conversation loaded from s3_path: {self.file_path}")
+            return self.message_list
+        except Exception as e:
+            raise ValueError(f"Error loading conversation from s3: {e}")
+
+    def _load_from_file(self, file_path: str) -> List[Any]:
+        """Load conversation from file"""            
+        try:
+            with open(file_path, 'r') as f:
+                lines = f.readlines()
+            for line in lines:
+                self.message_list.append(SystemMessage(content=line))
+            print(f"Conversation loaded from file: {file_path}")
+            return self.message_list
+        except Exception as e:
+            raise ValueError(f"Error loading conversation from file: {e}")
+
+class IPythonStreamHandler(StreamingStdOutCallbackHandler):
+    """Handler for IPython display"""
+    
+    def __init__(self):
+        self.output = ""
+        
+    def on_llm_new_token(self, token: str, **kwargs) -> None:
+        """Handle new token"""
+        self.output += token
+        display(HTML(self.output), clear=True)

mb_rag-1.0.124/mb_rag/chatbot/chains.py

@@ -0,0 +1,206 @@
+## file for chaining functions in chatbot
+
+from typing import Optional, List, Dict, Any, Union
+from dataclasses import dataclass
+from langchain.schema.output_parser import StrOutputParser
+from mb_rag.chatbot.prompts import invoke_prompt
+from langchain.schema.runnable import RunnableLambda, RunnableSequence
+from mb_rag.utils.extra import check_package
+
+__all__ = ['Chain', 'ChainConfig']
+
+def check_langchain_dependencies() -> None:
+    """
+    Check if required LangChain packages are installed
+    Raises:
+        ImportError: If any required package is missing
+    """
+    if not check_package("langchain"):
+        raise ImportError("LangChain package not found. Please install it using: pip install langchain")
+    if not check_package("langchain_core"):
+        raise ImportError("LangChain Core package not found. Please install it using: pip install langchain-core")
+
+# Check dependencies before importing
+check_langchain_dependencies()
+
+@dataclass
+class ChainConfig:
+    """Configuration for chain operations"""
+    prompt: Optional[str] = None
+    prompt_template: Optional[str] = None
+    input_dict: Optional[Dict[str, Any]] = None
+
+class Chain:
+    """
+    Class to chain functions in chatbot with improved OOP design
+    """
+    def __init__(self, model: Any, config: Optional[ChainConfig] = None, **kwargs):
+        """
+        Initialize chain
+        Args:
+            model: The language model to use
+            config: Chain configuration
+            **kwargs: Additional arguments
+        """
+        self.model = model
+        self._output_parser = StrOutputParser()
+        self._initialize_config(config, **kwargs)
+
+    @classmethod
+    def from_template(cls, model: Any, template: str, input_dict: Dict[str, Any], **kwargs) -> 'Chain':
+        """
+        Create chain from template
+        Args:
+            model: The language model
+            template: Prompt template
+            input_dict: Input dictionary for template
+            **kwargs: Additional arguments
+        Returns:
+            Chain: New chain instance
+        """
+        config = ChainConfig(
+            prompt_template=template,
+            input_dict=input_dict
+        )
+        return cls(model, config, **kwargs)
+
+    @classmethod
+    def from_prompt(cls, model: Any, prompt: str, **kwargs) -> 'Chain':
+        """
+        Create chain from direct prompt
+        Args:
+            model: The language model
+            prompt: Direct prompt
+            **kwargs: Additional arguments
+        Returns:
+            Chain: New chain instance
+        """
+        config = ChainConfig(prompt=prompt)
+        return cls(model, config, **kwargs)
+
+    def _initialize_config(self, config: Optional[ChainConfig], **kwargs) -> None:
+        """Initialize chain configuration"""
+        if config:
+            self.input_dict = config.input_dict
+            if config.prompt_template:
+                self.prompt = invoke_prompt(config.prompt_template, self.input_dict)
+            else:
+                self.prompt = config.prompt
+        else:
+            self.input_dict = kwargs.get('input_dict')
+            if prompt_template := kwargs.get('prompt_template'):
+                self.prompt = invoke_prompt(prompt_template, self.input_dict)
+            else:
+                self.prompt = kwargs.get('prompt')
+
+    @property
+    def output_parser(self) -> StrOutputParser:
+        """Get the output parser"""
+        return self._output_parser
+
+    @staticmethod
+    def _validate_chain_components(prompt: Any, middle_chain: Optional[List] = None) -> None:
+        """
+        Validate chain components
+        Args:
+            prompt: The prompt to validate
+            middle_chain: Optional middle chain to validate
+        Raises:
+            ValueError: If validation fails
+        """
+        if prompt is None:
+            raise ValueError("Prompt is not provided")
+        if middle_chain is not None and not isinstance(middle_chain, list):
+            raise ValueError("middle_chain should be a list")
+
+    def invoke(self) -> Any:
+        """
+        Invoke the chain
+        Returns:
+            Any: Output from the chain
+        Raises:
+            Exception: If prompt is not provided
+        """
+        self._validate_chain_components(self.prompt)
+        chain_output = self.prompt | self.model | self.output_parser
+        return chain_output
+
+    def chain_sequence_invoke(self, 
+                            middle_chain: Optional[List] = None,
+                            final_chain: Optional[RunnableLambda] = None) -> Any:
+        """
+        Chain invoke the sequence
+        Args:
+            middle_chain: List of functions/Prompts/RunnableLambda to chain
+            final_chain: Final chain to run
+        Returns:
+            Any: Output from the chain
+        """
+        self._validate_chain_components(self.prompt, middle_chain)
+        
+        final = final_chain if final_chain is not None else self.output_parser
+        
+        if middle_chain:
+            func_chain = RunnableSequence(self.prompt, middle_chain, final)
+            return func_chain.invoke()
+        return None
+
+    def chain_parallel_invoke(self, parallel_chain: List) -> Any:
+        """
+        Chain invoke in parallel
+        Args:
+            parallel_chain: List of chains to run in parallel
+        Returns:
+            Any: Output from the parallel chains
+        Raises:
+            ImportError: If LangChain is not installed
+        """
+        if not check_package("langchain"):
+            raise ImportError("LangChain package not found. Please install it using: pip install langchain")
+        return parallel_chain.invoke()
+
+    def chain_branch_invoke(self, branch_chain: Dict) -> Any:
+        """
+        Chain invoke with branching
+        Args:
+            branch_chain: Dictionary of branch chains
+        Returns:
+            Any: Output from the branch chain
+        Raises:
+            ImportError: If LangChain is not installed
+        """
+        if not check_package("langchain"):
+            raise ImportError("LangChain package not found. Please install it using: pip install langchain")
+        return branch_chain.invoke()
+
+    @staticmethod
+    def create_parallel_chain(prompt_template: str, model: Any, branches: Dict[str, Any]) -> Any:
+        """
+        Create a parallel chain
+        Args:
+            prompt_template: Template for the prompt
+            model: The language model
+            branches: Dictionary of branch configurations
+        Returns:
+            Any: Configured parallel chain
+        """
+        from langchain.schema.runnable import RunnableParallel
+        return (
+            prompt_template
+            | model
+            | StrOutputParser()
+            | RunnableParallel(branches=branches)
+        )
+
+    @staticmethod
+    def create_branch_chain(conditions: List[tuple], default_chain: Any) -> Any:
+        """
+        Create a branch chain
+        Args:
+            conditions: List of condition-chain tuples
+            default_chain: Default chain to use
+        Returns:
+            Any: Configured branch chain
+        """
+        from langchain.schema.runnable import RunnableBranch
+        return RunnableBranch(*conditions, default_chain)